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