mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): cold-start fast pass — defer only Kompress, not the whole pipeline (#2073)
## Description Since #1850, the freeze path forwards a session's provider-cached prefix byte-identical — so a session is permanently locked to whatever form its cold start put in the provider cache. That fix is correct (it stopped token-mode cache busting measured at +41% cost), but it interacts badly with off-path background compression (#1171): when `HEADROOM_BACKGROUND_COMPRESSION=1` defers a cold-start-large request (frozen=0, ≥50k tokens), the ENTIRE pipeline is deferred and the raw transcript is forwarded, cached, and frozen. The background job's results can never be applied afterward (doing so would rewrite the frozen prefix), so the session forfeits its compression savings for its lifetime. Field data (same day, same session, A/B across a version boundary): ~15k tokens/turn saved when the cold start compressed synchronously vs 0/turn forever when it deferred. Notably, the recurring savings came from `read_lifecycle` stale-read drops completing in ~300ms — deferral throws away sub-second lossless wins to avoid a 30s Kompress pass. Only the Kompress ML stage can blow the request budget (the #1171 cascade). This PR splits the two: - The deferral branch now runs the pipeline synchronously with a new `skip_kompress=True` per-call kwarg — everything except the ML stage — under a bounded budget, and forwards the pruned form. The provider caches (and #1850 freezes) the *compressed* transcript, so the cheap savings persist for the session's lifetime. - The full pipeline (Kompress included) still goes to the background job, unchanged, keyed against the original messages so its content-hash results remain reusable at future cache-miss boundaries. - Fail-open: on fast-pass timeout or error, the request forwards uncompressed exactly as before this change. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/transforms/content_router.py`: new per-call `skip_kompress` runtime kwarg (follows the existing `_runtime_force_kompress` pattern). Gates only the Kompress deep-path call site; units routed there take the identical fallback used when the model isn't ready. Wins over `force_kompress`. - `headroom/proxy/helpers.py`: `COLD_START_FAST_PASS_TIMEOUT_SECONDS` (env `HEADROOM_COLD_START_FAST_PASS_TIMEOUT_SECONDS`, default 10s), documented next to `COMPRESSION_TIMEOUT_SECONDS`. - `headroom/proxy/handlers/anthropic.py`: the background-deferral branch runs the fast pass synchronously, stores its result in the session `CompressionCache`, forwards the pruned messages, and tags `deferred:kompress_background` (or `deferred:dropped` when the enqueue was dropped). On failure it constructs the same `_DeferredCompressionResult` as before. The Anthropic handler is the only deferral site (OpenAI/Gemini handlers don't defer). - `tests/test_transforms/test_content_router.py`: `skip_kompress` never invokes the ML stage and wins over `force_kompress` (mirrors the existing `force_kompress` test). - `tests/test_cold_start_fast_pass.py`: handler-level tests — exactly one synchronous `skip_kompress=True` pass, the background job runs the full pipeline, the forwarded body carries the fast-pass form, fast-pass results land in the compression cache; and the fail-open path (executor timeout → original messages forwarded, background job still queued). - `CHANGELOG.md`: Bug Fixes entry. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uv run --frozen --extra dev pytest tests/test_cold_start_fast_pass.py -v tests/test_cold_start_fast_pass.py::test_cold_start_runs_fast_pass_and_defers_only_kompress PASSED tests/test_cold_start_fast_pass.py::test_fast_pass_failure_falls_back_to_full_deferral PASSED ============================== 2 passed in 0.28s =============================== $ uv run --frozen --extra dev pytest tests/test_proxy/test_background_compression.py \ tests/test_proxy/test_phase3_byte_identity.py tests/test_anthropic_stage_timings.py \ tests/test_transforms/test_content_router.py tests/test_anthropic_pre_upstream_backpressure.py ======================== 90 passed, 1 warning in 10.47s ======================== $ uv run --frozen --extra dev mypy headroom/proxy/handlers/anthropic.py headroom/proxy/helpers.py headroom/transforms/content_router.py Success: no issues found in 3 source files $ ruff check <changed files> && ruff format --check <changed files> All checks passed! / 5 files already formatted ``` ## Real Behavior Proof - Environment: macOS 15 (darwin 24.6.0), Python 3.10 via `uv run --frozen --extra dev`; field logs from a production desktop deployment (Python 3.12, `HEADROOM_MODE=token`, `HEADROOM_BACKGROUND_COMPRESSION=1`, subscription auth policy). - Exact command / steps: compared per-request PERF log lines for the same Claude Code session served by 0.30.0-lineage (sync cold start) vs 0.31.0-lineage (deferred cold start) on the same day. - Observed result: deferred-cold-start sessions log `tok_saved=0` on every subsequent turn with `Pipeline: freezing first 281/284 messages`; sync-cold-start sessions log `tok_saved=15526-18791` per turn with `read_lifecycle:stale` transforms at `opt_ms≈300`. - Not tested: this patch has not run against a live proxy yet (behavior verified at the handler-test level); `ruff`/`mypy` scoped to changed files. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A — proxy pipeline change, no UI. ## Additional Notes - Companion to #2057 (nested tool_result image token counting) and #2058 (new-content-relative savings rate) — all three came out of the same investigation into near-zero reported savings on long 1M-context Claude Code sessions. - Deliberate scope cuts: the OpenAI/Gemini handlers don't have a deferral branch, so nothing to change there; the background job is left keyed to original messages (not the fast-pass output) so its cached results match client-resent bytes at future cache-miss boundaries. - Timeout leak caveat is documented in code: a fast-pass timeout briefly leaks an executor worker, but without the ML stage the pass is bounded by routing + statistical crushers (observed 5-8s worst case on multi-M-token counted transcripts). 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Tejas Chopra <chopratejas@gmail.com> Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
fa330f3e2b
commit
fd9ddaa238
6 changed files with 434 additions and 7 deletions
|
|
@ -109,6 +109,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
|
||||
### Bug Fixes
|
||||
|
||||
* **proxy:** run a cold-start fast pass before background-compression deferral so byte-identical freeze doesn't lock sessions to the uncompressed transcript. Since #1850, a session's provider-cached prefix is frozen in whatever form its cold start forwarded; deferring the WHOLE pipeline (`HEADROOM_BACKGROUND_COMPRESSION=1`, frozen=0, ≥50k tokens) therefore cached the raw transcript and forfeited the session's compression savings for its lifetime — including sub-second lossless wins like `read_lifecycle` stale-read drops, observed in the field as sessions permanently stuck at 0 savings. The deferral branch now runs the pipeline synchronously with the new `skip_kompress=True` kwarg (everything except the Kompress ML stage — the only stage that can blow the request budget per #1171) under a bounded fast-pass budget (`HEADROOM_COLD_START_FAST_PASS_TIMEOUT_SECONDS`, default 10s), forwards the pruned form, and defers only Kompress to the background job (tagged `deferred:kompress_background`). Fail-open: on fast-pass timeout/error the request forwards uncompressed exactly as before. Units routed to Kompress under `skip_kompress` take the same fallback as when the model isn't ready.
|
||||
* **ccr:** detect `read_lifecycle` stale/superseded markers in the retrieve-tool injector so they stay redeemable. Those markers (`[Read content stale: … Retrieve original: hash=<hash>]`) store the original bytes in the CCR store under a valid hash, but none of `CCRToolInjector`'s patterns matched them — every pattern required the word "compressed" or the `<<ccr:` form. So on a frozen-prefix turn (both `read_lifecycle` and prefix freezing are on by default) the injector reported no compressed content, the `headroom_retrieve` tool was not injected, and the model was handed a marker advertising `Retrieve original: hash=X` with no tool to redeem it — silent data loss for stale reads, where retrieval is the only way to recover the original-at-read-time content (the exact case the #1006 guard exists to prevent). Added a pattern matching the load-bearing `Retrieve original: hash=` phrase, aligning the injector with the sibling `read_maturation` marker that was already (incidentally) detected.
|
||||
* **proxy:** strip output-only content blocks from request messages before forwarding. Anthropic's server-side refusal-fallback feature (`server-side-fallback-2026-06-01`) emits a `{"type":"fallback","from":{...},"to":{...}}` block inside the assistant response to signal that a refused request was re-served by the fallback model. That block is valid on the *response* path but rejected on the *request* path, so when a client replays the assistant turn the next request 400s (`invalid_request_error: messages.N.content.0: Input tag 'fallback' ...`) and the conversation gets permanently stuck through the proxy. `read_request_json_with_bytes` (Anthropic/OpenAI/Bedrock) and `_read_request_json` (Gemini) now drop such blocks — re-encoding the raw bytes so byte-faithful passthrough cannot leak the pre-strip body, backfilling a benign text block if a turn is emptied, and leaving requests without such blocks byte-identical (no cache churn).
|
||||
* **wrap/doctor:** make the Claude Remote Control gate warning accurate and stop it firing for users who never had the feature ([#1779](https://github.com/headroomlabs-ai/headroom/issues/1779)). Claude Code 2.1.196 added a client-side check that **deterministically** disables first-party Remote Control (`/remote-control` / `/rc`) whenever `ANTHROPIC_BASE_URL` points at a non-`api.anthropic.com` host — which Headroom always does. The old notice hedged ("may hide the Remote Control menu"); it now states the disable as fact, names the `/rc` command, and detects the installed Claude Code version so the wording is exact (`2.1.196` when known, `2.1.196+` when not). The gate is upstream and RC's control-plane talks to `claude.ai` (not the API host), so Headroom cannot restore it — the warning tells you to run Claude without Headroom for RC sessions. The warning is suppressed for auth modes that never had Remote Control (API-key/PAYG via `ANTHROPIC_API_KEY`/`ANTHROPIC_AUTH_TOKEN`, and Bedrock/Vertex/Foundry cloud IAM) and on Claude Code builds older than 2.1.196 where RC is unaffected by a custom base URL. Both the `headroom wrap claude` launch banner and `headroom doctor` co-report the sibling base-URL gates Headroom *does* restore — on-demand tool loading (#746, automatic) and the 1M context window (#1158, via `--1m`) — and the wrap-side co-report is session-accurate: it says "already restored via --1m" when the flag is in effect and reports tool deferral OFF (not falsely "kept on") when the user chose `--tool-search false`/`ENABLE_TOOL_SEARCH=false`; the `ENABLE_TOOL_SEARCH=...` banner line got the same accuracy fix. `is_custom_anthropic_base_url` now recognizes scheme-less values (`myproxy.local:8080`, `127.0.0.1:8787`) as custom hosts and degrades gracefully on malformed URLs instead of crashing `doctor`. `doctor` resolves the Claude Code version lazily, so runs with no custom base URL never pay the `claude --version` subprocess. No request bytes are touched (cache-safe); this is UX/notice-only.
|
||||
|
|
|
|||
|
|
@ -1314,16 +1314,79 @@ class AnthropicHandlerMixin:
|
|||
),
|
||||
)
|
||||
|
||||
class _DeferredCompressionResult:
|
||||
messages = working_messages
|
||||
transforms_applied = [
|
||||
"deferred:background_compression"
|
||||
# Cold-start fast pass: run everything EXCEPT the
|
||||
# Kompress ML stage synchronously before forwarding.
|
||||
# The byte-identical freeze (#1850) locks a session
|
||||
# to whatever form its cold start put in the provider
|
||||
# cache; deferring the WHOLE pipeline locks in the raw
|
||||
# transcript and forfeits the session's savings for
|
||||
# its lifetime — including sub-second wins like
|
||||
# read_lifecycle stale-read drops. Only Kompress can
|
||||
# blow the request budget (#1171), so only Kompress
|
||||
# stays deferred. Fail-open: on timeout/error the
|
||||
# request forwards exactly as before this pass
|
||||
# existed. On timeout the worker can't be cancelled
|
||||
# (Python can't preempt a running thread), but
|
||||
# _run_compression_in_executor runs it on the bounded
|
||||
# compression pool and tracks it via the leaked-thread
|
||||
# metric, so stragglers are capped and observable
|
||||
# rather than unbounded. The pass is also bounded by
|
||||
# routing + statistical crushers (observed seconds even
|
||||
# on multi-M-token counts).
|
||||
from headroom.proxy.helpers import (
|
||||
COLD_START_FAST_PASS_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
_fast_pass = None
|
||||
try:
|
||||
async with stage_timer.measure("compression_first_stage"):
|
||||
_fast_pass = await self._run_compression_in_executor(
|
||||
lambda: self.anthropic_pipeline.apply(
|
||||
messages=working_messages,
|
||||
model=model,
|
||||
model_limit=context_limit,
|
||||
context=extract_user_query(working_messages),
|
||||
frozen_message_count=frozen_message_count,
|
||||
idle_seconds=idle_seconds,
|
||||
biases=biases,
|
||||
request_id=request_id,
|
||||
compression_policy=compression_policy,
|
||||
skip_kompress=True,
|
||||
**proxy_pipeline_kwargs(self.config),
|
||||
),
|
||||
timeout=COLD_START_FAST_PASS_TIMEOUT_SECONDS,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.info(
|
||||
"[%s] Cold-start fast pass skipped (%s: %s); "
|
||||
"deferring full pipeline to background",
|
||||
request_id,
|
||||
type(e).__name__,
|
||||
e,
|
||||
)
|
||||
|
||||
if _fast_pass is not None:
|
||||
comp_cache.update_from_result(messages, _fast_pass.messages)
|
||||
_fast_pass.transforms_applied = list(
|
||||
_fast_pass.transforms_applied
|
||||
) + [
|
||||
"deferred:kompress_background"
|
||||
if accepted
|
||||
else "deferred:dropped"
|
||||
]
|
||||
timing = {}
|
||||
result = _fast_pass
|
||||
else:
|
||||
|
||||
result = _DeferredCompressionResult()
|
||||
class _DeferredCompressionResult:
|
||||
messages = working_messages
|
||||
transforms_applied = [
|
||||
"deferred:background_compression"
|
||||
if accepted
|
||||
else "deferred:dropped"
|
||||
]
|
||||
timing = {}
|
||||
|
||||
result = _DeferredCompressionResult()
|
||||
else:
|
||||
async with stage_timer.measure("compression_first_stage"):
|
||||
result = await self._run_compression_in_executor(
|
||||
|
|
|
|||
|
|
@ -620,6 +620,20 @@ try:
|
|||
except ValueError:
|
||||
COMPRESSION_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
# Cold-start fast-pass timeout in seconds. When background compression defers
|
||||
# a cold-start-large request, the handler still runs the pipeline synchronously
|
||||
# with skip_kompress=True (everything except the ML stage) under this budget so
|
||||
# the FORWARDED — and therefore provider-cached and byte-identically frozen —
|
||||
# form carries the cheap savings instead of the raw transcript. Without the ML
|
||||
# stage the pass is bounded by routing + statistical crushers (seconds, not the
|
||||
# 30s Kompress budget). Fail-open: on timeout the request forwards as before.
|
||||
try:
|
||||
COLD_START_FAST_PASS_TIMEOUT_SECONDS = float(
|
||||
os.environ.get("HEADROOM_COLD_START_FAST_PASS_TIMEOUT_SECONDS", "10")
|
||||
)
|
||||
except ValueError:
|
||||
COLD_START_FAST_PASS_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
# Eager startup preload timeout in seconds. The preload (compressor/parser models,
|
||||
# cache-only, allow_download=False) runs off the event loop during startup; this
|
||||
# bound only fires on a true hang or an uncatchable native stall so the proxy still
|
||||
|
|
|
|||
|
|
@ -2379,7 +2379,8 @@ class ContentRouter(Transform):
|
|||
# background (ensure_background_load) instead of blocking this request
|
||||
# thread on a 274MB download that races the compression timeout and
|
||||
# fails open. Until it is cached, route around the deep path.
|
||||
if self.config.enable_kompress:
|
||||
# skip_kompress (cold-start fast pass) takes the identical fallback.
|
||||
if self.config.enable_kompress and not getattr(self, "_runtime_skip_kompress", False):
|
||||
compressor = self._get_kompress()
|
||||
if compressor:
|
||||
if not compressor.is_ready():
|
||||
|
|
@ -3254,6 +3255,11 @@ class ContentRouter(Transform):
|
|||
self._runtime_force_kompress: bool = bool(
|
||||
kwargs.get("force_kompress", self.config.force_kompress_all)
|
||||
)
|
||||
# skip_kompress: run everything EXCEPT the Kompress ML stage this
|
||||
# call. Used by the cold-start fast pass so the request-path pass
|
||||
# stays sub-second; units routed to Kompress take the same fallback
|
||||
# they take when the model isn't ready. Wins over force_kompress.
|
||||
self._runtime_skip_kompress: bool = bool(kwargs.get("skip_kompress", False))
|
||||
self._runtime_kompress_model: str | None = kwargs.get("kompress_model")
|
||||
# F2.2: capture the per-request CompressionPolicy so
|
||||
# ``_record_to_toin`` can gate TOIN writes on
|
||||
|
|
|
|||
287
tests/test_cold_start_fast_pass.py
Normal file
287
tests/test_cold_start_fast_pass.py
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
"""Cold-start fast pass: when background compression defers a cold-start-large
|
||||
request, the handler still runs the pipeline synchronously with
|
||||
skip_kompress=True so the FORWARDED (and therefore provider-cached,
|
||||
byte-identically frozen) form carries the cheap savings. Only the Kompress ML
|
||||
stage stays deferred to the background job."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import anyio
|
||||
from fastapi import Request
|
||||
|
||||
from headroom.config import TransformResult
|
||||
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
|
||||
from headroom.proxy.models import ProxyConfig
|
||||
|
||||
_COMPRESSED_TEXT = "compressed tool output"
|
||||
|
||||
|
||||
class _DummyTokenizer:
|
||||
def count_messages(self, messages) -> int:
|
||||
return json.dumps(messages).count(" ") + 1
|
||||
|
||||
def count_text(self, text: str) -> int:
|
||||
return max(1, text.count(" ") + 1)
|
||||
|
||||
|
||||
class _DummyMetrics:
|
||||
async def record_request(self, **kwargs):
|
||||
return None
|
||||
|
||||
async def record_stage_timings(self, path, timings):
|
||||
return None
|
||||
|
||||
async def record_failed(self, **kwargs):
|
||||
return None
|
||||
|
||||
async def record_rate_limited(self, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
class _ResponseStub:
|
||||
status_code = 200
|
||||
headers: dict[str, str] = {}
|
||||
content = b'{"id":"msg_1","type":"message","role":"assistant","content":[],"usage":{"input_tokens":1,"output_tokens":1}}'
|
||||
|
||||
def json(self):
|
||||
return {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
}
|
||||
|
||||
|
||||
class _RecordingBackgroundCompressor:
|
||||
def __init__(self) -> None:
|
||||
self.enqueued: list[tuple[str, object, object]] = []
|
||||
|
||||
def enqueue(self, key, compress, store) -> bool:
|
||||
self.enqueued.append((key, compress, store))
|
||||
return True
|
||||
|
||||
|
||||
def _fake_pipeline_apply(messages, model, **kwargs):
|
||||
compressed = []
|
||||
for msg in messages:
|
||||
new = dict(msg)
|
||||
if msg.get("role") == "user" and isinstance(msg.get("content"), list):
|
||||
new["content"] = [
|
||||
{**part, "content": _COMPRESSED_TEXT}
|
||||
if isinstance(part, dict) and part.get("type") == "tool_result"
|
||||
else part
|
||||
for part in msg["content"]
|
||||
]
|
||||
compressed.append(new)
|
||||
return TransformResult(
|
||||
messages=compressed,
|
||||
tokens_before=1000,
|
||||
tokens_after=100,
|
||||
transforms_applied=["read_lifecycle:stale:test.py"],
|
||||
)
|
||||
|
||||
|
||||
class _DummyAnthropicHandler(AnthropicHandlerMixin):
|
||||
ANTHROPIC_API_URL = "https://api.anthropic.com"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.rate_limiter = None
|
||||
self.metrics = _DummyMetrics()
|
||||
self.config = ProxyConfig(
|
||||
optimize=True,
|
||||
image_optimize=False,
|
||||
retry_max_attempts=1,
|
||||
retry_base_delay_ms=1,
|
||||
retry_max_delay_ms=1,
|
||||
connect_timeout_seconds=10,
|
||||
mode="token",
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
fallback_enabled=False,
|
||||
fallback_provider=None,
|
||||
prefix_freeze_enabled=False,
|
||||
memory_enabled=False,
|
||||
)
|
||||
self.usage_reporter = None
|
||||
self.anthropic_provider = SimpleNamespace(get_context_limit=lambda model: 200_000)
|
||||
self.anthropic_pipeline = SimpleNamespace(apply=MagicMock(side_effect=_fake_pipeline_apply))
|
||||
self.anthropic_backend = None
|
||||
self.cost_tracker = None
|
||||
self.memory_handler = None
|
||||
self.cache = None
|
||||
self.security = None
|
||||
self.ccr_context_tracker = None
|
||||
self.ccr_injector = None
|
||||
self.ccr_response_handler = None
|
||||
self.ccr_feedback = None
|
||||
self.ccr_batch_processor = None
|
||||
self.ccr_mcp_server = None
|
||||
self.traffic_learner = None
|
||||
self.tool_injector = None
|
||||
self.read_lifecycle_manager = None
|
||||
self.logger = SimpleNamespace(log=lambda *a, **k: None)
|
||||
self.request_logger = self.logger
|
||||
self.usage_observer = None
|
||||
self.image_compressor = None
|
||||
self.session_tracker_store = SimpleNamespace(
|
||||
compute_session_id=lambda *a, **k: "sess-1",
|
||||
get_or_create=lambda *a, **k: SimpleNamespace(
|
||||
get_frozen_message_count=lambda: 0,
|
||||
get_last_original_messages=lambda: [],
|
||||
get_last_forwarded_messages=lambda: [],
|
||||
record_request=lambda *a, **k: None,
|
||||
),
|
||||
)
|
||||
# Cold-start deferral wiring under test.
|
||||
self._background_compression_enabled = True
|
||||
self._background_compression_min_tokens = 1
|
||||
self._background_compressor = _RecordingBackgroundCompressor()
|
||||
self.executor_calls: list[float] = []
|
||||
|
||||
async def _run_compression_in_executor(self, fn, timeout):
|
||||
self.executor_calls.append(timeout)
|
||||
return fn()
|
||||
|
||||
async def _next_request_id(self) -> str:
|
||||
return "req-fastpass-test"
|
||||
|
||||
def _extract_tags(self, headers):
|
||||
return {}
|
||||
|
||||
async def _retry_request(self, method, url, headers, body, **_kwargs):
|
||||
self.captured_body = body
|
||||
return _ResponseStub()
|
||||
|
||||
def _get_compression_cache(self, session_id):
|
||||
self.comp_cache_updates: list[tuple] = getattr(self, "comp_cache_updates", [])
|
||||
return SimpleNamespace(
|
||||
apply_cached=lambda m: m,
|
||||
compute_frozen_count=lambda m: 0,
|
||||
mark_stable_from_messages=lambda *a, **k: None,
|
||||
should_defer_compression=lambda h: False,
|
||||
mark_stable=lambda h: None,
|
||||
content_hash=lambda c: "h",
|
||||
update_from_result=lambda *a: self.comp_cache_updates.append(a),
|
||||
_cache={},
|
||||
_stable_hashes=set(),
|
||||
)
|
||||
|
||||
|
||||
def _build_request(body: dict) -> Request:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": payload, "more_body": False}
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"asgi": {"version": "3.0"},
|
||||
"http_version": "1.1",
|
||||
"method": "POST",
|
||||
"scheme": "https",
|
||||
"path": "/v1/messages",
|
||||
"raw_path": b"/v1/messages",
|
||||
"query_string": b"",
|
||||
"headers": [(b"authorization", b"Bearer sk-ant-api-test")],
|
||||
"client": ("127.0.0.1", 12345),
|
||||
"server": ("testserver", 443),
|
||||
}
|
||||
return Request(scope, receive)
|
||||
|
||||
|
||||
def test_cold_start_runs_fast_pass_and_defers_only_kompress(monkeypatch):
|
||||
import headroom.tokenizers as _tk
|
||||
|
||||
monkeypatch.setattr(_tk, "get_tokenizer", lambda model: _DummyTokenizer())
|
||||
|
||||
handler = _DummyAnthropicHandler()
|
||||
request = _build_request(
|
||||
{
|
||||
"model": "claude-3-5-sonnet-latest",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_1",
|
||||
"content": "verbose stale tool output " * 200,
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
anyio.run(handler.handle_anthropic_messages, request)
|
||||
|
||||
# The fast pass ran synchronously with the ML stage disabled.
|
||||
assert handler.executor_calls, "fast pass never ran through the executor"
|
||||
sync_calls = [
|
||||
c for c in handler.anthropic_pipeline.apply.call_args_list if c.kwargs.get("skip_kompress")
|
||||
]
|
||||
assert len(sync_calls) == 1, "expected exactly one synchronous skip_kompress pass"
|
||||
|
||||
# The full pipeline (kompress included) went to the background queue,
|
||||
# keyed against the ORIGINAL messages for content-hash reuse.
|
||||
assert len(handler._background_compressor.enqueued) == 1
|
||||
_key, bg_compress, _store = handler._background_compressor.enqueued[0]
|
||||
bg_compress()
|
||||
bg_calls = [
|
||||
c
|
||||
for c in handler.anthropic_pipeline.apply.call_args_list
|
||||
if not c.kwargs.get("skip_kompress")
|
||||
]
|
||||
assert len(bg_calls) == 1, "background job must run the full pipeline"
|
||||
|
||||
# The FORWARDED body carries the fast-pass form — that is what the
|
||||
# provider caches and the byte-identical freeze locks in.
|
||||
forwarded = handler.captured_body["messages"]
|
||||
assert forwarded[0]["content"][0]["content"] == _COMPRESSED_TEXT
|
||||
|
||||
# Fast-pass results were stored in the compression cache.
|
||||
assert handler.comp_cache_updates
|
||||
|
||||
|
||||
def test_fast_pass_failure_falls_back_to_full_deferral(monkeypatch):
|
||||
import headroom.tokenizers as _tk
|
||||
|
||||
monkeypatch.setattr(_tk, "get_tokenizer", lambda model: _DummyTokenizer())
|
||||
|
||||
handler = _DummyAnthropicHandler()
|
||||
|
||||
async def _boom(fn, timeout):
|
||||
raise TimeoutError("fast pass exceeded budget")
|
||||
|
||||
handler._run_compression_in_executor = _boom # type: ignore[method-assign]
|
||||
|
||||
original_text = "verbose stale tool output " * 200
|
||||
request = _build_request(
|
||||
{
|
||||
"model": "claude-3-5-sonnet-latest",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_1",
|
||||
"content": original_text,
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
anyio.run(handler.handle_anthropic_messages, request)
|
||||
|
||||
# Fail-open: original messages forwarded, background job still queued.
|
||||
forwarded = handler.captured_body["messages"]
|
||||
assert forwarded[0]["content"][0]["content"] == original_text
|
||||
assert len(handler._background_compressor.enqueued) == 1
|
||||
|
|
@ -204,6 +204,62 @@ def test_force_kompress_routes_anthropic_tool_result_to_targeted_kompress(
|
|||
assert captured["target_ratio"] == 0.10
|
||||
|
||||
|
||||
def test_skip_kompress_routes_around_ml_stage(router, tokenizer, monkeypatch):
|
||||
"""skip_kompress (cold-start fast pass) must never invoke the Kompress ML
|
||||
stage — units that would route there take the same fallback as when the
|
||||
model isn't ready — and wins over force_kompress."""
|
||||
calls: list[str] = []
|
||||
|
||||
class FakeKompress:
|
||||
def is_ready(self) -> bool:
|
||||
return True
|
||||
|
||||
def ensure_background_load(self) -> None:
|
||||
pass
|
||||
|
||||
def compress(self, content, **kwargs):
|
||||
calls.append(content)
|
||||
compressed = " ".join(content.split()[:20]) + " Retrieve more: hash=deadbeef"
|
||||
return SimpleNamespace(
|
||||
compressed=compressed,
|
||||
compressed_tokens=len(compressed.split()),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(router, "_get_kompress", lambda: FakeKompress())
|
||||
tool_content = " ".join(
|
||||
f'{{"file":"src/module_{i}.py","line":{i},"text":"repeated search payload"}}'
|
||||
for i in range(160)
|
||||
)
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_search_1",
|
||||
"content": tool_content,
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
result = router.apply(
|
||||
messages,
|
||||
tokenizer,
|
||||
force_kompress=True,
|
||||
skip_kompress=True,
|
||||
target_ratio=0.10,
|
||||
compress_user_messages=True,
|
||||
min_tokens_to_compress=10,
|
||||
read_protection_window=0,
|
||||
)
|
||||
|
||||
assert calls == []
|
||||
assert result.transforms_applied != ["router:tool_result:kompress"]
|
||||
# The pass still completes and returns a well-formed message list.
|
||||
assert result.messages[0]["content"][0]["content"]
|
||||
|
||||
|
||||
def test_anthropic_tool_result_lossy_without_marker_stays_verbatim(router, tokenizer, monkeypatch):
|
||||
"""Reversibility gate (#1307): a lossy Kompress result on a tool_result block
|
||||
with no CCR retrieval marker is unrecoverable, so the router must keep the
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue