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    pub system_id: Option<String>,
 21    pub metrics_id: Uuid,
 22    pub github_user_login: String,
 23    pub is_staff: bool,
 24    pub has_llm_closed_beta_feature_flag: bool,
 25    #[serde(default)]
 26    pub has_predict_edits_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    pub fn create(
 37        user: &user::Model,
 38        is_staff: bool,
 39        billing_preferences: Option<billing_preference::Model>,
 40        has_llm_closed_beta_feature_flag: bool,
 41        has_predict_edits_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: user.metrics_id,
 60            github_user_login: user.github_login.clone(),
 61            is_staff,
 62            has_llm_closed_beta_feature_flag,
 63            has_predict_edits_feature_flag,
 64            has_llm_subscription,
 65            max_monthly_spend_in_cents: billing_preferences
 66                .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
 67                    preferences.max_monthly_llm_usage_spending_in_cents as u32
 68                }),
 69            custom_llm_monthly_allowance_in_cents: user
 70                .custom_llm_monthly_allowance_in_cents
 71                .map(|allowance| allowance as u32),
 72            plan,
 73        };
 74
 75        Ok(jsonwebtoken::encode(
 76            &Header::default(),
 77            &claims,
 78            &EncodingKey::from_secret(secret.as_ref()),
 79        )?)
 80    }
 81
 82    pub fn validate(token: &str, config: &Config) -> Result<LlmTokenClaims, ValidateLlmTokenError> {
 83        let secret = config
 84            .llm_api_secret
 85            .as_ref()
 86            .ok_or_else(|| anyhow!("no LLM API secret"))?;
 87
 88        match jsonwebtoken::decode::<Self>(
 89            token,
 90            &DecodingKey::from_secret(secret.as_ref()),
 91            &Validation::default(),
 92        ) {
 93            Ok(token) => Ok(token.claims),
 94            Err(e) => {
 95                if e.kind() == &jsonwebtoken::errors::ErrorKind::ExpiredSignature {
 96                    Err(ValidateLlmTokenError::Expired)
 97                } else {
 98                    Err(ValidateLlmTokenError::JwtError(e))
 99                }
100            }
101        }
102    }
103
104    pub fn free_tier_monthly_spending_limit(&self) -> Cents {
105        self.custom_llm_monthly_allowance_in_cents
106            .map(Cents)
107            .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT)
108    }
109}
110
111#[derive(Error, Debug)]
112pub enum ValidateLlmTokenError {
113    #[error("access token is expired")]
114    Expired,
115    #[error("access token validation error: {0}")]
116    JwtError(#[from] jsonwebtoken::errors::Error),
117    #[error("{0}")]
118    Other(#[from] anyhow::Error),
119}