headroom/tests/test_provider_codex_install.py
gglucass 8c00f7103c
fix(codex): poll /wham/usage for subscription limits (handshake no longer sends x-codex-* headers) (#924)
## Description

Codex's subscription usage window (primary/secondary rate-limit gauges)
stopped populating for ChatGPT-OAuth sessions. This PR restores it by
polling Codex's dedicated usage endpoint instead of relying on response
headers that are no longer sent.

### Why the previous approach no longer works

The existing code populates `CodexRateLimitState` from `x-codex-*`
rate-limit headers captured on the `/v1/responses` WebSocket handshake
(`update_from_headers` at WS accept). That worked when OpenAI returned
`x-codex-primary-used-percent`, `x-codex-primary-window-minutes`, etc.
on the handshake response.

OpenAI has since stopped sending those headers on the ChatGPT WebSocket
handshake. I confirmed this by faithfully replaying a real Plus-account
handshake (both `prewarm` and regular `request_kind`): no `x-codex-*`
headers come back on either. This matches OpenAI's own move to a
dedicated usage endpoint (`GET /backend-api/codex/usage` in CodexApi
mode) and reports such as openai/codex#14728. So `update_from_headers`
now runs on every accept but finds nothing to parse, and the window
silently stays empty.

The headers aren't coming back, so there is nothing to fix in the
parsing path. The data now lives behind a request we have to make
ourselves.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `subscription/codex_rate_limits.py`:
- `parse_codex_usage_payload()` / `update_from_usage_payload()` — map
the `GET /backend-api/wham/usage` JSON (`rate_limit.primary_window` /
`secondary_window` with `used_percent`, `limit_window_seconds`,
`reset_at`; `credits`; `rate_limit_reached_type`) into the existing
`CodexRateLimitState`. `limit_window_seconds` is converted to
window-minutes with the same round-up codex-rs uses (`(secs + 59) //
60`).
- `maybe_schedule_usage_poll()` — fire-and-forget, throttled to one
request per 60s, scoped to ChatGPT sessions (requires both a Bearer
token and `ChatGPT-Account-Id`; API-key traffic is skipped). Uses an
in-flight guard so concurrent accepts don't stack polls. Endpoint is
overridable via `HEADROOM_CODEX_USAGE_URL`.
- `proxy/handlers/openai.py`:
- At the Codex WS accept site, after the now-usually-empty
`update_from_headers` block, schedule the usage poll. Wrapped in
`contextlib.suppress` and fully non-blocking so it can never delay or
fail the WebSocket accept.

The old header-capture path is intentionally left in place as a no-cost
fallback in case OpenAI restores the headers.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed (live `/wham/usage` replay against a Plus
account returned HTTP 200 with the expected schema; payload fixture in
tests mirrors that real shape)

## Test Output

```
$ uv run pytest tests/test_codex_rate_limits.py -q
........................................                                 [100%]
41 passed in 0.16s

$ uv run ruff check headroom/subscription/codex_rate_limits.py headroom/proxy/handlers/openai.py tests/test_codex_rate_limits.py
All checks passed!

$ uv run mypy headroom/subscription/codex_rate_limits.py
Success: no issues found in 1 source file
```

## Additional Notes

- New tests cover: full-payload mapping, window-minutes round-up,
credits balance kept only when `has_credits`, promo object vs string,
empty payload returns `None`, missing `used_percent` skipped,
header-gating (requires Bearer + account-id), poll throttling, and
no-event-loop safety.
- Scoping to `ChatGPT-Account-Id` keeps the poll off API-key traffic,
and the 60s throttle plus in-flight guard bound it to at most one
lightweight GET per minute per running proxy.

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-12 17:03:14 -05:00

87 lines
3.1 KiB
Python

from __future__ import annotations
from pathlib import Path
from headroom.providers.codex.install import build_provider_section, codex_uses_chatgpt_auth
def test_codex_provider_section_omits_requires_openai_auth_by_default() -> None:
"""#406: the flag must default off (API-key users), and only on for OAuth.
Setting requires_openai_auth on a custom [model_providers.headroom] block
forces codex to demand OpenAI OAuth login for every headroom-routed request,
which breaks API-key users; so callers opt in explicitly for ChatGPT users.
"""
section = build_provider_section(port=8787, name="OpenAI via Headroom proxy")
assert 'name = "OpenAI via Headroom proxy"' in section
assert 'base_url = "http://127.0.0.1:8787/v1"' in section
assert "requires_openai_auth" not in section, (
f"requires_openai_auth must be absent by default; got:\n{section}"
)
assert "supports_websockets = true" in section
assert 'env_key = "OPENAI_API_KEY"' not in section
def test_codex_provider_section_emits_requires_openai_auth_when_flagged() -> None:
section = build_provider_section(
port=8787, name="OpenAI via Headroom proxy", requires_openai_auth=True
)
assert "requires_openai_auth = true" in section
def test_codex_uses_chatgpt_auth_true_for_chatgpt_mode(tmp_path: Path) -> None:
auth = tmp_path / "auth.json"
auth.write_text('{"auth_mode": "chatgpt"}', encoding="utf-8")
assert codex_uses_chatgpt_auth(auth) is True
def test_codex_uses_chatgpt_auth_true_for_account_id_without_mode(tmp_path: Path) -> None:
auth = tmp_path / "auth.json"
auth.write_text('{"tokens": {"account_id": "acct_1"}}', encoding="utf-8")
assert codex_uses_chatgpt_auth(auth) is True
def test_codex_uses_chatgpt_auth_false_for_api_key(tmp_path: Path) -> None:
auth = tmp_path / "auth.json"
auth.write_text('{"auth_mode": "apikey", "OPENAI_API_KEY": "sk-x"}', encoding="utf-8")
assert codex_uses_chatgpt_auth(auth) is False
def test_codex_uses_chatgpt_auth_false_for_missing_or_malformed(tmp_path: Path) -> None:
assert codex_uses_chatgpt_auth(tmp_path / "absent.json") is False
bad = tmp_path / "auth.json"
bad.write_text("not json", encoding="utf-8")
assert codex_uses_chatgpt_auth(bad) is False
def test_codex_uses_chatgpt_auth_false_for_non_dict_json(tmp_path: Path) -> None:
auth = tmp_path / "auth.json"
auth.write_text("[]", encoding="utf-8")
assert codex_uses_chatgpt_auth(auth) is False
def test_codex_uses_chatgpt_auth_false_for_empty_object(tmp_path: Path) -> None:
auth = tmp_path / "auth.json"
auth.write_text("{}", encoding="utf-8")
assert codex_uses_chatgpt_auth(auth) is False
def test_codex_provider_section_supports_custom_markers() -> None:
section = build_provider_section(
port=9100,
name="Headroom init proxy",
marker_start="# --- start ---",
marker_end="# --- end ---",
)
assert section.startswith("# --- start ---\n")
assert section.endswith("# --- end ---\n")
assert 'base_url = "http://127.0.0.1:9100/v1"' in section
assert 'env_key = "OPENAI_API_KEY"' not in section