From 7f166298db8f76fa6e2c64ead20e892821709de3 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Tue, 25 Feb 2025 16:45:01 -0500 Subject: [PATCH] collab: Adjust maximum spending limit check (#25596) 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 --- crates/collab/src/llm.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs index c8eadeacd8d49e7081834bb90bd6b22c1d31f896..c365882efb4a17c499ac525af2d3efb012567ba1 100644 --- a/crates/collab/src/llm.rs +++ b/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,