headroom/tests/test_codex_client_stamp.py
gglucass b0cd0329c7
fix(proxy): stamp X-Client: codex on Responses endpoint for unidentified callers (#1036)
## Description

Codex Desktop (OpenAI's Codex GUI/IDE app) sends a `User-Agent` of the
form `Codex Desktop/<ver> (...)`, which is not in `CLIENT_UA_MAP`, so
`classify_client` returns `None`. On a compression timeout the backend
only takes the codex fail-open path when the client classifies as
`codex`; for an unidentified client it refuses with HTTP 413
(`compression_refused`), which Codex treats as a hard connection
failure.

This stamps `X-Client: codex` on requests to the Responses endpoint
(`/v1/responses`) only when the caller does not otherwise classify. The
stamp is scoped to the Responses endpoint and skipped for any caller
that already classifies through a recognized user-agent or explicit
`X-Client`, so non-Codex traffic is not relabeled.

## Type of Change

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

## Changes Made

- Added `should_stamp_codex_client(path, headers)` in
`headroom.proxy.auth_mode` for narrow Responses-endpoint client
stamping.
- Applied the stamp in HTTP middleware before downstream request
classification.
- Applied the same stamp in the Responses WebSocket handler, which
bypasses HTTP middleware.
- Added unit coverage for the stamp/skip matrix, including Codex
Desktop, explicit clients, recognized user-agents, and WebSocket
behavior.
- Merged current `main` and kept both the new stamp coverage and the
existing Codex WebSocket image-generation regression coverage.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`)
- [x] Formatting passes (`ruff format --check`)
- [x] New tests added for new functionality

### Test Output

```text
python -m pytest tests/test_codex_client_stamp.py tests/test_auth_mode.py tests/test_openai_codex_ws_lifecycle.py -q
48 passed in 1.13s

ruff check headroom/proxy/auth_mode.py headroom/proxy/server.py headroom/proxy/handlers/openai.py tests/test_codex_client_stamp.py tests/test_auth_mode.py tests/test_openai_codex_ws_lifecycle.py
All checks passed!

ruff format --check headroom/proxy/auth_mode.py headroom/proxy/server.py headroom/proxy/handlers/openai.py tests/test_codex_client_stamp.py tests/test_auth_mode.py tests/test_openai_codex_ws_lifecycle.py
6 files already formatted
```

## Real Behavior Proof

- Environment: local Windows 11 development checkout, Python 3.13.13,
branch updated from `upstream/main`.
- Exact command / steps: Ran the focused unit suite for the new
client-stamp behavior and the overlapping Codex WebSocket lifecycle
tests, then ran `ruff check` and `ruff format --check` on the changed
modules and tests.
- Observed result: The focused suite passed with 48 tests, lint passed,
and formatting passed. The tests assert that unidentified
`/v1/responses` callers classify as Codex after stamping while explicit
or already-recognized clients are preserved.
- Not tested: a live end-to-end Codex Desktop session through a running
`headroom wrap codex` instance; verification is at the unit/integration
boundary for classification and request routing.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
2026-06-17 11:45:20 -05:00

56 lines
2.1 KiB
Python

"""Tests for ``should_stamp_codex_client`` — the path-based ``X-Client: codex``
stamp on the Responses endpoint.
The stamp fires only for an unidentified caller on the Responses endpoint, so
Codex Desktop (whose User-Agent isn't a known codex UA) takes the codex
fail-open path instead of being refused with a 413 on a compression timeout.
"""
from __future__ import annotations
from headroom.proxy.auth_mode import classify_client, should_stamp_codex_client
CODEX_DESKTOP_UA = (
"Codex Desktop/0.140.0-alpha.2 (Mac OS 15.7.7; arm64) unknown (Codex Desktop; 26.609.71450)"
)
def test_unidentified_codex_desktop_on_responses_is_stamped() -> None:
assert should_stamp_codex_client("/v1/responses", {"user-agent": CODEX_DESKTOP_UA})
def test_stamp_then_classify_yields_codex() -> None:
# End-to-end of what the HTTP middleware and the WS handler both do:
# stamp the header, after which classify_client must read "codex".
headers = {"user-agent": CODEX_DESKTOP_UA}
assert should_stamp_codex_client("/v1/responses", headers)
headers["x-client"] = "codex"
assert classify_client(headers) == "codex"
def test_no_user_agent_on_responses_is_stamped() -> None:
assert should_stamp_codex_client("/v1/responses", {})
def test_responses_subpath_is_stamped() -> None:
assert should_stamp_codex_client("/v1/responses/foo", {"user-agent": CODEX_DESKTOP_UA})
def test_other_path_is_not_stamped() -> None:
# Scoped to the Responses endpoint; unknown callers elsewhere are untouched.
assert not should_stamp_codex_client("/v1/chat/completions", {"user-agent": CODEX_DESKTOP_UA})
def test_recognized_non_codex_client_is_not_stamped() -> None:
assert not should_stamp_codex_client("/v1/responses", {"user-agent": "claude-code/1.2.3"})
def test_recognized_codex_cli_is_not_stamped() -> None:
# Already classifies as codex via UA; no stamp needed.
assert not should_stamp_codex_client("/v1/responses", {"user-agent": "codex-cli/0.5"})
def test_explicit_x_client_is_not_stamped() -> None:
assert not should_stamp_codex_client(
"/v1/responses", {"x-client": "aider", "user-agent": CODEX_DESKTOP_UA}
)