fix(proxy): keep codex image-generation WS turns alive through the relay (#1000)

## Description

Image generation through the proxy fails. Driving Codex (`/v1/responses`
over
WebSocket) through Headroom, an image-generation turn never returns an
image —
the client retries (`Reconnecting… n/5`) and gives up, while the same
prompt
works when Codex talks to ChatGPT directly.

Root cause: two independent defects on the upstream
`websockets.connect()`, both
specific to how image generation behaves on the wire:

1. **Pong deadline kills the silent render.** An image turn emits a
single
`response.image_generation_call.generating` event and then goes silent
for
   20–60s while the model renders (no data frames). The hard-coded
`ping_timeout=20` treats that healthy-but-quiet connection as dead and
tears
   it down as `upstream_error` mid-render, before the image is ready.
2. **1 MiB frame cap drops the image.** The finished image comes back
inline as
a single base64 frame that exceeds the `websockets` default
`max_size=2**20`
   (1 MiB), raising `PayloadTooBig` exactly as the image lands.

They compound: with only ping fixed, the session survives the silent
phase
(observed ~20s → ~54s) but then dies on the oversized image frame.
Normal
text/tool turns stream tokens continuously and stay well under 1 MiB, so
neither
defect affects them — which is why this only ever bit image generation.

Closes: N/A (no tracking issue)

## 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

- `headroom/proxy/handlers/openai.py`: on the upstream `/v1/responses`
connect,
set `ping_timeout=None` (keep `ping_interval=20` for NAT keepalive) so a
long
silent render is not torn down on a missing pong, and `max_size=None` so
the
inline base64 image frame is accepted instead of raising
`PayloadTooBig`.
- `tests/test_openai_codex_ws_lifecycle.py`: add
`test_ws_upstream_connect_allows_large_frames_and_no_pong_deadline`,
which
  captures the upstream connect kwargs and pins `ping_timeout is None` /
  `max_size is None`.

## Testing

<!-- Check what you actually ran, then paste the real command output
below. -->

- [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

### Test Output

```text
$ ruff check .
All checks passed!

$ mypy headroom --ignore-missing-imports
Success: no issues found in 358 source files

$ pytest tests/test_openai_codex_ws_lifecycle.py -q
collected 15 items
tests/test_openai_codex_ws_lifecycle.py ..............                   [100%]
============================== 15 passed in 0.72s ==============================

$ pytest tests/test_openai_codex_ws_lifecycle.py -k large_frames_and_no_pong -q
collected 15 items / 14 deselected / 1 selected
tests/test_openai_codex_ws_lifecycle.py .                                [100%]
======================= 1 passed, 14 deselected in 0.35s =======================
```

End-to-end (managed Codex image generation through the running proxy):

```text
# BEFORE fix: fails at ~20s (unpatched) / ~54s (ping-only)
#   WS /v1/responses completed (cause=upstream_error,
#       last_upstream_type=response.image_generation_call.generating)
#   -> client "Reconnecting… n/5", no image produced

# AFTER fix:
[codex] Image ready; stopping the turn.
Saved image: /tmp/headroom-imagegen-test.png
$ file /tmp/headroom-imagegen-test.png
PNG image data, 1254 x 1254, 8-bit/color RGB, non-interlaced   (908 KB)
# proxy session count +1 -> the turn DID traverse the proxy and completed.
```

## Real Behavior Proof

- Environment: macOS, headroom 0.23.0 running as the Codex
`model_provider`
  (proxy on `127.0.0.1:8787`), Codex CLI 0.139.0 driving a managed
  `/v1/responses` image-generation turn through the proxy.
- Exact command / steps: trigger a Codex image-generation turn
(gpt-image-2)
with the proxy in front; observe the upstream `/v1/responses` WS session
in
  `proxy.log` and whether a PNG is returned.
- Observed result: before the change the session dies with
`upstream_error`
while `last_upstream_type=response.image_generation_call.generating` and
no
image is produced; after the change a valid 1254×1254 PNG is returned
and the
  session traverses the proxy normally.
- Not tested: the full `pytest` suite was not run locally — this machine
has no
Rust toolchain to rebuild the matching `_core` extension, so the
complete
  suite (incl. the pyo3 tests) is left to CI. The affected
`test_openai_codex_ws_lifecycle.py` module was run against the installed
extension and passes 15/15; `ruff check .` and `mypy headroom` were run
in
  full and pass.

## Review Readiness

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

## 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

## Screenshots (if applicable)

N/A.

## Additional Notes

- `ruff check .` (whole repo) and `mypy headroom
--ignore-missing-imports`
(358 source files) were run locally and pass. The only `pytest` not run
locally is the full suite, because the Rust `_core` cannot be rebuilt
here
without a toolchain; the directly affected lifecycle module passes 15/15
and
  CI runs the rest.
- Documentation / CHANGELOG left unchecked — this is a focused two-line
behavioral fix on the upstream WS connect; happy to add a CHANGELOG
entry if
  preferred.
- `ping_timeout=None` keeps `ping_interval` for NAT keepalive; if you'd
rather
bound it, a generous finite value (e.g. 300s) would also fix the render
case —
  happy to switch.

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

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

headroom/proxy/handlers/openai.py

@@ -3759,7 +3759,20 @@ class OpenAIHandlerMixin:
                         open_timeout=max(30, self.config.connect_timeout_seconds * 3),
                         close_timeout=10,
                         ping_interval=20,
-                        ping_timeout=20,
+                        # Image-generation turns go silent for 20-60s while the
+                        # model renders (a single ``image_generation_call`` event,
+                        # then a long quiet gap with no data frames). A 20s pong
+                        # deadline false-kills the still-healthy upstream
+                        # mid-render with ``upstream_error`` before the image
+                        # lands. Keep ``ping_interval`` for NAT keepalive but do
+                        # not tear the session down on a missing pong.
+                        ping_timeout=None,
+                        # The finished image arrives inline as a single base64
+                        # frame that exceeds the websockets default 1 MiB cap,
+                        # raising ``PayloadTooBig`` exactly as the image lands.
+                        # The relay must accept frames as large as the endpoints
+                        # do, so do not cap the upstream payload size.
+                        max_size=None,
                     )
                     ws_connected = True
                     if not _upstream_connect_recorded:

tests/test_openai_codex_ws_lifecycle.py

@@ -815,3 +815,41 @@ async def test_many_concurrent_sessions_cleanly_drained():
         if (t.get_name() or "").startswith("codex-ws-") and not t.done()
     ]
     assert leaked == []
+
+
+@pytest.mark.asyncio
+async def test_ws_upstream_connect_allows_large_frames_and_no_pong_deadline():
+    """The upstream WS must accept arbitrarily large frames and never impose a
+    pong deadline.
+
+    Image-generation turns expose two failure modes the relay was previously
+    blind to: (1) the render phase goes silent for 20-60s with no data frames,
+    so a 20s pong deadline false-kills the healthy upstream mid-render; and
+    (2) the finished image arrives inline as a single base64 frame larger than
+    the websockets default 1 MiB cap, raising ``PayloadTooBig`` just as it
+    lands. Pin the connect kwargs so neither regresses.
+    """
+    upstream_events = [
+        json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
+        json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
+    ]
+    upstream = _FakeUpstream(upstream_events)
+    fake_ws_mod = _make_fake_websockets_module(upstream)
+
+    captured: dict = {}
+    inner_connect = fake_ws_mod.connect
+
+    async def _capturing_connect(*args, **kwargs):
+        captured.update(kwargs)
+        return await inner_connect(*args, **kwargs)
+
+    fake_ws_mod.connect = _capturing_connect
+
+    client_ws = _FakeWebSocket(frames=[_first_frame()])
+    handler = _DummyOpenAIHandler()
+
+    with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
+        await handler.handle_openai_responses_ws(client_ws)
+
+    assert captured.get("max_size") is None, "upstream frame size must be uncapped"
+    assert captured.get("ping_timeout") is None, "upstream must not impose a pong deadline"