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