starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(observability): align dashboard and quota window usage

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-08-14 00:11:56 +08:00
co-authored by Codex
parent 819f428ac3
commit 82ba4fcbba
6 changed files with 210 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs
+1 -5
View File
@@ -101,12 +101,8 @@ impl ObservabilityRepository for PgObservabilityRepository {
observed_at: DateTime<Utc>,
) -> StoreResult<DashboardObservation> {
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,
backend/crates/gateway-store/src/postgres/provider_accounts/admin_adapter.rs
+51 -14
View File
@@ -97,22 +97,59 @@ impl PgAdminAccountStore {
.iter()
.map(|window| window.account_id.clone())
.collect::<Vec<_>>();
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<AccountCost>>::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<AccountModelUsage>>::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::<AdminStoreResult<Vec<_>>>()?;
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(
backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs
+63
View File
@@ -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::<String>(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,
backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs
+68
View File
@@ -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,
backend/crates/gateway-store/tests/postgres/observability.rs
+5 -3
View File
@@ -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
backend/crates/gateway-store/tests/postgres/provider_accounts.rs
+22 -4
View File
@@ -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::<Vec<_>>(),
[
("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;
}