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 2026-07-31 12:31:38 -07:00 committed by GitHub
parent e0ce4b1d48
commit f2c48e26c6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 883 additions and 41 deletions

View file

@ -1207,6 +1207,12 @@ RESPONSES_CONTEXT_SEARCH_TIMEOUT_SECONDS = 2.0
# accept) but short enough to bound the damage from a hung peer. # accept) but short enough to bound the damage from a hung peer.
WS_FIRST_FRAME_TIMEOUT_SECONDS = 60.0 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]]: def _extract_codex_handshake_headers(upstream: Any) -> list[tuple[str, str]]:
"""Return the ``x-codex-*`` headers from an upstream WS handshake response. """Return the ``x-codex-*`` headers from an upstream WS handshake response.
@ -8252,21 +8258,27 @@ class OpenAIHandlerMixin:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
await websocket.close() await websocket.close()
def _lossy_inline_pipeline(self) -> Any: def _derived_compress_pipeline(self, key: str, **overrides: Any) -> Any:
"""Cached pipeline for ``/v1/compress`` ``config.mode="lossy_inline"``. """Cached ``/v1/compress`` pipeline derived from the live OpenAI router.
Runs the lossless byte/data fold first, then Kompresses the folded ``overrides`` are applied to the live ContentRouter's config, so the
remainder (``lossless_then_lossy``). ``ccr_inject_marker=False`` makes derived pipeline inherits every operator setting (Kompress on/off,
every compressor (Kompress, SmartCrusher, search/log/config) emit inline exclusions, profile knobs) and differs only in what the caller needs.
lossy output with NO ``<<ccr:>>`` / ``Retrieve more: hash=`` marker and
NO CCR store write, so the result is safe to forward straight to a The derived router deliberately does NOT share the base router's
provider with no retrieval round-trip. Derived once from the live OpenAI compression cache. ``ContentRouter._cache`` is per-instance, and its
router's config and reused read-only across requests. 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 ponytail: a first-request race just builds it twice both are
equivalent and Kompress weights are cached at module level, so no lock. 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: if cached is not None:
return cached return cached
@ -8277,29 +8289,71 @@ class OpenAIHandlerMixin:
base = find_content_router(self.openai_pipeline) base = find_content_router(self.openai_pipeline)
if base is None: # ponytail: nothing to derive from — use default pipeline if base is None: # ponytail: nothing to derive from — use default pipeline
return self.openai_pipeline return self.openai_pipeline
cfg = replace( pipeline = TransformPipeline(
base.config, 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=False, # lossy mode (not lossless-only)
lossless_then_lossy=True, # fold first, then Kompress the remainder lossless_then_lossy=True, # fold first, then Kompress the remainder
ccr_inject_marker=False, # inline, marker-free everywhere ccr_inject_marker=False, # inline, marker-free everywhere
ccr_enabled=False, # no CCR store writes ccr_enabled=False, # no CCR store writes
smart_crusher_lossless_only=False, # keep SmartCrusher lossy 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: async def handle_compress(self, request: Request) -> JSONResponse:
"""Compress messages without calling an LLM. """Compress messages without calling an LLM.
POST /v1/compress POST /v1/compress
Body: {"messages": [...], "model": "...", "config": {}} 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 ``config.mode`` selects the pipeline:
retrieval round-trip the mode to use behind a gateway/sidecar.
* 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. Returns compressed messages + metrics.
""" """
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse
@ -8389,10 +8443,27 @@ class OpenAIHandlerMixin:
# Allow optional token_budget to override model's context limit # Allow optional token_budget to override model's context limit
# (used by OpenClaw compact() and other callers that need tighter budgets) # (used by OpenClaw compact() and other callers that need tighter budgets)
token_budget = body.get("token_budget") 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 = ( context_limit = (
token_budget token_budget
if token_budget and isinstance(token_budget, int) 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 # Extract CompressConfig options from request body
compress_config = body.get("config", {}) compress_config = body.get("config", {})
@ -8402,14 +8473,30 @@ class OpenAIHandlerMixin:
target_ratio = compress_config.get("target_ratio") target_ratio = compress_config.get("target_ratio")
protect_recent = compress_config.get("protect_recent") protect_recent = compress_config.get("protect_recent")
protect_analysis_context = compress_config.get("protect_analysis_context") protect_analysis_context = compress_config.get("protect_analysis_context")
# Marker-free lossless-then-lossy mode: safe to forward downstream # Mode selection. Default is marker-free (see _no_ccr_pipeline):
# with no CCR retrieval round-trip (see _lossy_inline_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") mode = compress_config.get("mode")
pipeline = ( if mode is not None and mode not in COMPRESS_MODES:
self._lossy_inline_pipeline() return JSONResponse(
if mode in ("lossy_inline", "lossless_then_lossy") status_code=400,
else self.openai_pipeline 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 = { pipeline_kwargs: dict = {
"model_limit": context_limit, "model_limit": context_limit,

View file

@ -956,6 +956,21 @@ class HeadroomProxy(
transforms=[*_intercept_prefix, cache_aligner, openai_router], transforms=[*_intercept_prefix, cache_aligner, openai_router],
provider=self.openai_provider, 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 # Initialize components
self.cache = ( self.cache = (
@ -1580,7 +1595,12 @@ class HeadroomProxy(
eager_status: dict[str, str] = {} eager_status: dict[str, str] = {}
transform_statuses: list[dict[str, str]] = [] transform_statuses: list[dict[str, str]] = []
seen_transform_ids: set[int] = set() 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: for transform in pipeline.transforms:
if id(transform) in seen_transform_ids: if id(transform) in seen_transform_ids:
continue continue

View file

@ -78,7 +78,11 @@ from .content_detector import (
) )
from .content_detector import detect_content_type as _regex_detect_content_type from .content_detector import detect_content_type as _regex_detect_content_type
from .error_detection import content_has_strong_error_indicators 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 .mixed_content import ContentSection, mixed_content_indicators
from .relevance_split import build_relevance_query, plan_relevance_split 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() _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: def _estimate_tokens(text: str) -> int:
"""Size-proportional token estimate for section ratio decisions. """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. # (fewer lossy chains, safer); 0 = keep the lossy pass on any improvement.
_DEFAULT_LOSSY_MIN_EXTRA_SAVINGS = 0.05 _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__( def __init__(
self, self,
config: ContentRouterConfig | None = None, config: ContentRouterConfig | None = None,
@ -1764,6 +1785,23 @@ class ContentRouter(Transform):
self._relevance_scorer: Any = None self._relevance_scorer: Any = None
self._relevance_scorer_tried: bool = False self._relevance_scorer_tried: bool = False
self._relevance_prewarm_started: 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. # tool_call_id → compact args text, populated by _build_tool_name_map.
self._tool_call_args: dict[str, str] = {} self._tool_call_args: dict[str, str] = {}
# tool_call_id → raw shell command (bash-search fold), same population. # 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 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: # Apply losslessness to the OUTPUT structure, not to the classification:
# try the fold implied by the detected strategy first, then the others. # try the fold implied by the detected strategy first, then the others.
# Each compact_lossless call is self-verifying (exact inverse or returns # Each compact_lossless call is self-verifying (exact inverse or returns
@ -2498,8 +2556,132 @@ class ContentRouter(Transform):
continue continue
if len(cand) < len(best): if len(cand) < len(best):
best, best_label = cand, f"lossless_{kind}" 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 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 @staticmethod
def _looks_like_diff(content: str) -> bool: def _looks_like_diff(content: str) -> bool:
"""Cheap structural sniff for unified/git-diff content. """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 Kompressing hunks corrupts ``git apply``. This is defense-in-depth beyond the
DIFF-strategy and ``lossless_diff``-label checks: a diff can be folded 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``) 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 ( return (
"diff --git " in content "diff --git " in content

View file

@ -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 Excluded ("protected") tool results are kept out of *lossy* compression for
accuracy; the content router still applies reversible/data-preserving folds to 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 * The provider MUST be deterministic and depend only on ``content`` (no
cross-message state), so the proxy's prefix cache stays byte-stable across cross-message state), so the proxy's prefix cache stays byte-stable across
turns. 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. Default is ``None`` the router uses its built-in folds, unchanged.
""" """
@ -26,16 +56,47 @@ from collections.abc import Callable
#: ``content -> (compacted, kind) | None``. #: ``content -> (compacted, kind) | None``.
LosslessProvider = Callable[[str], "tuple[str, str] | None"] LosslessProvider = Callable[[str], "tuple[str, str] | None"]
#: ``(original, compacted) -> is-it-really-lossless``.
LosslessVerifier = Callable[[str, str], bool]
_provider: LosslessProvider | None = None _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: def set_lossless_provider(
"""Register (or clear, with ``None``) the lossless compaction provider.""" provider: LosslessProvider | None,
global _provider *,
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 _provider = provider
_verifier = verifier if provider is not None else None
_generation += 1
def get_lossless_provider() -> LosslessProvider | None: def get_lossless_provider() -> LosslessProvider | None:
"""Return the registered provider, or ``None`` if the built-in should run.""" """Return the registered provider, or ``None`` if the built-in should run."""
return _provider 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

View file

@ -385,6 +385,9 @@ def test_v1_compress_then_v1_retrieve_resolves_marker_hash() -> None:
] ]
req = { req = {
"model": "gpt-4o", "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": [ "messages": [
{"role": "user", "content": "Get items"}, {"role": "user", "content": "Get items"},
{ {

View file

@ -9,8 +9,17 @@ its byte-exact fold — the lossless floor is never discarded by a later lossy
stage. 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_compaction import search_unheading
from headroom.transforms.lossless_provider import (
get_lossless_provider,
get_lossless_verifier,
set_lossless_provider,
)
def _grep_block() -> str: 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 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(): 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. # 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 # 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 was is False
assert rc.get("lossless_noop", 0) >= 1 assert rc.get("lossless_noop", 0) >= 1
assert rc.get("ratio_too_high", 0) == 0 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)

View file

@ -68,7 +68,10 @@ def test_v1_compress_success_reports_actual_metrics(monkeypatch) -> None:
markers_inserted=["marker-1"], 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: with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client:
response = client.post( response = client.post(

View file

@ -428,6 +428,211 @@ class TestCompressEndpointLossyInlineMode:
assert response.json()["ccr_hashes"] == [] 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: class TestCompressEndpointDoesNotBlockLoop:
"""/v1/compress must offload to the compression executor so a slow/large """/v1/compress must offload to the compression executor so a slow/large
payload cannot freeze the single event loop (#718).""" payload cannot freeze the single event loop (#718)."""
@ -477,6 +682,9 @@ class TestCompressEndpointDoesNotBlockLoop:
json={ json={
"messages": [{"role": "user", "content": "hello world"}], "messages": [{"role": "user", "content": "hello world"}],
"model": "gpt-4", "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"},
}, },
) )
) )

View file

@ -36,7 +36,14 @@ def _make_proxy(*, optimize: bool):
image_optimize=False, image_optimize=False,
subscription_tracking_enabled=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: class _FastTransform: