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    #[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_predict_edits_feature_flag: bool,
 43        has_llm_subscription: bool,
 44        plan: rpc::proto::Plan,
 45        system_id: Option<String>,
 46        config: &Config,
 47    ) -> Result<String> {
 48        let secret = config
 49            .llm_api_secret
 50            .as_ref()
 51            .ok_or_else(|| anyhow!("no LLM API secret"))?;
 52
 53        let now = Utc::now();
 54        let claims = Self {
 55            iat: now.timestamp() as u64,
 56            exp: (now + LLM_TOKEN_LIFETIME).timestamp() as u64,
 57            jti: uuid::Uuid::new_v4().to_string(),
 58            user_id: user.id.to_proto(),
 59            system_id,
 60            metrics_id: user.metrics_id,
 61            github_user_login: user.github_login.clone(),
 62            is_staff,
 63            has_llm_closed_beta_feature_flag,
 64            has_predict_edits_feature_flag,
 65            has_llm_subscription,
 66            max_monthly_spend_in_cents: billing_preferences
 67                .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
 68                    preferences.max_monthly_llm_usage_spending_in_cents as u32
 69                }),
 70            custom_llm_monthly_allowance_in_cents: user
 71                .custom_llm_monthly_allowance_in_cents
 72                .map(|allowance| allowance as u32),
 73            plan,
 74        };
 75
 76        Ok(jsonwebtoken::encode(
 77            &Header::default(),
 78            &claims,
 79            &EncodingKey::from_secret(secret.as_ref()),
 80        )?)
 81    }
 82
 83    pub fn validate(token: &str, config: &Config) -> Result<LlmTokenClaims, ValidateLlmTokenError> {
 84        let secret = config
 85            .llm_api_secret
 86            .as_ref()
 87            .ok_or_else(|| anyhow!("no LLM API secret"))?;
 88
 89        match jsonwebtoken::decode::<Self>(
 90            token,
 91            &DecodingKey::from_secret(secret.as_ref()),
 92            &Validation::default(),
 93        ) {
 94            Ok(token) => Ok(token.claims),
 95            Err(e) => {
 96                if e.kind() == &jsonwebtoken::errors::ErrorKind::ExpiredSignature {
 97                    Err(ValidateLlmTokenError::Expired)
 98                } else {
 99                    Err(ValidateLlmTokenError::JwtError(e))
100                }
101            }
102        }
103    }
104
105    pub fn free_tier_monthly_spending_limit(&self) -> Cents {
106        self.custom_llm_monthly_allowance_in_cents
107            .map(Cents)
108            .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT)
109    }
110}
111
112#[derive(Error, Debug)]
113pub enum ValidateLlmTokenError {
114    #[error("access token is expired")]
115    Expired,
116    #[error("access token validation error: {0}")]
117    JwtError(#[from] jsonwebtoken::errors::Error),
118    #[error("{0}")]
119    Other(#[from] anyhow::Error),
120}