1use crate::Cents;
2use crate::db::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 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 account_created_at: NaiveDateTime,
24 pub is_staff: bool,
25 pub has_llm_closed_beta_feature_flag: bool,
26 pub bypass_account_age_check: bool,
27 pub has_predict_edits_feature_flag: 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}
33
34const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
35
36impl LlmTokenClaims {
37 pub fn create(
38 user: &user::Model,
39 is_staff: bool,
40 billing_preferences: Option<billing_preference::Model>,
41 feature_flags: &Vec<String>,
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 account_created_at: user.account_created_at(),
62 is_staff,
63 has_llm_closed_beta_feature_flag: feature_flags
64 .iter()
65 .any(|flag| flag == "llm-closed-beta"),
66 bypass_account_age_check: feature_flags
67 .iter()
68 .any(|flag| flag == "bypass-account-age-check"),
69 has_predict_edits_feature_flag: feature_flags
70 .iter()
71 .any(|flag| flag == "predict-edits"),
72 has_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 };
82
83 Ok(jsonwebtoken::encode(
84 &Header::default(),
85 &claims,
86 &EncodingKey::from_secret(secret.as_ref()),
87 )?)
88 }
89
90 pub fn validate(token: &str, config: &Config) -> Result<LlmTokenClaims, ValidateLlmTokenError> {
91 let secret = config
92 .llm_api_secret
93 .as_ref()
94 .ok_or_else(|| anyhow!("no LLM API secret"))?;
95
96 match jsonwebtoken::decode::<Self>(
97 token,
98 &DecodingKey::from_secret(secret.as_ref()),
99 &Validation::default(),
100 ) {
101 Ok(token) => Ok(token.claims),
102 Err(e) => {
103 if e.kind() == &jsonwebtoken::errors::ErrorKind::ExpiredSignature {
104 Err(ValidateLlmTokenError::Expired)
105 } else {
106 Err(ValidateLlmTokenError::JwtError(e))
107 }
108 }
109 }
110 }
111
112 pub fn free_tier_monthly_spending_limit(&self) -> Cents {
113 self.custom_llm_monthly_allowance_in_cents
114 .map(Cents)
115 .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT)
116 }
117}
118
119#[derive(Error, Debug)]
120pub enum ValidateLlmTokenError {
121 #[error("access token is expired")]
122 Expired,
123 #[error("access token validation error: {0}")]
124 JwtError(#[from] jsonwebtoken::errors::Error),
125 #[error("{0}")]
126 Other(#[from] anyhow::Error),
127}