Remove RPC messages pertaining to the LLM token (#36252)

This PR removes the RPC messages pertaining to the LLM token. We now retrieve the LLM token from Cloud. Release Notes: - N/A
2025-08-15 13:26:21 -04:00
parent 257e0991d8
commit 75b832029a
8 changed files with 4 additions and 263 deletions
@@ -1,146 +0,0 @@
-use crate::db::billing_subscription::SubscriptionKind;
-use crate::db::{billing_customer, billing_subscription, user};
-use crate::llm::{AGENT_EXTENDED_TRIAL_FEATURE_FLAG, BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG};
-use crate::{Config, db::billing_preference};
-use anyhow::{Context as _, Result};
-use chrono::{NaiveDateTime, Utc};
-use cloud_llm_client::Plan;
-use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
-use serde::{Deserialize, Serialize};
-use std::time::Duration;
-use thiserror::Error;
-use uuid::Uuid;
-
-#[derive(Clone, Debug, Default, Serialize, Deserialize)]
-#[serde(rename_all = "camelCase")]
-pub struct LlmTokenClaims {
-    pub iat: u64,
-    pub exp: u64,
-    pub jti: String,
-    pub user_id: u64,
-    pub system_id: Option<String>,
-    pub metrics_id: Uuid,
-    pub github_user_login: String,
-    pub account_created_at: NaiveDateTime,
-    pub is_staff: bool,
-    pub has_llm_closed_beta_feature_flag: bool,
-    pub bypass_account_age_check: bool,
-    pub use_llm_request_queue: bool,
-    pub plan: Plan,
-    pub has_extended_trial: bool,
-    pub subscription_period: (NaiveDateTime, NaiveDateTime),
-    pub enable_model_request_overages: bool,
-    pub model_request_overages_spend_limit_in_cents: u32,
-    pub can_use_web_search_tool: bool,
-    #[serde(default)]
-    pub has_overdue_invoices: bool,
-}
-
-const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
-
-impl LlmTokenClaims {
-    pub fn create(
-        user: &user::Model,
-        is_staff: bool,
-        billing_customer: billing_customer::Model,
-        billing_preferences: Option<billing_preference::Model>,
-        feature_flags: &Vec<String>,
-        subscription: billing_subscription::Model,
-        system_id: Option<String>,
-        config: &Config,
-    ) -> Result<String> {
-        let secret = config
-            .llm_api_secret
-            .as_ref()
-            .context("no LLM API secret")?;
-
-        let plan = if is_staff {
-            Plan::ZedPro
-        } else {
-            subscription.kind.map_or(Plan::ZedFree, |kind| match kind {
-                SubscriptionKind::ZedFree => Plan::ZedFree,
-                SubscriptionKind::ZedPro => Plan::ZedPro,
-                SubscriptionKind::ZedProTrial => Plan::ZedProTrial,
-            })
-        };
-        let subscription_period =
-            billing_subscription::Model::current_period(Some(subscription), is_staff)
-                .map(|(start, end)| (start.naive_utc(), end.naive_utc()))
-                .context("A plan is required to use Zed's hosted models or edit predictions. Visit https://zed.dev/account to get started.")?;
-
-        let now = Utc::now();
-        let claims = Self {
-            iat: now.timestamp() as u64,
-            exp: (now + LLM_TOKEN_LIFETIME).timestamp() as u64,
-            jti: uuid::Uuid::new_v4().to_string(),
-            user_id: user.id.to_proto(),
-            system_id,
-            metrics_id: user.metrics_id,
-            github_user_login: user.github_login.clone(),
-            account_created_at: user.account_created_at(),
-            is_staff,
-            has_llm_closed_beta_feature_flag: feature_flags
-                .iter()
-                .any(|flag| flag == "llm-closed-beta"),
-            bypass_account_age_check: feature_flags
-                .iter()
-                .any(|flag| flag == BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG),
-            can_use_web_search_tool: true,
-            use_llm_request_queue: feature_flags.iter().any(|flag| flag == "llm-request-queue"),
-            plan,
-            has_extended_trial: feature_flags
-                .iter()
-                .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG),
-            subscription_period,
-            enable_model_request_overages: billing_preferences
-                .as_ref()
-                .map_or(false, |preferences| {
-                    preferences.model_request_overages_enabled
-                }),
-            model_request_overages_spend_limit_in_cents: billing_preferences
-                .as_ref()
-                .map_or(0, |preferences| {
-                    preferences.model_request_overages_spend_limit_in_cents as u32
-                }),
-            has_overdue_invoices: billing_customer.has_overdue_invoices,
-        };
-
-        Ok(jsonwebtoken::encode(
-            &Header::default(),
-            &claims,
-            &EncodingKey::from_secret(secret.as_ref()),
-        )?)
-    }
-
-    pub fn validate(token: &str, config: &Config) -> Result<LlmTokenClaims, ValidateLlmTokenError> {
-        let secret = config
-            .llm_api_secret
-            .as_ref()
-            .context("no LLM API secret")?;
-
-        match jsonwebtoken::decode::<Self>(
-            token,
-            &DecodingKey::from_secret(secret.as_ref()),
-            &Validation::default(),
-        ) {
-            Ok(token) => Ok(token.claims),
-            Err(e) => {
-                if e.kind() == &jsonwebtoken::errors::ErrorKind::ExpiredSignature {
-                    Err(ValidateLlmTokenError::Expired)
-                } else {
-                    Err(ValidateLlmTokenError::JwtError(e))
-                }
-            }
-        }
-    }
-}
-
-#[derive(Error, Debug)]
-pub enum ValidateLlmTokenError {
-    #[error("access token is expired")]
-    Expired,
-    #[error("access token validation error: {0}")]
-    JwtError(#[from] jsonwebtoken::errors::Error),
-    #[error("{0}")]
-    Other(#[from] anyhow::Error),
-}