From c709b66f35d7e6d07b2754933bf962aa7d01ecef Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Fri, 11 Oct 2024 17:54:10 -0400
Subject: [PATCH] collab: Don't record billing events if billing is not enabled
 (#19102)

This PR adjusts the billing logic to not write any records to
`billing_events` if:

- The user is staff, as we don't want to bill staff members
- Billing is disabled (we currently enable billing based on the presence
of the Stripe API key)

Release Notes:

- N/A
---
 crates/collab/src/llm.rs                   | 10 +++++++++-
 crates/collab/src/llm/db/queries/usages.rs |  3 ++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
index 70a309367c53850be3fca8a18d4ecb98b3a8ebd3..583cfdb7b82ea017d08330ca61709b2561d1ad60 100644
--- a/crates/collab/src/llm.rs
+++ b/crates/collab/src/llm.rs
@@ -639,6 +639,7 @@ 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);
@@ -652,7 +653,14 @@ impl<S> Drop for TokenCountingStream<S> {
                     provider,
                     &model,
                     tokens,
-                    claims.has_llm_subscription,
+                    // 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
+                    },
                     Cents(claims.max_monthly_spend_in_cents),
                     Utc::now(),
                 )
diff --git a/crates/collab/src/llm/db/queries/usages.rs b/crates/collab/src/llm/db/queries/usages.rs
index c7f8fd163a4272b1701166b8c6295ddc9c6ac5fc..1e9204b02b21c83eb573108513394077fdc76a5d 100644
--- a/crates/collab/src/llm/db/queries/usages.rs
+++ b/crates/collab/src/llm/db/queries/usages.rs
@@ -409,7 +409,8 @@ impl LlmDatabase {
                 monthly_usage.output_tokens as usize,
             );
 
-            if spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT
+            if !is_staff
+                && spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT
                 && has_llm_subscription
                 && spending_this_month <= max_monthly_spend
             {