diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs index 70a309367c5385..583cfdb7b82ea0 100644 --- a/crates/collab/src/llm.rs +++ b/crates/collab/src/llm.rs @@ -639,6 +639,7 @@ where impl Drop for TokenCountingStream { 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 Drop for TokenCountingStream { 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 c7f8fd163a4272..1e9204b02b21c8 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 {