diff --git a/.gitignore b/.gitignore index fb8cd89f0..4b843b290 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,4 @@ -# Private scripts (contain credentials) +# Private scripts (contain credentials). Allowlist checked-in helpers below. scripts/ !scripts/ scripts/* @@ -8,6 +8,10 @@ scripts/* !scripts/changelog-gen.py !scripts/verify-versions.py !scripts/tests/ +!scripts/README.md +!scripts/repro_codex_replay.py +!scripts/fixtures/ +!scripts/fixtures/*.json # Swift SDK (separate repo) swift/ diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 000000000..ad74e91f2 --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,77 @@ +# scripts/ + +Utility scripts bundled with the Headroom repo. Most are one-off operator +tools; a few are runnable as part of development workflows. + +## Reproducing the reconnect storm + +`repro_codex_replay.py` reproduces the multi-agent Codex reconnect/retry storm +against a local Headroom proxy (default `http://127.0.0.1:8787`), as described +in `wiki/plans/2026-04-17-codex-proxy-runtime-analysis.md` under "Latest +Correction". Use it to: + +- Regression-check that `/livez` stays responsive under a cold-start storm. +- Empirically tune the Unit 4 pre-upstream semaphore default + (`HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY`). +- Exercise the Codex WS lifecycle + Anthropic HTTP path simultaneously + without needing to replay captured production traffic. + +### Run + +```bash +# Default: 8 WS + 4 HTTP clients, 30s storm, p99 /livez must stay <= 500ms. +python scripts/repro_codex_replay.py + +# Tighter budget, shorter run: +python scripts/repro_codex_replay.py \ + --url http://127.0.0.1:8787 \ + --ws-clients 16 \ + --anthropic-clients 8 \ + --duration 60 \ + --livez-threshold-ms 100 + +# Dump the full summary as JSON for downstream tooling: +python scripts/repro_codex_replay.py --json +``` + +Exit code: + +- `0` — warmup succeeded (or was skipped), storm ran for the requested + duration, and `/livez` p99 stayed under `--livez-threshold-ms`. +- `1` — soft assertion failed, proxy unreachable, or unhandled exception. + Proxy-unreachable is detected and reported within ~5 seconds. + +### Fixtures + +The script loads two hand-crafted, fully synthetic JSON fixtures: + +- `scripts/fixtures/anthropic_replay_body.json` — shape of a large agent + reconnect replay `/v1/messages?beta=true` POST body. +- `scripts/fixtures/codex_response_create_frame.json` — first Codex WS frame + with the `{"type": "response.create", "response": {...}}` envelope. + +Override via `--ws-frame-fixture` / `--anthropic-body-fixture` if you have +captured traffic to replay instead. + +### Interpretation + +- `/livez p99` under threshold means the event loop is not starved during the + storm. If it rises with the semaphore unbounded + (`HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY=10000`) and drops back under + the default, Unit 4's backpressure is working. +- `Codex WS: opened` should equal `--ws-clients`. `response.completed` + typically stays low when upstream auth isn't configured locally — the goal + is handshake + relay wiring, not real upstream traffic. +- `Anthropic HTTP: ok_2xx + non_2xx + timed_out + errors` should roughly equal + `attempted`. Sustained non-zero `timed_out` during the storm is the failure + signal the plan targets. + +A smoke test at `tests/test_scripts/test_repro_codex_replay_smoke.py` +exercises the script against a mock FastAPI server on every PR. + +## Install scripts + +- `install.sh` — POSIX installer. +- `install.ps1` — Windows PowerShell installer. + +These are generated by the release pipeline; edit with care. diff --git a/scripts/fixtures/anthropic_replay_body.json b/scripts/fixtures/anthropic_replay_body.json new file mode 100644 index 000000000..455c7492e --- /dev/null +++ b/scripts/fixtures/anthropic_replay_body.json @@ -0,0 +1,133 @@ +{ + "model": "claude-opus-4-7", + "max_tokens": 4096, + "stream": true, + "system": "You are Claude Code, Anthropic's official CLI for disciplined software engineering. You are helping the operator triage a resilience bug in a Python/FastAPI LLM proxy that handles both Anthropic /v1/messages traffic and OpenAI Codex /v1/responses WebSocket traffic on the same process.\n\nWorking style: Read code with symbolic tools when available; prefer small, reviewable diffs; never bypass compression or add fast paths that skip transforms; keep /livez trivial and IO-free; treat the WebSocket relay lifecycle as load-bearing. When you are uncertain about upstream behavior, instrument first and mitigate second. Multi-agent reconnect storms arrive as a burst of /v1/messages?beta=true POSTs immediately after process restart; these are the failure mode under investigation.\n\nInvariants: compression stays on; /livez must respond under 100ms even during cold-start replay storms; memory-context lookups are wrapped in wait_for with a bounded timeout; upstream WS handshake has retry/open-timeout hardening; debug endpoints are loopback-only.", + "tools": [ + { + "name": "read_file", + "description": "Read a file from the local filesystem. Returns file contents as a string. Use when you need to inspect exact text of a source file before editing.", + "input_schema": { + "type": "object", + "properties": { + "path": {"type": "string", "description": "Absolute file path."}, + "start_line": {"type": "integer", "description": "Optional 1-indexed start line."}, + "end_line": {"type": "integer", "description": "Optional inclusive end line."} + }, + "required": ["path"] + } + }, + { + "name": "search_code", + "description": "Search the codebase for a regex pattern. Returns matching file paths and surrounding lines. Prefer over read_file when you don't know the file location yet.", + "input_schema": { + "type": "object", + "properties": { + "pattern": {"type": "string"}, + "glob": {"type": "string", "description": "Optional glob filter, e.g. '**/*.py'."}, + "max_results": {"type": "integer", "default": 50} + }, + "required": ["pattern"] + } + }, + { + "name": "run_tests", + "description": "Run the pytest suite or a subset. Returns summary output including passed/failed/errored counts and tracebacks for failures.", + "input_schema": { + "type": "object", + "properties": { + "expression": {"type": "string", "description": "pytest -k expression."}, + "paths": {"type": "array", "items": {"type": "string"}, "description": "Specific test files or directories."} + } + } + }, + { + "name": "edit_file", + "description": "Apply an exact-string replacement to a file. The old_string must be unique within the file; otherwise use replace_all. Prefer over write_file for modifications.", + "input_schema": { + "type": "object", + "properties": { + "path": {"type": "string"}, + "old_string": {"type": "string"}, + "new_string": {"type": "string"}, + "replace_all": {"type": "boolean", "default": false} + }, + "required": ["path", "old_string", "new_string"] + } + }, + { + "name": "run_command", + "description": "Execute a shell command and return stdout, stderr, and exit code. Use for git, build, and verification commands.", + "input_schema": { + "type": "object", + "properties": { + "command": {"type": "string"}, + "timeout_seconds": {"type": "integer", "default": 120} + }, + "required": ["command"] + } + } + ], + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "The shared proxy on 127.0.0.1:8787 started acting like a zombie this morning after a restart. New /livez probes time out, new /v1/responses WebSocket handshakes hang for ~20 seconds before the client gives up, but the existing agent sessions that were already streaming when the restart happened keep working fine. Two of us were connected at the time so there was a reconnect burst of Anthropic /v1/messages?beta=true POSTs carrying multi-turn history (~30 KB each). I suspect the replay storm is interacting badly with cold-start ONNX/embedder loads. Start by reading headroom/proxy/handlers/anthropic.py:handle_anthropic_messages and headroom/proxy/server.py:startup — I want to know whether the Anthropic pre-upstream stages (deep_copy, first compression pass, memory-context lookup) run before or after the Kompress/Magika preload settles, and whether there is any concurrency limit in front of that path. Don't propose a fix yet; just map the stages and tell me where the event loop would plausibly starve." + } + ] + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "I've read both files. Mapping the pre-upstream stages on the Anthropic path: (1) handle_anthropic_messages reads the raw request bytes, (2) _read_request_json decodes and json.loads — this is synchronous and runs on the event loop thread, (3) a deep copy of the body is taken for logging, (4) the first compression stage fires through anthropic_pipeline.transforms, which on cold start is the moment Kompress ONNX and Magika get lazy-loaded, (5) memory-context lookup runs via MemoryHandler.search_and_format_context inside asyncio.wait_for, (6) only then does the handler reach http_client.stream() for upstream.\n\nStages 3 and 4 are the choke: deep_copy is synchronous Python on 30KB bodies (fine individually, but N concurrent replays mean N deep copies on one event loop), and the first compression stage blocks on Kompress load when the eager preload hasn't finished yet. There is currently no asyncio.Semaphore in front of stages 1-5. /livez does not go through this path, but the event loop is shared — if stages 3-5 hog it long enough, /livez's response coroutine doesn't get scheduled until they yield.\n\nI have not instrumented timings yet. Do you want me to add a StageTimer pass first, or jump straight to proposing a pre-upstream semaphore?" + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Instrument first — that's the whole point of not guessing. Add per-stage timings for read_request_json, deep_copy, compression_first_stage, memory_context, upstream_connect, upstream_first_byte on the Anthropic HTTP path, and accept -> first_client_frame -> upstream_connect -> upstream_first_event -> total_session on the Codex WS path. Both should share a single StageTimer utility, log one structured line per request with request_id + session_id + all stage fields, and emit Prometheus histograms. Don't change any behavior. Once that's in, we can let the harness run and read the real numbers before touching concurrency." + } + ] + }, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_anthropic_replay_01", + "name": "search_code", + "input": { + "pattern": "class PrometheusMetrics", + "glob": "headroom/proxy/**/*.py" + } + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_anthropic_replay_01", + "content": "headroom/proxy/prometheus_metrics.py:41:class PrometheusMetrics:\n Holds in-process counters/gauges/histograms plus async record_* helpers. Backend-agnostic; OTLP export is deferred. Pattern: register each histogram once in __init__, update via record_request / record_stage_timings." + } + ] + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Plan: create headroom/proxy/stage_timer.py with a StageTimer context manager exposing .measure(stage_name) as both a sync and async context manager, accumulating into a dict[str, float] keyed by stage with millisecond values. Extend PrometheusMetrics with one histogram per stage, keyed by path ('anthropic' | 'codex_ws'). Extend request_logger with a stage_timings field. Then plumb stage_timer through handle_anthropic_messages and handle_openai_responses_ws without changing any existing control flow — the timer is a pure observer. I'll land the util + unit tests first, then thread it through the two handlers in separate commits so the diff stays reviewable." + } + ] + } + ] +} diff --git a/scripts/fixtures/codex_response_create_frame.json b/scripts/fixtures/codex_response_create_frame.json new file mode 100644 index 000000000..ab41d8386 --- /dev/null +++ b/scripts/fixtures/codex_response_create_frame.json @@ -0,0 +1,52 @@ +{ + "type": "response.create", + "response": { + "model": "gpt-5.3-codex", + "stream": true, + "store": false, + "parallel_tool_calls": true, + "tools": [], + "input": [ + { + "type": "message", + "role": "system", + "content": [ + { + "type": "input_text", + "text": "You are Codex, a proxy-aware coding agent. The upstream connection is a WebSocket proxied through a local Headroom process. Keep responses incremental and stream-friendly; the proxy compresses large tool outputs on the way back and expects well-formed response.created / response.output_text.delta / response.completed envelopes." + } + ] + }, + { + "type": "message", + "role": "user", + "content": [ + { + "type": "input_text", + "text": "We're reproducing a reconnect storm against the local proxy at 127.0.0.1:8787. Open a WebSocket session, send one response.create frame, and stay connected. The goal is not to generate useful output — it's to make sure the proxy's WS lifecycle (accept, upstream connect, relay-task wiring, registry register/deregister) survives N concurrent sessions plus a parallel burst of large /v1/messages POSTs without /livez going dark. Do not expect a real model response; upstream auth will fail and the proxy should gracefully close. Just hold the slot for the storm duration." + } + ] + }, + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Acknowledged. I'll keep this session open, avoid generating extra frames, and let the harness measure /livez latency and registry drain time. If the upstream closes with an auth error I'll expect a clean response.completed or error frame from the proxy, not a dangling task." + } + ] + }, + { + "type": "message", + "role": "user", + "content": [ + { + "type": "input_text", + "text": "Confirm the session. No further work required — the harness is driving." + } + ] + } + ] + } +} diff --git a/scripts/repro_codex_replay.py b/scripts/repro_codex_replay.py new file mode 100755 index 000000000..1824b9c92 --- /dev/null +++ b/scripts/repro_codex_replay.py @@ -0,0 +1,636 @@ +#!/usr/bin/env python3 +"""Repro harness: multi-agent Codex reconnect storm against a local proxy. + +This script reproduces the failure class described in +``wiki/plans/2026-04-17-codex-proxy-runtime-analysis.md`` ("Latest Correction"): +a burst of concurrent Codex WebSocket sessions plus a parallel burst of large +Anthropic ``/v1/messages`` replays immediately after a fresh proxy restart, with +``/livez`` probed continuously to detect event-loop starvation. + +Usage:: + + python scripts/repro_codex_replay.py \\ + --url http://127.0.0.1:8787 \\ + --ws-clients 8 \\ + --anthropic-clients 4 \\ + --duration 30 + +Exit code is ``0`` iff the warmup phase succeeded (or was skipped), the storm +phase ran for the requested duration, and ``/livez`` p99 stayed at or below the +configured threshold (``--livez-threshold-ms``, default 500ms). + +The harness adds no new pip dependencies: it uses ``asyncio`` + ``websockets`` ++ ``httpx`` only, all already available in the Headroom dev environment. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import random +import statistics +import sys +import time +from contextlib import suppress +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any +from urllib.parse import urlparse, urlunparse + +import httpx + +try: + from websockets.asyncio.client import connect as ws_connect +except ImportError: # pragma: no cover - older websockets fallback + from websockets.client import connect as ws_connect # type: ignore[no-redef] + +from websockets.exceptions import ConnectionClosed, InvalidStatus, WebSocketException + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_WS_FRAME_FIXTURE = SCRIPT_DIR / "fixtures" / "codex_response_create_frame.json" +DEFAULT_ANTHROPIC_BODY_FIXTURE = SCRIPT_DIR / "fixtures" / "anthropic_replay_body.json" + + +def _load_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +# --------------------------------------------------------------------------- +# Stats +# --------------------------------------------------------------------------- + + +@dataclass +class LatencyHistogram: + """Minimal fixed-list histogram with p50/p95/p99/max computation.""" + + samples_ms: list[float] = field(default_factory=list) + + def record(self, value_ms: float) -> None: + self.samples_ms.append(value_ms) + + @property + def count(self) -> int: + return len(self.samples_ms) + + def percentile(self, p: float) -> float: + if not self.samples_ms: + return 0.0 + ordered = sorted(self.samples_ms) + if p <= 0: + return ordered[0] + if p >= 100: + return ordered[-1] + idx = max(0, min(len(ordered) - 1, int(round((p / 100.0) * (len(ordered) - 1))))) + return ordered[idx] + + def as_summary(self) -> dict[str, float]: + if not self.samples_ms: + return {"count": 0, "p50": 0.0, "p95": 0.0, "p99": 0.0, "max": 0.0} + return { + "count": self.count, + "p50": self.percentile(50), + "p95": self.percentile(95), + "p99": self.percentile(99), + "max": max(self.samples_ms), + } + + +@dataclass +class CodexWsStats: + opened: int = 0 + response_completed: int = 0 + errors: dict[str, int] = field(default_factory=dict) + + def record_error(self, kind: str) -> None: + self.errors[kind] = self.errors.get(kind, 0) + 1 + + +@dataclass +class AnthropicHttpStats: + attempted: int = 0 + ok_2xx: int = 0 + non_2xx: int = 0 + timed_out: int = 0 + errors: int = 0 + first_byte_latency_ms: list[float] = field(default_factory=list) + + @property + def avg_first_byte_ms(self) -> float: + if not self.first_byte_latency_ms: + return 0.0 + return statistics.mean(self.first_byte_latency_ms) + + +# --------------------------------------------------------------------------- +# URL helpers +# --------------------------------------------------------------------------- + + +def _http_to_ws_url(http_url: str, path: str) -> str: + parsed = urlparse(http_url) + scheme = "wss" if parsed.scheme == "https" else "ws" + # Preserve host:port; override scheme + path. Ensure path starts with a single "/". + normalized_path = "/" + path.lstrip("/") if path else "" + return urlunparse((scheme, parsed.netloc, normalized_path, "", "", "")) + + +# --------------------------------------------------------------------------- +# Warmup phase +# --------------------------------------------------------------------------- + + +async def warmup_probe( + url: str, + ws_frame: dict[str, Any], + timeout_s: float = 10.0, +) -> tuple[bool, float, str]: + """Open one WS, send the response.create frame, wait briefly. + + Returns ``(success, elapsed_ms, note)``. Success is true if the proxy + accepts the WebSocket handshake and the initial frame is sent without + error. We do not require ``response.completed`` — upstream auth will + often fail locally and the proxy will send an error frame; the important + signal is that the handshake + relay wiring worked. + """ + ws_url = _http_to_ws_url(url, "/v1/responses") + start = time.perf_counter() + handshake_ok = False + frame_sent = False + note = "unknown" + try: + async with asyncio.timeout(timeout_s): + async with ws_connect( + ws_url, + additional_headers={"Authorization": "Bearer repro-harness"}, + ) as ws: + handshake_ok = True + try: + await ws.send(json.dumps(ws_frame)) + frame_sent = True + except (ConnectionClosed, WebSocketException) as exc: + note = f"send_failed:{type(exc).__name__}" + return False, (time.perf_counter() - start) * 1000.0, note + # Drain until terminal event, close, or deadline. + deadline_local = time.perf_counter() + timeout_s + while time.perf_counter() < deadline_local: + remaining = deadline_local - time.perf_counter() + if remaining <= 0: + break + try: + raw = await asyncio.wait_for(ws.recv(), timeout=remaining) + except asyncio.TimeoutError: + note = "timeout_waiting_for_completion" + break + except ConnectionClosed: + # Upstream (or mock) closed cleanly after accepting our + # frame — this still counts as a working pre-upstream + # path; the warmup's purpose is handshake + send, not + # demanding real upstream output. + note = "upstream_closed_after_send" + break + try: + evt = json.loads(raw) if isinstance(raw, str) else {} + except json.JSONDecodeError: + continue + etype = evt.get("type", "") + if etype == "response.completed": + note = "response.completed" + break + if etype in {"error", "response.failed"}: + note = f"terminal:{etype}" + break + elapsed_ms = (time.perf_counter() - start) * 1000.0 + return frame_sent, elapsed_ms, note + except (OSError, WebSocketException, InvalidStatus, asyncio.TimeoutError) as exc: + elapsed_ms = (time.perf_counter() - start) * 1000.0 + # If the handshake completed and we sent the frame, a subsequent close + # error from the context manager is not a warmup failure. + if frame_sent: + return True, elapsed_ms, f"post_send_close:{type(exc).__name__}" + if handshake_ok: + return False, elapsed_ms, f"post_handshake_error:{type(exc).__name__}: {exc}" + return False, elapsed_ms, f"{type(exc).__name__}: {exc}" + + +# --------------------------------------------------------------------------- +# Storm phase +# --------------------------------------------------------------------------- + + +async def _ws_client( + idx: int, + url: str, + ws_frame: dict[str, Any], + deadline: float, + stats: CodexWsStats, +) -> None: + ws_url = _http_to_ws_url(url, "/v1/responses") + try: + async with ws_connect( + ws_url, + additional_headers={"Authorization": f"Bearer repro-harness-ws-{idx}"}, + open_timeout=10, + ) as ws: + stats.opened += 1 + try: + await ws.send(json.dumps(ws_frame)) + except (ConnectionClosed, WebSocketException) as exc: + stats.record_error(f"send:{type(exc).__name__}") + return + # Hold the session open, draining events until deadline or close. + while True: + remaining = deadline - time.perf_counter() + if remaining <= 0: + break + try: + raw = await asyncio.wait_for(ws.recv(), timeout=min(remaining, 5.0)) + except asyncio.TimeoutError: + continue + except ConnectionClosed: + break + try: + evt = json.loads(raw) if isinstance(raw, str) else {} + except json.JSONDecodeError: + continue + etype = evt.get("type", "") + if etype == "response.completed": + stats.response_completed += 1 + break + if etype in {"error", "response.failed"}: + stats.record_error(f"upstream:{etype}") + break + except (OSError, InvalidStatus) as exc: + stats.record_error(f"connect:{type(exc).__name__}") + except WebSocketException as exc: + stats.record_error(f"ws:{type(exc).__name__}") + except Exception as exc: # pragma: no cover - defensive + stats.record_error(f"unexpected:{type(exc).__name__}") + + +async def _anthropic_client( + idx: int, + url: str, + body: dict[str, Any], + deadline: float, + stats: AnthropicHttpStats, +) -> None: + endpoint = url.rstrip("/") + "/v1/messages?beta=true" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer repro-harness-anthropic-{idx}", + "anthropic-version": "2023-06-01", + "x-api-key": f"repro-harness-anthropic-{idx}", + } + retry_cutoff = min(deadline, time.perf_counter() + 60.0) + attempt = 0 + async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, connect=5.0)) as client: + while time.perf_counter() < deadline: + attempt += 1 + stats.attempted += 1 + start = time.perf_counter() + try: + resp = await client.post(endpoint, headers=headers, json=body) + first_byte_ms = (time.perf_counter() - start) * 1000.0 + stats.first_byte_latency_ms.append(first_byte_ms) + # Drain body so the connection is released. + with suppress(Exception): + _ = resp.content + if 200 <= resp.status_code < 300: + stats.ok_2xx += 1 + return + stats.non_2xx += 1 + if resp.status_code < 500: + # Client error — no retry, agent would surface this. + return + except httpx.TimeoutException: + stats.timed_out += 1 + except (httpx.HTTPError, OSError): + stats.errors += 1 + # Retry loop — mimic agent behavior with bounded wall clock. + if time.perf_counter() >= retry_cutoff: + return + # Small jitter between 50ms and 250ms. + await asyncio.sleep(0.05 + random.random() * 0.2) + + +async def _livez_prober( + url: str, + histogram: LatencyHistogram, + deadline: float, + interval_ms: int = 250, +) -> None: + endpoint = url.rstrip("/") + "/livez" + async with httpx.AsyncClient(timeout=httpx.Timeout(5.0, connect=2.0)) as client: + while time.perf_counter() < deadline: + start = time.perf_counter() + try: + resp = await client.get(endpoint) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + if resp.status_code == 200: + histogram.record(elapsed_ms) + else: + # Treat non-200 as worst-case for our threshold. + histogram.record(max(elapsed_ms, 5000.0)) + except (httpx.HTTPError, OSError): + histogram.record(5000.0) + # Tick every interval_ms regardless of probe latency. + sleep_for = interval_ms / 1000.0 + next_wake = time.perf_counter() + sleep_for + remaining = max(0.0, next_wake - time.perf_counter()) + if remaining > 0: + await asyncio.sleep(remaining) + + +# --------------------------------------------------------------------------- +# Orchestration +# --------------------------------------------------------------------------- + + +async def _check_reachable(url: str, timeout_s: float = 5.0) -> tuple[bool, str]: + endpoint = url.rstrip("/") + "/livez" + try: + async with httpx.AsyncClient(timeout=httpx.Timeout(timeout_s, connect=timeout_s)) as client: + resp = await client.get(endpoint) + return (True, f"HTTP {resp.status_code}") + except (httpx.ConnectError, ConnectionRefusedError, OSError) as exc: + return (False, f"{type(exc).__name__}: {exc}") + except httpx.HTTPError as exc: + return (True, f"reachable-but-error: {type(exc).__name__}: {exc}") + + +async def run_harness(args: argparse.Namespace) -> dict[str, Any]: + ws_frame = _load_json(Path(args.ws_frame_fixture)) + anthropic_body = _load_json(Path(args.anthropic_body_fixture)) + + # Reachability gate — quick, clear failure if the proxy is not up. + reachable, reach_note = await _check_reachable(args.url, timeout_s=5.0) + if not reachable: + return { + "ok": False, + "reason": "proxy_unreachable", + "detail": reach_note, + "url": args.url, + } + + # Phase 1: warmup + warmup_result: dict[str, Any] + if args.no_warmup: + warmup_result = {"skipped": True} + else: + success, elapsed_ms, note = await warmup_probe( + args.url, ws_frame, timeout_s=args.warmup_timeout + ) + warmup_result = { + "skipped": False, + "success": success, + "elapsed_ms": round(elapsed_ms, 2), + "note": note, + } + + # Phase 2: storm + livez_hist = LatencyHistogram() + ws_stats = CodexWsStats() + http_stats = AnthropicHttpStats() + + storm_start = time.perf_counter() + deadline = storm_start + args.duration + + livez_task = asyncio.create_task( + _livez_prober(args.url, livez_hist, deadline, interval_ms=args.livez_interval_ms), + name="repro-livez-prober", + ) + ws_tasks = [ + asyncio.create_task( + _ws_client(i, args.url, ws_frame, deadline, ws_stats), + name=f"repro-ws-{i}", + ) + for i in range(args.ws_clients) + ] + http_tasks = [ + asyncio.create_task( + _anthropic_client(i, args.url, anthropic_body, deadline, http_stats), + name=f"repro-http-{i}", + ) + for i in range(args.anthropic_clients) + ] + + # Let the storm run. Gather storm tasks but keep probing /livez for the + # full requested duration regardless of when the clients exit — the whole + # point is to observe event-loop health across the window. + storm_tasks = ws_tasks + http_tasks + try: + await asyncio.wait_for( + asyncio.gather(*storm_tasks, return_exceptions=True), + timeout=args.duration + 10.0, + ) + except asyncio.TimeoutError: + for t in storm_tasks: + if not t.done(): + t.cancel() + await asyncio.gather(*storm_tasks, return_exceptions=True) + + # Let the /livez prober run to its own deadline so the histogram covers + # the full window, not just the interval where storm tasks were alive. + remaining = deadline - time.perf_counter() + if remaining > 0: + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(asyncio.shield(livez_task), timeout=remaining + 2.0) + if not livez_task.done(): + livez_task.cancel() + with suppress(asyncio.CancelledError, Exception): + await livez_task + + storm_duration_s = time.perf_counter() - storm_start + livez_summary = livez_hist.as_summary() + + # Soft assertion: livez p99 under threshold. + threshold_ms = args.livez_threshold_ms + livez_ok = livez_summary["p99"] <= threshold_ms + + return { + "ok": (warmup_result.get("skipped", False) or warmup_result.get("success", False)) + and livez_ok, + "warmup": warmup_result, + "storm": { + "ws_clients": args.ws_clients, + "anthropic_clients": args.anthropic_clients, + "requested_duration_s": args.duration, + "actual_duration_s": round(storm_duration_s, 3), + }, + "livez": { + **livez_summary, + "threshold_ms": threshold_ms, + "threshold_ok": livez_ok, + }, + "codex_ws": { + "opened": ws_stats.opened, + "response_completed": ws_stats.response_completed, + "errors": dict(ws_stats.errors), + }, + "anthropic_http": { + "attempted": http_stats.attempted, + "ok_2xx": http_stats.ok_2xx, + "non_2xx": http_stats.non_2xx, + "timed_out": http_stats.timed_out, + "errors": http_stats.errors, + "avg_first_byte_ms": round(http_stats.avg_first_byte_ms, 2), + }, + } + + +# --------------------------------------------------------------------------- +# Printing +# --------------------------------------------------------------------------- + + +def format_summary(result: dict[str, Any]) -> str: + if result.get("reason") == "proxy_unreachable": + return ( + "Proxy unreachable at {url}.\n" + " {detail}\n" + " Hint: start the proxy with `headroom proxy` and retry." + ).format(**result) + + lines: list[str] = [] + lines.append("=" * 72) + lines.append("Codex proxy reconnect-storm repro harness — summary") + lines.append("=" * 72) + warm = result.get("warmup", {}) + if warm.get("skipped"): + lines.append("Warmup: skipped") + else: + lines.append( + "Warmup: success={success} elapsed_ms={elapsed_ms} note={note}".format( + **warm + ) + ) + storm = result.get("storm", {}) + lines.append( + "Storm: ws_clients={ws_clients} anthropic_clients={anthropic_clients} " + "requested={requested_duration_s}s actual={actual_duration_s}s".format(**storm) + ) + livez = result.get("livez", {}) + lines.append( + "/livez: count={count} p50={p50:.2f}ms p95={p95:.2f}ms " + "p99={p99:.2f}ms max={max:.2f}ms (threshold={threshold_ms}ms, " + "ok={threshold_ok})".format(**livez) + ) + ws = result.get("codex_ws", {}) + lines.append( + "Codex WS: opened={opened} response.completed={response_completed} " + "errors={errors}".format(**ws) + ) + http = result.get("anthropic_http", {}) + lines.append( + "Anthropic HTTP: attempted={attempted} ok_2xx={ok_2xx} non_2xx={non_2xx} " + "timed_out={timed_out} errors={errors} avg_first_byte_ms={avg_first_byte_ms}".format(**http) + ) + lines.append("=" * 72) + lines.append("RESULT: {}".format("OK" if result.get("ok") else "FAIL")) + lines.append("=" * 72) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Reproducibly exercise the multi-agent Codex reconnect/retry " + "storm against a local Headroom proxy.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--url", + default="http://127.0.0.1:8787", + help="Base URL of the running proxy.", + ) + parser.add_argument( + "--ws-clients", + type=int, + default=8, + help="Number of concurrent Codex WS connections to open during the storm.", + ) + parser.add_argument( + "--anthropic-clients", + type=int, + default=4, + help="Number of concurrent Anthropic /v1/messages POST clients.", + ) + parser.add_argument( + "--duration", + type=float, + default=30.0, + help="Total storm phase length, in seconds.", + ) + parser.add_argument( + "--livez-threshold-ms", + type=float, + default=500.0, + help="Soft assertion threshold for /livez p99 (ms). Exit non-zero if exceeded.", + ) + parser.add_argument( + "--livez-interval-ms", + type=int, + default=250, + help="Interval between /livez probes, in milliseconds.", + ) + parser.add_argument( + "--warmup-timeout", + type=float, + default=10.0, + help="Max seconds to wait for the warmup WS session to complete.", + ) + parser.add_argument( + "--no-warmup", + action="store_true", + help="Skip the warmup probe phase.", + ) + parser.add_argument( + "--ws-frame-fixture", + default=str(DEFAULT_WS_FRAME_FIXTURE), + help="Path to the Codex response.create frame JSON fixture.", + ) + parser.add_argument( + "--anthropic-body-fixture", + default=str(DEFAULT_ANTHROPIC_BODY_FIXTURE), + help="Path to the Anthropic /v1/messages request body JSON fixture.", + ) + parser.add_argument( + "--json", + action="store_true", + help="Also print the full summary as JSON on stdout (after human summary).", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + try: + result = asyncio.run(run_harness(args)) + except KeyboardInterrupt: + print("\nInterrupted by user.", file=sys.stderr) + return 130 + except Exception as exc: # noqa: BLE001 - top-level guard + print(f"Harness crashed: {type(exc).__name__}: {exc}", file=sys.stderr) + return 1 + + print(format_summary(result)) + if args.json: + print(json.dumps(result, indent=2, sort_keys=True)) + + return 0 if result.get("ok") else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_scripts/__init__.py b/tests/test_scripts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_scripts/test_repro_codex_replay_smoke.py b/tests/test_scripts/test_repro_codex_replay_smoke.py new file mode 100644 index 000000000..2e8f85ffb --- /dev/null +++ b/tests/test_scripts/test_repro_codex_replay_smoke.py @@ -0,0 +1,228 @@ +"""Smoke test for scripts/repro_codex_replay.py. + +Spins up a minimal FastAPI + websockets mock proxy that answers ``/livez``, +``/v1/messages``, and ``/v1/responses`` (WS), then invokes the harness' +``main()`` in-process and verifies it exits 0 with the expected summary shape. + +The mock does *not* implement Codex semantics — it just accepts the WS +handshake, consumes the first frame, and closes. That's enough to exercise +the harness end-to-end in < 10s. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import io +import socket +import sys +import threading +from collections.abc import Iterator +from pathlib import Path + +import pytest +import uvicorn +from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect + +# Make sure `scripts/` is importable when running via pytest from repo root. +ROOT = Path(__file__).resolve().parents[2] +SCRIPTS_DIR = ROOT / "scripts" +if str(SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPTS_DIR)) + +import repro_codex_replay # type: ignore[import-not-found] # noqa: E402 + +# --------------------------------------------------------------------------- +# Mock proxy server +# --------------------------------------------------------------------------- + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _build_app() -> FastAPI: + app = FastAPI() + + @app.get("/livez") + async def livez() -> dict[str, str]: + return {"status": "ok"} + + @app.post("/v1/messages") + async def messages(request: Request) -> dict[str, object]: + # Drain body so the client gets a clean completion. + _ = await request.body() + return { + "id": "msg_mock_001", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "ok"}], + "model": "claude-mock", + "stop_reason": "end_turn", + } + + @app.websocket("/v1/responses") + async def responses_ws(websocket: WebSocket) -> None: + await websocket.accept() + try: + # Accept one frame and then close. Harness treats this as + # "handshake worked" — no need to emulate Codex events. + with contextlib.suppress(WebSocketDisconnect): + await asyncio.wait_for(websocket.receive_text(), timeout=2.0) + except asyncio.TimeoutError: + pass + with contextlib.suppress(Exception): + await websocket.close() + + return app + + +class _ServerThread: + """Run uvicorn in a background thread bound to a dedicated port.""" + + def __init__(self, app: FastAPI, port: int) -> None: + config = uvicorn.Config( + app, + host="127.0.0.1", + port=port, + log_level="warning", + loop="asyncio", + lifespan="on", + ) + self.server = uvicorn.Server(config) + self.port = port + self.thread = threading.Thread(target=self.server.run, name="mock-proxy", daemon=True) + + def start(self) -> None: + self.thread.start() + # Wait for the socket to accept connections. + deadline = 5.0 + import time + + t0 = time.perf_counter() + while time.perf_counter() - t0 < deadline: + if self.server.started: + return + time.sleep(0.05) + raise RuntimeError("mock proxy failed to start within 5s") + + def stop(self) -> None: + self.server.should_exit = True + self.thread.join(timeout=5.0) + + +@pytest.fixture +def mock_proxy() -> Iterator[str]: + port = _free_port() + srv = _ServerThread(_build_app(), port) + srv.start() + try: + yield f"http://127.0.0.1:{port}" + finally: + srv.stop() + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def _run_harness(argv: list[str]) -> tuple[int, str, str]: + stdout = io.StringIO() + stderr = io.StringIO() + old_out, old_err = sys.stdout, sys.stderr + sys.stdout, sys.stderr = stdout, stderr + try: + rc = repro_codex_replay.main(argv) + finally: + sys.stdout, sys.stderr = old_out, old_err + return rc, stdout.getvalue(), stderr.getvalue() + + +def test_harness_runs_against_mock_proxy_and_exits_zero(mock_proxy: str) -> None: + rc, out, err = _run_harness( + [ + "--url", + mock_proxy, + "--ws-clients", + "2", + "--anthropic-clients", + "2", + "--duration", + "2", + "--warmup-timeout", + "3", + "--livez-threshold-ms", + "2000", + "--livez-interval-ms", + "100", + "--json", + ] + ) + combined = out + err + assert rc == 0, f"expected exit 0, got {rc}\nSTDOUT:\n{out}\nSTDERR:\n{err}" + # Human summary shape. + for needle in ( + "Codex proxy reconnect-storm repro harness", + "Warmup:", + "Storm:", + "/livez:", + "Codex WS:", + "Anthropic HTTP:", + "RESULT: OK", + ): + assert needle in combined, f"missing '{needle}' in output:\n{combined}" + # JSON payload shape — --json appends one indent=2 JSON object at the end. + import json as _json + + # Find the last top-level JSON object in stdout. With indent=2 the closing + # brace is on a line by itself (no leading whitespace), so we search for + # the final "\n}" boundary, then walk back to the matching "{". + end_brace = out.rfind("\n}") + assert end_brace != -1, f"no JSON payload found in stdout:\n{out}" + # The matching opening brace must be preceded by a newline and start a line. + # Scan backwards for a line that is exactly "{". + lines = out[: end_brace + 2].splitlines() + # Find the last line equal to "{" — that's the start of the JSON object. + start_line_idx = None + for i in range(len(lines) - 1, -1, -1): + if lines[i] == "{": + start_line_idx = i + break + assert start_line_idx is not None, f"could not locate JSON start in stdout:\n{out}" + payload_text = "\n".join(lines[start_line_idx:]) + payload = _json.loads(payload_text) + for key in ("ok", "warmup", "storm", "livez", "codex_ws", "anthropic_http"): + assert key in payload, f"summary missing key {key!r}: {payload}" + assert payload["ok"] is True + assert payload["storm"]["ws_clients"] == 2 + assert payload["storm"]["anthropic_clients"] == 2 + assert payload["livez"]["count"] > 0 + # The mock accepts /v1/messages with 200 — at least one should succeed. + assert payload["anthropic_http"]["ok_2xx"] >= 1 + # Every WS client should have opened (handshake works against the mock). + assert payload["codex_ws"]["opened"] == 2 + + +def test_harness_reports_unreachable_proxy_fast() -> None: + import time + + t0 = time.perf_counter() + rc, out, err = _run_harness( + [ + "--url", + "http://127.0.0.1:1", # deliberately closed port + "--ws-clients", + "1", + "--anthropic-clients", + "1", + "--duration", + "1", + ] + ) + elapsed = time.perf_counter() - t0 + assert rc == 1, f"expected exit 1, got {rc}. stdout={out} stderr={err}" + assert "unreachable" in (out + err).lower() + assert elapsed < 10.0, f"unreachable detection too slow: {elapsed:.2f}s"