mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
## 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>
87 lines
3.1 KiB
Python
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
|