diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 5e7f26b3e..e9e07912c 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -78,6 +78,16 @@ _detect_panic_warned = False _detect_native_unhealthy = False # circuit breaker: native detect hung once (#575) +def _compression_deadline_seconds() -> float: + try: + return max( + 0.0, + float(os.environ.get("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")) / 1000.0, + ) + except ValueError: + return 20.0 + + def _router_debug_dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False, default=str, separators=(",", ":")) @@ -3723,8 +3733,48 @@ class ContentRouter(Transform): task_results = [] for _, task_content, task_ctx, task_bias, _, _ in pending_tasks: t0 = time.perf_counter() - r = self.compress(task_content, context=task_ctx, bias=task_bias) - task_results.append((r, (time.perf_counter() - t0) * 1000)) + deadline_s = _compression_deadline_seconds() if len(pending_tasks) == 1 else 0.0 + if deadline_s: + box: dict[str, Any] = {} + + def _run( + _box: dict[str, Any] = box, + _content: str = task_content, + _context: str = task_ctx, + _bias: float = task_bias, + ) -> None: + try: + _box["result"] = self.compress( + _content, context=_context, bias=_bias + ) + except BaseException as exc: # noqa: BLE001 + _box["error"] = exc + + # ponytail: daemon watchdog cannot stop native GIL holds; native layer owns that fix. + worker = threading.Thread( + target=_run, name="headroom-single-compress-watchdog", daemon=True + ) + worker.start() + worker.join(deadline_s) + if worker.is_alive(): + logger.warning( + "ContentRouter single-cache-miss compression exceeded %.1fs; " + "failing open via PASSTHROUGH", + deadline_s, + ) + r = RouterCompressionResult( + compressed=task_content, + original=task_content, + strategy_used=CompressionStrategy.PASSTHROUGH, + ) + elif "error" in box: + raise box["error"] + else: + r = box["result"] + else: + r = self.compress(task_content, context=task_ctx, bias=task_bias) + compress_ms = (time.perf_counter() - t0) * 1000 + task_results.append((r, compress_ms)) else: # Parallel compression via thread pool with ThreadPoolExecutor(max_workers=max_workers) as executor: diff --git a/pyproject.toml b/pyproject.toml index b48a5cdb8..86ef1a514 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,7 @@ dependencies = [ # ImportError-guarded. Marking it 3.14-optional lets headroom install on Python 3.14 # (core compression + the Anthropic proxy path never import litellm). See GH #956. "litellm>=1.86.2,<2.0; python_version < '3.14'", # model registry, pricing, providers (lazy) - "click>=8.1.0", # CLI framework + "click>=8.3.3", # CLI framework "rich>=13.0.0", # Rich terminal output "opentelemetry-api>=1.24.0", # Safe no-op OTEL API for instrumentation "ast-grep-cli>=0.30.0", # AST-aware code slicing (CodeCompressor); binary wheel @@ -165,7 +165,7 @@ relevance = [ # `headroom/image/compressor.py` adapts both API shapes at runtime via # a try/except cascade. See issue #372 for context. image = [ - "pillow>=10.0.0", + "pillow>=12.3.0", "sentencepiece>=0.1.99", # Required by SigLIP tokenizer (SiglipTokenizer) # Python 3.6–3.12: keep the proven ORT-bundled package directly. # ~15 MB ONNX models auto-downloaded on first use. diff --git a/tests/test_content_router_single_item_deadline.py b/tests/test_content_router_single_item_deadline.py new file mode 100644 index 000000000..258468f83 --- /dev/null +++ b/tests/test_content_router_single_item_deadline.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import time + +from headroom.transforms.content_detector import ContentType +from headroom.transforms.content_router import ( + CompressionStrategy, + ContentRouter, + ContentRouterConfig, + RouterCompressionResult, + RoutingDecision, +) + + +class _Tokenizer: + def count_text(self, content: str) -> int: + return len(content.split()) + + +def _compression_result(content: str, compressed: str) -> RouterCompressionResult: + return RouterCompressionResult( + compressed=compressed, + original=content, + strategy_used=CompressionStrategy.TEXT, + routing_log=[ + RoutingDecision( + content_type=ContentType.PLAIN_TEXT, + strategy=CompressionStrategy.TEXT, + original_tokens=len(content.split()), + compressed_tokens=len(compressed.split()), + ) + ], + ) + + +def _router() -> ContentRouter: + return ContentRouter( + ContentRouterConfig( + protect_recent_code=0, + protect_analysis_context=False, + skip_user_messages=False, + ) + ) + + +def _messages() -> list[dict[str, str]]: + return [ + {"role": "assistant", "content": "frozen prefix content remains unchanged"}, + { + "role": "assistant", + "content": "pending cache miss content takes the inline compression branch", + }, + ] + + +def test_single_cache_miss_fails_open_at_deadline(monkeypatch, caplog): + router = _router() + + def slow_compress(content, *, context="", bias=1.0): + time.sleep(0.2) + return _compression_result(content, "compressed output") + + monkeypatch.setattr(router, "compress", slow_compress) + monkeypatch.setenv("HEADROOM_COMPRESSION_DEADLINE_MS", "10") + + started = time.perf_counter() + result = router.apply( + _messages(), + _Tokenizer(), + frozen_message_count=1, + min_tokens_to_compress=1, + ) + + assert time.perf_counter() - started < 0.12 + assert result.messages[1]["content"] == _messages()[1]["content"] + assert "failing open via PASSTHROUGH" in caplog.text + + +def test_single_cache_miss_preserves_under_deadline_output(monkeypatch): + router = _router() + monkeypatch.setattr( + router, + "compress", + lambda content, *, context="", bias=1.0: _compression_result(content, "compressed output"), + ) + monkeypatch.setenv("HEADROOM_COMPRESSION_DEADLINE_MS", "1000") + + result = router.apply( + _messages(), + _Tokenizer(), + frozen_message_count=1, + min_tokens_to_compress=1, + ) + + assert result.messages[1]["content"] == "compressed output" + + +def test_single_cache_miss_preserves_disabled_deadline(monkeypatch): + router = _router() + monkeypatch.setattr( + router, + "compress", + lambda content, *, context="", bias=1.0: _compression_result(content, "compressed output"), + ) + monkeypatch.setenv("HEADROOM_COMPRESSION_DEADLINE_MS", "0") + + result = router.apply( + _messages(), + _Tokenizer(), + frozen_message_count=1, + min_tokens_to_compress=1, + ) + + assert result.messages[1]["content"] == "compressed output" diff --git a/uv.lock b/uv.lock index 6a72283ed..c5fa3fbd3 100644 --- a/uv.lock +++ b/uv.lock @@ -1737,7 +1737,7 @@ requires-dist = [ { name = "ast-grep-cli", specifier = ">=0.30.0" }, { name = "boto3", marker = "extra == 'bedrock'", specifier = ">=1.41.0" }, { name = "botocore", extras = ["crt"], marker = "extra == 'bedrock'", specifier = ">=1.41.0" }, - { name = "click", specifier = ">=8.1.0" }, + { name = "click", specifier = ">=8.3.3" }, { name = "datasets", marker = "extra == 'evals'", specifier = ">=2.14.0" }, { name = "datasets", marker = "extra == 'voice-train'", specifier = ">=2.14.0" }, { name = "fastapi", marker = "extra == 'dev'", specifier = ">=0.100.0" }, @@ -1785,7 +1785,7 @@ requires-dist = [ { name = "opentelemetry-sdk", marker = "extra == 'dev'", specifier = ">=1.24.0" }, { name = "opentelemetry-sdk", marker = "extra == 'otel'", specifier = ">=1.24.0" }, { name = "orjson", marker = "platform_python_implementation != 'PyPy' and extra == 'proxy'", specifier = ">=3.9.14" }, - { name = "pillow", marker = "extra == 'image'", specifier = ">=10.0.0" }, + { name = "pillow", marker = "extra == 'image'", specifier = ">=12.3.0" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.0.0" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" },