Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

collab: Remove endpoints for issuing notifications from Cloud (#36249)

Browse source 
This PR removes the `POST /users/:id/refresh_llm_tokens` and `POST
/users/:id/update_plan` endpoints from Collab.

These endpoints were added to be called by Cloud in order to push down
notifications over the Collab RPC connection.

Cloud now sends down notifications to clients directly, so we no longer
need these endpoints.

All calls to these endpoints have already been removed in production.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-08-15 10:46:06 -04:00 • committed by GitHub
parent 1e41d86b31
commit 485802b9e5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 0 additions and 139 deletions
Download patch file Download diff file Expand all files Collapse all files

92
crates/collab/src/api.rs
View file

@ -11,9 +11,7 @@ use crate::{
db::{User, UserId},
rpc,
};
use ::rpc::proto;
use anyhow::Context as _;
use axum::extract;
use axum::{
Extension, Json, Router,
body::Body,
@ -25,7 +23,6 @@ use axum::{
routing::{get, post},
};
use axum_extra::response::ErasedJson;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::sync::{Arc, OnceLock};
use tower::ServiceBuilder;
@ -102,8 +99,6 @@ pub fn routes(rpc_server: Arc<rpc::Server>) -> Router<(), Body> {
Router::new()
.route("/users/look_up", get(look_up_user))
.route("/users/:id/access_tokens", post(create_access_token))
.route("/users/:id/refresh_llm_tokens", post(refresh_llm_tokens))
.route("/users/:id/update_plan", post(update_plan))
.route("/rpc_server_snapshot", get(get_rpc_server_snapshot))
.merge(contributors::router())
.layer(
@ -295,90 +290,3 @@ async fn create_access_token(
encrypted_access_token,
}))
}
#[derive(Serialize)]
struct RefreshLlmTokensResponse {}
async fn refresh_llm_tokens(
Path(user_id): Path<UserId>,
Extension(rpc_server): Extension<Arc<rpc::Server>>,
) -> Result<Json<RefreshLlmTokensResponse>> {
rpc_server.refresh_llm_tokens_for_user(user_id).await;
Ok(Json(RefreshLlmTokensResponse {}))
}
#[derive(Debug, Serialize, Deserialize)]
struct UpdatePlanBody {
pub plan: cloud_llm_client::Plan,
pub subscription_period: SubscriptionPeriod,
pub usage: cloud_llm_client::CurrentUsage,
pub trial_started_at: Option<DateTime<Utc>>,
pub is_usage_based_billing_enabled: bool,
pub is_account_too_young: bool,
pub has_overdue_invoices: bool,
}
#[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)]
struct SubscriptionPeriod {
pub started_at: DateTime<Utc>,
pub ended_at: DateTime<Utc>,
}
#[derive(Serialize)]
struct UpdatePlanResponse {}
async fn update_plan(
Path(user_id): Path<UserId>,
Extension(rpc_server): Extension<Arc<rpc::Server>>,
extract::Json(body): extract::Json<UpdatePlanBody>,
) -> Result<Json<UpdatePlanResponse>> {
let plan = match body.plan {
cloud_llm_client::Plan::ZedFree => proto::Plan::Free,
cloud_llm_client::Plan::ZedPro => proto::Plan::ZedPro,
cloud_llm_client::Plan::ZedProTrial => proto::Plan::ZedProTrial,
};
let update_user_plan = proto::UpdateUserPlan {
plan: plan.into(),
trial_started_at: body
.trial_started_at
.map(|trial_started_at| trial_started_at.timestamp() as u64),
is_usage_based_billing_enabled: Some(body.is_usage_based_billing_enabled),
usage: Some(proto::SubscriptionUsage {
model_requests_usage_amount: body.usage.model_requests.used,
model_requests_usage_limit: Some(usage_limit_to_proto(body.usage.model_requests.limit)),
edit_predictions_usage_amount: body.usage.edit_predictions.used,
edit_predictions_usage_limit: Some(usage_limit_to_proto(
body.usage.edit_predictions.limit,
)),
}),
subscription_period: Some(proto::SubscriptionPeriod {
started_at: body.subscription_period.started_at.timestamp() as u64,
ended_at: body.subscription_period.ended_at.timestamp() as u64,
}),
account_too_young: Some(body.is_account_too_young),
has_overdue_invoices: Some(body.has_overdue_invoices),
};
rpc_server
.update_plan_for_user(user_id, update_user_plan)
.await?;
Ok(Json(UpdatePlanResponse {}))
}
fn usage_limit_to_proto(limit: cloud_llm_client::UsageLimit) -> proto::UsageLimit {
proto::UsageLimit {
variant: Some(match limit {
cloud_llm_client::UsageLimit::Limited(limit) => {
proto::usage_limit::Variant::Limited(proto::usage_limit::Limited {
limit: limit as u32,
})
}
cloud_llm_client::UsageLimit::Unlimited => {
proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {})
}
}),
}
}

47
crates/collab/src/rpc.rs
View file

@ -1081,53 +1081,6 @@ impl Server {
Ok(())
}
pub async fn update_plan_for_user(
self: &Arc<Self>,
user_id: UserId,
update_user_plan: proto::UpdateUserPlan,
) -> Result<()> {
let pool = self.connection_pool.lock();
for connection_id in pool.user_connection_ids(user_id) {
self.peer
.send(connection_id, update_user_plan.clone())
.trace_err();
}
Ok(())
}
/// This is the legacy way of updating the user's plan, where we fetch the data to construct the `UpdateUserPlan`
/// message on the Collab server.
///
/// The new way is to receive the data from Cloud via the `POST /users/:id/update_plan` endpoint.
pub async fn update_plan_for_user_legacy(self: &Arc<Self>, user_id: UserId) -> Result<()> {
let user = self
.app_state
.db
.get_user_by_id(user_id)
.await?
.context("user not found")?;
let update_user_plan = make_update_user_plan_message(
&user,
user.admin,
&self.app_state.db,
self.app_state.llm_db.clone(),
)
.await?;
self.update_plan_for_user(user_id, update_user_plan).await
}
pub async fn refresh_llm_tokens_for_user(self: &Arc<Self>, user_id: UserId) {
let pool = self.connection_pool.lock();
for connection_id in pool.user_connection_ids(user_id) {
self.peer
.send(connection_id, proto::RefreshLlmToken {})
.trace_err();
}
}
pub async fn snapshot(self: &Arc<Self>) -> ServerSnapshot<'_> {
ServerSnapshot {
connection_pool: ConnectionPoolGuard {