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

fix: expose compression latency bottlenecks

Browse source 
Add Codex WS unit-level timing and bounded parallel compression, clarify context-tool session savings, and avoid costly diff/log fallbacks to Kompress.
This commit is contained in:
Tejas Chopra 2026-05-12 13:34:08 -07:00
parent 0f6df1fef0
commit e9cae0131b
14 changed files with 1523 additions and 367 deletions
Download patch file Download diff file Expand all files Collapse all files

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

File diff suppressed because it is too large Load diff

309
headroom/proxy/helpers.py
View file

@ -571,7 +571,7 @@ _CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL"
_CONTEXT_TOOL_RTK = "rtk"
_CONTEXT_TOOL_LEAN_CTX = "lean-ctx"
RTK_STATS_CACHE_TTL_SECONDS = 5.0
RTK_STATS_CACHE_TTL_SECONDS = float(os.environ.get("HEADROOM_CONTEXT_TOOL_STATS_TTL_SECONDS", "60"))
CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS = RTK_STATS_CACHE_TTL_SECONDS
_context_tool_stats_cache_lock = threading.Lock()
_context_tool_stats_cache: dict[str, Any] = {
@ -584,7 +584,11 @@ _context_tool_session_baseline: dict[str, Any] = {
"initialized": False,
"tool": None,
"total_commands": 0,
"input_tokens": 0,
"output_tokens": 0,
"tokens_saved": 0,
"total_time_ms": 0,
"captured_at": 0.0,
}
_rtk_stats_cache_lock = _context_tool_stats_cache_lock
_rtk_stats_cache = _context_tool_stats_cache
@ -849,6 +853,107 @@ def _first_value(mapping: dict[str, Any], keys: tuple[str, ...], default: Any =
return default
def _context_tool_summary_payload(
*,
tool: str,
installed: bool,
summary: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Normalize RTK/lean-ctx lifetime gain output into one schema.
Both tools expose cumulative counters, but field names vary slightly.
Headroom computes session values by subtracting a startup baseline, so
keeping raw input/output counters is necessary for a truthful session
savings percentage.
"""
summary = summary or {}
input_tokens = _coerce_int(
_first_value(
summary,
(
"total_input",
"total_input_tokens",
"input_tokens",
"tokens_input",
"totalBefore",
),
)
)
output_tokens = _coerce_int(
_first_value(
summary,
(
"total_output",
"total_output_tokens",
"output_tokens",
"tokens_output",
"totalAfter",
),
)
)
tokens_saved = _coerce_int(
_first_value(
summary,
(
"total_saved",
"tokens_saved",
"total_tokens_saved",
"saved_tokens",
"totalSaved",
),
)
)
if tokens_saved <= 0 and input_tokens > 0 and output_tokens >= 0:
tokens_saved = max(input_tokens - output_tokens, 0)
if input_tokens <= 0 and tokens_saved > 0 and output_tokens >= 0:
input_tokens = tokens_saved + output_tokens
lifetime_savings_pct = _coerce_float(
_first_value(
summary,
(
"avg_savings_pct",
"average_savings_pct",
"savings_pct",
"savings_percent",
"avgSavingsPct",
),
0.0,
)
)
if lifetime_savings_pct <= 0 and input_tokens > 0:
lifetime_savings_pct = (tokens_saved / input_tokens) * 100.0
return {
"tool": tool,
"label": _context_tool_label(tool),
"installed": installed,
"scope": "project" if tool == _CONTEXT_TOOL_RTK else "local",
"total_commands": _coerce_int(
_first_value(
summary,
(
"total_commands",
"commands",
"command_count",
"totalCommandCount",
),
)
),
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"tokens_saved": tokens_saved,
# Backward-compatible name. See `lifetime_avg_savings_pct` and
# `session_savings_pct` below for explicit scopes.
"avg_savings_pct": lifetime_savings_pct,
"lifetime_avg_savings_pct": lifetime_savings_pct,
"total_time_ms": _coerce_int(
_first_value(summary, ("total_time_ms", "time_ms", "totalTimeMs"))
),
}
def _read_rtk_lifetime_stats() -> dict[str, Any] | None:
"""Read rtk's current project-level lifetime stats."""
@ -860,9 +965,14 @@ def _read_rtk_lifetime_stats() -> dict[str, Any] | None:
"tool": _CONTEXT_TOOL_RTK,
"label": _context_tool_label(_CONTEXT_TOOL_RTK),
"installed": False,
"scope": "project",
"total_commands": 0,
"input_tokens": 0,
"output_tokens": 0,
"tokens_saved": 0,
"avg_savings_pct": 0.0,
"lifetime_avg_savings_pct": 0.0,
"total_time_ms": 0,
}
try:
@ -875,31 +985,38 @@ def _read_rtk_lifetime_stats() -> dict[str, Any] | None:
if result.returncode == 0 and result.stdout.strip():
data = json.loads(result.stdout)
summary = data.get("summary", {})
payload = {
"tool": _CONTEXT_TOOL_RTK,
"label": _context_tool_label(_CONTEXT_TOOL_RTK),
"installed": True,
"total_commands": _coerce_int(summary.get("total_commands", 0)),
"tokens_saved": _coerce_int(summary.get("total_saved", 0)),
"avg_savings_pct": _coerce_float(summary.get("avg_savings_pct", 0.0)),
}
payload = _context_tool_summary_payload(
tool=_CONTEXT_TOOL_RTK,
installed=True,
summary=summary if isinstance(summary, dict) else {},
)
else:
return {
"tool": _CONTEXT_TOOL_RTK,
"label": _context_tool_label(_CONTEXT_TOOL_RTK),
"installed": True,
"scope": "project",
"total_commands": 0,
"input_tokens": 0,
"output_tokens": 0,
"tokens_saved": 0,
"avg_savings_pct": 0.0,
"lifetime_avg_savings_pct": 0.0,
"total_time_ms": 0,
}
except Exception:
return {
"tool": _CONTEXT_TOOL_RTK,
"label": _context_tool_label(_CONTEXT_TOOL_RTK),
"installed": True,
"scope": "project",
"total_commands": 0,
"input_tokens": 0,
"output_tokens": 0,
"tokens_saved": 0,
"avg_savings_pct": 0.0,
"lifetime_avg_savings_pct": 0.0,
"total_time_ms": 0,
}
return payload
@ -916,18 +1033,28 @@ def _read_lean_ctx_lifetime_stats() -> dict[str, Any] | None:
"tool": _CONTEXT_TOOL_LEAN_CTX,
"label": _context_tool_label(_CONTEXT_TOOL_LEAN_CTX),
"installed": False,
"scope": "local",
"total_commands": 0,
"input_tokens": 0,
"output_tokens": 0,
"tokens_saved": 0,
"avg_savings_pct": 0.0,
"lifetime_avg_savings_pct": 0.0,
"total_time_ms": 0,
}
base_payload = {
"tool": _CONTEXT_TOOL_LEAN_CTX,
"label": _context_tool_label(_CONTEXT_TOOL_LEAN_CTX),
"installed": True,
"scope": "local",
"total_commands": 0,
"input_tokens": 0,
"output_tokens": 0,
"tokens_saved": 0,
"avg_savings_pct": 0.0,
"lifetime_avg_savings_pct": 0.0,
"total_time_ms": 0,
}
try:
@ -945,44 +1072,11 @@ def _read_lean_ctx_lifetime_stats() -> dict[str, Any] | None:
if not isinstance(summary, dict):
return dict(base_payload)
return {
**base_payload,
"total_commands": _coerce_int(
_first_value(
summary,
(
"total_commands",
"commands",
"command_count",
"totalCommandCount",
),
)
),
"tokens_saved": _coerce_int(
_first_value(
summary,
(
"total_saved",
"tokens_saved",
"total_tokens_saved",
"saved_tokens",
"totalSaved",
),
)
),
"avg_savings_pct": _coerce_float(
_first_value(
summary,
(
"avg_savings_pct",
"average_savings_pct",
"avgSavingsPct",
"savings_percent",
),
0.0,
)
),
}
return _context_tool_summary_payload(
tool=_CONTEXT_TOOL_LEAN_CTX,
installed=True,
summary=summary,
)
except Exception:
return dict(base_payload)
@ -1004,7 +1098,11 @@ def initialize_context_tool_session_baseline() -> None:
"initialized": True,
"tool": tool,
"total_commands": int((payload or {}).get("total_commands", 0) or 0),
"input_tokens": int((payload or {}).get("input_tokens", 0) or 0),
"output_tokens": int((payload or {}).get("output_tokens", 0) or 0),
"tokens_saved": int((payload or {}).get("tokens_saved", 0) or 0),
"total_time_ms": int((payload or {}).get("total_time_ms", 0) or 0),
"captured_at": time.time(),
}
)
_context_tool_stats_cache.update(
@ -1053,25 +1151,124 @@ def _get_context_tool_stats() -> dict[str, Any] | None:
"initialized": True,
"tool": tool,
"total_commands": int((payload or {}).get("total_commands", 0) or 0),
"input_tokens": int((payload or {}).get("input_tokens", 0) or 0),
"output_tokens": int((payload or {}).get("output_tokens", 0) or 0),
"tokens_saved": int((payload or {}).get("tokens_saved", 0) or 0),
"total_time_ms": int((payload or {}).get("total_time_ms", 0) or 0),
"captured_at": time.time(),
}
)
if payload is not None:
lifetime_total_commands = int(payload.get("total_commands", 0) or 0)
lifetime_input_tokens = int(payload.get("input_tokens", 0) or 0)
lifetime_output_tokens = int(payload.get("output_tokens", 0) or 0)
lifetime_tokens_saved = int(payload.get("tokens_saved", 0) or 0)
lifetime_total_time_ms = int(payload.get("total_time_ms", 0) or 0)
baseline_total_commands = int(_context_tool_session_baseline["total_commands"])
baseline_input_tokens = int(_context_tool_session_baseline["input_tokens"])
baseline_output_tokens = int(_context_tool_session_baseline["output_tokens"])
baseline_tokens_saved = int(_context_tool_session_baseline["tokens_saved"])
baseline_total_time_ms = int(_context_tool_session_baseline["total_time_ms"])
counter_reset_detected = (
lifetime_total_commands < baseline_total_commands
or lifetime_input_tokens < baseline_input_tokens
or lifetime_output_tokens < baseline_output_tokens
or lifetime_tokens_saved < baseline_tokens_saved
or lifetime_total_time_ms < baseline_total_time_ms
)
if counter_reset_detected:
baseline_total_commands = lifetime_total_commands
baseline_input_tokens = lifetime_input_tokens
baseline_output_tokens = lifetime_output_tokens
baseline_tokens_saved = lifetime_tokens_saved
baseline_total_time_ms = lifetime_total_time_ms
_context_tool_session_baseline.update(
{
"total_commands": baseline_total_commands,
"input_tokens": baseline_input_tokens,
"output_tokens": baseline_output_tokens,
"tokens_saved": baseline_tokens_saved,
"total_time_ms": baseline_total_time_ms,
"captured_at": time.time(),
}
)
session_total_commands = max(lifetime_total_commands - baseline_total_commands, 0)
session_input_tokens = max(lifetime_input_tokens - baseline_input_tokens, 0)
session_output_tokens = max(lifetime_output_tokens - baseline_output_tokens, 0)
session_tokens_saved = max(lifetime_tokens_saved - baseline_tokens_saved, 0)
session_total_time_ms = max(lifetime_total_time_ms - baseline_total_time_ms, 0)
session_savings_pct = (
round(session_tokens_saved / session_input_tokens * 100.0, 4)
if session_input_tokens > 0
else None
)
session_avg_time_ms = (
round(session_total_time_ms / session_total_commands, 2)
if session_total_commands > 0 and session_total_time_ms > 0
else None
)
lifetime_savings_pct = float(payload.get("lifetime_avg_savings_pct", 0.0) or 0.0)
payload = {
**payload,
"tool": tool,
"label": _context_tool_label(tool),
"total_commands": max(
int(payload.get("total_commands", 0) or 0)
- int(_context_tool_session_baseline["total_commands"]),
0,
),
"tokens_saved": max(
int(payload.get("tokens_saved", 0) or 0)
- int(_context_tool_session_baseline["tokens_saved"]),
0,
# Backward-compatible session-delta fields.
"total_commands": session_total_commands,
"input_tokens": session_input_tokens,
"output_tokens": session_output_tokens,
"tokens_saved": session_tokens_saved,
"total_time_ms": session_total_time_ms,
"session_savings_pct": session_savings_pct,
"session_avg_time_ms": session_avg_time_ms,
# Keep old field for compatibility, but declare its scope.
"avg_savings_pct": lifetime_savings_pct,
"avg_savings_pct_scope": "lifetime",
"lifetime_avg_savings_pct": lifetime_savings_pct,
"lifetime_total_commands": lifetime_total_commands,
"lifetime_input_tokens": lifetime_input_tokens,
"lifetime_output_tokens": lifetime_output_tokens,
"lifetime_tokens_saved": lifetime_tokens_saved,
"lifetime_total_time_ms": lifetime_total_time_ms,
"session_baseline_total_commands": baseline_total_commands,
"session_baseline_input_tokens": baseline_input_tokens,
"session_baseline_output_tokens": baseline_output_tokens,
"session_baseline_tokens_saved": baseline_tokens_saved,
"session_baseline_total_time_ms": baseline_total_time_ms,
"session_baseline_captured_at": _context_tool_session_baseline.get(
"captured_at", 0.0
),
"session": {
"commands": session_total_commands,
"input_tokens": session_input_tokens,
"output_tokens": session_output_tokens,
"tokens_saved": session_tokens_saved,
"savings_pct": session_savings_pct,
"total_time_ms": session_total_time_ms,
"avg_time_ms": session_avg_time_ms,
},
"lifetime": {
"commands": lifetime_total_commands,
"input_tokens": lifetime_input_tokens,
"output_tokens": lifetime_output_tokens,
"tokens_saved": lifetime_tokens_saved,
"savings_pct": lifetime_savings_pct,
"total_time_ms": lifetime_total_time_ms,
},
"baseline": {
"commands": baseline_total_commands,
"input_tokens": baseline_input_tokens,
"output_tokens": baseline_output_tokens,
"tokens_saved": baseline_tokens_saved,
"total_time_ms": baseline_total_time_ms,
"captured_at": _context_tool_session_baseline.get("captured_at", 0.0),
},
"sampled_at": time.time(),
"sample_ttl_seconds": CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS,
"refresh_interval_seconds": CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS,
"counter_reset_detected": counter_reset_detected,
}
_context_tool_stats_cache.update(

134
headroom/proxy/prometheus_metrics.py
View file

@ -106,6 +106,36 @@ class PrometheusMetrics:
self.compressions_by_strategy: dict[str, int] = defaultdict(int)
self.tokens_saved_by_strategy: dict[str, int] = defaultdict(int)
# Codex WebSocket compression observability. These are intentionally
# aggregate counters/sums, not per-unit storage, so /stats can answer
# routing questions without growing with traffic volume.
self.codex_ws_units_total = 0
self.codex_ws_units_modified_total = 0
self.codex_ws_units_to_kompress_total = 0
self.codex_ws_units_kompress_attempted_total = 0
self.codex_ws_units_by_strategy: dict[str, int] = defaultdict(int)
self.codex_ws_units_by_category: dict[str, int] = defaultdict(int)
self.codex_ws_units_by_content_type: dict[str, int] = defaultdict(int)
self.codex_ws_units_by_text_shape: dict[str, int] = defaultdict(int)
self.codex_ws_unit_elapsed_ms_sum = 0.0
self.codex_ws_unit_elapsed_ms_max = 0.0
self.codex_ws_unit_bytes_sum = 0
self.codex_ws_unit_tokens_before_sum = 0
self.codex_ws_unit_tokens_after_sum = 0
self.codex_ws_unit_tokens_saved_sum = 0
self.codex_ws_frames_attempted_total = 0
self.codex_ws_frames_compressed_total = 0
self.codex_ws_frames_failed_total = 0
self.codex_ws_frames_to_kompress_total = 0
self.codex_ws_frames_kompress_attempted_total = 0
self.codex_ws_frame_elapsed_ms_sum = 0.0
self.codex_ws_frame_elapsed_ms_max = 0.0
self.codex_ws_frame_bytes_before_sum = 0
self.codex_ws_frame_bytes_after_sum = 0
self.codex_ws_frame_attempted_tokens_sum = 0
self.codex_ws_frame_tokens_saved_sum = 0
self.latency_sum_ms = 0.0
self.latency_min_ms = float("inf")
self.latency_max_ms = 0.0
@ -242,6 +272,33 @@ class PrometheusMetrics:
self.compressions_by_strategy.clear()
self.tokens_saved_by_strategy.clear()
self.codex_ws_units_total = 0
self.codex_ws_units_modified_total = 0
self.codex_ws_units_to_kompress_total = 0
self.codex_ws_units_kompress_attempted_total = 0
self.codex_ws_units_by_strategy.clear()
self.codex_ws_units_by_category.clear()
self.codex_ws_units_by_content_type.clear()
self.codex_ws_units_by_text_shape.clear()
self.codex_ws_unit_elapsed_ms_sum = 0.0
self.codex_ws_unit_elapsed_ms_max = 0.0
self.codex_ws_unit_bytes_sum = 0
self.codex_ws_unit_tokens_before_sum = 0
self.codex_ws_unit_tokens_after_sum = 0
self.codex_ws_unit_tokens_saved_sum = 0
self.codex_ws_frames_attempted_total = 0
self.codex_ws_frames_compressed_total = 0
self.codex_ws_frames_failed_total = 0
self.codex_ws_frames_to_kompress_total = 0
self.codex_ws_frames_kompress_attempted_total = 0
self.codex_ws_frame_elapsed_ms_sum = 0.0
self.codex_ws_frame_elapsed_ms_max = 0.0
self.codex_ws_frame_bytes_before_sum = 0
self.codex_ws_frame_bytes_after_sum = 0
self.codex_ws_frame_attempted_tokens_sum = 0
self.codex_ws_frame_tokens_saved_sum = 0
self.latency_sum_ms = 0.0
self.latency_min_ms = float("inf")
self.latency_max_ms = 0.0
@ -383,6 +440,83 @@ class PrometheusMetrics:
if count > 0:
self.router_route_counts[category] += int(count)
def record_codex_ws_unit(
self,
*,
strategy: str,
reason_category: str,
elapsed_ms: float,
text_bytes: int,
tokens_before: int,
tokens_after: int,
tokens_saved: int,
modified: bool,
strategy_chain: list[str] | None = None,
content_type: str = "unknown",
text_shape: str = "unknown",
) -> None:
"""Record one Codex WS compression unit decision."""
strategy = strategy or "unknown"
reason_category = reason_category or "unknown"
chain = strategy_chain or []
self.codex_ws_units_total += 1
self.codex_ws_units_by_strategy[strategy] += 1
self.codex_ws_units_by_category[reason_category] += 1
self.codex_ws_units_by_content_type[content_type or "unknown"] += 1
self.codex_ws_units_by_text_shape[text_shape or "unknown"] += 1
if modified:
self.codex_ws_units_modified_total += 1
if strategy == "kompress":
self.codex_ws_units_to_kompress_total += 1
if "kompress" in chain or strategy == "kompress":
self.codex_ws_units_kompress_attempted_total += 1
elapsed_ms = max(0.0, float(elapsed_ms))
self.codex_ws_unit_elapsed_ms_sum += elapsed_ms
self.codex_ws_unit_elapsed_ms_max = max(self.codex_ws_unit_elapsed_ms_max, elapsed_ms)
self.codex_ws_unit_bytes_sum += max(0, int(text_bytes))
self.codex_ws_unit_tokens_before_sum += max(0, int(tokens_before))
self.codex_ws_unit_tokens_after_sum += max(0, int(tokens_after))
self.codex_ws_unit_tokens_saved_sum += max(0, int(tokens_saved))
def record_codex_ws_frame(
self,
*,
elapsed_ms: float,
bytes_before: int,
bytes_after: int = 0,
attempted_tokens: int = 0,
tokens_saved: int = 0,
modified: bool = False,
failed: bool = False,
strategy_chain: list[str] | None = None,
final_strategies: list[str] | None = None,
) -> None:
"""Record one Codex WS response.create compression attempt."""
chain = strategy_chain or []
strategies = final_strategies or []
self.codex_ws_frames_attempted_total += 1
if modified:
self.codex_ws_frames_compressed_total += 1
if failed:
self.codex_ws_frames_failed_total += 1
if "kompress" in strategies:
self.codex_ws_frames_to_kompress_total += 1
if "kompress" in chain or "kompress" in strategies:
self.codex_ws_frames_kompress_attempted_total += 1
elapsed_ms = max(0.0, float(elapsed_ms))
self.codex_ws_frame_elapsed_ms_sum += elapsed_ms
self.codex_ws_frame_elapsed_ms_max = max(self.codex_ws_frame_elapsed_ms_max, elapsed_ms)
self.codex_ws_frame_bytes_before_sum += max(0, int(bytes_before))
self.codex_ws_frame_bytes_after_sum += max(0, int(bytes_after))
self.codex_ws_frame_attempted_tokens_sum += max(0, int(attempted_tokens))
self.codex_ws_frame_tokens_saved_sum += max(0, int(tokens_saved))
def record_inbound_request(self, *, method: str, path: str) -> None:
self.inbound_requests_total += 1
self.inbound_requests_active += 1

154
headroom/proxy/server.py
View file

@ -498,9 +498,16 @@ class HeadroomProxy(
# Gauge: currently-running compression tasks. Mutated under
# ``_compression_metrics_lock`` from worker threads + the asyncio
# event loop.
self._compression_queued: int = 0
self._compression_queued_max: int = 0
self._compression_queue_timeouts: int = 0
self._compression_queue_wait_seconds_total: float = 0.0
self._compression_queue_wait_seconds_max: float = 0.0
self._compression_in_flight: int = 0
# High-water mark for in-flight count.
self._compression_in_flight_max: int = 0
self._compression_run_seconds_total: float = 0.0
self._compression_run_seconds_max: float = 0.0
# Counter: threads that finished AFTER their asyncio future hit the
# timeout. Stuck-thread leak indicator.
self._compression_leaked_threads: int = 0
@ -683,10 +690,12 @@ class HeadroomProxy(
actually cancel a thread that has started — Python has no way to
preempt running CPython bytecode or in-flight Rust calls. The
worker keeps running to completion, ignored. We detect this by
recording the start timestamp and incrementing
marking the call timed out on the asyncio side and incrementing
``_compression_leaked_threads`` from the worker's ``finally``
block when ``elapsed > timeout``. Operators can see leaked-thread
rate climbing in ``/stats`` before the pool fills up.
block after it eventually finishes. Jobs that time out before a
worker starts are removed from the queued gauge instead. Operators
can see leaked-thread rate and queue pressure climbing in
``/stats`` before the pool fills up.
Args:
fn: A no-arg sync callable that runs the compression. Must not
@ -706,24 +715,49 @@ class HeadroomProxy(
unchanged.
"""
loop = asyncio.get_running_loop()
start = time.monotonic()
queued_at = time.monotonic()
state = {"queued": True, "timed_out": False}
with self._compression_metrics_lock:
self._compression_in_flight += 1
if self._compression_in_flight > self._compression_in_flight_max:
self._compression_in_flight_max = self._compression_in_flight
self._compression_queued += 1
if self._compression_queued > self._compression_queued_max:
self._compression_queued_max = self._compression_queued
def _wrapped(): # noqa: ANN202
started_at = time.monotonic()
queue_wait = started_at - queued_at
with self._compression_metrics_lock:
if state["queued"]:
self._compression_queued -= 1
state["queued"] = False
self._compression_queue_wait_seconds_total += queue_wait
if queue_wait > self._compression_queue_wait_seconds_max:
self._compression_queue_wait_seconds_max = queue_wait
self._compression_in_flight += 1
if self._compression_in_flight > self._compression_in_flight_max:
self._compression_in_flight_max = self._compression_in_flight
try:
return fn()
finally:
elapsed = time.monotonic() - start
elapsed = time.monotonic() - started_at
with self._compression_metrics_lock:
self._compression_in_flight -= 1
if elapsed > timeout:
self._compression_run_seconds_total += elapsed
if elapsed > self._compression_run_seconds_max:
self._compression_run_seconds_max = elapsed
if state["timed_out"]:
self._compression_leaked_threads += 1
future = loop.run_in_executor(self._compression_executor, _wrapped)
return await asyncio.wait_for(future, timeout=timeout)
try:
return await asyncio.wait_for(future, timeout=timeout)
except asyncio.TimeoutError:
with self._compression_metrics_lock:
state["timed_out"] = True
if state["queued"]:
self._compression_queued -= 1
state["queued"] = False
self._compression_queue_timeouts += 1
raise
def _get_compression_cache(self, session_id: str) -> CompressionCache:
"""Get or create a CompressionCache for a session.
@ -1491,8 +1525,15 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
# Snapshot compression executor metrics under their lock (gauges
# mutated by worker threads; not safe to read without).
with proxy._compression_metrics_lock:
_comp_queued = proxy._compression_queued
_comp_queued_max = proxy._compression_queued_max
_comp_queue_timeouts = proxy._compression_queue_timeouts
_comp_queue_wait_total = proxy._compression_queue_wait_seconds_total
_comp_queue_wait_max = proxy._compression_queue_wait_seconds_max
_comp_in_flight = proxy._compression_in_flight
_comp_in_flight_max = proxy._compression_in_flight_max
_comp_run_total = proxy._compression_run_seconds_total
_comp_run_max = proxy._compression_run_seconds_max
_comp_leaked = proxy._compression_leaked_threads
return {
"anthropic_pre_upstream": {
@ -1510,8 +1551,16 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
},
"compression_executor": {
"max_workers": proxy.compression_max_workers,
"queued": _comp_queued,
"queued_max": _comp_queued_max,
"queue_timeouts_total": _comp_queue_timeouts,
"queue_wait_seconds_total": _comp_queue_wait_total,
"queue_wait_seconds_max": _comp_queue_wait_max,
"running": _comp_in_flight,
"in_flight": _comp_in_flight,
"in_flight_max": _comp_in_flight_max,
"run_seconds_total": _comp_run_total,
"run_seconds_max": _comp_run_max,
"leaked_threads_total": _comp_leaked,
"source": ("auto" if config.compression_max_workers is None else "explicit"),
},
@ -1782,6 +1831,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
)
max_ttfb_ms = round(m.ttfb_max_ms, 2) if m.ttfb_count > 0 else 0
def _pct(part: int | float, whole: int | float) -> float:
return round((float(part) / float(whole)) * 100.0, 2) if whole else 0.0
# Get compression store stats
store = get_compression_store()
compression_stats = store.get_stats()
@ -1809,6 +1861,12 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
cli_tokens_avoided = (
cli_filtering_stats.get("tokens_saved", 0) if cli_filtering_stats else 0
)
cli_filtering_session = (
cli_filtering_stats.get("session", {}) if cli_filtering_stats else {}
)
cli_filtering_lifetime = (
cli_filtering_stats.get("lifetime", {}) if cli_filtering_stats else {}
)
rtk_tokens_avoided = cli_tokens_avoided if cli_filtering_tool == "rtk" else 0
lean_ctx_tokens_avoided = cli_tokens_avoided if cli_filtering_tool == "lean-ctx" else 0
@ -1890,6 +1948,23 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
"label": cli_filtering_label,
"tokens": cli_tokens_avoided,
"tokens_saved": cli_tokens_avoided,
"session": cli_filtering_session,
"lifetime": cli_filtering_lifetime,
"session_savings_pct": (
cli_filtering_stats.get("session_savings_pct")
if cli_filtering_stats
else None
),
"lifetime_savings_pct": (
cli_filtering_stats.get("lifetime_avg_savings_pct")
if cli_filtering_stats
else None
),
"refresh_interval_seconds": (
cli_filtering_stats.get("refresh_interval_seconds")
if cli_filtering_stats
else None
),
"included_in": "tokens.saved",
"description": (
f"Tokens avoided by CLI output filtering ({cli_filtering_label}) "
@ -2009,6 +2084,65 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
else {},
"compressions_by_strategy": dict(m.compressions_by_strategy),
"tokens_saved_by_strategy": dict(m.tokens_saved_by_strategy),
"codex_ws": {
"units_total": m.codex_ws_units_total,
"units_modified_total": m.codex_ws_units_modified_total,
"units_by_strategy": dict(m.codex_ws_units_by_strategy),
"units_by_category": dict(m.codex_ws_units_by_category),
"units_by_content_type": dict(m.codex_ws_units_by_content_type),
"units_by_text_shape": dict(m.codex_ws_units_by_text_shape),
"units_to_kompress_total": m.codex_ws_units_to_kompress_total,
"units_kompress_attempted_total": m.codex_ws_units_kompress_attempted_total,
"units_to_kompress_percent": _pct(
m.codex_ws_units_to_kompress_total,
m.codex_ws_units_total,
),
"units_kompress_attempted_percent": _pct(
m.codex_ws_units_kompress_attempted_total,
m.codex_ws_units_total,
),
"unit_elapsed_ms": {
"average": round(
m.codex_ws_unit_elapsed_ms_sum / m.codex_ws_units_total,
2,
)
if m.codex_ws_units_total
else 0.0,
"max": round(m.codex_ws_unit_elapsed_ms_max, 2),
},
"unit_bytes_sum": m.codex_ws_unit_bytes_sum,
"unit_tokens_before_sum": m.codex_ws_unit_tokens_before_sum,
"unit_tokens_after_sum": m.codex_ws_unit_tokens_after_sum,
"unit_tokens_saved_sum": m.codex_ws_unit_tokens_saved_sum,
"frames_attempted_total": m.codex_ws_frames_attempted_total,
"frames_compressed_total": m.codex_ws_frames_compressed_total,
"frames_failed_total": m.codex_ws_frames_failed_total,
"frames_to_kompress_total": m.codex_ws_frames_to_kompress_total,
"frames_kompress_attempted_total": (
m.codex_ws_frames_kompress_attempted_total
),
"frames_to_kompress_percent": _pct(
m.codex_ws_frames_to_kompress_total,
m.codex_ws_frames_attempted_total,
),
"frames_kompress_attempted_percent": _pct(
m.codex_ws_frames_kompress_attempted_total,
m.codex_ws_frames_attempted_total,
),
"frame_elapsed_ms": {
"average": round(
m.codex_ws_frame_elapsed_ms_sum / m.codex_ws_frames_attempted_total,
2,
)
if m.codex_ws_frames_attempted_total
else 0.0,
"max": round(m.codex_ws_frame_elapsed_ms_max, 2),
},
"frame_bytes_before_sum": m.codex_ws_frame_bytes_before_sum,
"frame_bytes_after_sum": m.codex_ws_frame_bytes_after_sum,
"frame_attempted_tokens_sum": m.codex_ws_frame_attempted_tokens_sum,
"frame_tokens_saved_sum": m.codex_ws_frame_tokens_saved_sum,
},
"waste_signals": dict(m.waste_signals_total) if m.waste_signals_total else {},
# ContentRouter protection categories aggregated across the
# session. Lets operators see, e.g., that 80% of messages

6
headroom/transforms/compression_units.py
View file

@ -13,7 +13,11 @@ from collections.abc import Iterable
from dataclasses import dataclass, field, replace
from typing import Protocol
from .content_router import CompressionStrategy, ContentRouter, RouterCompressionResult
from .content_router import (
CompressionStrategy,
ContentRouter,
RouterCompressionResult,
)
class TokenCounterLike(Protocol):

1
headroom/transforms/content_detector.py
View file

@ -108,6 +108,7 @@ _LOG_PATTERNS = [
re.compile(r"^\s*PASSED|^\s*FAILED|^\s*SKIPPED"), # test results
re.compile(r"^npm ERR!|^yarn error|^cargo error"), # build tools
re.compile(r"Traceback \(most recent call last\)"), # Python traceback
re.compile(r"^\w*(Error|Exception):"), # Python exception final line
re.compile(r"^\s*at\s+[\w.$]+\("), # JS/Java stack trace
]

222
headroom/transforms/content_router.py
View file

@ -49,6 +49,7 @@ from ..config import DEFAULT_EXCLUDE_TOOLS, ReadLifecycleConfig, TransformResult
from ..tokenizer import Tokenizer
from .base import Transform
from .content_detector import ContentType, DetectionResult
from .content_detector import detect_content_type as _regex_detect_content_type
logger = logging.getLogger(__name__)
@ -58,8 +59,10 @@ def _router_debug_dumps(value: Any) -> str:
def _log_router_debug(event: str, **payload: Any) -> None:
if not logger.isEnabledFor(logging.DEBUG):
return
payload = {"event": event, **payload}
logger.info("event=%s %s", event, _router_debug_dumps(payload))
logger.debug("event=%s %s", event, _router_debug_dumps(payload))
def _json_shape(content: str) -> dict[str, Any]:
@ -125,6 +128,10 @@ def _detect_content(content: str) -> DetectionResult:
# "json_array"); translate to the Python `ContentType` enum so
# downstream mapping keys match.
content_type = ContentType(rust_result.content_type)
if content_type is ContentType.PLAIN_TEXT:
regex_result = _regex_detect_content_type(content)
if regex_result.content_type is not ContentType.PLAIN_TEXT:
return regex_result
return DetectionResult(
content_type=content_type,
confidence=rust_result.confidence,
@ -837,25 +844,31 @@ class ContentRouter(Transform):
Returns:
RouterCompressionResult with compressed content and routing metadata.
"""
request_debug = {
"chars": len(content),
"bytes": len(content.encode("utf-8", errors="replace")),
"tokens_estimate": len(content.split()),
"json_shape": _json_shape(content),
"mixed_indicators": _mixed_indicators(content),
"context_chars": len(context),
"question": question,
"bias": bias,
"content": content,
"context": context,
}
debug_enabled = logger.isEnabledFor(logging.DEBUG)
request_debug = (
{
"chars": len(content),
"bytes": len(content.encode("utf-8", errors="replace")),
"tokens_estimate": len(content.split()),
"json_shape": _json_shape(content),
"mixed_indicators": _mixed_indicators(content),
"context_chars": len(context),
"question": question,
"bias": bias,
"content": content,
"context": context,
}
if debug_enabled
else {}
)
if not content or not content.strip():
_log_router_debug(
"content_router_input",
**request_debug,
selected_strategy=CompressionStrategy.PASSTHROUGH.value,
selection_reason="empty_or_whitespace",
)
if debug_enabled:
_log_router_debug(
"content_router_input",
**request_debug,
selected_strategy=CompressionStrategy.PASSTHROUGH.value,
selection_reason="empty_or_whitespace",
)
result = RouterCompressionResult(
compressed=content,
original=content,
@ -867,14 +880,15 @@ class ContentRouter(Transform):
mixed = is_mixed_content(content)
detection = _detect_content(content)
strategy = self._determine_strategy(content)
_log_router_debug(
"content_router_input",
**request_debug,
detected_content_type=detection.content_type.value,
detection_confidence=detection.confidence,
selected_strategy=strategy.value,
selection_reason="mixed_content" if mixed else "content_detection",
)
if debug_enabled:
_log_router_debug(
"content_router_input",
**request_debug,
detected_content_type=detection.content_type.value,
detection_confidence=detection.confidence,
selected_strategy=strategy.value,
selection_reason="mixed_content" if mixed else "content_detection",
)
if strategy == CompressionStrategy.MIXED:
result = self._compress_mixed(content, context, question, bias=bias)
@ -885,30 +899,31 @@ class ContentRouter(Transform):
# forcing function for catching strategy-level regressions.
# Empty routing_log (passthrough fast path) → no calls.
self._observe(result)
_log_router_debug(
"content_router_output",
selected_strategy=result.strategy_used.value,
sections_processed=result.sections_processed,
total_original_tokens=result.total_original_tokens,
total_compressed_tokens=result.total_compressed_tokens,
tokens_saved=result.tokens_saved,
savings_percentage=result.savings_percentage,
compression_ratio=result.compression_ratio,
routing_log=[
{
"content_type": decision.content_type.value,
"strategy": decision.strategy.value,
"original_tokens": decision.original_tokens,
"compressed_tokens": decision.compressed_tokens,
"confidence": decision.confidence,
"section_index": decision.section_index,
"compression_ratio": decision.compression_ratio,
}
for decision in result.routing_log
],
original=result.original,
compressed=result.compressed,
)
if debug_enabled:
_log_router_debug(
"content_router_output",
selected_strategy=result.strategy_used.value,
sections_processed=result.sections_processed,
total_original_tokens=result.total_original_tokens,
total_compressed_tokens=result.total_compressed_tokens,
tokens_saved=result.tokens_saved,
savings_percentage=result.savings_percentage,
compression_ratio=result.compression_ratio,
routing_log=[
{
"content_type": decision.content_type.value,
"strategy": decision.strategy.value,
"original_tokens": decision.original_tokens,
"compressed_tokens": decision.compressed_tokens,
"confidence": decision.confidence,
"section_index": decision.section_index,
"compression_ratio": decision.compression_ratio,
}
for decision in result.routing_log
],
original=result.original,
compressed=result.compressed,
)
return result
def _observe(self, result: RouterCompressionResult) -> None:
@ -997,12 +1012,13 @@ class ContentRouter(Transform):
RouterCompressionResult with reassembled content.
"""
sections = split_into_sections(content)
_log_router_debug(
"content_router_mixed_sections",
section_count=len(sections),
sections=[_section_debug(section, idx) for idx, section in enumerate(sections)],
content=content,
)
if logger.isEnabledFor(logging.DEBUG):
_log_router_debug(
"content_router_mixed_sections",
section_count=len(sections),
sections=[_section_debug(section, idx) for idx, section in enumerate(sections)],
content=content,
)
if not sections:
return RouterCompressionResult(
@ -1241,6 +1257,12 @@ class ContentRouter(Transform):
compressor_name = "KompressCompressor"
decision_reason = "text_uses_kompress"
elif strategy == CompressionStrategy.PASSTHROUGH:
compressed = content
compressed_tokens = original_tokens
compressor_name = "Passthrough"
decision_reason = "explicit_passthrough"
except Exception as e:
error = f"{type(e).__name__}: {e}"
decision_reason = "compression_exception"
@ -1251,8 +1273,6 @@ class ContentRouter(Transform):
fallback_eligible_strategy = strategy in {
CompressionStrategy.SMART_CRUSHER,
CompressionStrategy.CODE_AWARE,
CompressionStrategy.DIFF,
CompressionStrategy.LOG,
}
fallback_no_savings = compressed == content or compressed_tokens >= original_tokens
if fallback_eligible_strategy and fallback_no_savings:
@ -1294,25 +1314,28 @@ class ContentRouter(Transform):
f"{decision_reason}_fallback_log_after_no_savings"
)
_log_router_debug(
"content_router_strategy_result",
requested_strategy=requested_strategy.value,
actual_strategy=actual_strategy.value,
strategy_chain=strategy_chain,
compressor=compressor_name,
reason=decision_reason,
language=language,
question=question,
bias=bias,
original_tokens=original_tokens,
compressed_tokens=compressed_tokens,
tokens_saved=max(0, original_tokens - compressed_tokens),
compression_ratio=compressed_tokens / original_tokens if original_tokens else 1.0,
json_shape=_json_shape(content),
input=content,
output=compressed,
error=error,
)
if logger.isEnabledFor(logging.DEBUG):
_log_router_debug(
"content_router_strategy_result",
requested_strategy=requested_strategy.value,
actual_strategy=actual_strategy.value,
strategy_chain=strategy_chain,
compressor=compressor_name,
reason=decision_reason,
language=language,
question=question,
bias=bias,
original_tokens=original_tokens,
compressed_tokens=compressed_tokens,
tokens_saved=max(0, original_tokens - compressed_tokens),
compression_ratio=compressed_tokens / original_tokens
if original_tokens
else 1.0,
json_shape=_json_shape(content),
input=content,
output=compressed,
error=error,
)
self._record_to_toin(
strategy=strategy,
content=content,
@ -1326,25 +1349,26 @@ class ContentRouter(Transform):
# Fallback: return unchanged
strategy_chain.append(CompressionStrategy.PASSTHROUGH.value)
_log_router_debug(
"content_router_strategy_result",
requested_strategy=requested_strategy.value,
actual_strategy=CompressionStrategy.PASSTHROUGH.value,
strategy_chain=strategy_chain,
compressor=None,
reason=decision_reason,
language=language,
question=question,
bias=bias,
original_tokens=original_tokens,
compressed_tokens=original_tokens,
tokens_saved=0,
compression_ratio=1.0,
json_shape=_json_shape(content),
input=content,
output=content,
error=error,
)
if logger.isEnabledFor(logging.DEBUG):
_log_router_debug(
"content_router_strategy_result",
requested_strategy=requested_strategy.value,
actual_strategy=CompressionStrategy.PASSTHROUGH.value,
strategy_chain=strategy_chain,
compressor=None,
reason=decision_reason,
language=language,
question=question,
bias=bias,
original_tokens=original_tokens,
compressed_tokens=original_tokens,
tokens_saved=0,
compression_ratio=1.0,
json_shape=_json_shape(content),
input=content,
output=content,
error=error,
)
return content, original_tokens, strategy_chain
def _try_ml_compressor(
@ -1529,8 +1553,10 @@ class ContentRouter(Transform):
if self.config.enable_kompress:
compressor = self._get_kompress()
if compressor:
logger.info("Kompress model pre-loaded at startup")
backend = compressor.preload() if hasattr(compressor, "preload") else "unknown"
logger.info("Kompress model pre-loaded at startup backend=%s", backend)
status["kompress"] = "enabled"
status["kompress_backend"] = str(backend)
else:
status["kompress"] = "unavailable"

6
headroom/transforms/kompress_compressor.py
View file

@ -403,6 +403,12 @@ class KompressCompressor(Transform):
def __init__(self, config: KompressConfig | None = None):
self.config = config or KompressConfig()
def preload(self) -> str:
"""Load the backing model/tokenizer and return the selected backend."""
_model, _tokenizer, backend = _load_kompress(self.config.model_id, self.config.device)
return backend
def compress(
self,
content: str,

67
tests/test_compression_observability.py
View file

@ -273,6 +273,73 @@ def test_prometheus_metrics_accumulates_per_strategy_counters():
assert m.tokens_saved_by_strategy == {"smart_crusher": 210}
def test_prometheus_metrics_accumulates_codex_ws_unit_and_frame_counters():
from headroom.proxy.prometheus_metrics import PrometheusMetrics
m = PrometheusMetrics()
m.record_codex_ws_unit(
strategy="mixed",
reason_category="applied",
elapsed_ms=1250,
text_bytes=10_000,
tokens_before=2500,
tokens_after=1000,
tokens_saved=1500,
modified=True,
strategy_chain=["mixed", "kompress"],
content_type="text",
text_shape="jsonl_like",
)
m.record_codex_ws_unit(
strategy="passthrough",
reason_category="size_floor",
elapsed_ms=2,
text_bytes=100,
tokens_before=20,
tokens_after=20,
tokens_saved=0,
modified=False,
strategy_chain=["passthrough"],
content_type="unknown",
text_shape="plain_text_like",
)
m.record_codex_ws_frame(
elapsed_ms=1260,
bytes_before=20_000,
bytes_after=8_000,
attempted_tokens=2500,
tokens_saved=1500,
modified=True,
strategy_chain=["mixed", "kompress"],
final_strategies=["mixed"],
)
m.record_codex_ws_frame(
elapsed_ms=30_000,
bytes_before=426_318,
failed=True,
)
assert m.codex_ws_units_total == 2
assert m.codex_ws_units_modified_total == 1
assert m.codex_ws_units_by_strategy == {"mixed": 1, "passthrough": 1}
assert m.codex_ws_units_by_category == {"applied": 1, "size_floor": 1}
assert m.codex_ws_units_by_content_type == {"text": 1, "unknown": 1}
assert m.codex_ws_units_by_text_shape == {"jsonl_like": 1, "plain_text_like": 1}
assert m.codex_ws_units_to_kompress_total == 0
assert m.codex_ws_units_kompress_attempted_total == 1
assert m.codex_ws_unit_elapsed_ms_max == 1250
assert m.codex_ws_unit_tokens_saved_sum == 1500
assert m.codex_ws_frames_attempted_total == 2
assert m.codex_ws_frames_compressed_total == 1
assert m.codex_ws_frames_failed_total == 1
assert m.codex_ws_frames_to_kompress_total == 0
assert m.codex_ws_frames_kompress_attempted_total == 1
assert m.codex_ws_frame_elapsed_ms_max == 30_000
assert m.codex_ws_frame_tokens_saved_sum == 1500
def test_prometheus_export_does_not_leak_per_strategy_metrics():
"""Per-strategy state is tracked in-process only. The Prometheus
scrape output deliberately must NOT emit new metric names — the

9
tests/test_openai_codex_ws_lifecycle.py
View file

@ -360,7 +360,10 @@ async def test_ws_session_metrics_include_dashboard_performance_timings():
handler = _DummyOpenAIHandler()
handler.config.optimize = True
def _noop_compress(payload, *, model, request_id):
def _noop_compress(payload, *, model, request_id, timing=None):
if timing is not None:
timing["compression_live_unit_extraction"] = 2.0
timing["compression_unit_router_strategy_passthrough"] = 3.0
return payload, False, 0, [], "test_noop", 10, 10, 0
handler._compress_openai_responses_payload = _noop_compress # type: ignore[method-assign]
@ -374,6 +377,10 @@ async def test_ws_session_metrics_include_dashboard_performance_timings():
assert recorded["ttfb_ms"] > 0
assert recorded["pipeline_timing"]["codex_ws.compression"] > 0
assert recorded["pipeline_timing"]["codex_ws.upstream_first_event"] > 0
assert recorded["pipeline_timing"]["codex_ws.compression_preflight_serialization"] > 0
assert recorded["pipeline_timing"]["codex_ws.compression_executor_wait_run"] > 0
assert recorded["pipeline_timing"]["codex_ws.compression_live_unit_extraction"] == 2.0
assert recorded["pipeline_timing"]["codex_ws.compression_unit_router_strategy_passthrough"] == 3.0
@pytest.mark.asyncio

18
tests/test_openai_codex_ws_timings.py
View file

@ -11,6 +11,7 @@ from unittest.mock import MagicMock, patch
import anyio
import pytest
import headroom.proxy.handlers.openai as openai_handler
from headroom.proxy.handlers.openai import OpenAIHandlerMixin
@ -278,3 +279,20 @@ def test_codex_ws_request_id_and_session_id_present_in_log(stage_log_capture):
assert payload["request_id"] == "req-ws-test"
assert isinstance(payload["session_id"], str)
assert len(payload["session_id"]) >= 16
def test_codex_compression_debug_noop_skips_expensive_payload_debug(monkeypatch):
handler = _DummyOpenAIHandler()
def _fail_context_budget(_payload):
raise AssertionError("debug context budget should not be built")
monkeypatch.setattr(openai_handler, "_openai_responses_context_budget", _fail_context_budget)
result = handler._compress_openai_responses_payload(
{"model": "gpt-5.4", "input": "hello"},
model="gpt-5.4",
request_id="req-ws-test",
)
assert result[1] is False
assert result[4] == "router_no_compression"

62
tests/test_proxy_compression_executor.py
View file

@ -13,8 +13,9 @@ Locks the following invariants:
``TimeoutError`` — but the worker thread keeps running (Python cannot
preempt running CPython bytecode or in-flight Rust calls), and when the
work eventually completes, ``compression_leaked_threads`` increments.
4. ``/stats runtime.compression_executor`` surfaces the gauge + counter so
operators can see leaked-thread rate.
4. Jobs that time out while still queued do not leak the running gauge.
5. ``/stats runtime.compression_executor`` surfaces the gauges + counters so
operators can see leaked-thread rate and queue pressure.
These tests also serve as documentation: anyone reading them sees that
"timeout fired" does not mean "compression was cancelled" — it means "we
@ -205,6 +206,58 @@ def test_timeout_fires_and_leaked_thread_is_counted() -> None:
assert proxy._compression_in_flight == 0
def test_timeout_before_worker_start_does_not_leak_in_flight() -> None:
"""If a queued job times out before a worker starts, queued accounting
is cleaned up without touching the running gauge.
"""
proxy = _make_proxy(compression_max_workers=1)
first_started = threading.Event()
release_first = threading.Event()
second_started = threading.Event()
def _blocking_compression():
first_started.set()
release_first.wait(timeout=5.0)
return "first"
def _queued_compression():
second_started.set()
return "second"
async def _drive():
first_task = asyncio.create_task(
proxy._run_compression_in_executor(_blocking_compression, timeout=10.0)
)
for _ in range(50):
if first_started.is_set():
break
await asyncio.sleep(0.01)
assert first_started.is_set()
with pytest.raises(asyncio.TimeoutError):
await proxy._run_compression_in_executor(_queued_compression, timeout=0.05)
with proxy._compression_metrics_lock:
mid_queued = proxy._compression_queued
mid_in_flight = proxy._compression_in_flight
queue_timeouts = proxy._compression_queue_timeouts
release_first.set()
assert await first_task == "first"
return mid_queued, mid_in_flight, queue_timeouts
mid_queued, mid_in_flight, queue_timeouts = asyncio.run(_drive())
assert not second_started.is_set()
assert mid_queued == 0
assert mid_in_flight == 1
assert queue_timeouts == 1
with proxy._compression_metrics_lock:
assert proxy._compression_queued == 0
assert proxy._compression_in_flight == 0
assert proxy._compression_leaked_threads == 0
def test_compression_executor_metrics_appear_in_runtime_payload() -> None:
"""``/stats runtime.compression_executor`` surfaces the new gauges."""
from fastapi.testclient import TestClient
@ -232,7 +285,12 @@ def test_compression_executor_metrics_appear_in_runtime_payload() -> None:
assert "compression_executor" in runtime
ce = runtime["compression_executor"]
assert ce["max_workers"] == 5
assert ce["queued"] == 0
assert ce["running"] == 0
assert ce["in_flight"] == 0
assert ce["queue_timeouts_total"] == 0
assert ce["queue_wait_seconds_total"] == 0.0
assert ce["run_seconds_total"] == 0.0
assert ce["leaked_threads_total"] == 0
assert ce["source"] == "explicit"

144
tests/test_proxy_dashboard_stats_cache.py
View file

@ -36,7 +36,16 @@ def _reset_rtk_stats_cache(monkeypatch: pytest.MonkeyPatch) -> None:
{"expires_at": 0.0, "has_value": False, "tool": None, "value": None}
)
proxy_helpers._rtk_session_baseline.update(
{"initialized": False, "tool": None, "total_commands": 0, "tokens_saved": 0}
{
"initialized": False,
"tool": None,
"total_commands": 0,
"input_tokens": 0,
"output_tokens": 0,
"tokens_saved": 0,
"total_time_ms": 0,
"captured_at": 0.0,
}
)
@ -45,8 +54,22 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch
now = {"value": 100.0}
calls = {"run": 0}
totals = [
{"total_commands": 7, "total_saved": 1234},
{"total_commands": 9, "total_saved": 1500},
{
"total_commands": 7,
"total_input": 2000,
"total_output": 766,
"total_saved": 1234,
"avg_savings_pct": 61.7,
"total_time_ms": 700,
},
{
"total_commands": 9,
"total_input": 2600,
"total_output": 1100,
"total_saved": 1500,
"avg_savings_pct": 57.69,
"total_time_ms": 1000,
},
]
def _fake_run(*args, **kwargs):
@ -65,26 +88,57 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch
second = proxy_helpers._get_rtk_stats()
assert first == second
assert first == {
"tool": "rtk",
"label": "RTK",
"installed": True,
"total_commands": 0,
"tokens_saved": 0,
"avg_savings_pct": 0.0,
}
assert first["tool"] == "rtk"
assert first["label"] == "RTK"
assert first["installed"] is True
assert first["total_commands"] == 0
assert first["input_tokens"] == 0
assert first["output_tokens"] == 0
assert first["tokens_saved"] == 0
assert first["session_savings_pct"] is None
assert first["avg_savings_pct"] == 61.7
assert first["avg_savings_pct_scope"] == "lifetime"
assert first["lifetime_total_commands"] == 7
assert first["lifetime_input_tokens"] == 2000
assert first["lifetime_output_tokens"] == 766
assert first["lifetime_tokens_saved"] == 1234
assert first["session_baseline_total_commands"] == 7
assert first["session_baseline_input_tokens"] == 2000
assert first["session_baseline_output_tokens"] == 766
assert first["session_baseline_tokens_saved"] == 1234
assert first["session"]["tokens_saved"] == 0
assert first["lifetime"]["savings_pct"] == 61.7
assert first["sample_ttl_seconds"] == proxy_helpers.CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS
assert calls["run"] == 1
now["value"] += proxy_helpers.RTK_STATS_CACHE_TTL_SECONDS + 0.1
third = proxy_helpers._get_rtk_stats()
assert third == {
"tool": "rtk",
"label": "RTK",
"installed": True,
"total_commands": 2,
assert third["tool"] == "rtk"
assert third["label"] == "RTK"
assert third["installed"] is True
assert third["total_commands"] == 2
assert third["input_tokens"] == 600
assert third["output_tokens"] == 334
assert third["tokens_saved"] == 266
assert third["session_savings_pct"] == pytest.approx(44.3333)
assert third["session_avg_time_ms"] == 150.0
assert third["lifetime_total_commands"] == 9
assert third["lifetime_input_tokens"] == 2600
assert third["lifetime_output_tokens"] == 1100
assert third["lifetime_tokens_saved"] == 1500
assert third["session_baseline_total_commands"] == 7
assert third["session_baseline_input_tokens"] == 2000
assert third["session_baseline_output_tokens"] == 766
assert third["session_baseline_tokens_saved"] == 1234
assert third["session"] == {
"commands": 2,
"input_tokens": 600,
"output_tokens": 334,
"tokens_saved": 266,
"avg_savings_pct": 0.0,
"savings_pct": pytest.approx(44.3333),
"total_time_ms": 300,
"avg_time_ms": 150.0,
}
assert calls["run"] == 2
@ -94,8 +148,20 @@ def test_get_context_tool_stats_reads_lean_ctx_gain(monkeypatch: pytest.MonkeyPa
now = {"value": 100.0}
calls = {"run": 0}
totals = [
{"total_commands": 3, "tokens_saved": 400, "avg_savings_pct": 12.5},
{"total_commands": 5, "tokens_saved": 475, "avg_savings_pct": 15.0},
{
"total_commands": 3,
"total_input_tokens": 1000,
"total_output_tokens": 600,
"tokens_saved": 400,
"avg_savings_pct": 40.0,
},
{
"total_commands": 5,
"total_input_tokens": 1250,
"total_output_tokens": 775,
"tokens_saved": 475,
"avg_savings_pct": 38.0,
},
]
def _fake_run(args, **kwargs):
@ -115,27 +181,35 @@ def test_get_context_tool_stats_reads_lean_ctx_gain(monkeypatch: pytest.MonkeyPa
second = proxy_helpers._get_context_tool_stats()
assert first == second
assert first == {
"tool": "lean-ctx",
"label": "lean-ctx",
"installed": True,
"total_commands": 0,
"tokens_saved": 0,
"avg_savings_pct": 12.5,
}
assert first["tool"] == "lean-ctx"
assert first["label"] == "lean-ctx"
assert first["installed"] is True
assert first["total_commands"] == 0
assert first["tokens_saved"] == 0
assert first["avg_savings_pct"] == 40.0
assert first["session_savings_pct"] is None
assert first["lifetime_total_commands"] == 3
assert first["lifetime_input_tokens"] == 1000
assert first["lifetime_output_tokens"] == 600
assert first["lifetime_tokens_saved"] == 400
assert calls["run"] == 1
now["value"] += proxy_helpers.CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS + 0.1
third = proxy_helpers._get_context_tool_stats()
assert third == {
"tool": "lean-ctx",
"label": "lean-ctx",
"installed": True,
"total_commands": 2,
"tokens_saved": 75,
"avg_savings_pct": 15.0,
}
assert third["tool"] == "lean-ctx"
assert third["label"] == "lean-ctx"
assert third["installed"] is True
assert third["total_commands"] == 2
assert third["input_tokens"] == 250
assert third["output_tokens"] == 175
assert third["tokens_saved"] == 75
assert third["avg_savings_pct"] == 38.0
assert third["avg_savings_pct_scope"] == "lifetime"
assert third["session_savings_pct"] == 30.0
assert third["lifetime_total_commands"] == 5
assert third["lifetime_tokens_saved"] == 475
assert third["session"]["savings_pct"] == 30.0
assert calls["run"] == 2

56
tests/test_transforms_content_router.py
View file

@ -302,6 +302,62 @@ def test_content_router_mixed_pure_apply_and_toin(monkeypatch: pytest.MonkeyPatc
assert len(calls) == 1
def test_diff_strategy_does_not_fallback_to_kompress_when_diff_is_noop(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = ContentRouter()
diff = "diff --git a/a.txt b/a.txt\n--- a/a.txt\n+++ b/a.txt\n@@ -1 +1 @@\n-a\n+a"
class NoopDiffCompressor:
def compress(self, content: str, context: str = "") -> SimpleNamespace:
return SimpleNamespace(compressed=content)
monkeypatch.setattr(router, "_get_diff_compressor", lambda: NoopDiffCompressor())
def fail_kompress(*_args: object, **_kwargs: object) -> tuple[str, int]:
raise AssertionError("Diff compression must not fallback to Kompress")
monkeypatch.setattr(router, "_try_ml_compressor", fail_kompress)
compressed, compressed_tokens, strategy_chain = router._apply_strategy_to_content(
diff,
CompressionStrategy.DIFF,
context="",
)
assert compressed == diff
assert compressed_tokens == len(diff.split())
assert strategy_chain == ["diff"]
def test_log_strategy_does_not_fallback_to_kompress_when_log_is_noop(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router = ContentRouter()
log = "ERROR one\nERROR two\nERROR three"
class NoopLogCompressor:
def compress(self, content: str, bias: float = 1.0) -> SimpleNamespace:
return SimpleNamespace(compressed=content)
monkeypatch.setattr(router, "_get_log_compressor", lambda: NoopLogCompressor())
def fail_kompress(*_args: object, **_kwargs: object) -> tuple[str, int]:
raise AssertionError("Log compression must not fallback to Kompress")
monkeypatch.setattr(router, "_try_ml_compressor", fail_kompress)
compressed, compressed_tokens, strategy_chain = router._apply_strategy_to_content(
log,
CompressionStrategy.LOG,
context="",
)
assert compressed == log
assert compressed_tokens == len(log.split())
assert strategy_chain == ["log"]
# ---------------------------------------------------------------------------
# Cache-safety tests for _process_content_blocks. These pin down the
# block-level invariants that protect upstream prefix caches: