From 69711660ab7a007ed60675f42da4cd447828e706 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Wed, 9 Oct 2024 18:42:22 -0400
Subject: [PATCH] collab: Make LLM billing fields required in `LlmTokenClaims`
 (#18959)

This PR makes the `has_llm_subscription` and
`max_monthly_spend_in_cents` fields in the `LlmTokenClaims` required.

This change will be safe to deploy in ~45 minutes.

Release Notes:

- N/A
---
 crates/collab/src/llm.rs       |  2 +-
 crates/collab/src/llm/token.rs | 21 +++++----------------
 2 files changed, 6 insertions(+), 17 deletions(-)

diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
index efe28a81c73ba880e4036ad90868162d5d4ee65c..c389a4de62fed0bbabc8cfa0e8c2d3c523159234 100644
--- a/crates/collab/src/llm.rs
+++ b/crates/collab/src/llm.rs
@@ -470,7 +470,7 @@ async fn check_usage_limit(
 
     if state.config.is_llm_billing_enabled() {
         if usage.spending_this_month >= FREE_TIER_MONTHLY_SPENDING_LIMIT {
-            if !claims.has_llm_subscription.unwrap_or(false) {
+            if !claims.has_llm_subscription {
                 return Err(Error::http(
                     StatusCode::PAYMENT_REQUIRED,
                     "Maximum spending limit reached for this month.".to_string(),
diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs
index 17f6a8c22660b6c38b37ed3f7aad880a2d84ead4..28f52b5164832d63ae221d7d1336f868ef1be959 100644
--- a/crates/collab/src/llm/token.rs
+++ b/crates/collab/src/llm/token.rs
@@ -20,18 +20,8 @@ pub struct LlmTokenClaims {
     pub github_user_login: String,
     pub is_staff: bool,
     pub has_llm_closed_beta_feature_flag: 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 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 has_llm_subscription: bool,
+    pub max_monthly_spend_in_cents: u32,
     pub plan: rpc::proto::Plan,
 }
 
@@ -63,12 +53,11 @@ impl LlmTokenClaims {
             github_user_login,
             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| {
+            has_llm_subscription,
+            max_monthly_spend_in_cents: billing_preferences
+                .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
                     preferences.max_monthly_llm_usage_spending_in_cents as u32
                 }),
-            ),
             plan,
         };