1use crate::db::billing_subscription::SubscriptionKind;
2use crate::db::{billing_subscription, user};
3use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
4use crate::{Config, db::billing_preference};
5use anyhow::{Context as _, Result};
6use chrono::{NaiveDateTime, Utc};
7use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
8use serde::{Deserialize, Serialize};
9use std::time::Duration;
10use thiserror::Error;
11use uuid::Uuid;
12use zed_llm_client::Plan;
13
14#[derive(Clone, Debug, Default, Serialize, Deserialize)]
15#[serde(rename_all = "camelCase")]
16pub struct LlmTokenClaims {
17 pub iat: u64,
18 pub exp: u64,
19 pub jti: String,
20 pub user_id: u64,
21 pub system_id: Option<String>,
22 pub metrics_id: Uuid,
23 pub github_user_login: String,
24 pub account_created_at: NaiveDateTime,
25 pub is_staff: bool,
26 pub has_llm_closed_beta_feature_flag: bool,
27 pub bypass_account_age_check: bool,
28 pub use_llm_request_queue: bool,
29 pub plan: Plan,
30 pub has_extended_trial: bool,
31 pub subscription_period: (NaiveDateTime, NaiveDateTime),
32 pub enable_model_request_overages: bool,
33 pub model_request_overages_spend_limit_in_cents: u32,
34 pub can_use_web_search_tool: bool,
35}
36
37const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
38
39impl LlmTokenClaims {
40 pub fn create(
41 user: &user::Model,
42 is_staff: bool,
43 billing_preferences: Option<billing_preference::Model>,
44 feature_flags: &Vec<String>,
45 subscription: billing_subscription::Model,
46 system_id: Option<String>,
47 config: &Config,
48 ) -> Result<String> {
49 let secret = config
50 .llm_api_secret
51 .as_ref()
52 .context("no LLM API secret")?;
53
54 let plan = if is_staff {
55 Plan::ZedPro
56 } else {
57 subscription.kind.map_or(Plan::ZedFree, |kind| match kind {
58 SubscriptionKind::ZedFree => Plan::ZedFree,
59 SubscriptionKind::ZedPro => Plan::ZedPro,
60 SubscriptionKind::ZedProTrial => Plan::ZedProTrial,
61 })
62 };
63 let subscription_period =
64 billing_subscription::Model::current_period(Some(subscription), is_staff)
65 .map(|(start, end)| (start.naive_utc(), end.naive_utc()))
66 .context("A plan is required to use Zed's hosted models or edit predictions. Visit https://zed.dev/account to get started.")?;
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: true,
86 use_llm_request_queue: feature_flags.iter().any(|flag| flag == "llm-request-queue"),
87 plan,
88 has_extended_trial: feature_flags
89 .iter()
90 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG),
91 subscription_period,
92 enable_model_request_overages: billing_preferences
93 .as_ref()
94 .map_or(false, |preferences| {
95 preferences.model_request_overages_enabled
96 }),
97 model_request_overages_spend_limit_in_cents: billing_preferences
98 .as_ref()
99 .map_or(0, |preferences| {
100 preferences.model_request_overages_spend_limit_in_cents as u32
101 }),
102 };
103
104 Ok(jsonwebtoken::encode(
105 &Header::default(),
106 &claims,
107 &EncodingKey::from_secret(secret.as_ref()),
108 )?)
109 }
110
111 pub fn validate(token: &str, config: &Config) -> Result<LlmTokenClaims, ValidateLlmTokenError> {
112 let secret = config
113 .llm_api_secret
114 .as_ref()
115 .context("no LLM API secret")?;
116
117 match jsonwebtoken::decode::<Self>(
118 token,
119 &DecodingKey::from_secret(secret.as_ref()),
120 &Validation::default(),
121 ) {
122 Ok(token) => Ok(token.claims),
123 Err(e) => {
124 if e.kind() == &jsonwebtoken::errors::ErrorKind::ExpiredSignature {
125 Err(ValidateLlmTokenError::Expired)
126 } else {
127 Err(ValidateLlmTokenError::JwtError(e))
128 }
129 }
130 }
131 }
132}
133
134#[derive(Error, Debug)]
135pub enum ValidateLlmTokenError {
136 #[error("access token is expired")]
137 Expired,
138 #[error("access token validation error: {0}")]
139 JwtError(#[from] jsonwebtoken::errors::Error),
140 #[error("{0}")]
141 Other(#[from] anyhow::Error),
142}