Unit 3 of the Codex proxy resilience plan. Eliminates the "aged process
has leaked relay tasks" hypothesis by making every WS session explicitly
tracked and both relay tasks deterministically cancelled when either
exits.
- New headroom/proxy/ws_session_registry.py: dict-backed
WebSocketSessionRegistry + WSSessionHandle with register /
deregister / attach_tasks / snapshot. Deregister is idempotent and
clears relay-task references so coroutine frames are not retained
past session end.
- HeadroomProxy exposes proxy.ws_sessions so /debug/ws-sessions
(Unit 5) can read the live snapshot.
- handle_openai_responses_ws now registers on websocket.accept()
success and deregisters in the outermost finally so no leak can
survive handshake-phase, mid-stream, or upstream-error paths. The
session_id / termination_cause is threaded through both relay
halves and both sides raise asyncio.CancelledError cleanly.
- Replaced asyncio.gather(_client_to_upstream(), _upstream_to_client(),
return_exceptions=True) with explicit asyncio.create_task(...)
(named codex-ws-c2u-<sid> / codex-ws-u2c-<sid>) +
asyncio.wait(FIRST_COMPLETED) + cancel-and-await on the survivor.
Termination cause is classified as client_disconnect /
client_error / upstream_disconnect / upstream_error /
response_completed from which task completed first plus inline
error captures from the halves.
- Prometheus metrics: new active_ws_sessions and active_relay_tasks
gauges plus ws_session_duration_ms_{sum,count,max} histogram
bucketed by termination cause. Mirrors the Unit 2 stage_timing_*
shape.
Preserved: upstream WS retry loop, WS→HTTP fallback, memory-context
timeout, compression pipeline, Unit 2 stage timings. Memory-tool
execution inside _upstream_to_client still runs when the client task
exits first; however, if the client disconnects *before* the upstream
emits response.completed, pending memory writes in `pending_fcs` are
dropped (unchanged from prior behavior — a crashing upstream has the
same effect). Note: handle_openai_responses (HTTP, line ~800) is a
single-shot HTTP request; lifecycle tracking isn't added there
(scope boundary).
Tests:
- tests/test_ws_session_registry.py: 8 registry unit tests
(register/deregister idempotency, snapshot shape, attach merging,
reference release).
- tests/test_openai_codex_ws_lifecycle.py: 6 integration tests
using real relay tasks (only upstream WS endpoint mocked):
happy-path, failing-test-first "client disconnect cancels upstream
relay within 100 ms", upstream-closes-first, upstream-error mid-
stream, handshake-failure deregister, 50 concurrent sessions.
- Regression: test_openai_codex_ws_timings, test_openai_codex_routing,
test_proxy_codex_route_aliases, test_ws_memory_relay all pass.
- Tests pass under python -W error::RuntimeWarning (no "coroutine
was never awaited").
