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