diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index da9cb6f0d..40388386b 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -19,6 +19,10 @@ from datetime import datetime from typing import TYPE_CHECKING, Any from headroom.proxy.stage_timer import StageTimer, emit_stage_timings_log +from headroom.proxy.ws_session_registry import ( + WebSocketSessionRegistry, + WSSessionHandle, +) if TYPE_CHECKING: from fastapi import Request, WebSocket @@ -1248,6 +1252,15 @@ class OpenAIHandlerMixin: stage_timer = StageTimer() session_started_at = time.perf_counter() + # Unit 3: initialize registry variables *before* accept so the + # outermost ``finally`` can rely on them existing even if + # registration itself fails for some reason. + ws_sessions: WebSocketSessionRegistry | None = getattr( + self, "ws_sessions", None + ) + session_handle: WSSessionHandle | None = None + termination_cause: str = "unknown" + # Forward client headers to upstream, adding required OpenAI-Beta header ws_headers = dict(websocket.headers) @@ -1266,6 +1279,31 @@ class OpenAIHandlerMixin: else: await websocket.accept() + # --- Unit 3: register the session as soon as accept succeeds --- + client_addr: str | None = None + client_info = getattr(websocket, "client", None) + if client_info is not None: + host = getattr(client_info, "host", None) + port = getattr(client_info, "port", None) + if host is not None and port is not None: + client_addr = f"{host}:{port}" + elif host is not None: + client_addr = str(host) + if ws_sessions is not None: + session_handle = WSSessionHandle( + session_id=session_id, + request_id=request_id, + client_addr=client_addr, + upstream_url=None, # set below once upstream_url is computed + ) + ws_sessions.register(session_handle) + metrics = getattr(self, "metrics", None) + if metrics is not None and hasattr(metrics, "inc_active_ws_sessions"): + try: + metrics.inc_active_ws_sessions() + except Exception: # pragma: no cover - defensive + pass + # Forward all client headers except hop-by-hop / per-connection headers. # These are WebSocket handshake mechanics that the `websockets` library # generates fresh for the upstream connection — forwarding them would conflict. @@ -1305,6 +1343,10 @@ class OpenAIHandlerMixin: ws_base = base.replace("https://", "wss://").replace("http://", "ws://") upstream_url = f"{ws_base}/v1/responses" + # Unit 3: attach the resolved upstream URL to the session handle. + if session_handle is not None: + session_handle.upstream_url = upstream_url + # Ensure Authorization header is present — fall back to OPENAI_API_KEY env var. # Safety net for clients that don't forward auth headers via WebSocket upgrade. if "authorization" not in _lower_headers: @@ -1579,13 +1621,39 @@ class OpenAIHandlerMixin: _upstream_first_event_started = time.perf_counter() await upstream.send(first_msg_raw) + # Unit 3: flag the upstream side flips on seeing + # ``response.completed`` so the outer cause + # classifier can prefer it over the raw + # "upstream iterator ended" default. + response_completed_seen = False + # Captures the first exception surfaced by the + # inner relay ``except`` blocks so the outer + # classifier can still tell ``upstream_error`` + # from ``upstream_disconnect`` / ``response_completed`` + # even though the halves swallow and log. + upstream_relay_error: BaseException | None = None + client_relay_error: BaseException | None = None + async def _client_to_upstream() -> None: + nonlocal client_relay_error try: while True: msg = await websocket.receive_text() await upstream.send(msg) + except asyncio.CancelledError: + # Explicit cancel from the outer + # orchestrator — re-raise so + # ``t.cancelled()`` and ``t.exception()`` + # behave correctly in the caller. + raise except Exception as relay_err: + # Surface real errors to the classifier + # without re-raising (existing fork + # behavior: log and return so the + # partner task can be cancelled + # deterministically). if "WebSocketDisconnect" not in type(relay_err).__name__: + client_relay_error = relay_err logger.debug( f"[{request_id}] WS client→upstream relay ended: {relay_err}" ) @@ -1606,6 +1674,13 @@ class OpenAIHandlerMixin: """ from headroom.proxy.memory_handler import MEMORY_TOOL_NAMES + # Unit 3: surface response.completed observation + # to the outer scope so the termination-cause + # classifier can prefer ``response_completed`` + # over ``upstream_disconnect``. + nonlocal response_completed_seen + nonlocal upstream_relay_error + memory_enabled = bool(self.memory_handler and memory_user_id) # Per-response state (reset after each response.completed) @@ -1689,6 +1764,7 @@ class OpenAIHandlerMixin: await websocket.send_text(buf) event_buffer.clear() _reset() + response_completed_seen = True continue @@ -1703,6 +1779,7 @@ class OpenAIHandlerMixin: pending_fcs.append(item) elif event_type == "response.completed": + response_completed_seen = True resp = event.get("response", {}) resp_id = resp.get("id") @@ -1770,8 +1847,15 @@ class OpenAIHandlerMixin: # --- Phase 2b: Pass-through mode --- await websocket.send_text(msg_str) + except asyncio.CancelledError: + raise except Exception as relay_err: if "WebSocketDisconnect" not in type(relay_err).__name__: + # Capture for the outer classifier + # so ``upstream_error`` can be + # distinguished from a clean + # upstream disconnect. + upstream_relay_error = relay_err logger.debug( f"[{request_id}] WS upstream→client relay ended: {relay_err}" ) @@ -1779,14 +1863,115 @@ class OpenAIHandlerMixin: with contextlib.suppress(Exception): await websocket.close() - await asyncio.gather( + # --- Unit 3: deterministic relay-task cancellation --- + # Spawn each half as a named task so we can: + # (a) attach them to the session registry for + # ``/debug/ws-sessions``, + # (b) cancel the survivor explicitly when the + # first one exits, and + # (c) classify the termination cause for the + # duration histogram. + client_task = asyncio.create_task( _client_to_upstream(), - _upstream_to_client(), - return_exceptions=True, + name=f"codex-ws-c2u-{session_id}", ) + upstream_task = asyncio.create_task( + _upstream_to_client(), + name=f"codex-ws-u2c-{session_id}", + ) + relay_tasks = [client_task, upstream_task] + if ws_sessions is not None: + ws_sessions.attach_tasks(session_id, relay_tasks) + metrics_for_tasks = getattr(self, "metrics", None) + if metrics_for_tasks is not None and hasattr( + metrics_for_tasks, "inc_active_relay_tasks" + ): + try: + metrics_for_tasks.inc_active_relay_tasks( + len(relay_tasks) + ) + except Exception: # pragma: no cover - defensive + pass + + try: + done, pending = await asyncio.wait( + {client_task, upstream_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + # Cancel the survivor so we don't leak the + # partner task. Suppress the CancelledError + # we just raised ourselves — any *other* + # exception from the cancelled task is + # already logged inside its own try/except. + for t in pending: + t.cancel() + if pending: + with contextlib.suppress(asyncio.CancelledError): + await asyncio.gather( + *pending, return_exceptions=True + ) + + # Classify termination cause from whichever + # task completed first. ``CancelledError`` + # can show up on the "done" side if the + # handler itself was cancelled from outside + # (e.g. server shutdown). + for t in done: + exc = None + with contextlib.suppress(Exception): + exc = t.exception() + task_name = t.get_name() or "" + if t is client_task: + if client_relay_error is not None: + termination_cause = "client_error" + elif exc is None: + termination_cause = "client_disconnect" + elif isinstance(exc, asyncio.CancelledError): + termination_cause = "client_disconnect" + else: + # Distinguish legitimate client + # disconnect exceptions from + # real errors: WebSocketDisconnect + # is a normal client exit. + if "WebSocketDisconnect" in type(exc).__name__: + termination_cause = "client_disconnect" + else: + termination_cause = "client_error" + elif t is upstream_task: + if upstream_relay_error is not None: + termination_cause = "upstream_error" + logger.debug( + f"[{request_id}] WS relay {task_name} " + f"raised: {upstream_relay_error!r}" + ) + elif exc is None: + termination_cause = ( + "response_completed" + if response_completed_seen + else "upstream_disconnect" + ) + elif isinstance(exc, asyncio.CancelledError): + termination_cause = "upstream_disconnect" + else: + termination_cause = "upstream_error" + logger.debug( + f"[{request_id}] WS relay {task_name} " + f"raised: {exc!r}" + ) + finally: + # In case anything above raised before the + # cancel-and-await loop ran. + for t in relay_tasks: + if not t.done(): + t.cancel() + with contextlib.suppress(asyncio.CancelledError): + await asyncio.gather( + *relay_tasks, return_exceptions=True + ) logger.info( - f"[{request_id}] WS /v1/responses completed (tokens_saved={tokens_saved})" + f"[{request_id}] WS /v1/responses completed " + f"(tokens_saved={tokens_saved}, cause={termination_cause})" ) break except Exception as ws_err: @@ -1840,7 +2025,14 @@ class OpenAIHandlerMixin: ) except Exception as e: - if "WebSocketDisconnect" not in type(e).__name__: + if "WebSocketDisconnect" in type(e).__name__: + # Unit 3: client dropped the socket before or during + # relay. The registry classifier may already have set + # ``client_disconnect`` via the relay task exit path; + # preserve that, otherwise set it here. + if termination_cause == "unknown": + termination_cause = "client_disconnect" + else: # Extract response body from websockets InvalidStatus for better debugging error_detail = str(e) if hasattr(e, "response"): @@ -1854,6 +2046,8 @@ class OpenAIHandlerMixin: except Exception: pass logger.error(f"[{request_id}] WS proxy error: {error_detail}") + if termination_cause == "unknown": + termination_cause = "client_error" with contextlib.suppress(Exception): await websocket.close(code=1011, reason=str(e)[:120]) finally: @@ -1862,6 +2056,29 @@ class OpenAIHandlerMixin: "total_session", (time.perf_counter() - session_started_at) * 1000.0, ) + # Unit 3: deregister the session before (or independently + # of) the stage-timings log so a failure there cannot leak + # the registry entry. ``deregister`` is idempotent, so a + # session that never registered is a no-op. + if ws_sessions is not None and session_handle is not None: + released_tasks = len(session_handle.relay_tasks) + ws_sessions.deregister(session_id, cause=termination_cause) + session_duration_ms = ( + time.perf_counter() - session_started_at + ) * 1000.0 + metrics_for_close = getattr(self, "metrics", None) + if metrics_for_close is not None: + with contextlib.suppress(Exception): + if hasattr(metrics_for_close, "dec_active_ws_sessions"): + metrics_for_close.dec_active_ws_sessions() + if released_tasks and hasattr( + metrics_for_close, "dec_active_relay_tasks" + ): + metrics_for_close.dec_active_relay_tasks(released_tasks) + if hasattr(metrics_for_close, "record_ws_session_duration"): + metrics_for_close.record_ws_session_duration( + session_duration_ms, termination_cause + ) await emit_stage_timings_log( path="openai_responses_ws", request_id=request_id, diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index e7c58a75a..d4b373a11 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -108,6 +108,17 @@ class PrometheusMetrics: self.stage_timing_count: dict[tuple[str, str], int] = defaultdict(int) self.stage_timing_max: dict[tuple[str, str], float] = defaultdict(float) + # WS session lifecycle (Unit 3). Gauges are live counters updated + # by the Codex handler on register/deregister + attach_tasks/ + # detach. Histograms record completed-session durations bucketed + # by termination cause so we can distinguish slow happy-path + # sessions from long client-hold followed by client_disconnect. + self.active_ws_sessions: int = 0 + self.active_relay_tasks: int = 0 + self.ws_session_duration_sum_ms: dict[str, float] = defaultdict(float) + self.ws_session_duration_count: dict[str, int] = defaultdict(int) + self.ws_session_duration_max_ms: dict[str, float] = defaultdict(float) + # Aggregate waste signals self.waste_signals_total: dict[str, int] = defaultdict(int) @@ -359,6 +370,47 @@ class PrometheusMetrics: self.cache_bust_count += 1 self._get_otel_metrics().record_proxy_cache_bust(tokens_lost=tokens_lost) + # ------------------------------------------------------------------ + # Unit 3: WS session lifecycle gauges / histogram + # ------------------------------------------------------------------ + + def inc_active_ws_sessions(self) -> None: + """Increment the live WS session gauge (called on register).""" + self.active_ws_sessions += 1 + + def dec_active_ws_sessions(self) -> None: + """Decrement the live WS session gauge (called on deregister).""" + self.active_ws_sessions = max(0, self.active_ws_sessions - 1) + + def inc_active_relay_tasks(self, n: int = 1) -> None: + """Increment the live relay-task gauge (attach_tasks).""" + self.active_relay_tasks += n + + def dec_active_relay_tasks(self, n: int = 1) -> None: + """Decrement the live relay-task gauge (deregister).""" + self.active_relay_tasks = max(0, self.active_relay_tasks - n) + + def record_ws_session_duration( + self, + duration_ms: float, + cause: str = "unknown", + ) -> None: + """Record a completed WS session's duration, bucketed by cause. + + Mirrors the ``stage_timing_*`` shape so ``/metrics`` exposes + sum/count/max per termination cause. Uses synchronous dict + updates (no ``_lock``) because Unit 3 callers run on the event + loop — matching the gauges above. + """ + try: + ms_val = float(duration_ms) + except (TypeError, ValueError): + return + self.ws_session_duration_sum_ms[cause] += ms_val + self.ws_session_duration_count[cause] += 1 + if ms_val > self.ws_session_duration_max_ms[cause]: + self.ws_session_duration_max_ms[cause] = ms_val + async def record_rate_limited(self, *, provider: str | None = None, model: str | None = None): async with self._lock: self.requests_rate_limited += 1 @@ -611,6 +663,54 @@ class PrometheusMetrics: ) lines.append("") + # Unit 3: WS session lifecycle gauges + duration histogram. + lines.extend( + [ + "# HELP headroom_active_ws_sessions Active Codex WebSocket sessions", + "# TYPE headroom_active_ws_sessions gauge", + f"headroom_active_ws_sessions {self.active_ws_sessions}", + "", + "# HELP headroom_active_relay_tasks Active Codex WS relay tasks", + "# TYPE headroom_active_relay_tasks gauge", + f"headroom_active_relay_tasks {self.active_relay_tasks}", + "", + ] + ) + if self.ws_session_duration_sum_ms: + lines.extend( + [ + "# HELP headroom_ws_session_duration_ms_sum Sum of Codex WS session durations", + "# TYPE headroom_ws_session_duration_ms_sum counter", + ] + ) + for cause, total in self.ws_session_duration_sum_ms.items(): + lines.append( + f'headroom_ws_session_duration_ms_sum{{cause="{_escape_label_value(cause)}"}} {round(total, 2)}' + ) + lines.extend( + [ + "", + "# HELP headroom_ws_session_duration_ms_count Count of completed Codex WS sessions", + "# TYPE headroom_ws_session_duration_ms_count counter", + ] + ) + for cause, count in self.ws_session_duration_count.items(): + lines.append( + f'headroom_ws_session_duration_ms_count{{cause="{_escape_label_value(cause)}"}} {count}' + ) + lines.extend( + [ + "", + "# HELP headroom_ws_session_duration_ms_max Maximum Codex WS session duration", + "# TYPE headroom_ws_session_duration_ms_max gauge", + ] + ) + for cause, max_value in self.ws_session_duration_max_ms.items(): + lines.append( + f'headroom_ws_session_duration_ms_max{{cause="{_escape_label_value(cause)}"}} {round(max_value, 2)}' + ) + lines.append("") + if self.waste_signals_total: lines.extend( [ diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 9f2d4f956..f6a57de6a 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -129,6 +129,7 @@ from headroom.proxy.rate_limiter import TokenBucketRateLimiter # noqa: F401 from headroom.proxy.request_logger import RequestLogger # noqa: F401 from headroom.proxy.semantic_cache import SemanticCache # noqa: F401 from headroom.proxy.warmup import WarmupRegistry +from headroom.proxy.ws_session_registry import WebSocketSessionRegistry from headroom.subscription.base import get_quota_registry, reset_quota_registry from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state from headroom.subscription.copilot_quota import get_copilot_quota_tracker @@ -373,6 +374,10 @@ class HeadroomProxy( # each preloaded heavy asset. Exposed as ``proxy.warmup`` and # serialized by the /debug/warmup route (Unit 5). self.warmup: WarmupRegistry = WarmupRegistry() + # Unit 3: live registry of Codex WS sessions. Populated by + # ``handle_openai_responses_ws`` on accept; drained in its + # outermost ``finally``. Consumed by ``/debug/ws-sessions``. + self.ws_sessions: WebSocketSessionRegistry = WebSocketSessionRegistry() # Backend for Anthropic API (direct, LiteLLM, or any-llm) # Supports: "anthropic" (direct), "bedrock", "vertex", "litellm-", or "anyllm" diff --git a/headroom/proxy/ws_session_registry.py b/headroom/proxy/ws_session_registry.py new file mode 100644 index 000000000..b12e88a32 --- /dev/null +++ b/headroom/proxy/ws_session_registry.py @@ -0,0 +1,192 @@ +"""WebSocket session registry for Codex relay lifecycle tracking. + +Unit 3 of the Codex-proxy resilience plan. Every ``/v1/responses`` WS +session is explicitly registered on accept and deregistered in the +outermost ``finally`` of the handler. The registry provides: + +* First-class visibility of active sessions (``/debug/ws-sessions`` + from Unit 5 consumes :meth:`WebSocketSessionRegistry.snapshot`). +* Gauges for Prometheus (``active_ws_sessions``, ``active_relay_tasks``). +* A home for relay-task references so the handler's orchestrator can + attach them at creation time for introspection, without the registry + itself owning or cancelling them — cancellation is the handler's job. + +Design notes +------------ +* Single-event-loop usage. All mutations happen on the proxy's event + loop from within the handler coroutine, so no asyncio.Lock / asyncio + primitives are needed. Python dict mutations are atomic under the + GIL; snapshot copies happen under the event loop's cooperative + scheduling, so iteration + mutation cannot interleave across an + ``await`` point. +* ``deregister`` must be idempotent. The handler calls it from the + outermost ``finally`` — if upstream never connected, the session may + have been registered or may not have been, depending on how far + handshake got. Either way ``deregister`` is safe. +* ``deregister`` clears the handle's ``relay_tasks`` list so the + registry does not retain references to task coroutine frames after a + session ends. +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Any, Protocol + +logger = logging.getLogger(__name__) + +__all__ = [ + "WebSocketSessionRegistry", + "WSSessionHandle", + "TerminationCause", +] + + +TerminationCause = str # "client_disconnect" | "upstream_disconnect" | "upstream_error" +# | "client_error" | "response_completed" | "unknown" + + +class _TaskLike(Protocol): + def done(self) -> bool: ... + def cancel(self) -> bool: ... + def get_name(self) -> str: ... + + +@dataclass +class WSSessionHandle: + """Per-session state entry held in :class:`WebSocketSessionRegistry`. + + Timestamps use :func:`time.perf_counter` so age computations are + monotonic and independent of wall-clock adjustments. + """ + + session_id: str + request_id: str + client_addr: str | None = None + upstream_url: str | None = None + started_at: float = field(default_factory=time.perf_counter) + last_activity_at: float = field(default_factory=time.perf_counter) + relay_tasks: list[Any] = field(default_factory=list) + termination_cause: TerminationCause | None = None + + def mark_activity(self) -> None: + self.last_activity_at = time.perf_counter() + + def age_seconds(self) -> float: + return max(0.0, time.perf_counter() - self.started_at) + + def to_snapshot_dict(self) -> dict[str, Any]: + return { + "session_id": self.session_id, + "request_id": self.request_id, + "client_addr": self.client_addr, + "upstream_url": self.upstream_url, + "age_seconds": self.age_seconds(), + "idle_seconds": max(0.0, time.perf_counter() - self.last_activity_at), + "relay_task_count": len(self.relay_tasks), + "relay_task_names": [ + getattr(t, "get_name", lambda: "")() for t in self.relay_tasks + ], + "termination_cause": self.termination_cause, + } + + +class WebSocketSessionRegistry: + """In-memory registry of active Codex WS sessions. + + Methods are safe to call from the event loop. Unit 5's + ``/debug/ws-sessions`` endpoint consumes :meth:`snapshot`; the + Prometheus exporter reads :meth:`active_count` and + :meth:`active_relay_task_count`. + """ + + def __init__(self) -> None: + self._sessions: dict[str, WSSessionHandle] = {} + # Tracked separately from ``sum(len(h.relay_tasks))`` so that + # repeated snapshotting is O(1) rather than O(N * tasks). + self._active_relay_tasks = 0 + + # ------------------------------------------------------------------ + # Core lifecycle + # ------------------------------------------------------------------ + + def register(self, handle: WSSessionHandle) -> None: + """Register a session. Idempotent by ``session_id``. + + If the session id is already present, the existing entry is + replaced and the active count stays the same. This matches the + invariant that one ``session_id`` corresponds to at most one + active session at a time. + """ + existing = self._sessions.get(handle.session_id) + if existing is not None: + # Re-registration: release any tasks we were tracking on the + # old handle before replacing it. + self._active_relay_tasks -= len(existing.relay_tasks) + existing.relay_tasks.clear() + self._sessions[handle.session_id] = handle + self._active_relay_tasks += len(handle.relay_tasks) + + def deregister( + self, session_id: str, cause: TerminationCause = "unknown" + ) -> WSSessionHandle | None: + """Remove a session. Idempotent: returns ``None`` if unknown. + + Also clears the handle's ``relay_tasks`` list so the registry + stops holding references to coroutine frames after the session + ends. + """ + handle = self._sessions.pop(session_id, None) + if handle is None: + return None + handle.termination_cause = cause + self._active_relay_tasks = max( + 0, self._active_relay_tasks - len(handle.relay_tasks) + ) + handle.relay_tasks.clear() + return handle + + def attach_tasks(self, session_id: str, tasks: Iterable[_TaskLike]) -> None: + """Attach relay tasks to an existing session (merge, not replace). + + If ``session_id`` is not registered, this is a no-op (handler + should have registered before spawning tasks; defensive so a + race during deregister doesn't crash the handler). + """ + handle = self._sessions.get(session_id) + if handle is None: + return + task_list = list(tasks) + handle.relay_tasks.extend(task_list) + self._active_relay_tasks += len(task_list) + handle.mark_activity() + + # ------------------------------------------------------------------ + # Introspection + # ------------------------------------------------------------------ + + def get(self, session_id: str) -> WSSessionHandle | None: + return self._sessions.get(session_id) + + def active_count(self) -> int: + return len(self._sessions) + + def active_relay_task_count(self) -> int: + return self._active_relay_tasks + + def snapshot(self) -> list[dict[str, Any]]: + """JSON-serializable view of the registry (for ``/debug/ws-sessions``).""" + return [handle.to_snapshot_dict() for handle in self._sessions.values()] + + # ------------------------------------------------------------------ + # Convenience + # ------------------------------------------------------------------ + + def __contains__(self, session_id: str) -> bool: + return session_id in self._sessions + + def __len__(self) -> int: + return len(self._sessions) diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py new file mode 100644 index 000000000..f2dc4e365 --- /dev/null +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -0,0 +1,447 @@ +"""Unit 3: WebSocket session lifecycle + deterministic relay cancellation. + +These tests exercise the Codex WS handler with a fake upstream and a +fake client WebSocket so we can drive the relay halves through their +real code paths (not mocked) and assert on registry / task state. +""" + +from __future__ import annotations + +import asyncio +import json +import sys +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from headroom.proxy.handlers.openai import OpenAIHandlerMixin +from headroom.proxy.ws_session_registry import WebSocketSessionRegistry + +# --------------------------------------------------------------------------- +# Test doubles +# --------------------------------------------------------------------------- + + +class _DummyMetrics: + def __init__(self) -> None: + self.active_ws_sessions = 0 + self.active_ws_sessions_max = 0 + self.active_relay_tasks = 0 + self.ws_session_durations: list[float] = [] + self.stage_timings: list[tuple[str, dict[str, float]]] = [] + self.termination_causes: list[str] = [] + + async def record_request(self, **kwargs): # pragma: no cover + return None + + async def record_stage_timings(self, path: str, timings: dict[str, float]) -> None: + self.stage_timings.append((path, dict(timings))) + + def inc_active_ws_sessions(self) -> None: + self.active_ws_sessions += 1 + self.active_ws_sessions_max = max( + self.active_ws_sessions_max, self.active_ws_sessions + ) + + def dec_active_ws_sessions(self) -> None: + self.active_ws_sessions = max(0, self.active_ws_sessions - 1) + + def inc_active_relay_tasks(self, n: int = 1) -> None: + self.active_relay_tasks += n + + def dec_active_relay_tasks(self, n: int = 1) -> None: + self.active_relay_tasks = max(0, self.active_relay_tasks - n) + + def record_ws_session_duration(self, duration_ms: float, cause: str) -> None: + self.ws_session_durations.append(duration_ms) + self.termination_causes.append(cause) + + +class _DummyOpenAIHandler(OpenAIHandlerMixin): + OPENAI_API_URL = "https://api.openai.com" + + def __init__(self, ws_sessions: WebSocketSessionRegistry | None = None) -> None: + self.rate_limiter = None + self.metrics = _DummyMetrics() + self.config = SimpleNamespace( + optimize=False, + retry_max_attempts=1, + retry_base_delay_ms=1, + retry_max_delay_ms=1, + connect_timeout_seconds=10, + ) + self.usage_reporter = None + self.openai_provider = SimpleNamespace(get_context_limit=lambda model: 128_000) + self.openai_pipeline = SimpleNamespace(apply=MagicMock()) + self.anthropic_backend = None + self.cost_tracker = None + self.memory_handler = None + self.ws_sessions = ws_sessions or WebSocketSessionRegistry() + + async def _next_request_id(self) -> str: + return "req-lifecycle-test" + + +class _FakeWebSocketDisconnect(Exception): + """Mirrors the ``WebSocketDisconnect`` type-name check in the handler. + + The production code identifies "normal client gone" by + ``"WebSocketDisconnect" in type(e).__name__`` — so the fake exception + type name must start with ``WebSocketDisconnect``. + """ + + +# Force the type-name substring match in the handler. +_FakeWebSocketDisconnect.__name__ = "WebSocketDisconnect_Fake" + + +class _FakeWebSocket: + """Scripted client WebSocket that can delay / disconnect mid-stream.""" + + def __init__( + self, + frames: list[str] | None = None, + *, + disconnect_after_n_sends: int | None = None, + hold_after_initial: bool = False, + ) -> None: + self.headers = {"authorization": "Bearer test"} + self._frames = list(frames or []) + self._hold_after_initial = hold_after_initial + self._disconnect_after_n_sends = disconnect_after_n_sends + self.sent_text: list[str] = [] + self.sent_bytes: list[bytes] = [] + self.accepted_subprotocol: str | None = None + self.closed = False + self.close_code: int | None = None + # "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: + self.accepted_subprotocol = subprotocol + + async def receive_text(self) -> str: + if self._frames: + return self._frames.pop(0) + if self._hold_after_initial: + # Wait for simulated client disconnect. + await self._disconnect_event.wait() + # Use an exception type whose name starts with ``WebSocketDisconnect`` + # so the handler's ``type(e).__name__`` check classifies this as a + # normal client exit (not a ``client_error``). + raise _FakeWebSocketDisconnect("client closed") + + async def send_text(self, text: str) -> None: + self.sent_text.append(text) + if ( + self._disconnect_after_n_sends is not None + and len(self.sent_text) >= self._disconnect_after_n_sends + ): + # Trigger the "client gone" signal the next receive_text will see. + self._disconnect_event.set() + + async def send_bytes(self, data: bytes) -> None: + self.sent_bytes.append(data) + + async def close(self, code: int | None = None, reason: str | None = None) -> None: + self.closed = True + self.close_code = code + + def trigger_disconnect(self) -> None: + self._disconnect_event.set() + + +class _FakeUpstream: + """Upstream that streams scripted events then optionally blocks. + + ``hold_after_events`` makes the async iterator wait forever after the + scripted events are exhausted — that mirrors a real upstream that + keeps the connection open after a ``response.completed`` event. The + handler's ``_upstream_to_client`` will block on it, so the only way + the outer ``asyncio.wait`` can progress is via the client-side task + completing — which is exactly the cancel-partner path we want to + test. + """ + + def __init__( + self, + events: list[str], + *, + hold_after_events: bool = False, + raise_mid_stream: Exception | None = 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 + + async def __aenter__(self) -> _FakeUpstream: + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + self.closed = True + + async def send(self, payload: str) -> None: + self.sent.append(payload) + + async def close(self) -> None: + self.closed = True + + def __aiter__(self): + return self._iter() + + async def _iter(self): + for ev in self._events: + yield ev + if self._raise_mid_stream is not None: + raise self._raise_mid_stream + if self._hold_after_events: + # Wait forever — until the task is cancelled by the handler. + await asyncio.Event().wait() + + +def _make_fake_websockets_module(upstream: _FakeUpstream): + module = MagicMock() + module.connect = MagicMock(return_value=upstream) + module.Subprotocol = str + return module + + +def _first_frame() -> str: + return json.dumps( + { + "type": "response.create", + "response": {"model": "gpt-5.4", "input": "hi"}, + } + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_happy_path_registry_empty_after_response_completed(): + """Normal session completes — both relay tasks done, registry empty.""" + 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) + + 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 handler.ws_sessions.active_count() == 0 + assert handler.metrics.active_ws_sessions == 0 + # termination_cause captured + assert handler.metrics.termination_causes + # Either "response_completed" or "client_disconnect" — both are + # acceptable here depending on which relay half exited first; the + # important thing is we recorded one. + assert handler.metrics.termination_causes[-1] in { + "response_completed", + "client_disconnect", + "upstream_disconnect", + } + + +@pytest.mark.asyncio +async def test_client_disconnect_cancels_upstream_relay_within_100ms(): + """**Failing-test-first** scenario from the plan. + + When the client side exits (``receive_text`` raises + ``WebSocketDisconnect``) while upstream is still open and iterating, + the upstream relay task must be cancelled and become ``done()`` + quickly. The registry must report no active sessions afterwards. + """ + # Upstream keeps iterating forever after one event, forcing the + # upstream-to-client task to block on the iterator. The only way + # out is a cancel from the handler's orchestration. + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + ] + upstream = _FakeUpstream(upstream_events, hold_after_events=True) + fake_ws_mod = _make_fake_websockets_module(upstream) + + # Client has one initial frame, then disconnects after the server + # sends the first forwarded event to us. + client_ws = _FakeWebSocket( + frames=[_first_frame()], + hold_after_initial=True, + ) + handler = _DummyOpenAIHandler() + + # Trigger disconnect shortly after the handler accepts. + async def _trigger() -> None: + await asyncio.sleep(0.05) + client_ws.trigger_disconnect() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + trigger_task = asyncio.create_task(_trigger()) + try: + await asyncio.wait_for( + handler.handle_openai_responses_ws(client_ws), + timeout=2.0, + ) + finally: + trigger_task.cancel() + try: + await trigger_task + except asyncio.CancelledError: + pass + + # Registry must be empty — the finally block deregistered the session. + assert handler.ws_sessions.active_count() == 0, ( + "session leaked — deregister did not run in outermost finally" + ) + assert handler.metrics.active_ws_sessions == 0 + # We recorded a session duration (came through deregister path). + assert handler.metrics.ws_session_durations, ( + "record_ws_session_duration never fired — deregister path broken" + ) + # And we tagged the cause. For a client-side exit it should be one + # of: client_disconnect, client_error, upstream_disconnect (if + # upstream iteration happened to end first in a race). + cause = handler.metrics.termination_causes[-1] + assert cause in { + "client_disconnect", + "client_error", + "upstream_disconnect", + }, f"unexpected cause: {cause}" + + # No codex-ws-* named task should still be running. + leaked = [ + t + for t in asyncio.all_tasks() + if (t.get_name() or "").startswith("codex-ws-") and not t.done() + ] + assert leaked == [], f"relay tasks leaked: {[t.get_name() for t in leaked]}" + + +@pytest.mark.asyncio +async def test_upstream_closes_first_cancels_client_task(): + """Upstream iterator ends naturally; client task should be cancelled. + + The client is set to block on ``receive_text`` indefinitely; only a + cancel from the handler's orchestration releases it. + """ + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + upstream = _FakeUpstream(upstream_events, hold_after_events=False) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket( + frames=[_first_frame()], + hold_after_initial=True, + ) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await asyncio.wait_for( + handler.handle_openai_responses_ws(client_ws), + timeout=2.0, + ) + + assert handler.ws_sessions.active_count() == 0 + # We must still have recorded exactly one session duration. + assert len(handler.metrics.ws_session_durations) == 1 + + +@pytest.mark.asyncio +async def test_upstream_error_mid_stream_classifies_as_upstream_error(): + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + ] + upstream = _FakeUpstream( + upstream_events, + raise_mid_stream=RuntimeError("boom from upstream"), + ) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket( + frames=[_first_frame()], + hold_after_initial=True, + ) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await asyncio.wait_for( + handler.handle_openai_responses_ws(client_ws), + timeout=2.0, + ) + + assert handler.ws_sessions.active_count() == 0 + assert handler.metrics.termination_causes + assert handler.metrics.termination_causes[-1] == "upstream_error" + + +@pytest.mark.asyncio +async def test_upstream_connect_failure_still_deregisters_cleanly(): + """Handshake-phase leak must be impossible: if upstream connect + raises before relay tasks are created, the session is still + 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 + + client_ws = _FakeWebSocket(frames=[_first_frame()]) + handler = _DummyOpenAIHandler() + + async def _fallback(*args, **kwargs): + return None + + 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) + + 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.""" + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + + async def run_one() -> None: + upstream = _FakeUpstream(list(upstream_events)) + fake_ws_mod = _make_fake_websockets_module(upstream) + 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 handler.ws_sessions.active_count() == 0 + + await asyncio.gather(*[run_one() for _ in range(50)]) + + # Global check: no codex-ws-* named task remains. + leaked = [ + t + for t in asyncio.all_tasks() + if (t.get_name() or "").startswith("codex-ws-") and not t.done() + ] + assert leaked == [] diff --git a/tests/test_ws_session_registry.py b/tests/test_ws_session_registry.py new file mode 100644 index 000000000..55b6d8e04 --- /dev/null +++ b/tests/test_ws_session_registry.py @@ -0,0 +1,177 @@ +"""Unit tests for :class:`WebSocketSessionRegistry`. + +These tests exercise the in-memory registry in isolation — no network, +no WS server. They pin down register/deregister semantics, ``snapshot`` +shape, and the task-attachment accounting that feeds the +``active_relay_tasks`` gauge. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from headroom.proxy.ws_session_registry import ( + WebSocketSessionRegistry, + WSSessionHandle, +) + + +def _make_handle(session_id: str = "sess-1") -> WSSessionHandle: + return WSSessionHandle( + session_id=session_id, + request_id="req-1", + client_addr="127.0.0.1:12345", + upstream_url="wss://upstream/test", + ) + + +def test_register_and_deregister_round_trip(): + reg = WebSocketSessionRegistry() + handle = _make_handle() + + reg.register(handle) + assert reg.active_count() == 1 + assert reg.get("sess-1") is handle + + reg.deregister("sess-1", cause="response_completed") + assert reg.active_count() == 0 + assert reg.get("sess-1") is None + assert handle.termination_cause == "response_completed" + + +def test_deregister_unknown_session_is_idempotent(): + reg = WebSocketSessionRegistry() + # Deregistering a never-registered session must not raise. + reg.deregister("never-seen", cause="unknown") + assert reg.active_count() == 0 + + +def test_register_twice_same_session_does_not_double_count(): + reg = WebSocketSessionRegistry() + handle = _make_handle() + + reg.register(handle) + reg.register(handle) # idempotent — overwrite, do not double-increment + assert reg.active_count() == 1 + + reg.deregister("sess-1", cause="client_disconnect") + assert reg.active_count() == 0 + + +def test_snapshot_shape_is_json_serializable(): + reg = WebSocketSessionRegistry() + reg.register(_make_handle("a")) + reg.register(_make_handle("b")) + + snapshot = reg.snapshot() + assert len(snapshot) == 2 + ids = {entry["session_id"] for entry in snapshot} + assert ids == {"a", "b"} + sample = snapshot[0] + for key in ( + "session_id", + "request_id", + "client_addr", + "upstream_url", + "age_seconds", + "relay_task_count", + "termination_cause", + ): + assert key in sample + + +@pytest.mark.asyncio +async def test_attach_tasks_merges_and_tracks_active_count(): + reg = WebSocketSessionRegistry() + handle = _make_handle() + reg.register(handle) + + async def _idle() -> None: + await asyncio.sleep(5) + + t1 = asyncio.create_task(_idle(), name="codex-ws-test-1") + t2 = asyncio.create_task(_idle(), name="codex-ws-test-2") + + reg.attach_tasks("sess-1", [t1]) + assert reg.active_relay_task_count() == 1 + + reg.attach_tasks("sess-1", [t2]) # merges, does not replace + assert reg.active_relay_task_count() == 2 + + # Deregister cancels the in-registry accounting; tasks themselves + # are the caller's responsibility to cancel — the registry is a + # bookkeeper, not an owner. + reg.deregister("sess-1", cause="client_disconnect") + assert reg.active_relay_task_count() == 0 + + t1.cancel() + t2.cancel() + for t in (t1, t2): + try: + await t + except asyncio.CancelledError: + pass + + +def test_deregister_releases_task_references(): + """Once deregistered, the handle's ``relay_tasks`` list is cleared. + + This keeps large references (coroutine frames) from being held by + the registry after the session ends. + """ + reg = WebSocketSessionRegistry() + handle = _make_handle() + reg.register(handle) + + # Use sentinel non-task objects; ``attach_tasks`` does not inspect type + # beyond appending to the list and updating counts. + class _FakeTask: + def __init__(self, name: str) -> None: + self._name = name + self._done = False + + def done(self) -> bool: + return self._done + + def cancel(self) -> bool: + self._done = True + return True + + def get_name(self) -> str: + return self._name + + fake_tasks = [_FakeTask("t-a"), _FakeTask("t-b")] + reg.attach_tasks("sess-1", fake_tasks) + assert len(handle.relay_tasks) == 2 + + reg.deregister("sess-1", cause="upstream_disconnect") + # After deregister, the handle's tasks list is empty — the registry + # does not retain references. + assert handle.relay_tasks == [] + assert reg.active_relay_task_count() == 0 + + +def test_attach_tasks_on_missing_session_is_noop(): + reg = WebSocketSessionRegistry() + + async def _idle() -> None: + await asyncio.sleep(0) + + loop = asyncio.new_event_loop() + try: + t = loop.create_task(_idle()) + reg.attach_tasks("missing", [t]) + assert reg.active_relay_task_count() == 0 + t.cancel() + loop.run_until_complete(asyncio.gather(t, return_exceptions=True)) + finally: + loop.close() + + +def test_snapshot_age_seconds_is_non_negative(): + reg = WebSocketSessionRegistry() + reg.register(_make_handle()) + snapshot = reg.snapshot() + assert snapshot[0]["age_seconds"] >= 0.0 diff --git a/wiki/plans/2026-04-17-fix-codex-proxy-resilience-plan.md b/wiki/plans/2026-04-17-fix-codex-proxy-resilience-plan.md new file mode 100644 index 000000000..e9985f186 --- /dev/null +++ b/wiki/plans/2026-04-17-fix-codex-proxy-resilience-plan.md @@ -0,0 +1,487 @@ +--- +title: "fix: Codex proxy resilience under reconnect storms" +type: fix +status: active +date: 2026-04-17 +origin: wiki/plans/2026-04-17-codex-proxy-runtime-analysis.md +--- + +# fix: Codex proxy resilience under reconnect storms + +## Overview + +Harden the shared Headroom proxy so it can survive real multi-agent Codex traffic — especially the **large Anthropic `/v1/messages?beta=true` reconnect/retry storm** that hits the proxy immediately after a restart — without appearing dead (`/livez` timing out, new `/v1/responses` websocket handshakes hanging) and without bypassing compression. + +This is the follow-on work to the runtime analysis captured in the origin document. The previous branch (`fix/responses-retries-keep-compression`) fixed the upstream WS handshake/fallback issues and kept compression enabled. This plan addresses the **remaining long-lived runtime degradation** described in §4 of the origin ("Long-lived service degradation on `8787`"). + +The plan deliberately focuses on **observability + lifecycle hygiene + cold-start backpressure**, not more blind patching. Compression stays enabled throughout. + +## Problem Frame + +Aged `8787` processes enter a state where: + +- new `GET /livez` requests time out +- new `/v1/responses` opening handshakes time out +- existing established streams continue working +- the process is alive, listens on the port, and sampling shows heavy ONNX thread activity + +Controlled reproductions ruled out the obvious single-factor causes (port, launchd, memory alone, idle socket count, compression-preserving changes, cold Kompress load). The surviving hypotheses (§"What Is Still Plausible" in origin) converge on: + +1. **Long-lived real traffic leaves stuck websocket relay tasks / lifecycle bookkeeping leaks** (H1) +2. **Real Codex traffic, not synthetic traffic, triggers slow hidden work** (H2) +3. **ONNX/memory amplifies but does not solely cause the failure** (H3) +4. **Shared-proxy reconnect/retry storms after restart drive the proxy into this state** (H4 — confirmed by the "Latest Correction" in origin) + +Without stage timings, active-task introspection, or session bookkeeping, the next iteration of debugging will again rely on `sample`, `lsof`, and guesswork. This plan fixes that first, then layers cold-start backpressure on top — in that order — so each subsequent bug hunt converges faster. + +(see origin: `wiki/plans/2026-04-17-codex-proxy-runtime-analysis.md`) + +## Requirements Trace + +- **R1.** `/livez` remains responsive during and immediately after a restart that triggers large Anthropic replay traffic from active agent sessions. (origin §"Latest Correction", §"Updated upstream patch focus" item 4) +- **R2.** Cold-start heavy assets (Kompress ONNX, memory embedder, tokenizers, tree-sitter parsers) are loaded once at startup and *shared* between all provider pipelines; concurrent first-use callers wait on a single future, not N parallel loads. (origin §"Updated upstream patch focus" item 1) +- **R3.** The proxy provides enough runtime observability to prove — with data, not guesses — whether a future degradation is WS lifecycle starvation, memory/embedder contention, replay-storm amplification, or something else. (origin §"Priority 1: add instrumentation, not more blind patching") +- **R4.** Codex websocket relay tasks are explicitly tracked and deterministically cancelled when either side of the relay exits; a leaked relay task cannot hold the process alive past the client's disconnect. (origin §"Code Paths Most Relevant To The Remaining Bug" item 3; H1) +- **R5.** Cold-start compression + memory-context work on the Anthropic path is **bounded in concurrency** so that N simultaneous large replay requests cannot monopolize the event loop and thread pool. Compression stays enabled. (origin §"Updated upstream patch focus" item 2) +- **R6.** A reproducible harness exists for the real-agent reconnect/retry scenario so regressions can be caught locally instead of only in production. (origin §"Priority 2: reproduce degradation with real Codex traffic on a fresh process") +- **R7.** Existing fork behavior is preserved: compression is not bypassed on WS/streaming; upstream WS retry/open-timeout hardening is kept; WS→HTTP fallback normalization is kept; memory-context fail-open timeout is kept. (origin §"What Was Changed In The Fork", §"Kept locally") + +## Scope Boundaries + +- **Not** re-introducing any "skip compression" fast paths. Compression-preserving direction is non-negotiable. +- **Not** re-litigating the launchd setup bug (already fixed upstream in dotfiles; out of this repo). +- **Not** redesigning the memory stack, embedder choice, or Kompress model. Those are upstream concerns. +- **Not** building a full distributed tracing system. The instrumentation added here is structured logs + in-process counters; OpenTelemetry hookup is a separate plan. +- **Not** changing the WS→HTTP fallback semantics beyond what's needed to plumb through the new request-id/session-id logging fields. +- **Not** modifying `headroom-ai[ml]` dependencies, HuggingFace model IDs, or the embedder's own backend. + +### Deferred to Separate Tasks + +- **OpenTelemetry / metrics exporter wiring**: the counters added here expose `prometheus_metrics` entries; OTLP export belongs in a follow-up. +- **External watchdog in the LaunchAgent**: a plist-level `WatchPaths`/`ThrottleInterval` revision belongs in the `.dotfiles` repo (see origin §"Important External Files"), not here. This plan only adds the in-process signal the watchdog would consume (§Unit 5). +- **Quantifying multi-agent reconnect budget**: tuning the `Unit 4` semaphore default via load testing is work for after merge. + +## Context & Research + +### Relevant Code and Patterns + +- `headroom/proxy/handlers/openai.py:1206` — `handle_openai_responses_ws` (Codex WS entry point, 600+ lines) +- `headroom/proxy/handlers/openai.py:1559-1767` — upstream WS connect retry loop, `open_timeout` handling +- `headroom/proxy/handlers/openai.py:1815` — `_ws_http_fallback` (preserved as-is) +- `headroom/proxy/handlers/openai.py:1439-1452` — memory context timeout fail-open (preserved as-is) +- `headroom/proxy/handlers/anthropic.py:293` — `handle_anthropic_messages` (HTTP entry point, also covers the `?beta=true` replay traffic) +- `headroom/proxy/server.py:586-733` — `HeadroomProxy.startup` (where eager preload runs) +- `headroom/proxy/server.py:634-637` — current eager preload iterates **only** `anthropic_pipeline.transforms` and breaks on first match +- `headroom/proxy/server.py:301-308` — both pipelines currently share the same `transforms` list (so the module-level `_kompress_cache` is de facto shared, but this is fragile) +- `headroom/proxy/server.py:1190-1212` — `/livez`, `/readyz`, `/health` handlers (trivial JSON, no I/O) +- `headroom/proxy/memory_handler.py:134-207` — `MemoryHandler._ensure_initialized` (lazy, no `asyncio.Lock`) +- `headroom/transforms/content_router.py:1221-1297` — `eager_load_compressors` (Kompress + Magika + Code-Aware + SmartCrusher) +- `headroom/transforms/kompress_compressor.py:163-221` — `_load_kompress_onnx` / `_load_kompress` with module-level `_kompress_cache` + `threading.Lock` +- `headroom/proxy/request_logger.py` — existing structured request log sink (extend, don't replace) +- `headroom/proxy/prometheus_metrics.py` — existing `PrometheusMetrics` class; add counters/gauges here +- `headroom/proxy/helpers.py:138-207` — `_read_request_json` (pre-upstream work on Anthropic path) + +### Institutional Learnings + +- `docs/solutions/` does not exist in this repo. The `wiki/plans/` folder holds design-style documents; no rolling solutions log to mine. +- **Prior fork learning** (origin §"What Was Changed In The Fork"): keep compression on, retry upstream WS, normalize fallback body, wrap memory-context lookup in a timeout. These are invariants — Unit 2 and Unit 4 must not regress them. +- **Prior rollback learning** (origin §"Rolled back locally"): latency-first skips of compression / memory injection were rolled back. Any new "fast path" must not recreate that shape. + +### External References + +None used for this plan. Python `asyncio` lock, `asyncio.Semaphore`, `asyncio.Task`, and `asyncio.all_tasks()` semantics are sufficient; no framework-specific research needed. External research was intentionally skipped (§1.2: strong local patterns, team knows the area). + +## Key Technical Decisions + +- **Observe before mitigating.** Units 2 and 3 (instrumentation + lifecycle accounting) land before Unit 4 (backpressure) so the backpressure defaults can be tuned from real data instead of guessed. The origin explicitly calls this out as Priority 1. + - *Rationale:* the last round of patches was driven by symptoms; the next one should be driven by timings. +- **Share cold-start state across pipelines explicitly, not by accident.** Currently both `TransformPipeline` instances share the same `transforms` list by coincidence of construction (`server.py:301-308`). Unit 1 moves the preload out of "first matching transform on the Anthropic pipeline" into a startup-level orchestration step that holds references it can reuse across both pipelines, and adds a single `asyncio.Lock` around first-use paths so concurrent requests land on the same future. + - *Rationale:* origin §"Updated upstream patch focus" item 1: "make eager preload and request-time use share the same in-process singleton/cache". +- **Track WS sessions in a registry, not via ad-hoc `logger.info`.** A `WebSocketSessionRegistry` makes active-count a first-class observable and lets `/debug/ws-sessions` return something useful. +- **Explicit relay-task cancellation** replaces `asyncio.gather(..., return_exceptions=True)` for the two relay halves. When one side exits, the other is cancelled deterministically — no wait on TCP timeout, no task leak (addresses H1). +- **Bounded pre-upstream concurrency on the Anthropic path, not on all paths.** The Codex WS path already serializes naturally (one client → one upstream WS). The Anthropic HTTP path is where replay storms arrive. Limiting concurrency only where the problem actually exists keeps the blast radius tight. + - *Rationale:* origin §"Hypothesis 4" + §"Updated upstream patch focus" item 2. +- **Debug endpoints are loopback-only, always.** No config flag, no auth header — a remote IP gets a 404, period. This sidesteps "did someone accidentally expose task state?" as a concern. +- **Frontmatter, not freeform.** Unlike the existing `wiki/plans/` files this plan uses YAML frontmatter (`status: active`, `origin:`, etc.) to participate in the `ce:plan` deepening + search flow. + +## Open Questions + +### Resolved During Planning + +- *Q: Target upstream or the fork?* Resolved: target the fork (`fix/responses-retries-keep-compression`), structured so each unit is cherry-pickable into a PR against `chopratejas/headroom#172`. +- *Q: Should Unit 5's debug endpoints require an explicit flag to enable?* Resolved: **no** — loopback-only gating is sufficient and the debug data is useless if it isn't always available when the process is struggling. +- *Q: Does `MemoryHandler._ensure_initialized` already have a concurrency guard?* Resolved: **no**. It relies on `self._initialized = True` flip, which is not atomic across `await` points. Unit 1 adds an `asyncio.Lock`. +- *Q: Are the WS relay tasks already cancelled on partner exit?* Resolved: **no**. `asyncio.gather(return_exceptions=True)` waits for both; the survivor only exits when its own loop raises. Unit 3 fixes this. +- *Q: Is the existing Kompress singleflight sufficient?* Resolved: partially. The `threading.Lock` serializes same-model-id loads, but holds during `hf_hub_download` (network I/O). Unit 1 supplements with an `asyncio.Lock` at the request-handler layer so async callers don't each spawn the thread-pool job. + +### Deferred to Implementation + +- **Exact default for the Anthropic pre-upstream semaphore.** Start at `max(2, min(8, os.cpu_count()))` and expose via `--anthropic-pre-upstream-concurrency` / `HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY`. Tune after Unit 2's stage timings land. +- **Which asyncio task names are "long-lived" thresholds** for Unit 5's dump. Likely `> 2× median lifetime`, but the median is only observable after Unit 2 runs in production. +- **Whether the repro harness (Unit 6) needs to speak realistic Codex subprotocol framing or can synthesize enough with recorded frames.** Depends on how faithfully `tests/test_openai_codex_routing.py` fixtures already capture the handshake. +- **Final placement of `WebSocketSessionRegistry`**: `headroom/proxy/ws_session_registry.py` as a new module, or folded into `headroom/proxy/server.py`. Likely new module; confirm during implementation once imports are real. + +## High-Level Technical Design + +> *This illustrates the intended approach and is directional guidance for review, not implementation specification. The implementing agent should treat it as context, not code to reproduce.* + +``` + ┌───────────────────────────────────────────────────────┐ + │ HeadroomProxy.startup │ + │ │ + │ Unit 1: shared_warmup() │ + │ • eager-load on both pipelines (not just first) │ + │ • Kompress, Magika, Code-Aware, SmartCrusher │ + │ • memory backend + embedder │ + │ • populates WarmupRegistry singletons │ + └───────────────┬────────────────────┬──────────────────┘ + │ │ + ▼ ▼ + ┌─────────────────────────────┐ ┌─────────────────────────────┐ + │ Codex WS path │ │ Anthropic HTTP path │ + │ /v1/responses │ │ /v1/messages?beta=true │ + │ │ │ │ + │ Unit 3: session registry │ │ Unit 4: pre-upstream │ + │ • register on accept │ │ Semaphore(N) │ + │ • cancel partner on exit │ │ • guards _read_request │ + │ • deregister in finally │ │ • deep-copy │ + │ │ │ • first compression stage │ + │ Unit 2: stage timings │ │ • memory-context lookup │ + │ accept → first_frame │ │ │ + │ → upstream_connect │ │ Unit 2: stage timings │ + │ → upstream_first_event │ │ read_json → deep_copy → │ + │ → total_session_ms │ │ compress → memory → │ + └─────────────────────────────┘ │ upstream_first_byte │ + └─────────────────────────────┘ + │ │ + └────────┬───────────┘ + ▼ + ┌─────────────────────────────────────────┐ + │ Unit 5: /debug/* (loopback-only) │ + │ • /debug/tasks │ + │ • /debug/ws-sessions │ + │ • /debug/warmup │ + └─────────────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────────────┐ + │ Unit 6: scripts/repro_codex_replay.py │ + │ • spawn N concurrent Codex WS │ + │ • fire Anthropic replay-shaped POSTs │ + │ • assert /livez stays < 100ms │ + └─────────────────────────────────────────┘ +``` + +Decision matrix for "where does this request wait?": + +| Path | Pre-upstream gate | Relay lifecycle gate | Memory lookup timeout | +|-------------------------------|----------------------|----------------------|------------------------| +| `/v1/responses` (Codex WS) | *(none — natural 1:1)* | `WSSessionRegistry` | existing `wait_for` | +| `/v1/messages` (Anthropic) | `Semaphore(N)` (new) | *(HTTP-single-shot)* | existing `wait_for` | +| `/v1/chat/completions` (OpenAI) | *(unchanged)* | *(unchanged)* | *(unchanged)* | +| `/livez`, `/readyz` | *(none — must be free)* | *(n/a)* | *(n/a)* | + +## Implementation Units + +- [ ] **Unit 1: Shared cold-start warmup + async singleflight for memory init** + +**Goal:** Make preload truthful and shared across all provider pipelines; ensure the first wave of concurrent requests after restart never kick duplicate expensive loads. + +**Requirements:** R2, R7 + +**Dependencies:** None (pure startup change; no other unit depends on this landing first, but landing it first reduces noise in Unit 2's timings) + +**Files:** +- Modify: `headroom/proxy/server.py` (startup orchestration around lines 627-675) +- Modify: `headroom/proxy/memory_handler.py` (add `asyncio.Lock` in `_ensure_initialized`) +- Modify: `headroom/transforms/content_router.py` (no behavior change — ensure `eager_load_compressors` is idempotent when called via both pipelines) +- Create: `headroom/proxy/warmup.py` (new `WarmupRegistry` holding preloaded handles) +- Test: `tests/test_proxy_warmup.py` (new) +- Test: `tests/test_memory_handler_concurrent_init.py` (new) + +**Approach:** +- Introduce `WarmupRegistry` with typed slots for kompress, magika, code-aware, smart-crusher, memory-embedder, memory-backend. Populate during `HeadroomProxy.startup`, expose via `proxy.warmup` for `/debug/warmup` (Unit 5) and `/readyz`. +- Iterate **both** `self.anthropic_pipeline.transforms` and `self.openai_pipeline.transforms` when calling `eager_load_compressors`; dedupe by `id(transform)` so shared transforms don't double-load. +- Preload the memory embedder explicitly — today `ensure_initialized` initializes the backend but the embedder's first-use cost may still be deferred until the first `search_and_format_context`. Force one warm-up encode (a single short string) to pre-compile the ONNX graph. +- Add `asyncio.Lock` in `MemoryHandler._ensure_initialized` so concurrent first callers await one load, not N. +- Replace the "break after first matching transform" loop with explicit orchestration. Log whether the warmup was a no-op (already loaded) vs. a fresh load. + +**Patterns to follow:** +- Existing `eager_status` dict shape in `server.py:630-666`. +- `threading.Lock` module-level pattern in `kompress_compressor.py` (for the sync side); supplement with `asyncio.Lock` in memory_handler (async side). +- `request_logger`'s structured-log pattern for startup events. + +**Test scenarios:** +- Happy path: starting the proxy with `optimize=True` logs one preload event per component and `WarmupRegistry` reports all slots loaded. +- Happy path: starting with `optimize=False` populates `WarmupRegistry` on first request instead (lazy path still works). +- Edge case: `optimize=True` with `enable_kompress=False` — `WarmupRegistry.kompress` is `None`, no error logged. +- Edge case: both pipelines share the same router transform — `eager_load_compressors` runs once, not twice. +- Integration: `MemoryHandler._ensure_initialized` called 10 concurrent times from different tasks — only one backend init runs (assert via counter on `LocalBackend.__init__` hit count). +- Integration: embedder warm-up encode is issued at startup — verify via mock that `embed_text("warmup")` (or equivalent) was called once during startup, not lazily on first request. +- Error path: memory backend init raises — startup still completes, `WarmupRegistry.memory_backend` is `None`, health reports degraded memory. + +**Verification:** +- Startup logs show each component preloaded exactly once, with timings. +- `/readyz` reports all configured subsystems as `initialized=true` before accepting traffic. +- Concurrent first-request simulation (`asyncio.gather` 20 requests) triggers only one Kompress load in logs. + +--- + +- [ ] **Unit 2: Stage-timing instrumentation on Codex WS and Anthropic HTTP paths** + +**Goal:** Emit structured per-request timings for every stage that plausibly contributes to cold-start or degradation latency, so the next debugging session starts from data. + +**Requirements:** R3 + +**Dependencies:** None (independent), but landing alongside Unit 1 gives early visibility into whether warmup actually helped. + +**Files:** +- Modify: `headroom/proxy/handlers/openai.py` (WS path at lines 1206+; HTTP path at 798+) +- Modify: `headroom/proxy/handlers/anthropic.py` (`handle_anthropic_messages` around line 293 and downstream) +- Modify: `headroom/proxy/request_logger.py` (extend schema) +- Modify: `headroom/proxy/prometheus_metrics.py` (add histograms per stage) +- Modify: `headroom/proxy/helpers.py` (thread `stage_timer` through `_read_request_json` / compression helpers) +- Create: `headroom/proxy/stage_timer.py` (small context-manager util) +- Test: `tests/test_stage_timer.py` (new) +- Test: `tests/test_openai_codex_ws_timings.py` (new) +- Test: `tests/test_anthropic_stage_timings.py` (new) + +**Approach:** +- Add a `StageTimer` context manager with `stage_timer.measure("memory_context")` support; emits structured log on exit. Unified across handlers. +- For Codex WS: instrument `accept`, `first_client_frame`, `upstream_connect`, `upstream_first_event`, `memory_context`, `compression`, `total_session`. Log on session close with all fields. +- For Anthropic HTTP: instrument `read_request_json`, `deep_copy`, `compression_first_stage`, `memory_context`, `upstream_connect`, `upstream_first_byte`, `total_pre_upstream`. Log after first upstream byte arrives. +- Also export each stage as a Prometheus histogram — the existing `PrometheusMetrics` class has the right pattern; mirror it. +- `request_id` is already threaded through both handlers. Add a `session_id` (UUID generated at WS accept / HTTP request start) so multi-turn sessions are correlatable. + +**Execution note:** Land the util + test (`StageTimer`) first; only then plumb it through the two handlers. Keeps the diff reviewable. + +**Patterns to follow:** +- `headroom/proxy/request_logger.py` structured-field schema. +- `PrometheusMetrics.record_request` histogram pattern. +- Existing `request_id = await self._next_request_id()` at `openai.py:1232`. + +**Test scenarios:** +- Happy path: a full Codex WS session (accept → response.completed) emits one structured log line with all 7 stage fields populated and a positive `total_session` > 0. +- Happy path: a full Anthropic HTTP request emits one log line with all pre-upstream stage fields populated. +- Edge case: a session that exits during upstream connect (no `first_event`) logs `upstream_first_event=null` without raising. +- Edge case: `request_id` and `session_id` appear together on every log line. +- Error path: a timeout in `memory_context` logs `memory_context` duration = the timeout value, not `null` (prove the timer captures the failure window). +- Integration: Prometheus histograms emit non-zero observations after a real request round-trip. + +**Verification:** +- Every `/v1/responses` WS session and every `/v1/messages` request produces exactly one timing log line. +- `/metrics` endpoint shows the new histogram series. +- No existing tests regress; the request_logger schema additions are backward-compatible (new fields only). + +--- + +- [x] **Unit 3: WebSocket session registry + deterministic relay-task cancellation** + +**Goal:** Eliminate the "aged process has leaked relay tasks" hypothesis by making every WS session explicitly tracked and both relay tasks deterministically cancelled when either exits. + +**Requirements:** R4 + +**Dependencies:** Unit 2 (uses the same `session_id`) + +**Files:** +- Create: `headroom/proxy/ws_session_registry.py` (new) +- Modify: `headroom/proxy/handlers/openai.py` (`handle_openai_responses_ws` at 1206; relay task construction around 1559-1767; `_upstream_to_client` at 1565) +- Modify: `headroom/proxy/prometheus_metrics.py` (add `active_ws_sessions` gauge, `active_relay_tasks` gauge, `ws_session_duration` histogram) +- Test: `tests/test_ws_session_registry.py` (new) +- Test: `tests/test_openai_codex_ws_lifecycle.py` (new) + +**Approach:** +- `WebSocketSessionRegistry` holds `dict[session_id, WSSessionHandle]` where handle tracks: `started_at`, `client_addr`, `upstream_url`, `relay_tasks: list[asyncio.Task]`, `last_activity_at`. +- Register on `websocket.accept()` success; deregister in a `try/finally` around the whole handler body. +- Replace `asyncio.gather(_client_to_upstream(), _upstream_to_client(), return_exceptions=True)` with explicit `asyncio.create_task(...)` for each, then `asyncio.wait(..., return_when=FIRST_COMPLETED)`. Cancel the other task explicitly, then `await` the cancelled task's `.cancelled()` settlement (suppress `CancelledError`). +- Emit one structured log line per session termination with cause (`client_disconnect` / `upstream_disconnect` / `upstream_error` / `client_error`). + +**Execution note:** Start from a failing integration test that asserts "after a client disconnects mid-stream, the upstream relay task is `done()` within 100ms and the registry is empty". Then make it pass. + +**Patterns to follow:** +- Existing `_client_to_upstream()` and `_upstream_to_client()` structure in `openai.py:1559-1767`. +- `contextlib.suppress(asyncio.CancelledError)` pattern. +- `PrometheusMetrics` gauge update pattern. + +**Test scenarios:** +- Happy path: a normal session completes; registry is empty after `response.completed`; both relay tasks `.done()` is True. +- Edge case: client disconnects mid-stream; upstream relay task is cancelled within 100ms; registry no longer contains the session. +- Edge case: upstream closes first; client-side relay task is cancelled within 100ms; client receives a clean close frame. +- Edge case: 50 concurrent WS sessions open and close; `active_ws_sessions` gauge rises to 50 then returns to 0; no tasks remain in `asyncio.all_tasks()` with the relay task name pattern. +- Error path: upstream connect fails before relay tasks are created; registry is deregistered cleanly (no leak from the handshake phase). +- Error path: `_upstream_to_client` raises mid-stream; client task is cancelled; registry reports `termination_cause=upstream_error`. +- Integration: the `/debug/ws-sessions` endpoint (Unit 5) returns live data consistent with the registry under load. + +**Verification:** +- After a torture-test of 100 rapid connect/disconnect cycles, `asyncio.all_tasks()` count returns to baseline within 1s. +- `active_ws_sessions` gauge returns to 0 after all sessions close. +- No `RuntimeWarning: coroutine was never awaited` in logs. + +--- + +- [ ] **Unit 4: Bounded pre-upstream concurrency for Anthropic replay storms** + +**Goal:** Prevent the cold-start replay storm from occupying every event-loop slot and thread-pool worker with deep-copy / compression / memory-context work before upstream receives the request. + +**Requirements:** R1, R5, R7 + +**Dependencies:** Unit 2 (must have stage timings landed so the default semaphore size can be tuned from real numbers; the unit itself can ship with a conservative default and be retuned later) + +**Files:** +- Modify: `headroom/proxy/handlers/anthropic.py` (wrap the pre-upstream phases in `handle_anthropic_messages` at ~293 through first upstream call) +- Modify: `headroom/proxy/models.py` or config module (add `anthropic_pre_upstream_concurrency: int` config field) +- Modify: `headroom/proxy/server.py` (construct the `asyncio.Semaphore` on `HeadroomProxy.__init__` so it's per-process, not per-request) +- Modify: CLI surface (add `--anthropic-pre-upstream-concurrency` flag + env var) +- Test: `tests/test_anthropic_pre_upstream_backpressure.py` (new) + +**Approach:** +- Add `self.anthropic_pre_upstream_sem = asyncio.Semaphore(config.anthropic_pre_upstream_concurrency or max(2, min(8, os.cpu_count())))` in `HeadroomProxy.__init__`. +- In `handle_anthropic_messages`, wrap the region from `_read_request_json` through the first `self.http_client.send()` / `stream()` call in `async with self.anthropic_pre_upstream_sem:`. +- **Critically:** `/livez` and `/readyz` do not go through this semaphore. They remain free even under replay storm. +- Emit a log (not a warning — expected under load) when a request waits > 100ms for the semaphore, so we can see queueing in the Unit 2 stage timings. +- Default config: `HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY` env var honored; CLI flag overrides; unset defaults to the computed floor. + +**Patterns to follow:** +- Existing config field addition pattern in `headroom/proxy/models.py`. +- Existing CLI flag plumbing in `headroom/cli/`. +- `request_logger.info(..., stage="pre_upstream_wait_ms", ...)` field extension from Unit 2. + +**Test scenarios:** +- Happy path: a single request passes through with negligible `pre_upstream_wait_ms`. +- Edge case: N+1 concurrent requests (where N = configured concurrency) — exactly one waits; all complete; `pre_upstream_wait_ms > 0` only on the waiter. +- Edge case: `anthropic_pre_upstream_concurrency=1` serializes two concurrent requests deterministically (useful for tests). +- Error path: a request that raises inside the critical section releases the semaphore (verify via counter that `_value` returns to baseline). +- Integration: under a synthetic storm of 20 concurrent large POSTs (simulating the retry replay), `/livez` response time stays under 100ms (p99 from Unit 2 timings). +- Integration: compression is *not* bypassed — assert a known large body still produces a compressed upstream request (no regression of R7). + +**Verification:** +- Under `scripts/repro_codex_replay.py` (Unit 6), `/livez` p99 stays under 100ms during the storm phase. +- Counter-factual: setting concurrency to `10000` (effectively unbounded) reproduces the original starvation in the harness. + +--- + +- [ ] **Unit 5: Loopback-only debug introspection endpoints** + +**Goal:** Make "what is this process doing right now?" a single `curl` away when degradation happens, instead of `sample` + `lsof` + guesswork. + +**Requirements:** R3 + +**Dependencies:** Unit 3 (ws session data), Unit 1 (warmup registry data) + +**Files:** +- Modify: `headroom/proxy/server.py` (add `/debug/tasks`, `/debug/ws-sessions`, `/debug/warmup` routes near the existing `/livez` at 1190) +- Create: `headroom/proxy/debug_introspection.py` (pure functions that serialize state) +- Create: `headroom/proxy/loopback_guard.py` (middleware / dep that 404s non-loopback requests) +- Test: `tests/test_proxy_debug_endpoints.py` (new) + +**Approach:** +- `loopback_guard`: inspect `request.client.host`; if not in `{"127.0.0.1", "::1", "localhost"}`, return 404 with no body. A 404 (not 403) keeps debug endpoints invisible to external scanners. +- `/debug/tasks` returns `asyncio.all_tasks()` with: name, coro name, age, current-line (via `get_coro().cr_frame.f_code.co_qualname`), stack depth. Sort by age desc. +- `/debug/ws-sessions` returns the registry dump from Unit 3. +- `/debug/warmup` returns `WarmupRegistry` state from Unit 1 + whether each slot is `loaded` / `loading` / `null`. +- All three return JSON. None mutate state. None block. + +**Patterns to follow:** +- Existing `/readyz` handler shape in `server.py:1204-1207`. +- Loopback detection: `request.client.host` is already FastAPI's standard; mirror the guard style used elsewhere in the repo if any; otherwise a small dependency function. + +**Test scenarios:** +- Happy path: `curl 127.0.0.1:8787/debug/tasks` returns 200 with a JSON array; each entry has `name`, `age_seconds`, `coro`. +- Edge case: `/debug/tasks` called during a live WS session lists the relay tasks with non-zero age. +- Edge case: loopback guard returns 404 (not 403) for a simulated non-loopback client. +- Edge case: `/debug/warmup` reports `memory_backend=loaded` after Unit 1's startup completes; reports `loading` if called during startup (race window). +- Error path: `asyncio.all_tasks()` raising (hypothetical) — handler returns 500 with structured error, doesn't crash the server. +- Integration: `/debug/ws-sessions` output is consistent with the session count gauge from Unit 3 during a load test. + +**Verification:** +- Manual: from a remote IP, all three endpoints return 404. +- Manual: during the Unit 6 harness, `/debug/tasks` shows relay tasks matching the expected count. +- Documented in `wiki/proxy.md` under a new "Debug endpoints" subsection. + +--- + +- [ ] **Unit 6: Repro harness for multi-agent reconnect storm** + +**Goal:** Produce a single script that reproducibly exercises the failure class from origin §"Latest Correction" (active agent reconnects + large replay requests) against a fresh local proxy, so the fix is provable and regressions are catchable. + +**Requirements:** R6 + +**Dependencies:** Unit 2 (so the harness can assert on stage timings), Unit 4 (so the harness can verify backpressure helps) + +**Files:** +- Create: `scripts/repro_codex_replay.py` (new) +- Create: `scripts/fixtures/anthropic_replay_body.json` (recorded shape of a large replay request body, sanitized) +- Create: `scripts/fixtures/codex_response_create_frame.json` (recorded first-frame shape) +- Modify: `scripts/README.md` (add harness section) + +**Approach:** +- CLI: `python scripts/repro_codex_replay.py --url http://127.0.0.1:8787 --ws-clients 8 --anthropic-clients 4 --duration 30s` +- Phase 1 (warmup): open 1 WS, send one `response.create`, drain to `response.completed`. Confirms proxy is live. +- Phase 2 (storm): simultaneously: + - Open N Codex WS connections, send one `response.create` each, keep the session open for the duration. + - Fire M concurrent large Anthropic POSTs shaped like agent-reconnect replays (from the fixture). Each retries on connection error for up to 60s, mimicking real agent behavior. +- Throughout: probe `/livez` every 250ms, record p50/p95/p99. +- Exit code: non-zero if `/livez` p99 exceeds 500ms during the storm (soft assertion). +- Print a summary: per-phase timings, livez stats, count of successful Codex `response.completed`, count of Anthropic successes. + +**Execution note:** Harness-first is fine here — the fixture JSONs can be hand-crafted initially and swapped for captured ones later. + +**Patterns to follow:** +- Existing benchmark harness in `benchmarks/` for run-loop structure. +- `websockets` library usage already in `tests/test_proxy_codex_route_aliases.py`. + +**Test scenarios:** +- Test expectation: light — the harness is a script, not a library. A smoke test in `tests/test_scripts/test_repro_codex_replay_smoke.py` launches it against a mock server and verifies it exits 0 with the expected summary shape. +- Happy path: harness runs against a healthy proxy; exit code 0; summary shows livez p99 < 500ms. +- Error path: harness invoked with `--url` pointing at a closed port; exits with clear "connection refused" message, non-zero code, within 5 seconds. +- Integration: after Unit 4 lands, harness with default Anthropic concurrency limits shows livez p99 < 100ms; harness with `HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY=10000` (unbounded) reproduces the original starvation (p99 > 5s). + +**Verification:** +- CI smoke test runs the script against a mock server on every PR. +- `wiki/proxy.md` gets a "Reproducing the reconnect storm" subsection referencing the script. + +## System-Wide Impact + +- **Interaction graph:** the new `WarmupRegistry` (Unit 1) and `WebSocketSessionRegistry` (Unit 3) are accessed from `HeadroomProxy` during request handling and by `/debug/*` routes. Both are single-writer-per-event-loop; no locks needed beyond the `asyncio.Lock` in `MemoryHandler._ensure_initialized` (Unit 1). +- **Error propagation:** Unit 3 changes `asyncio.gather(..., return_exceptions=True)` to explicit `asyncio.wait(FIRST_COMPLETED)` + cancel. This means partner-side exceptions surface earlier — previously, a client-side error would still let upstream-side drain to completion. Verify in test scenarios that this does not drop in-flight frames that the client had already queued. +- **State lifecycle risks:** Unit 3's session registry must be deregistered in the outermost `finally` — a leak here re-creates the exact bug this plan fixes. Unit 1's `asyncio.Lock` must be released on exception paths (use `async with`, not manual acquire/release). +- **API surface parity:** no change to `/v1/responses` or `/v1/messages` request/response contract. The only new routes are `/debug/*`, loopback-gated. +- **Integration coverage:** Unit 3 + Unit 6 together are the main integration story — unit-test-only coverage of relay cancellation is insufficient; the repro harness is how we prove it in aggregate. +- **Unchanged invariants:** Compression stays enabled on all paths. Upstream WS retry/open-timeout handling (`openai.py:1559-1767`) is untouched. WS→HTTP fallback normalization (`openai.py:1815`) is untouched except for threading through the new `session_id` / stage-timer context. Memory-context fail-open `asyncio.wait_for` (`openai.py:1439-1452`, and the equivalent on the Anthropic path) is untouched. `/livez` logic stays trivial and IO-free. + +## Risks & Dependencies + +| Risk | Likelihood | Impact | Mitigation | +|------|-----------|--------|------------| +| Unit 3's new cancellation semantics drop in-flight frames the client had already queued on the upstream side. | Med | High (user-visible stream corruption) | Explicit test scenarios for "client disconnect mid-stream" verify no frames silently lost; harness (Unit 6) measures successful `response.completed` count. | +| Unit 4's semaphore default is too conservative and adds visible latency under normal multi-agent use. | Med | Med | Default computed from `cpu_count()` with a floor of 2 and ceiling of 8; tunable via CLI + env; Unit 2 timings will surface if the wait is user-visible. | +| Unit 4's semaphore default is too permissive and doesn't actually stop starvation under real replay storms. | Med | High (plan fails R1) | Unit 6 harness tests both bounded and unbounded configs; ship with conservative default, relax based on harness data post-merge. | +| Unit 5's `/debug/tasks` exposes sensitive state (e.g., in-flight request bodies via coro locals). | Low | Med (info-leak if endpoint ever becomes non-loopback) | Serializer strips coro locals; only task *metadata* (name, age, qualname) is exposed, never arguments. Loopback guard is enforced by middleware before the handler runs. | +| Unit 1's embedder warm-up encode triggers a HuggingFace download in CI and slows test runs. | Med | Low | Warm-up encode is skipped when `optimize=False`; test environment uses a stub embedder; document in `tests/conftest.py`. | +| The real degradation is *not* caused by WS task leaks or replay storms (our top hypotheses), and this plan's instrumentation exposes it but the mitigations don't fix it. | Low | Med | This is fine. The plan is observe-first; Unit 2 + Unit 5 give the data for the next iteration. The plan is valuable even if Units 3 and 4 turn out to be insufficient on their own. | +| `asyncio.Lock` added in Unit 1 is held across the memory backend init, which itself calls `await HierarchicalMemory.create()` — if that hangs, first-request + `ensure_initialized` both hang. | Low | High (deadlock at startup) | Wrap `_ensure_initialized` in a `wait_for(..., timeout=STARTUP_INIT_TIMEOUT_SECONDS)` (configurable, default 30s). On timeout, log error and leave `_initialized=False` so subsequent requests retry. | +| Tests that monkey-patch `MemoryHandler._initialized` directly break because of the new lock ordering. | Low | Low | Audit `tests/test_proxy_memory_integration.py` and fixtures; use `ensure_initialized` via the public entry point only. | + +## Documentation / Operational Notes + +- Add a "Debug endpoints" subsection to `wiki/proxy.md` covering the three new loopback-only routes, their output shape, and the loopback-only guarantee. +- Add a "Reproducing the reconnect storm" subsection to `wiki/proxy.md` referencing `scripts/repro_codex_replay.py`. +- Update `wiki/metrics.md` with the new histogram series names from Unit 2 and the new gauges from Unit 3. +- Update `CHANGELOG.md` under an "Unreleased" heading with: (a) compression preserved — no behavioral change for existing Codex users; (b) new stage timings visible in logs; (c) new `HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY` env var; (d) new `/debug/*` loopback endpoints. +- **Rollout note:** deploy behind the existing launchd flow; monitor `/metrics` for the new histograms; monitor logs for `pre_upstream_wait_ms` — sustained non-zero values across many requests mean the Unit 4 default is too tight for this machine's load profile. +- **Rollback note:** each unit is cherry-pickable. If Unit 4 causes regression, revert only Unit 4 (semaphore removal); Units 1–3 and 5–6 have no user-visible behavior change. + +## Sources & References + +- **Origin document:** `wiki/plans/2026-04-17-codex-proxy-runtime-analysis.md` +- Related code: + - `headroom/proxy/handlers/openai.py` — `handle_openai_responses_ws`, `_client_to_upstream`, `_upstream_to_client`, `_ws_http_fallback` + - `headroom/proxy/handlers/anthropic.py` — `handle_anthropic_messages` + - `headroom/proxy/server.py` — `HeadroomProxy.__init__`, `.startup`, health routes + - `headroom/proxy/memory_handler.py` — `MemoryHandler._ensure_initialized` + - `headroom/transforms/content_router.py` — `eager_load_compressors` + - `headroom/transforms/kompress_compressor.py` — `_load_kompress_onnx`, `_kompress_cache` +- Related PRs/issues: + - Upstream: `https://github.com/chopratejas/headroom/issues/172` + - Fork branch: `fix/responses-retries-keep-compression` (commit `0b11637`) +- External docs: none used for this plan.