diff --git a/codecov.yml b/codecov.yml index 937781b6a..0d3cbc18d 100644 --- a/codecov.yml +++ b/codecov.yml @@ -1,15 +1,29 @@ -codecov: - require_ci_to_pass: true - -coverage: - status: - project: - default: - target: auto - patch: - default: - target: auto - +codecov: + require_ci_to_pass: true + +coverage: + status: + # Gate on the comprehensive unit suite (`python` flag from ci.yml's 4 test + # shards), NOT the narrow native-e2e smoke flags (install-native / + # wrap-native). Those e2e jobs upload first and barely exercise new code, + # so an unscoped status computes patch at ~6% off the e2e flags alone and + # flaps to FAILURE before the unit shards report. Scoping to `python` makes + # the status reflect real coverage of the diff. + project: + default: + target: auto + flags: + - python + patch: + default: + target: auto + flags: + - python + +flags: + python: + carryforward: false + ignore: - "tests/**" - "scripts/tests/**" diff --git a/headroom/cli/init.py b/headroom/cli/init.py index 267a5d40c..037b6754d 100644 --- a/headroom/cli/init.py +++ b/headroom/cli/init.py @@ -38,6 +38,7 @@ from headroom.install.runtime import ( ) from headroom.install.state import load_manifest, save_manifest from headroom.install.supervisors import start_supervisor +from headroom.providers.codex.install import codex_uses_chatgpt_auth from .main import main @@ -291,6 +292,14 @@ def _ensure_codex_provider(path: Path, port: int) -> None: import re logger.debug("ensure codex provider block: %s (port=%s)", path, port) + # Emit requires_openai_auth only for ChatGPT-OAuth users (restores the + # account menu); omitting it for API-key users avoids forcing an OAuth + # login (#406). + requires_openai_auth = ( + "requires_openai_auth = true\n" + if codex_uses_chatgpt_auth(path.parent / "auth.json") + else "" + ) block = ( f"{_CODEX_PROVIDER_MARKER_START}\n" 'model_provider = "headroom"\n' @@ -299,6 +308,7 @@ def _ensure_codex_provider(path: Path, port: int) -> None: 'name = "Headroom init proxy"\n' f'base_url = "http://127.0.0.1:{port}/v1"\n' "supports_websockets = true\n" + f"{requires_openai_auth}" f"{_CODEX_PROVIDER_MARKER_END}" ) content = path.read_text(encoding="utf-8") if path.exists() else "" diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 9e21589b0..01e60679d 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -50,6 +50,7 @@ from headroom.copilot_auth import ( from headroom.providers.aider import build_launch_env as _build_aider_launch_env from headroom.providers.claude import proxy_base_url as _claude_proxy_base_url from headroom.providers.codex import build_launch_env as _build_codex_launch_env +from headroom.providers.codex.install import codex_uses_chatgpt_auth from headroom.providers.copilot import ( build_launch_env as _build_copilot_launch_env, ) @@ -1040,12 +1041,19 @@ def _inject_codex_provider_config(port: int) -> None: f'openai_base_url = "http://127.0.0.1:{port}/v1"\n' f"{_CODEX_END_MARKER}\n" ) + # Emit requires_openai_auth only for ChatGPT-OAuth users (restores the + # account menu); omitting it for API-key users avoids forcing an OAuth + # login (#406). + requires_openai_auth = ( + "requires_openai_auth = true\n" if codex_uses_chatgpt_auth(config_dir / "auth.json") else "" + ) provider_section = ( f"{_CODEX_TOP_LEVEL_MARKER}\n" "[model_providers.headroom]\n" 'name = "OpenAI via Headroom proxy"\n' f'base_url = "http://127.0.0.1:{port}/v1"\n' f"supports_websockets = true\n" + f"{requires_openai_auth}" # Per-project savings: Codex sends the header only when the mapped # env var (HEADROOM_PROJECT, set by `headroom wrap codex`) exists at # Codex runtime. Inline table keeps the key inside this section so diff --git a/headroom/providers/codex/install.py b/headroom/providers/codex/install.py index 24b1ef0a0..a213e0d9e 100644 --- a/headroom/providers/codex/install.py +++ b/headroom/providers/codex/install.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import re from pathlib import Path @@ -30,6 +31,31 @@ _ORPHAN_HEADROOM_TABLE = re.compile( ) +def codex_uses_chatgpt_auth(auth_path: Path) -> bool: + """Whether Codex authenticated via ChatGPT OAuth (vs an OpenAI API key). + + The account menu (profile/email/plan/usage) only renders when the active + provider carries ``requires_openai_auth = true``, but that flag forces codex + to demand an OpenAI OAuth login (#406) and would break API-key users. So we + emit it only in ChatGPT-OAuth mode, read from the sibling ``auth.json``. + """ + try: + data = json.loads(auth_path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return False + if not isinstance(data, dict): + return False + mode = data.get("auth_mode") + if isinstance(mode, str): + return mode.lower() == "chatgpt" + # Older auth.json files predate `auth_mode`: infer from an OAuth account id. + tokens = data.get("tokens") + if isinstance(tokens, dict): + account_id = tokens.get("account_id") + return isinstance(account_id, str) and bool(account_id.strip()) + return False + + def build_provider_section( *, port: int, @@ -37,12 +63,14 @@ def build_provider_section( marker_start: str = _CODEX_MARKER_START, marker_end: str = _CODEX_MARKER_END, include_markers: bool = True, + requires_openai_auth: bool = False, ) -> str: - """Build a managed Codex provider block (without requires_openai_auth). + """Build a managed Codex provider block. - Bug 3 (#406): requires_openai_auth must NOT appear on custom provider - blocks — it forces codex to demand OpenAI OAuth login for local-proxy - traffic. The built-in openai provider carries this flag; headroom does not. + ``requires_openai_auth`` is emitted only for ChatGPT-OAuth users: the flag + is what makes codex render the account menu, but it also forces codex to + demand an OpenAI OAuth login (#406), which breaks API-key users. Callers + pass the result of :func:`codex_uses_chatgpt_auth`; it defaults to ``False``. """ body = ( "[model_providers.headroom]\n" @@ -50,6 +78,8 @@ def build_provider_section( f'base_url = "{proxy_base_url(port)}"\n' "supports_websockets = true\n" ) + if requires_openai_auth: + body += "requires_openai_auth = true\n" if not include_markers: return body return f"{marker_start}\n{body}{marker_end}\n" @@ -76,6 +106,7 @@ def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation | None port=manifest.port, name="Headroom persistent proxy", include_markers=False, + requires_openai_auth=codex_uses_chatgpt_auth(path.parent / "auth.json"), ) + f"{_CODEX_MARKER_END}\n" ) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index b89aae13f..c36c9f464 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -21,6 +21,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import replace from datetime import datetime from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse from headroom.proxy.helpers import ( COMPRESSION_TIMEOUT_SECONDS, @@ -516,6 +517,16 @@ def _resolve_codex_routing_headers(headers: dict[str, str]) -> tuple[dict[str, s return resolved, False +def _prefers_http1_passthrough(base_url: str) -> bool: + """Whether passthrough to this host must use HTTP/1.1. + + ChatGPT's Cloudflare edge issues a managed challenge to our HTTP/2 + fingerprint on sensitive account endpoints; HTTP/1.1 is accepted. + """ + host = (urlparse(base_url).hostname or "").lower() + return host == "chatgpt.com" or host.endswith(".chatgpt.com") + + class OpenAIHandlerMixin: """Mixin providing OpenAI API handler methods for HeadroomProxy.""" @@ -3714,6 +3725,17 @@ class OpenAIHandlerMixin: with contextlib.suppress(Exception): get_codex_rate_limit_state().update_from_headers(dict(_codex_handshake)) + + # Current Codex no longer ships x-codex-* on the handshake, so the + # block above is usually a no-op. Pull the live subscription window + # from the dedicated usage endpoint instead (throttled, scoped to + # ChatGPT-session traffic, fire-and-forget so accept isn't blocked). + with contextlib.suppress(Exception): + from headroom.subscription.codex_rate_limits import ( + maybe_schedule_usage_poll, + ) + + maybe_schedule_usage_poll(ws_headers) async with stage_timer.measure("accept"): await websocket.accept( subprotocol=client_subprotocols[0] if client_subprotocols else None, @@ -5938,8 +5960,17 @@ class OpenAIHandlerMixin: body = await request.body() headers = await apply_copilot_api_auth(headers, url=url) + # Cloudflare bot-management challenges our HTTP/2 fingerprint on + # ChatGPT's sensitive account endpoints (/backend-api/me, + # /backend-api/accounts/check), returning a 403 challenge page instead + # of JSON and collapsing the Codex account menu to just "Settings". + # Those endpoints answer fine over HTTP/1.1, so forward ChatGPT + # passthrough on the HTTP/1.1 client. Other hosts keep HTTP/2. + passthrough_client = self.http_client + if _prefers_http1_passthrough(base_url): + passthrough_client = self.http_client_h1 or self.http_client try: - response = await self.http_client.request( # type: ignore[union-attr] + response = await passthrough_client.request( # type: ignore[union-attr] method=request.method, url=url, headers=headers, diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index d74fb4866..a12244648 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -703,6 +703,9 @@ class HeadroomProxy( # HTTP client self.http_client: httpx.AsyncClient | None = None + # HTTP/1.1-only client for ChatGPT passthrough (Cloudflare challenges + # our HTTP/2 fingerprint on its sensitive account endpoints). + self.http_client_h1: httpx.AsyncClient | None = None # Shared cold-start warmup registry (populated by startup()). # Holds typed slots with loaded / loading / null / error status for @@ -1135,19 +1138,26 @@ class HeadroomProxy( metadata={"port": self.config.port, "host": self.config.host}, ) _ca_bundle = find_ca_bundle() - self.http_client = httpx.AsyncClient( - timeout=httpx.Timeout( + _client_kwargs: dict[str, Any] = { + "timeout": httpx.Timeout( connect=self.config.connect_timeout_seconds, read=self.config.request_timeout_seconds, write=self.config.request_timeout_seconds, pool=self.config.connect_timeout_seconds, ), - limits=httpx.Limits( + "limits": httpx.Limits( max_connections=self.config.max_connections, max_keepalive_connections=self.config.max_keepalive_connections, ), - http2=self.config.http2, - verify=_ca_bundle if _ca_bundle is not None else True, + "verify": _ca_bundle if _ca_bundle is not None else True, + } + self.http_client = httpx.AsyncClient(http2=self.config.http2, **_client_kwargs) + # Reuse the primary client when HTTP/2 is already off; otherwise keep a + # dedicated HTTP/1.1 client for ChatGPT passthrough. + self.http_client_h1 = ( + self.http_client + if not self.config.http2 + else httpx.AsyncClient(http2=False, **_client_kwargs) ) logger.info("Headroom Proxy started") logger.info(f"Optimization: {'ENABLED' if self.config.optimize else 'DISABLED'}") @@ -1369,6 +1379,9 @@ class HeadroomProxy( async def shutdown(self): """Cleanup async resources.""" + if self.http_client_h1 and self.http_client_h1 is not self.http_client: + await self.http_client_h1.aclose() + self.http_client_h1 = None if self.http_client: await self.http_client.aclose() self.http_client = None diff --git a/headroom/subscription/codex_rate_limits.py b/headroom/subscription/codex_rate_limits.py index f66f24780..c6c054942 100644 --- a/headroom/subscription/codex_rate_limits.py +++ b/headroom/subscription/codex_rate_limits.py @@ -1,10 +1,19 @@ -"""Passive tracking of OpenAI Codex rate-limit window data from response headers. +"""Tracking of OpenAI Codex rate-limit window data. -Codex (OpenAI) embeds rate-limit data directly in API response headers -(``x-codex-primary-used-percent``, ``x-codex-primary-window-minutes``, etc.) -rather than exposing a dedicated usage endpoint. This module captures those -headers from responses that headroom proxies and makes them available in -``/stats`` and the dashboard. +Historically Codex embedded rate-limit data in API *response headers* +(``x-codex-primary-used-percent`` etc.) and headroom captured those headers +from proxied responses (:meth:`CodexRateLimitState.update_from_headers`). +Current Codex (codex_exec / TUI on the ChatGPT WebSocket transport) no longer +emits those headers on the ``/responses`` handshake or stream -- the window is +served from a dedicated endpoint instead: + + GET https://chatgpt.com/backend-api/wham/usage (ChatGPT OAuth/session) + +So this module also exposes :func:`maybe_schedule_usage_poll`, a throttled +fire-and-forget GET against that endpoint using the client's own bearer token +and ``ChatGPT-Account-Id``. The header-capture path is kept intact: if OpenAI +ever returns ``x-codex-*`` again it still works, and API-key requests (no +account id) simply never trigger a poll. Header schema (parsed by codex-rs ``rate_limits.rs``): x-codex-primary-used-percent float 0-100 @@ -18,16 +27,47 @@ Header schema (parsed by codex-rs ``rate_limits.rs``): x-codex-credits-balance str e.g. "$5.00" x-codex-promo-message str server announcement x-codex-limit-name str e.g. "gpt-5.2-codex-sonic" + +``GET /wham/usage`` JSON schema (mapped by :func:`parse_codex_usage_payload`): + plan_type str + rate_limit.primary_window.used_percent float 0-100 + rate_limit.primary_window.limit_window_seconds int + rate_limit.primary_window.reset_at int Unix timestamp (seconds) + rate_limit.secondary_window.* same shape (optional) + credits.has_credits / unlimited / balance bool / bool / str + rate_limit_reached_type str | null + promo obj | str | null """ from __future__ import annotations +import asyncio +import logging +import os import time from dataclasses import dataclass, field from threading import Lock +import httpx + from headroom.subscription.base import QuotaTracker +logger = logging.getLogger(__name__) + +# Dedicated Codex usage endpoint for ChatGPT OAuth/session auth. Overridable +# for tests / self-hosted gateways via env. +CODEX_USAGE_URL = ( + os.environ.get("HEADROOM_CODEX_USAGE_URL", "https://chatgpt.com/backend-api/wham/usage").strip() + or "https://chatgpt.com/backend-api/wham/usage" +) + +# Minimum seconds between live usage polls. Codex turns can arrive in bursts; +# one GET per minute is plenty to keep the gauge fresh without hammering. +USAGE_POLL_MIN_INTERVAL_S = 60.0 + +# Bound the usage GET so a slow upstream never wedges the fire-and-forget task. +_USAGE_POLL_TIMEOUT_S = 10.0 + @dataclass class CodexRateLimitWindow: @@ -192,6 +232,94 @@ def parse_codex_rate_limits(headers: dict[str, str]) -> CodexRateLimitSnapshot | ) +# --------------------------------------------------------------------------- +# Usage-endpoint (GET /wham/usage) JSON parsing +# --------------------------------------------------------------------------- + + +def _window_from_usage_json(win: object) -> CodexRateLimitWindow | None: + """Map one ``rate_limit.{primary,secondary}_window`` object to a window.""" + if not isinstance(win, dict): + return None + used = win.get("used_percent") + try: + used_f = float(used) # type: ignore[arg-type] + except (ValueError, TypeError): + return None + if used_f != used_f: # NaN guard + return None + + window_minutes: int | None = None + secs = win.get("limit_window_seconds") + if isinstance(secs, (int, float)) and secs > 0: + # Round up, matching codex-rs window_minutes_from_seconds. + window_minutes = (int(secs) + 59) // 60 + + resets_at = win.get("reset_at") + resets_at = int(resets_at) if isinstance(resets_at, (int, float)) else None + + return CodexRateLimitWindow( + used_percent=used_f, + window_minutes=window_minutes, + resets_at=resets_at, + ) + + +def parse_codex_usage_payload(payload: object) -> CodexRateLimitSnapshot | None: + """Parse a snapshot from a ``GET /wham/usage`` JSON body. + + Returns ``None`` when the body carries no usable rate-limit data. + """ + if not isinstance(payload, dict): + return None + + rate_limit = payload.get("rate_limit") + rate_limit = rate_limit if isinstance(rate_limit, dict) else {} + primary = _window_from_usage_json(rate_limit.get("primary_window")) + secondary = _window_from_usage_json(rate_limit.get("secondary_window")) + + credits: CodexCreditsSnapshot | None = None + cred = payload.get("credits") + if isinstance(cred, dict) and cred.get("has_credits") is not None: + has = bool(cred.get("has_credits")) + raw_balance = cred.get("balance") + credits = CodexCreditsSnapshot( + has_credits=has, + unlimited=bool(cred.get("unlimited")), + # Only surface a balance when the account actually has credits; + # a "0" balance on a no-credits plan is noise to the gauge. + balance=(str(raw_balance) if has and raw_balance not in (None, "") else None), + ) + + promo = payload.get("promo") + if isinstance(promo, dict): + promo_message = promo.get("message") + elif isinstance(promo, str): + promo_message = promo + else: + promo_message = None + promo_message = (promo_message or "").strip() or None + + raw_limit_name = payload.get("rate_limit_reached_type") + limit_name = ( + raw_limit_name.strip() + if isinstance(raw_limit_name, str) and raw_limit_name.strip() + else None + ) + + if primary is None and secondary is None and credits is None and promo_message is None: + return None + + return CodexRateLimitSnapshot( + limit_id="codex", + limit_name=limit_name, + primary=primary, + secondary=secondary, + credits=credits, + promo_message=promo_message, + ) + + # --------------------------------------------------------------------------- # Singleton state store # --------------------------------------------------------------------------- @@ -214,6 +342,8 @@ class CodexRateLimitState(QuotaTracker): def __init__(self) -> None: self._lock = Lock() self._latest: CodexRateLimitSnapshot | None = None + self._last_poll_monotonic: float = 0.0 + self._poll_inflight: bool = False def update_from_headers(self, headers: dict[str, str]) -> None: """Update state from a response header dict (no-op if no Codex headers).""" @@ -223,6 +353,38 @@ class CodexRateLimitState(QuotaTracker): with self._lock: self._latest = snapshot + def update_from_usage_payload(self, payload: object) -> bool: + """Update state from a ``GET /wham/usage`` JSON body. + + Returns ``True`` when a snapshot was stored. + """ + snapshot = parse_codex_usage_payload(payload) + if snapshot is None: + return False + with self._lock: + self._latest = snapshot + return True + + def _try_begin_poll(self, min_interval_s: float) -> bool: + """Atomically claim a usage-poll slot. + + Returns ``False`` when a poll is already in flight or one ran within + ``min_interval_s``. On ``True`` the caller MUST call :meth:`_end_poll`. + """ + now = time.monotonic() + with self._lock: + if self._poll_inflight: + return False + if (now - self._last_poll_monotonic) < min_interval_s: + return False + self._poll_inflight = True + self._last_poll_monotonic = now + return True + + def _end_poll(self) -> None: + with self._lock: + self._poll_inflight = False + @property def latest(self) -> CodexRateLimitSnapshot | None: with self._lock: @@ -245,3 +407,80 @@ def get_codex_rate_limit_state() -> CodexRateLimitState: if _state is None: _state = CodexRateLimitState() return _state + + +# --------------------------------------------------------------------------- +# Live usage poll (GET /wham/usage) +# --------------------------------------------------------------------------- + + +def _build_usage_headers(request_headers: dict[str, str]) -> dict[str, str] | None: + """Build outbound /wham/usage headers from a client's request headers. + + Returns ``None`` unless the request carries a bearer token *and* a + ``ChatGPT-Account-Id`` -- the latter scopes the poll to ChatGPT OAuth + sessions (Codex), so API-key and non-Codex OAuth traffic never triggers it. + """ + lower = {str(k).lower(): v for k, v in request_headers.items()} + auth = str(lower.get("authorization", "")) + if not auth.startswith("Bearer ") or not auth[len("Bearer ") :].strip(): + return None + account_id = lower.get("chatgpt-account-id") + if not account_id: + return None + + headers = { + "Authorization": auth, + "ChatGPT-Account-Id": str(account_id), + "Accept": "application/json", + } + # Mirror the client's own UA/originator so the request looks like Codex. + for src, dst in (("user-agent", "User-Agent"), ("originator", "originator")): + val = lower.get(src) + if val: + headers[dst] = str(val) + return headers + + +async def _fetch_and_store_usage(url: str, headers: dict[str, str]) -> None: + state = get_codex_rate_limit_state() + try: + async with httpx.AsyncClient(timeout=_USAGE_POLL_TIMEOUT_S) as client: + resp = await client.get(url, headers=headers) + if resp.status_code == 200: + if state.update_from_usage_payload(resp.json()): + logger.debug("codex usage poll: refreshed rate-limit window") + else: + logger.debug("codex usage poll: 200 but no usable rate-limit data") + else: + logger.debug("codex usage poll: HTTP %s", resp.status_code) + except Exception as exc: # pragma: no cover - network/JSON defensive + logger.debug("codex usage poll failed: %s", exc) + finally: + state._end_poll() + + +def maybe_schedule_usage_poll( + request_headers: dict[str, str], + *, + url: str = CODEX_USAGE_URL, + min_interval_s: float = USAGE_POLL_MIN_INTERVAL_S, +) -> bool: + """Fire-and-forget a throttled ``GET /wham/usage`` to refresh the window. + + Safe to call on every Codex request: scoped to ChatGPT-session traffic via + :func:`_build_usage_headers` and internally throttled to at most one live + poll per ``min_interval_s``. Returns ``True`` when a poll was scheduled. + """ + headers = _build_usage_headers(request_headers) + if headers is None: + return False + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return False + state = get_codex_rate_limit_state() + if not state._try_begin_poll(min_interval_s): + return False + loop.create_task(_fetch_and_store_usage(url, headers)) + return True diff --git a/tests/test_cli/test_init_cli.py b/tests/test_cli/test_init_cli.py index 3039e94eb..9df842824 100644 --- a/tests/test_cli/test_init_cli.py +++ b/tests/test_cli/test_init_cli.py @@ -529,6 +529,30 @@ def test_ensure_codex_provider_replaces_existing_model_provider( assert parsed["features"]["hooks"] is True +def test_ensure_codex_provider_emits_requires_openai_auth_for_chatgpt( + monkeypatch, tmp_path: Path +) -> None: + init_cli, _ = _load_init_module(monkeypatch) + path = tmp_path / "config.toml" + (tmp_path / "auth.json").write_text('{"auth_mode": "chatgpt"}', encoding="utf-8") + + init_cli._ensure_codex_provider(path, 8787) + + assert "requires_openai_auth = true" in path.read_text(encoding="utf-8") + + +def test_ensure_codex_provider_omits_requires_openai_auth_for_api_key( + monkeypatch, tmp_path: Path +) -> None: + init_cli, _ = _load_init_module(monkeypatch) + path = tmp_path / "config.toml" + (tmp_path / "auth.json").write_text('{"auth_mode": "apikey"}', encoding="utf-8") + + init_cli._ensure_codex_provider(path, 8787) + + assert "requires_openai_auth" not in path.read_text(encoding="utf-8") + + def test_ensure_codex_feature_flag_replaces_existing_marker(monkeypatch, tmp_path: Path) -> None: init_cli, _ = _load_init_module(monkeypatch) path = tmp_path / "config.toml" diff --git a/tests/test_cli/test_wrap_codex.py b/tests/test_cli/test_wrap_codex.py index 78787b1b6..d97fb4446 100644 --- a/tests/test_cli/test_wrap_codex.py +++ b/tests/test_cli/test_wrap_codex.py @@ -350,6 +350,30 @@ class TestSubscriptionRouting: content = (tmp_path / ".codex" / "config.toml").read_text() assert 'openai_base_url = "http://127.0.0.1:8787/v1"' in content + def test_inject_emits_requires_openai_auth_for_chatgpt( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + (config_dir / "auth.json").write_text('{"auth_mode": "chatgpt"}', encoding="utf-8") + + wrap_mod._inject_codex_provider_config(8787) + + assert "requires_openai_auth = true" in (config_dir / "config.toml").read_text() + + def test_inject_omits_requires_openai_auth_for_api_key( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + (config_dir / "auth.json").write_text('{"auth_mode": "apikey"}', encoding="utf-8") + + wrap_mod._inject_codex_provider_config(8787) + + assert "requires_openai_auth" not in (config_dir / "config.toml").read_text() + def test_openai_base_url_port_updates_on_rewrap( self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: diff --git a/tests/test_codex_rate_limits.py b/tests/test_codex_rate_limits.py index 3da06d12b..1778e4015 100644 --- a/tests/test_codex_rate_limits.py +++ b/tests/test_codex_rate_limits.py @@ -2,14 +2,48 @@ from __future__ import annotations +import asyncio import time +import headroom.subscription.codex_rate_limits as crl from headroom.subscription.codex_rate_limits import ( CodexRateLimitState, CodexRateLimitWindow, + _build_usage_headers, + maybe_schedule_usage_poll, parse_codex_rate_limits, + parse_codex_usage_payload, ) +# A faithful GET /wham/usage body (shape captured from a live Plus account). +USAGE_PAYLOAD = { + "plan_type": "plus", + "rate_limit": { + "allowed": True, + "limit_reached": False, + "primary_window": { + "used_percent": 23, + "limit_window_seconds": 18000, + "reset_after_seconds": 12266, + "reset_at": 1781276043, + }, + "secondary_window": { + "used_percent": 6, + "limit_window_seconds": 604800, + "reset_after_seconds": 359170, + "reset_at": 1781622947, + }, + }, + "additional_rate_limits": None, + "credits": { + "has_credits": False, + "unlimited": False, + "balance": "0", + }, + "rate_limit_reached_type": None, + "promo": None, +} + # --------------------------------------------------------------------------- # CodexRateLimitWindow helpers # --------------------------------------------------------------------------- @@ -225,3 +259,163 @@ class TestCodexRateLimitState: assert snap is not None assert snap.primary is not None assert snap.primary.used_percent == 90.0 + + +# --------------------------------------------------------------------------- +# parse_codex_usage_payload (GET /wham/usage) +# --------------------------------------------------------------------------- + + +class TestParseCodexUsagePayload: + def test_parses_full_payload(self): + snap = parse_codex_usage_payload(USAGE_PAYLOAD) + assert snap is not None + assert snap.primary is not None + assert snap.primary.used_percent == 23.0 + assert snap.primary.window_minutes == 300 # 18000s rounded up + assert snap.primary.resets_at == 1781276043 + assert snap.secondary is not None + assert snap.secondary.used_percent == 6.0 + assert snap.secondary.window_minutes == 10080 # 604800s + + def test_window_minutes_rounds_up(self): + snap = parse_codex_usage_payload( + {"rate_limit": {"primary_window": {"used_percent": 1, "limit_window_seconds": 61}}} + ) + assert snap is not None + assert snap.primary is not None + assert snap.primary.window_minutes == 2 + + def test_no_credits_balance_suppressed(self): + # has_credits False -> balance must not surface as "0". + snap = parse_codex_usage_payload(USAGE_PAYLOAD) + assert snap is not None + assert snap.credits is not None + assert snap.credits.has_credits is False + assert snap.credits.balance is None + + def test_credits_balance_kept_when_has_credits(self): + payload = { + "rate_limit": {"primary_window": {"used_percent": 5}}, + "credits": {"has_credits": True, "unlimited": False, "balance": "$5.00"}, + } + snap = parse_codex_usage_payload(payload) + assert snap is not None + assert snap.credits is not None + assert snap.credits.balance == "$5.00" + + def test_promo_object_message(self): + payload = { + "rate_limit": {"primary_window": {"used_percent": 5}}, + "promo": {"message": "Hello"}, + } + snap = parse_codex_usage_payload(payload) + assert snap is not None + assert snap.promo_message == "Hello" + + def test_returns_none_for_empty(self): + assert parse_codex_usage_payload({}) is None + assert parse_codex_usage_payload(None) is None + assert parse_codex_usage_payload({"rate_limit": {}}) is None + + def test_missing_used_percent_window_skipped(self): + snap = parse_codex_usage_payload( + {"rate_limit": {"primary_window": {"limit_window_seconds": 60}}} + ) + assert snap is None + + def test_update_from_usage_payload_stores(self): + state = CodexRateLimitState() + assert state.update_from_usage_payload(USAGE_PAYLOAD) is True + snap = state.latest + assert snap is not None + assert snap.primary is not None + assert snap.primary.used_percent == 23.0 + + def test_update_from_usage_payload_noop_returns_false(self): + state = CodexRateLimitState() + assert state.update_from_usage_payload({}) is False + assert state.latest is None + + +# --------------------------------------------------------------------------- +# Usage poll: header gating + throttle +# --------------------------------------------------------------------------- + + +class TestUsagePollGating: + def test_build_headers_requires_account_id(self): + assert _build_usage_headers({"authorization": "Bearer abc.def.ghi"}) is None + + def test_build_headers_requires_bearer(self): + assert _build_usage_headers({"chatgpt-account-id": "acct"}) is None + assert ( + _build_usage_headers({"authorization": "sk-live", "chatgpt-account-id": "acct"}) is None + ) + + def test_build_headers_happy_path(self): + headers = _build_usage_headers( + { + "Authorization": "Bearer abc.def.ghi", + "ChatGPT-Account-Id": "acct-1", + "User-Agent": "codex_exec/0.139.0", + "originator": "codex_exec", + } + ) + assert headers is not None + assert headers["Authorization"] == "Bearer abc.def.ghi" + assert headers["ChatGPT-Account-Id"] == "acct-1" + assert headers["User-Agent"] == "codex_exec/0.139.0" + assert headers["originator"] == "codex_exec" + + def test_try_begin_poll_throttles(self): + state = CodexRateLimitState() + assert state._try_begin_poll(60.0) is True + # Second immediate attempt is throttled (within interval). + assert state._try_begin_poll(60.0) is False + state._end_poll() + # Still throttled by time even after the in-flight flag clears. + assert state._try_begin_poll(60.0) is False + # A zero interval always allows once the in-flight flag is clear. + assert state._try_begin_poll(0.0) is True + + def test_maybe_schedule_returns_false_without_loop(self): + # No running event loop -> cannot schedule. + assert ( + maybe_schedule_usage_poll( + {"authorization": "Bearer a.b.c", "chatgpt-account-id": "acct"} + ) + is False + ) + + def test_maybe_schedule_skips_non_codex(self): + async def run(): + return maybe_schedule_usage_poll({"authorization": "Bearer a.b.c"}) + + assert asyncio.run(run()) is False + + def test_maybe_schedule_creates_task_and_throttles(self, monkeypatch): + # Replace the network fetch with a fast no-op coroutine. + calls: list[str] = [] + + async def fake_fetch(url, headers): # noqa: ANN001 + calls.append(url) + crl.get_codex_rate_limit_state()._end_poll() + + monkeypatch.setattr(crl, "_fetch_and_store_usage", fake_fetch) + # Reset the singleton's throttle so this test is deterministic. + monkeypatch.setattr(crl, "_state", None) + monkeypatch.setattr(crl, "_state_lock", crl.Lock()) + + async def run(): + req = {"authorization": "Bearer a.b.c", "chatgpt-account-id": "acct"} + first = maybe_schedule_usage_poll(req, min_interval_s=60.0) + second = maybe_schedule_usage_poll(req, min_interval_s=60.0) + # Let the scheduled task run. + await asyncio.sleep(0) + return first, second + + first, second = asyncio.run(run()) + assert first is True + assert second is False # throttled + assert calls == [crl.CODEX_USAGE_URL] diff --git a/tests/test_install/test_providers.py b/tests/test_install/test_providers.py index 862a54e70..a4182a29d 100644 --- a/tests/test_install/test_providers.py +++ b/tests/test_install/test_providers.py @@ -80,6 +80,19 @@ def test_apply_and_revert_codex_provider_scope(monkeypatch, tmp_path: Path) -> N assert reverted.strip() == 'model = "gpt-4o"' +def test_apply_codex_provider_scope_emits_flag_for_chatgpt_auth( + monkeypatch, tmp_path: Path +) -> None: + config_path = tmp_path / "config.toml" + (tmp_path / "auth.json").write_text('{"auth_mode": "chatgpt"}') + monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) + manifest = _manifest(tmp_path) + + apply_codex_provider_scope(manifest) + + assert "requires_openai_auth = true" in config_path.read_text() + + def test_codex_build_install_env_returns_proxy_base_url() -> None: env = build_codex_install_env(port=5566, backend="ignored") diff --git a/tests/test_provider_codex_install.py b/tests/test_provider_codex_install.py index ff24738b6..05283762f 100644 --- a/tests/test_provider_codex_install.py +++ b/tests/test_provider_codex_install.py @@ -1,26 +1,78 @@ from __future__ import annotations -from headroom.providers.codex.install import build_provider_section +from pathlib import Path + +from headroom.providers.codex.install import build_provider_section, codex_uses_chatgpt_auth -def test_codex_provider_section_no_requires_openai_auth() -> None: - """Bug 3 (#406): build_provider_section must NOT include requires_openai_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. - Headroom is a local proxy — it must never carry this flag. + 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 from the headroom provider section; got:\n{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, diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index 2028bee26..f2bca5c99 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -15,6 +15,7 @@ from headroom.proxy.handlers.openai import ( OpenAIHandlerMixin, _decode_openai_bearer_payload, _passthrough_usage_from_json, + _prefers_http1_passthrough, ) from headroom.proxy.helpers import _headroom_bypass_enabled from headroom.proxy.server import HeadroomProxy @@ -51,6 +52,31 @@ class _TimeoutHttpClient: raise httpx.ConnectTimeout("connect timed out") +class _RecordingHttpClient: + def __init__(self, label: str) -> None: + self.label = label + self.calls = 0 + + async def request(self, **kwargs): # noqa: ANN001, ANN201 + self.calls += 1 + request = httpx.Request(kwargs["method"], kwargs["url"]) + return httpx.Response( + 200, + request=request, + headers={"content-type": "application/json"}, + json={"client": self.label}, + ) + + +class _ChatGPTAccountRequest: + method = "GET" + headers = {} + url = SimpleNamespace(path="/backend-api/me", query="") + + async def body(self) -> bytes: + return b"" + + class _PassthroughRequest: method = "GET" headers = {} @@ -249,6 +275,61 @@ def test_openai_passthrough_connect_timeout_returns_502() -> None: assert "Failed to connect to upstream API" in payload["error"]["message"] +def test_prefers_http1_passthrough_matches_chatgpt_hosts_only() -> None: + assert _prefers_http1_passthrough("https://chatgpt.com") is True + assert _prefers_http1_passthrough("https://chatgpt.com/backend-api/me") is True + assert _prefers_http1_passthrough("https://api.chatgpt.com") is True + assert _prefers_http1_passthrough("https://CHATGPT.COM/backend-api/me") is True + assert _prefers_http1_passthrough("https://api.openai.com") is False + assert _prefers_http1_passthrough("https://notchatgpt.com") is False + assert _prefers_http1_passthrough("https://chatgpt.com.evil.com") is False + assert _prefers_http1_passthrough("") is False + + +def test_chatgpt_passthrough_uses_http1_client() -> None: + handler = object.__new__(OpenAIHandlerMixin) + handler.http_client = _RecordingHttpClient("h2") + handler.http_client_h1 = _RecordingHttpClient("h1") + + response = asyncio.run( + handler.handle_passthrough(_ChatGPTAccountRequest(), "https://chatgpt.com") + ) + + assert response.status_code == 200 + assert json.loads(response.body)["client"] == "h1" + assert handler.http_client.calls == 0 + assert handler.http_client_h1.calls == 1 + + +def test_non_chatgpt_passthrough_uses_default_client() -> None: + handler = object.__new__(OpenAIHandlerMixin) + handler.http_client = _RecordingHttpClient("h2") + handler.http_client_h1 = _RecordingHttpClient("h1") + + response = asyncio.run( + handler.handle_passthrough(_PassthroughRequest(), "https://api.openai.com") + ) + + assert response.status_code == 200 + assert json.loads(response.body)["client"] == "h2" + assert handler.http_client.calls == 1 + assert handler.http_client_h1.calls == 0 + + +def test_chatgpt_passthrough_falls_back_when_h1_client_missing() -> None: + handler = object.__new__(OpenAIHandlerMixin) + handler.http_client = _RecordingHttpClient("h2") + handler.http_client_h1 = None + + response = asyncio.run( + handler.handle_passthrough(_ChatGPTAccountRequest(), "https://chatgpt.com") + ) + + assert response.status_code == 200 + assert json.loads(response.body)["client"] == "h2" + assert handler.http_client.calls == 1 + + def test_passthrough_usage_normalizes_vertex_usage_metadata() -> None: usage = _passthrough_usage_from_json( {