feat: use business traffic for Codex overdraft detection

This commit is contained in:
bmiddle1.backend.zhangfucheng committed 2026-08-16 09:31:21 +08:00
1 parent eb8eadd7a5
commit 9e4ba889ae
24 files changed
+900 -481

No files matched your search

+19 -22
View File
@@ -6,7 +6,7 @@
## 重要说明
- 本功能仅表示:当 Codex 5h 或 7d 页面额度达到 100% 后,通过真实请求确认上游是否仍允许继续调用。
- 本功能仅表示:Codex 5h 或 7d 用量达到 95% 后开始注入兼容请求形态,并在额度达到 100% 后通过真实业务结果或单次独立探测确认上游是否仍允许继续调用。
- `passed` 只代表探测时上游仍返回有效响应,不承诺之后始终可用,也不会绕过认证失效、账号禁用、网络故障或其他风控限制。
- 探测请求和透支期请求都可能产生真实上游用量和费用。
- 此行为可能不符合上游服务条款。部署者必须自行确认合规性,并承担账号限制、服务中断等风险。
@@ -20,7 +20,7 @@
- 对符合条件的 OpenAI OAuth 普通账号和 Agent Identity 账号启用透支调度,Shadow 子账号除外。
- 支持 `/v1/responses`、转换为 Responses 的 `/v1/chat/completions` 和 `/v1/messages`,以及 Responses WebSocket v2。
- 在适用请求的最后一条用户消息后注入一对无操作工具调用,用于保持与参考实现一致的请求形态。
- 当 5h 或 7d 额度首次达到 100%,或收到带明确额度证据的 429 时,最多执行 5 次真实探测。
- 用量达到 95% 后开始对适用业务请求注入;达到 100% 后,注入业务成功直接确认 `passed`,明确额度 429 直接确认 `failed`。没有业务结果可用时,同周期最多补充 1 次独立探测。
- 管理页面的 OpenAI OAuth 常规文本“测试账号连接”也使用相同请求形态,并接入同一套额度探测状态机。
- 分别维护 5h、7d 透支周期,并在管理页面显示状态、请求数、Token、账号金额和预计恢复时间。
- 将探测状态保存在现有 `accounts.extra` JSONB 字段中,不新增数据库表,不需要手动执行数据库迁移(migration)。
@@ -257,7 +257,7 @@ docker compose \
### 第三层:查看真实探测日志
只有账号 5h/7d 额度达到 100%,或者收到带明确额度信息的上游 429,才会启动探测。额度达到 100% 后,可以在管理页面对该 OpenAI OAuth 账号执行常规文本“测试账号连接”;测试请求本身会使用透支请求形态,并触发或接续同一额度周期的探测。额度未耗尽时没有相关日志是正常现象。
账号 5h/7d 用量达到 95% 后开始注入;只有额度达到 100%,或者收到带明确额度信息的上游 429,才会产生透支确认状态。额度达到 100% 后,可以在管理页面对该 OpenAI OAuth 账号执行常规文本“测试账号连接”;测试请求本身会使用透支请求形态,并直接参与同一额度周期的判定。额度未耗尽时没有探测日志是正常现象。
```bash
docker logs --since 30m sub2api 2>&1 | \
@@ -268,12 +268,12 @@ docker logs --since 30m sub2api 2>&1 | \
| 日志 | 含义 |
| --- | --- |
| `codex_quota_overdraft_business_passed` | 注入后的真实业务或账号测试成功,直接确认透支可用 |
| `codex_quota_overdraft_business_exhausted` | 注入后的业务请求收到明确额度 429,直接确认额度耗尽并暂停账号 |
| `codex_quota_overdraft_probe_passed` | 至少一次真实探测成功,透支成立,账号继续参与调度 |
| `codex_quota_overdraft_probe_failed` | 5 次都明确返回额度限制,账号暂停至额度恢复时间 |
| `codex_quota_overdraft_probe_failed` | 单次独立探测明确返回额度限制,账号暂停至额度恢复时间 |
| `codex_quota_overdraft_probe_attempt` | 单次探测结果,包含模型、HTTP 状态码、结果和原因,不记录凭据及完整响应正文 |
| `codex_quota_overdraft_probe_inconclusive` | 网络、超时、5xx 等导致无法确认,按 1、3、10 分钟退避重试 |
| `codex_quota_overdraft_retry_scheduled` | 已安排主动重试 |
| `codex_quota_overdraft_retry_started` | 主动重试开始;服务重启后也会由数据库到期扫描恢复 |
| `codex_quota_overdraft_probe_inconclusive` | 网络、超时、5xx、普通瞬时 429 等导致无法确认;同周期不自动重试 |
| `codex_quota_overdraft_probe_claim_failed` | 无法在 PostgreSQL 中原子领取探测任务,需要排查数据库或版本 |
| `codex_quota_overdraft_state_persist_failed` | 状态无法写入 `accounts.extra`,页面可能不显示最新结果 |
| `codex_quota_overdraft_pause_applied` | `failed` 状态、账号暂停和调度通知已经原子提交 |
@@ -282,10 +282,10 @@ docker logs --since 30m sub2api 2>&1 | \
真正“透支成功”的最直接证据是:
```text
codex_quota_overdraft_probe_passed
codex_quota_overdraft_business_passed
```
其中会包含 `account_id`、`attempts`、`model` 和 `quota_window`。
如果没有业务成功证据,也可以由 `codex_quota_overdraft_probe_passed` 确认。其中会包含 `account_id`、`model` 和 `quota_window`。
### 第四层:查看数据库持久化状态
@@ -306,29 +306,26 @@ docker compose \
| 状态 | 页面显示 | 含义 |
| --- | --- | --- |
| `pending` | 透支探测中 | 最多 5 次真实探测正在执行 |
| `passed` | 透支中 | 上游仍可调用,账号继续调度并开始统计透支用量 |
| `failed` | 已确认限额 | 5 次均明确受额度限制,暂停到恢复时间 |
| `inconclusive` | 探测无法确认 | 暂时无法判断,查看 `reason_code` 和 `retry_at` |
| `pending` | 透支探测中 | 单次独立探测正在执行 |
| `passed` | 透支中 | 注入业务或独立探测成功,账号继续调度并开始统计透支用量 |
| `failed` | 已确认限额 | 注入业务或独立探测明确受额度限制,暂停到恢复时间 |
| `inconclusive` | 探测无法确认 | 单次探测受网络、瞬时限流或上游异常干扰;同周期不自动重试 |
| `recovered` | 额度已恢复 | 原额度周期结束,相关暂停已清理 |
页面显示 `透支中`,数据库为 `passed`,并且后续业务请求返回成功,三者同时满足即可确认功能完整生效。
## 探测逻辑
每个额度周期的单轮计划最多探测 5 次,每次最长 20 秒。模型按以下顺序循环:
每个额度周期采用以下混合判定流程:
```text
首选模型 -> gpt-5.5 -> gpt-5.4-mini -> 首选模型 -> gpt-5.5
```
- 任意一次返回有效 `response.completed` 或 `response.output_item.done`,立即记为 `passed`。
- 只有 5 次都明确返回 `quota_limited`,才记为 `failed`。
- 用量低于 95% 时不注入,达到 95% 后对适用的普通 OAuth 文本请求持续注入。
- 达到 100% 后,注入业务返回成功即记为 `passed`;返回明确 `quota_limited` 即记为 `failed`,暂停账号并切换到其他账号。
- 如果额度刚从 95% 以下跳到 100%,当前业务请求没有注入,系统仅补充 1 次独立 Responses 探测,最长 20 秒,并优先使用当前业务模型。
- 网络错误、超时、5xx、普通瞬时 429、无效响应等记为 `inconclusive`,不会因此误停账号。
- 401/403 继续由原有认证失效逻辑处理。
- 账号禁用、过载、代理故障、模型冷却等非额度限制不会被本功能绕过。
无法确认时,`retry_count` 记录当前退避轮次。后台会依次等待 1 分钟、3 分钟和 10 分钟执行最多 3 轮自动重试;每分钟的数据库扫描会恢复因容器重启而丢失的内存定时器,并接续超过 2 分钟的遗留 `pending` 状态。达到自动重试上限后不会持续消耗上游请求,但后续额度刷新、账号测试或新的明确额度响应仍可继续观察。
独立探测无法确认后,同一额度周期不会自动重试,也不会持续消耗探测请求。后续已经注入的真实业务成功或明确额度 429 仍可把状态更新为 `passed` 或 `failed`;同周期 `failed` 为终态,成功结果不能覆盖它。
多实例部署通过 PostgreSQL 原子领取(atomic claim)保证同一账号、同一额度周期只有一个实例发起探测。
@@ -488,7 +485,7 @@ docker compose \
页面的“限流中”直接来自账号记录的 `rate_limit_reset_at`。旧实现存在竞态:探测开始后若另一个 429 更新了该时间,即使探测最终 `passed`,严格时间比较也可能不清理它,因此出现“实际透支可用但页面限流中”。当前版本在 `passed/recovered` 时会清除该残留状态,并记录 `codex_quota_overdraft_stale_rate_limit_cleared`;`failed` 不会清理。
其他真实限制仍会正常生效,包括账号禁用、401/403、5 次探测均确认额度耗尽、代理/网络错误、模型级冷却、并发限制或没有明确额度证据的其他上游 429。若状态持续存在,先按 `account_id` 检查上述日志、数据库中的 `extra->'codex_quota_overdraft_probe'` 和 `rate_limit_reset_at`,区分额度透支与普通限流。
其他真实限制仍会正常生效,包括账号禁用、401/403、注入业务或单次探测确认额度耗尽、代理/网络错误、模型级冷却、并发限制或没有明确额度证据的其他上游 429。若状态持续存在,先按 `account_id` 检查上述日志、数据库中的 `extra->'codex_quota_overdraft_probe'` 和 `rate_limit_reset_at`,区分额度透支与普通限流。
### `pq: could not determine data type of parameter $1`
+14 -15
View File
@@ -16,13 +16,13 @@
源码 Docker 构建会读取根目录的 `FORK_VERSION`,并以 `BuildType=source` 写入二进制。后台更新服务通过 GitHub Contents API 读取 Fork 分支上的同名文件,使用语义化版本比较判断是否有新版本。Redis 缓存同时记录仓库和构建类型,因此旧的官方更新缓存不会继续生效。
源码构建仅显示 `git pull` 更新提示,`PerformUpdate`、指定版本回退和在线回退列表均被禁用,防止官方二进制覆盖透支功能。维护者每次发布 Fork 更新都必须递增 `FORK_VERSION`;当前版本为 `0.1.177-overdraft.4`,同步下一上游版本后使用 `0.1.178-overdraft.1`。
源码构建仅显示 `git pull` 更新提示,`PerformUpdate`、指定版本回退和在线回退列表均被禁用,防止官方二进制覆盖透支功能。维护者每次发布 Fork 更新都必须递增 `FORK_VERSION`;当前版本为 `0.1.177-overdraft.5`,同步下一上游版本后使用 `0.1.178-overdraft.1`。
Sub2API `v0.1.177` 将原生 remote compaction v2 保留在 `/responses` 路径。本定制同时检查旧 Compact 路径和原生 v2 请求信号,二者都不会开启额度透支调度或注入透支请求形态。
## 后台测试连接与残留限流修复
管理页面的 OpenAI OAuth 常规文本“测试账号连接”使用与真实业务请求相同的透支请求形态。测试成功并取得 5h/7d 快照后会触发额度观察;测试返回带明确额度证据的 429 时交给透支协调器执行最多 5 次真实探测,不再直接写入普通账号级限流。API Key、Shadow、图片和 Compact 测试不启用该行为。
管理页面的 OpenAI OAuth 常规文本“测试账号连接”与真实业务请求使用相同策略:用量达到 95% 后注入;100% 时成功直接确认 `passed`,明确额度 429 直接确认 `failed`。API Key、Shadow、图片和 Compact 测试不启用该行为。
探测得到 `passed`,或额度周期进入 `recovered` 后,会清除仍残留的账号级 `rate_limit_reset_at` 以及本功能或额度阈值产生的暂停。这里不再要求当前重置时间与探测开始时记录的 `observed_rate_limit_reset_at` 完全相同,因为并发 429 可能在探测期间更新该字段;`failed` 状态不会清理限流。发生清理时记录 `codex_quota_overdraft_stale_rate_limit_cleared`。
@@ -57,19 +57,18 @@ gateway:
- Responses WebSocket v2 首轮和后续轮
- 管理页面 OpenAI OAuth 常规文本“测试账号连接”
每次最终发往上游前,如果 `input` 最后一项是用户消息,则追加一对与参考项目相同的无操作 `custom_tool_call` 和 `custom_tool_call_output`。请求体超过 32 MiB、JSON 无效、形状不匹配或已经注入时保持原样。
用量达到 95% 后,每次最终发往上游前,如果 `input` 最后一项是用户消息,则追加一对与参考项目相同的无操作 `custom_tool_call` 和 `custom_tool_call_output`。请求体超过 32 MiB、JSON 无效、形状不匹配或已经注入时保持原样。
5h 或 7d 额度首次达到 100%,或者收到包含明确额度证据的真实 429 时,按参考项目执行复核:
5h 或 7d 额度达到 100%,或者收到包含明确额度证据的真实 429 时,执行以下混合判定:
1. 同一额度周期只启动一个探测计划,状态持久化在 `accounts.extra.codex_quota_overdraft_probe`,多实例使用数据库原子 claim 去重。
2. 最多执行 5 次真实 Responses 请求,每次独立 20 秒超时。
3. 模型按“首选模型、`gpt-5.5`、`gpt-5.4-mini`”循环。例如首选为 `gpt-5.4` 时,顺序是 `gpt-5.4`、`gpt-5.5`、`gpt-5.4-mini`、`gpt-5.4`、`gpt-5.5`。
4. 任意一次返回有效的 `response.completed` 或 `response.output_item.done` 即判定 `passed`,账号继续调度并开始透支统计。
5. 只有 5 次均返回明确 `quota_limited` 才判定 `failed`,账号暂停到对应额度恢复时间;5h 和 7d 同时耗尽时取最晚恢复时间。`failed` 状态、临时暂停和调度 outbox 在同一数据库事务中提交,避免出现状态已失败但账号没有暂停。
6. 网络错误、超时、5xx、普通瞬时 429 和无效响应均判定 `inconclusive`,不会误停账号;后台按 1 分钟、3 分钟、10 分钟主动重试,并每分钟扫描数据库恢复服务重启前已到期的任务。重启时遗留超过 2 分钟的 `pending` 状态也会自动接续。每个额度周期最多自动重试 3 轮,之后仍可由额度刷新或新请求继续观察。
7. 401/403、账号停用等认证问题交给原有认证异常逻辑处理,不加入额度重试;400/404 会轮换全部探测模型后再判定为 `inconclusive`,不会暂停整个账号。
8. 账号禁用、过载、代理/传输故障、模型级冷却及其他临时不可调度原因不被绕过。
9. 额度恢复后状态改为 `recovered`,清理本功能或额度阈值产生的暂停;5h 与 7d 分别维护透支起点,不会因另一窗口后来耗尽而重置已有统计。
1. 实际注入的业务请求成功时直接判定 `passed`,账号继续调度并开始透支统计,不额外消耗探测请求。
2. 实际注入的业务请求返回明确额度 429 时直接判定 `failed`,账号暂停到对应额度恢复时间并切号;5h 和 7d 同时耗尽时取最晚恢复时间。
3. 如果没有可用的业务成功证据,同一额度周期仅启动 1 次独立 Responses 探测,最长 20 秒并优先使用当前业务模型。多实例使用 PostgreSQL 原子 claim 去重。
4. 独立探测成功即判定 `passed`,明确 `quota_limited` 即判定 `failed`;网络错误、超时、5xx、普通瞬时 429 和无效响应判定 `inconclusive`,同周期不自动重试。
5. `failed` 状态、临时暂停和调度 outbox 在同一数据库事务中提交。同周期 `failed` 为终态,晚到的业务成功或探测结果不能覆盖;业务明确额度 429 可以把并发产生的 `passed` 或 `inconclusive` 收敛为 `failed`。
6. 401/403、账号停用等认证问题交给原有认证异常逻辑处理;400/404 判定为 `inconclusive`,不会暂停整个账号。
7. 账号禁用、过载、代理/传输故障、模型级冷却及其他临时不可调度原因不被绕过。
8. 额度恢复后状态改为 `recovered`,清理本功能或额度阈值产生的暂停;5h 与 7d 分别维护透支起点,不会因另一窗口后来耗尽而重置已有统计。
状态值包括 `pending`、`passed`、`failed`、`inconclusive`、`recovered`。账号用量页面显示探测状态、尝试次数、额度周期、透支期成功请求数、Token、账号金额及预计恢复时间。
@@ -157,8 +156,8 @@ codex_quota_overdraft_probe_passed
codex_quota_overdraft_probe_failed
codex_quota_overdraft_probe_inconclusive
codex_quota_overdraft_probe_attempt
codex_quota_overdraft_retry_scheduled
codex_quota_overdraft_retry_started
codex_quota_overdraft_business_passed
codex_quota_overdraft_business_exhausted
codex_quota_overdraft_pause_applied
codex_quota_overdraft_stale_rate_limit_cleared
```
+1 -1
View File
@@ -1 +1 @@
0.1.177-overdraft.4
0.1.177-overdraft.5
+3 -2
View File
@@ -21,10 +21,11 @@ English | [中文](README_CN.md) | [日本語](README_JA.md)
## Fork Features
- Runs up to five real probes after the reported Codex 5h or 7d quota reaches 100%.
- Pre-arms ordinary OAuth text traffic with the overdraft payload at 95% usage and uses successful business traffic as direct evidence after usage reaches 100%.
- Marks an injected request's explicit quota 429 as terminal for that cycle; when business evidence is unavailable, runs at most one independent probe per quota cycle.
- Keeps an account schedulable after a successful probe and tracks overdraft requests, tokens, cost, and recovery for both windows independently.
- Exposes `pending`, `passed`, `failed`, `inconclusive`, and `recovered` states in the admin UI and PostgreSQL.
- Actively retries inconclusive probes after 1, 3, and 10 minutes, including recovery from persisted retry state after a restart.
- Treats transient 429s, timeouts, network failures, and 5xx responses as inconclusive without automatic background retries.
- Uses an atomic PostgreSQL claim for multi-instance deployments and atomically commits terminal failure, account pause, and scheduler notification without a schema migration.
- Supports an immediate configuration rollback to the upstream scheduling behavior.
+4 -3
View File
@@ -21,8 +21,9 @@
## 本 Fork 增加的功能
- Codex 5h / 7d 额度达到 100% 后,最多执行 5 次真实请求探测,判断账号是否仍可继续调用。
- 探测成功后继续参与账号调度,并分别统计 5h / 7d 透支期请求数、Token 和金额。
- Codex 5h / 7d 用量达到 95% 后,为普通 OAuth 文本业务请求注入透支请求形态;达到 100% 后直接使用真实业务结果确认透支状态。
- 注入业务请求成功即记为 `passed`;返回明确额度 429 即记为 `failed` 并切号冷却。缺少业务证据时,同一额度周期最多补充 1 次独立探测。
- 确认成功后继续参与账号调度,并分别统计 5h / 7d 透支期请求数、Token 和金额。
- 管理页面显示“透支探测中”“透支中”“已确认限额”“探测无法确认”和“额度已恢复”。
- 网络、超时、5xx 和普通瞬时 429 不会被误判为额度耗尽;401/403、账号禁用和其他风控仍使用原有策略。
- 多实例通过 PostgreSQL 原子领取(atomic claim)去重;状态保存在现有 `accounts.extra`,无需新增数据表。
@@ -70,7 +71,7 @@ docker logs --since 30m sub2api 2>&1 | \
grep -E 'codex_quota_overdraft_(probe|state|pause|stale_rate_limit)'
```
出现 `codex_quota_overdraft_probe_passed`,页面显示“透支中”,并且后续真实业务请求成功,即可确认透支功能完整生效。无法确认的探测会按 1、3、10 分钟主动重试,容器重启后由持久化状态继续;5 次明确额度失败时,探测状态、账号暂停和调度通知会原子提交。OpenAI OAuth 常规文本“测试账号连接”也会使用透支请求形态并接入同一探测状态机;API Key、Shadow、图片和 Compact 测试除外。额度未达到 100% 时没有探测日志是正常现象。
出现 `codex_quota_overdraft_business_passed` 或 `codex_quota_overdraft_probe_passed`,页面显示“透支中”,并且后续真实业务请求成功,即可确认透支功能完整生效。注入业务请求返回明确额度 429 时会出现 `codex_quota_overdraft_business_exhausted`,状态、账号暂停和调度通知会原子提交。网络错误、5xx、超时和普通瞬时 429 不会判定透支结束;独立探测为 `inconclusive` 后同周期不自动重试。OpenAI OAuth 常规文本“测试账号连接”也使用同一请求形态和状态机;API Key、Shadow、图片和 Compact 测试除外。额度未达到 95% 时没有透支注入,未达到 100% 时没有探测日志,均属正常现象。
## 来源、许可证与风险
+1 -1
View File
@@ -199,7 +199,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, configConfig, usageLogRepository, settingService)
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIGatewayService)
usageCache := service.NewUsageCache()
codexQuotaOverdraftCoordinator := service.ProvideCodexQuotaOverdraftCoordinator(accountRepository, httpUpstream, openAITokenProvider, tlsFingerprintProfileService, configConfig, tempUnschedCache, openAIGatewayService, rateLimitService, timingWheelService)
codexQuotaOverdraftCoordinator := service.ProvideCodexQuotaOverdraftCoordinator(accountRepository, httpUpstream, openAITokenProvider, tlsFingerprintProfileService, configConfig, tempUnschedCache, openAIGatewayService, rateLimitService)
accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService, codexQuotaOverdraftCoordinator)
accountTestService := service.ProvideAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService, openAIGatewayService, settingService, codexQuotaOverdraftCoordinator)
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
@@ -344,9 +344,9 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
}
}
if result != nil {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, result.FirstTokenMs, c.Request.Context())
} else {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, nil, c.Request.Context())
}
userAgent := c.GetHeader("User-Agent")
@@ -717,9 +717,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
if account.Type == service.AccountTypeOAuth && !account.IsShadow() {
h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(c.Request.Context(), account.ID, result.ResponseHeaders)
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), result.FirstTokenMs, c.Request.Context())
} else {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), nil, c.Request.Context())
}
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
@@ -1241,9 +1241,9 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
}
}
if result != nil {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, result.FirstTokenMs, c.Request.Context())
} else {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, nil)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, nil, c.Request.Context())
}
userAgent := c.GetHeader("User-Agent")
@@ -2204,7 +2204,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
if scheduleModel == "" {
scheduleModel = turnRequestedModel
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, scheduleModel, openAIForwardSucceededForScheduling(result), result.FirstTokenMs)
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, scheduleModel, openAIForwardSucceededForScheduling(result), result.FirstTokenMs, ctx)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
+103 -84
View File
@@ -2593,8 +2593,8 @@ func (r *accountRepository) UpdateExtra(ctx context.Context, id int64, updates m
return nil
}
// ClaimCodexQuotaOverdraftProbe atomically reserves one quota cycle. This
// prevents duplicate five-request probe plans across multiple sub2api replicas.
// ClaimCodexQuotaOverdraftProbe atomically reserves one quota cycle. A cycle is
// claimed at most once across all sub2api replicas.
func (r *accountRepository) ClaimCodexQuotaOverdraftProbe(
ctx context.Context,
id int64,
@@ -2607,6 +2607,39 @@ func (r *accountRepository) ClaimCodexQuotaOverdraftProbe(
if err != nil {
return false, err
}
result, err := r.sql.ExecContext(ctx, `
UPDATE accounts
SET extra = COALESCE(extra, '{}'::jsonb) || jsonb_build_object($1::text, $2::jsonb),
updated_at = NOW()
WHERE id = $3
AND deleted_at IS NULL
AND COALESCE(extra #>> '{codex_quota_overdraft_probe,cycle_key}', '') <> $4
`, service.CodexQuotaOverdraftProbeExtraKey, string(payload), id, state.CycleKey)
if err != nil {
return false, err
}
affected, err := result.RowsAffected()
if err != nil || affected == 0 {
return false, err
}
r.syncSchedulerAccountSnapshot(ctx, id)
return true, nil
}
// PersistCodexQuotaOverdraftProbeUnlessFailed stores a non-failure result while
// preserving a terminal failure already confirmed for the same quota cycle.
func (r *accountRepository) PersistCodexQuotaOverdraftProbeUnlessFailed(
ctx context.Context,
id int64,
state *service.CodexQuotaOverdraftProbeState,
) (bool, error) {
if state == nil || state.Status == "failed" || strings.TrimSpace(state.CycleKey) == "" {
return false, nil
}
payload, err := json.Marshal(state)
if err != nil {
return false, err
}
result, err := r.sql.ExecContext(ctx, `
UPDATE accounts
SET extra = COALESCE(extra, '{}'::jsonb) || jsonb_build_object($1::text, $2::jsonb),
@@ -2615,14 +2648,7 @@ func (r *accountRepository) ClaimCodexQuotaOverdraftProbe(
AND deleted_at IS NULL
AND (
COALESCE(extra #>> '{codex_quota_overdraft_probe,cycle_key}', '') <> $4
OR (
extra #>> '{codex_quota_overdraft_probe,status}' = 'inconclusive'
AND COALESCE(NULLIF(extra #>> '{codex_quota_overdraft_probe,retry_at}', '')::timestamptz, '1970-01-01'::timestamptz) <= NOW()
)
OR (
extra #>> '{codex_quota_overdraft_probe,status}' = 'pending'
AND COALESCE(NULLIF(extra #>> '{codex_quota_overdraft_probe,started_at}', '')::timestamptz, '1970-01-01'::timestamptz) <= NOW() - INTERVAL '2 minutes'
)
OR COALESCE(extra #>> '{codex_quota_overdraft_probe,status}', '') <> 'failed'
)
`, service.CodexQuotaOverdraftProbeExtraKey, string(payload), id, state.CycleKey)
if err != nil {
@@ -2636,6 +2662,68 @@ func (r *accountRepository) ClaimCodexQuotaOverdraftProbe(
return true, nil
}
// ClearCodexQuotaOverdraftPauseIfState clears stale scheduling state only when
// the persisted overdraft result is still the expected available state.
func (r *accountRepository) ClearCodexQuotaOverdraftPauseIfState(
ctx context.Context,
id int64,
cycleKey string,
status string,
expectedTempReason string,
) (bool, bool, error) {
if id <= 0 || strings.TrimSpace(cycleKey) == "" || (status != "passed" && status != "recovered") {
return false, false, nil
}
beginner, ok := r.sql.(codexQuotaOverdraftTxBeginner)
if !ok {
return false, false, errors.New("account repository does not support transactions")
}
tx, err := beginner.BeginTx(ctx, nil)
if err != nil {
return false, false, err
}
defer func() { _ = tx.Rollback() }()
var clearedRateLimit, clearedTemp bool
err = tx.QueryRowContext(ctx, `
WITH target AS (
SELECT id,
(rate_limited_at IS NOT NULL OR rate_limit_reset_at IS NOT NULL) AS clear_rate_limit,
($4::text <> '' AND COALESCE(temp_unschedulable_reason = $4, FALSE)) AS clear_temp
FROM accounts
WHERE id = $1
AND deleted_at IS NULL
AND extra #>> '{codex_quota_overdraft_probe,cycle_key}' = $2
AND extra #>> '{codex_quota_overdraft_probe,status}' = $3
FOR UPDATE
)
UPDATE accounts AS a
SET rate_limited_at = CASE WHEN target.clear_rate_limit THEN NULL ELSE a.rate_limited_at END,
rate_limit_reset_at = CASE WHEN target.clear_rate_limit THEN NULL ELSE a.rate_limit_reset_at END,
temp_unschedulable_until = CASE WHEN target.clear_temp THEN NULL ELSE a.temp_unschedulable_until END,
temp_unschedulable_reason = CASE WHEN target.clear_temp THEN NULL ELSE a.temp_unschedulable_reason END,
updated_at = NOW()
FROM target
WHERE a.id = target.id
AND (target.clear_rate_limit OR target.clear_temp)
RETURNING target.clear_rate_limit, target.clear_temp
`, id, cycleKey, status, expectedTempReason).Scan(&clearedRateLimit, &clearedTemp)
if errors.Is(err, sql.ErrNoRows) {
return false, false, nil
}
if err != nil {
return false, false, err
}
if err := enqueueSchedulerOutbox(ctx, tx, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
return false, false, err
}
if err := tx.Commit(); err != nil {
return false, false, err
}
r.syncSchedulerAccountSnapshotDetached(ctx, id)
return clearedRateLimit, clearedTemp, nil
}
type codexQuotaOverdraftTxBeginner interface {
BeginTx(context.Context, *sql.TxOptions) (*sql.Tx, error)
}
@@ -2681,8 +2769,11 @@ func (r *accountRepository) FinalizeCodexQuotaOverdraftProbeFailed(
WHERE id = $5
AND deleted_at IS NULL
AND extra #>> '{codex_quota_overdraft_probe,cycle_key}' = $6
AND extra #>> '{codex_quota_overdraft_probe,status}' IN ('pending', 'failed')
`, service.CodexQuotaOverdraftProbeExtraKey, string(payload), until, reason, id, state.CycleKey)
AND (
extra #>> '{codex_quota_overdraft_probe,status}' IN ('pending', 'failed')
OR ($7::boolean AND extra #>> '{codex_quota_overdraft_probe,status}' IN ('passed', 'inconclusive', 'recovered'))
)
`, service.CodexQuotaOverdraftProbeExtraKey, string(payload), until, reason, id, state.CycleKey, state.ReasonCode == "business_quota_limited")
if err != nil {
return false, err
}
@@ -2703,78 +2794,6 @@ func (r *accountRepository) FinalizeCodexQuotaOverdraftProbeFailed(
return true, nil
}
// ListDueCodexQuotaOverdraftProbes returns a bounded set of persisted
// inconclusive or stale pending probes so the retry sweep can resume work.
func (r *accountRepository) ListDueCodexQuotaOverdraftProbes(ctx context.Context, now time.Time, limit int) ([]*service.Account, error) {
if limit <= 0 {
limit = 100
}
if limit > 1000 {
limit = 1000
}
rows, err := r.sql.QueryContext(ctx, `
SELECT id
FROM accounts
WHERE deleted_at IS NULL
AND platform = $1
AND type = $2
AND status = $3
AND schedulable IS TRUE
AND parent_account_id IS NULL
AND (
(
extra #>> '{codex_quota_overdraft_probe,status}' = 'inconclusive'
AND COALESCE(NULLIF(extra #>> '{codex_quota_overdraft_probe,retry_at}', '')::timestamptz, 'infinity'::timestamptz) <= $4
AND COALESCE(NULLIF(extra #>> '{codex_quota_overdraft_probe,retry_count}', '')::integer, 0) <= $5
)
OR (
extra #>> '{codex_quota_overdraft_probe,status}' = 'pending'
AND COALESCE(NULLIF(extra #>> '{codex_quota_overdraft_probe,started_at}', '')::timestamptz, 'infinity'::timestamptz) <= $4 - INTERVAL '2 minutes'
)
)
AND COALESCE(NULLIF(extra #>> '{codex_quota_overdraft_probe,recover_at}', '')::timestamptz, 'infinity'::timestamptz) > $4
ORDER BY COALESCE(
NULLIF(extra #>> '{codex_quota_overdraft_probe,retry_at}', '')::timestamptz,
NULLIF(extra #>> '{codex_quota_overdraft_probe,started_at}', '')::timestamptz
) ASC, id ASC
LIMIT $6
`, service.PlatformOpenAI, service.AccountTypeOAuth, service.StatusActive, now, 3, limit)
if err != nil {
return nil, err
}
defer rows.Close()
ids := make([]int64, 0, limit)
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
ids = append(ids, id)
}
if err := rows.Err(); err != nil {
return nil, err
}
if len(ids) == 0 {
return nil, nil
}
accounts, err := r.GetByIDs(ctx, ids)
if err != nil {
return nil, err
}
byID := make(map[int64]*service.Account, len(accounts))
for i := range accounts {
account := accounts[i]
byID[account.ID] = account
}
ordered := make([]*service.Account, 0, len(accounts))
for _, id := range ids {
if account := byID[id]; account != nil {
ordered = append(ordered, account)
}
}
return ordered, nil
}
// UpdateUpstreamBillingProbeSnapshot stores a probe result only while the
// network identity used by that probe is still current.
func (r *accountRepository) UpdateUpstreamBillingProbeSnapshot(
@@ -23,8 +23,8 @@ func TestFinalizeCodexQuotaOverdraftProbeFailedIsAtomic(t *testing.T) {
Status: "failed",
CycleKey: "5h:quota-cycle",
QuotaWindow: "5h",
Attempts: 5,
Limit: 5,
Attempts: 1,
Limit: 1,
StartedAt: now,
RecoverAt: &until,
}
@@ -32,7 +32,7 @@ func TestFinalizeCodexQuotaOverdraftProbeFailedIsAtomic(t *testing.T) {
mock.ExpectBegin()
mock.ExpectExec(`(?s)UPDATE accounts.*temp_unschedulable_until.*codex_quota_overdraft_probe`).
WithArgs(service.CodexQuotaOverdraftProbeExtraKey, sqlmock.AnyArg(), until, reason, int64(77), state.CycleKey).
WithArgs(service.CodexQuotaOverdraftProbeExtraKey, sqlmock.AnyArg(), until, reason, int64(77), state.CycleKey, false).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload, dedup_key)")).
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(77), nil, nil, sqlmock.AnyArg()).
@@ -58,8 +58,8 @@ func TestFinalizeCodexQuotaOverdraftProbeFailedRollsBackWhenOutboxFails(t *testi
Status: "failed",
CycleKey: "5h:quota-cycle",
QuotaWindow: "5h",
Attempts: 5,
Limit: 5,
Attempts: 1,
Limit: 1,
StartedAt: now,
RecoverAt: &until,
}
@@ -79,20 +79,92 @@ func TestFinalizeCodexQuotaOverdraftProbeFailedRollsBackWhenOutboxFails(t *testi
require.NoError(t, mock.ExpectationsWereMet())
}
func TestListDueCodexQuotaOverdraftProbesIncludesInconclusiveAndStalePending(t *testing.T) {
func TestPersistCodexQuotaOverdraftProbeUnlessFailedPreservesTerminalFailure(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
now := time.Now().UTC()
mock.ExpectQuery(`(?s)SELECT id.*status}' = 'inconclusive'.*retry_count.*status}' = 'pending'.*INTERVAL '2 minutes'`).
WithArgs(service.PlatformOpenAI, service.AccountTypeOAuth, service.StatusActive, now, 3, 100).
WillReturnRows(sqlmock.NewRows([]string{"id"}))
state := &service.CodexQuotaOverdraftProbeState{
Status: "passed",
CycleKey: "5h:quota-cycle",
QuotaWindow: "5h",
Limit: 1,
StartedAt: now,
}
mock.ExpectExec(`(?s)UPDATE accounts.*status}', ''\) <> 'failed'`).
WithArgs(service.CodexQuotaOverdraftProbeExtraKey, sqlmock.AnyArg(), int64(77), state.CycleKey).
WillReturnResult(sqlmock.NewResult(0, 0))
repo := newAccountRepositoryWithSQL(nil, db, nil)
accounts, err := repo.ListDueCodexQuotaOverdraftProbes(context.Background(), now, 100)
persisted, err := repo.PersistCodexQuotaOverdraftProbeUnlessFailed(context.Background(), 77, state)
require.NoError(t, err)
require.Empty(t, accounts)
require.False(t, persisted)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestClearCodexQuotaOverdraftPauseIfStateIsAtomic(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
reason := service.BuildTempUnschedReasonPayload("codex_quota_overdraft", "quota exhausted")
mock.ExpectBegin()
mock.ExpectQuery(`(?s)WITH target AS.*codex_quota_overdraft_probe,cycle_key.*FOR UPDATE.*RETURNING target.clear_rate_limit, target.clear_temp`).
WithArgs(int64(77), "5h:quota-cycle", "passed", reason).
WillReturnRows(sqlmock.NewRows([]string{"clear_rate_limit", "clear_temp"}).AddRow(true, true))
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload, dedup_key)")).
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(77), nil, nil, sqlmock.AnyArg()).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectCommit()
repo := newAccountRepositoryWithSQL(nil, db, nil)
clearedRate, clearedTemp, err := repo.ClearCodexQuotaOverdraftPauseIfState(
context.Background(),
77,
"5h:quota-cycle",
"passed",
reason,
)
require.NoError(t, err)
require.True(t, clearedRate)
require.True(t, clearedTemp)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestFinalizeCodexQuotaOverdraftBusinessFailureCanReplacePassedState(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
now := time.Now().UTC()
until := now.Add(5 * time.Hour)
state := &service.CodexQuotaOverdraftProbeState{
Status: "failed",
CycleKey: "5h:quota-cycle",
QuotaWindow: "5h",
Limit: 1,
ReasonCode: "business_quota_limited",
StartedAt: now,
RecoverAt: &until,
}
reason := service.BuildTempUnschedReasonPayload("codex_quota_overdraft", "quota exhausted")
mock.ExpectBegin()
mock.ExpectExec(`(?s)UPDATE accounts.*'passed', 'inconclusive', 'recovered'`).
WithArgs(service.CodexQuotaOverdraftProbeExtraKey, sqlmock.AnyArg(), until, reason, int64(77), state.CycleKey, true).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload, dedup_key)")).
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(77), nil, nil, sqlmock.AnyArg()).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectCommit()
repo := newAccountRepositoryWithSQL(nil, db, nil)
finalized, err := repo.FinalizeCodexQuotaOverdraftProbeFailed(context.Background(), 77, state, until, reason)
require.NoError(t, err)
require.True(t, finalized)
require.NoError(t, mock.ExpectationsWereMet())
}
@@ -156,6 +156,7 @@ type AccountTestService struct {
type codexQuotaOverdraftAccountTestCoordinator interface {
ObserveAccount(account *Account, preferredModel string)
ObserveBusinessSuccess(account *Account, preferredModel string)
HandleQuota429(ctx context.Context, account *Account, headers http.Header, body []byte, preferredModel string) bool
}
@@ -706,9 +707,13 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
}
payload := createOpenAITestPayload(upstreamTestModelID, isOAuth)
payloadBytes, _ := json.Marshal(payload)
if s.codexQuotaOverdraftTestEnabled(account) {
overdraftInjected := false
if s.codexQuotaOverdraftTestEnabled(account) && codexQuotaOverdraftInjectionEligible(account, time.Now().UTC()) {
ctx = WithCodexQuotaOverdraftScheduling(ctx)
if updated, changed, err := injectCodexQuotaOverdraft(payloadBytes); err == nil && changed {
payloadBytes = updated
overdraftInjected = true
markCodexQuotaOverdraftInjected(ctx, account.ID)
}
}
@@ -812,7 +817,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
return err
}
if s.codexQuotaOverdraft != nil && s.codexQuotaOverdraftTestEnabled(account) {
s.codexQuotaOverdraft.ObserveAccount(account, upstreamTestModelID)
if overdraftInjected {
s.codexQuotaOverdraft.ObserveBusinessSuccess(account, upstreamTestModelID)
} else {
s.codexQuotaOverdraft.ObserveAccount(account, upstreamTestModelID)
}
}
return nil
}
@@ -75,6 +75,7 @@ type openAIAccountTestRepo struct {
type accountTestOverdraftCoordinatorStub struct {
observeCalls int
businessCalls int
observedAccount *Account
observedModel string
handleCalls int
@@ -89,6 +90,12 @@ func (s *accountTestOverdraftCoordinatorStub) ObserveAccount(account *Account, p
s.observedModel = preferredModel
}
func (s *accountTestOverdraftCoordinatorStub) ObserveBusinessSuccess(account *Account, preferredModel string) {
s.businessCalls++
s.observedAccount = account
s.observedModel = preferredModel
}
func (s *accountTestOverdraftCoordinatorStub) HandleQuota429(
_ context.Context,
account *Account,
@@ -190,6 +197,10 @@ func TestAccountTestService_OpenAIOAuthOverdraftTestInjectsAndObservesSnapshot(t
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
Extra: map[string]any{
"codex_5h_used_percent": 95,
"codex_5h_reset_at": time.Now().Add(5 * time.Hour).Format(time.RFC3339),
},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.4", "", "")
@@ -201,7 +212,8 @@ func TestAccountTestService_OpenAIOAuthOverdraftTestInjectsAndObservesSnapshot(t
require.Equal(t, "custom_tool_call", gjson.GetBytes(body, "input.1.type").String())
require.Equal(t, "custom_tool_call_output", gjson.GetBytes(body, "input.2.type").String())
require.Equal(t, gjson.GetBytes(body, "input.1.call_id").String(), gjson.GetBytes(body, "input.2.call_id").String())
require.Equal(t, 1, coordinator.observeCalls)
require.Zero(t, coordinator.observeCalls)
require.Equal(t, 1, coordinator.businessCalls)
require.Same(t, account, coordinator.observedAccount)
require.Equal(t, "gpt-5.4", coordinator.observedModel)
require.Equal(t, 100.0, account.Extra["codex_5h_used_percent"])
@@ -2299,9 +2299,12 @@ func (s *OpenAIGatewayService) isOpenAIAccountTransportCompatible(account *Accou
return s.getOpenAIWSProtocolResolver().Resolve(account).Transport == requiredTransport
}
func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(accountID int64, model string, success bool, firstTokenMs *int) {
func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(accountID int64, model string, success bool, firstTokenMs *int, requestCtx ...context.Context) {
if success {
s.clearOpenAIAccountModelTransientState(accountID, normalizeOpenAIAccountModelTransientModel(model))
if len(requestCtx) > 0 && s.codexQuotaOverdraft != nil && codexQuotaOverdraftWasInjected(requestCtx[0], accountID) {
s.codexQuotaOverdraft.ObserveBusinessSuccessByID(accountID, model)
}
}
scheduler := s.getOpenAIAccountScheduler(context.Background())
if scheduler == nil {
@@ -6,6 +6,7 @@ import (
"encoding/hex"
"encoding/json"
"strings"
"sync"
"sync/atomic"
"time"
@@ -13,9 +14,10 @@ import (
)
const (
codexQuotaOverdraftCallIDPrefix = "call_sub2api_overdraft_"
codexQuotaOverdraftExecInput = `const r = await tools.exec_command({"cmd":"true","yield_time_ms":1000,"max_output_tokens":1000}); text(r.output);`
codexQuotaOverdraftMaxBodyBytes = 32 << 20
codexQuotaOverdraftCallIDPrefix = "call_sub2api_overdraft_"
codexQuotaOverdraftExecInput = `const r = await tools.exec_command({"cmd":"true","yield_time_ms":1000,"max_output_tokens":1000}); text(r.output);`
codexQuotaOverdraftMaxBodyBytes = 32 << 20
codexQuotaOverdraftPrearmPercent = 95
)
var codexQuotaOverdraftEnabled atomic.Bool
@@ -40,6 +42,10 @@ func isCodexQuotaOverdraftAccount(account *Account) bool {
type codexQuotaOverdraftSchedulingCtxKey struct{}
type codexQuotaOverdraftRequestState struct {
injectedAccounts sync.Map
}
// WithCodexQuotaOverdraftScheduling marks normal text-generation requests as
// eligible for the experimental quota-overdraft behavior. The process-wide
// configuration switch is still checked at every scheduling and mutation gate.
@@ -47,7 +53,10 @@ func WithCodexQuotaOverdraftScheduling(ctx context.Context) context.Context {
if ctx == nil {
ctx = context.Background()
}
return context.WithValue(ctx, codexQuotaOverdraftSchedulingCtxKey{}, true)
if codexQuotaOverdraftRequestStateFromContext(ctx) != nil {
return ctx
}
return context.WithValue(ctx, codexQuotaOverdraftSchedulingCtxKey{}, &codexQuotaOverdraftRequestState{})
}
// CodexQuotaOverdraftSchedulingEnabled reports whether the global switch and
@@ -56,18 +65,70 @@ func CodexQuotaOverdraftSchedulingEnabled(ctx context.Context) bool {
if !CodexQuotaOverdraftEnabled() || ctx == nil {
return false
}
enabled, _ := ctx.Value(codexQuotaOverdraftSchedulingCtxKey{}).(bool)
return enabled
return codexQuotaOverdraftRequestStateFromContext(ctx) != nil
}
func codexQuotaOverdraftSchedulingEnabled(ctx context.Context) bool {
return CodexQuotaOverdraftSchedulingEnabled(ctx)
}
func codexQuotaOverdraftRequestStateFromContext(ctx context.Context) *codexQuotaOverdraftRequestState {
if ctx == nil {
return nil
}
state, _ := ctx.Value(codexQuotaOverdraftSchedulingCtxKey{}).(*codexQuotaOverdraftRequestState)
return state
}
func markCodexQuotaOverdraftInjected(ctx context.Context, accountID int64) {
if accountID <= 0 {
return
}
if state := codexQuotaOverdraftRequestStateFromContext(ctx); state != nil {
state.injectedAccounts.Store(accountID, struct{}{})
}
}
func codexQuotaOverdraftWasInjected(ctx context.Context, accountID int64) bool {
if accountID <= 0 {
return false
}
state := codexQuotaOverdraftRequestStateFromContext(ctx)
if state == nil {
return false
}
_, ok := state.injectedAccounts.Load(accountID)
return ok
}
func codexQuotaOverdraftInjectionEligible(account *Account, now time.Time) bool {
if !isCodexQuotaOverdraftAccount(account) {
return false
}
state, _ := codexQuotaOverdraftStateFromAccount(account)
if state != nil && state.RecoverAt != nil && state.RecoverAt.After(now) {
switch state.Status {
case codexQuotaOverdraftProbePending, codexQuotaOverdraftProbePassed, codexQuotaOverdraftProbeInconclusive:
return true
case codexQuotaOverdraftProbeFailed:
return false
}
}
windowEligible := func(usedKey, resetKey string) bool {
if parseExtraFloat64(account.Extra[usedKey]) < codexQuotaOverdraftPrearmPercent {
return false
}
resetAt := codexQuotaOverdraftResetAt(account.Extra[resetKey], now)
return resetAt == nil || resetAt.After(now)
}
return windowEligible("codex_5h_used_percent", "codex_5h_reset_at") ||
windowEligible("codex_7d_used_percent", "codex_7d_reset_at")
}
func (s *OpenAIGatewayService) shouldInjectCodexQuotaOverdraft(ctx context.Context, account *Account, compact bool) bool {
return codexQuotaOverdraftSchedulingEnabled(ctx) && !compact &&
s != nil && s.cfg != nil && s.cfg.Gateway.CodexQuotaOverdraftEnabled &&
isCodexQuotaOverdraftAccount(account)
codexQuotaOverdraftInjectionEligible(account, time.Now().UTC())
}
func (s *OpenAIGatewayService) prepareCodexQuotaOverdraftBody(ctx context.Context, account *Account, compact bool, body []byte) []byte {
@@ -75,10 +136,14 @@ func (s *OpenAIGatewayService) prepareCodexQuotaOverdraftBody(ctx context.Contex
return body
}
updated, changed, _ := injectCodexQuotaOverdraft(body)
if !changed {
return body
if changed {
markCodexQuotaOverdraftInjected(ctx, account.ID)
return updated
}
return updated
if codexQuotaOverdraftBodyHasInjection(body) {
markCodexQuotaOverdraftInjected(ctx, account.ID)
}
return body
}
func (s *OpenAIGatewayService) prepareCodexQuotaOverdraftPayload(ctx context.Context, account *Account, payload map[string]any) map[string]any {
@@ -90,6 +155,11 @@ func (s *OpenAIGatewayService) prepareCodexQuotaOverdraftPayload(ctx context.Con
return payload
}
updated, changed, _ := injectCodexQuotaOverdraft(raw)
if changed {
markCodexQuotaOverdraftInjected(ctx, account.ID)
} else if codexQuotaOverdraftBodyHasInjection(raw) {
markCodexQuotaOverdraftInjected(ctx, account.ID)
}
if !changed {
return payload
}
@@ -110,6 +180,26 @@ type codexQuotaOverdraftInputItem struct {
CallID string `json:"call_id"`
}
func codexQuotaOverdraftBodyHasInjection(body []byte) bool {
var document codexQuotaOverdraftDocument
if len(body) == 0 || json.Unmarshal(body, &document) != nil {
return false
}
return codexQuotaOverdraftInputHasInjection(document.Input)
}
func codexQuotaOverdraftInputHasInjection(input []json.RawMessage) bool {
for _, raw := range input {
var item codexQuotaOverdraftInputItem
if err := json.Unmarshal(raw, &item); err == nil &&
item.Type == "custom_tool_call" &&
strings.HasPrefix(item.CallID, codexQuotaOverdraftCallIDPrefix) {
return true
}
}
return false
}
// injectCodexQuotaOverdraft appends the same no-op custom tool call pair used by
// cpa-account-config-manager. Unsupported request shapes fail open unchanged.
func injectCodexQuotaOverdraft(body []byte) ([]byte, bool, error) {
@@ -125,13 +215,8 @@ func injectCodexQuotaOverdraft(body []byte) ([]byte, bool, error) {
return body, false, nil
}
for _, raw := range document.Input {
var item codexQuotaOverdraftInputItem
if err := json.Unmarshal(raw, &item); err == nil &&
item.Type == "custom_tool_call" &&
strings.HasPrefix(item.CallID, codexQuotaOverdraftCallIDPrefix) {
return body, false, nil
}
if codexQuotaOverdraftInputHasInjection(document.Input) {
return body, false, nil
}
var last codexQuotaOverdraftInputItem
@@ -27,15 +27,12 @@ const (
codexQuotaOverdraftProbeInconclusive = "inconclusive"
codexQuotaOverdraftProbeRecovered = "recovered"
codexQuotaOverdraftProbeAttemptLimit = 5
codexQuotaOverdraftProbeAttemptLimit = 1
codexQuotaOverdraftProbeAttemptTimeout = 20 * time.Second
codexQuotaOverdraftProbePlanTimeout = codexQuotaOverdraftProbeAttemptLimit * codexQuotaOverdraftProbeAttemptTimeout
codexQuotaOverdraftProbeBodyLimit = 256 << 10
codexQuotaOverdraftPauseSource = "codex_quota_overdraft"
codexQuotaOverdraftPersistAttempts = 3
codexQuotaOverdraftAutomaticRetries = 3
codexQuotaOverdraftRetrySweepInterval = time.Minute
codexQuotaOverdraftRetrySweepLimit = 100
codexQuotaOverdraftFallbackModel = "gpt-5.5"
codexQuotaOverdraftCompatibilityModel = "gpt-5.4-mini"
@@ -89,14 +86,12 @@ type codexQuotaOverdraftFailedFinalizer interface {
FinalizeCodexQuotaOverdraftProbeFailed(context.Context, int64, *CodexQuotaOverdraftProbeState, time.Time, string) (bool, error)
}
type codexQuotaOverdraftDueProbeLister interface {
ListDueCodexQuotaOverdraftProbes(context.Context, time.Time, int) ([]*Account, error)
type codexQuotaOverdraftNonFailedPersister interface {
PersistCodexQuotaOverdraftProbeUnlessFailed(context.Context, int64, *CodexQuotaOverdraftProbeState) (bool, error)
}
type codexQuotaOverdraftRetryScheduler interface {
Schedule(string, time.Duration, func())
ScheduleRecurring(string, time.Duration, func())
Cancel(string)
type codexQuotaOverdraftPauseClearer interface {
ClearCodexQuotaOverdraftPauseIfState(context.Context, int64, string, string, string) (bool, bool, error)
}
// CodexQuotaOverdraftCoordinator implements the bounded real-request gate used
@@ -113,7 +108,6 @@ type CodexQuotaOverdraftCoordinator struct {
agentIdentityWS agentIdentityWSConnectionInvalidator
agentIdentityTaskMu sync.Mutex
running sync.Map
retryScheduler codexQuotaOverdraftRetryScheduler
now func() time.Time
probeAttemptForTest func(context.Context, *Account, string) codexQuotaOverdraftProbeResult
}
@@ -142,17 +136,6 @@ func NewCodexQuotaOverdraftCoordinator(
coordinator.agentIdentityWS, _ = runtimeBlocker.(agentIdentityWSConnectionInvalidator)
return coordinator
}
// SetRetryScheduler enables active retry timers and a small persisted-state
// sweep so inconclusive or stale pending probes resume after a restart.
func (c *CodexQuotaOverdraftCoordinator) SetRetryScheduler(scheduler codexQuotaOverdraftRetryScheduler) {
if c == nil || scheduler == nil {
return
}
c.retryScheduler = scheduler
scheduler.ScheduleRecurring("codex-overdraft:retry-sweep", codexQuotaOverdraftRetrySweepInterval, c.resumeDueProbes)
}
func (c *CodexQuotaOverdraftCoordinator) enabled() bool {
return c != nil && c.cfg != nil && c.cfg.Gateway.CodexQuotaOverdraftEnabled &&
c.accountRepo != nil && c.httpUpstream != nil
@@ -168,6 +151,77 @@ func (c *CodexQuotaOverdraftCoordinator) ObserveAccount(account *Account, prefer
go c.observeAccount(accountCopy, preferredModel)
}
// ObserveBusinessSuccess treats a successful request that actually carried the
// overdraft payload as stronger evidence than a synthetic probe.
func (c *CodexQuotaOverdraftCoordinator) ObserveBusinessSuccess(account *Account, preferredModel string) {
if !c.enabled() || !isCodexQuotaOverdraftAccount(account) || account.ID <= 0 {
return
}
accountCopy := cloneCodexQuotaOverdraftAccount(account)
go c.observeBusinessSuccess(accountCopy, preferredModel)
}
func (c *CodexQuotaOverdraftCoordinator) ObserveBusinessSuccessByID(accountID int64, preferredModel string) {
if !c.enabled() || accountID <= 0 {
return
}
go func() {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
account, err := c.accountRepo.GetByID(ctx, accountID)
cancel()
if err != nil || !isCodexQuotaOverdraftAccount(account) {
return
}
c.observeBusinessSuccess(account, preferredModel)
}()
}
func (c *CodexQuotaOverdraftCoordinator) observeBusinessSuccess(account *Account, preferredModel string) {
if account == nil {
return
}
now := c.currentTime()
current, _ := codexQuotaOverdraftStateFromAccount(account)
signal, exhausted := codexQuotaOverdraftSignalFromAccount(account, current, now)
if !exhausted {
return
}
if current != nil && current.Status == codexQuotaOverdraftProbePassed && codexQuotaOverdraftStateCoversSignal(current, signal) {
return
}
if current != nil && current.Status == codexQuotaOverdraftProbeFailed && codexQuotaOverdraftStateCoversSignal(current, signal) {
return
}
state := &CodexQuotaOverdraftProbeState{
Status: codexQuotaOverdraftProbePassed,
QuotaWindow: signal.Window,
CycleKey: signal.CycleKey,
Attempts: 1,
Limit: codexQuotaOverdraftProbeAttemptLimit,
Model: strings.TrimSpace(preferredModel),
ReasonCode: "business_request_ok",
StartedAt: now,
TestedAt: codexQuotaOverdraftTimePtr(now),
RecoverAt: codexQuotaOverdraftTimePtr(signal.RecoverAt),
FiveHourRecoverAt: cloneTimePtr(signal.FiveHourRecoverAt),
SevenDayRecoverAt: cloneTimePtr(signal.SevenDayRecoverAt),
ObservedRateLimit: cloneTimePtr(account.RateLimitResetAt),
}
carryCodexQuotaOverdraftWindowStarts(state, current, signal, now)
startCodexQuotaOverdraftWindows(state, signal, now)
if !c.persistNonFailedState(account.ID, state) {
return
}
mergeAccountExtra(account, map[string]any{CodexQuotaOverdraftProbeExtraKey: state})
c.clearQuotaPause(account.ID, state)
slog.Info("codex_quota_overdraft_business_passed",
"account_id", account.ID,
"model", state.Model,
"quota_window", signal.Window,
"cycle_key", signal.CycleKey,
)
}
func (c *CodexQuotaOverdraftCoordinator) observeAccount(account *Account, preferredModel string) {
if account == nil {
return
@@ -216,10 +270,82 @@ func (c *CodexQuotaOverdraftCoordinator) HandleQuota429(
if signal.CycleKey == "" {
return false
}
if codexQuotaOverdraftWasInjected(ctx, account.ID) {
return c.finishBusinessQuotaFailure(accountCopy, signal, preferredModel)
}
c.startProbe(accountCopy, signal, preferredModel)
return true
}
func (c *CodexQuotaOverdraftCoordinator) finishBusinessQuotaFailure(
account *Account,
signal codexQuotaOverdraftSignal,
preferredModel string,
) bool {
if account == nil || signal.CycleKey == "" {
return false
}
now := c.currentTime()
current, hasCurrent := codexQuotaOverdraftStateFromAccount(account)
if hasCurrent && codexQuotaOverdraftStateCoversSignal(current, signal) && current.Status == codexQuotaOverdraftProbeFailed {
return c.ensureFailedPause(account, current)
}
if !hasCurrent || !codexQuotaOverdraftStateCoversSignal(current, signal) {
current = &CodexQuotaOverdraftProbeState{
Status: codexQuotaOverdraftProbePending,
QuotaWindow: signal.Window,
CycleKey: signal.CycleKey,
Limit: codexQuotaOverdraftProbeAttemptLimit,
StartedAt: now,
RecoverAt: codexQuotaOverdraftTimePtr(signal.RecoverAt),
FiveHourRecoverAt: cloneTimePtr(signal.FiveHourRecoverAt),
SevenDayRecoverAt: cloneTimePtr(signal.SevenDayRecoverAt),
ObservedRateLimit: cloneTimePtr(account.RateLimitResetAt),
}
claimed, err := c.claimProbe(account.ID, current)
if err != nil {
slog.Warn("codex_quota_overdraft_business_failure_claim_failed", "account_id", account.ID, "error", err)
return false
}
if !claimed {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
latest, loadErr := c.accountRepo.GetByID(ctx, account.ID)
cancel()
latestState, ok := codexQuotaOverdraftStateFromAccount(latest)
if loadErr != nil || !ok || !codexQuotaOverdraftStateCoversSignal(latestState, signal) {
return false
}
account = latest
current = latestState
if current.Status == codexQuotaOverdraftProbeFailed {
return c.ensureFailedPause(account, current)
}
}
}
state := *current
state.Status = codexQuotaOverdraftProbeFailed
if state.Attempts < 1 {
state.Attempts = 1
}
state.Model = strings.TrimSpace(preferredModel)
state.ReasonCode = "business_quota_limited"
state.TestedAt = codexQuotaOverdraftTimePtr(now)
state.RetryAt = nil
state.RecoverAt = codexQuotaOverdraftTimePtr(signal.RecoverAt)
state.FiveHourRecoverAt = cloneTimePtr(signal.FiveHourRecoverAt)
state.SevenDayRecoverAt = cloneTimePtr(signal.SevenDayRecoverAt)
if !c.ensureFailedPause(account, &state) {
return false
}
slog.Warn("codex_quota_overdraft_business_exhausted",
"account_id", account.ID,
"model", state.Model,
"quota_window", signal.Window,
"recover_at", state.RecoverAt,
)
return true
}
func (c *CodexQuotaOverdraftCoordinator) startProbe(account *Account, signal codexQuotaOverdraftSignal, preferredModel string) {
if account == nil || signal.CycleKey == "" {
return
@@ -232,14 +358,8 @@ func (c *CodexQuotaOverdraftCoordinator) startProbe(account *Account, signal cod
case codexQuotaOverdraftProbeFailed:
c.ensureFailedPause(account, current)
return
case codexQuotaOverdraftProbePending:
if c.currentTime().Sub(current.StartedAt) < 2*time.Minute {
return
}
case codexQuotaOverdraftProbeInconclusive:
if current.RetryAt != nil && c.currentTime().Before(*current.RetryAt) {
return
}
case codexQuotaOverdraftProbePending, codexQuotaOverdraftProbeInconclusive:
return
}
}
@@ -255,9 +375,6 @@ func (c *CodexQuotaOverdraftCoordinator) startProbe(account *Account, signal cod
SevenDayRecoverAt: cloneTimePtr(signal.SevenDayRecoverAt),
ObservedRateLimit: cloneTimePtr(account.RateLimitResetAt),
}
if current != nil && codexQuotaOverdraftStateCoversSignal(current, signal) {
state.RetryCount = current.RetryCount
}
carryCodexQuotaOverdraftWindowStarts(state, current, signal, startedAt)
claimed, err := c.claimProbe(account.ID, state)
if err != nil {
@@ -302,7 +419,7 @@ func (c *CodexQuotaOverdraftCoordinator) runProbePlan(
account, err := c.accountRepo.GetByID(planCtx, accountID)
if err != nil || !isCodexQuotaOverdraftAccount(account) {
c.finishInconclusive(accountID, signal, preferredModel, state, "upstream_unavailable")
c.finishInconclusive(accountID, state, "upstream_unavailable")
return
}
models := codexQuotaOverdraftProbeModels(preferredModel)
@@ -310,7 +427,7 @@ func (c *CodexQuotaOverdraftCoordinator) runProbePlan(
quotaLimitedAttempts := 0
for attempt := 0; attempt < codexQuotaOverdraftProbeAttemptLimit; attempt++ {
if planCtx.Err() != nil {
c.finishInconclusive(accountID, signal, preferredModel, state, "request_timeout")
c.finishInconclusive(accountID, state, "request_timeout")
return
}
model := models[attempt%len(models)]
@@ -336,40 +453,40 @@ func (c *CodexQuotaOverdraftCoordinator) runProbePlan(
"reason", result.ReasonCode,
"quota_window", signal.Window,
)
c.persistState(accountID, state)
if result.Status == "available" {
state.Status = codexQuotaOverdraftProbePassed
state.ReasonCode = "model_response_ok"
state.RetryAt = nil
state.RetryCount = 0
startCodexQuotaOverdraftWindows(state, signal, testedAt)
c.persistState(accountID, state)
c.cancelRetry(accountID)
if !c.persistNonFailedState(accountID, state) {
return
}
c.clearQuotaPause(accountID, state)
slog.Info("codex_quota_overdraft_probe_passed", "account_id", accountID, "attempts", state.Attempts, "model", state.Model, "quota_window", signal.Window)
return
}
if result.Status == "inconclusive" {
c.finishInconclusive(accountID, signal, preferredModel, state, result.ReasonCode)
c.finishInconclusive(accountID, state, result.ReasonCode)
return
}
if result.Status == "authentication_failed" {
if c.rateLimitService != nil {
c.rateLimitService.HandleUpstreamError(context.Background(), account, result.StatusCode, result.Headers, result.Body)
}
c.finishInconclusive(accountID, signal, preferredModel, state, result.ReasonCode)
c.finishInconclusive(accountID, state, result.ReasonCode)
return
}
}
if quotaLimitedAttempts < codexQuotaOverdraftProbeAttemptLimit {
c.finishInconclusive(accountID, signal, preferredModel, state, lastReason)
c.finishInconclusive(accountID, state, lastReason)
return
}
state.Status = codexQuotaOverdraftProbeFailed
state.ReasonCode = lastReason
state.RetryAt = nil
c.cancelRetry(accountID)
state.RetryCount = 0
if !c.ensureFailedPause(account, state) {
return
}
@@ -378,150 +495,25 @@ func (c *CodexQuotaOverdraftCoordinator) runProbePlan(
func (c *CodexQuotaOverdraftCoordinator) finishInconclusive(
accountID int64,
signal codexQuotaOverdraftSignal,
preferredModel string,
state *CodexQuotaOverdraftProbeState,
reason string,
) {
now := c.currentTime().UTC()
retryable := reason != "authentication_failed" && reason != "account_deactivated"
state.Status = codexQuotaOverdraftProbeInconclusive
state.ReasonCode = reason
state.TestedAt = &now
state.RetryAt = nil
if retryable && state.RetryCount < codexQuotaOverdraftAutomaticRetries {
state.RetryCount++
retryAt := now.Add(codexQuotaOverdraftRetryDelay(state.RetryCount))
state.RetryAt = &retryAt
}
c.persistState(accountID, state)
if state.RetryAt != nil {
c.scheduleRetry(accountID, signal.CycleKey, preferredModel, *state.RetryAt)
state.RetryCount = 0
if !c.persistNonFailedState(accountID, state) {
return
}
slog.Warn("codex_quota_overdraft_probe_inconclusive",
"account_id", accountID,
"attempts", state.Attempts,
"reason", reason,
"retry_count", state.RetryCount,
"retry_at", state.RetryAt,
)
}
func codexQuotaOverdraftRetryDelay(retryCount int) time.Duration {
switch retryCount {
case 1:
return time.Minute
case 2:
return 3 * time.Minute
default:
return 10 * time.Minute
}
}
func codexQuotaOverdraftRetryTaskName(accountID int64) string {
return fmt.Sprintf("codex-overdraft:retry:%d", accountID)
}
func (c *CodexQuotaOverdraftCoordinator) scheduleRetry(accountID int64, cycleKey, preferredModel string, retryAt time.Time) {
if c == nil || c.retryScheduler == nil || accountID <= 0 {
return
}
delay := retryAt.Sub(c.currentTime())
if delay < time.Second {
delay = time.Second
}
name := codexQuotaOverdraftRetryTaskName(accountID)
c.retryScheduler.Schedule(name, delay, func() {
c.resumeProbe(accountID, cycleKey, preferredModel)
})
slog.Info("codex_quota_overdraft_retry_scheduled",
"account_id", accountID,
"cycle_key", cycleKey,
"retry_at", retryAt,
)
}
func (c *CodexQuotaOverdraftCoordinator) cancelRetry(accountID int64) {
if c == nil || c.retryScheduler == nil || accountID <= 0 {
return
}
c.retryScheduler.Cancel(codexQuotaOverdraftRetryTaskName(accountID))
}
func (c *CodexQuotaOverdraftCoordinator) resumeProbe(accountID int64, cycleKey, preferredModel string) {
if !c.enabled() || accountID <= 0 {
return
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
account, err := c.accountRepo.GetByID(ctx, accountID)
cancel()
if err != nil || !isCodexQuotaOverdraftAccount(account) {
slog.Warn("codex_quota_overdraft_retry_load_failed", "account_id", accountID, "error", err)
return
}
state, ok := codexQuotaOverdraftStateFromAccount(account)
if !ok || (cycleKey != "" && state.CycleKey != cycleKey) {
return
}
switch state.Status {
case codexQuotaOverdraftProbeInconclusive:
if state.RetryAt == nil {
return
}
if c.currentTime().Before(*state.RetryAt) {
c.scheduleRetry(accountID, state.CycleKey, preferredModel, *state.RetryAt)
return
}
case codexQuotaOverdraftProbePending:
staleAt := state.StartedAt.Add(2 * time.Minute)
if c.currentTime().Before(staleAt) {
c.scheduleRetry(accountID, state.CycleKey, preferredModel, staleAt)
return
}
default:
return
}
if strings.TrimSpace(preferredModel) == "" {
preferredModel = state.Model
}
slog.Info("codex_quota_overdraft_retry_started",
"account_id", accountID,
"cycle_key", state.CycleKey,
"retry_count", state.RetryCount,
)
c.observeAccount(account, preferredModel)
}
func (c *CodexQuotaOverdraftCoordinator) resumeDueProbes() {
if !c.enabled() {
return
}
lister, ok := c.accountRepo.(codexQuotaOverdraftDueProbeLister)
if !ok {
return
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
accounts, err := lister.ListDueCodexQuotaOverdraftProbes(ctx, c.currentTime(), codexQuotaOverdraftRetrySweepLimit)
cancel()
if err != nil {
slog.Warn("codex_quota_overdraft_retry_sweep_failed", "error", err)
return
}
for _, account := range accounts {
if account == nil {
continue
}
state, _ := codexQuotaOverdraftStateFromAccount(account)
cycleKey := ""
preferredModel := ""
if state != nil {
cycleKey = state.CycleKey
preferredModel = state.Model
}
go c.resumeProbe(account.ID, cycleKey, preferredModel)
}
}
func (c *CodexQuotaOverdraftCoordinator) runProbeAttempt(ctx context.Context, account *Account, model string) codexQuotaOverdraftProbeResult {
if c.probeAttemptForTest != nil {
return c.probeAttemptForTest(ctx, account, model)
@@ -697,6 +689,39 @@ func (c *CodexQuotaOverdraftCoordinator) persistState(accountID int64, state *Co
return false
}
// persistNonFailedState keeps an already-confirmed quota failure terminal for
// the current cycle, even when a concurrent success or probe finishes later.
func (c *CodexQuotaOverdraftCoordinator) persistNonFailedState(accountID int64, state *CodexQuotaOverdraftProbeState) bool {
if state == nil || state.Status == codexQuotaOverdraftProbeFailed {
return false
}
persister, ok := c.accountRepo.(codexQuotaOverdraftNonFailedPersister)
if !ok {
return c.persistState(accountID, state)
}
var (
persisted bool
lastErr error
)
for attempt := 1; attempt <= codexQuotaOverdraftPersistAttempts; attempt++ {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
persisted, lastErr = persister.PersistCodexQuotaOverdraftProbeUnlessFailed(ctx, accountID, state)
cancel()
if lastErr == nil {
return persisted
}
if attempt < codexQuotaOverdraftPersistAttempts {
time.Sleep(time.Duration(attempt) * 100 * time.Millisecond)
}
}
slog.Warn("codex_quota_overdraft_state_persist_failed",
"account_id", accountID,
"attempts", codexQuotaOverdraftPersistAttempts,
"error", lastErr,
)
return false
}
func (c *CodexQuotaOverdraftCoordinator) ensureFailedPause(account *Account, state *CodexQuotaOverdraftProbeState) bool {
if account == nil || state == nil || state.RecoverAt == nil || !state.RecoverAt.After(c.currentTime()) {
return false
@@ -706,7 +731,7 @@ func (c *CodexQuotaOverdraftCoordinator) ensureFailedPause(account *Account, sta
account.TempUnschedulableUntil != nil && !account.TempUnschedulableUntil.Before(*state.RecoverAt) {
return true
}
reason := BuildTempUnschedReasonPayload(codexQuotaOverdraftPauseSource, "five real overdraft probes confirmed quota exhaustion")
reason := BuildTempUnschedReasonPayload(codexQuotaOverdraftPauseSource, "injected overdraft request confirmed quota exhaustion")
if finalizer, ok := c.accountRepo.(codexQuotaOverdraftFailedFinalizer); ok {
finalized := false
var lastErr error
@@ -758,7 +783,6 @@ func (c *CodexQuotaOverdraftCoordinator) ensureFailedPause(account *Account, sta
return false
}
}
c.cancelRetry(account.ID)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if c.tempUnschedCache != nil {
@@ -766,7 +790,7 @@ func (c *CodexQuotaOverdraftCoordinator) ensureFailedPause(account *Account, sta
UntilUnix: state.RecoverAt.Unix(),
TriggeredAtUnix: c.currentTime().Unix(),
StatusCode: http.StatusTooManyRequests,
ErrorMessage: "five real overdraft probes confirmed quota exhaustion",
ErrorMessage: "injected overdraft request confirmed quota exhaustion",
})
}
if c.runtimeBlocker != nil {
@@ -777,17 +801,46 @@ func (c *CodexQuotaOverdraftCoordinator) ensureFailedPause(account *Account, sta
}
func (c *CodexQuotaOverdraftCoordinator) clearQuotaPause(accountID int64, state *CodexQuotaOverdraftProbeState) {
c.cancelRetry(accountID)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
account, err := c.accountRepo.GetByID(ctx, accountID)
if err != nil || account == nil {
return
}
persisted, ok := codexQuotaOverdraftStateFromAccount(account)
if !ok || state == nil || persisted.CycleKey != state.CycleKey || persisted.Status != state.Status ||
(state.Status != codexQuotaOverdraftProbePassed && state.Status != codexQuotaOverdraftProbeRecovered) {
return
}
expectedTempReason := ""
if IsAccountSchedulingThresholdReason(account.TempUnschedulableReason) || codexQuotaOverdraftPauseReason(account.TempUnschedulableReason) {
expectedTempReason = account.TempUnschedulableReason
}
if clearer, ok := c.accountRepo.(codexQuotaOverdraftPauseClearer); ok {
clearedRateLimit, clearedTemp, clearErr := clearer.ClearCodexQuotaOverdraftPauseIfState(
ctx,
accountID,
state.CycleKey,
state.Status,
expectedTempReason,
)
if clearErr != nil {
slog.Warn("codex_quota_overdraft_stale_rate_limit_clear_failed", "account_id", accountID, "error", clearErr)
return
}
if clearedRateLimit {
slog.Info("codex_quota_overdraft_stale_rate_limit_cleared", "account_id", accountID, "probe_status", state.Status)
}
if clearedTemp && c.tempUnschedCache != nil {
_ = c.tempUnschedCache.DeleteTempUnsched(ctx, accountID)
}
if (clearedRateLimit || clearedTemp) && c.runtimeBlocker != nil {
c.runtimeBlocker.ClearAccountSchedulingBlock(accountID)
}
return
}
clearedSchedulingState := false
probeProvedAvailable := state != nil &&
(state.Status == codexQuotaOverdraftProbePassed || state.Status == codexQuotaOverdraftProbeRecovered)
if probeProvedAvailable && account.RateLimitResetAt != nil {
if account.RateLimitResetAt != nil {
staleResetAt := account.RateLimitResetAt.UTC()
if err := c.accountRepo.ClearRateLimit(ctx, accountID); err != nil {
slog.Warn("codex_quota_overdraft_stale_rate_limit_clear_failed", "account_id", accountID, "reset_at", staleResetAt, "error", err)
@@ -796,7 +849,7 @@ func (c *CodexQuotaOverdraftCoordinator) clearQuotaPause(accountID int64, state
slog.Info("codex_quota_overdraft_stale_rate_limit_cleared", "account_id", accountID, "reset_at", staleResetAt, "probe_status", state.Status)
}
}
if IsAccountSchedulingThresholdReason(account.TempUnschedulableReason) || codexQuotaOverdraftPauseReason(account.TempUnschedulableReason) {
if expectedTempReason != "" {
_ = c.accountRepo.ClearTempUnschedulable(ctx, accountID)
clearedSchedulingState = true
if c.tempUnschedCache != nil {
@@ -853,15 +906,9 @@ func codexQuotaOverdraftProbeModels(preferred string) []string {
}
func codexQuotaOverdraftResponseIsQuotaLimited(headers http.Header, body []byte) bool {
if snapshot := ParseCodexRateLimitHeaders(headers); snapshot != nil {
if normalized := snapshot.Normalize(); normalized != nil &&
(normalized.Used5hPercent != nil && *normalized.Used5hPercent >= 100 ||
normalized.Used7dPercent != nil && *normalized.Used7dPercent >= 100) {
return true
}
}
var payload any
if len(bytes.TrimSpace(body)) > 0 && json.Unmarshal(body, &payload) == nil && codexQuotaOverdraftJSONHasQuotaEvidence(payload, 0) {
parsedPayload := len(bytes.TrimSpace(body)) > 0 && json.Unmarshal(body, &payload) == nil
if parsedPayload && codexQuotaOverdraftJSONHasQuotaEvidence(payload, 0) {
return true
}
text := strings.ToLower(strings.Join(strings.Fields(string(body)), " "))
@@ -878,6 +925,16 @@ func codexQuotaOverdraftResponseIsQuotaLimited(headers http.Header, body []byte)
return true
}
}
if parsedPayload && codexQuotaOverdraftJSONHasTransientRateLimitEvidence(payload, 0) {
return false
}
if snapshot := ParseCodexRateLimitHeaders(headers); snapshot != nil {
if normalized := snapshot.Normalize(); normalized != nil &&
(normalized.Used5hPercent != nil && *normalized.Used5hPercent >= 100 ||
normalized.Used7dPercent != nil && *normalized.Used7dPercent >= 100) {
return true
}
}
return false
}
@@ -944,6 +1001,37 @@ func codexQuotaOverdraftQuotaCode(value string) bool {
}
}
func codexQuotaOverdraftJSONHasTransientRateLimitEvidence(value any, depth int) bool {
if depth > 6 {
return false
}
switch typed := value.(type) {
case map[string]any:
for key, raw := range typed {
normalizedKey := strings.ToLower(strings.ReplaceAll(strings.TrimSpace(key), "-", "_"))
if normalizedKey == "type" || normalizedKey == "code" || normalizedKey == "reason" || normalizedKey == "error_code" {
if marker, ok := raw.(string); ok {
marker = strings.NewReplacer("-", "_", " ", "_").Replace(strings.ToLower(strings.TrimSpace(marker)))
switch marker {
case "rate_limit_exceeded", "too_many_requests", "request_rate_limited", "token_rate_limited":
return true
}
}
}
if codexQuotaOverdraftJSONHasTransientRateLimitEvidence(raw, depth+1) {
return true
}
}
case []any:
for _, item := range typed {
if codexQuotaOverdraftJSONHasTransientRateLimitEvidence(item, depth+1) {
return true
}
}
}
return false
}
func codexQuotaOverdraftSignalFromAccount(account *Account, state *CodexQuotaOverdraftProbeState, now time.Time) (codexQuotaOverdraftSignal, bool) {
if account == nil || len(account.Extra) == 0 {
return codexQuotaOverdraftSignal{}, false
@@ -1057,7 +1145,7 @@ func codexQuotaOverdraftStateFromAccount(account *Account) (*CodexQuotaOverdraft
}
func codexQuotaOverdraftSchedulingAllowed(account *Account, now time.Time) bool {
if !isCodexQuotaOverdraftAccount(account) {
if !codexQuotaOverdraftInjectionEligible(account, now) {
return false
}
state, _ := codexQuotaOverdraftStateFromAccount(account)
@@ -1068,12 +1156,12 @@ func codexQuotaOverdraftSchedulingAllowed(account *Account, now time.Time) bool
return state.Status != codexQuotaOverdraftProbeFailed
}
func codexQuotaOverdraftSnapshotExhausted(updates map[string]any) bool {
func codexQuotaOverdraftSnapshotPrearmReached(updates map[string]any) bool {
if len(updates) == 0 {
return false
}
return parseExtraFloat64(updates["codex_5h_used_percent"]) >= 100 ||
parseExtraFloat64(updates["codex_7d_used_percent"]) >= 100
return parseExtraFloat64(updates["codex_5h_used_percent"]) >= codexQuotaOverdraftPrearmPercent ||
parseExtraFloat64(updates["codex_7d_used_percent"]) >= codexQuotaOverdraftPrearmPercent
}
func applyCodexQuotaOverdraftUsage(
@@ -10,6 +10,7 @@ import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
@@ -26,28 +27,16 @@ type codexOverdraftRuntimeBlockerStub struct {
clearCalls int
}
type codexOverdraftRetrySchedulerStub struct {
scheduledName string
scheduledDelay time.Duration
scheduledFn func()
recurringName string
recurringDelay time.Duration
canceledName string
type codexOverdraftClaimConflictRepoStub struct {
*codexOverdraftProbeRepoStub
}
func (s *codexOverdraftRetrySchedulerStub) Schedule(name string, delay time.Duration, fn func()) {
s.scheduledName = name
s.scheduledDelay = delay
s.scheduledFn = fn
}
func (s *codexOverdraftRetrySchedulerStub) ScheduleRecurring(name string, delay time.Duration, _ func()) {
s.recurringName = name
s.recurringDelay = delay
}
func (s *codexOverdraftRetrySchedulerStub) Cancel(name string) {
s.canceledName = name
func (r *codexOverdraftClaimConflictRepoStub) ClaimCodexQuotaOverdraftProbe(
context.Context,
int64,
*CodexQuotaOverdraftProbeState,
) (bool, error) {
return false, nil
}
func (b *codexOverdraftRuntimeBlockerStub) BlockAccountScheduling(*Account, time.Time, string) {}
@@ -108,7 +97,7 @@ func newCodexOverdraftProbeTestAccount(now time.Time) *Account {
}
}
func TestCodexQuotaOverdraftProbeModelsMatchPluginRotation(t *testing.T) {
func TestCodexQuotaOverdraftProbeUsesOnePreferredModel(t *testing.T) {
models := codexQuotaOverdraftProbeModels("gpt-5.4")
require.Equal(t, []string{"gpt-5.4", "gpt-5.5", "gpt-5.4-mini"}, models)
@@ -116,7 +105,7 @@ func TestCodexQuotaOverdraftProbeModelsMatchPluginRotation(t *testing.T) {
for attempt := 0; attempt < codexQuotaOverdraftProbeAttemptLimit; attempt++ {
got = append(got, models[attempt%len(models)])
}
require.Equal(t, []string{"gpt-5.4", "gpt-5.5", "gpt-5.4-mini", "gpt-5.4", "gpt-5.5"}, got)
require.Equal(t, []string{"gpt-5.4"}, got)
}
func TestCodexQuotaOverdraftSignalsKeepFiveHourAndSevenDayCyclesSeparate(t *testing.T) {
@@ -141,18 +130,15 @@ func TestCodexQuotaOverdraftSignalsKeepFiveHourAndSevenDayCyclesSeparate(t *test
require.Contains(t, multiple.CycleKey, "|7d:")
}
func TestCodexQuotaOverdraftProbePassesOnAnySuccessfulAttempt(t *testing.T) {
func TestCodexQuotaOverdraftSingleProbePassesOnSuccess(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
coordinator := &CodexQuotaOverdraftCoordinator{accountRepo: repo, now: func() time.Time { return now }}
models := make([]string, 0, 3)
models := make([]string, 0, 1)
coordinator.probeAttemptForTest = func(_ context.Context, _ *Account, model string) codexQuotaOverdraftProbeResult {
models = append(models, model)
if len(models) == 3 {
return codexQuotaOverdraftProbeResult{Status: "available", ReasonCode: "model_response_ok", StatusCode: http.StatusOK, Model: model}
}
return codexQuotaOverdraftProbeResult{Status: "retry", ReasonCode: "quota_limited", StatusCode: http.StatusTooManyRequests, Model: model}
return codexQuotaOverdraftProbeResult{Status: "available", ReasonCode: "model_response_ok", StatusCode: http.StatusOK, Model: model}
}
signal, exhausted := codexQuotaOverdraftSignalFromAccount(account, nil, now)
require.True(t, exhausted)
@@ -161,13 +147,13 @@ func TestCodexQuotaOverdraftProbePassesOnAnySuccessfulAttempt(t *testing.T) {
coordinator.runProbePlan(account.ID, signal, "gpt-5.4", state)
require.Equal(t, codexQuotaOverdraftProbePassed, state.Status)
require.Equal(t, 3, state.Attempts)
require.Equal(t, []string{"gpt-5.4", "gpt-5.5", "gpt-5.4-mini"}, models)
require.Equal(t, 1, state.Attempts)
require.Equal(t, []string{"gpt-5.4"}, models)
require.NotNil(t, state.FiveHourStartedAt)
require.Zero(t, repo.tempPauseCalls)
}
func TestCodexQuotaOverdraftProbeRequiresFiveExplicitQuotaFailures(t *testing.T) {
func TestCodexQuotaOverdraftSingleProbeQuotaFailurePauses(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
@@ -190,6 +176,129 @@ func TestCodexQuotaOverdraftProbeRequiresFiveExplicitQuotaFailures(t *testing.T)
require.False(t, codexQuotaOverdraftSchedulingAllowed(account, now), "同周期确认失败后不得继续绕过限额")
}
func TestCodexQuotaOverdraftBusinessSuccessAtExhaustionPassesWithoutProbe(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
coordinator := &CodexQuotaOverdraftCoordinator{accountRepo: repo, now: func() time.Time { return now }}
coordinator.observeBusinessSuccess(account, "gpt-5.4")
state, ok := codexQuotaOverdraftStateFromAccount(account)
require.True(t, ok)
require.Equal(t, codexQuotaOverdraftProbePassed, state.Status)
require.Equal(t, "business_request_ok", state.ReasonCode)
require.Equal(t, "gpt-5.4", state.Model)
require.Equal(t, 1, state.Attempts)
require.Equal(t, 1, state.Limit)
require.NotNil(t, state.FiveHourStartedAt)
}
func TestCodexQuotaOverdraftBusinessSuccessCannotReplaceFailedCycle(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
signal, _ := codexQuotaOverdraftSignalFromAccount(account, nil, now)
failed := newCodexOverdraftPendingState(signal, now)
failed.Status = codexQuotaOverdraftProbeFailed
failed.ReasonCode = "business_quota_limited"
account.Extra[CodexQuotaOverdraftProbeExtraKey] = failed
repo := &codexOverdraftProbeRepoStub{account: account}
coordinator := &CodexQuotaOverdraftCoordinator{accountRepo: repo, now: func() time.Time { return now }}
coordinator.observeBusinessSuccess(account, "gpt-5.4")
state, ok := codexQuotaOverdraftStateFromAccount(account)
require.True(t, ok)
require.Equal(t, codexQuotaOverdraftProbeFailed, state.Status)
require.Empty(t, repo.states)
}
func TestCodexQuotaOverdraftInjectedBusinessQuota429FailsImmediately(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
coordinator := &CodexQuotaOverdraftCoordinator{
accountRepo: repo,
httpUpstream: &queuedHTTPUpstream{},
cfg: &config.Config{Gateway: config.GatewayConfig{
CodexQuotaOverdraftEnabled: true,
}},
now: func() time.Time { return now },
}
ctx := WithCodexQuotaOverdraftScheduling(context.Background())
markCodexQuotaOverdraftInjected(ctx, account.ID)
handled := coordinator.HandleQuota429(
ctx,
account,
nil,
[]byte(`{"error":{"type":"usage_limit_reached"}}`),
"gpt-5.4",
)
require.True(t, handled)
state, ok := codexQuotaOverdraftStateFromAccount(account)
require.True(t, ok)
require.Equal(t, codexQuotaOverdraftProbeFailed, state.Status)
require.Equal(t, "business_quota_limited", state.ReasonCode)
require.Equal(t, 1, repo.tempPauseCalls)
}
func TestCodexQuotaOverdraftTransient429DoesNotEnterQuotaCooldown(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
coordinator := &CodexQuotaOverdraftCoordinator{
accountRepo: repo,
httpUpstream: &queuedHTTPUpstream{},
cfg: &config.Config{Gateway: config.GatewayConfig{
CodexQuotaOverdraftEnabled: true,
}},
now: func() time.Time { return now },
}
ctx := WithCodexQuotaOverdraftScheduling(context.Background())
markCodexQuotaOverdraftInjected(ctx, account.ID)
headers := http.Header{}
headers.Set("x-codex-primary-used-percent", "100")
headers.Set("x-codex-primary-reset-after-seconds", "3600")
headers.Set("x-codex-primary-window-minutes", "300")
handled := coordinator.HandleQuota429(
ctx,
account,
headers,
[]byte(`{"error":{"type":"rate_limit_exceeded","message":"too many requests"}}`),
"gpt-5.4",
)
require.False(t, handled)
require.Zero(t, repo.tempPauseCalls)
_, hasState := codexQuotaOverdraftStateFromAccount(account)
require.False(t, hasState)
}
func TestCodexQuotaOverdraftBusinessFailureReloadsAfterClaimConflict(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
latest := newCodexOverdraftProbeTestAccount(now)
signal, _ := codexQuotaOverdraftSignalFromAccount(latest, nil, now)
passed := newCodexOverdraftPendingState(signal, now)
passed.Status = codexQuotaOverdraftProbePassed
latest.Extra[CodexQuotaOverdraftProbeExtraKey] = passed
baseRepo := &codexOverdraftProbeRepoStub{account: latest}
repo := &codexOverdraftClaimConflictRepoStub{codexOverdraftProbeRepoStub: baseRepo}
coordinator := &CodexQuotaOverdraftCoordinator{accountRepo: repo, now: func() time.Time { return now }}
stale := newCodexOverdraftProbeTestAccount(now)
delete(stale.Extra, CodexQuotaOverdraftProbeExtraKey)
handled := coordinator.finishBusinessQuotaFailure(stale, signal, "gpt-5.4")
require.True(t, handled)
state, ok := codexQuotaOverdraftStateFromAccount(latest)
require.True(t, ok)
require.Equal(t, codexQuotaOverdraftProbeFailed, state.Status)
require.Equal(t, 1, baseRepo.tempPauseCalls)
}
func TestCodexQuotaOverdraftFailedPauseIsIdempotentForPersistedCycle(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
@@ -232,13 +341,14 @@ func TestCodexQuotaOverdraftProbeInconclusiveNeverPauses(t *testing.T) {
require.Equal(t, codexQuotaOverdraftProbeInconclusive, state.Status)
require.Equal(t, 1, state.Attempts)
require.NotNil(t, state.RetryAt)
require.Nil(t, state.RetryAt)
require.Zero(t, state.RetryCount)
require.Zero(t, repo.tempPauseCalls)
})
}
}
func TestCodexQuotaOverdraftProbeRotatesModelsBeforeInconclusive(t *testing.T) {
func TestCodexQuotaOverdraftModelNotFoundIsInconclusiveWithoutRetry(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
@@ -255,74 +365,10 @@ func TestCodexQuotaOverdraftProbeRotatesModelsBeforeInconclusive(t *testing.T) {
require.Equal(t, codexQuotaOverdraftProbeInconclusive, state.Status)
require.Equal(t, codexQuotaOverdraftProbeAttemptLimit, state.Attempts)
require.Equal(t, 1, state.RetryCount)
require.Equal(t, []string{"gpt-5.4", "gpt-5.5", "gpt-5.4-mini", "gpt-5.4", "gpt-5.5"}, models)
require.Zero(t, repo.tempPauseCalls)
}
func TestCodexQuotaOverdraftInconclusiveSchedulesBackoffRetry(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
scheduler := &codexOverdraftRetrySchedulerStub{}
coordinator := &CodexQuotaOverdraftCoordinator{
accountRepo: repo,
retryScheduler: scheduler,
now: func() time.Time { return now },
}
signal, _ := codexQuotaOverdraftSignalFromAccount(account, nil, now)
state := newCodexOverdraftPendingState(signal, now)
coordinator.finishInconclusive(account.ID, signal, "gpt-5.4", state, "upstream_unavailable")
require.Equal(t, codexQuotaOverdraftProbeInconclusive, state.Status)
require.Equal(t, 1, state.RetryCount)
require.WithinDuration(t, now.Add(time.Minute), *state.RetryAt, time.Second)
require.Equal(t, codexQuotaOverdraftRetryTaskName(account.ID), scheduler.scheduledName)
require.Equal(t, time.Minute, scheduler.scheduledDelay)
require.NotNil(t, scheduler.scheduledFn)
state.RetryCount = 1
coordinator.finishInconclusive(account.ID, signal, "gpt-5.4", state, "upstream_unavailable")
require.Equal(t, 2, state.RetryCount)
require.Equal(t, 3*time.Minute, scheduler.scheduledDelay)
}
func TestCodexQuotaOverdraftAuthenticationFailureDoesNotScheduleQuotaRetry(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
scheduler := &codexOverdraftRetrySchedulerStub{}
coordinator := &CodexQuotaOverdraftCoordinator{accountRepo: repo, retryScheduler: scheduler, now: func() time.Time { return now }}
signal, _ := codexQuotaOverdraftSignalFromAccount(account, nil, now)
state := newCodexOverdraftPendingState(signal, now)
coordinator.finishInconclusive(account.ID, signal, "gpt-5.4", state, "authentication_failed")
require.Empty(t, scheduler.scheduledName)
require.Nil(t, state.RetryAt)
require.Zero(t, state.RetryCount)
}
func TestCodexQuotaOverdraftInconclusiveStopsAtAutomaticRetryLimit(t *testing.T) {
now := time.Date(2026, time.August, 13, 14, 0, 0, 0, time.UTC)
account := newCodexOverdraftProbeTestAccount(now)
repo := &codexOverdraftProbeRepoStub{account: account}
scheduler := &codexOverdraftRetrySchedulerStub{}
coordinator := &CodexQuotaOverdraftCoordinator{
accountRepo: repo,
retryScheduler: scheduler,
now: func() time.Time { return now },
}
signal, _ := codexQuotaOverdraftSignalFromAccount(account, nil, now)
state := newCodexOverdraftPendingState(signal, now)
state.RetryCount = codexQuotaOverdraftAutomaticRetries
coordinator.finishInconclusive(account.ID, signal, "gpt-5.4", state, "upstream_unavailable")
require.Equal(t, codexQuotaOverdraftAutomaticRetries, state.RetryCount)
require.Nil(t, state.RetryAt)
require.Empty(t, scheduler.scheduledName)
require.Equal(t, []string{"gpt-5.4"}, models)
require.Zero(t, repo.tempPauseCalls)
}
func TestCodexQuotaOverdraftNewWindowPreservesExistingWindowBaseline(t *testing.T) {
@@ -416,10 +462,14 @@ func TestCodexQuotaOverdraftAvailableStateClearsMismatchedStaleRateLimit(t *test
repo := &codexOverdraftProbeRepoStub{account: account}
blocker := &codexOverdraftRuntimeBlockerStub{}
coordinator := &CodexQuotaOverdraftCoordinator{accountRepo: repo, runtimeBlocker: blocker}
signal, _ := codexQuotaOverdraftSignalFromAccount(account, nil, now)
state := &CodexQuotaOverdraftProbeState{
Status: status,
CycleKey: signal.CycleKey,
QuotaWindow: signal.Window,
ObservedRateLimit: codexQuotaOverdraftTimePtr(now.Add(time.Hour)),
}
account.Extra[CodexQuotaOverdraftProbeExtraKey] = state
coordinator.clearQuotaPause(account.ID, state)
@@ -470,6 +520,14 @@ func TestClassifyCodexQuotaOverdraftProbeResponses(t *testing.T) {
require.Equal(t, "inconclusive", status)
require.Equal(t, "transient_failure", reason)
headers := http.Header{}
headers.Set("x-codex-primary-used-percent", "100")
headers.Set("x-codex-primary-reset-after-seconds", "3600")
headers.Set("x-codex-primary-window-minutes", "300")
status, reason = classifyCodexQuotaOverdraftProbe(http.StatusTooManyRequests, headers, []byte(`{"error":{"type":"rate_limit_exceeded","message":"too many requests"}}`))
require.Equal(t, "inconclusive", status)
require.Equal(t, "transient_failure", reason)
status, reason = classifyCodexQuotaOverdraftProbe(http.StatusServiceUnavailable, nil, nil)
require.Equal(t, "inconclusive", status)
require.Equal(t, "upstream_unavailable", reason)
@@ -18,11 +18,20 @@ func TestCodexQuotaOverdraftInjection(t *testing.T) {
SetCodexQuotaOverdraftEnabled(true)
ctx := WithCodexQuotaOverdraftScheduling(context.Background())
svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{CodexQuotaOverdraftEnabled: true}}}
oauth := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
oauth := &Account{
ID: 77,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_5h_used_percent": 95,
"codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339),
},
}
body := []byte(`{"model":"gpt-5.4","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}]}`)
updated := svc.prepareCodexQuotaOverdraftBody(ctx, oauth, false, body)
require.NotEqual(t, string(body), string(updated))
require.True(t, codexQuotaOverdraftWasInjected(ctx, oauth.ID))
var document codexQuotaOverdraftDocument
require.NoError(t, json.Unmarshal(updated, &document))
@@ -42,7 +51,15 @@ func TestCodexQuotaOverdraftInjection(t *testing.T) {
func TestCodexQuotaOverdraftInjectionGuards(t *testing.T) {
t.Cleanup(func() { SetCodexQuotaOverdraftEnabled(false) })
body := []byte(`{"input":[{"type":"message","role":"user"}]}`)
oauth := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
oauth := &Account{
ID: 78,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_5h_used_percent": 95,
"codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339),
},
}
enabledCfg := &config.Config{Gateway: config.GatewayConfig{CodexQuotaOverdraftEnabled: true}}
svc := &OpenAIGatewayService{cfg: enabledCfg}
@@ -50,6 +67,8 @@ func TestCodexQuotaOverdraftInjectionGuards(t *testing.T) {
require.Equal(t, string(body), string(svc.prepareCodexQuotaOverdraftBody(WithCodexQuotaOverdraftScheduling(context.Background()), oauth, false, body)))
SetCodexQuotaOverdraftEnabled(true)
underPrearm := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Extra: map[string]any{"codex_5h_used_percent": 94}}
require.Equal(t, string(body), string(svc.prepareCodexQuotaOverdraftBody(WithCodexQuotaOverdraftScheduling(context.Background()), underPrearm, false, body)), "95% 以下不能注入")
require.Equal(t, string(body), string(svc.prepareCodexQuotaOverdraftBody(context.Background(), oauth, false, body)), "未标记的端点不能注入")
require.Equal(t, string(body), string(svc.prepareCodexQuotaOverdraftBody(WithCodexQuotaOverdraftScheduling(context.Background()), oauth, true, body)), "compact 不能注入")
@@ -60,7 +79,15 @@ func TestCodexQuotaOverdraftInjectionGuards(t *testing.T) {
} {
require.Equal(t, string(body), string(svc.prepareCodexQuotaOverdraftBody(ctx, account, false, body)))
}
agentIdentity := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{openAIAuthModeCredentialKey: OpenAIAuthModeAgentIdentity}}
agentIdentity := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{openAIAuthModeCredentialKey: OpenAIAuthModeAgentIdentity},
Extra: map[string]any{
"codex_7d_used_percent": 95,
"codex_7d_reset_at": time.Now().Add(24 * time.Hour).Format(time.RFC3339),
},
}
require.NotEqual(t, string(body), string(svc.prepareCodexQuotaOverdraftBody(ctx, agentIdentity, false, body)), "Agent Identity 必须支持透支请求注入")
notUserLast := []byte(`{"input":[{"type":"message","role":"assistant"}]}`)
@@ -71,6 +98,23 @@ func TestCodexQuotaOverdraftInjectionGuards(t *testing.T) {
require.Equal(t, oversized, svc.prepareCodexQuotaOverdraftBody(ctx, oauth, false, oversized))
}
func TestCodexQuotaOverdraftSchedulingDoesNotBypassThresholdBelowPrearm(t *testing.T) {
t.Cleanup(func() { SetCodexQuotaOverdraftEnabled(false) })
SetCodexQuotaOverdraftEnabled(true)
ctx := WithCodexQuotaOverdraftScheduling(context.Background())
account := &Account{
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Extra: map[string]any{
"codex_5h_used_percent": 94,
},
}
quotaCtx := withOpenAIQuotaAutoPauseSettings(ctx, OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.9})
paused, _ := shouldAutoPauseOpenAIAccountByQuota(quotaCtx, account)
require.True(t, paused)
}
func TestCodexQuotaOverdraftSchedulingOnlyBypassesQuotaThresholds(t *testing.T) {
t.Cleanup(func() { SetCodexQuotaOverdraftEnabled(false) })
SetCodexQuotaOverdraftEnabled(true)
@@ -1028,7 +1028,8 @@ func (s *OpenAIGatewayService) updateCodexUsageSnapshot(ctx context.Context, acc
if len(updates) == 0 {
return
}
persistSnapshot := codexQuotaOverdraftSnapshotExhausted(updates) || s.getCodexSnapshotThrottle().Allow(accountID, now)
persistSnapshot := codexQuotaOverdraftSnapshotPrearmReached(updates) || s.getCodexSnapshotThrottle().Allow(accountID, now)
businessSuccess := codexQuotaOverdraftWasInjected(ctx, accountID)
go func() {
updateCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
@@ -1059,7 +1060,11 @@ func (s *OpenAIGatewayService) updateCodexUsageSnapshot(ctx context.Context, acc
account = current
}
mergeAccountExtra(account, updates)
s.codexQuotaOverdraft.ObserveAccount(account, "")
if businessSuccess {
s.codexQuotaOverdraft.observeBusinessSuccess(account, "")
} else {
s.codexQuotaOverdraft.observeAccount(account, "")
}
}()
}
@@ -486,6 +486,41 @@ func TestOpenAIGatewayService_UpdateCodexUsageSnapshot_ThrottlesExtraWrites(t *t
}
}
func TestOpenAIGatewayService_UpdateCodexUsageSnapshot_PrearmBypassesThrottle(t *testing.T) {
repo := &openAICodexSnapshotAsyncRepo{
updateExtraCh: make(chan map[string]any, 2),
}
svc := &OpenAIGatewayService{
accountRepo: repo,
codexSnapshotThrottle: newAccountWriteThrottle(time.Hour),
}
belowPrearm := &OpenAICodexUsageSnapshot{
PrimaryUsedPercent: ptrFloat64WS(94),
PrimaryResetAfterSeconds: ptrIntWS(3600),
PrimaryWindowMinutes: ptrIntWS(10080),
SecondaryUsedPercent: ptrFloat64WS(22),
SecondaryResetAfterSeconds: ptrIntWS(1200),
SecondaryWindowMinutes: ptrIntWS(300),
}
atPrearm := *belowPrearm
atPrearm.PrimaryUsedPercent = ptrFloat64WS(95)
svc.updateCodexUsageSnapshot(context.Background(), 778, belowPrearm)
select {
case <-repo.updateExtraCh:
case <-time.After(2 * time.Second):
t.Fatal("等待首次 codex 快照落库超时")
}
svc.updateCodexUsageSnapshot(context.Background(), 778, &atPrearm)
select {
case updates := <-repo.updateExtraCh:
require.Equal(t, 95.0, updates["codex_7d_used_percent"])
case <-time.After(2 * time.Second):
t.Fatal("95% 预热快照不应被写入节流拦截")
}
}
func ptrFloat64WS(v float64) *float64 { return &v }
func ptrIntWS(v int) *int { return &v }
-2
View File
@@ -231,7 +231,6 @@ func ProvideCodexQuotaOverdraftCoordinator(
tempUnschedCache TempUnschedCache,
runtimeBlocker AccountRuntimeBlocker,
rateLimitService *RateLimitService,
timingWheel *TimingWheelService,
) *CodexQuotaOverdraftCoordinator {
coordinator := NewCodexQuotaOverdraftCoordinator(
accountRepo,
@@ -243,7 +242,6 @@ func ProvideCodexQuotaOverdraftCoordinator(
runtimeBlocker,
rateLimitService,
)
coordinator.SetRetryScheduler(timingWheel)
if gateway, ok := runtimeBlocker.(*OpenAIGatewayService); ok {
gateway.SetCodexQuotaOverdraftCoordinator(coordinator)
}
@@ -763,19 +763,15 @@ const codexOverdraftStatus = computed(() => {
if (!probe) return null
const attempts = Math.max(0, probe.attempts || 0)
const limit = Math.max(1, probe.limit || 5)
const limit = Math.max(1, probe.limit || 1)
const windowLabel = probe.quota_window === 'multiple' ? '5h / 7d' : probe.quota_window
const testedAt = probe.tested_at ? new Date(probe.tested_at).toLocaleString() : ''
const recoverAt = probe.recover_at ? new Date(probe.recover_at).toLocaleString() : ''
const retryAt = probe.retry_at ? new Date(probe.retry_at).toLocaleString() : ''
const retryCount = Math.max(0, probe.retry_count || 0)
const titleParts = [windowLabel]
if (probe.model) titleParts.push(probe.model)
if (probe.reason_code) titleParts.push(probe.reason_code)
if (testedAt) titleParts.push(`${t('usage.overdraftTestedAt')}: ${testedAt}`)
if (recoverAt) titleParts.push(`${t('usage.overdraftRecoverAt')}: ${recoverAt}`)
if (retryAt) titleParts.push(`${t('usage.overdraftRetryAt')}: ${retryAt}`)
if (retryCount > 0) titleParts.push(`${t('usage.overdraftRetryCount')}: ${retryCount}`)
switch (probe.status) {
case 'pending':
@@ -615,7 +615,7 @@ describe('AccountUsageCell', () => {
expect(wrapper.text()).toContain('7d|100|106540000')
})
it('OpenAI OAuth 会显示五次探测确认的限额状态', async () => {
it('OpenAI OAuth 会显示单次确认的限额状态', async () => {
getUsage.mockResolvedValue({
five_hour: {
utilization: 100,
@@ -627,8 +627,8 @@ describe('AccountUsageCell', () => {
status: 'failed',
quota_window: '5h',
cycle_key: '5h:4076577600',
attempts: 5,
limit: 5,
attempts: 1,
limit: 1,
model: 'gpt-5.5',
reason_code: 'quota_limited',
started_at: '2099-03-07T10:00:00Z',
@@ -657,10 +657,10 @@ describe('AccountUsageCell', () => {
await flushPromises()
expect(wrapper.text()).toContain('usage.overdraftProbeFailed')
expect(wrapper.text()).toContain('5/5 · 5h')
expect(wrapper.text()).toContain('1/1 · 5h')
})
it('OpenAI OAuth 无法确认探测时会显示重试轮次和时间', async () => {
it('OpenAI OAuth 无法确认探测时不会显示自动重试信息', async () => {
getUsage.mockResolvedValue({
five_hour: {
utilization: 100,
@@ -672,7 +672,7 @@ describe('AccountUsageCell', () => {
quota_window: '5h',
cycle_key: '5h:4076577600',
attempts: 1,
limit: 5,
limit: 1,
model: 'gpt-5.4',
reason_code: 'upstream_unavailable',
retry_count: 2,
@@ -702,9 +702,10 @@ describe('AccountUsageCell', () => {
await flushPromises()
const status = wrapper.get('[title*="usage.overdraftRetryCount: 2"]')
const status = wrapper.get('[title*="upstream_unavailable"]')
expect(status.text()).toContain('usage.overdraftProbeInconclusive')
expect(status.attributes('title')).toContain('usage.overdraftRetryAt')
expect(status.text()).toContain('1/1 · 5h')
expect(status.attributes('title')).not.toContain('usage.overdraftRetryAt')
})
it('Key 账号会展示 today stats 徽章并带 A/U 提示', async () => {
@@ -284,8 +284,6 @@ export default {
accountBilled: 'Account billed',
overdraftActive: 'Overdraft active',
overdraftRecoverAt: 'Expected recovery',
overdraftRetryAt: 'Next retry',
overdraftRetryCount: 'Retry round',
overdraftTestedAt: 'Probed at',
overdraftProbePending: 'Overdraft probe running',
overdraftProbeFailed: 'Quota limit confirmed',
@@ -289,8 +289,6 @@ export default {
accountBilled: '账号计费',
overdraftActive: '透支中',
overdraftRecoverAt: '预计恢复',
overdraftRetryAt: '下次重试',
overdraftRetryCount: '重试轮次',
overdraftTestedAt: '探测时间',
overdraftProbePending: '透支探测中',
overdraftProbeFailed: '已确认限额',