akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy): restore Codex usage headers on WS and streaming SSE transports (#577) (#794)

Browse source 
## Description

Codex's subscription/rate-limit window (the `x-codex-*` headers) was
being
**stripped on every transport Codex actually uses**, so session/weekly
usage
never reached the Codex CLI's own `/status` display, Headroom
`/stats`/dashboard,
or any consumer that sniffs the client-facing handshake. This PR
restores it on
**both** the WebSocket and streaming-SSE paths — the two halves of #577
— in one
place.

Fixes #577

**Supersedes #582 and #590.** This PR incorporates #582's SSE fix
(carried verbatim
with a `Co-authored-by` trailer) and additionally forwards the window
onto the client
`101` on the WS path, which #582/#590's capture-only WS code cannot do.
Both can be
closed as superseded once this merges — GitHub closing keywords only
auto-close
issues (hence `Fixes #577` above), not PRs, so #582/#590 need a manual
close.

### WebSocket (`gpt-5.4+`)

OpenAI delivers `x-codex-*` **only** on the upstream WS handshake
response, never
in data frames. `handle_openai_responses_ws` accepted the client WS
*before* it
connected upstream and never read `upstream.response.headers`, so the
window was
dropped. This reorders the handler to **connect upstream first**,
extract the
`x-codex-*` subset, then **accept the client WS with those headers
attached** to
the `101`, and refresh the Python state for `/stats` parity.

### Streaming SSE (incorporated from #582, @m16khb)

Codex CLI almost always streams. `streaming.py` neither captured
`x-codex-*` into
`CodexRateLimitState` nor forwarded it — the forwarded-header filter
matched only
the substring `"ratelimit"`, which `x-codex-*` does not contain. This
calls
`update_from_headers()` **before** the `>=400` early-return (so a
streaming 429/5xx
still refreshes the window, matching the non-streaming handlers) and
widens the
forward filter to pass `x-codex-*`.

> Credit: the SSE fix is @m16khb's work from #582, carried here verbatim
with a
> `Co-authored-by` trailer so the maintainer gets a single PR covering
both
> transports. This supersedes #582/#590's **WS** capture (which only
writes
> `/stats`); the connect-before-accept reorder additionally forwards the
window to
> the client `101`, which capture-only cannot do. #590's optional
snapshot
> persistence is intentionally left out (separable; hot-path sync write;
doesn't
> help the `101`-sniff consumers).

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `openai.py`: add `_extract_codex_handshake_headers()` (strictly
`x-codex-*`, via
`raw_items()` to avoid `MultipleValuesError`; never
`set-cookie`/`authorization`).
- `openai.py`: reorder `handle_openai_responses_ws` — connect-only retry
loop runs
before `accept()`; `accept(headers=...)` carries the forwarded window;
first
client frame read afterward. HTTP fallback preserved; it now also
refreshes
  `/stats` from the HTTP response headers.
- `streaming.py`: capture `x-codex-*` on all statuses + widen the
forwarded-header
  filter (from #582).

### Diff-size note

The bulk of the `openai.py` line count is **whitespace-only
relocation**: the relay
block dedents one level out of the old per-attempt `async with`. Logical
change is
~290 lines. **Review with `?w=1`.** In API-key mode the handshake
carries no
`x-codex-*`, so the accept-header list is empty and the path behaves
exactly as
before — the fix only activates for ChatGPT-subscription auth.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

- WS: `test_ws_connect_happens_before_accept`,
`test_ws_forwards_codex_headers_to_client_accept`
(only `x-codex-*` forwarded; `set-cookie`/`authorization` excluded;
`/stats` refreshed),
`test_ws_connect_failure_falls_back_to_http`,
`test_ws_first_frame_timeout_after_connect_closes_upstream`.
- Fallback: `test_fallback_refreshes_codex_rate_limit_state`.
- SSE:
`test_codex_rate_limit_headers_captured_and_forwarded_in_streaming`,
  `test_codex_rate_limit_captured_on_streaming_429` (from #582).
- Wire-level e2e: `tests/e2e_ws_codex_usage_headers.py` boots the real
proxy + fake
upstream + real `websockets` client and reads the client `101` — closes
the gap
the unit tests stub (that uvicorn/starlette actually write
`accept(headers=...)`).

## Test Output

```
$ uv run pytest tests/test_proxy_streaming_ratelimit_headers.py \
                tests/test_ws_http_fallback.py \
                tests/test_openai_codex_ws_lifecycle.py \
                tests/test_openai_codex_ws_timings.py \
                tests/test_codex_rate_limits.py -q
63 passed in 0.83s

$ .venv/bin/python tests/e2e_ws_codex_usage_headers.py
[codex-hdr-e2e] client 101 headers:
    x-codex-primary-used-percent: 42
    x-codex-primary-window-minutes: 300
    x-codex-secondary-used-percent: 7
    x-codex-secondary-window-minutes: 10080
[codex-hdr-e2e] /stats reflects codex window (primary-used=42)
=== CODEX-HDR E2E ALL GREEN ===

$ uv run ruff check . && uv run ruff format --check <touched files>
All checks passed!
```

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

- **Why connect-before-accept (not capture-only).** Once `accept()`
sends the `101`,
headers can no longer be added; the `x-codex-*` window only exists after
we connect
upstream. Capturing into Python state (as #582/#590's WS code does)
fixes `/stats`
but not the Codex CLI's native display or any `101`-sniffing consumer —
those need
  the headers *on the client handshake*, which requires the reorder.
- **Security.** Forwarding is filtered strictly to `x-codex-*`;
`set-cookie`,
`authorization`, and all other upstream headers are never forwarded to
the client
  (asserted by both the unit test and the e2e).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

## Contract Schemas

Per maintainer request: a JSON Schema (draft 2020-12) artifact
enshrining the OpenAI
interaction expectations this changeset relies on, so drift is
detectable later.

Committed following the repo's parity convention:
- schema:
`tests/parity/fixtures/codex_openai_contracts/codex-openai-interaction.schema.json`
- test: `tests/test_codex_openai_contract_parity.py` binds the schema to
the **live code**
in both directions, so drift fails CI rather than living only in this
description -
every declared `x-codex-*` header must be consumed by
`parse_codex_rate_limits`, and
`_extract_codex_handshake_headers` must forward exactly the declared
subset and never
`set-cookie`/`authorization`. No new dependency (does not pull in
`jsonschema`).

It covers, as `$defs`:

- `WSUpstreamHandshakeResponse` / `StreamingUpstreamResponseHeaders` -
the upstream
`x-codex-*` header family (full superset, with per-header wire pattern +
the parsed
semantic type) the WS and SSE captures read. Source of truth:
`parse_codex_rate_limits`.
- `ClientForwardedHandshakeHeaders` - the WS-101 **allow/deny**
contract: only
`x-codex-*` may be forwarded; `set-cookie`/`authorization` are
explicitly forbidden
  (`propertyNames` + `not`).
- `ClientForwardedStreamingHeaders` - the wider SSE forward set
(`*ratelimit*` OR `x-codex*`).
- `WSClientRequestFrame` / `WSRelayEvent` / `HTTPFallbackRequestBody` -
the WS frame
  envelopes and the unwrapped HTTP-fallback POST body.
- `CodexRateLimitStatsOutput` - the headroom `/stats` shape the parity
tests assert.

Validated with `jsonschema` (Draft202012 `check_schema` passes; positive
instances from
the e2e validate; negative instances - a leaked `set-cookie`, a fallback
body still
carrying a top-level `type` - are correctly rejected).

<details>
<summary><code>codex-openai-interaction.schema.json</code> (draft
2020-12)</summary>

```json
{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "$id": "https://github.com/chopratejas/headroom/contracts/codex-openai-interaction.schema.json",
  "title": "Codex <-> OpenAI interaction contracts (PR #794)",
  "description": "Enshrines the OpenAI interaction expectations this changeset depends on, so drift is detectable. Header values are transported as strings on the wire; the `x-headroom-parsed-type` annotation on each records the semantic type the parser (headroom/subscription/codex_rate_limits.py) coerces them to. Sources: codex_rate_limits.parse_codex_rate_limits (header family + gating), openai._extract_codex_handshake_headers (WS-101 forward filter), streaming.py (SSE forward filter).",
  "$defs": {
    "OpenAICodexWindowHeaders": {
      "title": "x-codex-*-{primary,secondary} window headers",
      "description": "A rolling rate-limit/subscription window. A window is materialized iff its `*-used-percent` header is present and numeric; `*-window-minutes` and `*-reset-at` are optional. `primary` and `secondary` are independent and either may be absent.",
      "type": "object",
      "properties": {
        "x-codex-primary-used-percent": {
          "type": "string",
          "pattern": "^\\d+(?:\\.\\d+)?$",
          "x-headroom-parsed-type": "float (0-100, NaN-guarded)",
          "description": "Percent of the primary window consumed. Gates creation of the primary window."
        },
        "x-codex-primary-window-minutes": {
          "type": "string",
          "pattern": "^\\d+$",
          "x-headroom-parsed-type": "int",
          "description": "Primary window size in minutes."
        },
        "x-codex-primary-reset-at": {
          "type": "string",
          "pattern": "^\\d+$",
          "x-headroom-parsed-type": "int (Unix epoch seconds)",
          "description": "Absolute reset time of the primary window."
        },
        "x-codex-secondary-used-percent": {
          "type": "string",
          "pattern": "^\\d+(?:\\.\\d+)?$",
          "x-headroom-parsed-type": "float (0-100, NaN-guarded)",
          "description": "Percent of the secondary window consumed. Gates creation of the secondary window."
        },
        "x-codex-secondary-window-minutes": {
          "type": "string",
          "pattern": "^\\d+$",
          "x-headroom-parsed-type": "int"
        },
        "x-codex-secondary-reset-at": {
          "type": "string",
          "pattern": "^\\d+$",
          "x-headroom-parsed-type": "int (Unix epoch seconds)"
        }
      },
      "additionalProperties": true
    },
    "OpenAICodexCreditsHeaders": {
      "title": "x-codex-credits-* headers",
      "description": "OpenAI credits balance. A credits snapshot is materialized iff `x-codex-credits-has-credits` is present; `unlimited` defaults to false; `balance` is optional.",
      "type": "object",
      "properties": {
        "x-codex-credits-has-credits": {
          "type": "string",
          "pattern": "^(?:[Tt][Rr][Uu][Ee]|[Ff][Aa][Ll][Ss][Ee]|[01])$",
          "x-headroom-parsed-type": "bool (true|false|1|0, case-insensitive)",
          "description": "Gates creation of the credits snapshot."
        },
        "x-codex-credits-unlimited": {
          "type": "string",
          "pattern": "^(?:[Tt][Rr][Uu][Ee]|[Ff][Aa][Ll][Ss][Ee]|[01])$",
          "x-headroom-parsed-type": "bool (defaults false when absent/unparseable)"
        },
        "x-codex-credits-balance": {
          "type": "string",
          "x-headroom-parsed-type": "str (empty -> null)",
          "description": "Free-form server string, e.g. \"$5.00\"."
        }
      },
      "additionalProperties": true
    },
    "OpenAICodexMetaHeaders": {
      "title": "x-codex meta headers",
      "type": "object",
      "properties": {
        "x-codex-limit-name": {
          "type": "string",
          "x-headroom-parsed-type": "str (empty -> null)",
          "description": "Active limit/model label, e.g. \"gpt-5.2-codex-sonic\"."
        },
        "x-codex-promo-message": {
          "type": "string",
          "x-headroom-parsed-type": "str (empty -> null)",
          "description": "Server announcement. Also gates snapshot creation when present."
        }
      },
      "additionalProperties": true
    },
    "OpenAICodexRateLimitHeaders": {
      "title": "Full x-codex-* header family OpenAI may emit",
      "description": "Superset of every x-codex-* header headroom reads. parse_codex_rate_limits returns a snapshot iff at least one of: a primary window, a secondary window, a credits snapshot, or a non-empty promo message is present; otherwise null (treated as a non-Codex response). All members are individually optional.",
      "type": "object",
      "allOf": [
        { "$ref": "#/$defs/OpenAICodexWindowHeaders" },
        { "$ref": "#/$defs/OpenAICodexCreditsHeaders" },
        { "$ref": "#/$defs/OpenAICodexMetaHeaders" }
      ],
      "additionalProperties": true
    },
    "WSUpstreamHandshakeResponse": {
      "title": "OpenAI WS handshake (101) response headers consumed by the WS fix",
      "description": "On the Codex WebSocket transport the x-codex-* window is delivered ONLY on the upstream handshake response (never in data frames). handle_openai_responses_ws reads upstream.response.headers here. This is the contract the connect-before-accept reorder depends on: if OpenAI ever moves these headers off the handshake (e.g. into a frame), the WS half of the fix goes stale.",
      "$ref": "#/$defs/OpenAICodexRateLimitHeaders"
    },
    "StreamingUpstreamResponseHeaders": {
      "title": "OpenAI streaming/HTTP response headers consumed by the SSE fix",
      "description": "On the streaming SSE/HTTP transport the same x-codex-* headers ride the HTTP response. streaming.py captures them on ALL statuses (including >=400) via update_from_headers, and forwards a wider set to the client (see ClientForwardedStreamingHeaders).",
      "$ref": "#/$defs/OpenAICodexRateLimitHeaders"
    },
    "ClientForwardedHandshakeHeaders": {
      "title": "Headers forwarded onto the CLIENT-facing WS 101 (allow/deny contract)",
      "description": "_extract_codex_handshake_headers forwards ONLY headers whose (lowercased) name starts with `x-codex-`. Every other upstream handshake header - notably set-cookie and authorization - MUST NOT appear on the client 101. Enforced by propertyNames below and asserted by the unit tests + tests/e2e_ws_codex_usage_headers.py.",
      "type": "object",
      "propertyNames": {
        "pattern": "^[Xx]-[Cc][Oo][Dd][Ee][Xx]-"
      },
      "not": {
        "anyOf": [
          { "required": ["set-cookie"] },
          { "required": ["Set-Cookie"] },
          { "required": ["authorization"] },
          { "required": ["Authorization"] }
        ]
      },
      "additionalProperties": { "type": "string" }
    },
    "ClientForwardedStreamingHeaders": {
      "title": "Headers forwarded to the client on the streaming SSE path",
      "description": "streaming.py forwards a header iff `\"ratelimit\" in name.lower()` OR `name.lower().startswith(\"x-codex\")`. This is a SUPERSET of the WS allow-list: it additionally passes generic *ratelimit* headers (e.g. the Anthropic streaming path) which do not contain the x-codex prefix.",
      "type": "object",
      "propertyNames": {
        "pattern": "(?:[Rr][Aa][Tt][Ee][Ll][Ii][Mm][Ii][Tt])|^[Xx]-[Cc][Oo][Dd][Ee][Xx]"
      },
      "additionalProperties": { "type": "string" }
    },
    "WSClientRequestFrame": {
      "title": "Client -> proxy WS data frame (Responses API over WS)",
      "description": "Codex sends the request as a response.create envelope. The HTTP fallback unwraps `.response` for the POST body, forces stream=true, and strips any top-level `type`. A flattened variant (no envelope, fields at top level) is also tolerated by the fallback.",
      "type": "object",
      "properties": {
        "type": { "const": "response.create" },
        "response": {
          "type": "object",
          "properties": {
            "model": { "type": "string", "description": "e.g. gpt-5.4" },
            "input": {
              "description": "String prompt or Responses-API structured input array.",
              "type": ["string", "array"]
            },
            "stream": { "type": "boolean" }
          },
          "required": ["model"],
          "additionalProperties": true
        }
      },
      "required": ["type", "response"],
      "additionalProperties": true
    },
    "WSRelayEvent": {
      "title": "proxy -> client WS data frame (relayed Responses API event)",
      "description": "SSE `data:` payloads relayed verbatim as WS text frames. `[DONE]` sentinels are dropped (not relayed). Every relayed event is a JSON object carrying a `type`. response.completed additionally carries usage under `response.usage`. anyOf (not oneOf): an error event also satisfies the looser lifecycle shape, which is fine.",
      "anyOf": [
        {
          "title": "lifecycle event",
          "type": "object",
          "properties": {
            "type": {
              "type": "string",
              "examples": [
                "response.created",
                "response.output_item.added",
                "response.completed"
              ]
            },
            "response": { "type": "object", "additionalProperties": true }
          },
          "required": ["type"],
          "additionalProperties": true
        },
        {
          "title": "error event",
          "type": "object",
          "properties": {
            "type": { "const": "error" },
            "error": {
              "type": "object",
              "properties": { "message": { "type": "string" } },
              "required": ["message"],
              "additionalProperties": true
            }
          },
          "required": ["type", "error"],
          "additionalProperties": true
        }
      ]
    },
    "HTTPFallbackRequestBody": {
      "title": "proxy -> OpenAI HTTP POST body on WS->HTTP fallback",
      "description": "Derived from WSClientRequestFrame: the inner `.response` object, with `stream` forced to true and any top-level `type` removed.",
      "type": "object",
      "properties": {
        "model": { "type": "string" },
        "stream": { "const": true },
        "input": { "type": ["string", "array"] }
      },
      "required": ["model", "stream"],
      "not": { "required": ["type"] },
      "additionalProperties": true
    },
    "CodexRateLimitStatsOutput": {
      "title": "headroom /stats output for the codex tracker (CodexRateLimitSnapshot.to_dict)",
      "description": "Internal (headroom-emitted) shape produced from the headers above; the WS and SSE update_from_headers parity tests assert this is refreshed. Included so drift in our own surface is also caught.",
      "type": "object",
      "properties": {
        "limit_id": { "const": "codex" },
        "limit_name": { "type": ["string", "null"] },
        "primary": { "$ref": "#/$defs/CodexWindowDict" },
        "secondary": { "$ref": "#/$defs/CodexWindowDict" },
        "credits": {
          "oneOf": [
            { "type": "null" },
            {
              "type": "object",
              "properties": {
                "has_credits": { "type": "boolean" },
                "unlimited": { "type": "boolean" },
                "balance": { "type": ["string", "null"] }
              },
              "required": ["has_credits", "unlimited", "balance"],
              "additionalProperties": false
            }
          ]
        },
        "promo_message": { "type": ["string", "null"] },
        "captured_at": { "type": "number", "description": "Unix epoch seconds (float)." }
      },
      "required": ["limit_id", "limit_name", "primary", "secondary", "credits", "promo_message", "captured_at"],
      "additionalProperties": false
    },
    "CodexWindowDict": {
      "oneOf": [
        { "type": "null" },
        {
          "type": "object",
          "properties": {
            "used_percent": { "type": "number" },
            "window_minutes": { "type": ["integer", "null"] },
            "window_label": { "type": "string", "description": "e.g. \"5h\", \"7d\"-style label; \"unknown\" when window_minutes is null." },
            "resets_at": { "type": ["integer", "null"], "description": "Unix epoch seconds." },
            "seconds_until_reset": { "type": ["integer", "null"] }
          },
          "required": ["used_percent", "window_minutes", "window_label", "resets_at", "seconds_until_reset"],
          "additionalProperties": false
        }
      ]
    }
  }
}
```

</details>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: m16khb <m16khb@gmail.com>
This commit is contained in:
gglucass 2026-06-09 22:55:53 +02:00 • committed by GitHub
parent 0b8b8d92de
commit 0ce68dedd7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 2115 additions and 975 deletions
Download patch file Download diff file Expand all files Collapse all files

1933
headroom/proxy/handlers/openai.py
View file

File diff suppressed because it is too large Load diff

26
headroom/proxy/handlers/streaming.py
View file

@ -967,6 +967,22 @@ class StreamingMixin:
return StreamingResponse(_error_gen(), media_type="text/event-stream")
# Capture Codex rate-limit window data from the upstream response
# headers, for *every* status. Codex (gpt-5.x) almost always streams, so
# without this the session/weekly windows surfaced in ``/stats`` and the
# dashboard would only refresh on the rare non-streaming reply. We do this
# *before* the error early-return below so a streaming 429/5xx — the moment
# usage is most relevant — still refreshes the windows, matching the
# non-streaming HTTP handlers which capture on all statuses.
# ``update_from_headers`` is a no-op when the response carries no
# ``x-codex-*`` headers (e.g. the Anthropic streaming path), so this is
# safe to call unconditionally.
from headroom.subscription.codex_rate_limits import (
get_codex_rate_limit_state,
)
get_codex_rate_limit_state().update_from_headers(dict(upstream_response.headers))
if upstream_response.status_code >= 400:
logger.warning(
"[%s] Forwarding upstream streaming error status=%s url=%s",
@ -1050,9 +1066,15 @@ class StreamingMixin:
headers=response_headers,
)
# Forward upstream ratelimit headers to the client
# Forward upstream rate-limit headers to the client. We pass both the
# generic ``*ratelimit*`` headers (Anthropic) and Codex's ``x-codex-*``
# window/credit headers — the latter do not contain the ``ratelimit``
# substring, so without the second clause the Codex CLI's own
# session/weekly display would stop updating on the streaming path.
forwarded_headers = {
k: v for k, v in upstream_response.headers.items() if "ratelimit" in k.lower()
k: v
for k, v in upstream_response.headers.items()
if "ratelimit" in k.lower() or k.lower().startswith("x-codex")
}
async def generate():

247
tests/e2e_ws_codex_usage_headers.py Normal file
View file

@ -0,0 +1,247 @@
"""End-to-end verification that Codex x-codex-* usage headers are forwarded
onto the client-facing WebSocket handshake (101).
Unit tests in tests/test_openai_codex_ws_lifecycle.py prove the handler
*logic* (it builds the right accept-header list), but they stub starlette's
``WebSocket.accept`` -- so they cannot prove the one risky assumption: that
starlette + uvicorn actually WRITE ``accept(headers=...)`` onto the real 101.
This e2e closes that gap with real wire traffic and no OpenAI quota:
1. Stand up a *fake upstream* WS server whose handshake response carries
several ``x-codex-*`` headers PLUS a ``set-cookie`` and an
``authorization`` header (which must NOT be forwarded).
2. Boot the real proxy pointed at the fake upstream via --openai-api-url.
3. Connect a real ``websockets`` client to the proxy and read
``client.response.headers`` -- i.e. the client-facing 101.
Asserts:
- every ``x-codex-*`` header from the upstream handshake is present on the
client 101 (original casing preserved),
- ``set-cookie`` and ``authorization`` are NOT forwarded,
- the proxy's /stats reflects the Codex window (update_from_headers parity).
Run via:
.venv/bin/python tests/e2e_ws_codex_usage_headers.py
"""
from __future__ import annotations
import asyncio
import json
import os
import socket
import subprocess
import sys
import time
import urllib.request
from pathlib import Path
import websockets
REPO_ROOT = Path(__file__).resolve().parent.parent
# The x-codex-* window the fake upstream advertises on its handshake.
UPSTREAM_CODEX_HEADERS = {
"x-codex-primary-used-percent": "42",
"x-codex-primary-window-minutes": "300",
"x-codex-secondary-used-percent": "7",
"x-codex-secondary-window-minutes": "10080",
}
# Sensitive headers that must NEVER reach the client 101.
UPSTREAM_LEAK_HEADERS = {
"set-cookie": "session=should-not-forward",
"authorization": "Bearer upstream-secret-should-not-forward",
}
def free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def wait_ready(port: int, timeout_s: float = 60.0) -> None:
deadline = time.time() + timeout_s
while time.time() < deadline:
try:
with urllib.request.urlopen(f"http://127.0.0.1:{port}/livez", timeout=2) as r:
if r.status == 200:
return
except Exception:
time.sleep(0.5)
raise TimeoutError("proxy not ready")
# ── Fake upstream WS server ─────────────────────────────────────────────
#
# The handshake response (101) carries the x-codex-* window plus sensitive
# headers, mirroring what OpenAI's Codex WS endpoint returns.
class FakeUpstream:
def __init__(self) -> None:
self.server: websockets.asyncio.server.Server | None = None
self.port: int = 0
async def _handler(self, ws):
try:
async for msg in ws:
if isinstance(msg, str):
await ws.send(
json.dumps(
{
"type": "response.completed",
"response": {"id": "fake_resp", "output": []},
}
)
)
except websockets.exceptions.ConnectionClosed:
pass
def _process_response(self, connection, request, response):
# Inject the x-codex-* window + sensitive headers onto the 101.
for name, value in {**UPSTREAM_CODEX_HEADERS, **UPSTREAM_LEAK_HEADERS}.items():
response.headers[name] = value
return response
async def start(self) -> int:
self.port = free_port()
self.server = await websockets.serve(
self._handler,
"127.0.0.1",
self.port,
process_response=self._process_response,
)
return self.port
async def stop(self) -> None:
if self.server:
self.server.close()
await self.server.wait_closed()
async def main_async() -> int:
fake = FakeUpstream()
upstream_port = await fake.start()
upstream_url = f"http://127.0.0.1:{upstream_port}"
proxy_port = free_port()
print(f"[codex-hdr-e2e] fake upstream at ws://127.0.0.1:{upstream_port}")
print(f"[codex-hdr-e2e] starting proxy on :{proxy_port}")
log_fp = open("/tmp/e2e_ws_codex_headers_proxy.log", "w")
proc = subprocess.Popen(
[
str(REPO_ROOT / ".venv/bin/headroom"),
"proxy",
"--port",
str(proxy_port),
"--no-telemetry",
"--openai-api-url",
upstream_url,
],
env={
**os.environ,
"OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "sk-fake-for-test"),
"ANTHROPIC_API_KEY": os.environ.get("ANTHROPIC_API_KEY", "sk-ant-fake"),
},
stdout=log_fp,
stderr=subprocess.STDOUT,
cwd=str(REPO_ROOT),
)
failures: list[str] = []
try:
wait_ready(proxy_port)
print("[codex-hdr-e2e] proxy ready")
proxy_ws_url = f"ws://127.0.0.1:{proxy_port}/v1/responses"
# NOTE: API-key auth (no ChatGPT-Account-ID) so the upstream routes to
# --openai-api-url (our fake upstream). The x-codex forwarding code is
# auth-mode-agnostic -- it forwards whatever x-codex-* headers the
# upstream handshake returns -- so this exercises the exact same path
# the real chatgpt.com subscription handshake would, without leaving
# localhost.
async with websockets.connect(
proxy_ws_url,
additional_headers={
"Authorization": "Bearer sk-fake",
"OpenAI-Beta": "responses_websockets=2026-02-06",
},
) as ws:
# The client-facing 101 response headers — the thing under test.
client_101 = {k.lower(): v for k, v in ws.response.headers.raw_items()}
print("[codex-hdr-e2e] client 101 headers:")
for k, v in sorted(client_101.items()):
if k.startswith("x-codex-") or k in ("set-cookie", "authorization"):
print(f" {k}: {v}")
# 1. Every x-codex-* header forwarded.
for name, value in UPSTREAM_CODEX_HEADERS.items():
if client_101.get(name) != value:
failures.append(
f"x-codex header not forwarded to client 101: "
f"{name} (got {client_101.get(name)!r}, want {value!r})"
)
# 2. Sensitive headers NOT forwarded.
for name in UPSTREAM_LEAK_HEADERS:
if name in client_101:
failures.append(f"sensitive header leaked to client 101: {name}")
# Drive one frame so the session is real (upstream echoes completed).
await ws.send(
json.dumps(
{
"type": "response.create",
"response": {"model": "gpt-5.4", "input": "hi"},
}
)
)
await asyncio.sleep(1.0)
# 3. /stats reflects the Codex window (update_from_headers parity).
# Secondary, shape-dependent signal: look for the actual forwarded
# value, not just the word "codex" (which can appear as a null key).
await asyncio.sleep(0.5)
try:
with urllib.request.urlopen(f"http://127.0.0.1:{proxy_port}/stats", timeout=5) as r:
stats_text = r.read().decode("utf-8", errors="replace")
want = UPSTREAM_CODEX_HEADERS["x-codex-primary-used-percent"]
if want in stats_text:
print(f"[codex-hdr-e2e] /stats reflects codex window (primary-used={want})")
else:
print(
"[codex-hdr-e2e] note: /stats did not surface the codex value "
f"({want!r}); shape may differ — 101 forwarding is the primary check"
)
except Exception as exc: # noqa: BLE001 - best-effort secondary check
print(f"[codex-hdr-e2e] /stats check skipped: {exc}")
finally:
print("[codex-hdr-e2e] terminating proxy")
proc.terminate()
try:
proc.wait(timeout=10)
except subprocess.TimeoutExpired:
proc.kill()
log_fp.close()
await fake.stop()
if failures:
print("\n=== CODEX-HDR E2E FAILURES ===")
for f in failures:
print(" -", f)
return 1
print("\n=== CODEX-HDR E2E ALL GREEN ===")
return 0
def main() -> int:
return asyncio.run(main_async())
if __name__ == "__main__":
sys.exit(main())

259
tests/parity/fixtures/codex_openai_contracts/codex-openai-interaction.schema.json Normal file
View file

@ -0,0 +1,259 @@
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"$id": "https://github.com/chopratejas/headroom/blob/main/tests/parity/fixtures/codex_openai_contracts/codex-openai-interaction.schema.json",
"title": "Codex <-> OpenAI interaction contracts (PR #794)",
"description": "Enshrines the OpenAI interaction expectations this changeset depends on, so drift is detectable. Header values are transported as strings on the wire; the `x-headroom-parsed-type` annotation on each records the semantic type the parser (headroom/subscription/codex_rate_limits.py) coerces them to. Sources: codex_rate_limits.parse_codex_rate_limits (header family + gating), openai._extract_codex_handshake_headers (WS-101 forward filter), streaming.py (SSE forward filter).",
"$defs": {
"OpenAICodexWindowHeaders": {
"title": "x-codex-*-{primary,secondary} window headers",
"description": "A rolling rate-limit/subscription window. A window is materialized iff its `*-used-percent` header is present and numeric; `*-window-minutes` and `*-reset-at` are optional. `primary` and `secondary` are independent and either may be absent.",
"type": "object",
"properties": {
"x-codex-primary-used-percent": {
"type": "string",
"pattern": "^\\d+(?:\\.\\d+)?$",
"x-headroom-parsed-type": "float (0-100, NaN-guarded)",
"description": "Percent of the primary window consumed. Gates creation of the primary window."
},
"x-codex-primary-window-minutes": {
"type": "string",
"pattern": "^\\d+$",
"x-headroom-parsed-type": "int",
"description": "Primary window size in minutes."
},
"x-codex-primary-reset-at": {
"type": "string",
"pattern": "^\\d+$",
"x-headroom-parsed-type": "int (Unix epoch seconds)",
"description": "Absolute reset time of the primary window."
},
"x-codex-secondary-used-percent": {
"type": "string",
"pattern": "^\\d+(?:\\.\\d+)?$",
"x-headroom-parsed-type": "float (0-100, NaN-guarded)",
"description": "Percent of the secondary window consumed. Gates creation of the secondary window."
},
"x-codex-secondary-window-minutes": {
"type": "string",
"pattern": "^\\d+$",
"x-headroom-parsed-type": "int"
},
"x-codex-secondary-reset-at": {
"type": "string",
"pattern": "^\\d+$",
"x-headroom-parsed-type": "int (Unix epoch seconds)"
}
},
"additionalProperties": true
},
"OpenAICodexCreditsHeaders": {
"title": "x-codex-credits-* headers",
"description": "OpenAI credits balance. A credits snapshot is materialized iff `x-codex-credits-has-credits` is present; `unlimited` defaults to false; `balance` is optional.",
"type": "object",
"properties": {
"x-codex-credits-has-credits": {
"type": "string",
"pattern": "^(?:[Tt][Rr][Uu][Ee]|[Ff][Aa][Ll][Ss][Ee]|[01])$",
"x-headroom-parsed-type": "bool (true|false|1|0, case-insensitive)",
"description": "Gates creation of the credits snapshot."
},
"x-codex-credits-unlimited": {
"type": "string",
"pattern": "^(?:[Tt][Rr][Uu][Ee]|[Ff][Aa][Ll][Ss][Ee]|[01])$",
"x-headroom-parsed-type": "bool (defaults false when absent/unparseable)"
},
"x-codex-credits-balance": {
"type": "string",
"x-headroom-parsed-type": "str (empty -> null)",
"description": "Free-form server string, e.g. \"$5.00\"."
}
},
"additionalProperties": true
},
"OpenAICodexMetaHeaders": {
"title": "x-codex meta headers",
"type": "object",
"properties": {
"x-codex-limit-name": {
"type": "string",
"x-headroom-parsed-type": "str (empty -> null)",
"description": "Active limit/model label, e.g. \"gpt-5.2-codex-sonic\"."
},
"x-codex-promo-message": {
"type": "string",
"x-headroom-parsed-type": "str (empty -> null)",
"description": "Server announcement. Also gates snapshot creation when present."
}
},
"additionalProperties": true
},
"OpenAICodexRateLimitHeaders": {
"title": "Full x-codex-* header family OpenAI may emit",
"description": "Superset of every x-codex-* header headroom reads. parse_codex_rate_limits returns a snapshot iff at least one of: a primary window, a secondary window, a credits snapshot, or a non-empty promo message is present; otherwise null (treated as a non-Codex response). All members are individually optional.",
"type": "object",
"allOf": [
{ "$ref": "#/$defs/OpenAICodexWindowHeaders" },
{ "$ref": "#/$defs/OpenAICodexCreditsHeaders" },
{ "$ref": "#/$defs/OpenAICodexMetaHeaders" }
],
"additionalProperties": true
},
"WSUpstreamHandshakeResponse": {
"title": "OpenAI WS handshake (101) response headers consumed by the WS fix",
"description": "On the Codex WebSocket transport the x-codex-* window is delivered ONLY on the upstream handshake response (never in data frames). handle_openai_responses_ws reads upstream.response.headers here. This is the contract the connect-before-accept reorder depends on: if OpenAI ever moves these headers off the handshake (e.g. into a frame), the WS half of the fix goes stale.",
"$ref": "#/$defs/OpenAICodexRateLimitHeaders"
},
"StreamingUpstreamResponseHeaders": {
"title": "OpenAI streaming/HTTP response headers consumed by the SSE fix",
"description": "On the streaming SSE/HTTP transport the same x-codex-* headers ride the HTTP response. streaming.py captures them on ALL statuses (including >=400) via update_from_headers, and forwards a wider set to the client (see ClientForwardedStreamingHeaders).",
"$ref": "#/$defs/OpenAICodexRateLimitHeaders"
},
"ClientForwardedHandshakeHeaders": {
"title": "Headers forwarded onto the CLIENT-facing WS 101 (allow/deny contract)",
"description": "_extract_codex_handshake_headers forwards ONLY headers whose (lowercased) name starts with `x-codex-`. Every other upstream handshake header - notably set-cookie and authorization - MUST NOT appear on the client 101. Enforced by propertyNames below and asserted by the unit tests + tests/e2e_ws_codex_usage_headers.py.",
"type": "object",
"propertyNames": {
"pattern": "^[Xx]-[Cc][Oo][Dd][Ee][Xx]-"
},
"not": {
"anyOf": [
{ "required": ["set-cookie"] },
{ "required": ["Set-Cookie"] },
{ "required": ["authorization"] },
{ "required": ["Authorization"] }
]
},
"additionalProperties": { "type": "string" }
},
"ClientForwardedStreamingHeaders": {
"title": "Headers forwarded to the client on the streaming SSE path",
"description": "streaming.py forwards a header iff `\"ratelimit\" in name.lower()` OR `name.lower().startswith(\"x-codex\")`. This is a SUPERSET of the WS allow-list: it additionally passes generic *ratelimit* headers (e.g. the Anthropic streaming path) which do not contain the x-codex prefix.",
"type": "object",
"propertyNames": {
"pattern": "(?:[Rr][Aa][Tt][Ee][Ll][Ii][Mm][Ii][Tt])|^[Xx]-[Cc][Oo][Dd][Ee][Xx]"
},
"additionalProperties": { "type": "string" }
},
"WSClientRequestFrame": {
"title": "Client -> proxy WS data frame (Responses API over WS)",
"description": "Codex sends the request as a response.create envelope. The HTTP fallback unwraps `.response` for the POST body, forces stream=true, and strips any top-level `type`. A flattened variant (no envelope, fields at top level) is also tolerated by the fallback.",
"type": "object",
"properties": {
"type": { "const": "response.create" },
"response": {
"type": "object",
"properties": {
"model": { "type": "string", "description": "e.g. gpt-5.4" },
"input": {
"description": "String prompt or Responses-API structured input array.",
"type": ["string", "array"]
},
"stream": { "type": "boolean" }
},
"required": ["model"],
"additionalProperties": true
}
},
"required": ["type", "response"],
"additionalProperties": true
},
"WSRelayEvent": {
"title": "proxy -> client WS data frame (relayed Responses API event)",
"description": "SSE `data:` payloads relayed verbatim as WS text frames. `[DONE]` sentinels are dropped (not relayed). Every relayed event is a JSON object carrying a `type`. response.completed additionally carries usage under `response.usage`. anyOf (not oneOf): an error event also satisfies the looser lifecycle shape, which is fine.",
"anyOf": [
{
"title": "lifecycle event",
"type": "object",
"properties": {
"type": {
"type": "string",
"examples": [
"response.created",
"response.output_item.added",
"response.completed"
]
},
"response": { "type": "object", "additionalProperties": true }
},
"required": ["type"],
"additionalProperties": true
},
{
"title": "error event",
"type": "object",
"properties": {
"type": { "const": "error" },
"error": {
"type": "object",
"properties": { "message": { "type": "string" } },
"required": ["message"],
"additionalProperties": true
}
},
"required": ["type", "error"],
"additionalProperties": true
}
]
},
"HTTPFallbackRequestBody": {
"title": "proxy -> OpenAI HTTP POST body on WS->HTTP fallback",
"description": "Derived from WSClientRequestFrame: the inner `.response` object, with `stream` forced to true and any top-level `type` removed.",
"type": "object",
"properties": {
"model": { "type": "string" },
"stream": { "const": true },
"input": { "type": ["string", "array"] }
},
"required": ["model", "stream"],
"not": { "required": ["type"] },
"additionalProperties": true
},
"CodexRateLimitStatsOutput": {
"title": "headroom /stats output for the codex tracker (CodexRateLimitSnapshot.to_dict)",
"description": "Internal (headroom-emitted) shape produced from the headers above; the WS and SSE update_from_headers parity tests assert this is refreshed. Included so drift in our own surface is also caught.",
"type": "object",
"properties": {
"limit_id": { "const": "codex" },
"limit_name": { "type": ["string", "null"] },
"primary": { "$ref": "#/$defs/CodexWindowDict" },
"secondary": { "$ref": "#/$defs/CodexWindowDict" },
"credits": {
"oneOf": [
{ "type": "null" },
{
"type": "object",
"properties": {
"has_credits": { "type": "boolean" },
"unlimited": { "type": "boolean" },
"balance": { "type": ["string", "null"] }
},
"required": ["has_credits", "unlimited", "balance"],
"additionalProperties": false
}
]
},
"promo_message": { "type": ["string", "null"] },
"captured_at": { "type": "number", "description": "Unix epoch seconds (float)." }
},
"required": ["limit_id", "limit_name", "primary", "secondary", "credits", "promo_message", "captured_at"],
"additionalProperties": false
},
"CodexWindowDict": {
"oneOf": [
{ "type": "null" },
{
"type": "object",
"properties": {
"used_percent": { "type": "number" },
"window_minutes": { "type": ["integer", "null"] },
"window_label": { "type": "string", "description": "e.g. \"5h\", \"7d\"-style label; \"unknown\" when window_minutes is null." },
"resets_at": { "type": ["integer", "null"], "description": "Unix epoch seconds." },
"seconds_until_reset": { "type": ["integer", "null"] }
},
"required": ["used_percent", "window_minutes", "window_label", "resets_at", "seconds_until_reset"],
"additionalProperties": false
}
]
}
}
}

161
tests/test_codex_openai_contract_parity.py Normal file
View file

@ -0,0 +1,161 @@
"""Parity test: the committed Codex<->OpenAI contract schema vs the live code.
The schema at ``tests/parity/fixtures/codex_openai_contracts/`` enshrines the
OpenAI interaction expectations the Codex usage-header fix depends on (PR #577):
the ``x-codex-*`` header family, and the WS-101 forward allow/deny rule.
Rather than validate golden instances against the schema (which would only check
the instances, and would pull in ``jsonschema`` as a new dep), this test binds
the schema to the *live code* in both directions, so drift in either the schema
or the parser/filter fails CI:
- every ``x-codex-*`` header the schema declares is actually consumed by
``parse_codex_rate_limits`` (rename/removal upstream -> this test fails ->
update schema + parser together);
- ``_extract_codex_handshake_headers`` forwards exactly the ``x-codex-*``
subset the schema's allow/deny ``$def`` permits, and never ``set-cookie`` /
``authorization`` (the security half of the contract).
"""
from __future__ import annotations
import json
from pathlib import Path
from types import SimpleNamespace
from headroom.proxy.handlers.openai import _extract_codex_handshake_headers
from headroom.subscription.codex_rate_limits import parse_codex_rate_limits
_SCHEMA_PATH = (
Path(__file__).parent
/ "parity"
/ "fixtures"
/ "codex_openai_contracts"
/ "codex-openai-interaction.schema.json"
)
def _load_schema() -> dict:
return json.loads(_SCHEMA_PATH.read_text())
def _declared_codex_header_names(schema: dict) -> set[str]:
"""Every ``x-codex-*`` header name declared across the header ``$defs``."""
names: set[str] = set()
for def_name in (
"OpenAICodexWindowHeaders",
"OpenAICodexCreditsHeaders",
"OpenAICodexMetaHeaders",
):
props = schema["$defs"][def_name].get("properties", {})
names.update(k for k in props if k.lower().startswith("x-codex-"))
return names
class _FakeHeaders:
def __init__(self, items: list[tuple[str, str]]):
self._items = items
def raw_items(self):
return list(self._items)
def _fake_upstream(items: list[tuple[str, str]]):
return SimpleNamespace(response=SimpleNamespace(headers=_FakeHeaders(items)))
# A valid wire value for every declared header (string-typed, as on the wire).
_VALID_VALUES = {
"x-codex-primary-used-percent": "42",
"x-codex-primary-window-minutes": "300",
"x-codex-primary-reset-at": "1900000000",
"x-codex-secondary-used-percent": "7",
"x-codex-secondary-window-minutes": "10080",
"x-codex-secondary-reset-at": "1900000000",
"x-codex-credits-has-credits": "true",
"x-codex-credits-unlimited": "false",
"x-codex-credits-balance": "$5.00",
"x-codex-limit-name": "gpt-5.2-codex-sonic",
"x-codex-promo-message": "hello",
}
def test_schema_is_wellformed_and_has_expected_defs():
"""Guard against accidental corruption/deletion of the committed schema."""
schema = _load_schema()
assert schema["$schema"].endswith("2020-12/schema")
for required_def in (
"OpenAICodexRateLimitHeaders",
"WSUpstreamHandshakeResponse",
"StreamingUpstreamResponseHeaders",
"ClientForwardedHandshakeHeaders",
"ClientForwardedStreamingHeaders",
"WSClientRequestFrame",
"WSRelayEvent",
"HTTPFallbackRequestBody",
"CodexRateLimitStatsOutput",
):
assert required_def in schema["$defs"], f"missing $def: {required_def}"
def test_every_declared_test_value_covers_the_schema():
"""The test's fixture values must cover exactly the declared header set,
so a header added to the schema without a value here is caught here rather
than silently skipped by the parity assertion below."""
declared = _declared_codex_header_names(_load_schema())
assert set(_VALID_VALUES) == declared, (
f"_VALID_VALUES drifted from the schema header set; "
f"missing={declared - set(_VALID_VALUES)} extra={set(_VALID_VALUES) - declared}"
)
def test_declared_headers_are_consumed_by_parser():
"""Every x-codex-* header the schema declares is actually parsed into the
snapshot. Catches an upstream rename/removal or a schema/parser divergence."""
declared = _declared_codex_header_names(_load_schema())
headers = {name: _VALID_VALUES[name] for name in declared}
snap = parse_codex_rate_limits(headers)
assert snap is not None
# Each schema group must materialize from its declared headers.
assert snap.primary is not None, "primary window not parsed"
assert snap.primary.used_percent == 42.0
assert snap.primary.window_minutes == 300
assert snap.primary.resets_at == 1900000000
assert snap.secondary is not None, "secondary window not parsed"
assert snap.secondary.used_percent == 7.0
assert snap.credits is not None, "credits not parsed"
assert snap.credits.has_credits is True
assert snap.credits.balance == "$5.00"
assert snap.limit_name == "gpt-5.2-codex-sonic"
assert snap.promo_message == "hello"
def test_non_codex_headers_do_not_form_a_snapshot():
"""A response with no recognized x-codex-* headers yields no snapshot,
matching the schema's gate (snapshot iff a window/credits/promo is present)."""
assert parse_codex_rate_limits({"content-type": "text/event-stream"}) is None
assert parse_codex_rate_limits({"x-codex-unknown-future-field": "1"}) is None
def test_handshake_forward_obeys_allow_deny_contract():
"""_extract_codex_handshake_headers forwards exactly the x-codex-* subset
and never set-cookie / authorization (ClientForwardedHandshakeHeaders)."""
declared = _declared_codex_header_names(_load_schema())
upstream_items = [(name, _VALID_VALUES[name]) for name in sorted(declared)]
upstream_items += [
("set-cookie", "session=should-not-forward"),
("authorization", "Bearer upstream-secret"),
("content-type", "application/json"),
]
forwarded = _extract_codex_handshake_headers(_fake_upstream(upstream_items))
forwarded_names = {name.lower() for name, _ in forwarded}
assert forwarded_names == declared, (
f"forwarded set != declared x-codex set; "
f"missing={declared - forwarded_names} extra={forwarded_names - declared}"
)
assert "set-cookie" not in forwarded_names
assert "authorization" not in forwarded_names
assert "content-type" not in forwarded_names

219
tests/test_openai_codex_ws_lifecycle.py
View file

@ -121,6 +121,7 @@ class _FakeWebSocket:
*,
disconnect_after_n_sends: int | None = None,
hold_after_initial: bool = False,
call_log: list[str] | None = None,
) -> None:
self.headers = {"authorization": "Bearer test"}
self._frames = list(frames or [])
@ -129,14 +130,19 @@ class _FakeWebSocket:
self.sent_text: list[str] = []
self.sent_bytes: list[bytes] = []
self.accepted_subprotocol: str | None = None
self.accepted_headers: list[tuple[bytes, bytes]] | None = None
self.closed = False
self.close_code: int | None = None
self._call_log = call_log
# "client" can trip this event to simulate mid-stream disconnect.
self._disconnect_event = asyncio.Event()
self.client = SimpleNamespace(host="127.0.0.1", port=12345)
async def accept(self, subprotocol=None) -> None:
async def accept(self, subprotocol=None, headers=None) -> None:
self.accepted_subprotocol = subprotocol
self.accepted_headers = list(headers) if headers is not None else None
if self._call_log is not None:
self._call_log.append("accept")
async def receive_text(self) -> str:
if self._frames:
@ -169,6 +175,26 @@ class _FakeWebSocket:
self._disconnect_event.set()
class _FakeHeaders:
"""Minimal stand-in for websockets' handshake ``Headers``.
Exposes both ``raw_items()`` (preferred by the production header
extractor to survive duplicate names like ``set-cookie``) and
``items()``.
"""
def __init__(self, pairs) -> None:
if isinstance(pairs, dict):
pairs = list(pairs.items())
self._pairs = [(str(k), str(v)) for k, v in pairs]
def raw_items(self):
return list(self._pairs)
def items(self):
return list(self._pairs)
class _FakeUpstream:
"""Upstream that streams scripted events then optionally blocks.
@ -187,12 +213,16 @@ class _FakeUpstream:
*,
hold_after_events: bool = False,
raise_mid_stream: Exception | None = None,
response_headers=None,
) -> None:
self._events = list(events)
self._hold_after_events = hold_after_events
self._raise_mid_stream = raise_mid_stream
self.sent: list[str] = []
self.closed = False
# Mirror websockets' ClientConnection.response.headers, which is the
# only place OpenAI delivers the Codex x-codex-* subscription window.
self.response = SimpleNamespace(headers=_FakeHeaders(response_headers or []))
async def __aenter__(self) -> _FakeUpstream:
return self
@ -219,9 +249,29 @@ class _FakeUpstream:
await asyncio.Event().wait()
def _make_fake_websockets_module(upstream: _FakeUpstream):
def _make_fake_websockets_module(
upstream: _FakeUpstream | None,
*,
call_log: list[str] | None = None,
connect_error: Exception | None = None,
):
"""Build a fake ``websockets`` module.
Production now does ``upstream = await websockets.connect(...)`` (then
``async with upstream``), so ``connect`` must return an awaitable that
resolves to the connection. ``connect_error`` makes the await raise to
simulate an upstream handshake failure.
"""
module = MagicMock()
module.connect = MagicMock(return_value=upstream)
async def _connect(*args, **kwargs):
if call_log is not None:
call_log.append("connect")
if connect_error is not None:
raise connect_error
return upstream
module.connect = _connect
module.Subprotocol = str
return module
@ -572,17 +622,7 @@ async def test_upstream_connect_failure_still_deregisters_cleanly():
registered+deregistered cleanly (or never registered). Either way,
no leak.
"""
class _BoomUpstream:
async def __aenter__(self):
raise RuntimeError("upstream refused")
async def __aexit__(self, exc_type, exc, tb):
return None
fake_ws_mod = MagicMock()
fake_ws_mod.connect = MagicMock(return_value=_BoomUpstream())
fake_ws_mod.Subprotocol = str
fake_ws_mod = _make_fake_websockets_module(None, connect_error=RuntimeError("upstream refused"))
client_ws = _FakeWebSocket(frames=[_first_frame()])
handler = _DummyOpenAIHandler()
@ -598,6 +638,157 @@ async def test_upstream_connect_failure_still_deregisters_cleanly():
assert handler.ws_sessions.active_count() == 0
@pytest.mark.asyncio
async def test_ws_connect_failure_falls_back_to_http():
"""When every upstream connect attempt fails, the client is still
accepted (with no x-codex-* headers, since there is no upstream
window) and the request is served via the HTTP POST fallback with
the client's first frame. Preserves the pre-reorder WS-upgrade-
failure behaviour.
"""
fake_ws_mod = _make_fake_websockets_module(
None, connect_error=RuntimeError("HTTP 500 from upstream")
)
first = _first_frame()
client_ws = _FakeWebSocket(frames=[first])
handler = _DummyOpenAIHandler()
fallback_calls: list[tuple] = []
async def _fallback(websocket, body, first_msg_raw, upstream_headers, request_id):
fallback_calls.append((body, first_msg_raw))
handler._ws_http_fallback = _fallback # type: ignore[assignment]
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
# Client was accepted with no upstream window to forward.
assert client_ws.accepted_headers is None
# Fallback ran with the first frame.
assert len(fallback_calls) == 1
_body, _first_raw = fallback_calls[0]
assert _first_raw == first
assert _body == json.loads(first)
# Clean teardown.
assert handler.ws_sessions.active_count() == 0
@pytest.mark.asyncio
async def test_ws_connect_happens_before_accept():
"""The upstream connect must complete before the client 101 is sent,
so OpenAI's x-codex-* handshake headers are available to attach.
"""
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
]
call_log: list[str] = []
upstream = _FakeUpstream(upstream_events)
fake_ws_mod = _make_fake_websockets_module(upstream, call_log=call_log)
client_ws = _FakeWebSocket(frames=[_first_frame()], call_log=call_log)
handler = _DummyOpenAIHandler()
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
assert "connect" in call_log and "accept" in call_log
assert call_log.index("connect") < call_log.index("accept"), (
f"connect must precede accept, got {call_log}"
)
@pytest.mark.asyncio
async def test_ws_forwards_codex_headers_to_client_accept():
"""OpenAI's x-codex-* subscription window from the upstream WS
handshake must be forwarded onto the client-facing 101 (and only
that subset — never set-cookie/authorization), and Python /stats
state must be refreshed.
"""
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
]
# Include duplicate set-cookie to ensure raw_items() is used (a plain
# dict-style .items() on real websockets Headers raises on dupes).
handshake_headers = [
("x-codex-primary-used-percent", "42"),
("X-Codex-Primary-Window-Minutes", "300"),
("set-cookie", "a=1"),
("set-cookie", "b=2"),
("authorization", "Bearer leak"),
]
upstream = _FakeUpstream(upstream_events, response_headers=handshake_headers)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(frames=[_first_frame()])
handler = _DummyOpenAIHandler()
captured: dict = {}
def _fake_state():
class _S:
def update_from_headers(self, headers):
captured.update(headers)
return _S()
with (
patch.dict(sys.modules, {"websockets": fake_ws_mod}),
patch(
"headroom.subscription.codex_rate_limits.get_codex_rate_limit_state",
_fake_state,
),
):
await handler.handle_openai_responses_ws(client_ws)
assert client_ws.accepted_headers is not None
names = {name.decode("latin-1").lower() for name, _ in client_ws.accepted_headers}
assert names == {"x-codex-primary-used-percent", "x-codex-primary-window-minutes"}
assert "set-cookie" not in names
assert "authorization" not in names
# Original-case names preserved on the wire.
sent = {name.decode("latin-1") for name, _ in client_ws.accepted_headers}
assert "X-Codex-Primary-Window-Minutes" in sent
# Python /stats state refreshed with the same x-codex-* subset.
assert captured == {
"x-codex-primary-used-percent": "42",
"X-Codex-Primary-Window-Minutes": "300",
}
@pytest.mark.asyncio
async def test_ws_first_frame_timeout_after_connect_closes_upstream():
"""If the client never sends its first frame after we connected, the
upstream WS must be closed (no leak) and the session deregistered.
"""
upstream = _FakeUpstream([], hold_after_events=True)
fake_ws_mod = _make_fake_websockets_module(upstream)
# No frames + hold => receive_text blocks until disconnect; we force a
# short first-frame timeout so the handler hits the timeout branch.
client_ws = _FakeWebSocket(frames=[], hold_after_initial=True)
handler = _DummyOpenAIHandler()
with (
patch.dict(sys.modules, {"websockets": fake_ws_mod}),
patch(
"headroom.proxy.handlers.openai.WS_FIRST_FRAME_TIMEOUT_SECONDS",
0.05,
),
):
await asyncio.wait_for(
handler.handle_openai_responses_ws(client_ws),
timeout=2.0,
)
assert upstream.closed, "upstream not closed on first-frame timeout"
assert client_ws.closed and client_ws.close_code == 1001
assert handler.ws_sessions.active_count() == 0
@pytest.mark.asyncio
async def test_many_concurrent_sessions_cleanly_drained():
"""50 concurrent sessions: all drain; registry and named tasks go to 0."""

35
tests/test_openai_codex_ws_timings.py
View file

@ -59,11 +59,13 @@ class _FakeWebSocket:
self.sent_text: list[str] = []
self.sent_bytes: list[bytes] = []
self.accepted_subprotocol = None
self.accepted_headers: list[tuple[bytes, bytes]] | None = None
self.closed = False
self.close_code: int | None = None
async def accept(self, subprotocol=None) -> None:
async def accept(self, subprotocol=None, headers=None) -> None:
self.accepted_subprotocol = subprotocol
self.accepted_headers = list(headers) if headers is not None else None
async def receive_text(self) -> str:
if not self._frames:
@ -112,7 +114,13 @@ class _FakeUpstream:
def _make_fake_websockets_module(upstream: _FakeUpstream):
module = MagicMock()
module.connect = MagicMock(return_value=upstream)
# Production now does ``upstream = await websockets.connect(...)`` then
# ``async with upstream`` — so connect must return an awaitable.
async def _connect(*args, **kwargs):
return upstream
module.connect = _connect
module.Subprotocol = str # the handler wraps client subprotocols if present
return module
@ -221,15 +229,12 @@ def test_codex_ws_upstream_connect_failure_still_logs_timings(stage_log_capture)
"""A session that never connects upstream still logs a timing line
with ``upstream_first_event`` absent (null)."""
class _BoomUpstream:
async def __aenter__(self):
raise RuntimeError("upstream refused")
async def __aexit__(self, exc_type, exc, tb):
return None
fake_ws_mod = MagicMock()
fake_ws_mod.connect = MagicMock(return_value=_BoomUpstream())
async def _boom_connect(*args, **kwargs):
raise RuntimeError("upstream refused")
fake_ws_mod.connect = _boom_connect
fake_ws_mod.Subprotocol = str
first_frame = json.dumps(
@ -251,12 +256,14 @@ def test_codex_ws_upstream_connect_failure_still_logs_timings(stage_log_capture)
payload = _parse_stage_log(stage_log_capture)
stages = payload["stages"]
# upstream_first_event never fired because connect failed on entry.
# upstream_first_event never fired because connect failed.
assert stages.get("upstream_first_event") is None
# upstream_connect is also None because we record it only after the
# context manager successfully enters.
# upstream_connect is also None because we record it only after a
# successful ``await websockets.connect(...)``.
assert stages.get("upstream_connect") is None
# But the envelope is still complete.
# But the envelope is still complete: the client is accepted and its
# first frame is read before falling back to HTTP, even on connect
# failure.
assert stages["accept"] is not None
assert stages["first_client_frame"] is not None
assert stages["total_session"] > 0.0

168
tests/test_proxy_streaming_ratelimit_headers.py
View file

@ -18,6 +18,24 @@ import headroom.proxy.handlers.streaming as streaming_module
from headroom.proxy.server import HeadroomProxy
@pytest.fixture(autouse=True)
def _reset_codex_rate_limit_singleton():
"""Isolate the process-global CodexRateLimitState across tests.
The tracker is a module singleton; save/restore ``_latest`` around every
test so a captured snapshot never leaks into (or depends on) another test.
"""
from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state
state = get_codex_rate_limit_state()
saved = state._latest
state._latest = None
try:
yield
finally:
state._latest = saved
class TestStreamingRatelimitHeaderForwarding:
"""Test that upstream ratelimit headers are forwarded in streaming responses."""
@ -397,3 +415,153 @@ class TestStreamingRatelimitHeaderForwarding:
assert attempts["count"] == 2
assert chunks
@pytest.mark.asyncio
async def test_codex_rate_limit_headers_captured_and_forwarded_in_streaming(self):
"""Codex x-codex-* headers must refresh /stats state AND reach the client.
Regression guard for the bug where Codex session/weekly usage never
updated on the streaming SSE transport: the proxy neither captured the
``x-codex-*`` headers into ``CodexRateLimitState`` nor forwarded them to
the client (the old ``"ratelimit" in k`` filter dropped them, so the
Codex CLI's own usage display also went stale).
"""
from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state
state = get_codex_rate_limit_state()
proxy = self._create_mock_proxy()
mock_response = self._create_mock_upstream_response(
extra_headers={
"x-codex-primary-used-percent": "42.0",
"x-codex-primary-window-minutes": "300",
"x-codex-secondary-used-percent": "8.0",
"x-codex-secondary-window-minutes": "10080",
"x-codex-limit-name": "gpt-5.4-codex",
}
)
mock_request = MagicMock()
proxy.http_client.build_request = MagicMock(return_value=mock_request)
proxy.http_client.send = AsyncMock(return_value=mock_response)
result = await proxy._stream_response(
url="https://chatgpt.com/backend-api/codex/responses",
headers={"authorization": "Bearer sk-test"},
body={"model": "gpt-5.4", "stream": True, "input": "hi"},
provider="openai",
model="gpt-5.4",
request_id="test-codex-sse",
original_tokens=10,
optimized_tokens=10,
tokens_saved=0,
transforms_applied=[],
tags={},
optimization_latency=0.0,
)
# 1. Rate-limit state refreshed from the *streaming* response.
snap = state.latest
assert snap is not None
assert snap.primary is not None
assert snap.primary.used_percent == 42.0
assert snap.primary.window_minutes == 300
assert snap.secondary is not None
assert snap.secondary.used_percent == 8.0
assert snap.secondary.window_minutes == 10080
assert snap.limit_name == "gpt-5.4-codex"
# 2. x-codex headers forwarded so the Codex CLI's native usage display
# keeps working through the proxy on the streaming path.
assert result.headers.get("x-codex-primary-used-percent") == "42.0"
assert result.headers.get("x-codex-limit-name") == "gpt-5.4-codex"
# 3. Generic ratelimit headers still forwarded; unrelated headers dropped.
assert result.headers.get("anthropic-ratelimit-tokens-limit") == "80000"
assert result.headers.get("x-request-id") is None
@pytest.mark.asyncio
async def test_codex_rate_limit_captured_on_streaming_429(self):
"""A streaming 429 carrying x-codex-* must still refresh /stats.
The capture runs *before* the >=400 early-return, matching the
non-streaming HTTP handlers (which capture on all statuses). A 429 is
exactly when the session/weekly windows are most worth surfacing, so the
previous success-only placement left the most important update missing.
"""
from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state
state = get_codex_rate_limit_state()
proxy = self._create_mock_proxy()
mock_response = self._create_mock_upstream_response()
mock_response.status_code = 429
mock_response.headers = httpx.Headers(
{
"content-type": "application/json",
"x-codex-primary-used-percent": "99.5",
"x-codex-primary-window-minutes": "300",
}
)
mock_response.aread = AsyncMock(return_value=b'{"error":{"message":"rate limited"}}')
mock_response.aclose = AsyncMock()
mock_request = MagicMock()
proxy.http_client.build_request = MagicMock(return_value=mock_request)
proxy.http_client.send = AsyncMock(return_value=mock_response)
result = await proxy._stream_response(
url="https://chatgpt.com/backend-api/codex/responses",
headers={"authorization": "Bearer sk-test"},
body={"model": "gpt-5.4", "stream": True, "input": "hi"},
provider="openai",
model="gpt-5.4",
request_id="test-codex-429",
original_tokens=10,
optimized_tokens=10,
tokens_saved=0,
transforms_applied=[],
tags={},
optimization_latency=0.0,
)
assert result.status_code == 429
snap = state.latest
assert snap is not None
assert snap.primary is not None
assert snap.primary.used_percent == 99.5
@pytest.mark.asyncio
async def test_anthropic_stream_leaves_codex_state_untouched(self):
"""The now-unconditional capture must be a no-op for non-Codex streams."""
from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state
state = get_codex_rate_limit_state()
proxy = self._create_mock_proxy()
mock_response = self._create_mock_upstream_response() # anthropic-ratelimit-* only
mock_request = MagicMock()
proxy.http_client.build_request = MagicMock(return_value=mock_request)
proxy.http_client.send = AsyncMock(return_value=mock_response)
await proxy._stream_response(
url="https://api.anthropic.com/v1/messages",
headers={"x-api-key": "sk-test"},
body={
"model": "claude-sonnet-4-20250514",
"max_tokens": 100,
"stream": True,
"messages": [{"role": "user", "content": "hi"}],
},
provider="anthropic",
model="claude-sonnet-4-20250514",
request_id="test-anthropic-noop",
original_tokens=10,
optimized_tokens=10,
tokens_saved=0,
transforms_applied=[],
tags={},
optimization_latency=0.0,
)
assert state.latest is None

42
tests/test_ws_http_fallback.py
View file

@ -31,9 +31,15 @@ class FakeWebSocket:
class FakeStreamResponse:
"""Mock httpx streaming response."""
def __init__(self, status_code: int = 200, sse_events: list[str] | None = None):
def __init__(
self,
status_code: int = 200,
sse_events: list[str] | None = None,
headers: dict[str, str] | None = None,
):
self.status_code = status_code
self._events = sse_events or []
self.headers = headers or {}
async def aiter_text(self):
for event in self._events:
@ -279,3 +285,37 @@ class TestWsHttpFallback:
asyncio.run(handler._ws_http_fallback(ws, body, json.dumps(body), headers, "req_6"))
assert "api.openai.com" in captured_url["url"]
def test_fallback_refreshes_codex_rate_limit_state(self, monkeypatch):
"""A successful fallback refreshes Codex /stats from response headers.
The fallback can't forward headers onto the (already-accepted) client
101, but it should still keep Python /stats in sync so the gauge does
not go stale when the WS upgrade fails and we drop to HTTP.
"""
handler = _make_handler()
ws = FakeWebSocket()
captured: dict[str, dict[str, str]] = {}
class _FakeState:
def update_from_headers(self, hdrs):
captured["headers"] = dict(hdrs)
import headroom.subscription.codex_rate_limits as crl
monkeypatch.setattr(crl, "get_codex_rate_limit_state", lambda: _FakeState())
response = FakeStreamResponse(
200,
['data: {"type":"response.completed"}\n\n', "data: [DONE]\n\n"],
headers={
"x-codex-primary-used-percent": "42",
"content-type": "text/event-stream",
},
)
handler.http_client = FakeHttpClient(response)
body = {"model": "gpt-5.4", "input": "hi"}
asyncio.run(handler._ws_http_fallback(ws, body, json.dumps(body), {}, "req_capture"))
assert captured["headers"]["x-codex-primary-used-percent"] == "42"