@@ -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 {})
- }
- }),
- }
-}
@@ -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 {