codex: verify dirty quota

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-15 10:19:58 +08:00
co-authored by Codex
parent 0ecd630923
commit 8e9e89aeff
8 changed files with 408 additions and 18 deletions
+3 -4
View File
@@ -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 存储;否则将其记录为双方未使用字段。
+8
View File
@@ -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,
+21
View File
@@ -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<bool> {
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,
+1 -1
View File
@@ -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)
+11
View File
@@ -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, &quota.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 {
+359 -9
View File
@@ -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<AcquiredAccount> {
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<DirtyQuotaVerification, CodexClientError> {
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(&quota);
persist_verified_quota_with_deps(deps, account, &quota).await;
if limit_reached {
apply_verified_quota_limit_with_deps(deps, account, &quota).await;
}
Ok(DirtyQuotaVerification { limit_reached })
}
async fn fetch_dirty_account_usage_with_deps(
deps: &CodexUpstreamDependencies,
account: &Account,
request_id: &str,
) -> Result<Value, CodexClientError> {
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, &quota.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<DateTime<Utc>> {
let reset_at = quota
.pointer("/rate_limit/reset_at")
.and_then(Value::as_i64)?;
DateTime::<Utc>::from_timestamp(reset_at, 0)
}
fn quota_primary_limit_window_seconds(quota: &Value) -> Option<u64> {
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,
}
}
}
+1 -1
View File
@@ -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();
+4 -3
View File
@@ -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<String>) = sqlx::query_as(
"select quota_json, quota_limit_reached, quota_cooldown_until from accounts where id = ?",
let stored: (String, i64, Option<String>, 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<String>, Option<String>, Option<i64>) =
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 = ?",