mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(compress): reach the lossless provider seam on the general path and default /v1/compress to marker-free output (#2691)
## Description Two related changes to the compression seams, plus the review fixes for both. Supersedes #2661 and #2662, which are closed in favour of this branch — the fixes are inseparable from the code they fix, so reviewing them together is cheaper than landing two PRs and patching them afterwards. **1. A registered lossless provider now competes on the general path.** The `headroom.transforms.lossless_provider` seam was only ever consulted from `_lossless_compact_excluded`, gated on `DEFAULT_EXCLUDE_TOOLS` (`config.py:216` — `Read/Grep/Glob/Write/Edit/WebSearch/WebFetch`). Gateway traffic carries the caller's own tool names — LiteLLM's `headroom` guardrail (https://docs.litellm.ai/docs/proxy/headroom) posts requests containing tools like `search_docs` / `run_ci` / `fetch_rows` — so a registered provider was structurally unreachable for every gateway/sidecar deployment. The seam existed; nothing could get to it. **2. `POST /v1/compress` is marker-free by default.** A CCR marker is only useful to a caller that also injects the `headroom_retrieve` tool AND can reach `/v1/retrieve`. Neither holds here: tool injection lives in the provider request handlers (`handlers/anthropic.py:1894`), never in `handle_compress`; and every `/v1/retrieve*` route is `Depends(_require_loopback)` (`server.py:4422, 4470, 4749, 4781`) with no remote opt-in — `HEADROOM_COMPRESS_ALLOW_REMOTE` drops the loopback dependency on `/v1/compress` only. So a gateway forwards a `Retrieve more: hash=…` pointer the model cannot follow, and the proxy pays a CCR store write nobody reads. `config.mode="ccr"` opts back in. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [x] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made ### Seam — `content_router.py`, `lossless_provider.py` - `_lossless_first` (STAGE 0, every block on every path) consults `get_lossless_provider()` and keeps whichever output is smaller. **Strict no-op when no provider is registered**, which is the default; and because it is best-of rather than authoritative, a provider can never do worse than the built-in folds. - Malformed provider output can no longer escape. Every shape check runs inside the `try`: result must be `None`, or a 2-element tuple/list of two `str`. Anything else is ignored at debug level. (Previously the unpack sat outside the `try`, so a 3-tuple raised `ValueError` up through `TransformPipeline.apply`, which re-raises.) - Empty / whitespace-only candidates are rejected rather than silently replacing block content. - **Providers are never offered diff content.** Diff folding is subtractive with no inverse check and a reflowed hunk breaks `git apply` — the same reason the built-in `diff` fold is restricted at `content_router.py:2481`. - The third-party `kind` label is sanitised against `^[a-z0-9_]{1,32}$` before reaching `transforms_applied` and the per-strategy metric dicts, so a caller-controlled string cannot explode Prometheus label cardinality. `fullmatch`, not `match`: `$` also matches before a trailing newline, which would put a newline in a label. - `set_lossless_provider(provider, *, verifier=None)` — in lossless-only mode, where STAGE 0's output is final and there is no marker to recover from, a registered verifier must confirm the fold or the candidate is dropped. No verifier registered = today's behaviour. `provider=None` clears both. - The provider is invoked once per block, not twice (`_has_lossless_fold` probes `_lossless_first` and discards the result, then STAGE 0 recomputes). Bounded memo, wholesale clear on overflow, no lock — a race costs one redundant fold. The memo keys on the provider registration generation, so registering or clearing a provider after a block was already folded takes effect. - The seam docstring now records that the provider runs on the general path and inside the parallel compression pool, so it must be thread-safe as well as deterministic. ### Route — `handlers/openai.py`, `server.py` - `_derived_compress_pipeline(key, **overrides)` replaces the copy-pasted pipeline-derivation block; `_no_ccr_pipeline` (the new default) and `_lossy_inline_pipeline` both use it. - The default pipeline is **built at startup** and included in `_eager_preload_transforms`, so a fresh pod does not pay ContentRouter construction and compressor load on its first request, inside the compression-executor budget. - An unrecognised `config.mode` returns 400 naming the valid values instead of silently falling back to the default. - Claude-family model names resolve their context limit from the Anthropic provider. Real divergence: `bedrock/anthropic.claude-3-5-sonnet` is 200000 there and 128000 on the OpenAI provider. The tokenizer still comes from the OpenAI pipeline's provider — a separate, larger change, noted in a comment. - Documents why the derived router deliberately does **not** share the base router's compression cache: keys do not encode CCR-marker mode, so sharing would leak marker-laden entries into the marker-free path. ### Behavior change A `/v1/compress` caller that relied on default markers now gets none. The only in-tree caller that can resolve them is the TypeScript SDK (`sdk/typescript/src/client.ts:398 retrieve`, `:422 handleToolCall`); it needs `config: {"mode": "ccr"}` to keep today's behaviour, and landing that SDK default in the same release would leave only gateway callers — for whom markers were never resolvable — seeing a difference. No `HEADROOM_COMPRESS_DEFAULT_MODE` compat env deliberately: a flag nobody sets becomes permanent debt, and the wire-level `mode` already covers the one caller that needs it. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ python -m pytest tests/test_lossless_first_dispatch.py tests/test_lossless_excluded_compaction.py \ tests/test_lossless_mode.py tests/test_lossless_then_lossy.py tests/test_lossless_diff_fold_guard.py \ tests/test_bash_search_lossless_fold.py tests/test_proxy_compress_endpoint.py \ tests/test_ccr_row_drop_store_bridge.py tests/test_gateway_sidecar_ports.py tests/test_compress_api.py \ tests/test_platform_stabilization_functional.py tests/test_proxy_eager_preload_bind.py \ tests/test_proxy_warmup.py tests/test_router_registry_smartcrusher.py -q 189 passed in 32.04s $ ruff check <all 9 changed files> All checks passed! $ ruff format --check <all 9 changed files> 9 files already formatted $ mypy headroom/transforms/content_router.py headroom/transforms/lossless_provider.py \ headroom/proxy/handlers/openai.py headroom/proxy/server.py Success: no issues found in 4 source files ``` New tests cover, one concern each: every malformed provider shape; empty and whitespace-only results; diff content never reaching a provider (call-recording); `kind` sanitisation including the trailing-newline case; the verifier accepting / rejecting / raising; clearing a provider clearing its verifier; single provider invocation per block; memo invalidation on re-registration; unknown and valid `mode` values; the default pipeline existing before any request; and Claude vs OpenAI context-limit resolution with `token_budget` precedence preserved. ## Real Behavior Proof - **Environment:** macOS arm64, Python 3.12.6, proxy built from `_proxy_config_from_env()` with the default `coding` savings profile; Kompress both disabled and offloaded to a remote `kompress-v2-base` endpoint; `HEADROOM_COMPRESSION_TIMEOUT_SECONDS=300`. - **Steps:** `POST /v1/compress` over `TestClient` with OpenAI-shaped payloads under non-excluded tool names (`run_ci`, `list_files`, `code_search`, `fetch_rows`) — a CI log with ANSI escapes and repeated lines, a 160-path listing, a 150-line grep dump, a 150-row JSON array; plus a second payload with a RAG user blob, a 200-row JSON tool result and a 300-line log. Ran with and without a provider registered via `set_lossless_provider`. - **Observed — seam reachability:** | Kompress | no provider registered | provider registered | |---|---|---| | off | 19,284 → 10,265 tokens (46.8%) | 19,284 → **7,879 (59.1%)** | | on (remote) | 19,284 → 9,366 tokens (51.4%) | 19,284 → **7,146 (62.9%)** | Before this change the right-hand column was identical to the left — the registered provider was never called on this payload. - **Observed — marker-free default costs nothing:** | Config | tokens | saved | |---|---|---| | markers on (previous default) | 37,791 → 24,415 | 35.4% | | markers off (new default) | 37,791 → 24,415 | **35.4% — identical** | | `mode="lossy_inline"` | 37,791 → 25,129 | 33.5% | | `--lossless` | 37,791 → 35,100 | 7.1% | - **Observed — memo staleness, before the fix:** registering a provider that folds a grep block to 5 bytes left the block at its 1496-byte built-in fold, and clearing a provider kept serving the provider's output. Both correct after keying on the registration generation. - **Observed — context limit:** `bedrock/anthropic.claude-3-5-sonnet` resolves 200000 via the Anthropic provider, 128000 via the OpenAI provider. - **Not tested:** `tests/test_transforms_content_router.py` was not run — it does not complete on this machine, wedging on its 5th test while that test passes in 5.7s alone. Verified pre-existing before this work: with the diff stashed, the clean tree stalled at the identical test, and it stalled the same way under `HF_HUB_OFFLINE=1 HEADROOM_OFFLINE=1`. The machine was also out of disk at the time, which may be the real cause rather than the suspected native-detector deadlock (#575) — worth a separate issue either way. CI should be the arbiter here. ## 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 - [ ] 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 did **not** edit `CHANGELOG.md` ## Additional Notes Two pre-existing test files needed adjusting, both direct consequences rather than scope creep: - `test_platform_stabilization_functional.py::test_v1_compress_success_reports_actual_metrics` patches `openai_pipeline.apply`, which the default mode no longer routes through. **This test was already failing on the marker-free-default commit** before any of the fixes — my original test selection missed it. It now patches the pipeline the route actually uses. - `test_proxy_eager_preload_bind.py` substitutes fake pipelines to control exactly what the preload walks; the eager build injected the real derived router's statuses into an exact-equality assertion. Its shared helper now clears the derived cache, preserving each test's intent without weakening an assertion. Docs unchecked — follow-ups worth doing in the same release: document `config.mode` values in `docs/content/docs/litellm.mdx` and `wiki/proxy.md`; the TS SDK `mode:"ccr"` default; and an operational note that `COMPRESSION_TIMEOUT_SECONDS` defaults to 30 (`helpers.py:687`) while a remote ML endpoint makes one sequential call per unit — on a large payload it trips the executor timeout and the handler fails open, returning `compression_skipped: true` with `tokens_before: 0`, which reads as "nothing to save" rather than "we gave up". Those zeroed counters are misleading and worth a separate fix.
This commit is contained in:
Tejas Chopra
GitHub
parent
e0ce4b1d48
commit
f2c48e26c6
9 changed files with 883 additions and 41 deletions
|
|
@ -1207,6 +1207,12 @@ RESPONSES_CONTEXT_SEARCH_TIMEOUT_SECONDS = 2.0
|
|||
# accept) but short enough to bound the damage from a hung peer.
|
||||
WS_FIRST_FRAME_TIMEOUT_SECONDS = 60.0
|
||||
|
||||
# Accepted values for ``config.mode`` on POST /v1/compress. Unset/None means
|
||||
# the default marker-free pipeline. Anything else is a 400 rather than a
|
||||
# silent fall-through to the default (a typo or an unsupported mode like
|
||||
# "lossless" would otherwise look like it worked).
|
||||
COMPRESS_MODES = ("ccr", "lossy_inline", "lossless_then_lossy")
|
||||
|
||||
|
||||
def _extract_codex_handshake_headers(upstream: Any) -> list[tuple[str, str]]:
|
||||
"""Return the ``x-codex-*`` headers from an upstream WS handshake response.
|
||||
|
|
@ -8252,21 +8258,27 @@ class OpenAIHandlerMixin:
|
|||
with contextlib.suppress(Exception):
|
||||
await websocket.close()
|
||||
|
||||
def _lossy_inline_pipeline(self) -> Any:
|
||||
"""Cached pipeline for ``/v1/compress`` ``config.mode="lossy_inline"``.
|
||||
def _derived_compress_pipeline(self, key: str, **overrides: Any) -> Any:
|
||||
"""Cached ``/v1/compress`` pipeline derived from the live OpenAI router.
|
||||
|
||||
Runs the lossless byte/data fold first, then Kompresses the folded
|
||||
remainder (``lossless_then_lossy``). ``ccr_inject_marker=False`` makes
|
||||
every compressor (Kompress, SmartCrusher, search/log/config) emit inline
|
||||
lossy output with NO ``<<ccr:…>>`` / ``Retrieve more: hash=`` marker and
|
||||
NO CCR store write, so the result is safe to forward straight to a
|
||||
provider with no retrieval round-trip. Derived once from the live OpenAI
|
||||
router's config and reused read-only across requests.
|
||||
``overrides`` are applied to the live ContentRouter's config, so the
|
||||
derived pipeline inherits every operator setting (Kompress on/off,
|
||||
exclusions, profile knobs) and differs only in what the caller needs.
|
||||
|
||||
The derived router deliberately does NOT share the base router's
|
||||
compression cache. ``ContentRouter._cache`` is per-instance, and its
|
||||
keys do not encode the CCR-marker mode, so a shared cache would serve
|
||||
marker-laden entries (written by the OpenAI request path) to this
|
||||
marker-free path and vice versa. Sharing would be a correctness bug,
|
||||
not an optimisation — the duplicate cache is the intended trade.
|
||||
|
||||
ponytail: a first-request race just builds it twice — both are
|
||||
equivalent and Kompress weights are cached at module level, so no lock.
|
||||
"""
|
||||
cached = getattr(self, "_lossy_inline_pipeline_cache", None)
|
||||
cache = getattr(self, "_compress_pipeline_cache", None)
|
||||
if cache is None:
|
||||
cache = self._compress_pipeline_cache = {}
|
||||
cached = cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
|
|
@ -8277,29 +8289,71 @@ class OpenAIHandlerMixin:
|
|||
base = find_content_router(self.openai_pipeline)
|
||||
if base is None: # ponytail: nothing to derive from — use default pipeline
|
||||
return self.openai_pipeline
|
||||
cfg = replace(
|
||||
base.config,
|
||||
pipeline = TransformPipeline(
|
||||
transforms=[ContentRouter(replace(base.config, **overrides), observer=self.metrics)],
|
||||
provider=self.openai_provider,
|
||||
)
|
||||
cache[key] = pipeline
|
||||
return pipeline
|
||||
|
||||
def _no_ccr_pipeline(self) -> Any:
|
||||
"""Default ``/v1/compress`` pipeline: compress, but emit no CCR markers.
|
||||
|
||||
Every caller of this route is a gateway/sidecar or SDK client that
|
||||
forwards the returned messages straight to a provider. A
|
||||
``Retrieve more: hash=`` marker is only useful to a caller that also
|
||||
injects the ``headroom_retrieve`` tool AND can reach ``/v1/retrieve``
|
||||
(loopback-only). LiteLLM's `headroom` guardrail — the main consumer —
|
||||
does neither: it swaps ``messages`` and forwards. So markers here are a
|
||||
dangling pointer for the model plus a pointless CCR store write.
|
||||
``config.mode="ccr"`` opts back in for callers that do run the loop
|
||||
(e.g. the TypeScript SDK's ``retrieve``/``handleToolCall``).
|
||||
|
||||
Nothing else changes: same compressors, same aggressiveness. Measured
|
||||
identical token savings to the marker-on default on OpenAI-shaped
|
||||
gateway traffic.
|
||||
"""
|
||||
return self._derived_compress_pipeline(
|
||||
"no_ccr",
|
||||
ccr_inject_marker=False, # no markers in returned content
|
||||
ccr_enabled=False, # no CCR store writes
|
||||
)
|
||||
|
||||
def _lossy_inline_pipeline(self) -> Any:
|
||||
"""Pipeline for ``/v1/compress`` ``config.mode="lossy_inline"``.
|
||||
|
||||
Runs the lossless byte/data fold first, then Kompresses the folded
|
||||
remainder (``lossless_then_lossy``), marker-free like
|
||||
:meth:`_no_ccr_pipeline`. Kept as a distinct mode because the
|
||||
fold-then-Kompress ordering is a different compression posture, not a
|
||||
different CCR setting.
|
||||
"""
|
||||
return self._derived_compress_pipeline(
|
||||
"lossy_inline",
|
||||
lossless=False, # lossy mode (not lossless-only)
|
||||
lossless_then_lossy=True, # fold first, then Kompress the remainder
|
||||
ccr_inject_marker=False, # inline, marker-free everywhere
|
||||
ccr_enabled=False, # no CCR store writes
|
||||
smart_crusher_lossless_only=False, # keep SmartCrusher lossy
|
||||
) # enable_kompress inherited: on by default, off if operator disabled it
|
||||
pipeline = TransformPipeline(
|
||||
transforms=[ContentRouter(cfg, observer=self.metrics)],
|
||||
provider=self.openai_provider,
|
||||
)
|
||||
self._lossy_inline_pipeline_cache = pipeline
|
||||
return pipeline
|
||||
|
||||
async def handle_compress(self, request: Request) -> JSONResponse:
|
||||
"""Compress messages without calling an LLM.
|
||||
|
||||
POST /v1/compress
|
||||
Body: {"messages": [...], "model": "...", "config": {}}
|
||||
``config.mode="lossy_inline"`` (alias ``"lossless_then_lossy"``) selects
|
||||
the marker-free lossless-then-lossy pipeline whose output needs no CCR
|
||||
retrieval round-trip — the mode to use behind a gateway/sidecar.
|
||||
|
||||
``config.mode`` selects the pipeline:
|
||||
|
||||
* unset (default) — marker-free; forward the result straight to a
|
||||
provider with no CCR retrieval round-trip (see _no_ccr_pipeline).
|
||||
* ``"ccr"`` — CCR markers + store writes, for a caller that injects the
|
||||
``headroom_retrieve`` tool and can reach loopback ``/v1/retrieve``.
|
||||
* ``"lossy_inline"`` (alias ``"lossless_then_lossy"``) — marker-free,
|
||||
lossless fold first then Kompress the remainder.
|
||||
|
||||
Any other ``config.mode`` value is a 400 (see ``COMPRESS_MODES``).
|
||||
|
||||
Returns compressed messages + metrics.
|
||||
"""
|
||||
from fastapi.responses import JSONResponse
|
||||
|
|
@ -8389,10 +8443,27 @@ class OpenAIHandlerMixin:
|
|||
# Allow optional token_budget to override model's context limit
|
||||
# (used by OpenClaw compact() and other callers that need tighter budgets)
|
||||
token_budget = body.get("token_budget")
|
||||
# Resolve the context limit against the model's own provider. This
|
||||
# route is OpenAI-shaped, but LiteLLM's `headroom` guardrail passes
|
||||
# Anthropic model names straight through (claude-sonnet-4-5-...,
|
||||
# bedrock/anthropic.claude-3-5-sonnet, anthropic/claude-opus-4), and
|
||||
# the OpenAI provider answers those with its 128K default instead of
|
||||
# 200K+. Substring match is enough here — no model registry.
|
||||
#
|
||||
# Known gap: the TOKENIZER still comes from the OpenAI pipeline's
|
||||
# provider, so token counts remain approximate for Claude models.
|
||||
# Fixing that means selecting the whole pipeline per model family,
|
||||
# which is a larger change and out of scope for this route.
|
||||
model_name = model if isinstance(model, str) else str(model)
|
||||
limit_provider = (
|
||||
self.anthropic_provider
|
||||
if ("claude" in model_name.lower() or "anthropic" in model_name.lower())
|
||||
else self.openai_provider
|
||||
)
|
||||
context_limit = (
|
||||
token_budget
|
||||
if token_budget and isinstance(token_budget, int)
|
||||
else self.openai_provider.get_context_limit(model)
|
||||
else limit_provider.get_context_limit(model)
|
||||
)
|
||||
# Extract CompressConfig options from request body
|
||||
compress_config = body.get("config", {})
|
||||
|
|
@ -8402,14 +8473,30 @@ class OpenAIHandlerMixin:
|
|||
target_ratio = compress_config.get("target_ratio")
|
||||
protect_recent = compress_config.get("protect_recent")
|
||||
protect_analysis_context = compress_config.get("protect_analysis_context")
|
||||
# Marker-free lossless-then-lossy mode: safe to forward downstream
|
||||
# with no CCR retrieval round-trip (see _lossy_inline_pipeline).
|
||||
# Mode selection. Default is marker-free (see _no_ccr_pipeline):
|
||||
# no caller of this route can resolve a CCR marker unless it opts in
|
||||
# with mode="ccr", which restores the full marker + store behaviour.
|
||||
mode = compress_config.get("mode")
|
||||
pipeline = (
|
||||
self._lossy_inline_pipeline()
|
||||
if mode in ("lossy_inline", "lossless_then_lossy")
|
||||
else self.openai_pipeline
|
||||
)
|
||||
if mode is not None and mode not in COMPRESS_MODES:
|
||||
return JSONResponse(
|
||||
status_code=400,
|
||||
content={
|
||||
"error": {
|
||||
"type": "invalid_request",
|
||||
"message": (
|
||||
f"Invalid config.mode: {mode!r}. "
|
||||
f"Valid values are: {', '.join(COMPRESS_MODES)} "
|
||||
"(or omit config.mode for the default marker-free mode)."
|
||||
),
|
||||
}
|
||||
},
|
||||
)
|
||||
if mode in ("lossy_inline", "lossless_then_lossy"):
|
||||
pipeline = self._lossy_inline_pipeline()
|
||||
elif mode == "ccr":
|
||||
pipeline = self.openai_pipeline
|
||||
else:
|
||||
pipeline = self._no_ccr_pipeline()
|
||||
|
||||
pipeline_kwargs: dict = {
|
||||
"model_limit": context_limit,
|
||||
|
|
|
|||
|
|
@ -956,6 +956,21 @@ class HeadroomProxy(
|
|||
transforms=[*_intercept_prefix, cache_aligner, openai_router],
|
||||
provider=self.openai_provider,
|
||||
)
|
||||
# Build the DEFAULT /v1/compress pipeline now, not on first request.
|
||||
# It is a ContentRouter derived from `openai_router` (marker-free), so
|
||||
# a lazy build would land inside the bounded compression executor on a
|
||||
# cold pod's very first gateway request — paying router construction
|
||||
# and, when the ML model runs in-process, model load against the 30 s
|
||||
# compression budget. Building it here also lets startup warmup see it
|
||||
# (`_eager_preload_transforms`). Never fatal: on failure the handler
|
||||
# falls back to the original lazy path.
|
||||
try:
|
||||
self._no_ccr_pipeline()
|
||||
except Exception as exc: # pragma: no cover - defensive
|
||||
logger.debug(
|
||||
"Eager /v1/compress pipeline build failed (%s); deferring to first request",
|
||||
exc,
|
||||
)
|
||||
|
||||
# Initialize components
|
||||
self.cache = (
|
||||
|
|
@ -1580,7 +1595,12 @@ class HeadroomProxy(
|
|||
eager_status: dict[str, str] = {}
|
||||
transform_statuses: list[dict[str, str]] = []
|
||||
seen_transform_ids: set[int] = set()
|
||||
for pipeline in (self.anthropic_pipeline, self.openai_pipeline):
|
||||
# The derived /v1/compress pipelines (built eagerly in __init__) own
|
||||
# their own ContentRouter instances, so dedup-by-id() does not cover
|
||||
# them via the two request pipelines — warm them explicitly or the
|
||||
# first gateway request still pays the compressor load.
|
||||
derived_pipelines = list(getattr(self, "_compress_pipeline_cache", {}).values())
|
||||
for pipeline in (self.anthropic_pipeline, self.openai_pipeline, *derived_pipelines):
|
||||
for transform in pipeline.transforms:
|
||||
if id(transform) in seen_transform_ids:
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -78,7 +78,11 @@ from .content_detector import (
|
|||
)
|
||||
from .content_detector import detect_content_type as _regex_detect_content_type
|
||||
from .error_detection import content_has_strong_error_indicators
|
||||
from .lossless_provider import get_lossless_provider
|
||||
from .lossless_provider import (
|
||||
get_lossless_generation,
|
||||
get_lossless_provider,
|
||||
get_lossless_verifier,
|
||||
)
|
||||
from .mixed_content import ContentSection, mixed_content_indicators
|
||||
from .relevance_split import build_relevance_query, plan_relevance_split
|
||||
|
||||
|
|
@ -101,6 +105,17 @@ _detect_native_verified = False # native detect has returned once -> skip the w
|
|||
_TOKEN_ESTIMATOR = EstimatingTokenCounter()
|
||||
|
||||
|
||||
# `kind` labels supplied by a third-party lossless provider flow into
|
||||
# `transforms_applied`, `transforms_summary` and the per-strategy metric/timing
|
||||
# dicts (which become Prometheus label values). An unbounded caller-controlled
|
||||
# string there is a label-cardinality explosion, so only a short, lowercase,
|
||||
# identifier-shaped label is accepted; anything else is replaced with
|
||||
# `_PROVIDER_KIND_FALLBACK` rather than raising (a bad label must not fail a
|
||||
# request, and the fold itself is still valid).
|
||||
_PROVIDER_KIND_RE = re.compile(r"^[a-z0-9_]{1,32}$")
|
||||
_PROVIDER_KIND_FALLBACK = "provider"
|
||||
|
||||
|
||||
def _estimate_tokens(text: str) -> int:
|
||||
"""Size-proportional token estimate for section ratio decisions.
|
||||
|
||||
|
|
@ -1700,6 +1715,12 @@ class ContentRouter(Transform):
|
|||
# (fewer lossy chains, safer); 0 = keep the lossy pass on any improvement.
|
||||
_DEFAULT_LOSSY_MIN_EXTRA_SAVINGS = 0.05
|
||||
|
||||
# Entry cap for the `_lossless_first` memo. Sized for "the blocks of a
|
||||
# handful of in-flight requests", not for cross-request reuse — on overflow
|
||||
# the whole dict is dropped rather than evicted entry-by-entry, which keeps
|
||||
# the memo O(1) and unlockable at the cost of an occasional cold start.
|
||||
_LOSSLESS_FIRST_MEMO_MAX = 256
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: ContentRouterConfig | None = None,
|
||||
|
|
@ -1764,6 +1785,23 @@ class ContentRouter(Transform):
|
|||
self._relevance_scorer: Any = None
|
||||
self._relevance_scorer_tried: bool = False
|
||||
self._relevance_prewarm_started: bool = False
|
||||
# STAGE 0 memo: `(hash(content), len(content), strategy) -> (best, label)`.
|
||||
# `_lossless_first` is a pure function of its arguments but is called
|
||||
# twice for the same block (`_has_lossless_fold` probes it, then STAGE 0
|
||||
# recomputes it), which also means a registered lossless provider gets
|
||||
# invoked twice. Bounded by `_LOSSLESS_FIRST_MEMO_MAX` with a wholesale
|
||||
# clear (no LRU bookkeeping — this is a within-request hit, not a
|
||||
# long-lived cache; `_cache` is the cross-request one).
|
||||
self._lossless_first_memo: dict[
|
||||
tuple[int, int, CompressionStrategy, int], tuple[str, str | None]
|
||||
] = {}
|
||||
# Companion memo for the third-party half of STAGE 0, keyed by content
|
||||
# (plus the provider generation) but NOT by strategy — the provider
|
||||
# contract takes no strategy, so the two probes above, which pass
|
||||
# different strategies for the same block, share one provider
|
||||
# invocation. See `_lossless_provider_result`.
|
||||
self._lossless_provider_memo: dict[tuple[int, int, int], tuple[str, str] | None] = {}
|
||||
|
||||
# tool_call_id → compact args text, populated by _build_tool_name_map.
|
||||
self._tool_call_args: dict[str, str] = {}
|
||||
# tool_call_id → raw shell command (bash-search fold), same population.
|
||||
|
|
@ -2462,6 +2500,26 @@ class ContentRouter(Transform):
|
|||
"""
|
||||
from headroom.transforms.lossless_compaction import compact_lossless
|
||||
|
||||
# Memo: `_has_lossless_fold` probes this method purely to decide whether a
|
||||
# small block is worth admitting and throws the result away, and the real
|
||||
# STAGE 0 pass then recomputes it. That doubles the built-in folds *and*
|
||||
# the third-party provider call for every block that takes both paths.
|
||||
# The method is a pure function of (content, strategy), so cache it.
|
||||
# Resolved with `getattr` so the method still works on a bare instance
|
||||
# (`object.__new__(ContentRouter)`, used by unit tests that exercise the
|
||||
# folds without paying for a full router init).
|
||||
memo = getattr(self, "_lossless_first_memo", None)
|
||||
if memo is None:
|
||||
memo = self._lossless_first_memo = {}
|
||||
# The provider generation is part of the key: this memo caches a value
|
||||
# that DEPENDS on the registered provider, so a provider registered (or
|
||||
# cleared) after a block was first folded must not be served the old
|
||||
# answer forever.
|
||||
memo_key = (hash(content), len(content), strategy, get_lossless_generation())
|
||||
memoized = memo.get(memo_key)
|
||||
if memoized is not None:
|
||||
return memoized
|
||||
|
||||
# Apply losslessness to the OUTPUT structure, not to the classification:
|
||||
# try the fold implied by the detected strategy first, then the others.
|
||||
# Each compact_lossless call is self-verifying (exact inverse or returns
|
||||
|
|
@ -2498,8 +2556,132 @@ class ContentRouter(Transform):
|
|||
continue
|
||||
if len(cand) < len(best):
|
||||
best, best_label = cand, f"lossless_{kind}"
|
||||
|
||||
# A registered lossless provider (headroom.transforms.lossless_provider)
|
||||
# competes on the GENERAL path too, not only on excluded-tool output via
|
||||
# `_lossless_compact_excluded`. Without this an external fold only ever
|
||||
# saw Read/Grep/Glob-style results, so it was inert for gateway traffic
|
||||
# (`/v1/compress`), where tool names are the caller's own. Same contract
|
||||
# (information-preserving + deterministic) and we keep whichever output
|
||||
# is smaller, so a provider can only improve on the built-in folds.
|
||||
#
|
||||
# Diffs are off limits for a provider, for the same reason the built-in
|
||||
# "diff" fold is restricted above: diff folding is subtractive with no
|
||||
# inverse check, and a Kompressed/reflowed hunk breaks `git apply`. The
|
||||
# built-in path can at least reason about its own fold; a third-party one
|
||||
# we cannot, so we simply never offer it diff content.
|
||||
if strategy is not CompressionStrategy.DIFF and not self._looks_like_diff(content):
|
||||
best, best_label = self._apply_lossless_provider(content, best, best_label)
|
||||
|
||||
# Plain dict operations only: this is a pure-function cache, so a racing
|
||||
# double-compute under the parallel compression pool (see
|
||||
# HEADROOM_COMPRESS_WORKERS) costs one redundant fold and nothing else.
|
||||
# A lock here would sit on the hot path of every block for no
|
||||
# correctness gain. Bounded by a wholesale clear rather than LRU
|
||||
# bookkeeping — the memo exists to collapse a within-request double
|
||||
# call, not to be a cross-request cache (`self._cache` is that).
|
||||
if len(memo) >= self._LOSSLESS_FIRST_MEMO_MAX:
|
||||
memo.clear()
|
||||
memo[memo_key] = (best, best_label)
|
||||
return best, best_label
|
||||
|
||||
def _apply_lossless_provider(
|
||||
self, content: str, best: str, best_label: str | None
|
||||
) -> tuple[str, str | None]:
|
||||
"""Let a registered lossless provider beat ``best``; never raise.
|
||||
|
||||
The provider only wins on a strictly smaller, non-blank candidate — and
|
||||
in lossless-only mode only if an optionally registered verifier confirms
|
||||
the fold is genuinely reversible.
|
||||
"""
|
||||
supplied = self._lossless_provider_result(content)
|
||||
if supplied is None:
|
||||
return best, best_label
|
||||
cand, kind = supplied
|
||||
|
||||
if len(cand) >= len(best):
|
||||
return best, best_label
|
||||
|
||||
# Lossless-only mode has no downstream recovery: STAGE 0's output IS the
|
||||
# answer and there is no CCR marker to retrieve the original from. The
|
||||
# built-in folds are self-verifying; a provider is trusted unless the
|
||||
# operator registered a verifier, in which case it must pass.
|
||||
if self.config.lossless:
|
||||
verifier = get_lossless_verifier()
|
||||
if verifier is not None:
|
||||
try:
|
||||
ok = verifier(content, cand)
|
||||
except Exception: # noqa: BLE001 - a raising verifier means "unverified"
|
||||
logger.debug("lossless verifier raised; rejecting candidate", exc_info=True)
|
||||
return best, best_label
|
||||
if not ok:
|
||||
logger.debug("lossless verifier rejected the provider candidate")
|
||||
return best, best_label
|
||||
|
||||
return cand, f"lossless_{kind}"
|
||||
|
||||
def _lossless_provider_result(self, content: str) -> tuple[str, str] | None:
|
||||
"""Validated + memoized ``(candidate, safe_kind)`` from the provider.
|
||||
|
||||
Everything a third-party provider hands back is treated as hostile: a
|
||||
malformed shape, an empty result, a metric-unsafe ``kind`` label or an
|
||||
outright exception all degrade to ``None`` ("provider ignored"). This
|
||||
runs on the request path and its caller (``TransformPipeline.apply``)
|
||||
re-raises, so nothing here may escape.
|
||||
|
||||
Memoized per ``(hash(content), len(content))`` — the provider contract is
|
||||
``content -> result`` with no other input, so its answer cannot depend on
|
||||
the routing strategy. That matters because one block reaches
|
||||
``_lossless_first`` twice under *different* strategies (the
|
||||
``_has_lossless_fold`` admission probe passes PASSTHROUGH, STAGE 0 passes
|
||||
the detected strategy), which would otherwise invoke third-party code
|
||||
twice per block. Same no-lock rationale as the ``_lossless_first`` memo.
|
||||
"""
|
||||
provider = get_lossless_provider()
|
||||
if provider is None:
|
||||
return None
|
||||
|
||||
memo = getattr(self, "_lossless_provider_memo", None)
|
||||
if memo is None:
|
||||
memo = self._lossless_provider_memo = {}
|
||||
memo_key = (hash(content), len(content), get_lossless_generation())
|
||||
if memo_key in memo:
|
||||
return memo[memo_key]
|
||||
|
||||
result: tuple[str, str] | None = None
|
||||
try:
|
||||
supplied = provider(content)
|
||||
if supplied is None:
|
||||
result = None
|
||||
# Validate defensively *inside* the try: unpacking a 3-tuple / a bare
|
||||
# string / a non-sequence raises, and that exception used to escape
|
||||
# all the way out of the request.
|
||||
elif not isinstance(supplied, tuple | list) or len(supplied) != 2:
|
||||
logger.debug("lossless provider returned a malformed result; ignoring")
|
||||
elif not isinstance(supplied[0], str) or not isinstance(supplied[1], str):
|
||||
logger.debug("lossless provider returned non-str members; ignoring")
|
||||
# An empty / whitespace-only "fold" wins every length comparison and
|
||||
# would silently delete the block's content. That is not lossless.
|
||||
elif not supplied[0].strip():
|
||||
logger.debug("lossless provider returned an empty result; ignoring")
|
||||
else:
|
||||
raw_kind = supplied[1]
|
||||
# fullmatch, not match: Python's `$` also matches just BEFORE a
|
||||
# trailing newline, so `re.match` would let "log\n" through and
|
||||
# put a newline into a metric label.
|
||||
kind = (
|
||||
raw_kind if _PROVIDER_KIND_RE.fullmatch(raw_kind) else _PROVIDER_KIND_FALLBACK
|
||||
)
|
||||
result = (supplied[0], kind)
|
||||
except Exception: # noqa: BLE001 - a broken provider must not break routing
|
||||
logger.debug("lossless provider failed in _lossless_first", exc_info=True)
|
||||
result = None
|
||||
|
||||
if len(memo) >= self._LOSSLESS_FIRST_MEMO_MAX:
|
||||
memo.clear()
|
||||
memo[memo_key] = result
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _looks_like_diff(content: str) -> bool:
|
||||
"""Cheap structural sniff for unified/git-diff content.
|
||||
|
|
@ -2508,7 +2690,8 @@ class ContentRouter(Transform):
|
|||
Kompressing hunks corrupts ``git apply``. This is defense-in-depth beyond the
|
||||
DIFF-strategy and ``lossless_diff``-label checks: a diff can be folded
|
||||
best under a non-diff label (e.g. blank-line collapse → ``lossless_text``)
|
||||
or mis-detected, and must still never reach the lossy stage.
|
||||
or mis-detected, and must still never reach the lossy stage. It is also
|
||||
what keeps a third-party lossless provider away from diffs entirely.
|
||||
"""
|
||||
return (
|
||||
"diff --git " in content
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
"""Pluggable provider for information-preserving compaction of protected output.
|
||||
"""Pluggable provider for information-preserving compaction of tool output.
|
||||
|
||||
Excluded ("protected") tool results are kept out of *lossy* compression for
|
||||
accuracy; the content router still applies reversible/data-preserving folds to
|
||||
|
|
@ -14,9 +14,39 @@ Contract — ``provider(content: str) -> tuple[compacted: str, kind: str] | None
|
|||
* The provider MUST be deterministic and depend only on ``content`` (no
|
||||
cross-message state), so the proxy's prefix cache stays byte-stable across
|
||||
turns.
|
||||
* The provider MUST be thread-safe. It is called from the router's parallel
|
||||
compression worker pool (``HEADROOM_COMPRESS_WORKERS``), so several blocks can
|
||||
be handed to it concurrently on different threads.
|
||||
* ``kind`` is a short label that ends up in ``transforms_applied`` and in
|
||||
per-strategy metric/timing keys. Only ``^[a-z0-9_]{1,32}$`` is accepted;
|
||||
anything else is replaced with a fixed fallback label so a third-party string
|
||||
cannot blow up Prometheus label cardinality.
|
||||
* A malformed return (not ``None``, not a 2-tuple of ``str``) or an empty /
|
||||
whitespace-only ``compacted`` is ignored, never fatal.
|
||||
|
||||
Where the provider is called:
|
||||
|
||||
* :meth:`ContentRouter._lossless_compact_excluded` — protected/excluded tool
|
||||
output. Here the provider is *authoritative*: when one is set the router does
|
||||
not run its built-in folds; it falls back to the built-in only if the provider
|
||||
raises.
|
||||
* :meth:`ContentRouter._lossless_first` — the GENERAL compression path, i.e.
|
||||
every block of every provider/harness (including ``/v1/compress`` gateway
|
||||
traffic, where tool names are the caller's own). Here the provider *competes*
|
||||
with the built-in folds and the smaller output wins, so it can only improve on
|
||||
them. Diff content is never handed to the provider — folding a diff can break
|
||||
``git apply``.
|
||||
|
||||
Verification — ``verifier(original: str, compacted: str) -> bool``:
|
||||
|
||||
In lossless-only mode (``--lossless``) the fold IS the final answer: nothing
|
||||
downstream can recover a mistake. The built-in folds are self-verifying
|
||||
(``compact_lossless`` returns its input when it cannot invert), a provider is
|
||||
not. Register an optional ``verifier`` alongside the provider and the router
|
||||
will run it on the provider's candidate in lossless-only mode, rejecting the
|
||||
candidate when it returns falsy *or* raises. With no verifier registered the
|
||||
behaviour is unchanged: the provider is trusted (documented trust).
|
||||
|
||||
A registered provider is *authoritative*: when one is set the router does not run
|
||||
its built-in folds — it falls back to the built-in only if the provider raises.
|
||||
Default is ``None`` → the router uses its built-in folds, unchanged.
|
||||
"""
|
||||
|
||||
|
|
@ -26,16 +56,47 @@ from collections.abc import Callable
|
|||
|
||||
#: ``content -> (compacted, kind) | None``.
|
||||
LosslessProvider = Callable[[str], "tuple[str, str] | None"]
|
||||
#: ``(original, compacted) -> is-it-really-lossless``.
|
||||
LosslessVerifier = Callable[[str, str], bool]
|
||||
|
||||
_provider: LosslessProvider | None = None
|
||||
_verifier: LosslessVerifier | None = None
|
||||
#: Bumped on every registration. Callers that memoize a provider's output must
|
||||
#: include this in their cache key, or a provider registered (or cleared) after
|
||||
#: a block was already folded would be ignored for that block forever. Normal
|
||||
#: deployments register once at extension install, but tests and any hot-reload
|
||||
#: path re-register at will.
|
||||
_generation: int = 0
|
||||
|
||||
|
||||
def set_lossless_provider(provider: LosslessProvider | None) -> None:
|
||||
"""Register (or clear, with ``None``) the lossless compaction provider."""
|
||||
global _provider
|
||||
def set_lossless_provider(
|
||||
provider: LosslessProvider | None,
|
||||
*,
|
||||
verifier: LosslessVerifier | None = None,
|
||||
) -> None:
|
||||
"""Register (or clear, with ``None``) the lossless compaction provider.
|
||||
|
||||
``verifier`` is optional and keyword-only so existing single-argument calls
|
||||
keep working. Passing ``provider=None`` clears *both* the provider and any
|
||||
previously registered verifier — a verifier without a provider is dead
|
||||
state, and leaving one behind would silently apply to the next provider.
|
||||
"""
|
||||
global _provider, _verifier, _generation
|
||||
_provider = provider
|
||||
_verifier = verifier if provider is not None else None
|
||||
_generation += 1
|
||||
|
||||
|
||||
def get_lossless_provider() -> LosslessProvider | None:
|
||||
"""Return the registered provider, or ``None`` if the built-in should run."""
|
||||
return _provider
|
||||
|
||||
|
||||
def get_lossless_generation() -> int:
|
||||
"""Registration counter — include it in any key that memoizes provider output."""
|
||||
return _generation
|
||||
|
||||
|
||||
def get_lossless_verifier() -> LosslessVerifier | None:
|
||||
"""Return the registered verifier, or ``None`` when the provider is trusted."""
|
||||
return _verifier
|
||||
|
|
|
|||
|
|
@ -385,6 +385,9 @@ def test_v1_compress_then_v1_retrieve_resolves_marker_hash() -> None:
|
|||
]
|
||||
req = {
|
||||
"model": "gpt-4o",
|
||||
# /v1/compress is marker-free by default (no gateway caller can resolve a
|
||||
# marker); mode="ccr" is the opt-in for callers that run the retrieve loop.
|
||||
"config": {"mode": "ccr"},
|
||||
"messages": [
|
||||
{"role": "user", "content": "Get items"},
|
||||
{
|
||||
|
|
|
|||
|
|
@ -9,8 +9,17 @@ its byte-exact fold — the lossless floor is never discarded by a later lossy
|
|||
stage.
|
||||
"""
|
||||
|
||||
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
|
||||
from headroom.transforms.content_router import (
|
||||
CompressionStrategy,
|
||||
ContentRouter,
|
||||
ContentRouterConfig,
|
||||
)
|
||||
from headroom.transforms.lossless_compaction import search_unheading
|
||||
from headroom.transforms.lossless_provider import (
|
||||
get_lossless_provider,
|
||||
get_lossless_verifier,
|
||||
set_lossless_provider,
|
||||
)
|
||||
|
||||
|
||||
def _grep_block() -> str:
|
||||
|
|
@ -105,6 +114,44 @@ def test_has_lossless_fold_admits_small_block_below_size_floor():
|
|||
assert router._has_lossless_fold("def f():\n return 1\n") is False
|
||||
|
||||
|
||||
def test_registered_provider_competes_on_the_general_path():
|
||||
# A registered lossless provider must be consulted for ANY block, not only
|
||||
# excluded-tool output — that gate made external folds inert for gateway
|
||||
# traffic (/v1/compress), where tool names are the caller's own. The smaller
|
||||
# output wins, so a provider can only improve on the built-in folds.
|
||||
from headroom.transforms.lossless_provider import set_lossless_provider
|
||||
|
||||
block = _grep_block()
|
||||
baseline, _, _ = _compress(block, lossless=True)
|
||||
better = "SHORTER-THAN-ANY-BUILTIN-FOLD\n"
|
||||
assert len(better) < len(baseline)
|
||||
|
||||
try:
|
||||
set_lossless_provider(lambda content: (better, "plugin"))
|
||||
out, was, tr = _compress(block + "\n", lossless=True) # fresh cache key
|
||||
assert was is True
|
||||
assert out == better
|
||||
assert tr == ["router:tool_result:lossless_plugin"]
|
||||
|
||||
# A provider that loses to the built-in fold is ignored, not adopted.
|
||||
set_lossless_provider(lambda content: (content + "padding" * 100, "plugin"))
|
||||
out2, _, tr2 = _compress(block + "\n\n", lossless=True)
|
||||
assert tr2 == ["router:tool_result:lossless_search"]
|
||||
assert len(out2) < len(block)
|
||||
|
||||
# A raising provider must not break routing.
|
||||
set_lossless_provider(_raise)
|
||||
out3, was3, _ = _compress(block + "\n\n\n", lossless=True)
|
||||
assert was3 is True
|
||||
assert len(out3) < len(block)
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
|
||||
|
||||
def _raise(content: str):
|
||||
raise RuntimeError("broken provider")
|
||||
|
||||
|
||||
def test_lossless_mode_non_foldable_is_lossless_noop_not_ratio_too_high():
|
||||
# In lossless-only mode, code with no byte-lossless fold is left verbatim.
|
||||
# That is NOT a rejected compression, so it must not be bucketed as
|
||||
|
|
@ -129,3 +176,226 @@ def test_lossless_mode_non_foldable_is_lossless_noop_not_ratio_too_high():
|
|||
assert was is False
|
||||
assert rc.get("lossless_noop", 0) >= 1
|
||||
assert rc.get("ratio_too_high", 0) == 0
|
||||
|
||||
|
||||
# ── Hostile / third-party provider input ──────────────────────────────────────
|
||||
# The provider is arbitrary out-of-tree code called on the request path, and
|
||||
# `_lossless_first`'s caller (`TransformPipeline.apply`) re-raises. Nothing a
|
||||
# provider returns may fail a request or silently destroy a block.
|
||||
|
||||
_SHORT = "SHORTER-THAN-ANY-BUILTIN-FOLD\n"
|
||||
|
||||
_MALFORMED_RESULTS = [
|
||||
(_SHORT, "plugin", "extra"), # 3-tuple -> would raise on unpack
|
||||
_SHORT, # bare string -> would unpack into 2 chars, or raise
|
||||
(_SHORT, 42), # non-str kind
|
||||
(42, "plugin"), # non-str candidate
|
||||
{"compacted": _SHORT}, # non-tuple entirely
|
||||
17, # not even iterable
|
||||
]
|
||||
|
||||
|
||||
def test_malformed_provider_result_is_ignored_not_raised():
|
||||
block = _grep_block()
|
||||
for i, bad in enumerate(_MALFORMED_RESULTS):
|
||||
try:
|
||||
set_lossless_provider(lambda content, bad=bad: bad)
|
||||
# Distinct content per shape so nothing is served from a cache/memo.
|
||||
out, was, tr = _compress(block + "\n" * (i + 1), lossless=True)
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
# No exception, and the built-in fold is still the answer.
|
||||
assert was is True, bad
|
||||
assert tr == ["router:tool_result:lossless_search"], bad
|
||||
assert len(out) < len(block), bad
|
||||
|
||||
|
||||
def test_empty_provider_result_is_rejected():
|
||||
# "" beats every candidate on length, so an unchecked empty result would
|
||||
# win and silently delete the block's content.
|
||||
block = _grep_block()
|
||||
for i, blank in enumerate(("", " ", "\n\t \n")):
|
||||
try:
|
||||
set_lossless_provider(lambda content, blank=blank: (blank, "plugin"))
|
||||
out, was, tr = _compress(block + "x" * (i + 1) + "\n", lossless=True)
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
assert was is True, repr(blank)
|
||||
assert tr == ["router:tool_result:lossless_search"], repr(blank)
|
||||
assert out.strip(), repr(blank)
|
||||
|
||||
|
||||
def test_provider_is_never_offered_diff_content():
|
||||
# Diff folding is subtractive with no inverse check and a reflowed hunk
|
||||
# breaks `git apply`, so a third-party fold must never see a diff — the
|
||||
# same reason the built-in "diff" fold is strategy-gated.
|
||||
diff = (
|
||||
"diff --git a/x b/x\nindex 1111111..2222222 100644\n--- a/x\n+++ b/x\n@@ -1 +1 @@\n-a\n+b\n"
|
||||
)
|
||||
router = ContentRouter(ContentRouterConfig(lossless=True))
|
||||
baseline_diff = router._lossless_first(diff, CompressionStrategy.DIFF)
|
||||
baseline_pass = router._lossless_first(diff, CompressionStrategy.PASSTHROUGH)
|
||||
|
||||
seen: list[str] = []
|
||||
|
||||
def recording(content: str):
|
||||
seen.append(content)
|
||||
return (_SHORT, "plugin")
|
||||
|
||||
try:
|
||||
set_lossless_provider(recording)
|
||||
fresh = ContentRouter(ContentRouterConfig(lossless=True))
|
||||
# Both the DIFF strategy and diff-*shaped* content under another
|
||||
# strategy must bypass the provider entirely.
|
||||
assert fresh._lossless_first(diff, CompressionStrategy.DIFF) == baseline_diff
|
||||
assert fresh._lossless_first(diff, CompressionStrategy.PASSTHROUGH) == baseline_pass
|
||||
assert seen == []
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
|
||||
|
||||
def test_provider_kind_label_is_sanitized():
|
||||
# `kind` reaches transforms_applied and the per-strategy metric/timing keys
|
||||
# (Prometheus label values), so an unbounded caller-controlled string is a
|
||||
# cardinality bomb. Only ^[a-z0-9_]{1,32}$ survives.
|
||||
block = _grep_block()
|
||||
# "log\n" is the subtle one: Python's `$` matches before a trailing newline,
|
||||
# so a `re.match`-based check would let a newline into a metric label.
|
||||
bogus = ["UPPER", "has space", "punct!", "x" * 200, "", "kebab-case", "log\n"]
|
||||
for i, kind in enumerate(bogus):
|
||||
try:
|
||||
set_lossless_provider(lambda content, kind=kind: (_SHORT, kind))
|
||||
out, was, tr = _compress(block + "\n" * (i + 1), lossless=True)
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
assert was is True, kind
|
||||
assert out == _SHORT, kind
|
||||
assert tr == ["router:tool_result:lossless_provider"], kind
|
||||
|
||||
# A clean label is preserved verbatim.
|
||||
try:
|
||||
set_lossless_provider(lambda content: (_SHORT, "log_fold_2"))
|
||||
_out, _was, tr = _compress(block + "\t", lossless=True)
|
||||
assert tr == ["router:tool_result:lossless_log_fold_2"]
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
|
||||
|
||||
def test_lossless_mode_verifier_gates_the_provider_candidate():
|
||||
# In --lossless mode STAGE 0's output IS the answer: no CCR marker, no
|
||||
# retrieval. The built-in folds self-verify; a provider only does if the
|
||||
# operator registered a verifier.
|
||||
block = _grep_block()
|
||||
provider = lambda content: (_SHORT, "plugin") # noqa: E731
|
||||
|
||||
def _run(verifier, suffix):
|
||||
try:
|
||||
set_lossless_provider(provider, verifier=verifier)
|
||||
return _compress(block + suffix, lossless=True)
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
|
||||
# No verifier -> documented trust, unchanged behaviour.
|
||||
_out, was, tr = _run(None, "\n")
|
||||
assert was is True
|
||||
assert tr == ["router:tool_result:lossless_plugin"]
|
||||
|
||||
# Verifier says yes -> adopted.
|
||||
out, was, tr = _run(lambda original, compacted: True, "\n\n")
|
||||
assert out == _SHORT
|
||||
assert tr == ["router:tool_result:lossless_plugin"]
|
||||
|
||||
# Verifier says no -> rejected, built-in fold stands.
|
||||
out, was, tr = _run(lambda original, compacted: False, "\n\n\n")
|
||||
assert tr == ["router:tool_result:lossless_search"]
|
||||
assert len(out) < len(block)
|
||||
|
||||
# Verifier raises -> "unverified" -> rejected, never fatal.
|
||||
def _boom(original, compacted):
|
||||
raise RuntimeError("verifier exploded")
|
||||
|
||||
out, was, tr = _run(_boom, "\n\n\n\n")
|
||||
assert tr == ["router:tool_result:lossless_search"]
|
||||
assert len(out) < len(block)
|
||||
|
||||
|
||||
def test_memo_does_not_go_stale_when_the_provider_changes():
|
||||
# The STAGE 0 memo caches a value that depends on the registered provider,
|
||||
# so its key includes the registration generation. Without that, a provider
|
||||
# registered (or cleared) after a block was already folded is ignored for
|
||||
# that exact block forever — invisible in production (extensions register at
|
||||
# startup) but wrong, and a trap for tests and any hot-reload path.
|
||||
from headroom.transforms.content_router import CompressionStrategy
|
||||
|
||||
router = ContentRouter(ContentRouterConfig(lossless=True))
|
||||
block = _grep_block()
|
||||
tiny = "TINY\n"
|
||||
|
||||
try:
|
||||
set_lossless_provider(None)
|
||||
baseline, baseline_label = router._lossless_first(block, CompressionStrategy.SEARCH)
|
||||
assert baseline_label == "lossless_search"
|
||||
|
||||
# Register AFTER the block was already folded once.
|
||||
set_lossless_provider(lambda content: (tiny, "plugin"))
|
||||
out, label = router._lossless_first(block, CompressionStrategy.SEARCH)
|
||||
assert out == tiny, "provider registered after first fold was ignored (stale memo)"
|
||||
assert label == "lossless_plugin"
|
||||
|
||||
# Clearing it must take effect too.
|
||||
set_lossless_provider(None)
|
||||
out, label = router._lossless_first(block, CompressionStrategy.SEARCH)
|
||||
assert out == baseline, "cleared provider still served from the memo"
|
||||
assert label == "lossless_search"
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
|
||||
|
||||
def test_clearing_the_provider_also_clears_the_verifier():
|
||||
# A verifier with no provider is dead state that would silently start
|
||||
# gating the *next* provider someone registers.
|
||||
try:
|
||||
set_lossless_provider(lambda content: None, verifier=lambda o, c: True)
|
||||
assert get_lossless_verifier() is not None
|
||||
set_lossless_provider(None)
|
||||
assert get_lossless_provider() is None
|
||||
assert get_lossless_verifier() is None
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
|
||||
|
||||
def test_provider_is_called_once_per_block_not_twice():
|
||||
# `_has_lossless_fold` (the small-block admission probe) and STAGE 0 both
|
||||
# run `_lossless_first` over the same block, under *different* strategies.
|
||||
# Memoizing the provider's answer by content collapses that to one call
|
||||
# into third-party code.
|
||||
calls: list[str] = []
|
||||
|
||||
def counting(content: str):
|
||||
calls.append(content)
|
||||
return None
|
||||
|
||||
small = "\n".join(f"pkg/mod/long_filename.py:{n}:value = {n}" for n in range(1, 8)) + "\n"
|
||||
try:
|
||||
set_lossless_provider(counting)
|
||||
router = ContentRouter(ContentRouterConfig(lossless=True))
|
||||
assert router._has_lossless_fold(small) is True
|
||||
assert len(calls) == 1
|
||||
out, was = router._compress_block_content(
|
||||
small,
|
||||
hash(small),
|
||||
"",
|
||||
1.0,
|
||||
1.0,
|
||||
None,
|
||||
[],
|
||||
{},
|
||||
[],
|
||||
"tool_result",
|
||||
"tool",
|
||||
True,
|
||||
)
|
||||
assert was is True
|
||||
assert len(calls) == 1
|
||||
finally:
|
||||
set_lossless_provider(None)
|
||||
|
|
|
|||
|
|
@ -68,7 +68,10 @@ def test_v1_compress_success_reports_actual_metrics(monkeypatch) -> None:
|
|||
markers_inserted=["marker-1"],
|
||||
)
|
||||
|
||||
monkeypatch.setattr(proxy.openai_pipeline, "apply", fake_apply)
|
||||
# The default /v1/compress mode runs a marker-free pipeline derived from
|
||||
# `openai_pipeline`, not `openai_pipeline` itself, so patch the one the
|
||||
# route actually uses. It is built eagerly at create_app() time.
|
||||
monkeypatch.setattr(proxy._compress_pipeline_cache["no_ccr"], "apply", fake_apply)
|
||||
|
||||
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client:
|
||||
response = client.post(
|
||||
|
|
|
|||
|
|
@ -428,6 +428,211 @@ class TestCompressEndpointLossyInlineMode:
|
|||
assert response.json()["ccr_hashes"] == []
|
||||
|
||||
|
||||
class TestCompressEndpointModeValidation:
|
||||
"""``config.mode`` must be validated, not silently ignored.
|
||||
|
||||
Before this, ``mode: "lossless"`` (a mode that does not exist) or any typo
|
||||
fell through to the default pipeline and the caller got a 200 describing a
|
||||
compression posture it never asked for.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad_mode",
|
||||
["lossless", "CCR", "ccr ", "no_ccr", "", 7, ["ccr"]],
|
||||
)
|
||||
def test_unknown_mode_returns_400(self, client, bad_mode):
|
||||
response = client.post(
|
||||
"/v1/compress",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"model": "gpt-4",
|
||||
"config": {"mode": bad_mode},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
data = response.json()
|
||||
assert data["error"]["type"] == "invalid_request"
|
||||
message = data["error"]["message"]
|
||||
# The message must tell the caller what IS valid.
|
||||
for valid in ("ccr", "lossy_inline", "lossless_then_lossy"):
|
||||
assert valid in message
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config",
|
||||
[
|
||||
{}, # mode unset -> default marker-free pipeline
|
||||
{"mode": None}, # explicit null is the same as unset
|
||||
{"mode": "ccr"},
|
||||
{"mode": "lossy_inline"},
|
||||
{"mode": "lossless_then_lossy"},
|
||||
],
|
||||
ids=["unset", "null", "ccr", "lossy_inline", "lossless_then_lossy"],
|
||||
)
|
||||
def test_valid_modes_return_200(self, client, config):
|
||||
response = client.post(
|
||||
"/v1/compress",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"model": "gpt-4",
|
||||
"config": config,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert isinstance(response.json()["messages"], list)
|
||||
|
||||
|
||||
class TestCompressDefaultPipelineBuiltAtStartup:
|
||||
"""The default (marker-free) /v1/compress pipeline must exist before the
|
||||
first request.
|
||||
|
||||
It used to be built lazily, so a cold pod paid ContentRouter construction —
|
||||
and in-process ML model load — inside the bounded compression executor on
|
||||
its first real gateway request.
|
||||
"""
|
||||
|
||||
def test_default_pipeline_exists_before_any_request(self):
|
||||
config = ProxyConfig(
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
)
|
||||
# No TestClient / no request: only create_app().
|
||||
proxy = create_app(config).state.proxy
|
||||
|
||||
cache = getattr(proxy, "_compress_pipeline_cache", None)
|
||||
assert cache, "default /v1/compress pipeline was not built at startup"
|
||||
assert "no_ccr" in cache
|
||||
# It must be a DERIVED pipeline, not the shared request pipeline.
|
||||
assert cache["no_ccr"] is not proxy.openai_pipeline
|
||||
|
||||
def test_startup_warmup_covers_the_derived_router(self):
|
||||
"""The eager compressor preload must walk the derived pipeline too."""
|
||||
config = ProxyConfig(
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
)
|
||||
proxy = create_app(config).state.proxy
|
||||
|
||||
derived = proxy._compress_pipeline_cache["no_ccr"]
|
||||
derived_ids = {id(t) for t in derived.transforms}
|
||||
|
||||
seen: list[int] = []
|
||||
for pipeline in (proxy.anthropic_pipeline, proxy.openai_pipeline):
|
||||
seen.extend(id(t) for t in pipeline.transforms)
|
||||
# Precondition: the derived router is NOT reachable via the request
|
||||
# pipelines, so dedup-by-id() cannot have covered it implicitly.
|
||||
assert derived_ids - set(seen)
|
||||
|
||||
_status, transform_statuses = proxy._eager_preload_transforms()
|
||||
# Base router + derived router both report a status dict.
|
||||
assert len(transform_statuses) >= 2
|
||||
|
||||
|
||||
class TestCompressContextLimitByModelFamily:
|
||||
"""LiteLLM's guardrail forwards Anthropic model names through this
|
||||
OpenAI-shaped route; the context limit must come from the Anthropic
|
||||
provider for those, not the OpenAI provider's 128K default."""
|
||||
|
||||
@staticmethod
|
||||
def _spy_providers(proxy, monkeypatch):
|
||||
anthropic_calls: list[str] = []
|
||||
openai_calls: list[str] = []
|
||||
|
||||
def anthropic_limit(model):
|
||||
anthropic_calls.append(model)
|
||||
return 987_654
|
||||
|
||||
def openai_limit(model):
|
||||
openai_calls.append(model)
|
||||
return 123_456
|
||||
|
||||
monkeypatch.setattr(proxy.anthropic_provider, "get_context_limit", anthropic_limit)
|
||||
monkeypatch.setattr(proxy.openai_provider, "get_context_limit", openai_limit)
|
||||
return anthropic_calls, openai_calls
|
||||
|
||||
@staticmethod
|
||||
def _spy_pipeline(proxy, monkeypatch):
|
||||
"""Capture the kwargs handed to the default (marker-free) pipeline."""
|
||||
seen: dict = {}
|
||||
|
||||
def fake_apply(**kwargs):
|
||||
seen.update(kwargs)
|
||||
return SimpleNamespace(
|
||||
messages=kwargs["messages"],
|
||||
tokens_before=10,
|
||||
tokens_after=10,
|
||||
transforms_applied=[],
|
||||
transforms_summary={},
|
||||
markers_inserted=[],
|
||||
)
|
||||
|
||||
monkeypatch.setattr(proxy._compress_pipeline_cache["no_ccr"], "apply", fake_apply)
|
||||
return seen
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"claude-sonnet-4-5-20250929",
|
||||
"bedrock/anthropic.claude-3-5-sonnet",
|
||||
"anthropic/claude-opus-4",
|
||||
"CLAUDE-Sonnet-4-5", # case-insensitive
|
||||
],
|
||||
)
|
||||
def test_claude_models_use_anthropic_context_limit(self, client, monkeypatch, model):
|
||||
proxy = client.app.state.proxy
|
||||
anthropic_calls, openai_calls = self._spy_providers(proxy, monkeypatch)
|
||||
seen = self._spy_pipeline(proxy, monkeypatch)
|
||||
|
||||
response = client.post(
|
||||
"/v1/compress",
|
||||
json={"messages": [{"role": "user", "content": "hello"}], "model": model},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert anthropic_calls == [model]
|
||||
assert openai_calls == []
|
||||
assert seen["model_limit"] == 987_654
|
||||
|
||||
def test_openai_models_still_use_openai_context_limit(self, client, monkeypatch):
|
||||
proxy = client.app.state.proxy
|
||||
anthropic_calls, openai_calls = self._spy_providers(proxy, monkeypatch)
|
||||
seen = self._spy_pipeline(proxy, monkeypatch)
|
||||
|
||||
response = client.post(
|
||||
"/v1/compress",
|
||||
json={"messages": [{"role": "user", "content": "hello"}], "model": "gpt-4o"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert openai_calls == ["gpt-4o"]
|
||||
assert anthropic_calls == []
|
||||
assert seen["model_limit"] == 123_456
|
||||
|
||||
def test_token_budget_still_overrides_for_claude_models(self, client, monkeypatch):
|
||||
"""token_budget precedence must survive the provider routing."""
|
||||
proxy = client.app.state.proxy
|
||||
anthropic_calls, openai_calls = self._spy_providers(proxy, monkeypatch)
|
||||
seen = self._spy_pipeline(proxy, monkeypatch)
|
||||
|
||||
response = client.post(
|
||||
"/v1/compress",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"model": "claude-sonnet-4-5-20250929",
|
||||
"token_budget": 4096,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert seen["model_limit"] == 4096
|
||||
# Neither provider is consulted when the caller pins a budget.
|
||||
assert anthropic_calls == []
|
||||
assert openai_calls == []
|
||||
|
||||
|
||||
class TestCompressEndpointDoesNotBlockLoop:
|
||||
"""/v1/compress must offload to the compression executor so a slow/large
|
||||
payload cannot freeze the single event loop (#718)."""
|
||||
|
|
@ -477,6 +682,9 @@ class TestCompressEndpointDoesNotBlockLoop:
|
|||
json={
|
||||
"messages": [{"role": "user", "content": "hello world"}],
|
||||
"model": "gpt-4",
|
||||
# mode="ccr" routes to `openai_pipeline` (the default is a
|
||||
# derived marker-free pipeline); this test patches that one.
|
||||
"config": {"mode": "ccr"},
|
||||
},
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -36,7 +36,14 @@ def _make_proxy(*, optimize: bool):
|
|||
image_optimize=False,
|
||||
subscription_tracking_enabled=False,
|
||||
)
|
||||
return create_app(config).state.proxy
|
||||
proxy = create_app(config).state.proxy
|
||||
# Every test here substitutes fake pipelines to control exactly what the
|
||||
# preload walks. The proxy also eagerly builds the default /v1/compress
|
||||
# pipeline (a derived ContentRouter, warmed alongside the request
|
||||
# pipelines), which would inject real transform statuses into those
|
||||
# assertions — drop it so the fakes remain the only input.
|
||||
proxy._compress_pipeline_cache = {}
|
||||
return proxy
|
||||
|
||||
|
||||
class _FastTransform:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue