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