From 3db789ed90f1cd7c39c5d9c5fc4b8fe414fa63e5 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Wed, 9 Oct 2024 17:39:34 -0400
Subject: [PATCH] collab: Include max monthly spend preference in LLM token
 (#18955)

This PR updates the LLM token claims to include the maximum monthly
spend.

Release Notes:

- N/A
---
 crates/collab/src/llm/token.rs | 19 ++++++++++++++++++-
 crates/collab/src/rpc.rs       |  4 ++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs
index 6098a1cd13..17f6a8c226 100644
--- a/crates/collab/src/llm/token.rs
+++ b/crates/collab/src/llm/token.rs
@@ -1,4 +1,8 @@
-use crate::{db::UserId, Config};
+use crate::llm::DEFAULT_MAX_MONTHLY_SPEND;
+use crate::{
+    db::{billing_preference, UserId},
+    Config,
+};
 use anyhow::{anyhow, Result};
 use chrono::Utc;
 use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
@@ -22,16 +26,24 @@ pub struct LlmTokenClaims {
     // this change has been deployed).
     #[serde(default)]
     pub has_llm_subscription: Option<bool>,
+    // This field is temporarily optional so it can be added
+    // in a backwards-compatible way. We can make it required
+    // once all of the LLM tokens have cycled (~1 hour after
+    // this change has been deployed).
+    #[serde(default)]
+    pub max_monthly_spend_in_cents: Option<u32>,
     pub plan: rpc::proto::Plan,
 }
 
 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,
         is_staff: bool,
+        billing_preferences: Option<billing_preference::Model>,
         has_llm_closed_beta_feature_flag: bool,
         has_llm_subscription: bool,
         plan: rpc::proto::Plan,
@@ -52,6 +64,11 @@ impl LlmTokenClaims {
             is_staff,
             has_llm_closed_beta_feature_flag,
             has_llm_subscription: Some(has_llm_subscription),
+            max_monthly_spend_in_cents: Some(
+                billing_preferences.map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
+                    preferences.max_monthly_llm_usage_spending_in_cents as u32
+                }),
+            ),
             plan,
         };
 
diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs
index e5e401e493..2b861dd403 100644
--- a/crates/collab/src/rpc.rs
+++ b/crates/collab/src/rpc.rs
@@ -4917,10 +4917,14 @@ async fn get_llm_api_token(
     if Utc::now().naive_utc() - account_created_at < MIN_ACCOUNT_AGE_FOR_LLM_USE {
         Err(anyhow!("account too young"))?
     }
+
+    let billing_preferences = db.get_billing_preferences(user.id).await?;
+
     let token = LlmTokenClaims::create(
         user.id,
         user.github_login.clone(),
         session.is_staff(),
+        billing_preferences,
         has_llm_closed_beta_feature_flag,
         session.has_llm_subscription(&db).await?,
         session.current_plan(&db).await?,