collab: Adjust maximum spending limit check (#25596)

Marshall Bowers created 10 months ago

This is a follow-up to https://github.com/zed-industries/zed/pull/25573.

We were still using the spend for a particular model when determining if
the user was over their maximum monthly spend instead of looking at the
usage across all models.

Release Notes:

- N/A

Change summary

crates/collab/src/llm.rs | 12 +++++++-----
1 file changed, 7 insertions(+), 5 deletions(-)

Detailed changes

crates/collab/src/llm.rs

@@ -453,10 +453,6 @@ async fn check_usage_limit(
 
     let user_id = UserId::from_proto(claims.user_id);
     let model = state.db.model(provider, model_name)?;
-    let usage = state
-        .db
-        .get_usage(user_id, provider, model_name, Utc::now())
-        .await?;
     let free_tier = claims.free_tier_monthly_spending_limit();
 
     let spending_this_month = state
@@ -471,7 +467,8 @@ async fn check_usage_limit(
             ));
         }
 
-        if (usage.spending_this_month - free_tier) >= Cents(claims.max_monthly_spend_in_cents) {
+        let monthly_spend = spending_this_month.saturating_sub(free_tier);
+        if monthly_spend >= Cents(claims.max_monthly_spend_in_cents) {
             return Err(Error::Http(
                 StatusCode::FORBIDDEN,
                 "Maximum spending limit reached for this month.".to_string(),
@@ -496,6 +493,11 @@ async fn check_usage_limit(
         model.max_tokens_per_minute as usize / users_in_recent_minutes;
     let per_user_max_tokens_per_day = model.max_tokens_per_day as usize / users_in_recent_days;
 
+    let usage = state
+        .db
+        .get_usage(user_id, provider, model_name, Utc::now())
+        .await?;
+
     let checks = [
         (
             usage.requests_this_minute,