collab: Remove legacy claims from LLM token (#30294)

Marshall Bowers created 7 months ago

This PR removes some legacy claims related to the old billing from the
LLM token.

We already stopped reading this in the LLM Worker.

Also removed an outdated feature flag check that restricted access to
obtaining an LLM token.

Release Notes:

- N/A

Change summary

crates/collab/src/llm/token.rs | 29 ++---------------------------
crates/collab/src/rpc.rs       | 20 --------------------
2 files changed, 2 insertions(+), 47 deletions(-)

Detailed changes

crates/collab/src/llm/token.rs

@@ -1,9 +1,6 @@
-use crate::Cents;
 use crate::db::billing_subscription::SubscriptionKind;
 use crate::db::{billing_subscription, user};
-use crate::llm::{
-    AGENT_EXTENDED_TRIAL_FEATURE_FLAG, DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT,
-};
+use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
 use crate::{Config, db::billing_preference};
 use anyhow::{Result, anyhow};
 use chrono::{NaiveDateTime, Utc};
@@ -28,13 +25,8 @@ pub struct LlmTokenClaims {
     pub is_staff: bool,
     pub has_llm_closed_beta_feature_flag: bool,
     pub bypass_account_age_check: bool,
-    pub has_llm_subscription: bool,
     #[serde(default)]
     pub use_llm_request_queue: bool,
-    pub max_monthly_spend_in_cents: u32,
-    pub custom_llm_monthly_allowance_in_cents: Option<u32>,
-    #[serde(default)]
-    pub use_new_billing: bool,
     pub plan: Plan,
     #[serde(default)]
     pub has_extended_trial: bool,
@@ -56,7 +48,6 @@ impl LlmTokenClaims {
         is_staff: bool,
         billing_preferences: Option<billing_preference::Model>,
         feature_flags: &Vec<String>,
-        has_legacy_llm_subscription: bool,
         subscription: Option<billing_subscription::Model>,
         system_id: Option<String>,
         config: &Config,
@@ -83,17 +74,7 @@ impl LlmTokenClaims {
             bypass_account_age_check: feature_flags
                 .iter()
                 .any(|flag| flag == "bypass-account-age-check"),
-            can_use_web_search_tool: feature_flags.iter().any(|flag| flag == "assistant2"),
-            has_llm_subscription: has_legacy_llm_subscription,
-            max_monthly_spend_in_cents: billing_preferences
-                .as_ref()
-                .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
-                    preferences.max_monthly_llm_usage_spending_in_cents as u32
-                }),
-            custom_llm_monthly_allowance_in_cents: user
-                .custom_llm_monthly_allowance_in_cents
-                .map(|allowance| allowance as u32),
-            use_new_billing: feature_flags.iter().any(|flag| flag == "new-billing"),
+            can_use_web_search_tool: true,
             use_llm_request_queue: feature_flags.iter().any(|flag| flag == "llm-request-queue"),
             plan: if is_staff {
                 Plan::ZedPro
@@ -155,12 +136,6 @@ impl LlmTokenClaims {
             }
         }
     }
-
-    pub fn free_tier_monthly_spending_limit(&self) -> Cents {
-        self.custom_llm_monthly_allowance_in_cents
-            .map(Cents)
-            .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT)
-    }
 }
 
 #[derive(Error, Debug)]

crates/collab/src/rpc.rs

@@ -166,19 +166,6 @@ impl Session {
         }
     }
 
-    pub async fn has_llm_subscription(
-        &self,
-        db: &MutexGuard<'_, DbHandle>,
-    ) -> anyhow::Result<bool> {
-        if self.is_staff() {
-            return Ok(true);
-        }
-
-        let user_id = self.user_id();
-
-        Ok(db.has_active_billing_subscription(user_id).await?)
-    }
-
     pub async fn current_plan(&self, db: &MutexGuard<'_, DbHandle>) -> anyhow::Result<proto::Plan> {
         if self.is_staff() {
             return Ok(proto::Plan::ZedPro);
@@ -4000,11 +3987,6 @@ async fn get_llm_api_token(
     let db = session.db().await;
 
     let flags = db.get_user_flags(session.user_id()).await?;
-    let has_language_models_feature_flag = flags.iter().any(|flag| flag == "language-models");
-
-    if !session.is_staff() && !has_language_models_feature_flag {
-        Err(anyhow!("permission denied"))?
-    }
 
     let user_id = session.user_id();
     let user = db
@@ -4016,7 +3998,6 @@ async fn get_llm_api_token(
         Err(anyhow!("terms of service not accepted"))?
     }
 
-    let has_legacy_llm_subscription = session.has_llm_subscription(&db).await?;
     let billing_subscription = db.get_active_billing_subscription(user.id).await?;
     let billing_preferences = db.get_billing_preferences(user.id).await?;
 
@@ -4025,7 +4006,6 @@ async fn get_llm_api_token(
         session.is_staff(),
         billing_preferences,
         &flags,
-        has_legacy_llm_subscription,
         billing_subscription,
         session.system_id.clone(),
         &session.app_state.config,