diff --git a/headroom/proxy/debug_introspection.py b/headroom/proxy/debug_introspection.py index 60254a8e2..1c3da366f 100644 --- a/headroom/proxy/debug_introspection.py +++ b/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 diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 4ce34ce5f..a89e8f828 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/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, diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 0240adca1..b3518ac40 100644 --- a/headroom/proxy/handlers/openai.py +++ b/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( diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index cb36efe36..0f213610d 100644 --- a/headroom/proxy/helpers.py +++ b/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. diff --git a/headroom/proxy/memory_handler.py b/headroom/proxy/memory_handler.py index 808e51087..d9b792bcb 100644 --- a/headroom/proxy/memory_handler.py +++ b/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: diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 5eb9dbd8e..4dfc6f585 100644 --- a/headroom/proxy/server.py +++ b/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( diff --git a/headroom/proxy/stage_timer.py b/headroom/proxy/stage_timer.py index 8c957e017..8714349c9 100644 --- a/headroom/proxy/stage_timer.py +++ b/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}") diff --git a/pyproject.toml b/pyproject.toml index 2960eb209..096e2cc4e 100644 --- a/pyproject.toml +++ b/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", diff --git a/scripts/repro_codex_replay.py b/scripts/repro_codex_replay.py index a0829a6c7..6525acd30 100755 --- a/scripts/repro_codex_replay.py +++ b/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( diff --git a/tests/test_anthropic_pre_upstream_backpressure.py b/tests/test_anthropic_pre_upstream_backpressure.py index 456228f59..09d312e3b 100644 --- a/tests/test_anthropic_pre_upstream_backpressure.py +++ b/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(): diff --git a/tests/test_memory_handler_concurrent_init.py b/tests/test_memory_handler_concurrent_init.py index b73ee044b..9100c544f 100644 --- a/tests/test_memory_handler_concurrent_init.py +++ b/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()) diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index f2dc4e365..ad3bac35f 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/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) diff --git a/tests/test_proxy_debug_endpoints.py b/tests/test_proxy_debug_endpoints.py index c64f4228e..deccbb064 100644 --- a/tests/test_proxy_debug_endpoints.py +++ b/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}" ) diff --git a/tests/test_proxy_warmup.py b/tests/test_proxy_warmup.py index a5db602db..fa09577cd 100644 --- a/tests/test_proxy_warmup.py +++ b/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: diff --git a/tests/test_scripts/test_repro_codex_replay_smoke.py b/tests/test_scripts/test_repro_codex_replay_smoke.py index 7fa4ac6db..826ad4640 100644 --- a/tests/test_scripts/test_repro_codex_replay_smoke.py +++ b/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" diff --git a/uv.lock b/uv.lock index d8509b574..0664e939c 100644 --- a/uv.lock +++ b/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" }, ]