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 c82968efec..78f514adb7 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 b782db5870..9ee31ab3d1 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(),
                 )