diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index d78f762fb..4f0c5e9ec 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -11,9 +11,12 @@ import json import logging import os import time +import uuid from datetime import datetime from typing import TYPE_CHECKING, Any +from headroom.proxy.stage_timer import StageTimer, emit_stage_timings_log + if TYPE_CHECKING: from fastapi import Request from fastapi.responses import Response, StreamingResponse @@ -313,10 +316,46 @@ class AnthropicHandlerMixin: start_time = time.time() request_id = await self._next_request_id() + trace_session_id = uuid.uuid4().hex + + # Unit 2: per-stage timings for the pre-upstream phase. The + # finalizer emits one structured log line + Prometheus + # observations even if the handler raises. + stage_timer = StageTimer() + pre_upstream_started_at = time.perf_counter() + _stage_timings_emitted = False + + async def _emit_pre_upstream_stage_timings() -> None: + nonlocal _stage_timings_emitted + if _stage_timings_emitted: + return + _stage_timings_emitted = True + if "total_pre_upstream" not in stage_timer: + stage_timer.record( + "total_pre_upstream", + (time.perf_counter() - pre_upstream_started_at) * 1000.0, + ) + await emit_stage_timings_log( + path="anthropic_messages", + request_id=request_id, + session_id=trace_session_id, + stage_timer=stage_timer, + expected_stages=( + "read_request_json", + "deep_copy", + "compression_first_stage", + "memory_context", + "upstream_connect", + "upstream_first_byte", + "total_pre_upstream", + ), + metrics=getattr(self, "metrics", None), + ) # Check request body size content_length = request.headers.get("content-length") if content_length and int(content_length) > MAX_REQUEST_BODY_SIZE: + await _emit_pre_upstream_stage_timings() return JSONResponse( status_code=413, content={ @@ -330,8 +369,10 @@ class AnthropicHandlerMixin: # Parse request try: - body = await _read_request_json(request) + async with stage_timer.measure("read_request_json"): + body = await _read_request_json(request) except (json.JSONDecodeError, ValueError) as e: + await _emit_pre_upstream_stage_timings() return JSONResponse( status_code=400, content={ @@ -344,7 +385,8 @@ class AnthropicHandlerMixin: ) model = body.get("model", "unknown") messages = body.get("messages", []) - original_client_messages = copy.deepcopy(messages) + with stage_timer.measure("deep_copy"): + original_client_messages = copy.deepcopy(messages) # Validate message array size if len(messages) > MAX_MESSAGE_ARRAY_LENGTH: @@ -609,19 +651,20 @@ class AnthropicHandlerMixin: frozen_message_count = ttl_frozen - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.anthropic_pipeline.apply( - messages=working_messages, - model=model, - model_limit=context_limit, - context=extract_user_query(working_messages), - frozen_message_count=frozen_message_count, - biases=biases, - ) - ), - timeout=COMPRESSION_TIMEOUT_SECONDS, - ) + async with stage_timer.measure("compression_first_stage"): + result = await asyncio.wait_for( + asyncio.to_thread( + lambda: self.anthropic_pipeline.apply( + messages=working_messages, + model=model, + model_limit=context_limit, + context=extract_user_query(working_messages), + frozen_message_count=frozen_message_count, + biases=biases, + ) + ), + timeout=COMPRESSION_TIMEOUT_SECONDS, + ) # Cache newly compressed messages (index-aligned diff) if result.messages != working_messages: @@ -636,19 +679,20 @@ class AnthropicHandlerMixin: # original_tokens was set at line ~2183 from uncompressed messages. optimized_tokens = result.tokens_after elif not is_cache_mode(self.config.mode): - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.anthropic_pipeline.apply( - messages=messages, - model=model, - model_limit=context_limit, - context=extract_user_query(messages), - frozen_message_count=frozen_message_count, - biases=biases, - ) - ), - timeout=COMPRESSION_TIMEOUT_SECONDS, - ) + async with stage_timer.measure("compression_first_stage"): + result = await asyncio.wait_for( + asyncio.to_thread( + lambda: self.anthropic_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + context=extract_user_query(messages), + frozen_message_count=frozen_message_count, + biases=biases, + ) + ), + timeout=COMPRESSION_TIMEOUT_SECONDS, + ) if result.messages != messages: optimized_messages = result.messages @@ -868,9 +912,10 @@ class AnthropicHandlerMixin: # Search and inject memory context if self.memory_handler.config.inject_context: try: - memory_context = await self.memory_handler.search_and_format_context( - memory_user_id, optimized_messages - ) + async with stage_timer.measure("memory_context"): + memory_context = await self.memory_handler.search_and_format_context( + memory_user_id, optimized_messages + ) if memory_context: if is_cache_mode(self.config.mode): logger.info( @@ -935,11 +980,19 @@ class AnthropicHandlerMixin: tools = self._sort_tools_deterministically(tools) body["tools"] = tools + # Unit 2: mark end of pre-upstream phase. Everything after this + # point is upstream I/O or post-response bookkeeping. + stage_timer.record( + "total_pre_upstream", + (time.perf_counter() - pre_upstream_started_at) * 1000.0, + ) + # Forward request - use Bedrock backend if configured, otherwise direct API if self.anthropic_backend is not None: # Route through Bedrock backend try: if stream: + await _emit_pre_upstream_stage_timings() return await self._stream_response_bedrock( body, headers, @@ -955,7 +1008,19 @@ class AnthropicHandlerMixin: pipeline_timing=pipeline_timing, ) else: - backend_response = await self.anthropic_backend.send_message(body, headers) + async with stage_timer.measure("upstream_connect"): + backend_response = await self.anthropic_backend.send_message( + body, headers + ) + # Non-stream: first-byte and connect are effectively + # the same horizon — ``send_message`` awaits until + # the response body is fully buffered. + if "upstream_first_byte" not in stage_timer and "upstream_connect" in stage_timer: + stage_timer.record( + "upstream_first_byte", + stage_timer.summary()["upstream_connect"], + ) + await _emit_pre_upstream_stage_timings() if backend_response.error: return JSONResponse( @@ -1031,6 +1096,7 @@ class AnthropicHandlerMixin: try: if stream: + await _emit_pre_upstream_stage_timings() return await self._stream_response( url, headers, @@ -1050,7 +1116,14 @@ class AnthropicHandlerMixin: original_messages=original_client_messages, ) else: - response = await self._retry_request("POST", url, headers, body) + async with stage_timer.measure("upstream_connect"): + response = await self._retry_request("POST", url, headers, body) + if "upstream_first_byte" not in stage_timer and "upstream_connect" in stage_timer: + stage_timer.record( + "upstream_first_byte", + stage_timer.summary()["upstream_connect"], + ) + await _emit_pre_upstream_stage_timings() # Full diagnostic dump on upstream errors. # Writes pre/post compression messages, tools, and error @@ -1503,6 +1576,10 @@ class AnthropicHandlerMixin: }, }, ) + finally: + # Unit 2: always emit pre-upstream stage timings exactly + # once per request, even on early/error paths. + await _emit_pre_upstream_stage_timings() async def handle_anthropic_batch_create( self, diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 593fd4d6c..da9cb6f0d 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -14,16 +14,18 @@ import logging import os import random import time +import uuid from datetime import datetime from typing import TYPE_CHECKING, Any +from headroom.proxy.stage_timer import StageTimer, emit_stage_timings_log + if TYPE_CHECKING: from fastapi import Request, WebSocket from fastapi.responses import JSONResponse, Response, StreamingResponse import httpx - logger = logging.getLogger("headroom.proxy") @@ -1239,6 +1241,12 @@ class OpenAIHandlerMixin: return request_id = await self._next_request_id() + session_id = uuid.uuid4().hex + + # Stage-timer — captures per-stage durations for the structured + # log emitted on session close. Unit 2 instrumentation. + stage_timer = StageTimer() + session_started_at = time.perf_counter() # Forward client headers to upstream, adding required OpenAI-Beta header ws_headers = dict(websocket.headers) @@ -1252,10 +1260,11 @@ class OpenAIHandlerMixin: client_subprotocols = [p.strip() for p in raw_protocol.split(",") if p.strip()] # Accept client connection with the requested subprotocol - if client_subprotocols: - await websocket.accept(subprotocol=client_subprotocols[0]) - else: - await websocket.accept() + async with stage_timer.measure("accept"): + if client_subprotocols: + await websocket.accept(subprotocol=client_subprotocols[0]) + else: + await websocket.accept() # Forward all client headers except hop-by-hop / per-connection headers. # These are WebSocket handshake mechanics that the `websockets` library @@ -1322,7 +1331,8 @@ class OpenAIHandlerMixin: try: # Receive the first message from client (the response.create request) - first_msg_raw = await websocket.receive_text() + async with stage_timer.measure("first_client_frame"): + first_msg_raw = await websocket.receive_text() # --- Optional: compress the input in the first message --- body: dict[str, Any] = {} @@ -1362,17 +1372,18 @@ class OpenAIHandlerMixin: original_tokens = tokenizer.count_messages(messages) context_limit = self.openai_provider.get_context_limit(model) - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.openai_pipeline.apply( - messages=messages, - model=model, - model_limit=context_limit, - context=extract_user_query(messages), - ) - ), - timeout=COMPRESSION_TIMEOUT_SECONDS, - ) + async with stage_timer.measure("compression"): + result = await asyncio.wait_for( + asyncio.to_thread( + lambda: self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + context=extract_user_query(messages), + ) + ), + timeout=COMPRESSION_TIMEOUT_SECONDS, + ) if result.messages != messages: opt = result.messages @@ -1446,12 +1457,13 @@ class OpenAIHandlerMixin: ws_msgs.extend(converted_msgs) try: - memory_context = await asyncio.wait_for( - self.memory_handler.search_and_format_context( - memory_user_id, ws_msgs - ), - timeout=RESPONSES_CONTEXT_SEARCH_TIMEOUT_SECONDS, - ) + async with stage_timer.measure("memory_context"): + memory_context = await asyncio.wait_for( + self.memory_handler.search_and_format_context( + memory_user_id, ws_msgs + ), + timeout=RESPONSES_CONTEXT_SEARCH_TIMEOUT_SECONDS, + ) except TimeoutError: memory_context = None logger.info( @@ -1537,6 +1549,9 @@ class OpenAIHandlerMixin: ws_connected = False ws_connect_attempts = max(1, getattr(self.config, "retry_max_attempts", 3)) ws_last_err: Exception | None = None + _upstream_connect_started = time.perf_counter() + _upstream_connect_recorded = False + _upstream_first_event_started: float | None = None for ws_attempt in range(ws_connect_attempts): try: @@ -1555,6 +1570,13 @@ class OpenAIHandlerMixin: ping_timeout=20, ) as upstream: ws_connected = True + if not _upstream_connect_recorded: + stage_timer.record( + "upstream_connect", + (time.perf_counter() - _upstream_connect_started) * 1000.0, + ) + _upstream_connect_recorded = True + _upstream_first_event_started = time.perf_counter() await upstream.send(first_msg_raw) async def _client_to_upstream() -> None: @@ -1601,8 +1623,22 @@ class OpenAIHandlerMixin: pending_fcs.clear() resp_id = None + # The retry-loop variable is safe to close over here: + # ``_upstream_to_client`` is defined and awaited within + # a single iteration and never escapes. + _first_event_started_at = _upstream_first_event_started # noqa: B023 + try: async for msg in upstream: + if ( + _first_event_started_at is not None + and "upstream_first_event" not in stage_timer + ): + stage_timer.record( + "upstream_first_event", + (time.perf_counter() - _first_event_started_at) + * 1000.0, + ) if isinstance(msg, bytes): await websocket.send_bytes(msg) continue @@ -1820,6 +1856,28 @@ class OpenAIHandlerMixin: logger.error(f"[{request_id}] WS proxy error: {error_detail}") with contextlib.suppress(Exception): await websocket.close(code=1011, reason=str(e)[:120]) + finally: + # Unit 2: emit structured per-session stage timings. + stage_timer.record( + "total_session", + (time.perf_counter() - session_started_at) * 1000.0, + ) + await emit_stage_timings_log( + path="openai_responses_ws", + request_id=request_id, + session_id=session_id, + stage_timer=stage_timer, + expected_stages=( + "accept", + "first_client_frame", + "upstream_connect", + "upstream_first_event", + "memory_context", + "compression", + "total_session", + ), + metrics=getattr(self, "metrics", None), + ) async def _ws_http_fallback( self, diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index 446bae552..51e609c5d 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -48,6 +48,23 @@ class RequestLog: response_content: str | None = None error: str | None = None + # Per-stage timings (Unit 2 — Codex WS + Anthropic HTTP paths) + # Populated from ``StageTimer.summary()``. Keys are stage names + # (e.g. ``accept``, ``upstream_connect``, ``memory_context``); + # values are durations in milliseconds. Absent when the handler did + # not thread a timer through (backward-compatible default). + stage_timings: dict[str, float] | None = None + + # Session id — a UUID generated at WS accept / HTTP request start, + # paired with ``request_id`` so multi-turn sessions can be + # correlated. Optional for backward compatibility. + session_id: str | None = None + + # Path key under which stage timings were recorded (e.g. + # ``openai_responses_ws`` or ``anthropic_messages``). Used by the + # Prometheus histogram series. + stage_timings_path: str | None = None + @dataclass class CacheEntry: diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 43ea17d09..e7c58a75a 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -100,6 +100,14 @@ class PrometheusMetrics: self.transform_timing_count: dict[str, int] = defaultdict(int) self.transform_timing_max: dict[str, float] = defaultdict(float) + # Per-stage timing (Unit 2). Keyed by ``(path, stage)`` tuples so + # a single metric name can distinguish between, e.g., + # ``openai_responses_ws`` ``upstream_connect`` and + # ``anthropic_messages`` ``upstream_connect``. + self.stage_timing_sum: dict[tuple[str, str], float] = defaultdict(float) + self.stage_timing_count: dict[tuple[str, str], int] = defaultdict(int) + self.stage_timing_max: dict[tuple[str, str], float] = defaultdict(float) + # Aggregate waste signals self.waste_signals_total: dict[str, int] = defaultdict(int) @@ -317,6 +325,33 @@ class PrometheusMetrics: uncached_input_tokens=uncached_input_tokens, ) + async def record_stage_timings( + self, + path: str, + timings: dict[str, float], + ) -> None: + """Record per-stage timings as histogram-style observations. + + ``path`` identifies the code path that emitted the timings (e.g. + ``openai_responses_ws`` or ``anthropic_messages``). ``timings`` + maps stage names to millisecond durations. Mirrors the + ``transform_timing_*`` aggregation pattern so the ``/metrics`` + endpoint exposes sum/count/max series per ``(path, stage)``. + """ + if not timings: + return + async with self._lock: + for stage, ms in timings.items(): + try: + ms_val = float(ms) + except (TypeError, ValueError): + continue + key = (path, stage) + self.stage_timing_sum[key] += ms_val + self.stage_timing_count[key] += 1 + if ms_val > self.stage_timing_max[key]: + self.stage_timing_max[key] = ms_val + async def record_cache_bust(self, tokens_lost: int) -> None: """Record tokens that lost their cache discount due to compression.""" async with self._lock: @@ -541,6 +576,41 @@ class PrometheusMetrics: ) lines.append("") + if self.stage_timing_sum: + lines.extend( + [ + "# HELP headroom_stage_timing_ms_sum Sum of per-stage handler timings in milliseconds", + "# TYPE headroom_stage_timing_ms_sum counter", + ] + ) + for (path_label, stage), total in self.stage_timing_sum.items(): + lines.append( + f'headroom_stage_timing_ms_sum{{path="{_escape_label_value(path_label)}",stage="{_escape_label_value(stage)}"}} {round(total, 2)}' + ) + lines.extend( + [ + "", + "# HELP headroom_stage_timing_ms_count Count of per-stage handler timing samples", + "# TYPE headroom_stage_timing_ms_count counter", + ] + ) + for (path_label, stage), count in self.stage_timing_count.items(): + lines.append( + f'headroom_stage_timing_ms_count{{path="{_escape_label_value(path_label)}",stage="{_escape_label_value(stage)}"}} {count}' + ) + lines.extend( + [ + "", + "# HELP headroom_stage_timing_ms_max Maximum per-stage handler timing in milliseconds", + "# TYPE headroom_stage_timing_ms_max gauge", + ] + ) + for (path_label, stage), max_value in self.stage_timing_max.items(): + lines.append( + f'headroom_stage_timing_ms_max{{path="{_escape_label_value(path_label)}",stage="{_escape_label_value(stage)}"}} {round(max_value, 2)}' + ) + lines.append("") + if self.waste_signals_total: lines.extend( [ diff --git a/headroom/proxy/stage_timer.py b/headroom/proxy/stage_timer.py new file mode 100644 index 000000000..ff4a44200 --- /dev/null +++ b/headroom/proxy/stage_timer.py @@ -0,0 +1,190 @@ +"""Stage-timing instrumentation for request handlers. + +Provides a lightweight, synchronous+async context-manager utility for +measuring per-stage durations within a single request or WebSocket +session. Timings are collected into a single dict that can be emitted +on the structured log line for the request. + +The design goals: + 1. Durations are captured even if the measured body raises (the + ``finally`` clause records the partial duration before + re-raising). + 2. A single ``StageTimer`` holds every stage for one request/session + and produces a ``dict[str, float]`` of millisecond durations on + :meth:`summary`. + 3. Concurrent :meth:`measure` calls are independent — each one owns + its own entry in the dict, so overlapping stages do not collide. + 4. Uses ``time.perf_counter()`` for monotonic, high-resolution + measurement. + +This module is intentionally free of any external dependencies so it +can be safely imported from both handler code paths. +""" + +from __future__ import annotations + +import json +import logging +import time +from collections.abc import Iterable +from contextlib import AbstractAsyncContextManager, AbstractContextManager +from types import TracebackType +from typing import Any + +logger = logging.getLogger("headroom.proxy") + +__all__ = ["StageTimer", "StageMeasurement", "emit_stage_timings_log"] + + +class StageMeasurement(AbstractContextManager["StageMeasurement"], AbstractAsyncContextManager["StageMeasurement"]): + """Context manager that measures a single named stage. + + Acts as both a synchronous (``with timer.measure(...):``) and an + asynchronous (``async with timer.measure(...):``) context manager. + The body's duration is captured in a ``finally`` clause so it is + recorded even if the body raises. + """ + + __slots__ = ("_timer", "_name", "_start") + + def __init__(self, timer: StageTimer, name: str) -> None: + self._timer = timer + self._name = name + self._start: float | None = None + + def __enter__(self) -> StageMeasurement: + self._start = time.perf_counter() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self._finalize() + + async def __aenter__(self) -> StageMeasurement: + self._start = time.perf_counter() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self._finalize() + + def _finalize(self) -> None: + if self._start is None: + # Defensive — should not happen in practice. + return + duration_ms = (time.perf_counter() - self._start) * 1000.0 + self._timer._record(self._name, duration_ms) + + +class StageTimer: + """Collect per-stage durations for one request/session. + + A single instance is created at the start of a request/session and + passed through the handler. Each ``measure(stage_name)`` call + returns a context manager that records its body's duration in + milliseconds under ``stage_name``. + + Stages that never run remain absent from :meth:`summary` — callers + that need ``null`` placeholders for unused stages should overlay + them explicitly on the returned dict. + """ + + __slots__ = ("_stages", "_created_at") + + def __init__(self) -> None: + self._stages: dict[str, float] = {} + self._created_at = time.perf_counter() + + def measure(self, name: str) -> StageMeasurement: + """Return a context manager that records the named stage's duration.""" + return StageMeasurement(self, name) + + def record(self, name: str, duration_ms: float) -> None: + """Record a pre-computed duration (e.g. from an existing timer). + + If the stage already has a recorded value, the new value + replaces it. This matches the semantics of ``measure`` — the + most recent measurement for a named stage wins. Callers that + need accumulation should aggregate externally before calling + :meth:`record`. + """ + self._stages[name] = float(duration_ms) + + def _record(self, name: str, duration_ms: float) -> None: + """Internal: record a duration from a ``StageMeasurement``.""" + self._stages[name] = duration_ms + + def elapsed_ms(self) -> float: + """Return total milliseconds since the timer was created.""" + return (time.perf_counter() - self._created_at) * 1000.0 + + def summary(self) -> dict[str, float]: + """Return a snapshot of the recorded stage durations (in ms).""" + return dict(self._stages) + + def __contains__(self, name: str) -> bool: + return name in self._stages + + +async def emit_stage_timings_log( + *, + path: str, + request_id: str, + session_id: str, + stage_timer: StageTimer, + expected_stages: Iterable[str], + metrics: Any | None = None, +) -> None: + """Emit one structured log line of stage timings + record Prometheus series. + + ``expected_stages`` guarantees that every stage the handler plans + to instrument shows up in the log line, even when it never ran + (``None`` placeholder). This makes the log stream trivial to parse + without knowing the full stage vocabulary per path up-front. + + ``metrics`` (optional) is anything with an + ``async record_stage_timings(path, timings)`` method — typically + ``HeadroomProxy.metrics``. Errors from the metrics sink are logged + at DEBUG and do not propagate. + """ + summary = stage_timer.summary() + padded: dict[str, float | None] = {s: summary.get(s) for s in expected_stages} + # Include any extra stages that were recorded but not listed in + # ``expected_stages`` — defensive, covers future additions. + for extra_stage, extra_value in summary.items(): + if extra_stage not in padded: + padded[extra_stage] = extra_value + + try: + payload = json.dumps( + { + "event": "stage_timings", + "path": path, + "request_id": request_id, + "session_id": session_id, + "stages": padded, + }, + default=str, + ) + 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}" + ) + + 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}" + ) diff --git a/tests/test_anthropic_stage_timings.py b/tests/test_anthropic_stage_timings.py new file mode 100644 index 000000000..e28fafea4 --- /dev/null +++ b/tests/test_anthropic_stage_timings.py @@ -0,0 +1,296 @@ +"""Unit 2: stage-timing instrumentation on the Anthropic HTTP path.""" + +from __future__ import annotations + +import json +import logging +from types import SimpleNamespace +from unittest.mock import MagicMock + +import anyio +import pytest +from fastapi import Request + +from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin +from headroom.proxy.models import ProxyConfig + + +class _DummyTokenizer: + def count_messages(self, messages) -> int: + return 1 + + +class _DummyMetrics: + def __init__(self) -> None: + self.stage_timings: list[tuple[str, dict[str, float]]] = [] + + async def record_request(self, **kwargs): + return None + + async def record_stage_timings(self, path: str, timings: dict[str, float]) -> None: + self.stage_timings.append((path, dict(timings))) + + async def record_failed(self, **kwargs): + return None + + async def record_rate_limited(self, **kwargs): + return None + + +class _ResponseStub: + status_code = 200 + headers: dict[str, str] = {} + content = b'{"id":"msg_1","type":"message","role":"assistant","content":[],"usage":{"input_tokens":1,"output_tokens":1}}' + + def json(self): + return { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + + +class _DummyAnthropicHandler(AnthropicHandlerMixin): + ANTHROPIC_API_URL = "https://api.anthropic.com" + + def __init__(self) -> None: + self.rate_limiter = None + self.metrics = _DummyMetrics() + self.config = ProxyConfig( + optimize=False, + image_optimize=False, + retry_max_attempts=1, + retry_base_delay_ms=1, + retry_max_delay_ms=1, + connect_timeout_seconds=10, + mode="token", + cache_enabled=False, + rate_limit_enabled=False, + fallback_enabled=False, + fallback_provider=None, + prefix_freeze_enabled=False, + memory_enabled=False, + ) + self.usage_reporter = None + self.anthropic_provider = SimpleNamespace(get_context_limit=lambda model: 200_000) + self.anthropic_pipeline = SimpleNamespace(apply=MagicMock()) + self.anthropic_backend = None + self.cost_tracker = None + self.memory_handler = None + self.cache = None + self.security = None + self.ccr_context_tracker = None + self.ccr_injector = None + self.ccr_response_handler = None + self.ccr_feedback = None + self.ccr_batch_processor = None + self.ccr_mcp_server = None + self.traffic_learner = None + self.tool_injector = None + self.read_lifecycle_manager = None + self.logger = SimpleNamespace(log=lambda *a, **k: None) + self.request_logger = self.logger + self.usage_observer = None + self.image_compressor = None + self.session_tracker_store = SimpleNamespace( + compute_session_id=lambda *a, **k: "sess-1", + get_or_create=lambda *a, **k: SimpleNamespace( + get_frozen_message_count=lambda: 0, + get_last_original_messages=lambda: [], + get_last_forwarded_messages=lambda: [], + record_request=lambda *a, **k: None, + ), + ) + + async def _next_request_id(self) -> str: + return "req-anth-test" + + def _extract_tags(self, headers): + return {} + + async def _retry_request(self, method: str, url: str, headers: dict, body: dict): + self.captured = (method, url, headers, body) + return _ResponseStub() + + def _get_compression_cache(self, session_id): + return SimpleNamespace( + apply_cached=lambda m: m, + compute_frozen_count=lambda m: 0, + mark_stable_from_messages=lambda *a, **k: None, + should_defer_compression=lambda h: False, + mark_stable=lambda h: None, + content_hash=lambda c: "h", + update_from_result=lambda *a, **k: None, + _cache={}, + _stable_hashes=set(), + ) + + +def _build_request(body: dict, headers: dict[str, str]) -> Request: + payload = json.dumps(body).encode("utf-8") + + async def receive(): + return {"type": "http.request", "body": payload, "more_body": False} + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "https", + "path": "/v1/messages", + "raw_path": b"/v1/messages", + "query_string": b"", + "headers": [ + (key.lower().encode("utf-8"), value.encode("utf-8")) for key, value in headers.items() + ], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 443), + } + return Request(scope, receive) + + +class _CapturingHandler(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.INFO) + self.records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.records.append(record) + + +@pytest.fixture +def stage_log_capture(): + target = logging.getLogger("headroom.proxy") + handler = _CapturingHandler() + previous_level = target.level + target.addHandler(handler) + target.setLevel(logging.INFO) + try: + yield handler + finally: + target.removeHandler(handler) + target.setLevel(previous_level) + + +def _parse_stage_log(handler: _CapturingHandler) -> dict: + for record in handler.records: + msg = record.getMessage() + if "STAGE_TIMINGS" in msg: + payload_start = msg.index("STAGE_TIMINGS ") + len("STAGE_TIMINGS ") + return json.loads(msg[payload_start:]) + raise AssertionError("no STAGE_TIMINGS log line captured") + + +def test_anthropic_http_happy_path_emits_stage_timings(stage_log_capture): + request = _build_request( + { + "model": "claude-3-5-sonnet-latest", + "messages": [{"role": "user", "content": "hello"}], + }, + {"authorization": "Bearer sk-ant-api-test"}, + ) + handler = _DummyAnthropicHandler() + + # Force tokenizer to a stub. + import headroom.tokenizers as _tk + + orig_get = _tk.get_tokenizer + _tk.get_tokenizer = lambda model: _DummyTokenizer() + try: + anyio.run(handler.handle_anthropic_messages, request) + finally: + _tk.get_tokenizer = orig_get + + payload = _parse_stage_log(stage_log_capture) + assert payload["event"] == "stage_timings" + assert payload["path"] == "anthropic_messages" + assert payload["request_id"] == "req-anth-test" + assert payload["session_id"] + stages = payload["stages"] + + for key in ( + "read_request_json", + "deep_copy", + "compression_first_stage", + "memory_context", + "upstream_connect", + "upstream_first_byte", + "total_pre_upstream", + ): + assert key in stages, f"missing stage: {key}" + + assert stages["read_request_json"] is not None + assert stages["deep_copy"] is not None + assert stages["upstream_connect"] is not None + assert stages["upstream_first_byte"] is not None + assert stages["total_pre_upstream"] is not None + # compression + memory_context were skipped (optimize=False, no memory) + assert stages["compression_first_stage"] is None + assert stages["memory_context"] is None + + # Metrics sink got the observation too. + assert handler.metrics.stage_timings + path, emitted = handler.metrics.stage_timings[-1] + assert path == "anthropic_messages" + assert "total_pre_upstream" in emitted + + +def test_anthropic_http_invalid_body_still_emits_stage_timings(stage_log_capture): + async def receive(): + # Invalid JSON — produces ``ValueError`` from ``_read_request_json``. + return {"type": "http.request", "body": b"not-json", "more_body": False} + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "https", + "path": "/v1/messages", + "raw_path": b"/v1/messages", + "query_string": b"", + "headers": [(b"authorization", b"Bearer sk-ant-api-test")], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 443), + } + request = Request(scope, receive) + handler = _DummyAnthropicHandler() + + anyio.run(handler.handle_anthropic_messages, request) + + payload = _parse_stage_log(stage_log_capture) + stages = payload["stages"] + # Even on invalid body, ``read_request_json`` duration is recorded + # (the measure wraps the call, including the raised exception path). + assert stages["read_request_json"] is not None + # Downstream stages never ran: + assert stages["upstream_connect"] is None + assert stages["upstream_first_byte"] is None + + +def test_anthropic_http_request_and_session_ids_present(stage_log_capture): + request = _build_request( + { + "model": "claude-3-5-sonnet-latest", + "messages": [{"role": "user", "content": "hi"}], + }, + {"authorization": "Bearer sk-ant-api-test"}, + ) + handler = _DummyAnthropicHandler() + + import headroom.tokenizers as _tk + + orig_get = _tk.get_tokenizer + _tk.get_tokenizer = lambda model: _DummyTokenizer() + try: + anyio.run(handler.handle_anthropic_messages, request) + finally: + _tk.get_tokenizer = orig_get + + payload = _parse_stage_log(stage_log_capture) + assert payload["request_id"] == "req-anth-test" + assert isinstance(payload["session_id"], str) + assert len(payload["session_id"]) >= 16 diff --git a/tests/test_openai_codex_ws_timings.py b/tests/test_openai_codex_ws_timings.py new file mode 100644 index 000000000..d0337798e --- /dev/null +++ b/tests/test_openai_codex_ws_timings.py @@ -0,0 +1,280 @@ +"""Unit 2: stage-timing instrumentation on the Codex WS path.""" + +from __future__ import annotations + +import json +import logging +import sys +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import anyio +import pytest + +from headroom.proxy.handlers.openai import OpenAIHandlerMixin + + +class _DummyMetrics: + def __init__(self) -> None: + self.stage_timings: list[tuple[str, dict[str, float]]] = [] + + async def record_request(self, **kwargs): # pragma: no cover - unused here + return None + + async def record_stage_timings(self, path: str, timings: dict[str, float]) -> None: + self.stage_timings.append((path, dict(timings))) + + +class _DummyOpenAIHandler(OpenAIHandlerMixin): + OPENAI_API_URL = "https://api.openai.com" + + def __init__(self) -> 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 + + async def _next_request_id(self) -> str: + return "req-ws-test" + + +class _FakeWebSocket: + """Minimal async WebSocket stub that delivers a scripted frame list.""" + + def __init__(self, frames: list[str] | None = None, headers: dict | None = None) -> None: + self.headers = headers or {"authorization": "Bearer test"} + self._frames = list(frames or []) + self.sent_text: list[str] = [] + self.sent_bytes: list[bytes] = [] + self.accepted_subprotocol = None + self.closed = False + self.close_code: int | None = None + + async def accept(self, subprotocol=None) -> None: + self.accepted_subprotocol = subprotocol + + async def receive_text(self) -> str: + if not self._frames: + # Simulate client disconnect: raise a WebSocketDisconnect-like error. + raise RuntimeError("WebSocketDisconnect: no more frames") + return self._frames.pop(0) + + async def send_text(self, text: str) -> None: + self.sent_text.append(text) + + 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 + + +class _FakeUpstream: + """Async context manager mirroring the websockets.connect API.""" + + def __init__(self, events: list[str]) -> None: + self._events = list(events) + 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 + + +def _make_fake_websockets_module(upstream: _FakeUpstream): + module = MagicMock() + module.connect = MagicMock(return_value=upstream) + module.Subprotocol = str # the handler wraps client subprotocols if present + return module + + +class _CapturingHandler(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.INFO) + self.records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.records.append(record) + + +@pytest.fixture +def stage_log_capture(): + """Attach a ``Handler`` directly to the ``headroom.proxy`` logger. + + Using a direct handler is more robust than ``caplog`` for this + logger because upstream configuration may set ``propagate=False`` + during module import, which bypasses pytest's root-logger capture. + """ + target = logging.getLogger("headroom.proxy") + handler = _CapturingHandler() + previous_level = target.level + target.addHandler(handler) + target.setLevel(logging.INFO) + try: + yield handler + finally: + target.removeHandler(handler) + target.setLevel(previous_level) + + +def _parse_stage_log(handler: _CapturingHandler) -> dict: + for record in handler.records: + msg = record.getMessage() + if "STAGE_TIMINGS" in msg: + # msg format: "[req-id] STAGE_TIMINGS {json}" + payload_start = msg.index("STAGE_TIMINGS ") + len("STAGE_TIMINGS ") + return json.loads(msg[payload_start:]) + raise AssertionError("no STAGE_TIMINGS log line captured") + + +def test_codex_ws_happy_path_emits_all_stage_timings(stage_log_capture): + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "resp_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "resp_1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + + first_frame = json.dumps( + { + "type": "response.create", + "response": {"model": "gpt-5.4", "input": "hello"}, + } + ) + client_ws = _FakeWebSocket(frames=[first_frame]) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + anyio.run(handler.handle_openai_responses_ws, client_ws) + + # Upstream received the compressed (or unmodified) first frame + assert len(upstream.sent) == 1 + + # Structured log emitted with all expected stages + payload = _parse_stage_log(stage_log_capture) + assert payload["event"] == "stage_timings" + assert payload["path"] == "openai_responses_ws" + assert payload["request_id"] == "req-ws-test" + assert payload["session_id"] # non-empty UUID + + stages = payload["stages"] + # Every expected stage key appears in the dict (may be None when not run) + for key in ( + "accept", + "first_client_frame", + "upstream_connect", + "upstream_first_event", + "memory_context", + "compression", + "total_session", + ): + assert key in stages, f"missing stage: {key}" + + # Stages that actually ran are positive floats + assert stages["accept"] is not None and stages["accept"] >= 0.0 + assert stages["first_client_frame"] is not None + assert stages["upstream_connect"] is not None + assert stages["upstream_first_event"] is not None + assert stages["total_session"] > 0.0 + + # Stages that were skipped (no memory handler, optimize=False) are None. + assert stages["memory_context"] is None + assert stages["compression"] is None + + # Prometheus metric sink captured the same path + timings. + assert handler.metrics.stage_timings + path, emitted = handler.metrics.stage_timings[-1] + assert path == "openai_responses_ws" + assert "total_session" in emitted + + +def test_codex_ws_upstream_connect_failure_still_logs_timings(stage_log_capture): + """A session that never connects upstream still logs a timing line + with ``upstream_first_event`` absent (null).""" + + 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 + + first_frame = json.dumps( + {"type": "response.create", "response": {"model": "gpt-5.4", "input": "hi"}} + ) + client_ws = _FakeWebSocket(frames=[first_frame]) + handler = _DummyOpenAIHandler() + # With retry_max_attempts=1 we do not retry; fallback path attempts HTTP. + + # Stub the HTTP fallback so we don't need a network mock. + async def _fallback(*args, **kwargs): + return None + + handler._ws_http_fallback = _fallback # type: ignore[assignment] + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + anyio.run(handler.handle_openai_responses_ws, client_ws) + + payload = _parse_stage_log(stage_log_capture) + stages = payload["stages"] + + # upstream_first_event never fired because connect failed on entry. + assert stages.get("upstream_first_event") is None + # upstream_connect is also None because we record it only after the + # context manager successfully enters. + assert stages.get("upstream_connect") is None + # But the envelope is still complete. + assert stages["accept"] is not None + assert stages["first_client_frame"] is not None + assert stages["total_session"] > 0.0 + + +def test_codex_ws_request_id_and_session_id_present_in_log(stage_log_capture): + upstream = _FakeUpstream([]) + fake_ws_mod = _make_fake_websockets_module(upstream) + + first_frame = json.dumps( + {"type": "response.create", "response": {"model": "gpt-5.4", "input": "hi"}} + ) + client_ws = _FakeWebSocket(frames=[first_frame]) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + anyio.run(handler.handle_openai_responses_ws, client_ws) + + payload = _parse_stage_log(stage_log_capture) + assert payload["request_id"] == "req-ws-test" + assert isinstance(payload["session_id"], str) + assert len(payload["session_id"]) >= 16 diff --git a/tests/test_stage_timer.py b/tests/test_stage_timer.py new file mode 100644 index 000000000..1cde9f41a --- /dev/null +++ b/tests/test_stage_timer.py @@ -0,0 +1,139 @@ +"""Tests for ``headroom.proxy.stage_timer.StageTimer``.""" + +from __future__ import annotations + +import asyncio +import time + +import pytest + +from headroom.proxy.stage_timer import StageTimer + + +def test_sync_measure_records_duration(): + timer = StageTimer() + + with timer.measure("stage_a"): + time.sleep(0.01) + + summary = timer.summary() + assert "stage_a" in summary + assert summary["stage_a"] >= 10.0 # at least 10 ms + # Sanity: should not be wildly off (upper bound generous for slow CI). + assert summary["stage_a"] < 5000.0 + + +def test_sync_measure_records_even_when_body_raises(): + timer = StageTimer() + + with pytest.raises(RuntimeError, match="boom"): + with timer.measure("stage_err"): + time.sleep(0.005) + raise RuntimeError("boom") + + summary = timer.summary() + assert "stage_err" in summary + assert summary["stage_err"] >= 5.0 + + +@pytest.mark.asyncio +async def test_async_measure_records_duration(): + timer = StageTimer() + + async with timer.measure("async_stage"): + await asyncio.sleep(0.01) + + summary = timer.summary() + assert "async_stage" in summary + assert summary["async_stage"] >= 10.0 + + +@pytest.mark.asyncio +async def test_async_measure_records_even_when_body_raises(): + timer = StageTimer() + + with pytest.raises(ValueError, match="oops"): + async with timer.measure("async_err"): + await asyncio.sleep(0.002) + raise ValueError("oops") + + summary = timer.summary() + assert "async_err" in summary + assert summary["async_err"] >= 2.0 + + +@pytest.mark.asyncio +async def test_concurrent_measure_calls_are_independent(): + timer = StageTimer() + + async def _stage(name: str, delay: float) -> None: + async with timer.measure(name): + await asyncio.sleep(delay) + + await asyncio.gather( + _stage("fast", 0.005), + _stage("medium", 0.02), + _stage("slow", 0.04), + ) + + summary = timer.summary() + assert set(summary) == {"fast", "medium", "slow"} + # Each stage's duration reflects its own body, not the whole gather. + assert summary["fast"] < summary["medium"] + assert summary["medium"] < summary["slow"] + # All stages measured while scheduled concurrently should each be + # close to their own sleep, not to the serialized sum (~65 ms). + assert summary["slow"] < 200.0 + + +def test_summary_returns_independent_snapshot(): + timer = StageTimer() + + with timer.measure("one"): + pass + + snapshot = timer.summary() + snapshot["mutated"] = 999.0 + + # The timer itself does not retain the mutation. + assert "mutated" not in timer.summary() + + +def test_record_allows_preexisting_duration(): + timer = StageTimer() + timer.record("precomputed_ms", 42.0) + + summary = timer.summary() + assert summary["precomputed_ms"] == pytest.approx(42.0) + + +def test_elapsed_ms_is_monotonic(): + timer = StageTimer() + first = timer.elapsed_ms() + time.sleep(0.005) + second = timer.elapsed_ms() + assert second > first + assert second - first >= 5.0 + + +def test_contains_operator_checks_recorded_stages(): + timer = StageTimer() + + assert "absent" not in timer + + with timer.measure("present"): + pass + + assert "present" in timer + + +def test_unused_stages_absent_from_summary(): + timer = StageTimer() + + with timer.measure("only_one"): + pass + + summary = timer.summary() + assert summary == {"only_one": pytest.approx(summary["only_one"])} + # No sentinel placeholders — callers overlay ``null`` themselves. + assert "not_measured" not in summary