mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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>
This commit is contained in:
gglucass
GitHub
parent
7ced77b6e7
commit
8c00f7103c
13 changed files with 768 additions and 34 deletions
14
codecov.yml
14
codecov.yml
|
|
@ -3,12 +3,26 @@ codecov:
|
|||
|
||||
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/**"
|
||||
|
|
|
|||
|
|
@ -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 ""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
{
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue