From 6827bf114a4fdab89a2d64e89e813a6136213fe7 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Thu, 8 May 2025 14:32:54 -0400
Subject: [PATCH] collab: Remove legacy claims from LLM token (#30294)

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
---
 crates/collab/src/llm/token.rs | 29 ++---------------------------
 crates/collab/src/rpc.rs       | 20 --------------------
 2 files changed, 2 insertions(+), 47 deletions(-)

diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs
index 52c2acc5842d8a3c3f75003e311b4710568cb745..af3b70ed7500243420fcf0d5d6f5fc4d897147b0 100644
--- a/crates/collab/src/llm/token.rs
+++ b/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)]
diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs
index 22972e62e108bf433fe785b0981ed0430566bb24..e11abf45875eea4241e5432b1c5960cf90eeedfc 100644
--- a/crates/collab/src/rpc.rs
+++ b/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,