From 6422a80a58010da805d4001e83265300aa716d8a Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Mon, 3 Aug 2026 12:20:33 -0700 Subject: [PATCH] fix(compress): resolve the /v1/compress tokenizer per model, and document the real contract (#2743) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `/v1/compress` does no format conversion — callers send whichever wire shape they already use — but the pipeline pinned **one provider's token counter for the whole route**. `OpenAITokenCounter.count_message` walks list content for `text` and `image_url` only and has **no else branch**, so Anthropic content blocks contributed literally zero. A 599-token `tool_result` scored 8. A request that really removed 235 characters reported `tokens_saved: 0` — so a caller gating on `tokens_saved > 0` concludes compression is broken while it is working. Prompted by a Kong integration question ("do you support the Anthropic native format?"). The answer is that we already did — we just reported zeros for it, and the docs said otherwise. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [x] Documentation update ## Changes Made ### Tokenizer resolution (no hardcoded lists) Build the derived pipelines with `provider=None` so `TransformPipeline` resolves the tokenizer from the **per-model registry**. Every registry tokenizer derives from `BaseTokenizer`, whose `_count_content_parts` ends in a serialize-and-count catch-all, which means: - No block type counts as zero, and there is **no per-provider block-type list to keep in sync**. An enumerated set was the first thing I tried and it already missed `mcp_tool_result`, `web_search_tool_result`, `document`, and `thinking`. - Gemini / Mistral / DeepSeek / Kimi stop defaulting to a tiktoken count when the registry already has a calibrated counter for them. - Gateway aliases matching no vendor pattern still count correctly. `mode="ccr"` now runs a derived pipeline too, for the same reason — sharing `openai_pipeline` pinned its provider. Costs that mode its own cold compression cache; correct metrics win. ### Tokenizer selection stays separate from context-limit resolution Deliberately not welded together. `model_limit` feeds `context_pressure -> min_ratio`, so letting a tokenizer decision pick the limit table changes compression aggressiveness: `gpt-4-32k` answered by the Anthropic table is **8,192 instead of 32,768**, a 4× under-estimate. `test_tokenizer_choice_does_not_move_the_context_limit` pins the independence. ### Docs, rewritten from the code - **`proxy.mdx`** — the loopback-only default and **404-not-403** behavior, previously undocumented *anywhere* in `docs/` despite shipping in #2458 explicitly for gateway sidecars; `HEADROOM_COMPRESS_ALLOW_REMOTE`; all four request fields; the whole `config` object including every `mode` value and `frozen_message_count`; `transforms_summary`; the 400/401/404/503 contract; and the timeout fail-open shape (`compression_skipped` / `skip_reason`). - **Corrected "never calls an LLM"** — accurate about *generative* provider requests, misleading for a sidecar operator. Kompress (a ModernBERT **encoder**, classification not generation) and Magika run **in-process**, and `HEADROOM_KOMPRESS_ENDPOINT` offloads inference over HTTP — **real egress**. Now stated explicitly, with `HEADROOM_DISABLE_KOMPRESS=1` as the structural-only option. - **Both wire formats documented as accepted**, and removed `anthropic-sdk.mdx`'s claim that OpenAI format is "the compression engine's native format" — the exact misconception that prompted this work. The SDK's conversion is now framed as an SDK choice, not an API requirement. - **`litellm.mdx`** had no mention of the endpoint at all, despite the code naming LiteLLM's guardrail as its primary consumer. Added the HTTP deployment path, the `HEADROOM_COMPRESS_ALLOW_REMOTE` requirement, and why to leave `config.mode` unset. - **`index.mdx`** printed `compressionRatio * 100` labelled "Saved …%", so a 77% saving displayed as **23%**. `api-reference.mdx` already defined it correctly, so the docs contradicted each other. - `openai-sdk.mdx`, `wiki/proxy.md`, `wiki/typescript-sdk.md` — same corrections; dropped "any HTTP client", "Cloud", and a CacheAligner claim (it is detector-only). ## 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 $ .venv/bin/ruff check headroom/ tests/ --exclude headroom/dashboard/templates All checks passed! $ .venv/bin/mypy headroom/ Success: no issues found in 511 source files $ python -m pytest tests/test_compress_route_tokenizer_by_model.py \ tests/test_proxy_compress_endpoint.py tests/test_compress_api.py \ tests/test_platform_stabilization_functional.py tests/test_proxy_eager_preload_bind.py -q 99 passed, 2 warnings in 47.15s ``` Broader sweep (`-k "compress or litellm or gateway or guardrail"`): **1625 passed, 4 failed** — all 4 pre-existing, verified by stashing this diff and re-running on clean `main` (2 strands hook tests, 1 codex WS semaphore-tail timing test, 1 unrelated local WIP test). ## Real Behavior Proof - **Environment:** macOS 26.4 arm64, Python 3.12.6, repo `.venv`, branch rebased on `upstream/main`. **(1) Before → after, same request** (60-line grep payload in an Anthropic `tool_result`): | model | before | after | | --- | --- | --- | | `claude-sonnet-4-6` | `before=28 saved=0` | `before=1223 saved=58` | | `bedrock/anthropic.claude-3-5-sonnet` | `saved=0` | `before=1037 saved=59` | | `my-gateway/big-model` (alias) | `saved=0` | `before=1037 saved=59` | | `gemini-2.5-pro` | `saved=0` | `before=1036 saved=59` | | `gpt-4o` + OpenAI shape | `before=1225 saved=58` | `before=1225 saved=58` (unchanged) | All three `config.mode` values verified for each. Response shape preserved: `type=tool_result`, `tool_use_id` intact. **(2) Counter-level root cause**, 6.8 KB body, `count_message()`: ```text OpenAITokenCounter string-content -> 1406 tool_result block -> 5 registry (BaseTokenizer) claude tool_result=408 thinking=418 mcp_tool_result=421 web_search_tool_result=421 document=422 ``` **(3) Every documented behavior asserted against the running app** — 13 checks, all PASS: 400s for missing `messages`/`model`, invalid `config.mode`, and all four invalid `frozen_message_count` forms; 200 for valid ones; non-dict `config` ignored; bypass and empty-messages omit `transforms_summary`; success returns exactly the 8 documented keys. - **Not tested:** the docs site was not built (`docs/node_modules` absent) — MDX was checked for balanced `` tags only, so a reviewer with the site running should eyeball rendering. No live gateway/Kong request; verification is via `TestClient` against the real ASGI app. - **Note:** `HEADROOM_DISABLE_KOMPRESS` is read into `ProxyConfig` at `server.py:4919` and by the CLI, not by `create_app(ProxyConfig(...))` directly — I confirmed `disable_kompress=True` does reach the derived pipeline. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- docs/content/docs/anthropic-sdk.mdx | 6 +- docs/content/docs/index.mdx | 3 +- docs/content/docs/litellm.mdx | 21 ++ docs/content/docs/openai-sdk.mdx | 4 +- docs/content/docs/proxy.mdx | 146 ++++++++- headroom/proxy/handlers/openai.py | 66 +++- .../test_compress_route_tokenizer_by_model.py | 289 ++++++++++++++++++ tests/test_proxy_compress_endpoint.py | 7 +- wiki/proxy.md | 57 +++- wiki/typescript-sdk.md | 6 +- 10 files changed, 566 insertions(+), 39 deletions(-) create mode 100644 tests/test_compress_route_tokenizer_by_model.py diff --git a/docs/content/docs/anthropic-sdk.mdx b/docs/content/docs/anthropic-sdk.mdx index 150c0021c..8dc6fd65e 100644 --- a/docs/content/docs/anthropic-sdk.mdx +++ b/docs/content/docs/anthropic-sdk.mdx @@ -41,7 +41,7 @@ Every call to `client.messages.create()` compresses messages first. The response `withHeadroom()` returns a proxy around your Anthropic client that intercepts `messages.create()`: -1. Converts Anthropic-format messages to OpenAI format (the compression engine's native format) +1. Converts Anthropic-format messages to OpenAI format 2. Sends them to the Headroom proxy's `/v1/compress` endpoint 3. Converts the compressed messages back to Anthropic format 4. Forwards the request to Anthropic as normal @@ -58,6 +58,10 @@ The adapter handles the full Anthropic message format including content blocks: This conversion is lossless. Your request and response behave identically to an unwrapped client. + +`POST /v1/compress` does no format conversion and compresses Anthropic content blocks natively — see [Message format](/docs/proxy#message-format). If you are calling the endpoint directly (from a gateway, or LiteLLM's `headroom` guardrail), send Anthropic-shaped messages as-is; you get the same shape back. Only this TypeScript adapter converts, because it normalises on OpenAI types internally. + + ## Options Pass compression options as the second argument: diff --git a/docs/content/docs/index.mdx b/docs/content/docs/index.mdx index 75d3d2a84..3f2471687 100644 --- a/docs/content/docs/index.mdx +++ b/docs/content/docs/index.mdx @@ -17,7 +17,8 @@ const messages = [ ]; const result = await compress(messages, { model: 'gpt-4o' }); -console.log(`Saved ${result.tokensSaved} tokens (${(result.compressionRatio * 100).toFixed(0)}%)`); +// compressionRatio is tokensAfter / tokensBefore, so savings is 1 - ratio. +console.log(`Saved ${result.tokensSaved} tokens (${((1 - result.compressionRatio) * 100).toFixed(0)}%)`); ``` diff --git a/docs/content/docs/litellm.mdx b/docs/content/docs/litellm.mdx index aa10b5509..b7c174604 100644 --- a/docs/content/docs/litellm.mdx +++ b/docs/content/docs/litellm.mdx @@ -94,3 +94,24 @@ app.add_middleware(CompressionMiddleware) ``` Response headers include `x-headroom-compressed: true` and `x-headroom-tokens-saved: 1234`. + +## Over HTTP (guardrail / gateway) + +The options above run Headroom **in** the LiteLLM process. If instead LiteLLM runs as its own proxy and you want it to call Headroom over the network — the guardrail deployment — point it at [`POST /v1/compress`](/docs/proxy#post-v1compress). LiteLLM swaps `messages` for the compressed result and forwards to the provider. + +Two things this deployment needs: + +```bash +# Headroom is loopback-only by default and answers remote callers with 404. +HEADROOM_COMPRESS_ALLOW_REMOTE=1 headroom proxy +``` + + +Without `HEADROOM_COMPRESS_ALLOW_REMOTE=1` a remote caller gets `404`, not `403` — so a misconfigured guardrail looks exactly like a wrong URL. If you also set `HEADROOM_PROXY_TOKEN`, send it as `X-Headroom-Proxy-Token` or you get `401`. + + +Leave `config.mode` unset. The default pipeline is marker-free, which is what a forward-only caller wants: `mode: "ccr"` emits retrieval markers that are a dangling pointer unless you also inject the `headroom_retrieve` tool and can reach `/v1/retrieve`. + +Because LiteLLM passes model names through, send the real one — `claude-sonnet-4-6`, `bedrock/anthropic.claude-3-5-sonnet`, `gemini-2.5-pro` — so Headroom resolves the right tokenizer and context limit. Anthropic-shaped messages need no conversion; see [Message format](/docs/proxy#message-format). + +For multi-turn agent loops, set `config.frozen_message_count` to the number of messages the provider has already cached, **and send back the messages you previously forwarded rather than the pristine originals**. Getting this wrong silently destroys the provider's prefix cache — see [Multi-turn usage](/docs/proxy#multi-turn-usage-keeping-the-prefix-cache) for the loop. diff --git a/docs/content/docs/openai-sdk.mdx b/docs/content/docs/openai-sdk.mdx index 399367dfd..7a3bbefc3 100644 --- a/docs/content/docs/openai-sdk.mdx +++ b/docs/content/docs/openai-sdk.mdx @@ -42,10 +42,12 @@ That's it. Every call to `client.chat.completions.create()` compresses the messa `withHeadroom()` returns a proxy around your OpenAI client that intercepts `chat.completions.create()`: 1. Extracts `messages` from the request params -2. Sends them to the Headroom proxy's `/v1/compress` endpoint +2. Sends them to the Headroom proxy's [`POST /v1/compress`](/docs/proxy#post-v1compress) endpoint 3. Replaces the original messages with the compressed result 4. Forwards the request to OpenAI as normal +The SDK talks to a **local** proxy, which is why no extra configuration is needed: `/v1/compress` is loopback-only by default. If you move the proxy to another host, set `HEADROOM_COMPRESS_ALLOW_REMOTE=1` on it or requests come back `404`. + All other client methods are untouched: ```ts twoslash diff --git a/docs/content/docs/proxy.mdx b/docs/content/docs/proxy.mdx index 5de6a334d..aa9c1dc61 100644 --- a/docs/content/docs/proxy.mdx +++ b/docs/content/docs/proxy.mdx @@ -294,6 +294,7 @@ See [Metrics](/docs/metrics) for the Prometheus and Grafana setup. | Flag / env | Default | Effect | |---|---|---| | `HEADROOM_PROXY_TOKEN` | none | Require a bearer token (`X-Headroom-Proxy-Token`) from non-loopback callers. | +| `HEADROOM_COMPRESS_ALLOW_REMOTE` | `false` | Allow non-loopback callers to reach [`POST /v1/compress`](#post-v1compress). Required to run Headroom as a gateway/sidecar; without it remote callers get `404`. | | `--offline` / `HEADROOM_OFFLINE` | `false` | Air-gap mode: hard-disable **all** egress (telemetry, update checks, license reporting, model downloads). | | `--stateless` / `HEADROOM_STATELESS` | `false` | Keep all state in memory; no filesystem writes (disables logs, memory, TOIN). | | `HEADROOM_STRIP_INTERNAL_HEADERS` | `enabled` | Strip internal `x-headroom-*` headers before forwarding upstream. | @@ -410,17 +411,79 @@ The proxy also accepts: ### `POST /v1/compress` -Compression-only endpoint. Compresses messages without calling any LLM. Used by the TypeScript SDK. +Compression-only endpoint. Compresses messages and returns them without ever making a **completion request to an LLM provider** — no generation, no provider API key, no upstream chat call. Used by the TypeScript SDK, by LiteLLM's `headroom` guardrail, and by API gateways running Headroom as a sidecar. -**Request:** -```json -{ - "messages": [{ "role": "user", "content": "..." }], - "model": "gpt-4o" -} -``` + +"No LLM call" means no *generative* request to a provider. Compression itself is ML-backed: **Kompress** is a ModernBERT encoder that scores tokens for retention (classification, not generation), and Magika classifies content types. Both run in-process by default, so budget CPU and memory for the sidecar accordingly. + +If `HEADROOM_KOMPRESS_ENDPOINT` is set, Kompress inference is offloaded over HTTP to that model server — **real egress from the sidecar**, which matters if you deployed it expecting none. Only inference goes remote: the CCR store and retrieval markers stay proxy-local, and original content never persists off-box. Leave the variable unset to keep everything in-process, or run with `HEADROOM_DISABLE_KOMPRESS=1` for structural compression only. + + + +This route is restricted to loopback callers and answers everyone else with **`404`**, not `403` — deliberately, so it stays invisible to external scanners. A gateway calling it from another host or pod therefore sees what looks like a missing route. + +Both the client IP and the inbound `Host:` header must name loopback. To allow remote callers, set `HEADROOM_COMPRESS_ALLOW_REMOTE=1`. `HEADROOM_PROXY_TOKEN` still applies if set. + + +#### Message format + +The endpoint does **no format conversion**. Whatever shape you send in `messages` is the shape you get back, and both wire formats are compressed natively: + +- **OpenAI shape** — `role: "tool"` messages with `tool_call_id`, assistant `tool_calls` +- **Anthropic shape** — content-block lists with `tool_use` / `tool_result` / `thinking` blocks + +So an Anthropic-native caller does not need to convert to OpenAI format first. Block types, `tool_use_id`s and message order are all preserved. + +`model` selects the tokenizer (per-model, from Headroom's tokenizer registry) and the context limit. Send the real model name — including gateway-prefixed forms like `bedrock/anthropic.claude-3-5-sonnet` or `vertex_ai/claude-sonnet-4@20250514` — so token counts and compression aggressiveness are right. + +#### Request + +| Field | Type | Required | Description | +| --- | --- | --- | --- | +| `messages` | array | yes | Messages to compress, in either wire format. `400` if missing. Empty array returns immediately with zero metrics. | +| `model` | string | yes | Model name. Drives tokenizer + context-limit resolution. `400` if missing. | +| `token_budget` | integer | no | Overrides the model's context limit. Used by callers that need to fit a tighter budget. | +| `config` | object | no | Compression options, below. A non-object value is ignored rather than rejected. | + + +Only the four fields above are read. Anthropic sends `system` and `tools` **out of band**, alongside `messages` — this endpoint accepts them without complaint (you get a `200`, no warning) and returns neither, so neither is compressed. + +Keep carrying both yourself and send them upstream unchanged. Two consequences worth knowing: + +- An Anthropic system prompt is not compressed here, even though it is resent on every request. +- Tool-schema compaction and tool-search deferral are not reachable through this endpoint — on tool-heavy traffic those can be the largest share of available savings. Run Headroom as the proxy (rather than calling `/v1/compress`) if you need them. + + +`config` fields: + +| Field | Type | Default | Description | +| --- | --- | --- | --- | +| `mode` | string | unset | `ccr`, `lossy_inline`, or `lossless_then_lossy`. Unset selects the default marker-free pipeline. Any other value is a `400`. | +| `frozen_message_count` | integer | unset | Pin a prefix: the first N messages are returned byte-for-byte unchanged while staying visible to cross-message transforms like dedup. Set it to the number of messages the provider has already cached so compression cannot rewrite the prefix and bust that cache. Must be a non-negative integer; anything else is a `400`. | +| `compress_user_messages` | boolean | `false` | Also compress user-role messages. | +| `target_ratio` | number | unset | Target compression ratio. | +| `protect_recent` | integer | unset | Leave the last N messages uncompressed. | +| `protect_analysis_context` | boolean | unset | Preserve analysis context blocks. | + +**`config.mode` values:** + +- **unset (default)** — marker-free. Emits no `<>` retrieval markers and writes nothing to the CCR store, so you can forward the returned messages straight to a provider. This is the right mode for a gateway or guardrail that just swaps `messages` and forwards. +- **`ccr`** — emits CCR markers and writes to the store. Only for callers that also inject the `headroom_retrieve` tool *and* can reach `/v1/retrieve` (itself loopback-only). Markers are a dangling pointer for the model otherwise. +- **`lossy_inline`** (alias `lossless_then_lossy`) — runs the lossless byte/data fold first, then compresses the folded remainder. Marker-free. + +#### Response + +| Field | Type | Description | +| --- | --- | --- | +| `messages` | array | Compressed messages, in the shape you sent. | +| `tokens_before` | integer | Token count before compression. | +| `tokens_after` | integer | Token count after compression. | +| `tokens_saved` | integer | `tokens_before - tokens_after`. | +| `compression_ratio` | number | `tokens_after / tokens_before` — so **lower is better**. A ratio of `0.23` means a 77% reduction, not 23%. `1.0` when nothing was compressed. | +| `transforms_applied` | array | Transform labels that ran. | +| `transforms_summary` | object | Per-transform counts. | +| `ccr_hashes` | array | Retrieval hashes for markers inserted (empty unless `mode: "ccr"`). | -**Response:** ```json { "messages": [{ "role": "user", "content": "..." }], @@ -429,11 +492,72 @@ Compression-only endpoint. Compresses messages without calling any LLM. Used by "tokens_saved": 11500, "compression_ratio": 0.23, "transforms_applied": ["router:smart_crusher:0.35"], - "ccr_hashes": ["a1b2c3"] + "transforms_summary": { "router:smart_crusher:0.35": 1 }, + "ccr_hashes": [] } ``` -Set `x-headroom-bypass: true` to skip compression. +#### Headers + +`x-headroom-bypass: true` (case-insensitive) skips compression entirely and echoes your messages back with zeroed metrics. The bypass and empty-messages responses omit `transforms_summary`. + +#### Errors and fail-open + +| Status | Body | When | +| --- | --- | --- | +| `400` | `error.type = "invalid_request"` | Missing `messages` or `model`, malformed JSON, invalid `config.mode`, or invalid `config.frozen_message_count`. | +| `401` | — | `HEADROOM_PROXY_TOKEN` is set and the bearer token is missing or wrong. | +| `404` | — | Non-loopback caller without `HEADROOM_COMPRESS_ALLOW_REMOTE=1`. | +| `503` | `error.type = "compression_error"` | Compression failed unexpectedly. | + +Compression **fails open on timeout**: you get `200` with your original messages, zeroed metrics, plus `compression_skipped: true` and `skip_reason: "compression_timeout"`. Always check `compression_skipped` if you need to know whether compression actually ran. + +Requests are recorded under `provider="compress"` in `/stats` and `/metrics`. + +#### Multi-turn usage: keeping the prefix cache + +This is the single most important thing to get right, and the default is not safe for an agent loop. + +When Headroom proxies a request itself it watches the provider's cache hit rate turn over turn and freezes the already-cached prefix. `/v1/compress` **cannot do that — it is stateless.** It sees one isolated call and has no idea what the provider already cached. + +The provider caches the bytes you **forwarded**. Compression changed those bytes, so your original messages and the ones the provider cached are no longer the same thing — and it is the forwarded version you have to keep reproducing. Send the pristine originals again next turn and the provider sees a different prefix and re-reads it from scratch. On Anthropic a cache read is ~90% cheaper than fresh input, so that can easily cost more than the compression saves. + +Compression is also not uniform over a conversation: how hard a message is compressed depends partly on how far it now sits from the end, so an older tool result can fall outside the recent-read protection window as the conversation grows and be compressed harder than it was last turn. Another reason not to rely on re-compression reproducing earlier output. + +Two rules: + +1. **Pass `config.frozen_message_count`** — how many leading messages the provider has already cached. +2. **Send back your own previous output, not the original messages.** `frozen_message_count` returns those leading messages *exactly as you passed them in* — it pins whatever you hand it. Hand it pristine originals and you get pristine originals back, which is precisely the prefix the provider does not have. + +```python +# Keep what you FORWARDED, not what you started with. +forwarded: list[dict] = [] + +def next_turn(new_messages: list[dict]) -> list[dict]: + body = { + "messages": forwarded + new_messages, + "model": "claude-sonnet-4-6", + # Everything already forwarded is already cached upstream — pin it. + "config": {"frozen_message_count": len(forwarded)}, + } + result = requests.post(f"{proxy}/v1/compress", json=body).json() + forwarded[:] = result["messages"] # becomes next turn's frozen prefix + return forwarded +``` + + +Compressing the full original conversation on each turn looks correct — you get a `200` and a positive `tokens_saved` — but the leading messages come back different from the ones the provider cached. You pay for compression *and* for a cache miss. Nothing in the response tells you this happened; watch your provider's cache-read tokens. + + +Also for multi-turn callers: + +- **Leave `config.mode` unset.** The default is marker-free, which is what a forward-only caller wants. +- **Send the real model name** so the tokenizer and context limit resolve correctly — including gateway-prefixed forms. +- **`protect_recent` is not a substitute.** It guards the newest messages; `frozen_message_count` guards the oldest, which is the cached end. + + +`HEADROOM_KOMPRESS_ENDPOINT` points *outbound* at a remote Kompress ML model server that happens to expose a `/compress` path. It is unrelated to this inbound endpoint. + ## Agent wrapping diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index d2a2e39f1..736455386 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -8325,6 +8325,28 @@ class OpenAIHandlerMixin: marker-free path and vice versa. Sharing would be a correctness bug, not an optimisation — the duplicate cache is the intended trade. + Built with ``provider=None`` on purpose: ``TransformPipeline`` then + resolves the tokenizer from the per-model registry + (``headroom.tokenizers.get_tokenizer``) on every call instead of pinning + one provider's counter for the whole route. + + That matters because this route does no format conversion — callers send + whichever wire shape they already use. The provider counters walk only the + block types they own (``OpenAITokenCounter`` handles ``text`` and + ``image_url``; everything else falls through with NO else branch, so an + Anthropic ``tool_result`` contributed literally zero and ``tokens_saved`` + reported 0 on requests where compression really ran). Every registry + tokenizer derives from ``BaseTokenizer``, whose ``_count_content_parts`` + ends in a serialize-and-count catch-all, so no block type counts as zero + and there is no per-provider type list to keep in sync. It also stops + defaulting Gemini/Mistral/DeepSeek/Kimi traffic to a tiktoken count when + the registry already has a calibrated counter for them. + + Note the pipeline's own ``_provider_name()`` is therefore ``None`` here. + That is honest for a provider-agnostic route — it used to report + ``"openai"`` even for Claude payloads — and the request outcome still + records ``provider="compress"`` for /stats and /metrics. + ponytail: a first-request race just builds it twice — both are equivalent and Kompress weights are cached at module level, so no lock. """ @@ -8344,7 +8366,7 @@ class OpenAIHandlerMixin: return self.openai_pipeline pipeline = TransformPipeline( transforms=[ContentRouter(replace(base.config, **overrides), observer=self.metrics)], - provider=self.openai_provider, + provider=None, # per-model registry tokenizer — see docstring ) cache[key] = pipeline return pipeline @@ -8372,6 +8394,16 @@ class OpenAIHandlerMixin: ccr_enabled=False, # no CCR store writes ) + def _ccr_pipeline(self) -> Any: + """Pipeline for ``/v1/compress`` ``config.mode="ccr"``. + + No config overrides: markers and CCR store writes are exactly what this + mode asks for, so it inherits the live router's settings verbatim. It is + still a DERIVED pipeline rather than ``openai_pipeline`` so the tokenizer + comes from the per-model registry instead of a pinned provider counter. + """ + return self._derived_compress_pipeline("ccr") + def _lossy_inline_pipeline(self) -> Any: """Pipeline for ``/v1/compress`` ``config.mode="lossy_inline"``. @@ -8503,17 +8535,21 @@ 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. + # Resolve the CONTEXT LIMIT against the model's own provider. This + # route accepts either wire shape, and LiteLLM's `headroom` guardrail + # passes Anthropic model names straight through + # (claude-sonnet-4-5-..., bedrock/anthropic.claude-3-5-sonnet, + # anthropic/claude-opus-4), where the OpenAI provider answers with its + # 128K default instead of 200K+. Substring match is enough — 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. + # Deliberately separate from tokenizer selection, which the derived + # pipelines resolve per model from the tokenizer registry (see + # _derived_compress_pipeline). Welding the two together would let a + # tokenizer decision move `model_limit`, and `model_limit` feeds + # context_pressure -> min_ratio: e.g. gpt-4-32k answered by the + # Anthropic table is 8,192 instead of 32,768, a 4x under-estimate that + # silently changes compression aggressiveness. model_name = model if isinstance(model, str) else str(model) limit_provider = ( self.anthropic_provider @@ -8578,7 +8614,13 @@ class OpenAIHandlerMixin: if mode in ("lossy_inline", "lossless_then_lossy"): pipeline = self._lossy_inline_pipeline() elif mode == "ccr": - pipeline = self.openai_pipeline + # Markers + store writes wanted, so inherit the live router config + # unchanged — but as a DERIVED pipeline, not `openai_pipeline` + # itself, so the per-model registry tokenizer applies here too. + # Sharing the request path's instance pinned the OpenAI counter, + # which reports zero for Anthropic content blocks. Costs this mode + # its own (cold) compression cache; correct metrics win. + pipeline = self._ccr_pipeline() else: pipeline = self._no_ccr_pipeline() diff --git a/tests/test_compress_route_tokenizer_by_model.py b/tests/test_compress_route_tokenizer_by_model.py new file mode 100644 index 000000000..ab1ee9bcc --- /dev/null +++ b/tests/test_compress_route_tokenizer_by_model.py @@ -0,0 +1,289 @@ +"""`/v1/compress` must count tokens for ANY wire shape and ANY model family. + +The route does no format conversion — callers send whichever shape they already +use. Pinning one provider's token counter for the whole route silently reported +zero savings for Anthropic-shaped payloads: ``OpenAITokenCounter.count_message`` +walks list content for ``text`` and ``image_url`` only, and has no else branch, +so an Anthropic ``tool_result`` block contributed literally nothing. A real +request that removed 235 characters reported ``tokens_saved: 0``. + +The derived pipelines are built with ``provider=None`` so ``TransformPipeline`` +resolves the tokenizer from the per-model registry instead. Every registry +tokenizer derives from ``BaseTokenizer``, whose ``_count_content_parts`` ends in +a serialize-and-count catch-all, so no block type counts as zero and there is no +per-provider block-type list to keep in sync. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("fastapi") + +from fastapi.testclient import TestClient + +from headroom.proxy.server import ProxyConfig, create_app + +# Big enough that any real tokenizer must report hundreds of tokens, and +# compressible so the router actually folds it (repeated grep-shaped lines). +_GREP = "\n".join( + f"src/module_{i}.py:{i * 7}: result = compute_value(item_{i}, flag=True)" for i in range(60) +) + + +@pytest.fixture +def client(): + app = create_app( + ProxyConfig( + optimize=True, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + ) + # /v1/compress is loopback-gated (#1227). + with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as c: + yield c + + +def _anthropic_messages() -> list[dict]: + """Anthropic native shape: tool_use / tool_result content blocks.""" + return [ + {"role": "user", "content": [{"type": "text", "text": "find compute_value"}]}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "c1", "name": "grep", "input": {"pattern": "compute"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "c1", "content": _GREP}], + }, + ] + + +def _openai_messages() -> list[dict]: + """OpenAI native shape: tool_calls + role=tool.""" + return [ + {"role": "user", "content": "find compute_value"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "grep", "arguments": '{"pattern":"compute"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": _GREP}, + ] + + +# Model names spanning every routing path a gateway realistically sends, including +# a custom alias that matches no known vendor pattern. +_MODELS = [ + "claude-sonnet-4-6", + "bedrock/anthropic.claude-3-5-sonnet", + "vertex_ai/claude-sonnet-4@20250514", + "gemini-2.5-pro", + "deepseek/deepseek-v4", + "moonshotai/kimi-k2", + "my-gateway/big-model", + "gpt-4o", +] + + +@pytest.mark.parametrize("model", _MODELS) +def test_anthropic_shape_is_counted_for_every_model_family(client, model): + """No model name may produce a zero token count for Anthropic content blocks.""" + response = client.post("/v1/compress", json={"messages": _anthropic_messages(), "model": model}) + + assert response.status_code == 200, response.text + body = response.json() + # The payload is ~4 KB of text. Any honest tokenizer reports hundreds; the + # pinned OpenAI counter reported 28 for this exact request. + assert body["tokens_before"] > 500, f"{model} undercounted: {body['tokens_before']}" + assert body["tokens_saved"] > 0, f"{model} reported no savings: {body}" + assert body["compression_ratio"] < 1.0 + + +@pytest.mark.parametrize("model", _MODELS) +def test_openai_shape_still_counted_for_every_model_family(client, model): + """The OpenAI-shaped path must not regress while fixing the Anthropic one.""" + response = client.post("/v1/compress", json={"messages": _openai_messages(), "model": model}) + + assert response.status_code == 200, response.text + body = response.json() + assert body["tokens_before"] > 500, f"{model} undercounted: {body['tokens_before']}" + assert body["tokens_saved"] > 0, f"{model} reported no savings: {body}" + + +@pytest.mark.parametrize("mode", [None, "ccr", "lossy_inline", "lossless_then_lossy"]) +def test_every_mode_counts_anthropic_shape(client, mode): + """mode="ccr" used to share the request pipeline, which pinned the OpenAI counter.""" + payload: dict = {"messages": _anthropic_messages(), "model": "claude-sonnet-4-6"} + if mode is not None: + payload["config"] = {"mode": mode} + + response = client.post("/v1/compress", json=payload) + + assert response.status_code == 200, response.text + body = response.json() + assert body["tokens_before"] > 500, f"mode={mode} undercounted: {body['tokens_before']}" + assert body["tokens_saved"] > 0, f"mode={mode} reported no savings: {body}" + + +def test_response_preserves_the_anthropic_wire_shape(client): + """Passthrough contract: no format conversion in either direction.""" + response = client.post( + "/v1/compress", + json={"messages": _anthropic_messages(), "model": "claude-sonnet-4-6"}, + ) + + assert response.status_code == 200 + block = response.json()["messages"][2]["content"][0] + assert block["type"] == "tool_result" + assert block["tool_use_id"] == "c1" + # Content was folded, not dropped or restructured. + assert 0 < len(block["content"]) < len(_GREP) + + +def test_tokenizer_choice_does_not_move_the_context_limit(client): + """Regression guard: the two resolutions must stay independent. + + `model_limit` feeds context_pressure -> min_ratio, so letting a tokenizer + decision pick the limit table changes compression aggressiveness. gpt-4-32k + answered by the Anthropic table is 8,192 instead of 32,768 — a 4x + under-estimate — even though its payload needs a non-OpenAI tokenizer. + """ + seen: dict = {} + proxy = client.app.state.proxy + original = proxy._no_ccr_pipeline().apply + + def spy(**kwargs): + seen.update(kwargs) + return original(**kwargs) + + proxy._no_ccr_pipeline().apply = spy + try: + response = client.post( + "/v1/compress", + json={"messages": _anthropic_messages(), "model": "gpt-4-32k"}, + ) + finally: + proxy._no_ccr_pipeline().apply = original + + assert response.status_code == 200 + # OpenAI's table, because the MODEL is an OpenAI model — regardless of the + # Anthropic-shaped body that drives tokenizer selection. + assert seen["model_limit"] == 32_768 + + +# ── The documented multi-turn recipe ────────────────────────────────────────── +# The endpoint is stateless: unlike the proxy's own request path it runs no +# CacheAligner and tracks no provider cache state, so keeping the prefix stable is +# the caller's job. docs/content/docs/proxy.mdx documents the loop; these two tests +# pin both halves of it so the guidance cannot rot. + + +def _turn(i: int) -> list[dict]: + body = "\n".join(f"src/mod_{i}_{j}.py:{j}: match compute_value(x{j})" for j in range(40)) + return [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": f"c{i}", "name": "grep", "input": {"p": "x"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"c{i}", "content": body}], + }, + ] + + +def _compress(client, messages: list[dict], frozen: int | None = None) -> list[dict]: + body: dict = {"messages": messages, "model": "claude-sonnet-4-6"} + if frozen is not None: + body["config"] = {"frozen_message_count": frozen} + response = client.post("/v1/compress", json=body) + assert response.status_code == 200, response.text + return response.json()["messages"] + + +def test_system_and_tools_are_accepted_and_ignored(client): + """Documented contract: only messages/model/token_budget/config are read. + + Anthropic sends `system` and `tools` out of band. The endpoint takes them + without complaint and returns neither, so neither is compressed — callers must + keep carrying them. Pinned because the silence is the hazard: a caller has no + signal that the fields did nothing. If this ever starts returning them, the + contract changed and docs/content/docs/proxy.mdx needs updating with it. + """ + response = client.post( + "/v1/compress", + json={ + "messages": _anthropic_messages(), + "model": "claude-sonnet-4-6", + "system": "You are a coding agent. " * 200, + "tools": [ + { + "name": "read", + "description": "Read a file from disk. " * 20, + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + } + ], + }, + ) + + assert response.status_code == 200, response.text + body = response.json() + assert "system" not in body + assert "tools" not in body + # Messages are still compressed normally alongside the ignored fields. + assert body["tokens_saved"] > 0 + + +def test_documented_loop_keeps_the_cached_prefix_byte_identical(client): + """Feed back previous OUTPUT + frozen_message_count -> stable prefix every turn.""" + import json + + forwarded = _compress(client, [{"role": "user", "content": [{"type": "text", "text": "go"}]}]) + for i in range(1, 6): + previous = forwarded + forwarded = _compress(client, previous + _turn(i), frozen=len(previous)) + replayed = forwarded[: len(previous)] + assert [json.dumps(m, sort_keys=True) for m in replayed] == [ + json.dumps(m, sort_keys=True) for m in previous + ], f"turn {i} rewrote the cached prefix" + + +def test_frozen_prefix_replays_what_you_sent_not_what_you_forwarded(client): + """Why re-sending pristine originals busts the cache — the documented trap. + + `frozen_message_count` returns the leading messages *exactly as passed in*. So + the bytes you get back depend entirely on which version you sent: feed it your + previous OUTPUT and the prefix matches what the provider cached; feed it the + pristine ORIGINALS and you hand the provider different bytes for a message it + already cached, paying for compression and a cache miss at once. + """ + import json + + base = [{"role": "user", "content": [{"type": "text", "text": "go"}]}] + originals = base + _turn(1) + + forwarded = _compress(client, originals) + # Precondition: compression actually changed the prefix, so the two candidate + # inputs for next turn genuinely differ. + assert json.dumps(forwarded) != json.dumps(originals) + + # Correct: previous output in, same bytes back. + good = _compress(client, forwarded + _turn(2), frozen=len(forwarded)) + assert good[: len(forwarded)] == forwarded + + # The trap: pristine originals in, pristine originals back — which is NOT what + # was forwarded last turn, so the provider's cached prefix no longer matches. + trap = _compress(client, originals + _turn(2), frozen=len(originals)) + assert trap[: len(originals)] == originals + assert trap[: len(forwarded)] != forwarded diff --git a/tests/test_proxy_compress_endpoint.py b/tests/test_proxy_compress_endpoint.py index 31c6162f0..53aa74141 100644 --- a/tests/test_proxy_compress_endpoint.py +++ b/tests/test_proxy_compress_endpoint.py @@ -671,7 +671,7 @@ class TestCompressEndpointDoesNotBlockLoop: markers_inserted=[], ) - monkeypatch.setattr(proxy.openai_pipeline, "apply", blocking_apply) + monkeypatch.setattr(proxy._ccr_pipeline(), "apply", blocking_apply) # /v1/compress is loopback-gated (#1227) — present as 127.0.0.1. transport = httpx.ASGITransport(app=app, client=("127.0.0.1", 12345)) @@ -682,8 +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. + # Every mode now runs a DERIVED pipeline so the tokenizer + # comes from the per-model registry; mode="ccr" is the + # marker-on one, patched above. "config": {"mode": "ccr"}, }, ) diff --git a/wiki/proxy.md b/wiki/proxy.md index 71d7faefa..184640916 100644 --- a/wiki/proxy.md +++ b/wiki/proxy.md @@ -2,7 +2,7 @@ The Headroom proxy server is a production-ready HTTP server that applies context optimization to all requests passing through it. -> **New:** The proxy now supports the [TypeScript SDK](typescript-sdk.md) via the `POST /v1/compress` endpoint, enabling compression-as-a-service for any HTTP client without calling an LLM. +> The proxy exposes compression-as-a-service via the `POST /v1/compress` endpoint — used by the [TypeScript SDK](typescript-sdk.md), LiteLLM's `headroom` guardrail, and gateway sidecars. It is loopback-only by default; see the endpoint section below. ## Starting the Proxy @@ -266,13 +266,30 @@ POST /v1/chat/completions ### `POST /v1/compress` -Compression-only endpoint. Compresses messages without calling any LLM. Used by the [TypeScript SDK](typescript-sdk.md) and any HTTP client that wants compression as a service. +Compression-only endpoint. Compresses messages without ever making a **completion request to an LLM provider** — no generation, no provider API key. Used by the [TypeScript SDK](typescript-sdk.md), LiteLLM's `headroom` guardrail, and gateway sidecars. + +**It does run local ML models.** Compression is ML-backed: Kompress is a ModernBERT encoder scoring tokens for retention (classification, not generation) and Magika classifies content types, both in-process by default. If `HEADROOM_KOMPRESS_ENDPOINT` is set, Kompress inference is offloaded over HTTP to that model server — real egress from the sidecar. Only inference goes remote; the CCR store and markers stay proxy-local. `HEADROOM_DISABLE_KOMPRESS=1` gives structural compression only. + +**Loopback-only by default.** Non-loopback callers get **404** (not 403 — the route stays invisible to scanners). Set `HEADROOM_COMPRESS_ALLOW_REMOTE=1` to allow remote callers. + +**No format conversion.** `messages` may be OpenAI-shaped (`role: "tool"` + `tool_call_id`) or Anthropic-shaped (`tool_use` / `tool_result` content blocks); the same shape comes back. `model` selects the tokenizer and context limit — send the real name, including gateway-prefixed forms like `bedrock/anthropic.claude-3-5-sonnet`. + +**`system` and `tools` are ignored.** Anthropic sends both out of band. This endpoint accepts them without complaint (200, no warning) and returns neither, so neither is compressed — keep carrying them yourself. That means the Anthropic system prompt is not compressed here, and tool-schema compaction / tool-search deferral are not reachable through this route; run Headroom as the proxy if you need those. **Request:** ```json { - "messages": [...], // OpenAI chat format - "model": "gpt-4o" // model name (for token counting) + "messages": [...], // either wire format + "model": "gpt-4o", // tokenizer + context limit + "token_budget": 8000, // optional: override the context limit + "config": { // optional + "mode": "lossy_inline", // ccr | lossy_inline | lossless_then_lossy + "frozen_message_count": 12, // pin an already-cached prefix + "compress_user_messages": false, + "target_ratio": 0.5, + "protect_recent": 2, + "protect_analysis_context": true + } } ``` @@ -283,16 +300,40 @@ Compression-only endpoint. Compresses messages without calling any LLM. Used by "tokens_before": 15000, "tokens_after": 3500, "tokens_saved": 11500, - "compression_ratio": 0.23, + "compression_ratio": 0.23, // tokens_after / tokens_before — LOWER is better "transforms_applied": ["router:smart_crusher:0.35"], - "ccr_hashes": ["a1b2c3"] + "transforms_summary": {"router:smart_crusher:0.35": 1}, + "ccr_hashes": [] // non-empty only with mode="ccr" } ``` **Headers:** -- `x-headroom-bypass: true` — skip compression, return messages as-is +- `x-headroom-bypass: true` — skip compression, return messages as-is with zeroed metrics -**Error responses:** 400 (missing fields), 401 (bad API key), 503 (compression failed) +**Error responses:** 400 (missing/invalid fields, bad `config.mode` or `config.frozen_message_count`), 401 (bad `HEADROOM_PROXY_TOKEN`), 404 (non-loopback without `HEADROOM_COMPRESS_ALLOW_REMOTE=1`), 503 (compression failed) + +**Fail-open:** on timeout you get 200 with the original messages plus `compression_skipped: true` and `skip_reason: "compression_timeout"`. + +**Multi-turn callers — don't lose the prefix cache.** This endpoint is stateless: unlike the proxy's own request path (which runs a CacheAligner and tracks provider cache hits across turns), it has no idea what the provider already cached. + +The provider caches the bytes you *forwarded*, which compression already changed — so your originals and the cached prefix are no longer the same thing, and it is the forwarded version you must keep reproducing. Compression also varies with position: an older tool result can fall outside the recent-read protection window as the conversation grows and be compressed harder than last turn, so re-compression is not guaranteed to reproduce earlier output either. Two rules: + +1. Pass `config.frozen_message_count` = the number of leading messages already cached upstream. +2. Send back the messages you **previously forwarded**, not the pristine originals. `frozen_message_count` returns leading messages exactly as passed in, so feeding it originals hands the provider different bytes than last turn and busts the cache anyway. + +```python +forwarded = [] +def next_turn(new_messages): + r = requests.post(f"{proxy}/v1/compress", json={ + "messages": forwarded + new_messages, + "model": "claude-sonnet-4-6", + "config": {"frozen_message_count": len(forwarded)}, + }).json() + forwarded[:] = r["messages"] # next turn's frozen prefix + return forwarded +``` + +Note `protect_recent` is not a substitute — it guards the newest messages, while `frozen_message_count` guards the oldest, which is the cached end. ## Using with Claude Code diff --git a/wiki/typescript-sdk.md b/wiki/typescript-sdk.md index 1a35c2032..eb2c9bd3f 100644 --- a/wiki/typescript-sdk.md +++ b/wiki/typescript-sdk.md @@ -26,7 +26,9 @@ const response = await openai.chat.completions.create({ ## How It Works -The TypeScript SDK is an HTTP client. When you call `compress()`, it sends your messages to the Headroom proxy's `POST /v1/compress` endpoint. The proxy runs the full compression pipeline (SmartCrusher, ContentRouter, CacheAligner, etc.) and returns compressed messages. No compression logic runs in Node.js — all the heavy lifting happens in the proxy. +The TypeScript SDK is an HTTP client. When you call `compress()`, it sends your messages to the Headroom proxy's `POST /v1/compress` endpoint. The proxy runs the compression pipeline (ContentRouter and its compressors, including SmartCrusher) and returns compressed messages. No compression logic runs in Node.js — all the heavy lifting happens in the proxy. + +The proxy must be reachable on **loopback**: `/v1/compress` rejects remote callers with `404` unless it was started with `HEADROOM_COMPRESS_ALLOW_REMOTE=1`. ``` Your TypeScript App @@ -37,7 +39,7 @@ headroom-ai (npm) ← HTTP client │ │ POST /v1/compress ▼ -Headroom Proxy / Cloud ← compression pipeline (Python) +Headroom Proxy (loopback) ← compression pipeline (Python) │ │ compressed messages ▼