From 82ba4fcbba2e4b0eae368d3149452c4dc5e50c30 Mon Sep 17 00:00:00 2001 From: Zyy Date: Fri, 14 Aug 2026 00:11:56 +0800 Subject: [PATCH] fix(observability): align dashboard and quota window usage Co-authored-by: Codex --- .../postgres/observability/admin_adapter.rs | 6 +- .../provider_accounts/admin_adapter.rs | 65 ++++++++++++++---- .../src/postgres/provider_accounts/mapping.rs | 63 +++++++++++++++++ .../src/postgres/provider_accounts/rows.rs | 68 +++++++++++++++++++ .../tests/postgres/observability.rs | 8 ++- .../tests/postgres/provider_accounts.rs | 26 +++++-- 6 files changed, 210 insertions(+), 26 deletions(-) diff --git a/backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs b/backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs index ba3930cfd..6cd110ff7 100644 --- a/backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs +++ b/backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs @@ -101,12 +101,8 @@ impl ObservabilityRepository for PgObservabilityRepository { observed_at: DateTime, ) -> StoreResult { let filter = UsageRecordFilter::default(); - let account_usage_range = ObservabilityRange::new( - range.end - TimeDelta::hours(ACCOUNT_USAGE_TIMELINE_HOURS), - range.end, - )?; let account_usage_query = - ProviderAccountUsageQuery::recent(account_usage_range, DASHBOARD_ACCOUNT_LIMIT)? + ProviderAccountUsageQuery::recent(range, DASHBOARD_ACCOUNT_LIMIT)? .with_hourly_request_buckets()?; let recent_query = UsageRecordQuery { range, diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs index 1254b1b65..0f6081b61 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs @@ -97,22 +97,59 @@ impl PgAdminAccountStore { .iter() .map(|window| window.account_id.clone()) .collect::>(); - sqlx::query(ACCOUNT_USAGE_BY_WINDOWS_SQL) - .bind(account_ids) - .bind(keys) - .bind(starts) - .bind(ends) - .fetch_all(&self.pool) - .await - .map_err(|_| { - admin_store_error( - ENTITY, - postgres_unavailable("load provider account quota window usage"), - ) - })? + let (usage_rows, model_rows, model_cost_rows) = futures::try_join!( + sqlx::query(ACCOUNT_USAGE_BY_WINDOWS_SQL) + .bind(account_ids.clone()) + .bind(keys.clone()) + .bind(starts.clone()) + .bind(ends.clone()) + .fetch_all(&self.pool), + sqlx::query(ACCOUNT_USAGE_MODELS_BY_WINDOWS_SQL) + .bind(account_ids.clone()) + .bind(keys.clone()) + .bind(starts.clone()) + .bind(ends.clone()) + .fetch_all(&self.pool), + sqlx::query(ACCOUNT_USAGE_MODEL_COSTS_BY_WINDOWS_SQL) + .bind(account_ids) + .bind(keys) + .bind(starts) + .bind(ends) + .fetch_all(&self.pool), + ) + .map_err(|_| { + admin_store_error( + ENTITY, + postgres_unavailable("load provider account quota window usage"), + ) + })?; + let mut model_costs = BTreeMap::<(String, String, String), Vec>::new(); + for row in &model_cost_rows { + let (key, cost) = admin_account_usage_window_model_cost(row)?; + model_costs.entry(key).or_default().push(cost); + } + let mut models_by_window = BTreeMap::<(String, String), Vec>::new(); + for row in &model_rows { + let ((account_id, window_key, model), mut usage) = + admin_account_usage_window_model(row)?; + usage.costs = model_costs + .remove(&(account_id.clone(), window_key.clone(), model)) + .unwrap_or_default(); + models_by_window + .entry((account_id, window_key)) + .or_default() + .push(usage); + } + let mut results = usage_rows .iter() .map(admin_account_usage_window) - .collect() + .collect::>>()?; + for result in &mut results { + result.usage.models = models_by_window + .remove(&(result.account_id.clone(), result.key.clone())) + .unwrap_or_default(); + } + Ok(results) } async fn required_scope( diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs index bd5e221fa..61b7f00db 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs @@ -276,6 +276,69 @@ pub(crate) fn admin_account_usage_window( }) } +pub(crate) fn admin_account_usage_window_model( + row: &sqlx::postgres::PgRow, +) -> AdminStoreResult<((String, String, String), AccountModelUsage)> { + let account_id: String = window_usage_value(row, "account_id")?; + let window_key: String = window_usage_value(row, "window_key")?; + let model: String = window_usage_value(row, "model")?; + let provider_reported_count = window_usage_count(row, "provider_reported_count")?; + let calculated_count = window_usage_count(row, "calculated_count")?; + let unavailable_count = window_usage_count(row, "unavailable_count")?; + Ok(( + (account_id, window_key, model.clone()), + AccountModelUsage { + model, + request_count: window_usage_count(row, "request_count")?, + success_count: window_usage_count(row, "success_count")?, + input_tokens: optional_window_usage_count(row, "input_tokens")?, + output_tokens: optional_window_usage_count(row, "output_tokens")?, + cached_tokens: optional_window_usage_count(row, "cached_tokens")?, + cache_write_tokens: optional_window_usage_count(row, "cache_write_tokens")?, + reasoning_tokens: optional_window_usage_count(row, "reasoning_tokens")?, + image_input_tokens: optional_window_usage_count(row, "image_input_tokens")?, + image_output_tokens: optional_window_usage_count(row, "image_output_tokens")?, + image_request_count: window_usage_count(row, "image_request_count")?, + image_request_failed_count: window_usage_count(row, "image_request_failed_count")?, + total_tokens: optional_window_usage_count(row, "total_tokens")?, + cost_coverage: AdminCostCoverage { + provider_reported_count, + calculated_count, + partial_count: 0, + unavailable_count, + not_billable_count: 0, + }, + costs: Vec::new(), + last_used_at: window_usage_value(row, "last_used_at")?, + }, + )) +} + +pub(crate) fn admin_account_usage_window_model_cost( + row: &sqlx::postgres::PgRow, +) -> AdminStoreResult<((String, String, String), AccountCost)> { + let key = ( + window_usage_value(row, "account_id")?, + window_usage_value(row, "window_key")?, + window_usage_value(row, "model")?, + ); + let amount = window_usage_value::(row, "amount")?; + let amount = AdminDecimalAmount::from_str(&amount).map_err(|_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + ENTITY, + "persisted account model cost is invalid", + ) + })?; + Ok(( + key, + AccountCost { + currency: window_usage_value(row, "cost_currency")?, + amount, + }, + )) +} + pub(crate) fn window_usage_count( row: &sqlx::postgres::PgRow, column: &'static str, diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs index 9b2c674c0..1dfbf0c0a 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs @@ -51,6 +51,74 @@ pub(crate) const ACCOUNT_USAGE_BY_WINDOWS_SQL: &str = "with requested_windows as group by requested.account_id, requested.window_key order by requested.account_id, requested.window_key"; +pub(crate) const ACCOUNT_USAGE_MODELS_BY_WINDOWS_SQL: &str = "with requested_windows as ( + select * + from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[]) + as requested(account_id, window_key, window_start, window_end) + ) + select requested.account_id, + requested.window_key, + coalesce(mr.upstream_model_id, mr.requested_model_id) as model, + count(mr.id)::bigint as request_count, + count(mr.id) filter (where mr.outcome = 'succeeded')::bigint as success_count, + sum(mr.input_tokens)::bigint as input_tokens, + sum(mr.output_tokens)::bigint as output_tokens, + sum(mr.cached_tokens)::bigint as cached_tokens, + sum(mr.cache_write_tokens)::bigint as cache_write_tokens, + sum(mr.reasoning_tokens)::bigint as reasoning_tokens, + sum(mr.image_input_tokens)::bigint as image_input_tokens, + sum(mr.image_output_tokens)::bigint as image_output_tokens, + count(mr.id) filter (where mr.image_generation_succeeded is true)::bigint + as image_request_count, + count(mr.id) filter (where mr.image_generation_succeeded is false)::bigint + as image_request_failed_count, + coalesce(sum(coalesce( + mr.total_tokens, + coalesce(mr.input_tokens, 0) + coalesce(mr.output_tokens, 0) + )), 0)::bigint as total_tokens, + count(mr.id) filter (where mr.cost_source = 'provider_reported')::bigint + as provider_reported_count, + count(mr.id) filter (where mr.cost_source = 'calculated')::bigint + as calculated_count, + count(mr.id) filter (where mr.cost_source = 'unavailable')::bigint + as unavailable_count, + max(mr.started_at) as last_used_at + from requested_windows requested + join model_requests mr + on mr.provider_account_ref = requested.account_id + and mr.started_at >= requested.window_start + and mr.started_at < requested.window_end + and mr.outcome = 'succeeded' + and mr.downstream_committed_at is not null + and mr.client_status_code between 200 and 399 + group by requested.account_id, requested.window_key, + coalesce(mr.upstream_model_id, mr.requested_model_id) + order by requested.account_id, requested.window_key, request_count desc, model"; + +pub(crate) const ACCOUNT_USAGE_MODEL_COSTS_BY_WINDOWS_SQL: &str = "with requested_windows as ( + select * + from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[]) + as requested(account_id, window_key, window_start, window_end) + ) + select requested.account_id, + requested.window_key, + coalesce(mr.upstream_model_id, mr.requested_model_id) as model, + mr.cost_currency, + sum(mr.cost_amount)::text as amount + from requested_windows requested + join model_requests mr + on mr.provider_account_ref = requested.account_id + and mr.started_at >= requested.window_start + and mr.started_at < requested.window_end + and mr.outcome = 'succeeded' + and mr.downstream_committed_at is not null + and mr.client_status_code between 200 and 399 + and mr.cost_amount is not null + and mr.cost_currency is not null + group by requested.account_id, requested.window_key, + coalesce(mr.upstream_model_id, mr.requested_model_id), mr.cost_currency + order by requested.account_id, requested.window_key, model, mr.cost_currency"; + #[derive(Debug, Clone, PartialEq, Eq)] pub struct ProviderAccountAdminScope { pub provider_kind: String, diff --git a/backend/crates/gateway-store/tests/postgres/observability.rs b/backend/crates/gateway-store/tests/postgres/observability.rs index fa2877c82..d7aa797a1 100644 --- a/backend/crates/gateway-store/tests/postgres/observability.rs +++ b/backend/crates/gateway-store/tests/postgres/observability.rs @@ -433,7 +433,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det assert_eq!(dashboard.attempts.costs[0].amount.as_str(), "1.25"); assert_eq!(dashboard.provider_accounts.total, 1); assert_eq!(dashboard.account_usage[0].request_count, 1); - assert_eq!(dashboard.account_usage[0].request_buckets.len(), 24); + assert_eq!(dashboard.account_usage[0].request_buckets.len(), 2); assert_eq!( dashboard.account_usage[0] .request_buckets @@ -714,7 +714,7 @@ async fn dashboard_summary_totals_include_history_outside_selected_range() { .expect("seed observability facts"); sqlx::query( "update model_requests - set started_at = $1 - interval '2 days' + set started_at = $1 - interval '2 hours' where id = 'req_observe_success'", ) .bind(now) @@ -730,6 +730,8 @@ async fn dashboard_summary_totals_include_history_outside_selected_range() { .await .expect("dashboard summary"); assert_eq!(dashboard.requests.request_count, 2); + assert_eq!(dashboard.account_usage[0].request_count, 0); + assert_eq!(dashboard.account_usage[0].request_buckets.len(), 2); assert_eq!( ( dashboard.totals.request_count, @@ -813,7 +815,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact ); assert_eq!(dashboard.provider_accounts.total, 1); assert_eq!(dashboard.account_usage[0].request_count, 1); - assert_eq!(dashboard.account_usage[0].request_buckets.len(), 24); + assert_eq!(dashboard.account_usage[0].request_buckets.len(), 2); assert_eq!( dashboard.account_usage[0] .request_buckets diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs index 64b24112c..f8bafd810 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs @@ -656,9 +656,9 @@ async fn terminal_admin_quota_window_usage_prefers_provider_total_and_falls_back request_id: "req_quota_window_recent", account_id: "acct_quota_window", provider_kind: "openai", - model: "gpt-window", + model: "gpt-window-recent", total_tokens: 10, - cost_amount: "0", + cost_amount: "1.25", started_at: now - TimeDelta::minutes(30), }, ) @@ -670,9 +670,9 @@ async fn terminal_admin_quota_window_usage_prefers_provider_total_and_falls_back request_id: "req_quota_window_older", account_id: "acct_quota_window", provider_kind: "openai", - model: "gpt-window", + model: "gpt-window-older", total_tokens: 20, - cost_amount: "0", + cost_amount: "2.50", started_at: now - TimeDelta::hours(6), }, ) @@ -723,8 +723,26 @@ async fn terminal_admin_quota_window_usage_prefers_provider_total_and_falls_back assert_eq!(usage[0].key, "long"); assert_eq!(usage[0].usage.total_tokens, Some(720)); + assert_eq!(usage[0].usage.models.len(), 2); + assert_eq!( + usage[0] + .usage + .models + .iter() + .map(|model| (model.model.as_str(), model.total_tokens)) + .collect::>(), + [ + ("gpt-window-older", Some(20)), + ("gpt-window-recent", Some(700)), + ], + ); + assert_eq!(usage[0].usage.models[0].costs[0].amount.as_str(), "2.5"); + assert_eq!(usage[0].usage.models[1].costs[0].amount.as_str(), "1.25"); assert_eq!(usage[1].key, "short"); assert_eq!(usage[1].usage.total_tokens, Some(700)); + assert_eq!(usage[1].usage.models.len(), 1); + assert_eq!(usage[1].usage.models[0].model, "gpt-window-recent"); + assert_eq!(usage[1].usage.models[0].total_tokens, Some(700)); database.close().await; }