diff --git a/crates/collab/src/api/billing.rs b/crates/collab/src/api/billing.rs index fc854303be..624dfa2fc4 100644 --- a/crates/collab/src/api/billing.rs +++ b/crates/collab/src/api/billing.rs @@ -1090,9 +1090,17 @@ struct UsageCounts { pub remaining: Option, } +#[derive(Debug, Serialize)] +struct ModelRequestUsage { + pub model: String, + pub mode: CompletionMode, + pub requests: i32, +} + #[derive(Debug, Serialize)] struct GetCurrentUsageResponse { pub model_requests: UsageCounts, + pub model_request_usage: Vec, pub edit_predictions: UsageCounts, } @@ -1119,6 +1127,7 @@ async fn get_current_usage( limit: Some(0), remaining: Some(0), }, + model_request_usage: Vec::new(), edit_predictions: UsageCounts { used: 0, limit: Some(0), @@ -1163,12 +1172,30 @@ async fn get_current_usage( zed_llm_client::UsageLimit::Unlimited => None, }; + let subscription_usage_meters = llm_db + .get_current_subscription_usage_meters_for_user(user.id, Utc::now()) + .await?; + + let model_request_usage = subscription_usage_meters + .into_iter() + .filter_map(|(usage_meter, _usage)| { + let model = llm_db.model_by_id(usage_meter.model_id).ok()?; + + Some(ModelRequestUsage { + model: model.name.clone(), + mode: usage_meter.mode, + requests: usage_meter.requests, + }) + }) + .collect::>(); + Ok(Json(GetCurrentUsageResponse { model_requests: UsageCounts { used: usage.model_requests, limit: model_requests_limit, remaining: model_requests_limit.map(|limit| (limit - usage.model_requests).max(0)), }, + model_request_usage, edit_predictions: UsageCounts { used: usage.edit_predictions, limit: edit_prediction_limit, diff --git a/crates/collab/src/llm/db/queries/subscription_usage_meters.rs b/crates/collab/src/llm/db/queries/subscription_usage_meters.rs index 5666b1cda9..c0ce5d679b 100644 --- a/crates/collab/src/llm/db/queries/subscription_usage_meters.rs +++ b/crates/collab/src/llm/db/queries/subscription_usage_meters.rs @@ -1,3 +1,4 @@ +use crate::db::UserId; use crate::llm::db::queries::subscription_usages::convert_chrono_to_time; use super::*; @@ -34,4 +35,38 @@ impl LlmDatabase { }) .await } + + /// Returns all current subscription usage meters for the given user as of the given timestamp. + pub async fn get_current_subscription_usage_meters_for_user( + &self, + user_id: UserId, + now: DateTimeUtc, + ) -> Result> { + let now = convert_chrono_to_time(now)?; + + self.transaction(|tx| async move { + let result = subscription_usage_meter::Entity::find() + .inner_join(subscription_usage::Entity) + .filter(subscription_usage::Column::UserId.eq(user_id)) + .filter( + subscription_usage::Column::PeriodStartAt + .lte(now) + .and(subscription_usage::Column::PeriodEndAt.gte(now)), + ) + .select_also(subscription_usage::Entity) + .all(&*tx) + .await?; + + let result = result + .into_iter() + .filter_map(|(meter, usage)| { + let usage = usage?; + Some((meter, usage)) + }) + .collect(); + + Ok(result) + }) + .await + } }