1use crate::db::user;
  2use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
  3use crate::Cents;
  4use crate::{db::billing_preference, Config};
  5use anyhow::{anyhow, Result};
  6use chrono::Utc;
  7use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
  8use serde::{Deserialize, Serialize};
  9use std::time::Duration;
 10use thiserror::Error;
 11use uuid::Uuid;
 12
 13#[derive(Clone, Debug, Default, Serialize, Deserialize)]
 14#[serde(rename_all = "camelCase")]
 15pub struct LlmTokenClaims {
 16    pub iat: u64,
 17    pub exp: u64,
 18    pub jti: String,
 19    pub user_id: u64,
 20    #[serde(default)]
 21    pub system_id: Option<String>,
 22    #[serde(default)]
 23    pub metrics_id: Option<Uuid>,
 24    pub github_user_login: String,
 25    pub is_staff: bool,
 26    pub has_llm_closed_beta_feature_flag: bool,
 27    pub has_llm_subscription: bool,
 28    pub max_monthly_spend_in_cents: u32,
 29    pub custom_llm_monthly_allowance_in_cents: Option<u32>,
 30    pub plan: rpc::proto::Plan,
 31}
 32
 33const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
 34
 35impl LlmTokenClaims {
 36    #[allow(clippy::too_many_arguments)]
 37    pub fn create(
 38        user: &user::Model,
 39        is_staff: bool,
 40        billing_preferences: Option<billing_preference::Model>,
 41        has_llm_closed_beta_feature_flag: bool,
 42        has_llm_subscription: bool,
 43        plan: rpc::proto::Plan,
 44        system_id: Option<String>,
 45        config: &Config,
 46    ) -> Result<String> {
 47        let secret = config
 48            .llm_api_secret
 49            .as_ref()
 50            .ok_or_else(|| anyhow!("no LLM API secret"))?;
 51
 52        let now = Utc::now();
 53        let claims = Self {
 54            iat: now.timestamp() as u64,
 55            exp: (now + LLM_TOKEN_LIFETIME).timestamp() as u64,
 56            jti: uuid::Uuid::new_v4().to_string(),
 57            user_id: user.id.to_proto(),
 58            system_id,
 59            metrics_id: Some(user.metrics_id),
 60            github_user_login: user.github_login.clone(),
 61            is_staff,
 62            has_llm_closed_beta_feature_flag,
 63            has_llm_subscription,
 64            max_monthly_spend_in_cents: billing_preferences
 65                .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
 66                    preferences.max_monthly_llm_usage_spending_in_cents as u32
 67                }),
 68            custom_llm_monthly_allowance_in_cents: user
 69                .custom_llm_monthly_allowance_in_cents
 70                .map(|allowance| allowance as u32),
 71            plan,
 72        };
 73
 74        Ok(jsonwebtoken::encode(
 75            &Header::default(),
 76            &claims,
 77            &EncodingKey::from_secret(secret.as_ref()),
 78        )?)
 79    }
 80
 81    pub fn validate(token: &str, config: &Config) -> Result<LlmTokenClaims, ValidateLlmTokenError> {
 82        let secret = config
 83            .llm_api_secret
 84            .as_ref()
 85            .ok_or_else(|| anyhow!("no LLM API secret"))?;
 86
 87        match jsonwebtoken::decode::<Self>(
 88            token,
 89            &DecodingKey::from_secret(secret.as_ref()),
 90            &Validation::default(),
 91        ) {
 92            Ok(token) => Ok(token.claims),
 93            Err(e) => {
 94                if e.kind() == &jsonwebtoken::errors::ErrorKind::ExpiredSignature {
 95                    Err(ValidateLlmTokenError::Expired)
 96                } else {
 97                    Err(ValidateLlmTokenError::JwtError(e))
 98                }
 99            }
100        }
101    }
102
103    pub fn free_tier_monthly_spending_limit(&self) -> Cents {
104        self.custom_llm_monthly_allowance_in_cents
105            .map(Cents)
106            .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT)
107    }
108}
109
110#[derive(Error, Debug)]
111pub enum ValidateLlmTokenError {
112    #[error("access token is expired")]
113    Expired,
114    #[error("access token validation error: {0}")]
115    JwtError(#[from] jsonwebtoken::errors::Error),
116    #[error("{0}")]
117    Other(#[from] anyhow::Error),
118}