From f2c48e26c684a31e2802de9f49ce2075ef9cbf4b Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Fri, 31 Jul 2026 12:31:38 -0700 Subject: [PATCH] feat(compress): reach the lossless provider seam on the general path and default /v1/compress to marker-free output (#2691) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 checks passed! $ ruff format --check 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. --- headroom/proxy/handlers/openai.py | 145 ++++++++-- headroom/proxy/server.py | 22 +- headroom/transforms/content_router.py | 187 +++++++++++- headroom/transforms/lossless_provider.py | 73 ++++- tests/test_ccr_row_drop_store_bridge.py | 3 + tests/test_lossless_first_dispatch.py | 272 +++++++++++++++++- .../test_platform_stabilization_functional.py | 5 +- tests/test_proxy_compress_endpoint.py | 208 ++++++++++++++ tests/test_proxy_eager_preload_bind.py | 9 +- 9 files changed, 883 insertions(+), 41 deletions(-) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 62b46c03e..f17ecbba7 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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 ``<>`` / ``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, diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index acb769294..f7c34f4f3 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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 diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 75ee193c4..af806406f 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -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 diff --git a/headroom/transforms/lossless_provider.py b/headroom/transforms/lossless_provider.py index 36e3d3ece..f2b8e6d3b 100644 --- a/headroom/transforms/lossless_provider.py +++ b/headroom/transforms/lossless_provider.py @@ -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 diff --git a/tests/test_ccr_row_drop_store_bridge.py b/tests/test_ccr_row_drop_store_bridge.py index 9866ddceb..41a3d1926 100644 --- a/tests/test_ccr_row_drop_store_bridge.py +++ b/tests/test_ccr_row_drop_store_bridge.py @@ -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"}, { diff --git a/tests/test_lossless_first_dispatch.py b/tests/test_lossless_first_dispatch.py index 04ccd729b..0238bb95e 100644 --- a/tests/test_lossless_first_dispatch.py +++ b/tests/test_lossless_first_dispatch.py @@ -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) diff --git a/tests/test_platform_stabilization_functional.py b/tests/test_platform_stabilization_functional.py index f89b1a296..92ab78d7c 100644 --- a/tests/test_platform_stabilization_functional.py +++ b/tests/test_platform_stabilization_functional.py @@ -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( diff --git a/tests/test_proxy_compress_endpoint.py b/tests/test_proxy_compress_endpoint.py index 9ae358341..ac1be3c10 100644 --- a/tests/test_proxy_compress_endpoint.py +++ b/tests/test_proxy_compress_endpoint.py @@ -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"}, }, ) ) diff --git a/tests/test_proxy_eager_preload_bind.py b/tests/test_proxy_eager_preload_bind.py index 45b05c549..6f01a6981 100644 --- a/tests/test_proxy_eager_preload_bind.py +++ b/tests/test_proxy_eager_preload_bind.py @@ -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: