1use crate::Cents;
2use crate::db::billing_subscription::SubscriptionKind;
3use crate::db::{billing_subscription, user};
4use crate::llm::{
5 AGENT_EXTENDED_TRIAL_FEATURE_FLAG, DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT,
6};
7use crate::{Config, db::billing_preference};
8use anyhow::{Result, anyhow};
9use chrono::{Datelike, NaiveDate, NaiveDateTime, Utc};
10use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
11use serde::{Deserialize, Serialize};
12use std::time::Duration;
13use thiserror::Error;
14use util::maybe;
15use uuid::Uuid;
16use zed_llm_client::Plan;
17
18#[derive(Clone, Debug, Default, Serialize, Deserialize)]
19#[serde(rename_all = "camelCase")]
20pub struct LlmTokenClaims {
21 pub iat: u64,
22 pub exp: u64,
23 pub jti: String,
24 pub user_id: u64,
25 pub system_id: Option<String>,
26 pub metrics_id: Uuid,
27 pub github_user_login: String,
28 pub account_created_at: NaiveDateTime,
29 pub is_staff: bool,
30 pub has_llm_closed_beta_feature_flag: bool,
31 pub bypass_account_age_check: bool,
32 pub has_llm_subscription: bool,
33 #[serde(default)]
34 pub use_llm_request_queue: bool,
35 pub max_monthly_spend_in_cents: u32,
36 pub custom_llm_monthly_allowance_in_cents: Option<u32>,
37 #[serde(default)]
38 pub use_new_billing: bool,
39 pub plan: Plan,
40 #[serde(default)]
41 pub has_extended_trial: bool,
42 #[serde(default)]
43 pub subscription_period: Option<(NaiveDateTime, NaiveDateTime)>,
44 #[serde(default)]
45 pub enable_model_request_overages: bool,
46 #[serde(default)]
47 pub model_request_overages_spend_limit_in_cents: u32,
48 #[serde(default)]
49 pub can_use_web_search_tool: bool,
50}
51
52const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
53
54impl LlmTokenClaims {
55 pub fn create(
56 user: &user::Model,
57 is_staff: bool,
58 billing_preferences: Option<billing_preference::Model>,
59 feature_flags: &Vec<String>,
60 has_legacy_llm_subscription: bool,
61 subscription: Option<billing_subscription::Model>,
62 system_id: Option<String>,
63 config: &Config,
64 ) -> Result<String> {
65 let secret = config
66 .llm_api_secret
67 .as_ref()
68 .ok_or_else(|| anyhow!("no LLM API secret"))?;
69
70 let now = Utc::now();
71 let claims = Self {
72 iat: now.timestamp() as u64,
73 exp: (now + LLM_TOKEN_LIFETIME).timestamp() as u64,
74 jti: uuid::Uuid::new_v4().to_string(),
75 user_id: user.id.to_proto(),
76 system_id,
77 metrics_id: user.metrics_id,
78 github_user_login: user.github_login.clone(),
79 account_created_at: user.account_created_at(),
80 is_staff,
81 has_llm_closed_beta_feature_flag: feature_flags
82 .iter()
83 .any(|flag| flag == "llm-closed-beta"),
84 bypass_account_age_check: feature_flags
85 .iter()
86 .any(|flag| flag == "bypass-account-age-check"),
87 can_use_web_search_tool: feature_flags.iter().any(|flag| flag == "assistant2"),
88 has_llm_subscription: has_legacy_llm_subscription,
89 max_monthly_spend_in_cents: billing_preferences
90 .as_ref()
91 .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
92 preferences.max_monthly_llm_usage_spending_in_cents as u32
93 }),
94 custom_llm_monthly_allowance_in_cents: user
95 .custom_llm_monthly_allowance_in_cents
96 .map(|allowance| allowance as u32),
97 use_new_billing: feature_flags.iter().any(|flag| flag == "new-billing"),
98 use_llm_request_queue: feature_flags.iter().any(|flag| flag == "llm-request-queue"),
99 plan: if is_staff {
100 Plan::ZedPro
101 } else {
102 subscription
103 .as_ref()
104 .and_then(|subscription| subscription.kind)
105 .map_or(Plan::Free, |kind| match kind {
106 SubscriptionKind::ZedFree => Plan::Free,
107 SubscriptionKind::ZedPro => Plan::ZedPro,
108 SubscriptionKind::ZedProTrial => Plan::ZedProTrial,
109 })
110 },
111 has_extended_trial: feature_flags
112 .iter()
113 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG),
114 subscription_period: if is_staff {
115 maybe!({
116 let now = Utc::now();
117 let year = now.year();
118 let month = now.month();
119
120 let first_day_of_this_month =
121 NaiveDate::from_ymd_opt(year, month, 1)?.and_hms_opt(0, 0, 0)?;
122
123 let next_month = if month == 12 { 1 } else { month + 1 };
124 let next_month_year = if month == 12 { year + 1 } else { year };
125 let first_day_of_next_month =
126 NaiveDate::from_ymd_opt(next_month_year, next_month, 1)?
127 .and_hms_opt(23, 59, 59)?;
128
129 let last_day_of_this_month = first_day_of_next_month - chrono::Days::new(1);
130
131 Some((first_day_of_this_month, last_day_of_this_month))
132 })
133 } else {
134 maybe!({
135 let subscription = subscription?;
136 let period_start_at = subscription.current_period_start_at()?;
137 let period_end_at = subscription.current_period_end_at()?;
138
139 Some((period_start_at.naive_utc(), period_end_at.naive_utc()))
140 })
141 },
142 enable_model_request_overages: billing_preferences
143 .as_ref()
144 .map_or(false, |preferences| {
145 preferences.model_request_overages_enabled
146 }),
147 model_request_overages_spend_limit_in_cents: billing_preferences
148 .as_ref()
149 .map_or(0, |preferences| {
150 preferences.model_request_overages_spend_limit_in_cents as u32
151 }),
152 };
153
154 Ok(jsonwebtoken::encode(
155 &Header::default(),
156 &claims,
157 &EncodingKey::from_secret(secret.as_ref()),
158 )?)
159 }
160
161 pub fn validate(token: &str, config: &Config) -> Result<LlmTokenClaims, ValidateLlmTokenError> {
162 let secret = config
163 .llm_api_secret
164 .as_ref()
165 .ok_or_else(|| anyhow!("no LLM API secret"))?;
166
167 match jsonwebtoken::decode::<Self>(
168 token,
169 &DecodingKey::from_secret(secret.as_ref()),
170 &Validation::default(),
171 ) {
172 Ok(token) => Ok(token.claims),
173 Err(e) => {
174 if e.kind() == &jsonwebtoken::errors::ErrorKind::ExpiredSignature {
175 Err(ValidateLlmTokenError::Expired)
176 } else {
177 Err(ValidateLlmTokenError::JwtError(e))
178 }
179 }
180 }
181 }
182
183 pub fn free_tier_monthly_spending_limit(&self) -> Cents {
184 self.custom_llm_monthly_allowance_in_cents
185 .map(Cents)
186 .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT)
187 }
188}
189
190#[derive(Error, Debug)]
191pub enum ValidateLlmTokenError {
192 #[error("access token is expired")]
193 Expired,
194 #[error("access token validation error: {0}")]
195 JwtError(#[from] jsonwebtoken::errors::Error),
196 #[error("{0}")]
197 Other(#[from] anyhow::Error),
198}