From 9bfd8a59cd4c9ffdc0f09a03be324eb3e1b32d7e 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,