akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(proxy): add per-stage timings for Codex WS and Anthropic HTTP paths

Browse source 
Unit 2 of the Codex proxy resilience plan. Introduce a lightweight
StageTimer utility and thread it through handle_openai_responses_ws and
handle_anthropic_messages to capture per-stage durations
(accept/first_client_frame/upstream_connect/upstream_first_event/
memory_context/compression/total_session for Codex WS;
read_request_json/deep_copy/compression_first_stage/memory_context/
upstream_connect/upstream_first_byte/total_pre_upstream for Anthropic
HTTP).

Each request/session emits exactly one structured log line
(STAGE_TIMINGS ...) and records sum/count/max Prometheus histograms
under headroom_stage_timing_ms_{sum,count,max} keyed by
(path, stage). A new session_id UUID is generated per request/session
and paired with the existing request_id so multi-turn sessions are
correlatable.

Compression, WS retries, WS->HTTP fallback, and the memory-context
fail-open timeout are unchanged. RequestLog gains three optional
backward-compatible fields (stage_timings, session_id,
stage_timings_path).
This commit is contained in:
Adryan Eka Vandra 2026-04-17 12:35:12 +07:00
parent b0422cd27a
commit c2fee2f706
No known key found for this signature in database
GPG key ID: A46A577A26A97682
8 changed files with 1183 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

143
headroom/proxy/handlers/anthropic.py
View file

@ -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,

104
headroom/proxy/handlers/openai.py
View file

@ -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,

17
headroom/proxy/models.py
View file

@ -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:

70
headroom/proxy/prometheus_metrics.py
View file

@ -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(
[

190
headroom/proxy/stage_timer.py Normal file
View file

@ -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}"
)

296
tests/test_anthropic_stage_timings.py Normal file
View file

@ -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

280
tests/test_openai_codex_ws_timings.py Normal file
View file

@ -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

139
tests/test_stage_timer.py Normal file
View file

@ -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