mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat: compression safety rails — error-output protection, pipeline circuit breaker, library inflation guard (#851)
Closes #847 ## What Three safety rails, each of which only ever makes compression LESS aggressive — zero behavior change for content that compresses normally: 1. **Error-output protection** (`ContentRouter`) — failed tool calls pass through verbatim on both the OpenAI `role=tool` string path and the Anthropic `tool_result` block path. Triggered by the explicit `is_error: true` flag or the existing Rust error-indicator detector (`headroom._core.content_has_error_indicators`, previously only used for TOIN signatures). Capped by `error_protection_max_chars` (8000, ~2K tokens) so big error-laden CI logs still reach `LogCompressor`, which preserves error lines — the two features stay complementary. `protect_error_outputs=False` disables. 2. **Pipeline circuit breaker** (`TransformPipeline`) — after 3 consecutive transform failures, `apply()` passes messages through untouched for a 60s cooldown instead of re-running (and re-failing) transforms on every request. Env-tunable: `HEADROOM_PIPELINE_BREAKER_THRESHOLD` (0 disables), `HEADROOM_PIPELINE_BREAKER_COOLDOWN_S`. Passthrough results tagged `pipeline:circuit_open`; a clean run closes the breaker. Thread-safe (lock-guarded counters, `time.monotonic`). 3. **Library inflation guard** (`compress()`) — all four proxy handlers already revert when "optimization" inflates tokens; the public library path returned inflated messages as-is. Now mirrors the proxy guard and tags `inflation_guard:reverted`. ## Why Production agent research backs each rail: keeping error outputs verbatim measurably improves agent recovery (Manus context-engineering; JetBrains "Complexity Trap", arXiv:2508.21433); Claude Code added its consecutive-compaction-failure cap after telemetry showed failure loops; the inflation guard closes a library/proxy asymmetry. All three follow CONTRIBUTING's "Safety first: never drop user/assistant content, prefer false negatives." ## Tests `tests/test_compression_safety_rails.py` — 10 tests: - error protection: string path, `is_error` flag (with neutral text proving the flag alone triggers), indicator scan, size-cap fall-through, config-disable - circuit breaker: opens after threshold + passthrough, success resets count, cooldown expiry closes, env-disable - inflation guard: inflated result reverts to originals Regression: `test_transforms_content_router`, `test_pipeline`, `test_compress_api`, `test_compress_failure`, `test_canonical_pipeline`, `test_proxy_pipeline_lifecycle`, `test_observability_*`, `test_compression_policy` — 69 passed. `ruff check` + `ruff format --check` clean. --------- Co-authored-by: Ash Rhodes <ashley.rhodes@king.com>
This commit is contained in:
parent
dc95c6bb00
commit
c0cadccff9
5 changed files with 519 additions and 2 deletions
|
|
@ -252,6 +252,24 @@ def compress(
|
|||
tokens_after = result.tokens_after
|
||||
compressed_messages = result.messages
|
||||
|
||||
# Guard: if "optimization" inflated tokens, revert to originals.
|
||||
# Mirrors the inflation guards in the proxy handlers
|
||||
# (anthropic/openai/gemini/batch) — the library path had none.
|
||||
if tokens_after > tokens_before:
|
||||
logger.warning(
|
||||
"Optimization inflated tokens (%d -> %d); reverting to original messages",
|
||||
tokens_before,
|
||||
tokens_after,
|
||||
)
|
||||
return CompressResult(
|
||||
messages=messages,
|
||||
tokens_before=tokens_before,
|
||||
tokens_after=tokens_before,
|
||||
tokens_saved=0,
|
||||
compression_ratio=0.0,
|
||||
transforms_applied=["inflation_guard:reverted"],
|
||||
)
|
||||
|
||||
routing_markers = summarize_routing_markers(result.transforms_applied)
|
||||
if routing_markers:
|
||||
routed_event = pipeline_extensions.emit(
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ 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
|
||||
from .error_detection import content_has_strong_error_indicators
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -457,6 +458,14 @@ class ContentRouterConfig:
|
|||
protect_recent_code: int = 4 # Don't compress CODE in last N messages (0 = disabled)
|
||||
protect_analysis_context: bool = True # Detect "analyze/review" intent, protect code
|
||||
|
||||
# Protection: failed tool calls / error outputs stay verbatim (issue #847).
|
||||
# The model needs exact tracebacks and error text to recover; compressing
|
||||
# them measurably hurts agent recovery. Outputs above the size cap still
|
||||
# compress — LogCompressor preserves error lines in big logs, so the two
|
||||
# features stay complementary.
|
||||
protect_error_outputs: bool = True
|
||||
error_protection_max_chars: int = 8000 # ~2K tokens; larger errors compress
|
||||
|
||||
# Cache safety: assistant text-block compression.
|
||||
# Default OFF. Assistant content is echoed back by the client in
|
||||
# subsequent turns and becomes part of the upstream provider's
|
||||
|
|
@ -2120,6 +2129,24 @@ class ContentRouter(Transform):
|
|||
route_counts["small"] += 1
|
||||
continue
|
||||
|
||||
# Protection: failed tool calls / error outputs stay verbatim
|
||||
# (issue #847). The model needs exact tracebacks to recover.
|
||||
# Strong (>=2 distinct indicators) match only — a single
|
||||
# keyword false-positives on benign outputs that mention
|
||||
# errors. Above the size cap, fall through — LogCompressor
|
||||
# preserves error lines in big logs.
|
||||
if (
|
||||
self.config.protect_error_outputs
|
||||
and role == "tool"
|
||||
and len(content) <= self.config.error_protection_max_chars
|
||||
and content_has_strong_error_indicators(content)
|
||||
):
|
||||
result_slots[i] = message
|
||||
transforms_applied.append("router:protected:error_output")
|
||||
route_counts.setdefault("error_protected", 0)
|
||||
route_counts["error_protected"] += 1
|
||||
continue
|
||||
|
||||
# Detect content type for protection decisions
|
||||
detection = _detect_content(content)
|
||||
is_code = detection.content_type == ContentType.SOURCE_CODE
|
||||
|
|
@ -2272,6 +2299,8 @@ class ContentRouter(Transform):
|
|||
parts.append(f"{route_counts['analysis_ctx']} protected (analysis ctx)")
|
||||
if route_counts.get("already_compressed"):
|
||||
parts.append(f"{route_counts['already_compressed']} pinned (already compressed)")
|
||||
if route_counts.get("error_protected"):
|
||||
parts.append(f"{route_counts['error_protected']} protected (error output)")
|
||||
if route_counts["ratio_too_high"]:
|
||||
parts.append(f"{route_counts['ratio_too_high']} unchanged (ratio>={min_ratio:.2f})")
|
||||
if route_counts["content_blocks"]:
|
||||
|
|
@ -2455,6 +2484,29 @@ class ContentRouter(Transform):
|
|||
|
||||
tool_content = block.get("content", "")
|
||||
|
||||
# Protection: failed tool calls / error outputs stay verbatim
|
||||
# (issue #847). `is_error` is Anthropic's explicit failure
|
||||
# flag and suffices alone; the indicator scan catches error
|
||||
# text without the flag but requires >=2 distinct keywords
|
||||
# so benign outputs mentioning errors don't skip compression.
|
||||
# Above the size cap, fall through — LogCompressor preserves
|
||||
# error lines in big logs.
|
||||
if (
|
||||
self.config.protect_error_outputs
|
||||
and isinstance(tool_content, str)
|
||||
and len(tool_content) <= self.config.error_protection_max_chars
|
||||
and (
|
||||
block.get("is_error") is True
|
||||
or content_has_strong_error_indicators(tool_content)
|
||||
)
|
||||
):
|
||||
new_blocks.append(block)
|
||||
transforms_applied.append("router:protected:error_output")
|
||||
if route_counts is not None:
|
||||
route_counts.setdefault("error_protected", 0)
|
||||
route_counts["error_protected"] += 1
|
||||
continue
|
||||
|
||||
# Only process string content
|
||||
if isinstance(tool_content, str) and len(tool_content) > min_chars:
|
||||
# Compression pinning: skip already-compressed content
|
||||
|
|
|
|||
|
|
@ -145,6 +145,33 @@ def content_has_error_indicators(text: str) -> bool:
|
|||
return bool(_rust_content_has_error_indicators(text))
|
||||
|
||||
|
||||
def content_has_strong_error_indicators(text: str) -> bool:
|
||||
"""Stricter triage for compression-protection gates.
|
||||
|
||||
:func:`content_has_error_indicators` substring-matches a single
|
||||
keyword, which false-positives on benign outputs that merely
|
||||
mention errors — grep hits, ``"errors": []`` JSON fields,
|
||||
``error_handler.py`` filenames, ``except Exception`` in file
|
||||
reads. Protection gates exempt content from compression entirely,
|
||||
so a lax match there silently costs savings on the hot path.
|
||||
|
||||
Require at least two DISTINCT indicator keywords: genuine failure
|
||||
output nearly always pairs the failure kind with a second
|
||||
indicator (``Traceback`` + ``ValueError``, ``fatal`` +
|
||||
``crash``), while passing mentions rarely do. Misses here are
|
||||
safe — downstream compressors (LogCompressor) still preserve
|
||||
error lines.
|
||||
"""
|
||||
lowered = text.lower()
|
||||
hits = 0
|
||||
for keyword in ERROR_INDICATOR_KEYWORDS:
|
||||
if keyword in lowered:
|
||||
hits += 1
|
||||
if hits >= 2:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ERROR_KEYWORDS",
|
||||
"IMPORTANCE_KEYWORDS",
|
||||
|
|
@ -158,5 +185,6 @@ __all__ = [
|
|||
"PRIORITY_PATTERNS_DIFF",
|
||||
"PRIORITY_PATTERNS_TEXT",
|
||||
"content_has_error_indicators",
|
||||
"content_has_strong_error_indicators",
|
||||
"score_line",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -3,9 +3,12 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from contextlib import nullcontext
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
from ..config import (
|
||||
CacheAlignerConfig,
|
||||
|
|
@ -27,6 +30,24 @@ if TYPE_CHECKING:
|
|||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_N = TypeVar("_N", int, float)
|
||||
|
||||
|
||||
def _breaker_env(name: str, default: _N, cast: Callable[[str], _N]) -> _N:
|
||||
"""Parse a circuit-breaker env var, falling back on bad input.
|
||||
|
||||
The breaker is a safety net — a typo'd value must degrade to the
|
||||
default with a warning, not crash proxy startup.
|
||||
"""
|
||||
raw = os.environ.get(name)
|
||||
if raw is None:
|
||||
return default
|
||||
try:
|
||||
return cast(raw)
|
||||
except ValueError:
|
||||
logger.warning("Invalid %s=%r; using default %s", name, raw, default)
|
||||
return default
|
||||
|
||||
|
||||
class TransformPipeline:
|
||||
"""
|
||||
|
|
@ -65,6 +86,16 @@ class TransformPipeline:
|
|||
else:
|
||||
self.transforms = self._build_default_transforms()
|
||||
|
||||
# Circuit breaker (issue #847): after N consecutive pipeline
|
||||
# failures, pass messages through untouched for a cooldown window
|
||||
# instead of re-running (and re-failing) transforms on every
|
||||
# request. Threshold <= 0 disables the breaker.
|
||||
self._breaker_threshold = _breaker_env("HEADROOM_PIPELINE_BREAKER_THRESHOLD", 3, int)
|
||||
self._breaker_cooldown_s = _breaker_env("HEADROOM_PIPELINE_BREAKER_COOLDOWN_S", 60.0, float)
|
||||
self._breaker_lock = threading.Lock()
|
||||
self._breaker_failures = 0
|
||||
self._breaker_open_until = 0.0
|
||||
|
||||
def _build_default_transforms(self) -> list[Transform]:
|
||||
"""Build default transform pipeline from config."""
|
||||
transforms: list[Transform] = []
|
||||
|
|
@ -134,6 +165,36 @@ class TransformPipeline:
|
|||
|
||||
return self._provider.__class__.__name__.removesuffix("Provider").lower()
|
||||
|
||||
def _breaker_is_open(self) -> bool:
|
||||
"""True while the circuit breaker cooldown window is active."""
|
||||
if self._breaker_threshold <= 0:
|
||||
return False
|
||||
with self._breaker_lock:
|
||||
return time.monotonic() < self._breaker_open_until
|
||||
|
||||
def _breaker_record_failure(self) -> None:
|
||||
"""Count a pipeline failure; open the breaker at the threshold."""
|
||||
if self._breaker_threshold <= 0:
|
||||
return
|
||||
with self._breaker_lock:
|
||||
self._breaker_failures += 1
|
||||
if self._breaker_failures >= self._breaker_threshold:
|
||||
self._breaker_open_until = time.monotonic() + self._breaker_cooldown_s
|
||||
self._breaker_failures = 0
|
||||
logger.warning(
|
||||
"Pipeline circuit breaker OPEN after %d consecutive failures; "
|
||||
"passing messages through for %.0fs",
|
||||
self._breaker_threshold,
|
||||
self._breaker_cooldown_s,
|
||||
)
|
||||
|
||||
def _breaker_record_success(self) -> None:
|
||||
"""Reset the consecutive-failure count after a clean run."""
|
||||
if self._breaker_threshold <= 0:
|
||||
return
|
||||
with self._breaker_lock:
|
||||
self._breaker_failures = 0
|
||||
|
||||
def apply(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
|
|
@ -168,6 +229,16 @@ class TransformPipeline:
|
|||
)
|
||||
|
||||
# Start with original tokens
|
||||
# Circuit breaker open — pass through untouched (issue #847).
|
||||
if self._breaker_is_open():
|
||||
passthrough_tokens = tokenizer.count_messages(messages)
|
||||
return TransformResult(
|
||||
messages=messages,
|
||||
tokens_before=passthrough_tokens,
|
||||
tokens_after=passthrough_tokens,
|
||||
transforms_applied=["pipeline:circuit_open"],
|
||||
)
|
||||
|
||||
t_count = time.perf_counter()
|
||||
tokens_before = tokenizer.count_messages(messages)
|
||||
count_ms = (time.perf_counter() - t_count) * 1000
|
||||
|
|
@ -248,7 +319,11 @@ class TransformPipeline:
|
|||
with transform_span_context as transform_span:
|
||||
# Time the transform
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
result = transform.apply(current_messages, tokenizer, **kwargs)
|
||||
except Exception:
|
||||
self._breaker_record_failure()
|
||||
raise
|
||||
duration_ms = (time.perf_counter() - t0) * 1000
|
||||
|
||||
# Update messages for next transform
|
||||
|
|
@ -316,6 +391,9 @@ class TransformPipeline:
|
|||
)
|
||||
)
|
||||
|
||||
# All transforms ran without raising — reset the breaker.
|
||||
self._breaker_record_success()
|
||||
|
||||
# Single final token count — the only full recount in the pipeline.
|
||||
# Earlier per-transform counts come from each transform's own result.
|
||||
t_final_count = time.perf_counter()
|
||||
|
|
|
|||
341
tests/test_compression_safety_rails.py
Normal file
341
tests/test_compression_safety_rails.py
Normal file
|
|
@ -0,0 +1,341 @@
|
|||
"""Compression safety rails (issue #847).
|
||||
|
||||
Three rails, each of which only ever makes compression LESS aggressive:
|
||||
|
||||
1. Error-output protection — failed tool calls / error outputs pass
|
||||
through ``ContentRouter`` verbatim (string path and content-block path,
|
||||
including Anthropic ``is_error: true``), capped by
|
||||
``error_protection_max_chars`` so big error-laden logs still reach
|
||||
``LogCompressor`` (which preserves error lines).
|
||||
2. Pipeline circuit breaker — after N consecutive pipeline failures,
|
||||
``TransformPipeline.apply`` passes messages through untouched for a
|
||||
cooldown window instead of re-running failing transforms.
|
||||
3. Library inflation guard — ``headroom.compress()`` reverts to the
|
||||
original messages when "optimization" inflated tokens, mirroring the
|
||||
proxy handlers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom import OpenAIProvider, Tokenizer
|
||||
from headroom.compress import compress
|
||||
from headroom.config import HeadroomConfig, TransformResult
|
||||
from headroom.tokenizer import Tokenizer as TokenizerType
|
||||
from headroom.transforms.base import Transform
|
||||
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
|
||||
from headroom.transforms.pipeline import TransformPipeline
|
||||
|
||||
# ``headroom.compress`` the submodule is shadowed by the function of the
|
||||
# same name re-exported in ``headroom/__init__.py``.
|
||||
compress_module = importlib.import_module("headroom.compress")
|
||||
|
||||
_provider = OpenAIProvider()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tokenizer() -> Tokenizer:
|
||||
return Tokenizer(_provider.get_token_counter("gpt-4o"), "gpt-4o")
|
||||
|
||||
|
||||
# A realistic failed-tool-call output: error indicators, > min_tokens (50),
|
||||
# well under the 8000-char protection cap.
|
||||
_TRACEBACK = (
|
||||
"Traceback (most recent call last):\n"
|
||||
+ "".join(
|
||||
f' File "/app/services/worker_{i}.py", line {i * 17}, in handle_request\n'
|
||||
f" result = downstream.dispatch(payload, retries={i})\n"
|
||||
for i in range(12)
|
||||
)
|
||||
+ "ValueError: connection refused while dispatching payload to upstream "
|
||||
"service after 3 retries; check that the worker pool is initialized "
|
||||
"before the scheduler starts accepting jobs\n"
|
||||
)
|
||||
|
||||
# Error text with no error-indicator keywords — only the explicit
|
||||
# Anthropic ``is_error`` flag marks it as a failure.
|
||||
_NEUTRAL_TOOL_OUTPUT = (
|
||||
"The operation finished without producing the expected artifact. "
|
||||
"Output directory listing follows.\n"
|
||||
+ "\n".join(f"entry_{i}.txt 4096 bytes" for i in range(80))
|
||||
)
|
||||
|
||||
|
||||
# Benign outputs that merely MENTION errors — exactly one distinct
|
||||
# indicator keyword ("error"). A lax substring gate would exempt these
|
||||
# from compression (savings regression); the strong gate must not.
|
||||
_BENIGN_GREP_OUTPUT = (
|
||||
"src/error_handler.py:12:def handle_error(code):\n"
|
||||
"src/error_handler.py:48: log_error(code, context)\n"
|
||||
+ "\n".join(
|
||||
f"src/module_{i}.py:{i * 3}: error_count = metrics.get('error', 0)" for i in range(20)
|
||||
)
|
||||
)
|
||||
|
||||
_BENIGN_JSON_OUTPUT = (
|
||||
'{"status": "completed", "errors": [], "warnings": [], "items": ['
|
||||
+ ", ".join(f'{{"id": {i}, "name": "artifact_{i}", "size": {i * 1024}}}' for i in range(30))
|
||||
+ "]}"
|
||||
)
|
||||
|
||||
|
||||
def _filler_messages(n: int = 2) -> list[dict[str, Any]]:
|
||||
return [{"role": "user", "content": f"step {i}: please continue the task"} for i in range(n)]
|
||||
|
||||
|
||||
class TestErrorOutputProtection:
|
||||
def test_string_tool_message_with_error_protected(self, tokenizer: Tokenizer) -> None:
|
||||
router = ContentRouter()
|
||||
messages = _filler_messages() + [
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": _TRACEBACK},
|
||||
{"role": "user", "content": "what went wrong?"},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
assert "router:protected:error_output" in result.transforms_applied
|
||||
tool_msgs = [m for m in result.messages if m.get("role") == "tool"]
|
||||
assert tool_msgs[0]["content"] == _TRACEBACK
|
||||
|
||||
def test_tool_result_block_with_is_error_flag_protected(self, tokenizer: Tokenizer) -> None:
|
||||
router = ContentRouter()
|
||||
messages = _filler_messages() + [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_1",
|
||||
"is_error": True,
|
||||
"content": _NEUTRAL_TOOL_OUTPUT,
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
assert "router:protected:error_output" in result.transforms_applied
|
||||
block = result.messages[-1]["content"][0]
|
||||
assert block["content"] == _NEUTRAL_TOOL_OUTPUT
|
||||
|
||||
def test_tool_result_block_with_error_indicators_protected(self, tokenizer: Tokenizer) -> None:
|
||||
router = ContentRouter()
|
||||
messages = _filler_messages() + [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_2",
|
||||
"content": _TRACEBACK,
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
assert "router:protected:error_output" in result.transforms_applied
|
||||
block = result.messages[-1]["content"][0]
|
||||
assert block["content"] == _TRACEBACK
|
||||
|
||||
def test_single_indicator_string_output_not_protected(self, tokenizer: Tokenizer) -> None:
|
||||
"""Grep-style output mentioning "error" must not skip compression."""
|
||||
router = ContentRouter()
|
||||
messages = _filler_messages() + [
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": _BENIGN_GREP_OUTPUT},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
assert "router:protected:error_output" not in result.transforms_applied
|
||||
|
||||
def test_single_indicator_block_not_protected_without_flag(self, tokenizer: Tokenizer) -> None:
|
||||
"""`"errors": []` JSON without `is_error` must not skip compression."""
|
||||
router = ContentRouter()
|
||||
messages = _filler_messages() + [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_3",
|
||||
"content": _BENIGN_JSON_OUTPUT,
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
assert "router:protected:error_output" not in result.transforms_applied
|
||||
|
||||
def test_is_error_flag_alone_protects_single_indicator_block(
|
||||
self, tokenizer: Tokenizer
|
||||
) -> None:
|
||||
"""The explicit `is_error` flag needs no indicator corroboration."""
|
||||
router = ContentRouter()
|
||||
messages = _filler_messages() + [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_4",
|
||||
"is_error": True,
|
||||
"content": _BENIGN_JSON_OUTPUT,
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
assert "router:protected:error_output" in result.transforms_applied
|
||||
block = result.messages[-1]["content"][0]
|
||||
assert block["content"] == _BENIGN_JSON_OUTPUT
|
||||
|
||||
def test_oversized_error_output_falls_through(self, tokenizer: Tokenizer) -> None:
|
||||
config = ContentRouterConfig(error_protection_max_chars=100)
|
||||
router = ContentRouter(config=config)
|
||||
messages = _filler_messages() + [
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": _TRACEBACK},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
assert "router:protected:error_output" not in result.transforms_applied
|
||||
|
||||
def test_protection_disabled_via_config(self, tokenizer: Tokenizer) -> None:
|
||||
config = ContentRouterConfig(protect_error_outputs=False)
|
||||
router = ContentRouter(config=config)
|
||||
messages = _filler_messages() + [
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": _TRACEBACK},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_1",
|
||||
"is_error": True,
|
||||
"content": _TRACEBACK,
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
assert "router:protected:error_output" not in result.transforms_applied
|
||||
|
||||
|
||||
class _FailingTransform(Transform):
|
||||
name = "always_fails"
|
||||
|
||||
def apply(
|
||||
self, messages: list[dict[str, Any]], tokenizer: TokenizerType, **kwargs: Any
|
||||
) -> TransformResult:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
|
||||
class _FlakyTransform(Transform):
|
||||
"""Fails for the first ``fail_times`` calls, then succeeds."""
|
||||
|
||||
name = "flaky"
|
||||
|
||||
def __init__(self, fail_times: int) -> None:
|
||||
self.fail_times = fail_times
|
||||
self.calls = 0
|
||||
|
||||
def apply(
|
||||
self, messages: list[dict[str, Any]], tokenizer: TokenizerType, **kwargs: Any
|
||||
) -> TransformResult:
|
||||
self.calls += 1
|
||||
if self.calls <= self.fail_times:
|
||||
raise RuntimeError("boom")
|
||||
tokens = tokenizer.count_messages(messages)
|
||||
return TransformResult(
|
||||
messages=messages,
|
||||
tokens_before=tokens,
|
||||
tokens_after=tokens,
|
||||
transforms_applied=[],
|
||||
)
|
||||
|
||||
|
||||
_MESSAGES = [{"role": "user", "content": "hello there, please summarize the build log"}]
|
||||
|
||||
|
||||
class TestPipelineCircuitBreaker:
|
||||
def test_opens_after_threshold_and_passes_through(self) -> None:
|
||||
pipeline = TransformPipeline(HeadroomConfig(), transforms=[_FailingTransform()])
|
||||
for _ in range(3):
|
||||
with pytest.raises(RuntimeError):
|
||||
pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
result = pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
assert result.transforms_applied == ["pipeline:circuit_open"]
|
||||
assert result.messages == _MESSAGES
|
||||
assert result.tokens_before == result.tokens_after
|
||||
|
||||
def test_success_resets_consecutive_failures(self) -> None:
|
||||
flaky = _FlakyTransform(fail_times=2)
|
||||
pipeline = TransformPipeline(HeadroomConfig(), transforms=[flaky])
|
||||
for _ in range(2):
|
||||
with pytest.raises(RuntimeError):
|
||||
pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
# Third call succeeds — resets the consecutive-failure count.
|
||||
pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
# Two more failures still don't reach the threshold of 3.
|
||||
flaky.fail_times = flaky.calls + 2
|
||||
for _ in range(2):
|
||||
with pytest.raises(RuntimeError):
|
||||
pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
result = pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
assert result.transforms_applied != ["pipeline:circuit_open"]
|
||||
|
||||
def test_cooldown_expiry_closes_breaker(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PIPELINE_BREAKER_COOLDOWN_S", "0.05")
|
||||
flaky = _FlakyTransform(fail_times=3)
|
||||
pipeline = TransformPipeline(HeadroomConfig(), transforms=[flaky])
|
||||
for _ in range(3):
|
||||
with pytest.raises(RuntimeError):
|
||||
pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
assert pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024).transforms_applied == [
|
||||
"pipeline:circuit_open"
|
||||
]
|
||||
time.sleep(0.1)
|
||||
result = pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
assert result.transforms_applied != ["pipeline:circuit_open"]
|
||||
|
||||
def test_invalid_env_values_fall_back_to_defaults(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Typo'd breaker env vars must not crash proxy startup."""
|
||||
monkeypatch.setenv("HEADROOM_PIPELINE_BREAKER_THRESHOLD", "three")
|
||||
monkeypatch.setenv("HEADROOM_PIPELINE_BREAKER_COOLDOWN_S", "1m")
|
||||
pipeline = TransformPipeline(HeadroomConfig(), transforms=[_FailingTransform()])
|
||||
assert pipeline._breaker_threshold == 3
|
||||
assert pipeline._breaker_cooldown_s == 60.0
|
||||
|
||||
def test_disabled_via_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PIPELINE_BREAKER_THRESHOLD", "0")
|
||||
pipeline = TransformPipeline(HeadroomConfig(), transforms=[_FailingTransform()])
|
||||
for _ in range(5):
|
||||
with pytest.raises(RuntimeError):
|
||||
pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
# Breaker never opens — failures keep propagating.
|
||||
with pytest.raises(RuntimeError):
|
||||
pipeline.apply(_MESSAGES, model="gpt-4o", model_limit=1024)
|
||||
|
||||
|
||||
class _InflatingPipeline:
|
||||
"""Fake pipeline whose 'optimization' makes messages bigger."""
|
||||
|
||||
def apply(self, messages: list[dict[str, Any]], **kwargs: Any) -> TransformResult:
|
||||
bloated = [{**m, "content": str(m.get("content", "")) + " PADDING" * 50} for m in messages]
|
||||
return TransformResult(
|
||||
messages=bloated,
|
||||
tokens_before=100,
|
||||
tokens_after=250,
|
||||
transforms_applied=["fake:inflate"],
|
||||
)
|
||||
|
||||
|
||||
class TestLibraryInflationGuard:
|
||||
def test_inflated_result_reverts_to_originals(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(compress_module, "_pipeline", _InflatingPipeline())
|
||||
messages = [{"role": "user", "content": "compress this message please"}]
|
||||
result = compress(messages, model="gpt-4o")
|
||||
assert result.messages == messages
|
||||
assert result.transforms_applied == ["inflation_guard:reverted"]
|
||||
assert result.tokens_saved == 0
|
||||
assert result.compression_ratio == 0.0
|
||||
Loading…
Add table
Add a link
Reference in a new issue