fix(ci): restore repro harness test in dev installs

Add websockets to the dev extra so the repro harness smoke test can import
its websocket client dependency in the CI test matrix. Also apply ruff
formatting to the files the formatter check was rejecting so the 3.12 lint
job passes.

headroom/proxy/debug_introspection.py

@@ -103,7 +103,7 @@ def _age_for_named_task(
         return None
     for prefix in _CODEX_WS_RELAY_PREFIXES:
         if task_name.startswith(prefix):
-            session_id = task_name[len(prefix):]
+            session_id = task_name[len(prefix) :]
             handle = ws_registry.get(session_id)
             if handle is None:
                 return None

headroom/proxy/handlers/anthropic.py

@@ -481,8 +481,12 @@ class AnthropicHandlerMixin:
     
             # Subscription tracker: notify on OAuth requests (not API-key requests)
             _auth_header = headers.get("authorization", "")
-            if _auth_header.startswith("Bearer ") and not _auth_header.startswith("Bearer sk-ant-api"):
-                from headroom.subscription.tracker import get_subscription_tracker as _get_sub_tracker
+            if _auth_header.startswith("Bearer ") and not _auth_header.startswith(
+                "Bearer sk-ant-api"
+            ):
+                from headroom.subscription.tracker import (
+                    get_subscription_tracker as _get_sub_tracker,
+                )
     
                 _sub_tracker = _get_sub_tracker()
                 if _sub_tracker is not None:
@@ -1285,7 +1289,9 @@ class AnthropicHandlerMixin:
                         and response.status_code == 200
                         and self.ccr_response_handler.has_ccr_tool_calls(resp_json, "anthropic")
                     ):
-                        logger.info(f"[{request_id}] CCR: Detected retrieval tool call, handling...")
+                        logger.info(
+                            f"[{request_id}] CCR: Detected retrieval tool call, handling..."
+                        )
     
                         # Create API call function for continuation
                         # Use a fresh client to avoid potential decompression state issues
@@ -1387,7 +1393,9 @@ class AnthropicHandlerMixin:
                         and response.status_code == 200
                         and self.memory_handler.has_memory_tool_calls(resp_json, "anthropic")
                     ):
-                        logger.info(f"[{request_id}] Memory: Detected memory tool call, handling...")
+                        logger.info(
+                            f"[{request_id}] Memory: Detected memory tool call, handling..."
+                        )
     
                         try:
                             # Execute memory tool calls
@@ -1406,7 +1414,10 @@ class AnthropicHandlerMixin:
                                     "content": tool_results,
                                 }
     
-                                continuation_messages = optimized_messages + [assistant_msg, user_msg]
+                                continuation_messages = optimized_messages + [
+                                    assistant_msg,
+                                    user_msg,
+                                ]
     
                                 # Make continuation API call
                                 continuation_body = {**body, "messages": continuation_messages}
@@ -1580,7 +1591,9 @@ class AnthropicHandlerMixin:
                     # Remove compression headers since httpx already decompressed the response
                     response_headers = dict(response.headers)
                     response_headers.pop("content-encoding", None)
-                    response_headers.pop("content-length", None)  # Length changed after decompression
+                    response_headers.pop(
+                        "content-length", None
+                    )  # Length changed after decompression
     
                     # Inject Headroom compression metrics (for SaaS metering)
                     response_headers["x-headroom-tokens-before"] = str(original_tokens)
@@ -1609,7 +1622,9 @@ class AnthropicHandlerMixin:
                                 headers=response_headers,
                             )
                         except Exception as sec_err:
-                            logger.warning(f"[{request_id}] Security response scan error: {sec_err}")
+                            logger.warning(
+                                f"[{request_id}] Security response scan error: {sec_err}"
+                            )
     
                     return Response(
                         content=response.content,
@@ -1658,7 +1673,6 @@ class AnthropicHandlerMixin:
             # deep-copy) would otherwise leak the pre-upstream semaphore
             # permanently. The emit function is idempotent.
             await _finalize_pre_upstream()
-
     async def handle_anthropic_batch_create(
         self,
         request: Request,

headroom/proxy/handlers/openai.py

@@ -1264,9 +1264,7 @@ class OpenAIHandlerMixin:
         # 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
-        )
+        ws_sessions: WebSocketSessionRegistry | None = getattr(self, "ws_sessions", None)
         session_handle: WSSessionHandle | None = None
         termination_cause: TerminationCause = "unknown"
 
@@ -1414,7 +1412,9 @@ class OpenAIHandlerMixin:
                 body = json.loads(first_msg_raw)
                 tokens_saved = 0
                 ws_request_body = body.get("response", body)
-                input_data = ws_request_body.get("input") if isinstance(ws_request_body, dict) else None
+                input_data = (
+                    ws_request_body.get("input") if isinstance(ws_request_body, dict) else None
+                )
 
                 should_compress = (
                     self.config.optimize
@@ -1919,9 +1919,7 @@ class OpenAIHandlerMixin:
                                 metrics_for_tasks, "inc_active_relay_tasks"
                             ):
                                 try:
-                                    metrics_for_tasks.inc_active_relay_tasks(
-                                        len(relay_tasks)
-                                    )
+                                    metrics_for_tasks.inc_active_relay_tasks(len(relay_tasks))
                                 except Exception:  # pragma: no cover - defensive
                                     pass
 
@@ -1939,9 +1937,7 @@ class OpenAIHandlerMixin:
                                 t.cancel()
                             if pending:
                                 with contextlib.suppress(asyncio.CancelledError):
-                                    await asyncio.gather(
-                                        *pending, return_exceptions=True
-                                    )
+                                    await asyncio.gather(*pending, return_exceptions=True)
 
                             # Classify termination cause from whichever
                             # task completed first. ``CancelledError``
@@ -1998,8 +1994,7 @@ class OpenAIHandlerMixin:
                                     else:
                                         termination_cause = "upstream_error"
                                         logger.debug(
-                                            f"[{request_id}] WS relay {task_name} "
-                                            f"raised: {exc!r}"
+                                            f"[{request_id}] WS relay {task_name} raised: {exc!r}"
                                         )
                         finally:
                             # In case anything above raised before the
@@ -2008,9 +2003,7 @@ class OpenAIHandlerMixin:
                                 if not t.done():
                                     t.cancel()
                             with contextlib.suppress(asyncio.CancelledError):
-                                await asyncio.gather(
-                                    *relay_tasks, return_exceptions=True
-                                )
+                                await asyncio.gather(*relay_tasks, return_exceptions=True)
 
                         logger.info(
                             f"[{request_id}] WS /v1/responses completed "
@@ -2113,17 +2106,13 @@ class OpenAIHandlerMixin:
                 _deregistered, released_tasks = ws_sessions.deregister_and_count(
                     session_id, cause=termination_cause
                 )
-                session_duration_ms = (
-                    time.perf_counter() - session_started_at
-                ) * 1000.0
+                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"
-                        ):
+                        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(

headroom/proxy/helpers.py

@@ -36,6 +36,7 @@ COMPRESSION_TIMEOUT_SECONDS = 30
 # Maximum compression cache sessions (prevents unbounded memory growth)
 MAX_COMPRESSION_CACHE_SESSIONS = 500
 
+
 def jitter_delay_ms(base_ms: int, max_ms: int, attempt: int) -> float:
     """Exponential backoff with 50-150% jitter.
 

headroom/proxy/memory_handler.py

@@ -207,10 +207,7 @@ class MemoryHandler:
             # cancellation is a signal, not an error to swallow.
             self._backend = None
             self._initialized = False
-            logger.info(
-                "Memory: backend initialization cancelled "
-                f"(backend={self.config.backend})"
-            )
+            logger.info(f"Memory: backend initialization cancelled (backend={self.config.backend})")
             raise
 
     async def _init_backend_locked(self) -> None:

headroom/proxy/server.py

@@ -397,8 +397,8 @@ class HeadroomProxy(
             _pre_upstream_resolved = _pre_upstream_cfg
         self.anthropic_pre_upstream_concurrency: int = _pre_upstream_resolved
         if _pre_upstream_resolved > 0:
-            self.anthropic_pre_upstream_sem: asyncio.Semaphore | None = (
-                asyncio.Semaphore(_pre_upstream_resolved)
+            self.anthropic_pre_upstream_sem: asyncio.Semaphore | None = asyncio.Semaphore(
+                _pre_upstream_resolved
             )
         else:
             self.anthropic_pre_upstream_sem = None
@@ -659,10 +659,7 @@ class HeadroomProxy(
         # value (auto-detected vs. explicit) so operators can correlate
         # ``pre_upstream_wait_ms`` log lines with the configured cap.
         if self.anthropic_pre_upstream_sem is None:
-            logger.info(
-                "Anthropic pre-upstream concurrency: unbounded "
-                "(explicitly disabled)"
-            )
+            logger.info("Anthropic pre-upstream concurrency: unbounded (explicitly disabled)")
         else:
             _explicit = self.config.anthropic_pre_upstream_concurrency
             _origin = "auto-detected" if _explicit is None else "explicit"
@@ -750,9 +747,7 @@ class HeadroomProxy(
                 await self.memory_handler.ensure_initialized()
             except Exception as exc:  # pragma: no cover - defensive
                 self.warmup.memory_backend.mark_error(str(exc))
-                logger.warning(
-                    "Memory: backend initialization failed (startup continues): %s", exc
-                )
+                logger.warning("Memory: backend initialization failed (startup continues): %s", exc)
             memory_status = self.memory_handler.health_status()
             if memory_status.get("initialized"):
                 self.warmup.memory_backend.mark_loaded(

headroom/proxy/stage_timer.py

@@ -180,14 +180,11 @@ async def emit_stage_timings_log(
         logger.info(f"[{request_id}] STAGE_TIMINGS {payload}")
     except (TypeError, ValueError):
         logger.info(
-            f"[{request_id}] STAGE_TIMINGS path={path} session_id={session_id} "
-            f"stages={padded!r}"
+            f"[{request_id}] STAGE_TIMINGS path={path} session_id={session_id} stages={padded!r}"
         )
 
     if metrics is not None and hasattr(metrics, "record_stage_timings"):
         try:
             await metrics.record_stage_timings(path, summary)
         except Exception as metric_err:  # pragma: no cover - defensive
-            logger.debug(
-                f"[{request_id}] record_stage_timings failed for {path}: {metric_err}"
-            )
+            logger.debug(f"[{request_id}] record_stage_timings failed for {path}: {metric_err}")

pyproject.toml

@@ -187,6 +187,7 @@ dev = [
     "fastapi>=0.100.0",
     "uvicorn>=0.23.0",
     "httpx[http2]>=0.24.0",
+    "websockets>=13.0",
     "opentelemetry-sdk>=1.24.0",
     "opentelemetry-exporter-otlp-proto-http>=1.24.0",
     "ollama>=0.4.0",

scripts/repro_codex_replay.py

@@ -322,9 +322,7 @@ async def _anthropic_client(
             # so this script stays dependency-free of the proxy package.
             _base_ms = 250
             _max_ms = 5000
-            _delay_ms = min(_base_ms * (2 ** (attempt - 1)), _max_ms) * (
-                0.5 + random.random()
-            )
+            _delay_ms = min(_base_ms * (2 ** (attempt - 1)), _max_ms) * (0.5 + random.random())
             await asyncio.sleep(_delay_ms / 1000.0)
 
 
@@ -515,9 +513,7 @@ def format_summary(result: dict[str, Any]) -> str:
         lines.append("Warmup:         skipped")
     else:
         lines.append(
-            "Warmup:         success={success} elapsed_ms={elapsed_ms} note={note}".format(
-                **warm
-            )
+            "Warmup:         success={success} elapsed_ms={elapsed_ms} note={note}".format(**warm)
         )
     storm = result.get("storm", {})
     lines.append(

tests/test_anthropic_pre_upstream_backpressure.py

@@ -317,9 +317,7 @@ def test_n_plus_one_contention_only_waiter_has_nonzero_wait(stage_log_capture):
         sem = asyncio.Semaphore(2)
         # Each request hogs the semaphore for ~150 ms. With concurrency=2,
         # 3 concurrent requests mean exactly one waits ~150 ms.
-        handler = _DummyAnthropicHandler(
-            anthropic_pre_upstream_sem=sem, upstream_delay_s=0.15
-        )
+        handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem, upstream_delay_s=0.15)
         reqs = [
             _build_request(
                 {
@@ -355,9 +353,7 @@ def test_n_plus_one_contention_only_waiter_has_nonzero_wait(stage_log_capture):
 def test_concurrency_one_serializes_requests():
     async def _run() -> float:
         sem = asyncio.Semaphore(1)
-        handler = _DummyAnthropicHandler(
-            anthropic_pre_upstream_sem=sem, upstream_delay_s=0.10
-        )
+        handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem, upstream_delay_s=0.10)
         reqs = [
             _build_request(
                 {
@@ -402,9 +398,7 @@ def test_unbounded_mode_requests_run_concurrently():
     """With concurrency=0 (sem disabled), two slow requests overlap."""
 
     async def _run() -> float:
-        handler = _DummyAnthropicHandler(
-            anthropic_pre_upstream_sem=None, upstream_delay_s=0.10
-        )
+        handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=None, upstream_delay_s=0.10)
         reqs = [
             _build_request(
                 {
@@ -435,9 +429,7 @@ def test_exception_inside_critical_section_releases_semaphore():
     async def _run() -> None:
         sem = asyncio.Semaphore(2)
         baseline = sem._value
-        handler = _DummyAnthropicHandler(
-            anthropic_pre_upstream_sem=sem, raise_during_critical=True
-        )
+        handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem, raise_during_critical=True)
         # Drive several cycles to ensure we don't leak on any path.
         for i in range(5):
             req = _build_request(
@@ -606,9 +598,7 @@ def test_env_var_sets_pre_upstream_concurrency():
 
 def test_cli_flag_overrides_env_var():
     env = {"HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY": "4"}
-    config = _run_cli_capture(
-        ["--anthropic-pre-upstream-concurrency", "7"], env=env
-    )
+    config = _run_cli_capture(["--anthropic-pre-upstream-concurrency", "7"], env=env)
     assert config.anthropic_pre_upstream_concurrency == 7
 
 
@@ -735,13 +725,10 @@ def test_early_exit_paths_release_semaphore_under_contention(scenario):
                 )
             else:
                 # security returns a JSONResponse; cache returns a Response.
-                assert raised is None, (
-                    f"{scenario}: unexpected exception {raised!r}"
-                )
+                assert raised is None, f"{scenario}: unexpected exception {raised!r}"
                 assert result is not None
             assert sem._value == original_value, (
-                f"{scenario}: semaphore leak "
-                f"got={sem._value}, want={original_value}"
+                f"{scenario}: semaphore leak got={sem._value}, want={original_value}"
             )
 
     with _tokenizer_patch():

tests/test_memory_handler_concurrent_init.py

@@ -53,9 +53,7 @@ async def test_concurrent_ensure_initialized_runs_init_once(tmp_path, monkeypatc
     monkeypatch.setattr(local_mod, "LocalBackend", FakeLocalBackend)
 
     handler = MemoryHandler(
-        MemoryConfig(
-            enabled=True, backend="local", db_path=str(tmp_path / "mem.db")
-        )
+        MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
     )
 
     async def caller() -> None:
@@ -88,9 +86,7 @@ async def test_ensure_initialized_noop_when_disabled(tmp_path):
 
 
 @pytest.mark.asyncio
-async def test_ensure_initialized_timeout_leaves_handler_unready(
-    tmp_path, monkeypatch
-):
+async def test_ensure_initialized_timeout_leaves_handler_unready(tmp_path, monkeypatch):
     class HangingBackend:
         def __init__(self, config):
             self.config = config
@@ -107,9 +103,7 @@ async def test_ensure_initialized_timeout_leaves_handler_unready(
     monkeypatch.setattr(local_mod, "LocalBackend", HangingBackend)
 
     handler = MemoryHandler(
-        MemoryConfig(
-            enabled=True, backend="local", db_path=str(tmp_path / "mem.db")
-        )
+        MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
     )
 
     # Attach a handler directly to the module logger — caplog has trouble
@@ -129,9 +123,7 @@ async def test_ensure_initialized_timeout_leaves_handler_unready(
     mem_logger.setLevel(_logging.DEBUG)
     try:
         # Shrink the module-level timeout to keep the test fast.
-        with patch(
-            "headroom.proxy.memory_handler.STARTUP_INIT_TIMEOUT_SECONDS", 0.1
-        ):
+        with patch("headroom.proxy.memory_handler.STARTUP_INIT_TIMEOUT_SECONDS", 0.1):
             await handler._ensure_initialized()
     finally:
         mem_logger.removeHandler(handler_log)
@@ -141,8 +133,7 @@ async def test_ensure_initialized_timeout_leaves_handler_unready(
     assert handler._initialized is False
     found_timeout_log = any("timed out" in rec.getMessage().lower() for rec in captured)
     assert found_timeout_log, (
-        "expected 'timed out' log record; "
-        f"got: {[(r.levelname, r.getMessage()) for r in captured]}"
+        f"expected 'timed out' log record; got: {[(r.levelname, r.getMessage()) for r in captured]}"
     )
 
     # Confirm the default constant is unchanged (sanity).
@@ -175,14 +166,10 @@ async def test_ensure_initialized_timeout_nulls_partially_initialized_backend(
     monkeypatch.setattr(local_mod, "LocalBackend", SlowBackend)
 
     handler = MemoryHandler(
-        MemoryConfig(
-            enabled=True, backend="local", db_path=str(tmp_path / "mem.db")
-        )
+        MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
     )
 
-    with patch(
-        "headroom.proxy.memory_handler.STARTUP_INIT_TIMEOUT_SECONDS", 0.01
-    ):
+    with patch("headroom.proxy.memory_handler.STARTUP_INIT_TIMEOUT_SECONDS", 0.01):
         await handler._ensure_initialized()
 
     # Both must be consistent after timeout.
@@ -191,9 +178,7 @@ async def test_ensure_initialized_timeout_nulls_partially_initialized_backend(
 
 
 @pytest.mark.asyncio
-async def test_ensure_initialized_cancellation_propagates_and_resets_state(
-    tmp_path, monkeypatch
-):
+async def test_ensure_initialized_cancellation_propagates_and_resets_state(tmp_path, monkeypatch):
     """External cancellation of an in-flight ``_ensure_initialized`` must
     propagate (CancelledError is BaseException — not a swallowable error)
     and leave the handler in a clean state."""
@@ -213,9 +198,7 @@ async def test_ensure_initialized_cancellation_propagates_and_resets_state(
     monkeypatch.setattr(local_mod, "LocalBackend", HangingBackend)
 
     handler = MemoryHandler(
-        MemoryConfig(
-            enabled=True, backend="local", db_path=str(tmp_path / "mem.db")
-        )
+        MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
     )
 
     task = asyncio.create_task(handler._ensure_initialized())

tests/test_openai_codex_ws_lifecycle.py

@@ -40,9 +40,7 @@ class _DummyMetrics:
 
     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
-        )
+        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)

tests/test_proxy_debug_endpoints.py

@@ -278,12 +278,9 @@ def test_debug_tasks_stack_depth_is_gated_behind_query(client):
     # itself runs under a task). Some entries may still be None if
     # get_stack raised defensively — we only require that opting in
     # produces at least one integer result.
-    integer_depths = [
-        e["stack_depth"] for e in entries if isinstance(e["stack_depth"], int)
-    ]
+    integer_depths = [e["stack_depth"] for e in entries if isinstance(e["stack_depth"], int)]
     assert integer_depths, (
-        "expected at least one int stack_depth when ?stack=true; "
-        f"got entries={entries!r}"
+        f"expected at least one int stack_depth when ?stack=true; got entries={entries!r}"
     )
 
 

tests/test_proxy_warmup.py

@@ -209,9 +209,7 @@ async def test_startup_memory_embedder_warmup_encodes_once(tmp_path, monkeypatch
     # Swap in a hand-rolled MemoryHandler whose backend exposes a mock
     # embedder. We don't want real ONNX here — just a spy.
     handler = MemoryHandler(
-        MemoryConfig(
-            enabled=True, backend="local", db_path=str(tmp_path / "mem.db")
-        )
+        MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
     )
     handler._initialized = True
 
@@ -257,9 +255,7 @@ async def test_startup_memory_backend_error_surfaced_and_health_degraded(tmp_pat
     proxy = HeadroomProxy(config)
 
     handler = MemoryHandler(
-        MemoryConfig(
-            enabled=True, backend="local", db_path=str(tmp_path / "mem.db")
-        )
+        MemoryConfig(enabled=True, backend="local", db_path=str(tmp_path / "mem.db"))
     )
 
     async def _boom() -> None:

tests/test_scripts/test_repro_codex_replay_smoke.py

@@ -49,6 +49,7 @@ def _restore_real_websockets_module() -> Iterator[None]:
         if mod is not None:
             sys.modules[name] = mod
 
+
 # Make sure `scripts/` is importable when running via pytest from repo root.
 ROOT = Path(__file__).resolve().parents[2]
 SCRIPTS_DIR = ROOT / "scripts"

uv.lock

@@ -1401,6 +1401,7 @@ dev = [
     { name = "sentence-transformers" },
     { name = "sqlite-vec" },
     { name = "uvicorn" },
+    { name = "websockets" },
 ]
 evals = [
     { name = "anthropic" },
@@ -1572,6 +1573,7 @@ requires-dist = [
     { name = "uvicorn", marker = "extra == 'dev'", specifier = ">=0.23.0" },
     { name = "uvicorn", marker = "extra == 'proxy'", specifier = ">=0.23.0" },
     { name = "watchdog", marker = "extra == 'proxy'", specifier = ">=4.0.0" },
+    { name = "websockets", marker = "extra == 'dev'", specifier = ">=13.0" },
     { name = "websockets", marker = "extra == 'proxy'", specifier = ">=13.0" },
     { name = "zstandard", marker = "extra == 'proxy'", specifier = ">=0.20.0" },
 ]