From 9bbc2e0fb26c1ec144e05de132bedec53f9022c5 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Tue, 22 Apr 2025 15:44:16 -0400
Subject: [PATCH] collab: Set `plan` in LLM token based on subscription
 (#29231)

This PR updates the `plan` field in the LLM token to be based on the
subscription.

We weren't using this field anywhere outside of the new billing code, so
it is safe to change its meaning.

Release Notes:

- N/A
---
 crates/collab/src/llm/token.rs | 15 +++++++++------
 crates/collab/src/rpc.rs       |  1 -
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs
index 1f64551b28..29a2acc816 100644
--- a/crates/collab/src/llm/token.rs
+++ b/crates/collab/src/llm/token.rs
@@ -1,4 +1,5 @@
 use crate::Cents;
+use crate::db::billing_subscription::SubscriptionKind;
 use crate::db::{billing_subscription, user};
 use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
 use crate::{Config, db::billing_preference};
@@ -43,7 +44,6 @@ impl LlmTokenClaims {
         billing_preferences: Option<billing_preference::Model>,
         feature_flags: &Vec<String>,
         has_legacy_llm_subscription: bool,
-        plan: rpc::proto::Plan,
         subscription: Option<billing_subscription::Model>,
         system_id: Option<String>,
         config: &Config,
@@ -78,11 +78,14 @@ impl LlmTokenClaims {
             custom_llm_monthly_allowance_in_cents: user
                 .custom_llm_monthly_allowance_in_cents
                 .map(|allowance| allowance as u32),
-            plan: match plan {
-                rpc::proto::Plan::Free => Plan::Free,
-                rpc::proto::Plan::ZedPro => Plan::ZedPro,
-                rpc::proto::Plan::ZedProTrial => Plan::ZedProTrial,
-            },
+            plan: subscription
+                .as_ref()
+                .and_then(|subscription| subscription.kind)
+                .map_or(Plan::Free, |kind| match kind {
+                    SubscriptionKind::ZedFree => Plan::Free,
+                    SubscriptionKind::ZedPro => Plan::ZedPro,
+                    SubscriptionKind::ZedProTrial => Plan::ZedProTrial,
+                }),
             subscription_period: maybe!({
                 let subscription = subscription?;
                 let period_start_at = subscription.current_period_start_at()?;
diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs
index 2c8bba8b6b..0e6c386a86 100644
--- a/crates/collab/src/rpc.rs
+++ b/crates/collab/src/rpc.rs
@@ -4147,7 +4147,6 @@ async fn get_llm_api_token(
         billing_preferences,
         &flags,
         has_legacy_llm_subscription,
-        session.current_plan(&db).await?,
         billing_subscription,
         session.system_id.clone(),
         &session.app_state.config,