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