From 1101eee628caa8b9fc84d261fa9490742a7cd288 Mon Sep 17 00:00:00 2001 From: Zyy Date: Mon, 6 Jul 2026 17:49:21 +0800 Subject: [PATCH] fix: align quota window accounting Co-authored-by: Codex --- backend/src/admin/accounts/quota_view.rs | 14 +- backend/src/admin/accounts/routes.rs | 120 +++++++++++++ backend/src/runtime/tasks/quota_refresh.rs | 3 +- .../src/upstream/accounts/quota/runtime.rs | 28 ++- backend/tests/admin/accounts/list.rs | 138 +++++++++++++-- .../tests/upstream/accounts/quota_refresh.rs | 160 +++++++++++++++++- 6 files changed, 442 insertions(+), 21 deletions(-) diff --git a/backend/src/admin/accounts/quota_view.rs b/backend/src/admin/accounts/quota_view.rs index 6e973ceb5..42206fa4c 100644 --- a/backend/src/admin/accounts/quota_view.rs +++ b/backend/src/admin/accounts/quota_view.rs @@ -43,6 +43,16 @@ pub(crate) struct AdminAccountQuotaUsageWindow { pub(crate) key: String, pub(crate) start: DateTime, pub(crate) end: DateTime, + pub(crate) window_seconds: u64, +} + +impl AdminAccountQuotaUsageWindow { + pub(crate) fn duration_seconds(&self) -> u64 { + self.end + .signed_duration_since(self.start) + .num_seconds() + .max(0) as u64 + } } #[derive(Debug, Clone, Copy, Default)] @@ -101,12 +111,14 @@ impl AdminAccountQuotaData { impl AdminAccountQuotaWindowData { fn usage_window(&self) -> Option { let end = self.reset_at?; - let seconds = i64::try_from(self.window_seconds?).ok()?; + let window_seconds = self.window_seconds?; + let seconds = i64::try_from(window_seconds).ok()?; let start = end.checked_sub_signed(Duration::seconds(seconds))?; (start <= end).then(|| AdminAccountQuotaUsageWindow { key: self.key.clone(), start, end, + window_seconds, }) } } diff --git a/backend/src/admin/accounts/routes.rs b/backend/src/admin/accounts/routes.rs index 7081ee34f..17d1e5e60 100644 --- a/backend/src/admin/accounts/routes.rs +++ b/backend/src/admin/accounts/routes.rs @@ -303,6 +303,14 @@ struct AccountModelUsageRecord { last_used_at: Option>, } +#[derive(Debug, Clone, Copy, Default)] +struct AccountUsageWindowStats { + request_count: i64, + input_tokens: i64, + output_tokens: i64, + cached_tokens: i64, +} + #[derive(Debug, Clone)] struct AccountQuotaUsageWindow { account_id: String, @@ -873,6 +881,19 @@ async fn account_list_stats( .unwrap_or_default(); let mut quota_by_account = quota_by_account.clone(); apply_quota_window_token_usage(state, accounts, &mut quota_by_account).await; + let selected_quota_windows = selected_quota_windows_by_account(accounts, "a_by_account); + let quota_window_usage_stats = + account_usage_stats_by_quota_window(state, &selected_quota_windows).await; + let usage_records = usage_records + .into_iter() + .map(|usage| { + apply_selected_quota_window_usage( + usage, + &selected_quota_windows, + "a_window_usage_stats, + ) + }) + .collect::>(); let models_by_account = list_current_window_model_usage(state, &usage_records).await; AccountListStats { @@ -885,6 +906,105 @@ async fn account_list_stats( } } +fn selected_quota_windows_by_account( + accounts: &[AdminAccountMetadata], + quota_by_account: &HashMap, +) -> HashMap { + accounts + .iter() + .filter_map(|account| { + quota_by_account + .get(&account.id) + .and_then(selected_quota_window) + .map(|window| (account.id.clone(), window)) + }) + .collect() +} + +fn selected_quota_window(quota: &AdminAccountQuotaData) -> Option { + quota.usage_windows().into_iter().max_by(|left, right| { + left.duration_seconds() + .cmp(&right.duration_seconds()) + .then_with(|| left.end.cmp(&right.end)) + .then_with(|| left.key.cmp(&right.key)) + }) +} + +async fn account_usage_stats_by_quota_window( + state: &AppState, + windows_by_account: &HashMap, +) -> HashMap { + if windows_by_account.is_empty() { + return HashMap::new(); + } + + let mut builder = QueryBuilder::::new( + "select + account_id, + coalesce(sum(request_count), 0) as request_count, + coalesce(sum(input_tokens), 0) as input_tokens, + coalesce(sum(output_tokens), 0) as output_tokens, + coalesce(sum(cached_tokens), 0) as cached_tokens + from usage_time_buckets + where ", + ); + for (index, (account_id, window)) in windows_by_account.iter().enumerate() { + if index > 0 { + builder.push(" or "); + } + let bucket_start = china_quarter_hour_start(window.start); + builder.push("(account_id = "); + builder.push_bind(account_id); + builder.push(" and bucket_start >= "); + builder.push_bind(bucket_start.to_rfc3339()); + builder.push(" and bucket_start <= "); + builder.push_bind(window.end.to_rfc3339()); + builder.push(")"); + } + builder.push(" group by account_id"); + + let rows = builder + .build() + .fetch_all(state.services.background_tasks.accounts.pool()) + .await + .unwrap_or_default(); + rows.into_iter() + .map(|row| { + ( + row.get::("account_id"), + AccountUsageWindowStats { + request_count: row.get("request_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + }, + ) + }) + .collect() +} + +fn apply_selected_quota_window_usage( + mut usage: AdminUsageRecord, + selected_quota_windows: &HashMap, + quota_window_usage_stats: &HashMap, +) -> AdminUsageRecord { + let Some(window) = selected_quota_windows.get(&usage.account_id) else { + return usage; + }; + let stats = quota_window_usage_stats + .get(&usage.account_id) + .copied() + .unwrap_or_default(); + usage.window_request_count = stats.request_count; + usage.window_input_tokens = stats.input_tokens; + usage.window_output_tokens = stats.output_tokens; + usage.window_cached_tokens = stats.cached_tokens; + usage.window_started_at = Some(window.start); + usage.window_reset_at = Some(window.end); + usage.limit_window_seconds = Some(window.window_seconds); + usage +} + async fn apply_account_quota_window_token_usage( state: &AppState, account_id: &str, diff --git a/backend/src/runtime/tasks/quota_refresh.rs b/backend/src/runtime/tasks/quota_refresh.rs index 3402d44c6..f2280c0dc 100644 --- a/backend/src/runtime/tasks/quota_refresh.rs +++ b/backend/src/runtime/tasks/quota_refresh.rs @@ -17,7 +17,7 @@ use crate::{ use super::{ coordinator::SchedulerHandle, - periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner}, + periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task}, }; /// 主动配额刷新后台任务。 @@ -89,6 +89,7 @@ impl PeriodicTaskRunner for QuotaRefreshTask { Ok(summary) => { debug!( candidates = summary.candidates, + skipped_cooldown = summary.skipped_cooldown, skipped_recent = summary.skipped_recent, failed = summary.failed, "没有需要刷新的 quota 锁定或待验证账号" diff --git a/backend/src/upstream/accounts/quota/runtime.rs b/backend/src/upstream/accounts/quota/runtime.rs index 0fc3cce60..678aad811 100644 --- a/backend/src/upstream/accounts/quota/runtime.rs +++ b/backend/src/upstream/accounts/quota/runtime.rs @@ -2,8 +2,9 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; +use chrono::Utc; use thiserror::Error; -use tokio::time::{sleep, Instant}; +use tokio::time::{Instant, sleep}; use tracing::warn; use crate::upstream::accounts::{ @@ -22,6 +23,7 @@ use crate::upstream::{ const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60; const DEFAULT_REQUEST_SPACING_SECS: u64 = 3; +const COOLDOWN_REFRESH_GRACE_SECS: i64 = 30; /// 运行时 quota refresh 服务。 pub struct RuntimeQuotaRefreshService { @@ -109,6 +111,7 @@ impl RuntimeQuotaRefreshService { }; let min_interval = Duration::from_secs(self.min_refresh_interval_secs); let now = Instant::now(); + let now_utc = Utc::now(); let mut candidates = accounts .into_iter() @@ -121,9 +124,18 @@ impl RuntimeQuotaRefreshService { while let Some(account) = candidates.next() { summary.candidates += 1; - if last_refreshed - .get(&account.id) - .is_some_and(|last| now.duration_since(*last) < min_interval) + let cooldown_ready = match account.quota_cooldown_until { + Some(cooldown_until) if !quota_cooldown_ready(cooldown_until, now_utc) => { + summary.skipped_cooldown += 1; + continue; + } + Some(_) => true, + None => false, + }; + if !cooldown_ready + && last_refreshed + .get(&account.id) + .is_some_and(|last| now.duration_since(*last) < min_interval) { summary.skipped_recent += 1; continue; @@ -221,6 +233,12 @@ impl RuntimeQuotaRefreshService { } } +fn quota_cooldown_ready(cooldown_until: chrono::DateTime, now: chrono::DateTime) -> bool { + cooldown_until + .checked_add_signed(chrono::Duration::seconds(COOLDOWN_REFRESH_GRACE_SECS)) + .is_some_and(|ready_at| now >= ready_at) +} + /// 单次配额刷新摘要。 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct QuotaRefreshSummary { @@ -228,6 +246,8 @@ pub struct QuotaRefreshSummary { pub scanned: usize, /// 配额锁定或待验证候选账号数。 pub candidates: usize, + /// 配额冷却宽限期内跳过的账号数。 + pub skipped_cooldown: usize, /// 成功刷新账号数。 pub refreshed: usize, /// 最近刷新过而跳过的账号数。 diff --git a/backend/tests/admin/accounts/list.rs b/backend/tests/admin/accounts/list.rs index de43b378a..5f7e77e52 100644 --- a/backend/tests/admin/accounts/list.rs +++ b/backend/tests/admin/accounts/list.rs @@ -259,11 +259,11 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { assert_eq!(response.status(), StatusCode::OK); let body = response_json(response).await; let item = &body["data"]["items"][0]; - assert_eq!(item["usage"]["requestCount"], 7); - assert_eq!(item["usage"]["inputTokens"], 120_000); - assert_eq!(item["usage"]["outputTokens"], 4_000); - assert_eq!(item["usage"]["cachedTokens"], 40_000); - assert_eq!(item["usage"]["totalTokens"], 124_000); + assert_eq!(item["usage"]["requestCount"], 107); + assert_eq!(item["usage"]["inputTokens"], 1_022_000); + assert_eq!(item["usage"]["outputTokens"], 94_500); + assert_eq!(item["usage"]["cachedTokens"], 49_100); + assert_eq!(item["usage"]["totalTokens"], 1_116_500); assert_eq!(item["quota"]["windows"].as_array().unwrap().len(), 3); assert_eq!(item["quota"]["windows"][0]["labelDisplay"], "月限额"); assert_eq!(item["quota"]["windows"][0]["group"], "monthly"); @@ -277,12 +277,128 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { assert!(item["quota"]["windows"][2]["windowUsedDisplay"] .as_str() .is_some_and(|value| value.contains(" / 7.0d"))); - assert_eq!(item["usage"]["createdTokens"], 80_000); - assert_eq!(item["usage"]["readTokens"], 40_000); - assert_eq!(item["usage"]["models"].as_array().unwrap().len(), 2); - assert_eq!(item["usage"]["models"][0]["model"], "gpt-5.5"); - assert_eq!(item["usage"]["models"][0]["requestCount"], 4); - assert_eq!(item["usage"]["models"][1]["model"], "gpt-5"); + assert_eq!(item["usage"]["createdTokens"], 972_900); + assert_eq!(item["usage"]["readTokens"], 49_100); + assert_eq!(item["usage"]["models"].as_array().unwrap().len(), 4); + assert_eq!(item["usage"]["models"][0]["model"], "gpt-outside-window"); + assert_eq!(item["usage"]["models"][0]["requestCount"], 99); + assert_eq!(item["usage"]["models"][1]["model"], "gpt-5.5"); + assert_eq!(item["usage"]["models"][2]["model"], "gpt-5"); + assert_eq!(item["usage"]["models"][3]["model"], "gpt-five-hour-window"); +} + +#[tokio::test] +async fn admin_accounts_list_should_use_largest_quota_reset_interval_for_usage_stats() { + let (app, _state, pool, _dir) = + admin_accounts_test_app("admin-accounts-largest-quota-window.sqlite", 71).await; + seed_usage_account( + &pool, + UsageAccountSeed { + id: "acct_custom_window", + email: "custom-window@example.com", + label: "custom", + plan_type: "plus", + request_count: 10, + empty_response_count: 0, + input_tokens: 500, + output_tokens: 100, + cached_tokens: 50, + window_request_count: 1, + window_input_tokens: 100, + window_output_tokens: 10, + window_cached_tokens: 5, + window_started_at: "2026-07-13T19:00:00Z", + window_reset_at: "2026-07-14T00:00:00Z", + limit_window_seconds: 18_000, + last_used_at: "2026-06-30T01:10:00Z", + }, + ) + .await; + sqlx::query("update accounts set quota_json = ?, quota_fetched_at = ? where id = ?") + .bind( + json!({ + "plan_type": "plus", + "snapshots": [{ + "source": "custom", + "limit_name": "Codex", + "metered_feature": "codex", + "allowed": true, + "limit_reached": false, + "blocked": false, + "primary": { + "used_percent": 50, + "remaining_percent": 50, + "reset_at": 1783987200, + "window_minutes": 20160, + "limit_reached": false + }, + "secondary": { + "used_percent": 10, + "remaining_percent": 90, + "reset_at": 1783987200, + "window_minutes": 10080, + "limit_reached": false + } + }] + }) + .to_string(), + ) + .bind("2026-06-30T01:10:00Z") + .bind("acct_custom_window") + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind("2026-06-30T01:00:00+00:00") + .bind("acct_custom_window") + .bind("gpt-custom-14d") + .bind(5) + .bind(10_000) + .bind(2_000) + .bind(1_000) + .bind("2026-06-30T01:10:00Z") + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind("2026-06-29T23:45:00+00:00") + .bind("acct_custom_window") + .bind("gpt-before-14d") + .bind(99) + .bind(900_000) + .bind(90_000) + .bind(9_000) + .bind("2026-06-29T23:59:00Z") + .execute(&pool) + .await + .unwrap(); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts?page=1&pageSize=10") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_accounts_largest_quota_window") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + let item = &body["data"]["items"][0]; + assert_eq!(item["usage"]["requestCount"], 5); + assert_eq!(item["usage"]["inputTokens"], 10_000); + assert_eq!(item["usage"]["outputTokens"], 2_000); + assert_eq!(item["usage"]["cachedTokens"], 1_000); + assert_eq!(item["usage"]["models"].as_array().unwrap().len(), 1); + assert_eq!(item["usage"]["models"][0]["model"], "gpt-custom-14d"); } #[tokio::test] diff --git a/backend/tests/upstream/accounts/quota_refresh.rs b/backend/tests/upstream/accounts/quota_refresh.rs index a7f95dd9a..1edb3ab43 100644 --- a/backend/tests/upstream/accounts/quota_refresh.rs +++ b/backend/tests/upstream/accounts/quota_refresh.rs @@ -5,16 +5,16 @@ use codex_proxy_rs::infra::database::connect_sqlite; use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore; use codex_proxy_rs::upstream::accounts::model::AccountStatus; use codex_proxy_rs::upstream::accounts::quota::{ - quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, - quota_snapshot_reset_at, RuntimeQuotaRefreshService, + RuntimeQuotaRefreshService, quota_from_usage, quota_snapshot_limit_reached, + quota_snapshot_limit_window_seconds, quota_snapshot_reset_at, }; use codex_proxy_rs::upstream::accounts::store::{NewAccount, SqliteAccountStore}; use codex_proxy_rs::upstream::transport::CodexBackendClient; use secrecy::SecretString; -use serde_json::{json, Value}; +use serde_json::{Value, json}; use wiremock::{ - matchers::{method, path}, Mock, MockServer, ResponseTemplate, + matchers::{method, path}, }; #[test] @@ -488,6 +488,145 @@ async fn quota_refresh_service_should_refresh_quota_exhausted_accounts() { assert_eq!(stored, ("active".to_string(), 0)); } +#[tokio::test] +async fn quota_refresh_service_should_skip_recent_locked_account_before_cooldown_expires() { + let server = MockServer::start().await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir + .path() + .join("quota-refresh-recent-future-cooldown.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool.clone()); + insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await; + set_quota_cooldown_until(&pool, "acct-quota", Utc::now() + Duration::minutes(5)).await; + let codex = CodexBackendClient::new( + reqwest::Client::new(), + server.uri(), + crate::support::fingerprint::test_fingerprint(), + ); + let service = + RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800); + + let mut last_refreshed = HashMap::new(); + last_refreshed.insert("acct-quota".to_string(), tokio::time::Instant::now()); + let summary = service + .refresh_locked_accounts(&mut last_refreshed) + .await + .expect("quota refresh should succeed"); + + assert_eq!( + ( + summary.candidates, + summary.skipped_cooldown, + summary.skipped_recent, + summary.refreshed, + usage_request_count(&server).await, + ), + (1, 1, 0, 0, 0) + ); +} + +#[tokio::test] +async fn quota_refresh_service_should_skip_recent_locked_account_inside_cooldown_grace() { + let server = MockServer::start().await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir + .path() + .join("quota-refresh-recent-cooldown-grace.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool.clone()); + insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await; + set_quota_cooldown_until(&pool, "acct-quota", Utc::now() - Duration::seconds(20)).await; + let codex = CodexBackendClient::new( + reqwest::Client::new(), + server.uri(), + crate::support::fingerprint::test_fingerprint(), + ); + let service = + RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800); + + let mut last_refreshed = HashMap::new(); + last_refreshed.insert("acct-quota".to_string(), tokio::time::Instant::now()); + let summary = service + .refresh_locked_accounts(&mut last_refreshed) + .await + .expect("quota refresh should succeed"); + + assert_eq!( + ( + summary.candidates, + summary.skipped_cooldown, + summary.skipped_recent, + summary.refreshed, + usage_request_count(&server).await, + ), + (1, 1, 0, 0, 0) + ); +} + +#[tokio::test] +async fn quota_refresh_service_should_bypass_recent_skip_after_cooldown_grace() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "plan_type": "plus", + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": { + "used_percent": 8, + "reset_at": 1_800_000_000, + "limit_window_seconds": 18_000 + } + } + }))) + .expect(1) + .mount(&server) + .await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir + .path() + .join("quota-refresh-recent-expired-cooldown.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool.clone()); + insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await; + set_quota_cooldown_until(&pool, "acct-quota", Utc::now() - Duration::seconds(31)).await; + let codex = CodexBackendClient::new( + reqwest::Client::new(), + server.uri(), + crate::support::fingerprint::test_fingerprint(), + ); + let service = + RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800); + + let mut last_refreshed = HashMap::new(); + last_refreshed.insert("acct-quota".to_string(), tokio::time::Instant::now()); + let summary = service + .refresh_locked_accounts(&mut last_refreshed) + .await + .expect("quota refresh should succeed"); + + assert_eq!( + ( + summary.skipped_recent, + summary.skipped_cooldown, + summary.refreshed, + usage_request_count(&server).await, + ), + (0, 0, 1, 1) + ); +} + #[tokio::test] async fn quota_refresh_service_should_stagger_multiple_locked_account_requests() { let server = MockServer::start().await; @@ -573,6 +712,19 @@ async fn insert_quota_locked_account( .expect("quota lock should be set"); } +async fn set_quota_cooldown_until( + pool: &sqlx::SqlitePool, + id: &str, + cooldown_until: chrono::DateTime, +) { + sqlx::query("update accounts set quota_cooldown_until = ? where id = ?") + .bind(cooldown_until.to_rfc3339()) + .bind(id) + .execute(pool) + .await + .expect("quota cooldown should be set"); +} + async fn wait_for_usage_requests(server: &MockServer, expected_count: usize) { let deadline = tokio::time::Instant::now() + StdDuration::from_secs(2); loop {