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 1f64551b28af7fbe815ec1d64221225bf323e83d..29a2acc816dcd11e50c87850d6f1d7523dcd7dc4 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 2c8bba8b6b88004bb148c9018370f661b6200142..0e6c386a8698828caaa5ded5a4663690b86c087d 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,