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