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/db/queries/usages.rs

@@ -1,5 +1,5 @@
+use crate::db::UserId;
 use crate::llm::Cents;
-use crate::{db::UserId, llm::FREE_TIER_MONTHLY_SPENDING_LIMIT};
 use chrono::{Datelike, Duration};
 use futures::StreamExt as _;
 use rpc::LanguageModelProvider;
@@ -299,6 +299,7 @@ impl LlmDatabase {
         tokens: TokenUsage,
         has_llm_subscription: bool,
         max_monthly_spend: Cents,
+        free_tier_monthly_spending_limit: Cents,
         now: DateTimeUtc,
     ) -> Result<Usage> {
         self.transaction(|tx| async move {
@@ -410,9 +411,9 @@ impl LlmDatabase {
             );
 
             if !is_staff
-                && spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT
+                && spending_this_month > free_tier_monthly_spending_limit
                 && has_llm_subscription
-                && (spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT) <= max_monthly_spend
+                && (spending_this_month - free_tier_monthly_spending_limit) <= max_monthly_spend
             {
                 billing_event::ActiveModel {
                     id: ActiveValue::not_set(),

crates/collab/src/llm/db/tests/billing_tests.rs

@@ -66,6 +66,7 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
         usage,
         true,
         max_monthly_spend,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         now,
     )
     .await
@@ -103,6 +104,7 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
         usage_2,
         true,
         max_monthly_spend,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         now,
     )
     .await
@@ -132,6 +134,7 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
         model,
         usage_exceeding,
         true,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         max_monthly_spend,
         now,
     )

crates/collab/src/llm/db/tests/usage_tests.rs

@@ -1,3 +1,4 @@
+use crate::llm::FREE_TIER_MONTHLY_SPENDING_LIMIT;
 use crate::{
     db::UserId,
     llm::db::{
@@ -49,6 +50,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
         },
         false,
         Cents::ZERO,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         now,
     )
     .await
@@ -68,6 +70,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
         },
         false,
         Cents::ZERO,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         now,
     )
     .await
@@ -124,6 +127,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
         },
         false,
         Cents::ZERO,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         now,
     )
     .await
@@ -180,6 +184,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
         },
         false,
         Cents::ZERO,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         now,
     )
     .await
@@ -222,6 +227,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
         },
         false,
         Cents::ZERO,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         now,
     )
     .await
@@ -259,6 +265,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
         },
         false,
         Cents::ZERO,
+        FREE_TIER_MONTHLY_SPENDING_LIMIT,
         now,
     )
     .await

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)]