From e9cae0131bfe0c3b6db9096b2c5604259bda673d Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Tue, 12 May 2026 13:34:08 -0700 Subject: [PATCH] fix: expose compression latency bottlenecks Add Codex WS unit-level timing and bounded parallel compression, clarify context-tool session savings, and avoid costly diff/log fallbacks to Kompress. --- headroom/proxy/handlers/openai.py | 702 ++++++++++++++++----- headroom/proxy/helpers.py | 309 +++++++-- headroom/proxy/prometheus_metrics.py | 134 ++++ headroom/proxy/server.py | 154 ++++- headroom/transforms/compression_units.py | 6 +- headroom/transforms/content_detector.py | 1 + headroom/transforms/content_router.py | 222 ++++--- headroom/transforms/kompress_compressor.py | 6 + tests/test_compression_observability.py | 67 ++ tests/test_openai_codex_ws_lifecycle.py | 9 +- tests/test_openai_codex_ws_timings.py | 18 + tests/test_proxy_compression_executor.py | 62 +- tests/test_proxy_dashboard_stats_cache.py | 144 ++++- tests/test_transforms_content_router.py | 56 ++ 14 files changed, 1523 insertions(+), 367 deletions(-) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 484fd68f7..70596acf6 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -7,12 +7,14 @@ from __future__ import annotations import asyncio import base64 +import concurrent.futures import contextlib import copy import hashlib import json import logging import os +import threading import time import uuid from datetime import datetime @@ -42,6 +44,64 @@ from headroom.proxy.auth_mode import classify_auth_mode logger = logging.getLogger("headroom.proxy") +_CODEX_WS_UNIT_ROUTER_MAX_WORKERS = 10 +_CODEX_WS_UNIT_ROUTER_SEMAPHORE = threading.BoundedSemaphore( + _CODEX_WS_UNIT_ROUTER_MAX_WORKERS +) + + +def _codex_ws_unit_worker_count(unit_count: int) -> int: + if unit_count <= 1: + return 1 + raw = os.environ.get("HEADROOM_CODEX_WS_UNIT_WORKERS", "4") + try: + requested = int(raw) + except ValueError: + requested = 4 + return max(1, min(unit_count, requested, _CODEX_WS_UNIT_ROUTER_MAX_WORKERS)) + + +def _codex_ws_text_shape(text: str) -> str: + stripped = text.strip() + if not stripped: + return "empty" + if stripped.startswith("```"): + return "code_fence" + if stripped.startswith("<") and stripped.endswith(">"): + return "xml_or_html" + if stripped.startswith("["): + return "json_array_like" + if stripped.startswith("{"): + lines = [line for line in stripped.splitlines() if line.strip()] + if len(lines) > 1 and all(line.lstrip().startswith("{") for line in lines[:20]): + return "jsonl_like" + return "json_object_like" + if stripped.startswith("Traceback (most recent call last)"): + return "traceback" + lines = stripped.splitlines() + sample = lines[:50] + if sample: + timestamp_lines = sum( + 1 + for line in sample + if len(line) >= 10 and line[:4].isdigit() and line[4:5] == "-" and line[7:8] == "-" + ) + level_lines = sum( + 1 + for line in sample + if any(level in line for level in (" ERROR ", " WARN ", " WARNING ", " INFO ")) + ) + search_lines = sum( + 1 + for line in sample + if ":" in line and line.split(":", 2)[1:2] and line.split(":", 2)[1].isdigit() + ) + if timestamp_lines >= max(2, len(sample) // 5) or level_lines >= max(2, len(sample) // 5): + return "log_like" + if search_lines >= max(2, len(sample) // 3): + return "search_result_like" + return "plain_text_like" + def _json_debug_dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False, default=str, separators=(",", ":")) @@ -51,6 +111,13 @@ def _log_codex_compression_debug(_event: str, **_payload: Any) -> None: return +_CODEX_COMPRESSION_DEBUG_NOOP = _log_codex_compression_debug + + +def _codex_compression_debug_enabled() -> bool: + return _log_codex_compression_debug is not _CODEX_COMPRESSION_DEBUG_NOOP + + def _json_shape(value: str) -> dict[str, Any]: try: parsed = json.loads(value) @@ -362,6 +429,7 @@ class OpenAIHandlerMixin: model: str, request_id: str, pass_id: str | None = None, + timing: dict[str, float] | None = None, ) -> tuple[dict[str, Any], bool, int, list[str], dict[str, int], list[str], int]: """Run ContentRouter on OpenAI Responses text units. @@ -373,8 +441,17 @@ class OpenAIHandlerMixin: are intentionally not exposed as text units. """ + debug_enabled = _codex_compression_debug_enabled() + def _log(_event: str, **_fields: Any) -> None: - return + if debug_enabled: + _log_codex_compression_debug( + _event, + request_id=request_id, + pass_id=pass_id, + model=model, + **_fields, + ) input_items = payload.get("input") messages_items = payload.get("messages") @@ -385,7 +462,7 @@ class OpenAIHandlerMixin: from headroom.transforms.compression_units import ( CompressionUnit, RoutedCompressionUnit, - compress_units_with_router, + compress_unit_with_router, find_content_router, ) except Exception as exc: @@ -447,75 +524,89 @@ class OpenAIHandlerMixin: if isinstance(call_id, str) and call_id: headroom_retrieve_call_ids.add(call_id) + timing_sink: dict[str, float] = timing if timing is not None else {} + + def _add_timing(name: str, started_at: float) -> None: + timing_sink[name] = ( + timing_sink.get(name, 0.0) + (time.perf_counter() - started_at) * 1000.0 + ) + + extraction_started = time.perf_counter() candidates: list[tuple[int, tuple[str, int | None], str]] = [] extraction_debug: list[dict[str, Any]] = [] for idx, item in enumerate(items): if not isinstance(item, dict): - extraction_debug.append( - { - "index": idx, - "eligible": False, - "reason": "item_not_dict", - "item_type": type(item).__name__, - "item": item, - } - ) + if debug_enabled: + extraction_debug.append( + { + "index": idx, + "eligible": False, + "reason": "item_not_dict", + "item_type": type(item).__name__, + "item": item, + } + ) continue item_type = item.get("type") if item_type in self.OPENAI_RESPONSES_OUTPUT_TYPES: call_id = item.get("call_id") if isinstance(call_id, str) and call_id in headroom_retrieve_call_ids: - extraction_debug.append( - { - "index": idx, - "eligible": False, - "reason": "headroom_retrieve_output_protected", - "item_type": item_type, - "call_id": call_id, - "item": item, - } - ) + if debug_enabled: + extraction_debug.append( + { + "index": idx, + "eligible": False, + "reason": "headroom_retrieve_output_protected", + "item_type": item_type, + "call_id": call_id, + "item": item, + } + ) continue slot = _slot_text(item) if slot is not None: text, slot_ref = slot candidates.append((idx, slot_ref, text)) - extraction_debug.append( - { - "index": idx, - "eligible": True, - "item_type": item_type, - "role": item.get("role"), - "slot": slot_ref, - "text_chars": len(text), - "text_bytes": len(text.encode("utf-8", errors="replace")), - "text_json_shape": _json_shape(text), - "item": item, - "text": text, - } - ) + if debug_enabled: + extraction_debug.append( + { + "index": idx, + "eligible": True, + "item_type": item_type, + "role": item.get("role"), + "slot": slot_ref, + "text_chars": len(text), + "text_bytes": len(text.encode("utf-8", errors="replace")), + "text_json_shape": _json_shape(text), + "item": item, + "text": text, + } + ) else: + if debug_enabled: + extraction_debug.append( + { + "index": idx, + "eligible": False, + "reason": "output_type_without_text_slot", + "item_type": item_type, + "item": item, + } + ) + else: + if debug_enabled: extraction_debug.append( { "index": idx, "eligible": False, - "reason": "output_type_without_text_slot", + "reason": "unsupported_item_type", "item_type": item_type, + "role": item.get("role"), "item": item, } ) - else: - extraction_debug.append( - { - "index": idx, - "eligible": False, - "reason": "unsupported_item_type", - "item_type": item_type, - "role": item.get("role"), - "item": item, - } - ) + _add_timing("compression_live_unit_extraction", extraction_started) _log( "codex_compression_extraction", item_count=len(items), @@ -535,7 +626,9 @@ class OpenAIHandlerMixin: ) return payload, False, 0, [], {}, [], 0 + deepcopy_started = time.perf_counter() updated = copy.deepcopy(payload) + _add_timing("compression_payload_deepcopy", deepcopy_started) updated_input_items = updated.get("input") updated_messages_items = updated.get("messages") updated_items = ( @@ -558,6 +651,7 @@ class OpenAIHandlerMixin: transforms: list[str] = [] routed_units: list[RoutedCompressionUnit] = [] + unit_build_started = time.perf_counter() unit_debug: list[dict[str, Any]] = [] for item_idx, slot_ref, original_text in candidates: item = items[item_idx] if item_idx < len(items) else {} @@ -574,23 +668,25 @@ class OpenAIHandlerMixin: min_bytes=self.OPENAI_RESPONSES_ROUTER_MIN_BYTES, ) routed_units.append(RoutedCompressionUnit(unit=unit, slot=(item_idx, slot_ref))) - unit_debug.append( - { - "item_index": item_idx, - "slot": slot_ref, - "provider": unit.provider, - "endpoint": unit.endpoint, - "role": unit.role, - "item_type": unit.item_type, - "cache_zone": unit.cache_zone, - "mutable": unit.mutable, - "min_bytes": unit.min_bytes, - "text_chars": len(unit.text), - "text_bytes": len(unit.text.encode("utf-8", errors="replace")), - "text_json_shape": _json_shape(unit.text), - "text": unit.text, - } - ) + if debug_enabled: + unit_debug.append( + { + "item_index": item_idx, + "slot": slot_ref, + "provider": unit.provider, + "endpoint": unit.endpoint, + "role": unit.role, + "item_type": unit.item_type, + "cache_zone": unit.cache_zone, + "mutable": unit.mutable, + "min_bytes": unit.min_bytes, + "text_chars": len(unit.text), + "text_bytes": len(unit.text.encode("utf-8", errors="replace")), + "text_json_shape": _json_shape(unit.text), + "text": unit.text, + } + ) + _add_timing("compression_unit_build", unit_build_started) _log( "codex_compression_units", @@ -603,11 +699,83 @@ class OpenAIHandlerMixin: units_by_category: dict[str, int] = {} strategy_chain_union: list[str] = [] - for slot, result in compress_units_with_router( - routed_units, - router=router, - tokenizer=tokenizer, - ): + def _compress_routed_unit( + routed: RoutedCompressionUnit, + ) -> tuple[object, Any, float]: + unit_started = time.perf_counter() + with _CODEX_WS_UNIT_ROUTER_SEMAPHORE: + result = compress_unit_with_router(routed.unit, router=router, tokenizer=tokenizer) + elapsed_ms = (time.perf_counter() - unit_started) * 1000.0 + return routed.slot, result, elapsed_ms + + router_total_started = time.perf_counter() + worker_count = _codex_ws_unit_worker_count(len(routed_units)) + if worker_count <= 1: + routed_results = [_compress_routed_unit(routed) for routed in routed_units] + else: + with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as executor: + routed_results = list(executor.map(_compress_routed_unit, routed_units)) + + for _, result, elapsed_ms in routed_results: + router_chain = list(result.router_result.strategy_chain) if result.router_result else [] + router_content_type = ( + result.router_result.routing_log[0].content_type.value + if result.router_result and result.router_result.routing_log + else "unknown" + ) + timing_sink["compression_unit_router_total"] = ( + timing_sink.get("compression_unit_router_total", 0.0) + elapsed_ms + ) + timing_sink[f"compression_unit_router_strategy_{result.strategy}"] = ( + timing_sink.get(f"compression_unit_router_strategy_{result.strategy}", 0.0) + + elapsed_ms + ) + timing_sink[f"compression_unit_router_category_{result.reason_category}"] = ( + timing_sink.get( + f"compression_unit_router_category_{result.reason_category}", 0.0 + ) + + elapsed_ms + ) + record_unit = getattr(getattr(self, "metrics", None), "record_codex_ws_unit", None) + if record_unit is not None: + record_unit( + strategy=result.strategy, + reason_category=result.reason_category, + elapsed_ms=elapsed_ms, + text_bytes=result.text_bytes, + tokens_before=result.tokens_before, + tokens_after=result.tokens_after, + tokens_saved=result.tokens_saved, + modified=result.modified, + strategy_chain=router_chain, + content_type=router_content_type, + text_shape=_codex_ws_text_shape(result.original), + ) + if elapsed_ms >= 1000.0: + logger.info( + "[%s] WS /v1/responses slow compression unit " + "elapsed_ms=%.0f strategy=%s category=%s modified=%s " + "content_type=%s text_shape=%s bytes=%d min_bytes=%d " + "tokens_before=%d tokens_after=%d tokens_saved=%d " + "strategy_chain=%s", + request_id, + elapsed_ms, + result.strategy, + result.reason_category, + result.modified, + router_content_type, + _codex_ws_text_shape(result.original), + result.text_bytes, + result.min_bytes, + result.tokens_before, + result.tokens_after, + result.tokens_saved, + router_chain, + ) + _add_timing("compression_units_router_loop", router_total_started) + + apply_started = time.perf_counter() + for slot, result, _elapsed_ms in routed_results: item_idx, slot_ref = slot router_chain = list(result.router_result.strategy_chain) if result.router_result else [] for s in router_chain: @@ -621,32 +789,33 @@ class OpenAIHandlerMixin: # role-protected, or in a frozen cache_zone don't count. if result.router_result is not None or result.modified: attempted_input_tokens += result.tokens_before - _log( - "codex_compression_unit_result", - item_index=item_idx, - slot=slot_ref, - modified=result.modified, - reason=result.reason, - reason_category=cat, - text_bytes=result.text_bytes, - min_bytes=result.min_bytes, - strategy=result.strategy, - strategy_chain=router_chain, - tokens_before=result.tokens_before, - tokens_after=result.tokens_after, - tokens_saved=result.tokens_saved, - transforms_applied=result.transforms_applied, - router_strategy=( - result.router_result.strategy_used.value if result.router_result else None - ), - router_summary=result.router_result.summary() if result.router_result else None, - router_routing_log=_routing_log_debug(result.router_result), - router_cache_hit=( - result.router_result.cache_hit if result.router_result else False - ), - original=result.original, - compressed=result.compressed, - ) + if debug_enabled: + _log( + "codex_compression_unit_result", + item_index=item_idx, + slot=slot_ref, + modified=result.modified, + reason=result.reason, + reason_category=cat, + text_bytes=result.text_bytes, + min_bytes=result.min_bytes, + strategy=result.strategy, + strategy_chain=router_chain, + tokens_before=result.tokens_before, + tokens_after=result.tokens_after, + tokens_saved=result.tokens_saved, + transforms_applied=result.transforms_applied, + router_strategy=( + result.router_result.strategy_used.value if result.router_result else None + ), + router_summary=result.router_result.summary() if result.router_result else None, + router_routing_log=_routing_log_debug(result.router_result), + router_cache_hit=( + result.router_result.cache_hit if result.router_result else False + ), + original=result.original, + compressed=result.compressed, + ) if not result.modified: continue @@ -659,6 +828,7 @@ class OpenAIHandlerMixin: for transform in result.transforms_applied: if transform not in transforms: transforms.append(transform) + _add_timing("compression_unit_apply_results", apply_started) _log( "codex_compression_payload_result", @@ -687,6 +857,7 @@ class OpenAIHandlerMixin: *, model: str, request_id: str, + timing: dict[str, float] | None = None, ) -> tuple[dict[str, Any], bool, int, list[str], str | None, int, int, int]: """Compress an OpenAI Responses payload through the shared router. @@ -696,7 +867,16 @@ class OpenAIHandlerMixin: compressor, then splices accepted replacements back into the payload. """ + timing_sink: dict[str, float] = timing if timing is not None else {} + + def _add_timing(name: str, started_at: float) -> None: + timing_sink[name] = ( + timing_sink.get(name, 0.0) + (time.perf_counter() - started_at) * 1000.0 + ) + + input_serialization_started = time.perf_counter() input_bytes = json.dumps(payload).encode("utf-8") + _add_timing("compression_input_json_dump", input_serialization_started) # Codex/Responses requests can re-enter this method many times per # request_id (one per turn over the same websocket). Tag every # event in this single pass with a content-derived id so dashboards @@ -707,54 +887,63 @@ class OpenAIHandlerMixin: # Identical pass_ids within one request_id indicate idempotent # retries on the same input bytes and are the only thing that # should be deduped. - pass_id = hashlib.sha256(input_bytes).hexdigest()[:12] - input_context_budget = _openai_responses_context_budget(payload) - _log_codex_compression_debug( - "codex_compression_payload_input", - request_id=request_id, - pass_id=pass_id, - model=model, - input_bytes=len(input_bytes), - context_budget=input_context_budget, - input_top_level_keys=list(payload.keys()), - input_field_type=type(payload.get("input")).__name__, - messages_field_type=type(payload.get("messages")).__name__, - payload=payload, - ) + debug_enabled = _codex_compression_debug_enabled() + pass_id = hashlib.sha256(input_bytes).hexdigest()[:12] if debug_enabled else None + input_context_budget: dict[str, Any] | None = None + if debug_enabled: + input_context_budget = _openai_responses_context_budget(payload) + _log_codex_compression_debug( + "codex_compression_payload_input", + request_id=request_id, + pass_id=pass_id, + model=model, + input_bytes=len(input_bytes), + context_budget=input_context_budget, + input_top_level_keys=list(payload.keys()), + input_field_type=type(payload.get("input")).__name__, + messages_field_type=type(payload.get("messages")).__name__, + payload=payload, + ) working = payload modified = False tokens_saved = 0 transforms: list[str] = [] reason: str | None = None + tool_compaction_started = time.perf_counter() compacted_payload, tools_modified, tools_before_bytes, tools_after_bytes = ( _compact_openai_responses_tools(working) ) + _add_timing("compression_tool_schema_compaction", tool_compaction_started) if tools_modified: working = compacted_payload modified = True reason = None transforms.append("openai:responses:tool_schema_compaction") try: + tool_token_started = time.perf_counter() tokenizer = self.openai_provider.get_token_counter(model) tokens_saved += max( 0, tokenizer.count_text(_json_debug_dumps(payload.get("tools"))) - tokenizer.count_text(_json_debug_dumps(working.get("tools"))), ) + _add_timing("compression_tool_schema_token_count", tool_token_started) except Exception: pass - _log_codex_compression_debug( - "codex_tool_schema_compaction", - request_id=request_id, - pass_id=pass_id, - model=model, - modified=True, - tools_bytes_before=tools_before_bytes, - tools_bytes_after=tools_after_bytes, - tools_bytes_saved=tools_before_bytes - tools_after_bytes, - ) + if debug_enabled: + _log_codex_compression_debug( + "codex_tool_schema_compaction", + request_id=request_id, + pass_id=pass_id, + model=model, + modified=True, + tools_bytes_before=tools_before_bytes, + tools_bytes_after=tools_after_bytes, + tools_bytes_saved=tools_before_bytes - tools_after_bytes, + ) + live_units_started = time.perf_counter() ( router_payload, router_modified, @@ -768,7 +957,9 @@ class OpenAIHandlerMixin: model=model, request_id=request_id, pass_id=pass_id, + timing=timing_sink, ) + _add_timing("compression_live_units_total", live_units_started) if router_modified: working = router_payload modified = True @@ -788,20 +979,31 @@ class OpenAIHandlerMixin: attempted_input_tokens = int(router_attempted_tokens) if tools_modified: try: + attempted_token_started = time.perf_counter() tokenizer = self.openai_provider.get_token_counter(model) attempted_input_tokens += tokenizer.count_text( _json_debug_dumps(payload.get("tools")) ) + _add_timing( + "compression_tool_schema_attempted_token_count", + attempted_token_started, + ) except Exception: pass + dedupe_started = time.perf_counter() deduped: list[str] = [] for transform in transforms: if transform not in deduped: deduped.append(transform) + _add_timing("compression_transform_dedupe", dedupe_started) + output_serialization_started = time.perf_counter() output_bytes = json.dumps(working).encode("utf-8") - output_context_budget = _openai_responses_context_budget(working) + _add_timing("compression_output_json_dump", output_serialization_started) + output_context_budget = ( + _openai_responses_context_budget(working) if debug_enabled else None + ) # One-line summary at INFO — the single event a human reading # logs should scan first to understand "what happened on this # pass". All the verbose per-event debug data stays available @@ -825,41 +1027,42 @@ class OpenAIHandlerMixin: attempted_pct = ( (tokens_saved / attempted_input_tokens) * 100.0 if attempted_input_tokens > 0 else 0.0 ) - _log_codex_compression_debug( - "codex_compression_pass_summary", - request_id=request_id, - pass_id=pass_id, - model=model, - modified=modified, - reason=reason, - input_bytes=len(input_bytes), - output_bytes=len(output_bytes), - bytes_saved=len(input_bytes) - len(output_bytes), - savings_pct=round(savings_pct, 2), - tokens_saved=tokens_saved, - attempted_input_tokens=attempted_input_tokens, - attempted_pct=round(attempted_pct, 2), - strategy_chain=strategy_chain, - units_by_category=units_by_category, - transforms=deduped, - ) - _log_codex_compression_debug( - "codex_compression_payload_output", - request_id=request_id, - pass_id=pass_id, - model=model, - modified=modified, - reason=reason, - tokens_saved=tokens_saved, - attempted_input_tokens=attempted_input_tokens, - transforms=deduped, - input_bytes=len(input_bytes), - output_bytes=len(output_bytes), - context_budget_before=input_context_budget, - context_budget_after=output_context_budget, - input_payload=payload, - output_payload=working, - ) + if debug_enabled: + _log_codex_compression_debug( + "codex_compression_pass_summary", + request_id=request_id, + pass_id=pass_id, + model=model, + modified=modified, + reason=reason, + input_bytes=len(input_bytes), + output_bytes=len(output_bytes), + bytes_saved=len(input_bytes) - len(output_bytes), + savings_pct=round(savings_pct, 2), + tokens_saved=tokens_saved, + attempted_input_tokens=attempted_input_tokens, + attempted_pct=round(attempted_pct, 2), + strategy_chain=strategy_chain, + units_by_category=units_by_category, + transforms=deduped, + ) + _log_codex_compression_debug( + "codex_compression_payload_output", + request_id=request_id, + pass_id=pass_id, + model=model, + modified=modified, + reason=reason, + tokens_saved=tokens_saved, + attempted_input_tokens=attempted_input_tokens, + transforms=deduped, + input_bytes=len(input_bytes), + output_bytes=len(output_bytes), + context_budget_before=input_context_budget, + context_budget_after=output_context_budget, + input_payload=payload, + output_payload=working, + ) return ( working, modified, @@ -877,15 +1080,33 @@ class OpenAIHandlerMixin: *, model: str, request_id: str, - ) -> tuple[dict[str, Any], bool, int, list[str], str | None, int, int, int]: - return await self._run_compression_in_executor( - lambda: self._compress_openai_responses_payload( - payload, - model=model, - request_id=request_id, - ), + ) -> tuple[dict[str, Any], bool, int, list[str], str | None, int, int, int, dict[str, float]]: + timing: dict[str, float] = {} + + def _compress(): # noqa: ANN202 + try: + return self._compress_openai_responses_payload( + payload, + model=model, + request_id=request_id, + timing=timing, + ) + except TypeError as exc: + if "unexpected keyword argument 'timing'" not in str(exc): + raise + return self._compress_openai_responses_payload( + payload, + model=model, + request_id=request_id, + ) + + result = await self._run_compression_in_executor( + _compress, timeout=COMPRESSION_TIMEOUT_SECONDS, ) + if len(result) == 8: + return (*result, timing) + return result async def handle_openai_chat( self, @@ -2366,6 +2587,7 @@ class OpenAIHandlerMixin: _bytes_before, _bytes_after, _attempted_tokens, + _compression_timing, ) = await self._compress_openai_responses_payload_in_executor( body, model=model, @@ -3067,6 +3289,8 @@ class OpenAIHandlerMixin: ws_client_disconnect_seen = False ws_overhead_ms_total = 0.0 ws_recorded_overhead_ms_total = 0.0 + ws_compression_timing_totals: dict[str, float] = {} + ws_recorded_compression_timing_totals: dict[str, float] = {} ws_ttfb_ms: float | None = None ws_recorded_ttfb_ms = False _ws_bypass = self._headroom_bypass_enabled(ws_headers) @@ -3094,6 +3318,29 @@ class OpenAIHandlerMixin: if ws_overhead_ms_total > 0: stage_timer.record("compression", ws_overhead_ms_total) + def _record_ws_compression_timing(name: str, duration_ms: float) -> None: + ws_compression_timing_totals[name] = ( + ws_compression_timing_totals.get(name, 0.0) + + max(0.0, float(duration_ms)) + ) + + def _codex_ws_final_strategies(timing: dict[str, float]) -> list[str]: + prefix = "compression_unit_router_strategy_" + return [ + name.removeprefix(prefix) + for name, ms in timing.items() + if name.startswith(prefix) and ms > 0 + ] + + def _codex_ws_strategy_chain(transforms: list[str]) -> list[str]: + chain: list[str] = [] + for transform in transforms: + if ":" in transform: + continue + if transform not in chain: + chain.append(transform) + return chain + def _current_ws_overhead_ms() -> float: summary = stage_timer.summary() return ws_overhead_ms_total + max(0.0, float(summary.get("memory_context") or 0.0)) @@ -3109,6 +3356,12 @@ class OpenAIHandlerMixin: if ttfb_ms > 0: timing["codex_ws.ttfb"] = ttfb_ms + for stage_name, total_ms in ws_compression_timing_totals.items(): + recorded_ms = ws_recorded_compression_timing_totals.get(stage_name, 0.0) + delta_ms = max(0.0, total_ms - recorded_ms) + if delta_ms > 0: + timing[f"codex_ws.{stage_name}"] = delta_ms + summary = stage_timer.summary() for stage_name in ( "memory_context", @@ -3306,7 +3559,9 @@ class OpenAIHandlerMixin: # anyway so a JSON-shape edge case can never break the WS # session. if self.config.optimize and not _ws_bypass: + _first_frame_compression_elapsed_ms = 0.0 try: + _preflight_started = time.perf_counter() _ws_auth_mode = classify_auth_mode(ws_headers) try: _send_body = json.loads(first_msg_raw) @@ -3320,6 +3575,12 @@ class OpenAIHandlerMixin: _inner = _send_body["response"] if _wrapped else _send_body _model = (_inner.get("model") if isinstance(_inner, dict) else None) or "" + _preflight_ms = (time.perf_counter() - _preflight_started) * 1000.0 + _record_ws_compression_timing( + "compression_preflight_serialization", + _preflight_ms, + ) + _record_ws_compression_overhead(_preflight_ms) _compression_started = time.perf_counter() try: ( @@ -3331,22 +3592,55 @@ class OpenAIHandlerMixin: _bytes_before, _bytes_after, _ws_attempted_tokens, + _ws_compression_timing, ) = await self._compress_openai_responses_payload_in_executor( _inner, model=_model, request_id=request_id, ) + for _timing_name, _timing_ms in _ws_compression_timing.items(): + _record_ws_compression_timing(_timing_name, _timing_ms) finally: + _first_frame_compression_elapsed_ms = ( + time.perf_counter() - _compression_started + ) * 1000.0 + _record_ws_compression_timing( + "compression_executor_wait_run", + _first_frame_compression_elapsed_ms, + ) _record_ws_compression_overhead( - (time.perf_counter() - _compression_started) * 1000.0 + _first_frame_compression_elapsed_ms + ) + record_frame = getattr( + getattr(self, "metrics", None), "record_codex_ws_frame", None + ) + if record_frame is not None: + record_frame( + elapsed_ms=_first_frame_compression_elapsed_ms, + bytes_before=_bytes_before, + bytes_after=_bytes_after, + attempted_tokens=_ws_attempted_tokens, + tokens_saved=_ws_saved, + modified=_modified, + strategy_chain=_codex_ws_strategy_chain(_ws_transforms), + final_strategies=_codex_ws_final_strategies( + _ws_compression_timing + ), ) if _modified: if isinstance(_new_inner, dict): + _rewrite_started = time.perf_counter() if _wrapped: _send_body["response"] = _new_inner else: _send_body = _new_inner first_msg_raw = json.dumps(_send_body) + _rewrite_ms = (time.perf_counter() - _rewrite_started) * 1000.0 + _record_ws_compression_timing( + "compression_payload_rewrite_json_dump", + _rewrite_ms, + ) + _record_ws_compression_overhead(_rewrite_ms) tokens_saved += int(_ws_saved) attempted_input_tokens_total += int(_ws_attempted_tokens) for _t in _ws_transforms: @@ -3363,6 +3657,7 @@ class OpenAIHandlerMixin: _ws_auth_mode.value, transforms_applied, ) + ws_frames_compressed += 1 else: _log_ws_passthrough( _ws_reason or "no_compression", @@ -3379,6 +3674,18 @@ class OpenAIHandlerMixin: frame_type="unknown", ) except Exception as _ce: + if _first_frame_compression_elapsed_ms > 0: + record_frame = getattr( + getattr(self, "metrics", None), "record_codex_ws_frame", None + ) + if record_frame is not None: + record_frame( + elapsed_ms=_first_frame_compression_elapsed_ms, + bytes_before=len( + first_msg_raw.encode("utf-8", errors="replace") + ), + failed=True, + ) logger.warning( f"[{request_id}] WS /v1/responses compression failed; " f"forwarding original frame: {type(_ce).__name__}: {_ce}" @@ -3511,6 +3818,7 @@ class OpenAIHandlerMixin: raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), ) return raw_msg, False, "optimize_disabled" + _preflight_started = time.perf_counter() try: parsed_frame = json.loads(raw_msg) except json.JSONDecodeError: @@ -3547,9 +3855,18 @@ class OpenAIHandlerMixin: frame_type="response.create", ) return raw_msg, False, "invalid_inner_payload" + frame_compression_elapsed_ms = 0.0 try: model_for_frame = inner_payload.get("model") or "" _frame_auth_mode = classify_auth_mode(ws_headers) + _preflight_ms = ( + time.perf_counter() - _preflight_started + ) * 1000.0 + _record_ws_compression_timing( + "compression_preflight_serialization", + _preflight_ms, + ) + _record_ws_compression_overhead(_preflight_ms) _compression_started = time.perf_counter() try: ( @@ -3561,16 +3878,62 @@ class OpenAIHandlerMixin: bytes_before, bytes_after, frame_attempted_tokens, + frame_compression_timing, ) = await self._compress_openai_responses_payload_in_executor( inner_payload, model=model_for_frame, request_id=request_id, ) + for _timing_name, _timing_ms in ( + frame_compression_timing.items() + ): + _record_ws_compression_timing(_timing_name, _timing_ms) finally: + frame_compression_elapsed_ms = ( + time.perf_counter() - _compression_started + ) * 1000.0 + _record_ws_compression_timing( + "compression_executor_wait_run", + frame_compression_elapsed_ms, + ) _record_ws_compression_overhead( - (time.perf_counter() - _compression_started) * 1000.0 + frame_compression_elapsed_ms + ) + record_frame = getattr( + getattr(self, "metrics", None), + "record_codex_ws_frame", + None, + ) + if record_frame is not None: + record_frame( + elapsed_ms=frame_compression_elapsed_ms, + bytes_before=bytes_before, + bytes_after=bytes_after, + attempted_tokens=frame_attempted_tokens, + tokens_saved=frame_saved, + modified=modified, + strategy_chain=_codex_ws_strategy_chain( + frame_transforms + ), + final_strategies=_codex_ws_final_strategies( + frame_compression_timing + ), ) except Exception as _frame_err: + if frame_compression_elapsed_ms > 0: + record_frame = getattr( + getattr(self, "metrics", None), + "record_codex_ws_frame", + None, + ) + if record_frame is not None: + record_frame( + elapsed_ms=frame_compression_elapsed_ms, + bytes_before=len( + raw_msg.encode("utf-8", errors="replace") + ), + failed=True, + ) logger.warning( "[%s] WS /v1/responses frame compression " "failed; forwarding original: %s: %s", @@ -3606,10 +3969,18 @@ class OpenAIHandlerMixin: ) return raw_msg, False, "compressed_payload_not_dict" if wrapped_frame: + _rewrite_started = time.perf_counter() parsed_frame["response"] = new_inner rewritten = json.dumps(parsed_frame) else: + _rewrite_started = time.perf_counter() rewritten = json.dumps(new_inner) + _rewrite_ms = (time.perf_counter() - _rewrite_started) * 1000.0 + _record_ws_compression_timing( + "compression_payload_rewrite_json_dump", + _rewrite_ms, + ) + _record_ws_compression_overhead(_rewrite_ms) tokens_saved += int(frame_saved) attempted_input_tokens_total += int(frame_attempted_tokens) for t in frame_transforms: @@ -3889,6 +4260,9 @@ class OpenAIHandlerMixin: attempted_input_tokens_total ) ws_recorded_overhead_ms_total = _current_ws_overhead_ms() + ws_recorded_compression_timing_totals.update( + ws_compression_timing_totals + ) if ttfb_for_record_ms > 0: ws_recorded_ttfb_ms = True diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 4223f4611..f99170446 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -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( diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 4c5aee0d1..600b485f8 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -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 diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 33d73dd63..2432290d0 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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 diff --git a/headroom/transforms/compression_units.py b/headroom/transforms/compression_units.py index ab97426eb..62ab046b3 100644 --- a/headroom/transforms/compression_units.py +++ b/headroom/transforms/compression_units.py @@ -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): diff --git a/headroom/transforms/content_detector.py b/headroom/transforms/content_detector.py index cca6f0e93..324139726 100644 --- a/headroom/transforms/content_detector.py +++ b/headroom/transforms/content_detector.py @@ -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 ] diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 06539b52a..ef96b487d 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -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" diff --git a/headroom/transforms/kompress_compressor.py b/headroom/transforms/kompress_compressor.py index ba9dc99f9..cf9e724c4 100644 --- a/headroom/transforms/kompress_compressor.py +++ b/headroom/transforms/kompress_compressor.py @@ -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, diff --git a/tests/test_compression_observability.py b/tests/test_compression_observability.py index 89f495aeb..c1f2b5857 100644 --- a/tests/test_compression_observability.py +++ b/tests/test_compression_observability.py @@ -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 diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index 17916a878..5beb90ee7 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -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 diff --git a/tests/test_openai_codex_ws_timings.py b/tests/test_openai_codex_ws_timings.py index d0337798e..8dfb1d099 100644 --- a/tests/test_openai_codex_ws_timings.py +++ b/tests/test_openai_codex_ws_timings.py @@ -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" diff --git a/tests/test_proxy_compression_executor.py b/tests/test_proxy_compression_executor.py index 2c2d78ccc..bd74f22ff 100644 --- a/tests/test_proxy_compression_executor.py +++ b/tests/test_proxy_compression_executor.py @@ -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" diff --git a/tests/test_proxy_dashboard_stats_cache.py b/tests/test_proxy_dashboard_stats_cache.py index d9d970615..e4c0fc467 100644 --- a/tests/test_proxy_dashboard_stats_cache.py +++ b/tests/test_proxy_dashboard_stats_cache.py @@ -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 diff --git a/tests/test_transforms_content_router.py b/tests/test_transforms_content_router.py index bea4fa735..8ee570979 100644 --- a/tests/test_transforms_content_router.py +++ b/tests/test_transforms_content_router.py @@ -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: