From b44bed0115e9db6c86bd3a1e584a58846ca6abf6 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Fri, 18 Oct 2024 15:55:28 -0400
Subject: [PATCH] collab: Unconditionally execute billing checks (#19432)

This PR removes the conditional checks around the billing-related
enforcement for LLM completions.

These were just in place to prevent executing any billing code before we
had rolled it out. Now that it is rolled out, we don't need this
conditional execution anymore.

Release Notes:

- N/A
---
 crates/collab/src/lib.rs |  4 ----
 crates/collab/src/llm.rs | 52 ++++++++++++++++------------------------
 2 files changed, 21 insertions(+), 35 deletions(-)

diff --git a/crates/collab/src/lib.rs b/crates/collab/src/lib.rs
index c82968efecd9c1864422c669d4fa26ecd0a3cece..78f514adb734d9a2e785384018e6b37f5c9a99b5 100644
--- a/crates/collab/src/lib.rs
+++ b/crates/collab/src/lib.rs
@@ -198,10 +198,6 @@ impl Config {
         }
     }
 
-    pub fn is_llm_billing_enabled(&self) -> bool {
-        self.stripe_api_key.is_some()
-    }
-
     #[cfg(test)]
     pub fn test() -> Self {
         Self {
diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
index b782db58706dfa05321bb22c31b462111cf4fe5a..9ee31ab3d15cff9bb65f7e6eb13b033dbe75fc59 100644
--- a/crates/collab/src/llm.rs
+++ b/crates/collab/src/llm.rs
@@ -460,29 +460,27 @@ async fn check_usage_limit(
         )
         .await?;
 
-    if state.config.is_llm_billing_enabled() {
-        if usage.spending_this_month >= FREE_TIER_MONTHLY_SPENDING_LIMIT {
-            if !claims.has_llm_subscription {
-                return Err(Error::http(
-                    StatusCode::PAYMENT_REQUIRED,
-                    "Maximum spending limit reached for this month.".to_string(),
-                ));
-            }
+    if usage.spending_this_month >= FREE_TIER_MONTHLY_SPENDING_LIMIT {
+        if !claims.has_llm_subscription {
+            return Err(Error::http(
+                StatusCode::PAYMENT_REQUIRED,
+                "Maximum spending limit reached for this month.".to_string(),
+            ));
+        }
 
-            if (usage.spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT)
-                >= Cents(claims.max_monthly_spend_in_cents)
-            {
-                return Err(Error::Http(
-                    StatusCode::FORBIDDEN,
-                    "Maximum spending limit reached for this month.".to_string(),
-                    [(
-                        HeaderName::from_static(MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME),
-                        HeaderValue::from_static("true"),
-                    )]
-                    .into_iter()
-                    .collect(),
-                ));
-            }
+        if (usage.spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT)
+            >= Cents(claims.max_monthly_spend_in_cents)
+        {
+            return Err(Error::Http(
+                StatusCode::FORBIDDEN,
+                "Maximum spending limit reached for this month.".to_string(),
+                [(
+                    HeaderName::from_static(MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME),
+                    HeaderValue::from_static("true"),
+                )]
+                .into_iter()
+                .collect(),
+            ));
         }
     }
 
@@ -627,7 +625,6 @@ where
 impl<S> Drop for TokenCountingStream<S> {
     fn drop(&mut self) {
         let state = self.state.clone();
-        let is_llm_billing_enabled = state.config.is_llm_billing_enabled();
         let claims = self.claims.clone();
         let provider = self.provider;
         let model = std::mem::take(&mut self.model);
@@ -641,14 +638,7 @@ impl<S> Drop for TokenCountingStream<S> {
                     provider,
                     &model,
                     tokens,
-                    // We're passing `false` here if LLM billing is not enabled
-                    // so that we don't write any records to the
-                    // `billing_events` table until we're ready to bill users.
-                    if is_llm_billing_enabled {
-                        claims.has_llm_subscription
-                    } else {
-                        false
-                    },
+                    claims.has_llm_subscription,
                     Cents(claims.max_monthly_spend_in_cents),
                     Utc::now(),
                 )