Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

collab: Require subscription_period in LLM token claims (#30392)

Browse source 
This PR makes the `subscription_period` field in the LLM token claims
required.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-05-09 11:20:14 -04:00 • committed by GitHub
parent 08f516ce9a
commit d270f6b953
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 20 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

44
crates/collab/src/llm/token.rs
View file

@ -25,18 +25,12 @@ pub struct LlmTokenClaims {
pub is_staff: bool, pub is_staff: bool,
pub has_llm_closed_beta_feature_flag: bool, pub has_llm_closed_beta_feature_flag: bool,
pub bypass_account_age_check: bool, pub bypass_account_age_check: bool,
#[serde(default)]
pub use_llm_request_queue: bool, pub use_llm_request_queue: bool,
pub plan: Plan, pub plan: Plan,
#[serde(default)]
pub has_extended_trial: bool, pub has_extended_trial: bool,
#[serde(default)] pub subscription_period: (NaiveDateTime, NaiveDateTime),
pub subscription_period: Option<(NaiveDateTime, NaiveDateTime)>,
#[serde(default)]
pub enable_model_request_overages: bool, pub enable_model_request_overages: bool,
#[serde(default)]
pub model_request_overages_spend_limit_in_cents: u32, pub model_request_overages_spend_limit_in_cents: u32,
#[serde(default)]
pub can_use_web_search_tool: bool, pub can_use_web_search_tool: bool,
} }
@ -57,6 +51,23 @@ impl LlmTokenClaims {
.as_ref() .as_ref()
.ok_or_else(|| anyhow!("no LLM API secret"))?; .ok_or_else(|| anyhow!("no LLM API secret"))?;
let plan = if is_staff {
Plan::ZedPro
} else {
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,
})
};
let subscription_period =
billing_subscription::Model::current_period(subscription, is_staff)
.map(|(start, end)| (start.naive_utc(), end.naive_utc()))
.ok_or_else(|| anyhow!("missing subscription period"))?;
let now = Utc::now(); let now = Utc::now();
let claims = Self { let claims = Self {
iat: now.timestamp() as u64, iat: now.timestamp() as u64,
@ -76,26 +87,11 @@ impl LlmTokenClaims {
.any(|flag| flag == "bypass-account-age-check"), .any(|flag| flag == "bypass-account-age-check"),
can_use_web_search_tool: true, can_use_web_search_tool: true,
use_llm_request_queue: feature_flags.iter().any(|flag| flag == "llm-request-queue"), use_llm_request_queue: feature_flags.iter().any(|flag| flag == "llm-request-queue"),
plan: if is_staff { plan,
Plan::ZedPro
} else {
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,
})
},
has_extended_trial: feature_flags has_extended_trial: feature_flags
.iter() .iter()
.any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG), .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG),
subscription_period: billing_subscription::Model::current_period( subscription_period,
subscription,
is_staff,
)
.map(|(start, end)| (start.naive_utc(), end.naive_utc())),
enable_model_request_overages: billing_preferences enable_model_request_overages: billing_preferences
.as_ref() .as_ref()
.map_or(false, |preferences| { .map_or(false, |preferences| {