diff --git a/crates/collab/src/api.rs b/crates/collab/src/api.rs index 6cf3f68f54eda75ac19950c53cf535ff30a107a9..078a4469ae713de4c79929db09ed3522a52790a3 100644 --- a/crates/collab/src/api.rs +++ b/crates/collab/src/api.rs @@ -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) -> 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, - Extension(rpc_server): Extension>, -) -> Result> { - 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>, - 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, - pub ended_at: DateTime, -} - -#[derive(Serialize)] -struct UpdatePlanResponse {} - -async fn update_plan( - Path(user_id): Path, - Extension(rpc_server): Extension>, - extract::Json(body): extract::Json, -) -> Result> { - 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 {}) - } - }), - } -} diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs index 18eb1457dc336c5e2bf32a3d8430514b29bb6966..584970a4c678c3f0f3db2a2516010dec7a30c1ad 100644 --- a/crates/collab/src/rpc.rs +++ b/crates/collab/src/rpc.rs @@ -1081,53 +1081,6 @@ impl Server { Ok(()) } - pub async fn update_plan_for_user( - self: &Arc, - 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, 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, 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) -> ServerSnapshot<'_> { ServerSnapshot { connection_pool: ConnectionPoolGuard {