diff --git a/docs/OPENAI_CHAIN_PARITY_REAUDIT.md b/docs/OPENAI_CHAIN_PARITY_REAUDIT.md index 7581ac22c..3dac22e12 100644 --- a/docs/OPENAI_CHAIN_PARITY_REAUDIT.md +++ b/docs/OPENAI_CHAIN_PARITY_REAUDIT.md @@ -480,6 +480,7 @@ Rust 证据： - `src/codex/accounts/pool.rs:403-441` acquire 时更新运行时 `request_count`、`window_request_count`、`last_used_at`，并在已知窗口长度但无窗口开始时间时推导窗口 reset。 - `src/codex/accounts/pool.rs:502-571` 运行时会在窗口到期时重置窗口计数、清除过期 quota cooldown，并推进下一窗口 reset。 - 本轮修正后，`src/codex/accounts/pool.rs` 在本地解除 `quota_limit_reached` 时会设置 `quota_verify_required=true`，避免本地窗口推断后直接把账号当成干净账号使用；`acquire_should_refresh_expired_cooldowns_before_selecting_account` 覆盖该标记。 +- 本轮修正后，`CodexUpstreamService::acquire_account_for_request()` 会在 dirty 账号真正进入 `/codex/responses` 前调用 `/api/codex/usage`/`/codex/usage` 校验；若 upstream 仍返回 `limit_reached`，会清理 dirty、同步 quota/window/cooldown、释放当前 slot、排除该账号并最多重试 5 个账号；`acquire_account_for_request_should_verify_dirty_quota_and_try_next_account` 覆盖该行为。passive rate-limit headers 成功更新 quota 时也会清理 dirty 标记，`v1_responses_should_passively_cache_rate_limit_headers` 覆盖该行为。 - 本轮修正后，`src/codex/gateway/transport/usage_events.rs` 的 SSE usage 提取按原版 Responses passthrough 语义优先使用 `response.completed.response.usage`，不会把 `response.created`/中间事件中的重复 usage 累加；若没有 completed usage，才回退到最后一份可见 usage。 - `tests/codex_accounts/repository.rs:230-326` 覆盖 window usage 持久化与恢复。 - `tests/codex_serving/responses_http_sse.rs:279-350` 覆盖 passive rate-limit headers 缓存 quota、窗口字段和 cooldown。 @@ -499,18 +500,16 @@ Rust 证据： - account retry decision 的失败请求持久 request_count 已对齐原版 release 语义：进入 429/402/403/401/model unsupported/Cloudflare 等账号 fallback/hold 处理时会写入一次无 token 的 request attempt；`v1_responses_should_mark_quota_exhausted_after_402_and_retry_next_account` 覆盖 402 fallback 失败账号计数。 - 非 account retry decision 的最终失败请求持久 request_count 已对齐到原版 release 语义：未进入账号 fallback/hold 的最终错误会写入一次无 token 的 request attempt；`v1_responses_should_record_request_count_when_5xx_retries_are_exhausted` 覆盖 5xx 内部重试 3 次但持久 request_count 只计 1 次。 - transport error 的最终失败请求持久 request_count 已有覆盖；`v1_responses_should_record_request_count_when_http_transport_fails` 证明 HTTP transport 失败也会按原版 release 语义写入一次无 token 的 request attempt。 -- dirty quota verification 的状态层已补齐：`quota_verify_required` 进入 schema、runtime `Account` 和启动恢复路径；本地解除 quota 限制时会把账号标成 dirty，等待请求前 `/usage` 校验。 +- dirty quota verification 已对齐原版主语义：`quota_verify_required` 进入 schema、runtime `Account` 和启动恢复路径；本地解除 quota 限制时会把账号标成 dirty；请求前会用 upstream `/usage` 校验 dirty 账号，仍限流时释放并切换账号，最多 5 次，校验失败时保留 dirty 标记并继续当前账号。 未完全对齐： - image_generation usage 未对齐。原版记录 image input/output tokens、image request success/failure 和对应窗口计数；Rust `TokenUsage` 只有 input/output/cached/total，没有解析或持久化 `tool_usage.image_gen`。 -- dirty quota verification 请求前校验尚未对齐。状态字段和本地 dirty 标记已经存在，但 Rust 还未在下一次请求前最多 5 次调用 `/usage` 验证 dirty 账号。 - Cookie domain/path 语义不同。原版 CookieJar 是 per-account 简单 map，不按 domain/path 过滤；Rust 按 domain/path 存储和匹配。对 `chatgpt.com` 主链路更严格，但不是原版逐字行为。 -- quota window reset 细节不完全一致。原版 reset 后设置 `window_counters_reset_at` 并可能标记 quotaVerifyRequired；Rust 设置 `window_started_at`，没有 `window_counters_reset_at` 字段，也不触发 quotaVerifyRequired。 +- quota window reset 细节不完全一致。原版 reset 后设置 `window_counters_reset_at`；Rust 设置 `window_started_at` 并触发 `quota_verify_required`，但没有 `window_counters_reset_at` 字段。 - passive header 文档提到 `x-codex-active-limit`，原版和 Rust 当前都没有实际解析该 header；如果 Codex Desktop 当前依赖 active-limit，这两边都需要重新确认。 缺口/后续动作： - 补 `TokenUsage` 的 image_generation 字段与数据库/运行时窗口字段，或明确 Rust 不支持 image_generation usage 统计。 -- 继续补 dirty quota verification 请求前校验：调用 usage endpoint，失败时保留 dirty 标记并限制放大次数。 - 明确 cookie domain/path 精细化是 Rust 新架构选择还是需要回到原版 per-account map；如果保留，应补 domain/path 行为测试。 - 若 OpenAI/Codex 当前返回 `x-codex-active-limit`，补解析和 quota_json 存储；否则将其记录为双方未使用字段。 diff --git a/src/codex/accounts/pool.rs b/src/codex/accounts/pool.rs index 90a6be782..e135ec455 100644 --- a/src/codex/accounts/pool.rs +++ b/src/codex/accounts/pool.rs @@ -208,6 +208,14 @@ impl AccountPool { true } + pub fn set_quota_verify_required(&mut self, account_id: &str, required: bool) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + account.quota_verify_required = required; + true + } + pub fn sync_rate_limit_window( &mut self, account_id: &str, diff --git a/src/codex/accounts/repository/quotas.rs b/src/codex/accounts/repository/quotas.rs index 7b5a3a6c9..136b48ef2 100644 --- a/src/codex/accounts/repository/quotas.rs +++ b/src/codex/accounts/repository/quotas.rs @@ -36,6 +36,13 @@ set updated_at = ? where id = ?"; +const SET_QUOTA_VERIFY_REQUIRED_SQL: &str = r" +update accounts +set + quota_verify_required = ?, + updated_at = ? +where id = ?"; + const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r" update accounts set @@ -98,6 +105,20 @@ impl AccountRepository { Ok(result.rows_affected() > 0) } + pub async fn set_quota_verify_required( + &self, + id: &str, + required: bool, + ) -> AccountRepositoryResult { + let result = sqlx::query(SET_QUOTA_VERIFY_REQUIRED_SQL) + .bind(if required { 1_i64 } else { 0 }) + .bind(Utc::now().to_rfc3339()) + .bind(id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + pub async fn set_cloudflare_cooldown_until( &self, id: &str, diff --git a/src/codex/accounts/service/quota.rs b/src/codex/accounts/service/quota.rs index 117dce8d4..69d8fcd80 100644 --- a/src/codex/accounts/service/quota.rs +++ b/src/codex/accounts/service/quota.rs @@ -150,7 +150,7 @@ fn max_optional_datetime( } } -pub(super) fn quota_from_usage(usage: &Value) -> Value { +pub(crate) fn quota_from_usage(usage: &Value) -> Value { let additional = usage .get("additional_rate_limits") .and_then(Value::as_array) diff --git a/src/codex/serving/dispatch/limits.rs b/src/codex/serving/dispatch/limits.rs index c2c9b9846..54f92cb85 100644 --- a/src/codex/serving/dispatch/limits.rs +++ b/src/codex/serving/dispatch/limits.rs @@ -17,6 +17,10 @@ pub(super) async fn apply_rate_limit_headers_with_deps( let existing_quota = existing_quota_json(deps, account_id).await; let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref()); + deps.account_pool + .lock() + .await + .set_quota_verify_required(account_id, false); if let Some(repo) = deps.account_repository.as_ref() { if let Err(error) = repo.update_quota_json(account_id, "a.to_string()).await { tracing::warn!( @@ -25,6 +29,13 @@ pub(super) async fn apply_rate_limit_headers_with_deps( "被动同步 quota 缓存失败" ); } + if let Err(error) = repo.set_quota_verify_required(account_id, false).await { + tracing::warn!( + error = %error, + account_id = %account_id, + "清理被动 quota dirty 标记失败" + ); + } } let Some(reset_at) = rate_limits.primary_reset_at() else { diff --git a/src/codex/serving/dispatch/mod.rs b/src/codex/serving/dispatch/mod.rs index 0a9994a02..8bce30180 100644 --- a/src/codex/serving/dispatch/mod.rs +++ b/src/codex/serving/dispatch/mod.rs @@ -37,6 +37,7 @@ use crate::{ model::Account, pool::{AccountAcquireRequest, AccountPool, AcquiredAccount}, repository::AccountRepository, + service::quota::quota_from_usage, }, codex::events::{ event::{EventLevel, EventLog}, @@ -114,6 +115,7 @@ struct CodexUpstreamDependencies { const MAX_UPSTREAM_5XX_RETRIES: u8 = 2; const UPSTREAM_5XX_RETRY_BASE_DELAY_MS: u64 = 1_000; +const MAX_DIRTY_QUOTA_VERIFY_ATTEMPTS: usize = 5; #[derive(Clone)] pub(crate) struct CodexUpstreamService { @@ -195,6 +197,7 @@ impl CodexUpstreamService { pub(crate) async fn acquire_account_for_request( &self, request: &CodexResponsesRequest, + request_id: &str, ) -> Option { let preferred_account_id = match request.previous_response_id.as_deref() { Some(previous_response_id) => { @@ -205,15 +208,60 @@ impl CodexUpstreamService { } None => None, }; - let mut acquire_request = AccountAcquireRequest::new(&request.model, Utc::now()); - if let Some(preferred_account_id) = preferred_account_id { - acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); + let mut excluded_account_ids = Vec::new(); + let mut verify_attempts = 0; + + loop { + let mut acquire_request = AccountAcquireRequest::new(&request.model, Utc::now()) + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + if let Some(preferred_account_id) = preferred_account_id.as_ref() { + acquire_request = + acquire_request.with_preferred_account_id(preferred_account_id.clone()); + } + let mut acquired = self + .deps + .account_pool + .lock() + .await + .acquire_with(acquire_request)?; + if !acquired.account.quota_verify_required { + return Some(acquired); + } + + match verify_dirty_quota_with_deps(&self.deps, &acquired.account, request_id).await { + Ok(verification) if verification.limit_reached => { + let account_id = acquired.account.id.clone(); + self.deps.account_pool.lock().await.release(&account_id); + excluded_account_ids.push(account_id.clone()); + verify_attempts += 1; + tracing::warn!( + account_id = %account_id, + verify_attempts, + max_verify_attempts = MAX_DIRTY_QUOTA_VERIFY_ATTEMPTS, + "dirty quota 校验确认账号仍被限流，改用备用账号" + ); + if verify_attempts >= MAX_DIRTY_QUOTA_VERIFY_ATTEMPTS { + tracing::warn!( + max_verify_attempts = MAX_DIRTY_QUOTA_VERIFY_ATTEMPTS, + "dirty quota 校验达到最大次数，停止放大 upstream usage 请求" + ); + return None; + } + } + Ok(_) => { + acquired.account.quota_verify_required = false; + return Some(acquired); + } + Err(error) => { + tracing::warn!( + error = %error, + account_id = %acquired.account.id, + "dirty quota 校验失败，保留 dirty 标记并继续当前账号" + ); + return Some(acquired); + } + } } - self.deps - .account_pool - .lock() - .await - .acquire_with(acquire_request) } pub(crate) async fn release_account(&self, account_id: &str) { @@ -396,6 +444,154 @@ async fn record_failed_request_attempt_with_deps( } } +struct DirtyQuotaVerification { + limit_reached: bool, +} + +async fn verify_dirty_quota_with_deps( + deps: &CodexUpstreamDependencies, + account: &Account, + request_id: &str, +) -> Result { + let usage = fetch_dirty_account_usage_with_deps(deps, account, request_id).await?; + let quota = quota_from_usage(&usage); + let limit_reached = quota_primary_limit_reached("a); + persist_verified_quota_with_deps(deps, account, "a).await; + if limit_reached { + apply_verified_quota_limit_with_deps(deps, account, "a).await; + } + Ok(DirtyQuotaVerification { limit_reached }) +} + +async fn fetch_dirty_account_usage_with_deps( + deps: &CodexUpstreamDependencies, + account: &Account, + request_id: &str, +) -> Result { + let request_domain = request_domain(&deps.config.api.base_url); + let cookie_header = match (deps.cookie_repository.as_ref(), request_domain.as_deref()) { + (Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(), + _ => None, + }; + let client = CodexBackendClient::new( + build_reqwest_client(deps.config.tls.force_http11)?, + deps.config.api.base_url.clone(), + deps.fingerprint.clone(), + ); + client + .fetch_usage(CodexRequestContext { + access_token: &account.access_token, + account_id: account.account_id.as_deref(), + request_id, + turn_state: None, + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: cookie_header.as_deref(), + installation_id: None, + session_id: None, + }) + .await +} + +async fn persist_verified_quota_with_deps( + deps: &CodexUpstreamDependencies, + account: &Account, + quota: &Value, +) { + deps.account_pool + .lock() + .await + .set_quota_verify_required(&account.id, false); + if let Some(repo) = deps.account_repository.as_ref() { + if let Err(error) = repo + .update_quota_json(&account.id, "a.to_string()) + .await + { + tracing::warn!( + error = %error, + account_id = %account.id, + "持久化 dirty quota 校验结果失败" + ); + } + if let Err(error) = repo.set_quota_verify_required(&account.id, false).await { + tracing::warn!( + error = %error, + account_id = %account.id, + "清理 dirty quota 持久化标记失败" + ); + } + } +} + +async fn apply_verified_quota_limit_with_deps( + deps: &CodexUpstreamDependencies, + account: &Account, + quota: &Value, +) { + let Some(reset_at) = quota_primary_reset_at(quota) else { + return; + }; + let limit_window_seconds = quota_primary_limit_window_seconds(quota); + { + let mut pool = deps.account_pool.lock().await; + pool.sync_rate_limit_window(&account.id, reset_at, limit_window_seconds); + if reset_at > Utc::now() { + pool.mark_quota_limited_until(&account.id, reset_at); + } + } + if let Some(repo) = deps.account_repository.as_ref() { + if let Err(error) = repo + .sync_rate_limit_window(&account.id, reset_at, limit_window_seconds) + .await + { + tracing::warn!( + error = %error, + account_id = %account.id, + window_reset_at = %reset_at, + "持久化 dirty quota window 失败" + ); + } + if reset_at > Utc::now() { + if let Err(error) = repo.set_quota_cooldown_until(&account.id, reset_at).await { + tracing::warn!( + error = %error, + account_id = %account.id, + cooldown_until = %reset_at, + "持久化 dirty quota cooldown 失败" + ); + } + } + } + if reset_at > Utc::now() { + deps.websocket_pool.evict_account(&account.id).await; + } +} + +fn quota_primary_limit_reached(quota: &Value) -> bool { + quota + .pointer("/rate_limit/limit_reached") + .and_then(Value::as_bool) + .unwrap_or(false) +} + +fn quota_primary_reset_at(quota: &Value) -> Option> { + let reset_at = quota + .pointer("/rate_limit/reset_at") + .and_then(Value::as_i64)?; + DateTime::::from_timestamp(reset_at, 0) +} + +fn quota_primary_limit_window_seconds(quota: &Value) -> Option { + quota + .pointer("/rate_limit/limit_window_seconds") + .and_then(Value::as_u64) + .filter(|seconds| *seconds > 0) +} + async fn apply_implicit_resume_with_deps( deps: &CodexUpstreamDependencies, request: &mut CodexResponsesRequest, @@ -1653,7 +1849,37 @@ async fn log_codex_upstream_response_with_deps( #[cfg(test)] mod tests { - use super::{append_premature_close_failed_event, jitter_request_interval_ms_with_factor}; + use std::{collections::BTreeMap, sync::Arc}; + + use serde_json::json; + use tokio::sync::Mutex; + use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, + }; + + use super::{ + append_premature_close_failed_event, jitter_request_interval_ms_with_factor, + CodexUpstreamRepositories, CodexUpstreamService, + }; + use crate::{ + codex::{ + accounts::{ + model::{Account, AccountStatus}, + pool::AccountPool, + }, + events::service::LogService, + gateway::{ + fingerprint::model::Fingerprint, + transport::{types::CodexResponsesRequest, websocket::CodexWebSocketPool}, + }, + }, + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, + ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, + TlsConfig, UsageStatsConfig, + }, + }; #[test] fn jitter_request_interval_ms_with_factor_should_match_original_bounds() { @@ -1675,4 +1901,128 @@ mod tests { let metadata = failure.unwrap().metadata(true); assert_eq!(metadata["upstreamCode"], "stream_disconnected"); } + + #[tokio::test] + async fn acquire_account_for_request_should_verify_dirty_quota_and_try_next_account() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .and(header("authorization", "Bearer dirty-access")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "plan_type": "plus", + "rate_limit": { + "allowed": true, + "limit_reached": true, + "primary_window": { + "used_percent": 100, + "reset_at": chrono::Utc::now().timestamp() + 300, + "limit_window_seconds": 300 + } + } + }))) + .mount(&server) + .await; + + let mut dirty = Account::test("acct_a", AccountStatus::Active); + dirty.access_token = "dirty-access".to_string(); + dirty.account_id = Some("chatgpt-dirty".to_string()); + dirty.quota_verify_required = true; + let mut clean = Account::test("acct_b", AccountStatus::Active); + clean.access_token = "clean-access".to_string(); + clean.account_id = Some("chatgpt-clean".to_string()); + let mut pool = AccountPool::default(); + pool.insert(dirty); + pool.insert(clean); + let service = test_upstream_service(server.uri(), pool); + let request = CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()); + + let acquired = service + .acquire_account_for_request(&request, "req_dirty_quota") + .await + .unwrap(); + + assert_eq!(acquired.account.id, "acct_b"); + } + + fn test_upstream_service(base_url: String, account_pool: AccountPool) -> CodexUpstreamService { + CodexUpstreamService::new( + Arc::new(test_app_config(base_url)), + Arc::new(Mutex::new(account_pool)), + CodexUpstreamRepositories { + account: None, + cookie: None, + session_affinity: None, + }, + LogService::new(test_logging_config(), None), + Fingerprint::default_for_tests(), + Arc::new(CodexWebSocketPool::default()), + ) + } + + fn test_app_config(base_url: String) -> AppConfig { + AppConfig { + server: ServerConfig { + host: "127.0.0.1".to_string(), + port: 0, + }, + api: ApiConfig { base_url }, + model: ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: BTreeMap::new(), + }, + auth: AuthConfig { + refresh_margin_seconds: 300, + refresh_enabled: true, + refresh_concurrency: 2, + max_concurrent_per_account: 3, + request_interval_ms: 50, + rotation_strategy: "least_used".to_string(), + tier_priority: Vec::new(), + oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), + oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), + oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), + }, + quota: QuotaConfig { + refresh_interval_minutes: 5, + warning_thresholds: QuotaWarningThresholds { + primary: vec![80, 90], + secondary: vec![80, 90], + }, + skip_exhausted: true, + }, + usage_stats: UsageStatsConfig { + history_retention_days: None, + }, + database: DatabaseConfig { + url: "sqlite::memory:".to_string(), + }, + security: SecurityConfig { + master_key_file: "data/master.key".to_string(), + api_key_pepper_file: "data/api-key-pepper.key".to_string(), + }, + tls: TlsConfig { + force_http11: false, + }, + ws_pool: Default::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: test_logging_config(), + } + } + + fn test_logging_config() -> LoggingConfig { + LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + } + } } diff --git a/src/codex/serving/responses.rs b/src/codex/serving/responses.rs index 9f09f85fc..293c0205b 100644 --- a/src/codex/serving/responses.rs +++ b/src/codex/serving/responses.rs @@ -154,7 +154,7 @@ impl ResponsesService { .await; let Some(mut acquired) = self .upstream - .acquire_account_for_request(&codex_request) + .acquire_account_for_request(&codex_request, request_id) .await else { return no_available_accounts_response().into_response(); diff --git a/tests/codex_serving/responses_http_sse.rs b/tests/codex_serving/responses_http_sse.rs index 8083f3973..196af4805 100644 --- a/tests/codex_serving/responses_http_sse.rs +++ b/tests/codex_serving/responses_http_sse.rs @@ -354,7 +354,7 @@ async fn v1_responses_should_passively_cache_rate_limit_headers() { .await; let imported = build_imported_app(server.uri()).await; sqlx::query( - r#"update accounts set quota_json = '{"credits":{"has_credits":true,"unlimited":false,"balance":12}}' where id = ?"#, + r#"update accounts set quota_json = '{"credits":{"has_credits":true,"unlimited":false,"balance":12}}', quota_verify_required = 1 where id = ?"#, ) .bind("acct_imported") .execute(&imported.pool) @@ -381,8 +381,8 @@ async fn v1_responses_should_passively_cache_rate_limit_headers() { .unwrap(); assert_eq!(response.status(), StatusCode::OK); - let stored: (String, i64, Option) = sqlx::query_as( - "select quota_json, quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let stored: (String, i64, Option, i64) = sqlx::query_as( + "select quota_json, quota_limit_reached, quota_cooldown_until, quota_verify_required from accounts where id = ?", ) .bind("acct_imported") .fetch_one(&imported.pool) @@ -394,6 +394,7 @@ async fn v1_responses_should_passively_cache_rate_limit_headers() { assert_eq!(quota["credits"]["balance"], 12); assert_eq!(stored.1, 1); assert!(stored.2.is_some()); + assert_eq!(stored.3, 0); let window: (i64, i64, i64, i64, Option, Option, Option) = sqlx::query_as( "select window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds from account_usage where account_id = ?",