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
This commit is contained in:
Marshall Bowers authored and GitHub committed 2025-04-22 19:44:16 +00:00
1 parent b588186a92
commit 9bfd8a59cd
2 files changed
+9 -7

No files matched your search

+9 -6
View File
@@ -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()?;
-1
View File
@@ -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,