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

codex: track image usage

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-15 10:47:40 +08:00
co-authored by Codex
parent 8e9e89aeff
commit 44c3306023
25 changed files with 697 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/OPENAI_CHAIN_PARITY_REAUDIT.md
+6 -3
View File
@@ -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 存储；否则将其记录为双方未使用字段。
src/admin/api/usage.rs
+16
View File
@@ -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<String>,
}
@@ -52,6 +56,10 @@ impl From<AccountUsageListRecord> 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<AccountUsageSummary> for AdminUsageStatsSummaryData {
@@ -77,6 +89,10 @@ impl From<AccountUsageSummary> 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,
}
}
}
src/codex/accounts/model.rs
+16
View File
@@ -29,10 +29,18 @@ pub struct Account {
pub cloudflare_cooldown_until: Option<DateTime<Utc>>,
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<DateTime<Utc>>,
pub window_reset_at: Option<DateTime<Utc>>,
pub limit_window_seconds: Option<u64>,
@@ -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,
src/codex/accounts/pool.rs
+60 -6
View File
@@ -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(
src/codex/accounts/repository/accounts.rs
+53 -1
View File
@@ -88,11 +88,15 @@ pub struct StoredAccountMetadata {
pub updated_at: DateTime<Utc>,
}
#[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<DateTime<Utc>>,
pub window_reset_at: Option<DateTime<Utc>>,
pub limit_window_seconds: Option<u64>,
@@ -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<DateTime<Utc>>,
}
@@ -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::<i64, _>("usage_request_count").max(0) as u64,
empty_response_count: row.get::<i64, _>("usage_empty_response_count").max(0) as u64,
image_input_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_image_input_tokens"),
),
image_output_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_image_output_tokens"),
),
image_request_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_image_request_count"),
),
image_request_failed_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_image_request_failed_count"),
),
window_request_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_request_count"),
),
@@ -551,6 +591,18 @@ impl AccountRepository {
window_cached_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_cached_tokens"),
),
window_image_input_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_image_input_tokens"),
),
window_image_output_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_image_output_tokens"),
),
window_image_request_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_image_request_count"),
),
window_image_request_failed_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_image_request_failed_count"),
),
window_started_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("usage_window_started_at"),
)?,
src/codex/accounts/repository/usage.rs
+88 -4
View File
@@ -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<Acco
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"),
window_request_count: row.get("window_request_count"),
window_input_tokens: row.get("window_input_tokens"),
window_output_tokens: row.get("window_output_tokens"),
window_cached_tokens: row.get("window_cached_tokens"),
window_image_input_tokens: row.get("window_image_input_tokens"),
window_image_output_tokens: row.get("window_image_output_tokens"),
window_image_request_count: row.get("window_image_request_count"),
window_image_request_failed_count: row.get("window_image_request_failed_count"),
window_started_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("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::<Option<String>, _>("last_used_at"))?,
})
}
src/codex/accounts/service/import.rs
+8
View File
@@ -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,
src/codex/accounts/service/pool_sync.rs
+8
View File
@@ -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,
src/codex/gateway/transport/types.rs
+8
View File
@@ -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)]
src/codex/gateway/transport/usage_events.rs
+16 -1
View File
@@ -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<TokenUsage> {
.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<TokenUsage> {
.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,
})
}
src/codex/serving/chat.rs
+3 -2
View File
@@ -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()
{
src/codex/serving/dispatch/fallback.rs
+5 -2
View File
@@ -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<String>,
image_generation_requested: bool,
) -> Option<AcquiredAccount> {
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,
src/codex/serving/dispatch/mod.rs
+57 -13
View File
@@ -308,6 +308,7 @@ impl CodexUpstreamService {
retry: UpstreamAccountRetry,
model: &str,
excluded_account_ids: &mut Vec<String>,
image_generation_requested: bool,
) -> Option<AcquiredAccount> {
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,
src/codex/serving/dispatch/stream_audit.rs
+16 -6
View File
@@ -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 存储失败";
src/codex/serving/dispatch/usage.rs
+37 -5
View File
@@ -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
}
}
src/codex/serving/responses.rs
+35 -8
View File
@@ -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()
{
src/platform/storage/schema.sql
+8
View File
@@ -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),
tests/codex_accounts/pool_scheduling.rs
+13 -2
View File
@@ -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();
tests/codex_accounts/repository.rs
+69 -1
View File
@@ -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));
tests/codex_gateway/http_client.rs
+2
View File
@@ -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,
})
);
tests/codex_gateway/usage_events.rs
+46
View File
@@ -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);
}
tests/codex_serving/responses_http_sse.rs
+93
View File
@@ -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;
tests/fixtures/responses/http_sse/completed_image_usage.sse
Vendored
+3
View File
@@ -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}}}}
tests/platform/storage_schema.rs
+29
View File
@@ -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::<Vec<_>>(),
[
"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();
tests/runtime/startup.rs
+2
View File
@@ -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