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