collab: Don't record billing events if billing is not enabled (#19102)

Marshall Bowers created 1 year ago

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

Change summary

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

Detailed changes

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(),
                 )

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
             {