From 44c33060236fac83a600dfafff2f2b357e141529 Mon Sep 17 00:00:00 2001 From: zyycn Date: Mon, 15 Jun 2026 10:47:40 +0800 Subject: [PATCH] codex: track image usage Co-authored-by: Codex --- docs/OPENAI_CHAIN_PARITY_REAUDIT.md | 9 +- src/admin/api/usage.rs | 16 ++++ src/codex/accounts/model.rs | 16 ++++ src/codex/accounts/pool.rs | 66 +++++++++++-- src/codex/accounts/repository/accounts.rs | 54 ++++++++++- src/codex/accounts/repository/usage.rs | 92 +++++++++++++++++- src/codex/accounts/service/import.rs | 8 ++ src/codex/accounts/service/pool_sync.rs | 8 ++ src/codex/gateway/transport/types.rs | 8 ++ src/codex/gateway/transport/usage_events.rs | 17 +++- src/codex/serving/chat.rs | 5 +- src/codex/serving/dispatch/fallback.rs | 7 +- src/codex/serving/dispatch/mod.rs | 70 +++++++++++--- src/codex/serving/dispatch/stream_audit.rs | 22 +++-- src/codex/serving/dispatch/usage.rs | 42 ++++++++- src/codex/serving/responses.rs | 43 +++++++-- src/platform/storage/schema.sql | 8 ++ tests/codex_accounts/pool_scheduling.rs | 15 ++- tests/codex_accounts/repository.rs | 70 +++++++++++++- tests/codex_gateway/http_client.rs | 2 + tests/codex_gateway/usage_events.rs | 46 +++++++++ tests/codex_serving/responses_http_sse.rs | 93 +++++++++++++++++++ .../http_sse/completed_image_usage.sse | 3 + tests/platform/storage_schema.rs | 29 ++++++ tests/runtime/startup.rs | 2 + 25 files changed, 697 insertions(+), 54 deletions(-) create mode 100644 tests/fixtures/responses/http_sse/completed_image_usage.sse diff --git a/docs/OPENAI_CHAIN_PARITY_REAUDIT.md b/docs/OPENAI_CHAIN_PARITY_REAUDIT.md index 3dac22e12..c9faff1ae 100644 --- a/docs/OPENAI_CHAIN_PARITY_REAUDIT.md +++ b/docs/OPENAI_CHAIN_PARITY_REAUDIT.md @@ -30,7 +30,7 @@ Rust 仓库：`/home/zyy/Codes/codex-proxy-rs` - `/v1/responses` 请求构造仍存在非 `/v1` alias route、`use_websocket:false` 本地扩展等差异；`/v1/responses/review` 强制 review subagent 和 `/v1/responses/compact` 已补齐。 - response 转换和错误恢复仍缺最终错误外壳对齐；tuple schema request conversion/reconvert、streaming premature close 合成、reasoning replay、implicit resume failure restore、`response.failed`/`error` 的 429/402/403 分类、账号 fallback、`previous_response_not_found`/unanswered function call strip-and-retry、5xx same-account retry、CF path-block 404、model unsupported fallback、401 token invalid fallback 已补齐。 - session affinity 的显式续链、implicit resume 主路径、implicit resume failure restore、reasoning replay cache、function call continuation 预检、variant_hash 基础算法、variant identity 和 OpenAI chat `user` conversation identity 已接入。 -- rate-limit/usage/cookie 持久化主路径已接近，但 image_generation usage、dirty quota verification 和失败请求持久计数仍有差异。 +- rate-limit/usage/cookie 持久化主路径已接近；dirty quota verification、失败请求持久计数和 image_generation usage 已补齐，剩余差异集中在 Cookie domain/path、quota window reset 细节和双方都未解析的 active-limit header。 明确不纳入缺口：IP 代理、VPN、本地代理探测、`HttpsProxyAgent`、账号代理池。 @@ -482,6 +482,10 @@ Rust 证据： - 本轮修正后，`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。 +- 本阶段重新核对 image_generation usage：原版 `src/types/codex-events.ts:288-298` 和 `src/routes/responses-passthrough.ts:147-156` 从 `response.tool_usage.image_gen.input_tokens/output_tokens` 提取图片工具 token，并明确这部分与 host-model `usage.input_tokens/output_tokens` 分开统计。 +- 原版 `src/routes/shared/proxy-handler-utils.ts:12-29` 只在请求声明 `tools: [{type: "image_generation"}]` 时标记 image request attempt；`image_output_tokens > 0` 记为成功，否则记为失败。`src/auth/account-registry.ts:414-480` release 时同时累计 total 与 window 的 image token、image request success/failure。 +- 本阶段修正后，`src/codex/gateway/transport/usage_events.rs` 的 `TokenUsage` 已增加 `image_input_tokens/image_output_tokens`，`extract_usage()` 和 `extract_sse_usage()` 会从 `response.tool_usage.image_gen` 读取图片工具 token，且 `total_tokens` 仍只表示 host-model token，符合原版“图片工具单独计费”的语义。 +- 本阶段修正后，`src/platform/storage/schema.sql`、`src/codex/accounts/repository/usage.rs`、`src/codex/accounts/repository/accounts.rs` 和运行时 `Account` 已加入 total/window image token 与 image request success/failure 字段；`src/codex/serving/dispatch/usage.rs` 按请求是否声明 `image_generation` tool 和 `image_output_tokens > 0` 记录成功或失败。 - `tests/codex_accounts/repository.rs:230-326` 覆盖 window usage 持久化与恢复。 - `tests/codex_serving/responses_http_sse.rs:279-350` 覆盖 passive rate-limit headers 缓存 quota、窗口字段和 cooldown。 - `tests/codex_accounts/pool_scheduling.rs:150-466` 覆盖窗口过期重置、least_used 的 quota/window/request_count 排序和 runtime request count。 @@ -501,15 +505,14 @@ Rust 证据： - 非 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；请求前会用 upstream `/usage` 校验 dirty 账号，仍限流时释放并切换账号，最多 5 次，校验失败时保留 dirty 标记并继续当前账号。 +- image_generation usage 已对齐原版主语义：`tool_usage.image_gen` token 与 host-model token 分开累计；请求声明 `tools: [{type: "image_generation"}]` 且 `image_output_tokens > 0` 记成功，否则记失败；total/window 字段都能持久化并恢复到运行时账号。覆盖测试包括 `extract_usage_reads_image_generation_tool_usage_separately`、`extract_sse_usage_should_read_completed_image_generation_tool_usage`、`sqlite_schema_should_persist_image_generation_usage_columns`、`account_repository_should_accumulate_image_generation_usage_counters`、`account_repository_should_restore_window_usage_into_runtime_pool_accounts`、`v1_responses_should_record_image_generation_usage_when_tool_succeeds`、`v1_responses_should_record_failed_image_generation_attempt_when_tool_has_no_output`。 未完全对齐： -- image_generation usage 未对齐。原版记录 image input/output tokens、image request success/failure 和对应窗口计数；Rust `TokenUsage` 只有 input/output/cached/total，没有解析或持久化 `tool_usage.image_gen`。 - Cookie domain/path 语义不同。原版 CookieJar 是 per-account 简单 map，不按 domain/path 过滤；Rust 按 domain/path 存储和匹配。对 `chatgpt.com` 主链路更严格，但不是原版逐字行为。 - 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 统计。 - 明确 cookie domain/path 精细化是 Rust 新架构选择还是需要回到原版 per-account map；如果保留，应补 domain/path 行为测试。 - 若 OpenAI/Codex 当前返回 `x-codex-active-limit`，补解析和 quota_json 存储；否则将其记录为双方未使用字段。 diff --git a/src/admin/api/usage.rs b/src/admin/api/usage.rs index ada03ef37..780b3a754 100644 --- a/src/admin/api/usage.rs +++ b/src/admin/api/usage.rs @@ -37,6 +37,10 @@ pub struct AdminUsageStatsData { pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, + pub image_input_tokens: i64, + pub image_output_tokens: i64, + pub image_request_count: i64, + pub image_request_failed_count: i64, pub last_used_at: Option, } @@ -52,6 +56,10 @@ impl From for AdminUsageStatsData { input_tokens: usage.input_tokens, output_tokens: usage.output_tokens, cached_tokens: usage.cached_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_count: usage.image_request_count, + image_request_failed_count: usage.image_request_failed_count, last_used_at: usage.last_used_at.map(|value| value.to_rfc3339()), } } @@ -66,6 +74,10 @@ pub struct AdminUsageStatsSummaryData { pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, + pub image_input_tokens: i64, + pub image_output_tokens: i64, + pub image_request_count: i64, + pub image_request_failed_count: i64, } impl From for AdminUsageStatsSummaryData { @@ -77,6 +89,10 @@ impl From for AdminUsageStatsSummaryData { input_tokens: summary.input_tokens, output_tokens: summary.output_tokens, cached_tokens: summary.cached_tokens, + image_input_tokens: summary.image_input_tokens, + image_output_tokens: summary.image_output_tokens, + image_request_count: summary.image_request_count, + image_request_failed_count: summary.image_request_failed_count, } } } diff --git a/src/codex/accounts/model.rs b/src/codex/accounts/model.rs index f29964d7a..9ecb0b82b 100644 --- a/src/codex/accounts/model.rs +++ b/src/codex/accounts/model.rs @@ -29,10 +29,18 @@ pub struct Account { pub cloudflare_cooldown_until: Option>, pub request_count: u64, pub empty_response_count: u64, + pub image_input_tokens: u64, + pub image_output_tokens: u64, + pub image_request_count: u64, + pub image_request_failed_count: u64, pub window_request_count: u64, pub window_input_tokens: u64, pub window_output_tokens: u64, pub window_cached_tokens: u64, + pub window_image_input_tokens: u64, + pub window_image_output_tokens: u64, + pub window_image_request_count: u64, + pub window_image_request_failed_count: u64, pub window_started_at: Option>, pub window_reset_at: Option>, pub limit_window_seconds: Option, @@ -59,10 +67,18 @@ impl Account { cloudflare_cooldown_until: None, request_count: 0, empty_response_count: 0, + image_input_tokens: 0, + image_output_tokens: 0, + image_request_count: 0, + image_request_failed_count: 0, window_request_count: 0, window_input_tokens: 0, window_output_tokens: 0, window_cached_tokens: 0, + window_image_input_tokens: 0, + window_image_output_tokens: 0, + window_image_request_count: 0, + window_image_request_failed_count: 0, window_started_at: None, window_reset_at: None, limit_window_seconds: None, diff --git a/src/codex/accounts/pool.rs b/src/codex/accounts/pool.rs index e135ec455..45d4fff26 100644 --- a/src/codex/accounts/pool.rs +++ b/src/codex/accounts/pool.rs @@ -86,6 +86,17 @@ pub struct AccountCapacitySummary { pub available_slots: usize, } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct AccountWindowUsageDelta { + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub image_input_tokens: u64, + pub image_output_tokens: u64, + pub image_request_succeeded: bool, + pub image_request_failed: bool, +} + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct AccountPoolStatusSummary { pub total: usize, @@ -171,10 +182,18 @@ impl AccountPool { account.last_used_at = None; account.request_count = 0; account.empty_response_count = 0; + account.image_input_tokens = 0; + account.image_output_tokens = 0; + account.image_request_count = 0; + account.image_request_failed_count = 0; account.window_request_count = 0; account.window_input_tokens = 0; account.window_output_tokens = 0; account.window_cached_tokens = 0; + account.window_image_input_tokens = 0; + account.window_image_output_tokens = 0; + account.window_image_request_count = 0; + account.window_image_request_failed_count = 0; account.window_started_at = None; account.window_reset_at = None; true @@ -228,16 +247,47 @@ impl AccountPool { pub fn record_window_token_usage( &mut self, account_id: &str, - input_tokens: u64, - output_tokens: u64, - cached_tokens: u64, + usage: AccountWindowUsageDelta, ) -> bool { let Some(account) = self.accounts.get_mut(account_id) else { return false; }; - account.window_input_tokens = account.window_input_tokens.saturating_add(input_tokens); - account.window_output_tokens = account.window_output_tokens.saturating_add(output_tokens); - account.window_cached_tokens = account.window_cached_tokens.saturating_add(cached_tokens); + account.image_input_tokens = account + .image_input_tokens + .saturating_add(usage.image_input_tokens); + account.image_output_tokens = account + .image_output_tokens + .saturating_add(usage.image_output_tokens); + if usage.image_request_succeeded { + account.image_request_count = account.image_request_count.saturating_add(1); + } + if usage.image_request_failed { + account.image_request_failed_count = + account.image_request_failed_count.saturating_add(1); + } + account.window_input_tokens = account + .window_input_tokens + .saturating_add(usage.input_tokens); + account.window_output_tokens = account + .window_output_tokens + .saturating_add(usage.output_tokens); + account.window_cached_tokens = account + .window_cached_tokens + .saturating_add(usage.cached_tokens); + account.window_image_input_tokens = account + .window_image_input_tokens + .saturating_add(usage.image_input_tokens); + account.window_image_output_tokens = account + .window_image_output_tokens + .saturating_add(usage.image_output_tokens); + if usage.image_request_succeeded { + account.window_image_request_count = + account.window_image_request_count.saturating_add(1); + } + if usage.image_request_failed { + account.window_image_request_failed_count = + account.window_image_request_failed_count.saturating_add(1); + } true } @@ -601,6 +651,10 @@ fn reset_window_counters(account: &mut Account) { account.window_input_tokens = 0; account.window_output_tokens = 0; account.window_cached_tokens = 0; + account.window_image_input_tokens = 0; + account.window_image_output_tokens = 0; + account.window_image_request_count = 0; + account.window_image_request_failed_count = 0; } fn next_window_reset_at( diff --git a/src/codex/accounts/repository/accounts.rs b/src/codex/accounts/repository/accounts.rs index 525857e34..2f485e5b4 100644 --- a/src/codex/accounts/repository/accounts.rs +++ b/src/codex/accounts/repository/accounts.rs @@ -88,11 +88,15 @@ pub struct StoredAccountMetadata { pub updated_at: DateTime, } -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Default)] pub struct UsageDelta { pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, + pub image_input_tokens: i64, + pub image_output_tokens: i64, + pub image_request_count: i64, + pub image_request_failed_count: i64, pub empty_response_count: i64, } @@ -104,10 +108,18 @@ pub struct AccountUsageRecord { pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, + pub image_input_tokens: i64, + pub image_output_tokens: i64, + pub image_request_count: i64, + pub image_request_failed_count: i64, pub window_request_count: i64, pub window_input_tokens: i64, pub window_output_tokens: i64, pub window_cached_tokens: i64, + pub window_image_input_tokens: i64, + pub window_image_output_tokens: i64, + pub window_image_request_count: i64, + pub window_image_request_failed_count: i64, pub window_started_at: Option>, pub window_reset_at: Option>, pub limit_window_seconds: Option, @@ -125,6 +137,10 @@ pub struct AccountUsageListRecord { pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, + pub image_input_tokens: i64, + pub image_output_tokens: i64, + pub image_request_count: i64, + pub image_request_failed_count: i64, pub last_used_at: Option>, } @@ -144,6 +160,10 @@ pub struct AccountUsageSummary { pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, + pub image_input_tokens: i64, + pub image_output_tokens: i64, + pub image_request_count: i64, + pub image_request_failed_count: i64, } #[derive(Clone)] @@ -329,10 +349,18 @@ select cloudflare_cooldown_until, coalesce(account_usage.request_count, 0) as usage_request_count, coalesce(account_usage.empty_response_count, 0) as usage_empty_response_count, + coalesce(account_usage.image_input_tokens, 0) as usage_image_input_tokens, + coalesce(account_usage.image_output_tokens, 0) as usage_image_output_tokens, + coalesce(account_usage.image_request_count, 0) as usage_image_request_count, + coalesce(account_usage.image_request_failed_count, 0) as usage_image_request_failed_count, coalesce(account_usage.window_request_count, 0) as usage_window_request_count, coalesce(account_usage.window_input_tokens, 0) as usage_window_input_tokens, coalesce(account_usage.window_output_tokens, 0) as usage_window_output_tokens, coalesce(account_usage.window_cached_tokens, 0) as usage_window_cached_tokens, + coalesce(account_usage.window_image_input_tokens, 0) as usage_window_image_input_tokens, + coalesce(account_usage.window_image_output_tokens, 0) as usage_window_image_output_tokens, + coalesce(account_usage.window_image_request_count, 0) as usage_window_image_request_count, + coalesce(account_usage.window_image_request_failed_count, 0) as usage_window_image_request_failed_count, account_usage.window_started_at as usage_window_started_at, account_usage.window_reset_at as usage_window_reset_at, account_usage.limit_window_seconds as usage_limit_window_seconds, @@ -539,6 +567,18 @@ impl AccountRepository { )?, request_count: row.get::("usage_request_count").max(0) as u64, empty_response_count: row.get::("usage_empty_response_count").max(0) as u64, + image_input_tokens: nonnegative_i64_to_u64( + row.get::("usage_image_input_tokens"), + ), + image_output_tokens: nonnegative_i64_to_u64( + row.get::("usage_image_output_tokens"), + ), + image_request_count: nonnegative_i64_to_u64( + row.get::("usage_image_request_count"), + ), + image_request_failed_count: nonnegative_i64_to_u64( + row.get::("usage_image_request_failed_count"), + ), window_request_count: nonnegative_i64_to_u64( row.get::("usage_window_request_count"), ), @@ -551,6 +591,18 @@ impl AccountRepository { window_cached_tokens: nonnegative_i64_to_u64( row.get::("usage_window_cached_tokens"), ), + window_image_input_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_image_input_tokens"), + ), + window_image_output_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_image_output_tokens"), + ), + window_image_request_count: nonnegative_i64_to_u64( + row.get::("usage_window_image_request_count"), + ), + window_image_request_failed_count: nonnegative_i64_to_u64( + row.get::("usage_window_image_request_failed_count"), + ), window_started_at: parse_optional_rfc3339( row.get::, _>("usage_window_started_at"), )?, diff --git a/src/codex/accounts/repository/usage.rs b/src/codex/accounts/repository/usage.rs index 9a6a098bf..6cbb8bdc8 100644 --- a/src/codex/accounts/repository/usage.rs +++ b/src/codex/accounts/repository/usage.rs @@ -17,22 +17,38 @@ insert into account_usage ( input_tokens, output_tokens, cached_tokens, + image_input_tokens, + image_output_tokens, + image_request_count, + image_request_failed_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, + window_image_input_tokens, + window_image_output_tokens, + window_image_request_count, + window_image_request_failed_count, last_used_at -) values (?, 1, ?, ?, ?, ?, 1, ?, ?, ?, ?) +) values (?, 1, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?, ?, ?) on conflict(account_id) do update set request_count = request_count + 1, empty_response_count = empty_response_count + excluded.empty_response_count, input_tokens = input_tokens + excluded.input_tokens, output_tokens = output_tokens + excluded.output_tokens, cached_tokens = cached_tokens + excluded.cached_tokens, + image_input_tokens = image_input_tokens + excluded.image_input_tokens, + image_output_tokens = image_output_tokens + excluded.image_output_tokens, + image_request_count = image_request_count + excluded.image_request_count, + image_request_failed_count = image_request_failed_count + excluded.image_request_failed_count, window_request_count = window_request_count + 1, window_input_tokens = window_input_tokens + excluded.window_input_tokens, window_output_tokens = window_output_tokens + excluded.window_output_tokens, window_cached_tokens = window_cached_tokens + excluded.window_cached_tokens, + window_image_input_tokens = window_image_input_tokens + excluded.window_image_input_tokens, + window_image_output_tokens = window_image_output_tokens + excluded.window_image_output_tokens, + window_image_request_count = window_image_request_count + excluded.window_image_request_count, + window_image_request_failed_count = window_image_request_failed_count + excluded.window_image_request_failed_count, window_started_at = case when account_usage.window_started_at is null and (account_usage.window_reset_at is not null or account_usage.limit_window_seconds is not null) @@ -49,10 +65,18 @@ select input_tokens, output_tokens, cached_tokens, + image_input_tokens, + image_output_tokens, + image_request_count, + image_request_failed_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, + window_image_input_tokens, + window_image_output_tokens, + window_image_request_count, + window_image_request_failed_count, window_started_at, window_reset_at, limit_window_seconds, @@ -74,15 +98,23 @@ insert into account_usage ( window_input_tokens, window_output_tokens, window_cached_tokens, + window_image_input_tokens, + window_image_output_tokens, + window_image_request_count, + window_image_request_failed_count, window_started_at, window_reset_at, limit_window_seconds -) values (?, 0, 0, 0, 0, ?, ?, ?) +) values (?, 0, 0, 0, 0, 0, 0, 0, 0, ?, ?, ?) on conflict(account_id) do update set window_request_count = 0, window_input_tokens = 0, window_output_tokens = 0, window_cached_tokens = 0, + window_image_input_tokens = 0, + window_image_output_tokens = 0, + window_image_request_count = 0, + window_image_request_failed_count = 0, window_started_at = excluded.window_started_at, window_reset_at = excluded.window_reset_at, limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)"; @@ -105,24 +137,40 @@ insert into account_usage ( input_tokens, output_tokens, cached_tokens, + image_input_tokens, + image_output_tokens, + image_request_count, + image_request_failed_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, + window_image_input_tokens, + window_image_output_tokens, + window_image_request_count, + window_image_request_failed_count, window_started_at, window_reset_at, last_used_at -) values (?, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null, null) +) values (?, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null, null) on conflict(account_id) do update set request_count = 0, empty_response_count = 0, input_tokens = 0, output_tokens = 0, cached_tokens = 0, + image_input_tokens = 0, + image_output_tokens = 0, + image_request_count = 0, + image_request_failed_count = 0, window_request_count = 0, window_input_tokens = 0, window_output_tokens = 0, window_cached_tokens = 0, + window_image_input_tokens = 0, + window_image_output_tokens = 0, + window_image_request_count = 0, + window_image_request_failed_count = 0, window_started_at = null, window_reset_at = null, last_used_at = null"; @@ -138,6 +186,10 @@ select account_usage.input_tokens, account_usage.output_tokens, account_usage.cached_tokens, + account_usage.image_input_tokens, + account_usage.image_output_tokens, + account_usage.image_request_count, + account_usage.image_request_failed_count, account_usage.last_used_at from account_usage left join accounts on accounts.id = account_usage.account_id @@ -160,6 +212,10 @@ select account_usage.input_tokens, account_usage.output_tokens, account_usage.cached_tokens, + account_usage.image_input_tokens, + account_usage.image_output_tokens, + account_usage.image_request_count, + account_usage.image_request_failed_count, account_usage.last_used_at from account_usage left join accounts on accounts.id = account_usage.account_id @@ -173,7 +229,11 @@ select coalesce(sum(empty_response_count), 0) as empty_response_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 + coalesce(sum(cached_tokens), 0) as cached_tokens, + coalesce(sum(image_input_tokens), 0) as image_input_tokens, + coalesce(sum(image_output_tokens), 0) as image_output_tokens, + coalesce(sum(image_request_count), 0) as image_request_count, + coalesce(sum(image_request_failed_count), 0) as image_request_failed_count from account_usage"; impl AccountRepository { @@ -189,9 +249,17 @@ impl AccountRepository { .bind(usage.input_tokens) .bind(usage.output_tokens) .bind(usage.cached_tokens) + .bind(usage.image_input_tokens) + .bind(usage.image_output_tokens) + .bind(usage.image_request_count) + .bind(usage.image_request_failed_count) .bind(usage.input_tokens) .bind(usage.output_tokens) .bind(usage.cached_tokens) + .bind(usage.image_input_tokens) + .bind(usage.image_output_tokens) + .bind(usage.image_request_count) + .bind(usage.image_request_failed_count) .bind(now) .execute(&self.pool) .await?; @@ -321,6 +389,10 @@ impl AccountUsageRepository { input_tokens: row.get("input_tokens"), output_tokens: row.get("output_tokens"), cached_tokens: row.get("cached_tokens"), + image_input_tokens: row.get("image_input_tokens"), + image_output_tokens: row.get("image_output_tokens"), + image_request_count: row.get("image_request_count"), + image_request_failed_count: row.get("image_request_failed_count"), }) } @@ -341,10 +413,18 @@ fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> AccountRepositoryResult, _>("window_started_at"), )?, @@ -367,6 +447,10 @@ fn usage_list_from_row( input_tokens: row.get("input_tokens"), output_tokens: row.get("output_tokens"), cached_tokens: row.get("cached_tokens"), + image_input_tokens: row.get("image_input_tokens"), + image_output_tokens: row.get("image_output_tokens"), + image_request_count: row.get("image_request_count"), + image_request_failed_count: row.get("image_request_failed_count"), last_used_at: parse_optional_rfc3339(row.get::, _>("last_used_at"))?, }) } diff --git a/src/codex/accounts/service/import.rs b/src/codex/accounts/service/import.rs index 204e19cd6..1e9d9116a 100644 --- a/src/codex/accounts/service/import.rs +++ b/src/codex/accounts/service/import.rs @@ -201,10 +201,18 @@ impl AccountService { cloudflare_cooldown_until: None, request_count: 0, empty_response_count: 0, + image_input_tokens: 0, + image_output_tokens: 0, + image_request_count: 0, + image_request_failed_count: 0, window_request_count: 0, window_input_tokens: 0, window_output_tokens: 0, window_cached_tokens: 0, + window_image_input_tokens: 0, + window_image_output_tokens: 0, + window_image_request_count: 0, + window_image_request_failed_count: 0, window_started_at: None, window_reset_at: None, limit_window_seconds: None, diff --git a/src/codex/accounts/service/pool_sync.rs b/src/codex/accounts/service/pool_sync.rs index e074881e9..465386663 100644 --- a/src/codex/accounts/service/pool_sync.rs +++ b/src/codex/accounts/service/pool_sync.rs @@ -57,10 +57,18 @@ pub(super) fn pool_account_from_stored(account: StoredAccount) -> Account { cloudflare_cooldown_until: None, request_count: 0, empty_response_count: 0, + image_input_tokens: 0, + image_output_tokens: 0, + image_request_count: 0, + image_request_failed_count: 0, window_request_count: 0, window_input_tokens: 0, window_output_tokens: 0, window_cached_tokens: 0, + window_image_input_tokens: 0, + window_image_output_tokens: 0, + window_image_request_count: 0, + window_image_request_failed_count: 0, window_started_at: None, window_reset_at: None, limit_window_seconds: None, diff --git a/src/codex/gateway/transport/types.rs b/src/codex/gateway/transport/types.rs index c1a293fb2..e57e04a3c 100644 --- a/src/codex/gateway/transport/types.rs +++ b/src/codex/gateway/transport/types.rs @@ -93,6 +93,14 @@ impl CodexResponsesRequest { parent_thread_id: None, } } + + pub fn expects_image_generation(&self) -> bool { + self.tools.as_deref().is_some_and(|tools| { + tools + .iter() + .any(|tool| tool.get("type").and_then(Value::as_str) == Some("image_generation")) + }) + } } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/codex/gateway/transport/usage_events.rs b/src/codex/gateway/transport/usage_events.rs index 0dfc42391..7989d1137 100644 --- a/src/codex/gateway/transport/usage_events.rs +++ b/src/codex/gateway/transport/usage_events.rs @@ -9,6 +9,8 @@ pub struct TokenUsage { pub input_tokens: u64, pub output_tokens: u64, pub cached_tokens: u64, + pub image_input_tokens: u64, + pub image_output_tokens: u64, pub total_tokens: u64, } @@ -18,6 +20,8 @@ impl TokenUsage { input_tokens, output_tokens, cached_tokens, + image_input_tokens: 0, + image_output_tokens: 0, total_tokens: input_tokens + output_tokens, } } @@ -27,6 +31,8 @@ impl TokenUsage { input_tokens: self.input_tokens + other.input_tokens, output_tokens: self.output_tokens + other.output_tokens, cached_tokens: self.cached_tokens + other.cached_tokens, + image_input_tokens: self.image_input_tokens + other.image_input_tokens, + image_output_tokens: self.image_output_tokens + other.image_output_tokens, total_tokens: self.total_tokens + other.total_tokens, } } @@ -48,6 +54,11 @@ pub fn extract_usage(body: &Value) -> Option { .or_else(|| nested_number_field(usage, &["prompt_tokens_details", "cached_tokens"])) .or_else(|| number_field(usage, "cached_tokens")) .unwrap_or_default(); + let image_input_tokens = + nested_number_field(body, &["tool_usage", "image_gen", "input_tokens"]).unwrap_or_default(); + let image_output_tokens = + nested_number_field(body, &["tool_usage", "image_gen", "output_tokens"]) + .unwrap_or_default(); let total_tokens = number_field(usage, "total_tokens").unwrap_or(input_tokens + output_tokens); let has_usage = [ @@ -61,11 +72,15 @@ pub fn extract_usage(body: &Value) -> Option { .iter() .any(|field| usage.get(*field).is_some()) || usage.get("input_tokens_details").is_some() - || usage.get("prompt_tokens_details").is_some(); + || usage.get("prompt_tokens_details").is_some() + || image_input_tokens > 0 + || image_output_tokens > 0; has_usage.then_some(TokenUsage { input_tokens, output_tokens, cached_tokens, + image_input_tokens, + image_output_tokens, total_tokens, }) } diff --git a/src/codex/serving/chat.rs b/src/codex/serving/chat.rs index 4a2c5b81f..1b88cd121 100644 --- a/src/codex/serving/chat.rs +++ b/src/codex/serving/chat.rs @@ -143,6 +143,7 @@ impl ChatService { retry, &codex_request.model, &mut excluded_account_ids, + false, ) .await; self.upstream @@ -183,7 +184,7 @@ impl ChatService { } if self .upstream - .record_request_attempt(&acquired.account.id) + .record_request_attempt(&acquired.account.id, false) .await .is_err() { @@ -229,7 +230,7 @@ impl ChatService { if let Some(usage) = response.usage { if self .upstream - .record_usage(&acquired.account.id, usage) + .record_usage(&acquired.account.id, usage, false) .await .is_err() { diff --git a/src/codex/serving/dispatch/fallback.rs b/src/codex/serving/dispatch/fallback.rs index a6ae23f3c..62b2edda2 100644 --- a/src/codex/serving/dispatch/fallback.rs +++ b/src/codex/serving/dispatch/fallback.rs @@ -381,8 +381,9 @@ pub(super) async fn apply_upstream_retry_and_acquire_fallback_with_deps( retry: UpstreamAccountRetry, model: &str, excluded_account_ids: &mut Vec, + image_generation_requested: bool, ) -> Option { - apply_upstream_account_retry_with_deps(deps, account, retry).await; + apply_upstream_account_retry_with_deps(deps, account, retry, image_generation_requested).await; excluded_account_ids.push(account.id.clone()); deps.account_pool.lock().await.acquire_with( AccountAcquireRequest::new(model, Utc::now()) @@ -394,8 +395,10 @@ pub(super) async fn apply_upstream_account_retry_with_deps( deps: &CodexUpstreamDependencies, account: &Account, retry: UpstreamAccountRetry, + image_generation_requested: bool, ) { - if let Err(error) = record_request_attempt(deps, &account.id).await { + if let Err(error) = record_request_attempt(deps, &account.id, image_generation_requested).await + { tracing::warn!( error = ?error, account_id = %account.id, diff --git a/src/codex/serving/dispatch/mod.rs b/src/codex/serving/dispatch/mod.rs index 8bce30180..717482d0c 100644 --- a/src/codex/serving/dispatch/mod.rs +++ b/src/codex/serving/dispatch/mod.rs @@ -308,6 +308,7 @@ impl CodexUpstreamService { retry: UpstreamAccountRetry, model: &str, excluded_account_ids: &mut Vec, + image_generation_requested: bool, ) -> Option { apply_upstream_retry_and_acquire_fallback_with_deps( &self.deps, @@ -315,6 +316,7 @@ impl CodexUpstreamService { retry, model, excluded_account_ids, + image_generation_requested, ) .await } @@ -323,8 +325,19 @@ impl CodexUpstreamService { fallback_exhausted_message_with_deps(&self.deps, message).await } - pub(crate) async fn apply_account_retry(&self, account: &Account, retry: UpstreamAccountRetry) { - apply_upstream_account_retry_with_deps(&self.deps, account, retry).await; + pub(crate) async fn apply_account_retry( + &self, + account: &Account, + retry: UpstreamAccountRetry, + image_generation_requested: bool, + ) { + apply_upstream_account_retry_with_deps( + &self.deps, + account, + retry, + image_generation_requested, + ) + .await; } pub(crate) async fn responses_stream( @@ -360,16 +373,29 @@ impl CodexUpstreamService { persist_upstream_cookies_with_deps(&self.deps, account_id, set_cookie_headers).await } - pub(crate) async fn record_usage(&self, account_id: &str, usage: TokenUsage) -> Result<(), ()> { - record_usage_with_deps(&self.deps, account_id, usage).await + pub(crate) async fn record_usage( + &self, + account_id: &str, + usage: TokenUsage, + image_generation_requested: bool, + ) -> Result<(), ()> { + record_usage_with_deps(&self.deps, account_id, usage, image_generation_requested).await } - pub(crate) async fn record_empty_response(&self, account_id: &str) -> Result<(), ()> { - record_empty_response_with_deps(&self.deps, account_id).await + pub(crate) async fn record_empty_response( + &self, + account_id: &str, + image_generation_requested: bool, + ) -> Result<(), ()> { + record_empty_response_with_deps(&self.deps, account_id, image_generation_requested).await } - pub(crate) async fn record_request_attempt(&self, account_id: &str) -> Result<(), ()> { - record_request_attempt(&self.deps, account_id).await + pub(crate) async fn record_request_attempt( + &self, + account_id: &str, + image_generation_requested: bool, + ) -> Result<(), ()> { + record_request_attempt(&self.deps, account_id, image_generation_requested).await } pub(crate) async fn record_response_affinity( @@ -434,8 +460,9 @@ async fn fallback_exhausted_message_with_deps( async fn record_failed_request_attempt_with_deps( deps: &CodexUpstreamDependencies, account_id: &str, + image_generation_requested: bool, ) { - if let Err(error) = record_request_attempt(deps, account_id).await { + if let Err(error) = record_request_attempt(deps, account_id, image_generation_requested).await { tracing::warn!( error = ?error, account_id = %account_id, @@ -794,6 +821,7 @@ async fn responses_http_sse_stream( retry, &request.model, &mut excluded_account_ids, + request.expects_image_generation(), ) .await; log_codex_upstream_response_with_deps( @@ -829,7 +857,12 @@ async fn responses_http_sse_stream( .await; return error_response.into_response(); } - record_failed_request_attempt_with_deps(&deps, &acquired.account.id).await; + record_failed_request_attempt_with_deps( + &deps, + &acquired.account.id, + request.expects_image_generation(), + ) + .await; let error_response = codex_client_error_response(error); log_codex_upstream_response_with_deps( &deps, @@ -1541,8 +1574,13 @@ async fn responses_websocket_stream( && retry.preserve_history_account_affinity() { // previous_response_id 的历史由上游账号持有，换账号会静默丢失会话上下文。 - apply_upstream_account_retry_with_deps(&deps, &acquired.account, retry) - .await; + apply_upstream_account_retry_with_deps( + &deps, + &acquired.account, + retry, + request.expects_image_generation(), + ) + .await; log_codex_upstream_response_with_deps( &deps, &log_context, @@ -1559,6 +1597,7 @@ async fn responses_websocket_stream( retry, &request.model, &mut excluded_account_ids, + request.expects_image_generation(), ) .await; log_codex_upstream_response_with_deps( @@ -1596,7 +1635,12 @@ async fn responses_websocket_stream( return error_response.into_response(); } } - record_failed_request_attempt_with_deps(&deps, &acquired.account.id).await; + record_failed_request_attempt_with_deps( + &deps, + &acquired.account.id, + request.expects_image_generation(), + ) + .await; let error_response = codex_client_error_response(error); log_codex_upstream_response_with_deps( &deps, diff --git a/src/codex/serving/dispatch/stream_audit.rs b/src/codex/serving/dispatch/stream_audit.rs index c8954be7f..7d3d6addf 100644 --- a/src/codex/serving/dispatch/stream_audit.rs +++ b/src/codex/serving/dispatch/stream_audit.rs @@ -69,9 +69,14 @@ impl StreamAudit { let mut metadata = match extract_sse_usage(&body) { Ok(usage) => { if let Some(usage) = usage { - if record_usage_with_deps(&self.deps, &self.context.account_id, usage) - .await - .is_err() + if record_usage_with_deps( + &self.deps, + &self.context.account_id, + usage, + self.request.expects_image_generation(), + ) + .await + .is_err() { level = EventLevel::Warn; message = "v1 responses stream 已完成但 usage 存储失败"; @@ -196,9 +201,14 @@ impl WebSocketStreamAudit { let mut message = "v1 responses WebSocket stream 已完成"; let mut metadata = match usage_result { Ok(Some(usage)) => { - if record_usage_with_deps(&self.deps, &self.context.account_id, usage) - .await - .is_err() + if record_usage_with_deps( + &self.deps, + &self.context.account_id, + usage, + self.request.expects_image_generation(), + ) + .await + .is_err() { level = EventLevel::Warn; message = "v1 responses WebSocket stream 已完成但 usage 存储失败"; diff --git a/src/codex/serving/dispatch/usage.rs b/src/codex/serving/dispatch/usage.rs index 7d2105c58..145db47b7 100644 --- a/src/codex/serving/dispatch/usage.rs +++ b/src/codex/serving/dispatch/usage.rs @@ -1,5 +1,6 @@ -use crate::{ - codex::accounts::repository::UsageDelta, codex::gateway::transport::usage_events::TokenUsage, +use crate::codex::{ + accounts::{pool::AccountWindowUsageDelta, repository::UsageDelta}, + gateway::transport::usage_events::TokenUsage, }; use super::CodexUpstreamDependencies; @@ -8,12 +9,21 @@ pub(super) async fn record_usage_with_deps( deps: &CodexUpstreamDependencies, account_id: &str, usage: TokenUsage, + image_generation_requested: bool, ) -> Result<(), ()> { + let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0; + let image_request_failed = image_generation_requested && !image_request_succeeded; deps.account_pool.lock().await.record_window_token_usage( account_id, - usage.input_tokens, - usage.output_tokens, - usage.cached_tokens, + AccountWindowUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_succeeded, + image_request_failed, + }, ); let Some(repo) = deps.account_repository.as_ref() else { return Ok(()); @@ -24,6 +34,10 @@ pub(super) async fn record_usage_with_deps( input_tokens: u64_to_i64_saturating(usage.input_tokens), output_tokens: u64_to_i64_saturating(usage.output_tokens), cached_tokens: u64_to_i64_saturating(usage.cached_tokens), + image_input_tokens: u64_to_i64_saturating(usage.image_input_tokens), + image_output_tokens: u64_to_i64_saturating(usage.image_output_tokens), + image_request_count: bool_to_i64(image_request_succeeded), + image_request_failed_count: bool_to_i64(image_request_failed), empty_response_count: 0, }, ) @@ -34,6 +48,7 @@ pub(super) async fn record_usage_with_deps( pub(super) async fn record_request_attempt( deps: &CodexUpstreamDependencies, account_id: &str, + image_generation_requested: bool, ) -> Result<(), ()> { let Some(repo) = deps.account_repository.as_ref() else { return Ok(()); @@ -44,6 +59,10 @@ pub(super) async fn record_request_attempt( input_tokens: 0, output_tokens: 0, cached_tokens: 0, + image_input_tokens: 0, + image_output_tokens: 0, + image_request_count: 0, + image_request_failed_count: bool_to_i64(image_generation_requested), empty_response_count: 0, }, ) @@ -54,6 +73,7 @@ pub(super) async fn record_request_attempt( pub(super) async fn record_empty_response_with_deps( deps: &CodexUpstreamDependencies, account_id: &str, + image_generation_requested: bool, ) -> Result<(), ()> { let Some(repo) = deps.account_repository.as_ref() else { return Ok(()); @@ -64,6 +84,10 @@ pub(super) async fn record_empty_response_with_deps( input_tokens: 0, output_tokens: 0, cached_tokens: 0, + image_input_tokens: 0, + image_output_tokens: 0, + image_request_count: 0, + image_request_failed_count: bool_to_i64(image_generation_requested), empty_response_count: 1, }, ) @@ -74,3 +98,11 @@ pub(super) async fn record_empty_response_with_deps( fn u64_to_i64_saturating(value: u64) -> i64 { value.min(i64::MAX as u64) as i64 } + +fn bool_to_i64(value: bool) -> i64 { + if value { + 1 + } else { + 0 + } +} diff --git a/src/codex/serving/responses.rs b/src/codex/serving/responses.rs index 293c0205b..f55d701b7 100644 --- a/src/codex/serving/responses.rs +++ b/src/codex/serving/responses.rs @@ -204,7 +204,10 @@ impl ResponsesService { Ok(CollectedResponse::Empty) => { if self .upstream - .record_empty_response(&acquired.account.id) + .record_empty_response( + &acquired.account.id, + codex_request.expects_image_generation(), + ) .await .is_err() { @@ -292,7 +295,11 @@ impl ResponsesService { && retry.preserve_history_account_affinity() { self.upstream - .apply_account_retry(&acquired.account, retry) + .apply_account_retry( + &acquired.account, + retry, + codex_request.expects_image_generation(), + ) .await; self.upstream .log_response( @@ -311,6 +318,7 @@ impl ResponsesService { retry, &codex_request.model, &mut excluded_account_ids, + codex_request.expects_image_generation(), ) .await; self.upstream @@ -355,7 +363,10 @@ impl ResponsesService { } if self .upstream - .record_request_attempt(&acquired.account.id) + .record_request_attempt( + &acquired.account.id, + codex_request.expects_image_generation(), + ) .await .is_err() { @@ -410,7 +421,11 @@ impl ResponsesService { { // previous_response_id history is account-affine upstream. self.upstream - .apply_account_retry(&acquired.account, retry) + .apply_account_retry( + &acquired.account, + retry, + codex_request.expects_image_generation(), + ) .await; self.upstream .log_response( @@ -429,6 +444,7 @@ impl ResponsesService { retry, &codex_request.model, &mut excluded_account_ids, + codex_request.expects_image_generation(), ) .await; self.upstream @@ -471,7 +487,10 @@ impl ResponsesService { } if self .upstream - .record_request_attempt(&acquired.account.id) + .record_request_attempt( + &acquired.account.id, + codex_request.expects_image_generation(), + ) .await .is_err() { @@ -526,7 +545,11 @@ impl ResponsesService { if let Some(usage) = response.usage { if self .upstream - .record_usage(&acquired.account.id, usage) + .record_usage( + &acquired.account.id, + usage, + codex_request.expects_image_generation(), + ) .await .is_err() { @@ -575,7 +598,10 @@ impl ResponsesService { Ok(CollectedResponse::Empty) => { if self .upstream - .record_empty_response(&acquired.account.id) + .record_empty_response( + &acquired.account.id, + codex_request.expects_image_generation(), + ) .await .is_err() { @@ -782,6 +808,7 @@ impl ResponsesService { retry, &compact_request.model, &mut excluded_account_ids, + false, ) .await; self.upstream @@ -824,7 +851,7 @@ impl ResponsesService { } if self .upstream - .record_request_attempt(&acquired.account.id) + .record_request_attempt(&acquired.account.id, false) .await .is_err() { diff --git a/src/platform/storage/schema.sql b/src/platform/storage/schema.sql index fbdfb5823..6d0422e1b 100644 --- a/src/platform/storage/schema.sql +++ b/src/platform/storage/schema.sql @@ -87,10 +87,18 @@ create table if not exists account_usage ( input_tokens integer not null default 0 check (input_tokens >= 0), output_tokens integer not null default 0 check (output_tokens >= 0), cached_tokens integer not null default 0 check (cached_tokens >= 0), + image_input_tokens integer not null default 0 check (image_input_tokens >= 0), + image_output_tokens integer not null default 0 check (image_output_tokens >= 0), + image_request_count integer not null default 0 check (image_request_count >= 0), + image_request_failed_count integer not null default 0 check (image_request_failed_count >= 0), window_request_count integer not null default 0 check (window_request_count >= 0), window_input_tokens integer not null default 0 check (window_input_tokens >= 0), window_output_tokens integer not null default 0 check (window_output_tokens >= 0), window_cached_tokens integer not null default 0 check (window_cached_tokens >= 0), + window_image_input_tokens integer not null default 0 check (window_image_input_tokens >= 0), + window_image_output_tokens integer not null default 0 check (window_image_output_tokens >= 0), + window_image_request_count integer not null default 0 check (window_image_request_count >= 0), + window_image_request_failed_count integer not null default 0 check (window_image_request_failed_count >= 0), window_started_at text, window_reset_at text, limit_window_seconds integer check (limit_window_seconds is null or limit_window_seconds > 0), diff --git a/tests/codex_accounts/pool_scheduling.rs b/tests/codex_accounts/pool_scheduling.rs index d7be7e511..f741ecad0 100644 --- a/tests/codex_accounts/pool_scheduling.rs +++ b/tests/codex_accounts/pool_scheduling.rs @@ -3,7 +3,10 @@ use std::collections::BTreeMap; use chrono::{Duration, TimeZone, Utc}; use codex_proxy_rs::codex::accounts::{ model::{Account, AccountStatus}, - pool::{AccountAcquireRequest, AccountPool, AccountPoolOptions, RotationStrategy}, + pool::{ + AccountAcquireRequest, AccountPool, AccountPoolOptions, AccountWindowUsageDelta, + RotationStrategy, + }, }; use crate::support::admin_accounts::test_jwt; @@ -473,7 +476,15 @@ fn record_window_token_usage_should_accumulate_runtime_window_tokens() { let mut pool = AccountPool::default(); pool.insert(Account::test("acct_a", AccountStatus::Active)); - pool.record_window_token_usage("acct_a", 7, 4, 2); + pool.record_window_token_usage( + "acct_a", + AccountWindowUsageDelta { + input_tokens: 7, + output_tokens: 4, + cached_tokens: 2, + ..AccountWindowUsageDelta::default() + }, + ); let acquired = pool .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); diff --git a/tests/codex_accounts/repository.rs b/tests/codex_accounts/repository.rs index e91dd58d6..ce9a742ff 100644 --- a/tests/codex_accounts/repository.rs +++ b/tests/codex_accounts/repository.rs @@ -168,6 +168,7 @@ async fn account_repository_should_accumulate_usage_counters() { output_tokens: 4, cached_tokens: 3, empty_response_count: 0, + ..UsageDelta::default() }, ) .await @@ -179,6 +180,7 @@ async fn account_repository_should_accumulate_usage_counters() { output_tokens: 2, cached_tokens: 1, empty_response_count: 1, + ..UsageDelta::default() }, ) .await @@ -223,6 +225,7 @@ async fn account_repository_should_accumulate_usage_window_counters() { output_tokens: 4, cached_tokens: 3, empty_response_count: 0, + ..UsageDelta::default() }, ) .await @@ -234,6 +237,7 @@ async fn account_repository_should_accumulate_usage_window_counters() { output_tokens: 2, cached_tokens: 1, empty_response_count: 0, + ..UsageDelta::default() }, ) .await @@ -250,6 +254,61 @@ async fn account_repository_should_accumulate_usage_window_counters() { assert_eq!(usage, (2, 18, 6, 4)); } +#[tokio::test] +async fn account_repository_should_accumulate_image_generation_usage_counters() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("accounts.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let repo = AccountRepository::new(pool.clone(), SecretBox::new([35u8; 32])); + + repo.insert(NewAccount { + id: "acct_a".to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .unwrap(); + + repo.record_usage( + "acct_a", + UsageDelta { + image_input_tokens: 31, + image_output_tokens: 9, + image_request_count: 1, + ..UsageDelta::default() + }, + ) + .await + .unwrap(); + repo.record_usage( + "acct_a", + UsageDelta { + image_request_failed_count: 1, + ..UsageDelta::default() + }, + ) + .await + .unwrap(); + + let usage = repo.get_usage("acct_a").await.unwrap().unwrap(); + assert_eq!(usage.image_input_tokens, 31); + assert_eq!(usage.image_output_tokens, 9); + assert_eq!(usage.image_request_count, 1); + assert_eq!(usage.image_request_failed_count, 1); + assert_eq!(usage.window_image_input_tokens, 31); + assert_eq!(usage.window_image_output_tokens, 9); + assert_eq!(usage.window_image_request_count, 1); + assert_eq!(usage.window_image_request_failed_count, 1); +} + #[tokio::test] async fn account_repository_should_load_usage_count_into_runtime_pool_accounts() { let dir = tempfile::tempdir().unwrap(); @@ -279,6 +338,7 @@ async fn account_repository_should_load_usage_count_into_runtime_pool_accounts() output_tokens: 2, cached_tokens: 3, empty_response_count: 0, + ..UsageDelta::default() }, ) .await @@ -302,7 +362,7 @@ async fn account_repository_should_restore_window_usage_into_runtime_pool_accoun seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; sqlx::query( - "insert into account_usage (account_id, request_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds) values (?, 9, 7, 11, 13, 17, ?, ?, 300)", + "insert into account_usage (account_id, request_count, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count, window_started_at, window_reset_at, limit_window_seconds) values (?, 9, 31, 9, 2, 1, 7, 11, 13, 17, 19, 5, 1, 1, ?, ?, 300)", ) .bind("acct_a") .bind(window_started_at.to_rfc3339()) @@ -314,9 +374,17 @@ async fn account_repository_should_restore_window_usage_into_runtime_pool_accoun let account = repo.list_pool_accounts().await.unwrap().remove(0); assert_eq!(account.window_request_count, 7); + assert_eq!(account.image_input_tokens, 31); + assert_eq!(account.image_output_tokens, 9); + assert_eq!(account.image_request_count, 2); + assert_eq!(account.image_request_failed_count, 1); assert_eq!(account.window_input_tokens, 11); assert_eq!(account.window_output_tokens, 13); assert_eq!(account.window_cached_tokens, 17); + assert_eq!(account.window_image_input_tokens, 19); + assert_eq!(account.window_image_output_tokens, 5); + assert_eq!(account.window_image_request_count, 1); + assert_eq!(account.window_image_request_failed_count, 1); assert_eq!(account.window_started_at, Some(window_started_at)); assert_eq!(account.window_reset_at, Some(window_reset_at)); assert_eq!(account.limit_window_seconds, Some(300)); diff --git a/tests/codex_gateway/http_client.rs b/tests/codex_gateway/http_client.rs index adb5a7999..9b9322ce2 100644 --- a/tests/codex_gateway/http_client.rs +++ b/tests/codex_gateway/http_client.rs @@ -90,6 +90,8 @@ async fn codex_backend_client_should_send_desktop_headers_and_capture_response_m input_tokens: 2, output_tokens: 3, cached_tokens: 1, + image_input_tokens: 0, + image_output_tokens: 0, total_tokens: 5, }) ); diff --git a/tests/codex_gateway/usage_events.rs b/tests/codex_gateway/usage_events.rs index 769800e78..c3701e6e4 100644 --- a/tests/codex_gateway/usage_events.rs +++ b/tests/codex_gateway/usage_events.rs @@ -113,11 +113,40 @@ fn extract_usage_reads_codex_token_usage_shape() { input_tokens: 12, output_tokens: 5, cached_tokens: 3, + image_input_tokens: 0, + image_output_tokens: 0, total_tokens: 17, } ); } +#[test] +fn extract_usage_reads_image_generation_tool_usage_separately() { + let body = json!({ + "usage": { + "input_tokens": 12, + "output_tokens": 5, + "input_tokens_details": { + "cached_tokens": 3 + } + }, + "tool_usage": { + "image_gen": { + "input_tokens": 31, + "output_tokens": 9 + } + } + }); + + let usage = extract_usage(&body).unwrap(); + let serialized = serde_json::to_value(usage).unwrap(); + + assert_eq!(serialized["imageInputTokens"], 31); + assert_eq!(serialized["imageOutputTokens"], 9); + assert_eq!(serialized["inputTokens"], 12); + assert_eq!(serialized["outputTokens"], 5); +} + #[test] fn extract_usage_reads_openai_token_usage_shape_and_merges() { let first = extract_usage(&json!({ @@ -134,6 +163,8 @@ fn extract_usage_reads_openai_token_usage_shape_and_merges() { input_tokens: 1, output_tokens: 2, cached_tokens: 0, + image_input_tokens: 0, + image_output_tokens: 0, total_tokens: 3, }; @@ -143,6 +174,8 @@ fn extract_usage_reads_openai_token_usage_shape_and_merges() { input_tokens: 9, output_tokens: 6, cached_tokens: 2, + image_input_tokens: 0, + image_output_tokens: 0, total_tokens: 15, } ); @@ -165,7 +198,20 @@ fn extract_sse_usage_should_use_completed_response_usage_without_merging_earlier input_tokens: 3, output_tokens: 5, cached_tokens: 1, + image_input_tokens: 0, + image_output_tokens: 0, total_tokens: 8, } ); } + +#[test] +fn extract_sse_usage_should_read_completed_image_generation_tool_usage() { + let body = include_str!("../fixtures/responses/http_sse/completed_image_usage.sse"); + + let usage = extract_sse_usage(body).unwrap().unwrap(); + let serialized = serde_json::to_value(usage).unwrap(); + + assert_eq!(serialized["imageInputTokens"], 31); + assert_eq!(serialized["imageOutputTokens"], 9); +} diff --git a/tests/codex_serving/responses_http_sse.rs b/tests/codex_serving/responses_http_sse.rs index 196af4805..b12d8e48d 100644 --- a/tests/codex_serving/responses_http_sse.rs +++ b/tests/codex_serving/responses_http_sse.rs @@ -35,6 +35,8 @@ use crate::support::{ const COMPLETED_USAGE_SSE: &str = include_str!("../fixtures/responses/http_sse/completed_usage.sse"); +const COMPLETED_IMAGE_USAGE_SSE: &str = + include_str!("../fixtures/responses/http_sse/completed_image_usage.sse"); const COMPLETED_FIELDS_SSE: &str = include_str!("../fixtures/responses/http_sse/completed_fields.sse"); const COMPLETED_REASONING_INCLUDE_SSE: &str = @@ -336,6 +338,97 @@ async fn v1_responses_should_use_imported_account_and_record_usage() { assert_eq!(event.4["usage"]["inputTokens"], 7); } +#[tokio::test] +async fn v1_responses_should_record_image_generation_usage_when_tool_succeeds() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(COMPLETED_IMAGE_USAGE_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).await; + + let response = imported + .app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header( + "authorization", + format!("Bearer {}", imported.client_api_key), + ) + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false,"tools":[{"type":"image_generation"}]}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let usage: (i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64) = + sqlx::query_as( + "select request_count, input_tokens, output_tokens, cached_tokens, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count from account_usage where account_id = ?", + ) + .bind("acct_imported") + .fetch_one(&imported.pool) + .await + .unwrap(); + + assert_eq!(usage, (1, 7, 4, 2, 31, 9, 1, 0, 31, 9, 1, 0)); +} + +#[tokio::test] +async fn v1_responses_should_record_failed_image_generation_attempt_when_tool_has_no_output() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(COMPLETED_USAGE_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).await; + + let response = imported + .app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header( + "authorization", + format!("Bearer {}", imported.client_api_key), + ) + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false,"tools":[{"type":"image_generation"}]}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let usage: (i64, i64, i64, i64) = sqlx::query_as( + "select image_request_count, image_request_failed_count, window_image_request_count, window_image_request_failed_count from account_usage where account_id = ?", + ) + .bind("acct_imported") + .fetch_one(&imported.pool) + .await + .unwrap(); + + assert_eq!(usage, (0, 1, 0, 1)); +} + #[tokio::test] async fn v1_responses_should_passively_cache_rate_limit_headers() { let server = MockServer::start().await; diff --git a/tests/fixtures/responses/http_sse/completed_image_usage.sse b/tests/fixtures/responses/http_sse/completed_image_usage.sse new file mode 100644 index 000000000..eef4c791e --- /dev/null +++ b/tests/fixtures/responses/http_sse/completed_image_usage.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_image","object":"response","status":"completed","usage":{"input_tokens":7,"output_tokens":4,"input_tokens_details":{"cached_tokens":2}},"tool_usage":{"image_gen":{"input_tokens":31,"output_tokens":9}}}} + diff --git a/tests/platform/storage_schema.rs b/tests/platform/storage_schema.rs index 29ca4e159..5ad1fd72a 100644 --- a/tests/platform/storage_schema.rs +++ b/tests/platform/storage_schema.rs @@ -48,6 +48,35 @@ async fn sqlite_schema_should_persist_account_usage_window_columns() { ); } +#[tokio::test] +async fn sqlite_schema_should_persist_image_generation_usage_columns() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("test.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let rows: Vec<(String,)> = sqlx::query_as( + "select name from pragma_table_info('account_usage') where name in ('image_input_tokens', 'image_output_tokens', 'image_request_count', 'image_request_failed_count', 'window_image_input_tokens', 'window_image_output_tokens', 'window_image_request_count', 'window_image_request_failed_count') order by name", + ) + .fetch_all(&pool) + .await + .unwrap(); + + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "image_input_tokens", + "image_output_tokens", + "image_request_count", + "image_request_failed_count", + "window_image_input_tokens", + "window_image_output_tokens", + "window_image_request_count", + "window_image_request_failed_count", + ] + ); +} + #[tokio::test] async fn sqlite_schema_should_persist_quota_verify_required_flag() { let dir = tempfile::tempdir().unwrap(); diff --git a/tests/runtime/startup.rs b/tests/runtime/startup.rs index 803991999..f94984799 100644 --- a/tests/runtime/startup.rs +++ b/tests/runtime/startup.rs @@ -129,6 +129,7 @@ async fn app_state_should_restore_account_pool_from_sqlite_accounts() { output_tokens: 3, cached_tokens: 2, empty_response_count: 0, + ..UsageDelta::default() }, ) .await @@ -254,6 +255,7 @@ async fn app_state_should_restore_session_affinity_from_sqlite() { output_tokens: 0, cached_tokens: 0, empty_response_count: 0, + ..UsageDelta::default() }, ) .await