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

fix: align quota window accounting

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-07-06 17:49:21 +08:00
co-authored by Codex
parent 0e952c9d83
commit 1101eee628
6 changed files with 442 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/admin/accounts/quota_view.rs
+13 -1
View File
@@ -43,6 +43,16 @@ pub(crate) struct AdminAccountQuotaUsageWindow {
pub(crate) key: String,
pub(crate) start: DateTime<Utc>,
pub(crate) end: DateTime<Utc>,
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<AdminAccountQuotaUsageWindow> {
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,
})
}
}
backend/src/admin/accounts/routes.rs
+120
View File
@@ -303,6 +303,14 @@ struct AccountModelUsageRecord {
last_used_at: Option<DateTime<Utc>>,
}
#[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, &quota_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,
&quota_window_usage_stats,
)
})
.collect::<Vec<_>>();
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<String, AdminAccountQuotaData>,
) -> HashMap<String, AdminAccountQuotaUsageWindow> {
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<AdminAccountQuotaUsageWindow> {
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<String, AdminAccountQuotaUsageWindow>,
) -> HashMap<String, AccountUsageWindowStats> {
if windows_by_account.is_empty() {
return HashMap::new();
}
let mut builder = QueryBuilder::<Sqlite>::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::<String, _>("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<String, AdminAccountQuotaUsageWindow>,
quota_window_usage_stats: &HashMap<String, AccountUsageWindowStats>,
) -> 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,
backend/src/runtime/tasks/quota_refresh.rs
+2 -1
View File
@@ -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 锁定或待验证账号"
backend/src/upstream/accounts/quota/runtime.rs
+24 -4
View File
@@ -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<Utc>, now: chrono::DateTime<Utc>) -> 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,
/// 最近刷新过而跳过的账号数。
backend/tests/admin/accounts/list.rs
+127 -11
View File
@@ -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]
backend/tests/upstream/accounts/quota_refresh.rs
+156 -4
View File
@@ -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<Utc>,
) {
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 {