collab: Add support for a custom monthly allowance for LLM usage (#19525)

This PR adds support for setting a monthly LLM usage allowance for
certain users.

Release Notes:

- N/A

crates/collab/src/llm/token.rs

@@ -1,8 +1,7 @@
-use crate::llm::DEFAULT_MAX_MONTHLY_SPEND;
-use crate::{
-    db::{billing_preference, UserId},
-    Config,
-};
+use crate::db::user;
+use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
+use crate::Cents;
+use crate::{db::billing_preference, Config};
 use anyhow::{anyhow, Result};
 use chrono::Utc;
 use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
@@ -22,6 +21,7 @@ pub struct LlmTokenClaims {
     pub has_llm_closed_beta_feature_flag: bool,
     pub has_llm_subscription: bool,
     pub max_monthly_spend_in_cents: u32,
+    pub custom_llm_monthly_allowance_in_cents: Option<u32>,
     pub plan: rpc::proto::Plan,
 }
 
@@ -30,8 +30,7 @@ const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
 impl LlmTokenClaims {
     #[allow(clippy::too_many_arguments)]
     pub fn create(
-        user_id: UserId,
-        github_user_login: String,
+        user: &user::Model,
         is_staff: bool,
         billing_preferences: Option<billing_preference::Model>,
         has_llm_closed_beta_feature_flag: bool,
@@ -49,8 +48,8 @@ impl LlmTokenClaims {
             iat: now.timestamp() as u64,
             exp: (now + LLM_TOKEN_LIFETIME).timestamp() as u64,
             jti: uuid::Uuid::new_v4().to_string(),
-            user_id: user_id.to_proto(),
-            github_user_login,
+            user_id: user.id.to_proto(),
+            github_user_login: user.github_login.clone(),
             is_staff,
             has_llm_closed_beta_feature_flag,
             has_llm_subscription,
@@ -58,6 +57,9 @@ impl LlmTokenClaims {
                 .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
                     preferences.max_monthly_llm_usage_spending_in_cents as u32
                 }),
+            custom_llm_monthly_allowance_in_cents: user
+                .custom_llm_monthly_allowance_in_cents
+                .map(|allowance| allowance as u32),
             plan,
         };
 
@@ -89,6 +91,12 @@ impl LlmTokenClaims {
             }
         }
     }
+
+    pub fn free_tier_monthly_spending_limit(&self) -> Cents {
+        self.custom_llm_monthly_allowance_in_cents
+            .map(Cents)
+            .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT)
+    }
 }
 
 #[derive(Error, Debug)]