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