From 63269650462c1a565c259cfe3fc3bd67ac942f5a Mon Sep 17 00:00:00 2001 From: gglucass Date: Wed, 30 Sep 2026 15:40:48 +0200 Subject: [PATCH 1/4] fix(transforms): judge the code_aware Kompress fallback in tokens (#3881) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `ContentRouter._apply_strategy_to_content` compares two different units in the `code_aware` path: - `code_aware`'s result is stored as a **word** count: `compressed_tokens = len(output.content.split())` (introduced with the registry dispatch in #2400, after #1857 had moved every other count to `_estimate_tokens`). - The no-savings Kompress fallback is then kept only `if fallback_tokens < compressed_tokens:`, where `fallback_tokens` is a **token** count (cl100k payload tokens for a real Kompress result). Code runs about 2 tokens per whitespace word, so an unchanged code block looks like a ~55% compression and the fallback can essentially never win. The common case is a Claude Code `Read` of source: the `N\t` line-number column stops `code_aware` from parsing, it hands the text back unchanged with `compressed=True`, the fallback runs Kompress over the whole block (up to ~20 s of ONNX inference on a 1,800-line file, holding the Kompress slot), and the result is always discarded. The phantom ~0.45 ratio is then refused by the lossy-unrecoverable guard and logged as `lossy_unrecoverable_skipped`, 0 tokens saved. It is the same class as the CONFIG word-count bug fixed in #2759. Fix: 1. Measure `code_aware`'s output with `_estimate_tokens`, the unit of `original_tokens` and of every other branch. 2. Measure the Kompress fallback's result with the same estimator instead of the count Kompress returns (a Kompress passthrough reports **words**, which would otherwise let an unchanged block "win"). 3. The lossless-then-lossy inline Kompress attempt in the same branch is measured the same way and now recorded in `strategy_chain` even when it loses, so the no-savings fallback below does not run the identical inference a second time. **This is a compression behaviour change, stated plainly.** The fallback can now win, so a `code_aware` no-op block (typically a large line-numbered code `Read`) can now be Kompress-compressed with a CCR retrieval marker, where before it went out verbatim after the inference was thrown away. Savings go up; the Kompress inference that already ran now pays for itself. Real `code_aware` compressions are now scored in honest tokens rather than an optimistic word ratio, so a marginal one could in principle land above `min_ratio` where it passed before (none did on the corpus below). ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/transforms/content_router.py`: `code_aware` result counted with `_estimate_tokens`; the no-savings Kompress fallback and the lossless-then-lossy inline attempt measured with `_estimate_tokens`; the inline attempt recorded in the chain before it runs. - `tests/test_transforms_content_router.py` (new, no tree-sitter needed): - `test_code_aware_noop_keeps_a_smaller_kompress_fallback`: a no-op `code_aware` on 120 line-numbered Rust lines, with a Kompress result smaller in tokens but larger than the word count, keeps the Kompress result. - `test_code_aware_noop_reports_tokens_when_kompress_passes_through`: when Kompress passes through, the block is reported at its token count, not its word count. - `test_code_aware_noop_runs_kompress_once_when_it_passes_through`: with lossless-then-lossy on, a losing inline attempt is not retried by the fallback. - `test_code_aware_noop_keeps_a_smaller_inline_kompress`: with lossless-then-lossy on, a winning inline attempt is kept, with one Kompress call. - `tests/test_code_aware_regressions.py`, `tests/test_router_registry_dispatch.py`: two existing assertions pinned the word-count metric (`len(shrunk.split())`); they now assert `_estimate_tokens`, the unit `original_tokens` uses. Their intent (the compressor's own estimate must not leak in, a real compression must not trigger a fallback) is unchanged. - `tests/test_router_registry_smartcrusher.py`, `tests/test_openai_responses_context_compaction.py`: two SMART_CRUSHER no-op tests go through the same shared fallback and asserted the count a mocked `_try_ml_compressor` reported (one mock "shrank" by returning longer text than it was given). The mocks now return shorter text and the assertions use `_estimate_tokens` of it. This also means a Kompress *passthrough* after a SMART_CRUSHER or CONFIG no-op, which Kompress reports in words, can no longer be adopted as a phantom saving. The first two new tests fail on `main` ffc6edb4 (the Kompress result is discarded; the no-op is reported at 720 "tokens", its word count, against 2,136 estimator tokens). The third fails when the chain entry is removed (Kompress runs twice). ## 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 $ uv run --frozen --extra dev --extra code pytest tests/test_transforms_content_router.py -q 72 passed, 2 skipped $ uv run --frozen --extra dev --extra code pytest tests/test_code_aware_regressions.py -q 5 passed $ uv run --frozen --extra dev --extra code pytest tests/test_router_registry_dispatch.py -q 11 passed $ uv run --frozen --extra dev --extra code pytest tests/test_transforms/test_content_router.py -q 63 passed $ uv run --frozen --extra dev --extra code pytest tests/test_lossless_mode.py -q 26 passed $ uv run --frozen --extra dev --extra code pytest tests/test_builtin_compressor_adapters.py -q 17 passed, 1 skipped $ uv run --frozen --extra dev --extra code pytest tests/test_content_router_tool_role_reversibility.py tests/test_agent_savings.py -q 10 passed / 39 passed $ uv run --frozen --extra dev --extra code pytest tests/test_router_registry_smartcrusher.py -q 7 passed $ uv run --frozen --extra dev --extra code pytest tests/test_openai_responses_context_compaction.py -q 10 passed # plus every other test file that mocks _try_ml_compressor (kompress_size_gate, lossless_then_lossy, ccr_tag_placeholder_regression, # relevance_split, kompress_request_nonblocking, bash_search_lossless_fold, text_crusher_routing, transforms_tabular): all pass # new tests with the headroom/ change stashed (main ffc6edb4): FAILED tests/test_transforms_content_router.py::test_code_aware_noop_keeps_a_smaller_kompress_fallback FAILED tests/test_transforms_content_router.py::test_code_aware_noop_reports_tokens_when_kompress_passes_through $ uvx ruff@0.16.8 check <6 touched files> All checks passed! $ uvx ruff@0.16.8 format --check <6 touched files> 6 files already formatted $ uv run --frozen --extra dev mypy headroom/transforms/content_router.py --ignore-missing-imports Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: macOS 15 (16 GB, under load from a concurrent Rust release build), Python 3.11 worktree venv with `--extra dev --extra proxy --extra code` (tree-sitter + ONNX Kompress, `chopratejas/kompress-v2-base` from the local HF cache, warm), upstream main ffc6edb4 with and without this commit. A real `HeadroomProxy` after `seed_proxy_env_defaults()`, with `savings_profile="coding"` and `code_aware_enabled=True` (what the `headroom proxy` CLI resolves by default), driven through `anthropic_pipeline.apply(...)` with `proxy_pipeline_kwargs(config)`. - Exact command / steps: 23 tool results built from a pristine `git archive ffc6edb4` of this repo, each sent as `[user, assistant tool_use, user tool_result]`: 10 Claude Code `Read`s in the `N\t` format (Rust, Python, TypeScript, Markdown, TOML, YAML), 10 `Bash` outputs (`cat -n` / `nl -ba` / `sed -n` / `cat` reads, `rg`, `ls -la`, JSON, `git log --stat`, `git diff`) and 3 raw-code results from a non-excluded MCP tool. Arms run fix, main, main, fix; `_try_ml_compressor` wrapped only to count calls and time them. - Observed result: exactly 4 of 23 items change, all large code `Read`s where `code_aware` was a no-op and Kompress had already run on main: | item | tokens | main saved | this PR saved | |---|---:|---:|---:| | `Read crates/headroom-proxy/src/proxy.rs` (1,709 lines) | 21,040 | 0 (`lossy_unrecoverable_skipped`, Kompress ~15-20 s discarded) | 4,847 (CCR marker) | | `Read crates/headroom-proxy/src/config.rs` | 10,623 | 0 (`lossy_unrecoverable_skipped`) | 2,072 | | `Read crates/headroom-proxy/src/responses_items.rs` | 6,614 | 0 (`lossy_unrecoverable_skipped`) | 1,422 | | `Read sdk/typescript/src/client.ts` | 5,876 | 0 (`lossy_unrecoverable_skipped`) | 799 | Corpus total 169,870 tokens: saved 13,674 on main vs 22,814 with this change (+9,140). Kompress calls are the same in both arms (12 per round; one main round made 13), so this does not reduce inference time; it keeps the result of inference main already paid for. The other 19 items are byte-identical between arms. A first full run of the same corpus was deterministic in both arms with the identical totals. In the second run the machine's concurrent build pushed a few Kompress calls past the 20 s request deadline (27.6 s on `proxy.rs`), so in one fix round `proxy.rs` and `config.rs` came back `ratio_too_high`, and in one main round `server.py` flipped; that is load, not the change. With `ProxyConfig` defaults (`code_aware` off, no savings profile) a 20-item run was byte-identical between arms. - Not tested: the model-facing effect of a Kompress-compressed code `Read` in a live Claude Code session (whether it is retrieved via CCR or re-read); Kompress on GPU backends; the latency effect of skipping a fallback that provably cannot win (not implemented, see Additional Notes). ### Update: measured on top of #3885 The four gains above were measured with each `Read` as the NEWEST message. Under `mode=token` with `savings_profile="coding"` (or `"general"`), a newest `Read` is only compressible because of a separate bug: the profile's `protect_recent=0` reaches the router as `read_protection_window=0` and replaces token mode's `max(4, 0.3*n)` window, so a `Read` has no protection at any position. #3885 fixes that. Without #3885, this PR turns the newest file read into a lossy one-line rewrite (`analyzer.py` 13,461 -> 11,139 tokens, newlines 1,082 -> 1); with it, that read stays byte-exact. Re-measured on the installed 0.39.0 wheel with the hunks exec'd in-process: a 23-item corpus of 147,527 tokens built from `git archive ffc6edb4` (10 `Read`s including the four above, 10 Bash outputs, 3 MCP code results), a fresh `HeadroomProxy` per arm, real ONNX Kompress with `HEADROOM_KOMPRESS_TIME_BUDGET_SECONDS=0` so load cannot flip a result. "aged" = the item followed by 40 messages. Both arms with this PR were run twice; 0 of 92 rows differed. | #3885 | this PR | newest: saved | aged: saved | |---|---|---:|---:| | no | no (main) | 19,972 (13.5%) | 19,972 (13.5%) | | no | yes | 38,240 (25.9%) | 38,240 (25.9%) | | yes | no | 9,696 (6.6%) | 19,972 (13.5%) | | yes | yes | 9,696 (6.6%) | 38,240 (25.9%) | With #3885 in place this PR adds nothing on a newest `Read` (routed `excluded_tool`) and +18,268 tokens once the same four code `Read`s age out of the window (proxy.rs 10,131, config.rs 4,487, responses_items.rs 2,145, client.ts 1,505; the other 19 items are identical). Per-item savings are about 2x the table above because this ran on the 0.39.0 wheel rather than main. Two limits: in live traffic an aged `Read` already inside the frozen prefix is never recompressed, so this is an upper bound; and under the default 20 s Kompress budget, proxy.rs (14-65 s of inference on a loaded 16 GB laptop) often comes back unchanged on slower machines. Recommend merging #3885 first or together. ## Runtime Rollout Safety - Rollout-managed feature(s): None. - Minimum rollout channel: Stable. - Stable/default behavior changed: Yes, when `code_aware` is enabled (the CLI default): a `code_aware` no-op block whose Kompress fallback is smaller in tokens is now Kompress-compressed with a CCR marker instead of forwarded verbatim; `code_aware` results are reported in estimator tokens instead of words. With `code_aware` disabled, nothing changes. - Kill switch / disable path: `--no-code-aware` / `HEADROOM_CODE_AWARE_ENABLED=0` (no-op blocks then take the plain Kompress route), or `HEADROOM_DISABLE_KOMPRESS=1` to turn Kompress off entirely. No dedicated switch for this fix. - Unsafe override required: None. - Qualification impact: Savings rise on large code reads (+9,140 tokens, +67% relative, on the corpus above). Those blocks now carry a lossy Kompress body plus a CCR retrieval marker, so it is worth watching retrieval and re-read rates on the release that carries it. - Rollback path: Revert this PR. ## 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` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) N/A ## Additional Notes - Documentation: N/A, internal accounting. - For reviewers: under the coding profile (`protect_recent=0`), a Claude Code `Read` result is already compressible on `main` (`codex.py` and `server.py` above are compressed by `code_aware` on both arms). If the intent is that `Read` results are never lossy-compressed, the word-count bug was masking a gap for no-op code reads, and that deserves an explicit protection rather than this accident; this PR does not change which blocks reach the router. #3885 is that protection for recent reads in token mode (see the update above). - Deliberately left out: skipping the Kompress call when it provably cannot win. In tokens, the only hard bound (the words Kompress always keeps, plus its marker) rarely reaches the block's token count, and on this corpus the fallback won 4 of the 5 times it ran. A heuristic skip is a latency optimisation with its own savings trade and belongs in its own PR. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 (1M context) --- headroom/transforms/content_router.py | 28 ++-- tests/test_code_aware_regressions.py | 14 +- ...est_openai_responses_context_compaction.py | 4 +- tests/test_router_registry_dispatch.py | 4 +- tests/test_router_registry_smartcrusher.py | 9 +- tests/test_transforms_content_router.py | 120 ++++++++++++++++++ 6 files changed, 155 insertions(+), 24 deletions(-) diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 482e1cb810..18514253f8 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -3704,7 +3704,11 @@ def _apply_strategy_to_content( ) if output is not None and output.compressed: compressed = output.content - compressed_tokens = len(output.content.split()) + # Tokens, like ``original_tokens`` and every other + # branch: a word count here (code runs ~2.2 + # tokens/word) made an unchanged block look like a + # 55% compression and no fallback could beat it. + compressed_tokens = _estimate_tokens(output.content) decision_reason = "code_aware" if compressed is None: # Fallback to Kompress @@ -3728,19 +3732,17 @@ def _apply_strategy_to_content( # lossless has no savings. Reads are protected upstream, so # only NON-read code reaches here. Keep Kompress ONLY if it # actually shrinks (never inflate). - _k, _kt = self._try_ml_compressor(content, context, question) - if ( - _k is not None - and _kt is not None - and _kt < original_tokens - and len(_k) < len(content) - ): + # Recorded as tried even if it loses, so the no-savings + # fallback below does not run the same inference again. + strategy_chain.append(CompressionStrategy.KOMPRESS.value) + _k, _ = self._try_ml_compressor(content, context, question) + _kt = _estimate_tokens(_k) + if _kt < original_tokens and len(_k) < len(content): compressed, compressed_tokens = _k, _kt strategy = CompressionStrategy.KOMPRESS actual_strategy = strategy compressor_name = "KompressCompressor" decision_reason = "code_aware_no_shrink_fallback_kompress" - strategy_chain.append(CompressionStrategy.KOMPRESS.value) elif strategy == CompressionStrategy.SMART_CRUSHER: # SmartCrusher handles its own TOIN recording @@ -3965,9 +3967,11 @@ def _apply_strategy_to_content( already_tried_kompress = CompressionStrategy.KOMPRESS.value in strategy_chain if not already_tried_kompress: strategy_chain.append(CompressionStrategy.KOMPRESS.value) - fallback_compressed, fallback_tokens = self._try_ml_compressor( - content, context, question - ) + fallback_compressed, _ = self._try_ml_compressor(content, context, question) + # Measure with the router's estimator, the unit of + # ``compressed_tokens``. Kompress reports a passthrough in + # WORDS, which would let an unchanged block "win". + fallback_tokens = _estimate_tokens(fallback_compressed) else: fallback_compressed = compressed fallback_tokens = compressed_tokens diff --git a/tests/test_code_aware_regressions.py b/tests/test_code_aware_regressions.py index 632d8fc808..2dbcc2c11b 100644 --- a/tests/test_code_aware_regressions.py +++ b/tests/test_code_aware_regressions.py @@ -7,8 +7,9 @@ 2. Multi-line Go signature (`) error {`) — the brace-bearing signature line survives truncation instead of being silently dropped. 3. ContentRouter CODE_AWARE token accounting — compressed_tokens must come - from the same word-count metric as original_tokens, not the compressor's - own (differently-scaled) estimator, or real savings get misread as none. + from the same metric as original_tokens (the router's token estimator), + not the compressor's own (differently-scaled) estimator, or real savings + get misread as none. 4. prefer_code_aware_for_code defaults to True, both on the dataclass and via the HEADROOM_PREFER_CODE_AWARE_FOR_CODE env var. """ @@ -27,6 +28,7 @@ CompressionStrategy, ContentRouter, ContentRouterConfig, + _estimate_tokens, ) pytestmark = pytest.mark.skipif( @@ -96,13 +98,13 @@ class _Result: r.compressed = self._compressed # Deliberately inflated/differently-scaled "own" token estimate — # simulates a compressor whose internal counter isn't comparable to - # the router's len(text.split()) word count. + # the router's own token estimate. r.compressed_tokens = 10_000_000 return r -def test_content_router_code_aware_token_accounting_matches_word_count() -> None: - """compressed_tokens must be derived from len(result.compressed.split()), +def test_content_router_code_aware_token_accounting_matches_router_estimate() -> None: + """compressed_tokens must be derived from _estimate_tokens(result.compressed), the same metric as original_tokens, not the compressor's own estimator. Using the mismatched estimator made a real compression look like "no savings" and forced a needless fallback to Kompress.""" @@ -117,7 +119,7 @@ def test_content_router_code_aware_token_accounting_matches_word_count() -> None ) assert compressed == compressed_text - assert compressed_tokens == len(compressed_text.split()) + assert compressed_tokens == _estimate_tokens(compressed_text) assert strategy_chain == ["code_aware"], ( f"real compression must not trigger a Kompress fallback, got {strategy_chain}" ) diff --git a/tests/test_openai_responses_context_compaction.py b/tests/test_openai_responses_context_compaction.py index 74dc99cba0..48f7c11ca6 100644 --- a/tests/test_openai_responses_context_compaction.py +++ b/tests/test_openai_responses_context_compaction.py @@ -12,6 +12,7 @@ CompressionStrategy, ContentRouter, ContentRouterConfig, + _estimate_tokens, ) @@ -432,7 +433,8 @@ def crush(self, value: str, query: str = "", bias: float = 1.0): ) assert compressed == "short summary" - assert compressed_tokens == 2 + # Measured with the router's estimator, the unit of the strategy result. + assert compressed_tokens == _estimate_tokens("short summary") # The fallback chain must record both strategies it tried. assert strategy_chain == ["smart_crusher", "kompress"] diff --git a/tests/test_router_registry_dispatch.py b/tests/test_router_registry_dispatch.py index 1846f631a7..92257a1053 100644 --- a/tests/test_router_registry_dispatch.py +++ b/tests/test_router_registry_dispatch.py @@ -244,8 +244,8 @@ def _fake_compress(content: str, language: object = None, context: str = "") -> content, CompressionStrategy.CODE_AWARE, "q", language="python", bias=1.0 ) assert out == shrunk - # CODE_AWARE's historical token metric is len(compressed.split()), NOT _estimate_tokens. - assert tokens == len(shrunk.split()) + # Measured in the router's token estimate, the unit of original_tokens. + assert tokens == _estimate_tokens(shrunk) assert chain == [CompressionStrategy.CODE_AWARE.value] # The flip forwarded content, language, and context through the registry adapter. assert seen == {"content": content, "language": "python", "context": "q"} diff --git a/tests/test_router_registry_smartcrusher.py b/tests/test_router_registry_smartcrusher.py index 1a820b7a05..d80684a317 100644 --- a/tests/test_router_registry_smartcrusher.py +++ b/tests/test_router_registry_smartcrusher.py @@ -144,12 +144,15 @@ def test_smart_crusher_kompress_fallback_matches_direct(monkeypatch: pytest.Monk lambda: SimpleNamespace(crush=lambda c, query="", bias=1.0: SimpleNamespace(compressed=c)), ) # Kompress shrinks (fallback_tokens < compressed_tokens) so it is adopted. - monkeypatch.setattr(router, "_try_ml_compressor", lambda c, ctx, q: ("KOMPRESSED::" + c, 3)) + # The fallback is measured with the router's estimator, not the count the + # ML compressor reports, so the mock must actually return less text. + kompressed = "KOMPRESSED::" + _JSON[:40] + monkeypatch.setattr(router, "_try_ml_compressor", lambda c, ctx, q: (kompressed, 3)) out, tokens, chain = router._apply_strategy_to_content( _JSON, CompressionStrategy.SMART_CRUSHER, "ctx", bias=1.0 ) - assert out == "KOMPRESSED::" + _JSON - assert tokens == 3 + assert out == kompressed + assert tokens == _estimate_tokens(kompressed) assert chain == [ CompressionStrategy.SMART_CRUSHER.value, CompressionStrategy.KOMPRESS.value, diff --git a/tests/test_transforms_content_router.py b/tests/test_transforms_content_router.py index 7ff3619b60..946710f6b3 100644 --- a/tests/test_transforms_content_router.py +++ b/tests/test_transforms_content_router.py @@ -2251,3 +2251,123 @@ def test_html_extraction_miss_still_gets_deferred_json_splice() -> None: assert len(text) < len(shell) json_marker = '"id":' assert json_marker not in text # the array was compressed away + + +# A Claude Code Read of Rust: line-numbered code that code_aware hands back +# unchanged (it cannot parse the number column) while reporting compressed=True. +# Code runs ~2 tokens per whitespace word, so a word count badly undercounts it. +_NUMBERED_RUST = "\n".join( + f"{i}\tlet value_{i} = compute(&items[{i}..], Some(opts.clone()))?;" for i in range(1, 121) +) + + +def _noop_code_aware_router(monkeypatch: pytest.MonkeyPatch, kompress_out: str) -> ContentRouter: + monkeypatch.delenv("HEADROOM_LOSSLESS_THEN_LOSSY", raising=False) + router = ContentRouter(ContentRouterConfig(enable_code_aware=True)) + + class NoopCodeCompressor: + def compress(self, content: str, language=None, context: str = "") -> SimpleNamespace: + return SimpleNamespace(compressed=content, compressed_tokens=0) + + monkeypatch.setattr(router, "_get_code_compressor", lambda: NoopCodeCompressor()) + + def kompress(content: str, *_args: object, **_kwargs: object) -> tuple[str, int]: + # What Kompress reports: payload tokens for a real compression, the + # WORD count for a passthrough. + if kompress_out == content: + return content, len(content.split()) + return kompress_out, _estimate_tokens(kompress_out) + + monkeypatch.setattr(router, "_try_ml_compressor", kompress) + return router + + +def test_code_aware_noop_keeps_a_smaller_kompress_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The fallback is judged in tokens, the unit of the code_aware result. + + With code_aware's result counted in words, a Kompress result at ~60% of the + tokens still "lost" to the unchanged block, so the inference always ran and + was always thrown away. + """ + kompressed = "\n".join(f"{i}\tvalue_{i} compute(&items[{i}..])" for i in range(1, 121)) + assert _estimate_tokens(kompressed) < _estimate_tokens(_NUMBERED_RUST) + assert _estimate_tokens(kompressed) > len(_NUMBERED_RUST.split()) + router = _noop_code_aware_router(monkeypatch, kompressed) + + compressed, compressed_tokens, strategy_chain = router._apply_strategy_to_content( + _NUMBERED_RUST, CompressionStrategy.CODE_AWARE, context="" + ) + + assert compressed == kompressed + assert compressed_tokens == _estimate_tokens(kompressed) + assert strategy_chain == ["code_aware", "kompress"] + + +def test_code_aware_noop_reports_tokens_when_kompress_passes_through( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An unchanged block is reported at its own token count, never its word count.""" + router = _noop_code_aware_router(monkeypatch, _NUMBERED_RUST) + + compressed, compressed_tokens, _chain = router._apply_strategy_to_content( + _NUMBERED_RUST, CompressionStrategy.CODE_AWARE, context="" + ) + + assert compressed == _NUMBERED_RUST + assert compressed_tokens == _estimate_tokens(_NUMBERED_RUST) + + +def test_code_aware_noop_runs_kompress_once_when_it_passes_through( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Lossless-then-lossy tries Kompress inline; a passthrough must not be retried.""" + router = _noop_code_aware_router(monkeypatch, _NUMBERED_RUST) + router._lossless_then_lossy = True + calls: list[str] = [] + kompress = router._try_ml_compressor + + def counting(content: str, *args: object, **kwargs: object) -> tuple[str, int]: + calls.append(content) + return kompress(content, *args, **kwargs) + + monkeypatch.setattr(router, "_try_ml_compressor", counting) + + compressed, compressed_tokens, _chain = router._apply_strategy_to_content( + _NUMBERED_RUST, CompressionStrategy.CODE_AWARE, context="" + ) + + assert compressed == _NUMBERED_RUST + assert compressed_tokens == _estimate_tokens(_NUMBERED_RUST) + assert len(calls) == 1 + + +def test_code_aware_noop_keeps_a_smaller_inline_kompress( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Lossless-then-lossy's inline attempt keeps a Kompress result that wins. + + Measured in tokens, a code_aware no-op no longer looks like a shrink, so the + inline attempt fires and its smaller result is kept without a second call. + """ + kompressed = "\n".join(f"{i}\tvalue_{i} compute(&items[{i}..])" for i in range(1, 121)) + router = _noop_code_aware_router(monkeypatch, kompressed) + router._lossless_then_lossy = True + calls: list[str] = [] + kompress = router._try_ml_compressor + + def counting(content: str, *args: object, **kwargs: object) -> tuple[str, int]: + calls.append(content) + return kompress(content, *args, **kwargs) + + monkeypatch.setattr(router, "_try_ml_compressor", counting) + + compressed, compressed_tokens, strategy_chain = router._apply_strategy_to_content( + _NUMBERED_RUST, CompressionStrategy.CODE_AWARE, context="" + ) + + assert compressed == kompressed + assert compressed_tokens == _estimate_tokens(kompressed) + assert strategy_chain == ["code_aware", "kompress"] + assert len(calls) == 1 From d1dc027e455474c9eb0d35b74f8c68b5f274d8a2 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 30 Sep 2026 08:43:18 -0500 Subject: [PATCH 2/4] chore(deps): bump accelerate from 1.12.0 to 1.15.0 in the uv group across 1 directory (#3856) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps the uv group with 1 update in the / directory: [accelerate](https://github.com/huggingface/accelerate). Updates `accelerate` from 1.12.0 to 1.15.0
Release notes

Sourced from accelerate's releases.

v1.15.0: FSDP2 activation memory, dtensor improvements

FSDP2

A large batch of FSDP2 work this release: two fixes that cut activation memory at long sequence lengths, tied-embedding support on torch >= 2.13, and a round of checkpointing correctness and scale fixes.

Activation checkpointing was wrapping each child of the matched layer (self_attn, mlp, the norms) instead of the layer itself, so every inter-child activation stayed saved for backward. It now wraps the layer.

There's also a new FSDP2-only activation_checkpointing_offload, which moves the remaining per-layer checkpoint inputs to pinned CPU memory. Gradients are exactly those of plain activation checkpointing:

# fsdp2.yaml
fsdp_config:
  fsdp_version: 2
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_activation_checkpointing: true
  fsdp_activation_checkpointing_offload: true

accelerate launch --config_file fsdp2.yaml
train.py

DTensor

Two fixes for DTensor-sharded models, which you hit with FSDP2, tensor parallelism, or any N-D parallelism setup: gradient clipping no longer fails on the foreach op when plain tensors and DTensors are mixed, and prepare_model leaves an already-sharded model where it is:

Offloading & Quantization

An entire model can now be dispatched to disk, including tied weights — useful for tools like llm-compressor that compress large models on machines that can't hold them:

Trackers

Custom trackers can be registered by name and then selected from log_with= like any built-in one:

from accelerate import Accelerator
from accelerate.tracking import register_tracker_class
</tr></table>

... (truncated)

Commits
  • 6afc1e5 Release: v1.15.0
  • 636a934 Fix nightly CI: tensor parallel size detection and DeepSpeed warmup steps (#4...
  • 385d9fb docs: fix three dead links left by the docs restructure (#4230)
  • c424d0b docs: fix garbled sentence in CONTRIBUTING.md (#4231)
  • 0f7e35f Clip grad norm support for dtensors (#4219)
  • b2ac509 Fix FSDP2/ PEFT/FULL_STATE_DICT dropping every rank's adapter shard except ...
  • 33e8bc4 Fix load_accelerator_state only restoring one RNG backend (#4217)
  • 50de3ad Treat MPS out-of-memory errors as OOM in find_executable_batch_size (#4227)
  • ab5fe8e feat: add the neuron device branch in state (#4218)
  • 9baf95b Fix MLFLOW_NESTED_RUN never being able to turn nesting off (#4223)
  • Additional commits viewable in compare view

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=accelerate&package-manager=uv&previous-version=1.12.0&new-version=1.15.0)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore major version` will close this group update PR and stop Dependabot creating any more for the specific dependency's major version (unless you unignore this specific dependency's major version or upgrade to it yourself) - `@dependabot ignore minor version` will close this group update PR and stop Dependabot creating any more for the specific dependency's minor version (unless you unignore this specific dependency's minor version or upgrade to it yourself) - `@dependabot ignore ` will close this group update PR and stop Dependabot creating any more for the specific dependency (unless you unignore this specific dependency or upgrade to it yourself) - `@dependabot unignore ` will remove all of the ignore conditions of the specified dependency - `@dependabot unignore ` will remove the ignore condition of the specified dependency and ignore conditions You can disable automated security fix PRs for this repo from the [Security Alerts page](https://github.com/headroomlabs-ai/headroom/network/alerts).
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: JerrettDavis --- uv.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/uv.lock b/uv.lock index e3b242384f..991f849c3c 100644 --- a/uv.lock +++ b/uv.lock @@ -32,7 +32,7 @@ constraints = [ [[package]] name = "accelerate" -version = "1.12.0" +version = "1.15.0" source = { registry = "https://pypi.org/simple/" } dependencies = [ { name = "huggingface-hub" }, @@ -44,9 +44,9 @@ dependencies = [ { name = "safetensors" }, { name = "torch" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/4a/8e/ac2a9566747a93f8be36ee08532eb0160558b07630a081a6056a9f89bf1d/accelerate-1.12.0.tar.gz", hash = "sha256:70988c352feb481887077d2ab845125024b2a137a5090d6d7a32b57d03a45df6", size = 398399, upload-time = "2025-11-21T11:27:46.973Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/b5/1d3ed029ac71d3f2961346829a268da923698e9fd63f218f78841f216bfd/accelerate-1.15.0.tar.gz", hash = "sha256:5654f8c5eaa0d4fa68b33e287a97765da6849bf6d51dcac874e73fbbddfb6134", size = 422615, upload-time = "2026-09-09T13:04:49.078Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9f/d2/c581486aa6c4fbd7394c23c47b83fa1a919d34194e16944241daf9e762dd/accelerate-1.12.0-py3-none-any.whl", hash = "sha256:3e2091cd341423207e2f084a6654b1efcd250dc326f2a37d6dde446e07cabb11", size = 380935, upload-time = "2025-11-21T11:27:44.522Z" }, + { url = "https://files.pythonhosted.org/packages/8a/4c/34f0450479d01195027260da68d8a3880683f1640c3ca5adf64acb3185f1/accelerate-1.15.0-py3-none-any.whl", hash = "sha256:97eacca0b73e45cb867dbf8c5d5d4dc32219544300e0c8992c7334dc2ef33cec", size = 394295, upload-time = "2026-09-09T13:04:47.331Z" }, ] [[package]] From 1512c062350631c1b17e8ecb139b18b7fc38ae61 Mon Sep 17 00:00:00 2001 From: gglucass Date: Wed, 30 Sep 2026 15:44:48 +0200 Subject: [PATCH 3/4] fix(transforms): keep file-read protection in token mode under the coding profile (#3885) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description In token mode, the `coding` and `general` savings profiles strip excluded-tool output of all protection, so a Claude Code `Read` is lossy-compressed even as the newest message. The chain: - `proxy/server.py` sets `router_config.protect_recent_reads_fraction = 0.3` in token mode: excluded-tool output in the last `max(4, 0.3 * n)` messages is protected, older output may compress. - `proxy_pipeline_kwargs()` passes `read_protection_window = profile.protect_recent` on every request. `coding` and `general` set `protect_recent=0`. That 0 is a positional guard meant for cache mode (#2145); the coding profile's own comment says the code working set "stays byte-exact via protect_reads below - a TYPE guard on file reads", and "NEVER lossy-compress a file read (the agent patches exact bytes)". - `ContentRouter.apply` lets the runtime window replace the fraction window whenever the fraction is > 0 (#2105). The window becomes 0, `messages_from_end <= 0` is never true (the newest message is 1), and every excluded-tool result falls through to compression at any position. Effect in token mode with `coding` or `general`: - `Read`, `Skill` and `read_file` (`DEFAULT_BYTE_EXACT_EXCLUDE_TOOLS` minus the verbatim set) lose byte-exact protection entirely. - `Grep`, `Glob`, `Edit` and `Write` lose recent-output protection (they were meant to get lossless compaction only while recent). - `WebFetch` / `WebSearch` are unaffected (`DEFAULT_VERBATIM_EXCLUDE_TOOLS`, protected at any age). - `HEADROOM_PROTECT_READS` does not cover this: it only protects SHELL reads (`_is_read_command` on the Bash command), so `cat NOTES.md` via Bash stays protected while `Read NOTES.md` does not. `general` seeds `HEADROOM_MODE=token`, so every `general` deployment hits this by default. `coding` defaults to cache mode and hits it with `--mode token`. `agent-90` (2) and `balanced` (4) pass a positive window and are unchanged. Fix: a runtime `read_protection_window` of 0 means "no override", not "protect nothing". Only a positive window narrows the fraction window; 0 falls back to `max(4, 0.3 * n)`. This keeps #2105's precedence (a fraction of 0.0 still ignores the kwarg) and leaves `protect_recent=0` alone, since it still governs the positional guard for all tool output in cache mode. **This is a compression behaviour change that LOWERS savings** on recent `Read`/`Skill`/`read_file` output (now byte-exact) and recent `Grep`/`Glob`/`Edit`/`Write` output (now lossless compaction or verbatim) in token mode with the `coding` or `general` profile. That is the documented intent of both token mode's fraction window and the coding profile. Excluded-tool output beyond the window still ages out to compression, as token mode intends. This is a prerequisite for #3881 (judge the `code_aware` Kompress fallback in tokens). On main, the unprotected code `Read`s in token mode are only left intact by accident: the word-vs-token bug #3881 fixes makes their Kompress fallback lose. With #3881 alone they get flattened (see proof below). #3881's measured gains were on newest-message `Read`s that this PR protects. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/transforms/content_router.py` (`ContentRouter.apply`): the runtime `read_protection_window` only replaces the fraction window when it is > 0. A window of 0 (or less) is ignored. - `tests/test_content_router_exclude_tools.py`: `test_token_mode_coding_profile_keeps_recent_read_byte_exact`, parametrized. It builds a real `HeadroomProxy(mode="token", savings_profile="coding")`, takes kwargs from `proxy_pipeline_kwargs(config)` (asserts `read_protection_window == 0`), and mocks Kompress (no ONNX). A newest-message `Read` must stay byte-exact with `router:excluded:tool`; the same `Read` with 40 messages after it must still compress. ## 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 $ uv run --frozen --extra dev pytest tests/test_content_router_exclude_tools.py \ tests/test_transforms/test_content_router.py tests/test_transforms_content_router.py \ tests/test_agent_savings.py tests/test_content_router_concurrency.py \ tests/test_force_kompress_all.py tests/test_lossless_excluded_compaction.py \ tests/test_observed_wire_shapes.py tests/test_openai_responses_read_protection.py \ tests/test_proxy/test_header_safe_transforms.py tests/test_websearch_tool_result_protection.py -q 274 passed, 2 skipped # new test with headroom/ change reverted (main ffc6edb4): FAILED tests/test_content_router_exclude_tools.py::test_token_mode_coding_profile_keeps_recent_read_byte_exact[0] E AssertionError: newest Read must stay byte-exact in token mode 1 failed, 7 passed $ uvx ruff@0.16.8 check headroom/transforms/content_router.py tests/test_content_router_exclude_tools.py All checks passed! $ uvx ruff@0.16.8 format --check headroom/transforms/content_router.py tests/test_content_router_exclude_tools.py 2 files already formatted $ uv run --frozen --extra dev mypy headroom/transforms/content_router.py --ignore-missing-imports Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: macOS 15 (16 GB), Python 3.11, `uv run --frozen --extra dev --extra proxy --extra code`, real ONNX Kompress (`chopratejas/kompress-v2-base` from the local HF cache), tree-sitter installed. No mocks. - Exact command / steps: `seed_proxy_env_defaults()`, then `HeadroomProxy(ProxyConfig(mode=, savings_profile="coding", code_aware_enabled=True))` and `proxy.anthropic_pipeline.apply(messages, model="claude-sonnet-4-5", model_limit=200000, frozen_message_count=0, prefix_replay_guaranteed=True, **proxy_pipeline_kwargs(cfg))`, the handler's call shape. Messages are `[user, assistant tool_use Read, user tool_result]`, the file in Claude Code's `N\t` line-number format, optionally followed by 40 plain messages ("aged"). Arms: main `ffc6edb4`, this PR, and each with #3881's `content_router.py` diff applied. Tokens are cl100k counts of the `Read` block. - Observed result: on main the newest token-mode `Read` takes the compression route and is rewritten to one line with a CCR marker (CSV below, and the code file once #3881 lands); with this PR it routes `excluded_tool` and stays byte-exact, while the aged copy still compresses. Cache mode is identical across arms. `Read sbom/headroom-python-licenses.csv` (9,166 tokens, 331 lines): | arm | placement | route | block after | byte-exact | |---|---|---|---|---| | main | newest | compression | one line + `Retrieve more` marker (9,769 tokens) | no | | this PR | newest | `excluded_tool` | 9,166 | yes | | this PR | aged (40 msgs after) | compression | one line + marker (9,769 tokens) | no | | cache mode, both arms | newest | `excluded_tool` | 9,166 | yes | The same newest read under the `general` profile (token mode is its default): main rewrites it the same way, this PR keeps it byte-exact via `excluded_tool`. `Read headroom/learn/analyzer.py` (13,461 tokens, 1,083 lines): | arm | placement | route | block after | byte-exact | |---|---|---|---|---| | main | newest | compression (not excluded) | 13,461 (Kompress fallback loses on the #3881 unit bug) | yes, by accident | | main + #3881 | newest | compression | 11,139, newlines 1,082 -> 1, CCR marker | no | | this PR + #3881 | newest | `excluded_tool` | 13,461 | yes | | this PR + #3881 | aged (40 msgs after) | compression | 11,139, CCR marker | no | | cache mode, all arms | newest | `excluded_tool` | 13,461 | yes | The flattened output drops every newline and line-number column, which is what the agent copies into `Edit(old_string=...)`. - Not tested: live provider calls and a full Claude Code session; the OpenAI `role: "tool"` shape (it reads the same `read_protection_window` variable from the same `apply` scope, so the fix covers it, but no separate proof was run); fleet-level savings impact. ## Runtime Rollout Safety - Rollout-managed feature(s): None. - Minimum rollout channel: Stable. - Stable/default behavior changed: Yes, in token mode with a profile whose `protect_recent` is 0 (`coding` with `--mode token`, and `general`, which defaults to token mode): excluded-tool output in the last `max(4, 0.3 * n)` messages is protected again (`Read`/`Skill`/`read_file` byte-exact, `Grep`/`Glob`/`Edit`/`Write` lossless compaction or verbatim). Savings on those recent blocks drop. Cache mode, `--protect-tool-results`, and profiles with a positive `protect_recent` are unchanged. - Kill switch / disable path: none dedicated. `--mode cache` was never affected; there is no config that restores "compress the newest Read" short of reverting. - Unsafe override required: None. - Qualification impact: Lower token-mode savings on recent file reads and recent Grep/Glob/Edit/Write output for `coding`/`general`; fewer CCR retrievals and re-reads of lossy-compressed reads expected. Worth watching token-mode savings rate on the release that carries it. - Rollback path: Revert this PR. ## 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` - it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) N/A, backend routing change. ## Additional Notes - Documentation: no doc describes `read_protection_window` or its interaction with `protect_recent`, so there was nothing to update. - Placement: the alternative was to stop `proxy_pipeline_kwargs()` from emitting `read_protection_window` when `protect_recent` is 0. That leaves any other caller that passes 0 still exposed and changes the kwargs contract `tests/test_agent_savings.py` pins, so the guard sits in `ContentRouter.apply`, where every caller routes through. - A negative window used to clamp to 0 (protect nothing); it now also falls back to the fraction window. - Found while measuring #3881: its gains were on newest-message code `Read`s in token mode, which should not have been compressible at all. Merging this first keeps #3881 from turning an accidental passthrough into a lossy rewrite of the newest file read. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 (1M context) --- headroom/transforms/content_router.py | 7 +- tests/test_content_router_exclude_tools.py | 91 ++++++++++++++++++++++ 2 files changed, 97 insertions(+), 1 deletion(-) diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 18514253f8..d672688dcc 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -5509,6 +5509,7 @@ def apply( runtime_read_protection_window = kwargs.get("read_protection_window") if ( runtime_read_protection_window is not None + and int(runtime_read_protection_window) > 0 and self.config.protect_recent_reads_fraction > 0 ): # A profile-derived window may only narrow protection when the @@ -5517,7 +5518,11 @@ def apply( # See #1374's documented contract: protected tool output must never # lossy-compress "regardless of conversation depth" -- a per-request # savings-profile kwarg must not silently weaken that. - read_protection_window = max(0, int(runtime_read_protection_window)) + # A window of 0 is "no override", not "protect nothing": it comes + # from a profile's protect_recent=0 (coding, general), a positional + # guard for cache mode, and would otherwise strip file reads of all + # protection in token mode, down to the newest message. + read_protection_window = int(runtime_read_protection_window) # Adaptive compression ratio: scale with context pressure if model_limit > 0: diff --git a/tests/test_content_router_exclude_tools.py b/tests/test_content_router_exclude_tools.py index e1fedc4ae2..ad1d44497f 100644 --- a/tests/test_content_router_exclude_tools.py +++ b/tests/test_content_router_exclude_tools.py @@ -214,6 +214,97 @@ def test_protect_tool_results_survives_runtime_read_protection_window_kwarg() -> assert "router:excluded:tool" in result.transforms_applied +# --------------------------------------------------------------------------- +# Test 5: token mode + coding profile keeps recent file reads byte-exact +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("trailing_turns", "window", "protected"), + [ + (0, None, True), # newest Read, the profile's window of 0 + (1, None, True), # 3 from the end, inside max(4, 0.3*n) + (1, 2, False), # a positive profile window still narrows it + (20, None, False), # past the 0.3 window: ages out + ], +) +def test_token_mode_coding_profile_keeps_recent_read_byte_exact( + trailing_turns: int, window: int | None, protected: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """The coding profile's protect_recent=0 reaches the router as + read_protection_window=0. In token mode (fraction 0.3) that used to replace + the max(4, 0.3*n) window with 0, so a `Read` result lossy-compressed even + as the newest message. Recent reads must stay byte-exact; reads past the + fraction window still age out to compression, and a positive profile + window still narrows the fraction window.""" + pytest.importorskip("tiktoken") + + from types import SimpleNamespace + + from headroom.agent_savings import proxy_pipeline_kwargs + from headroom.providers import OpenAIProvider + from headroom.tokenizer import Tokenizer + + tokenizer = Tokenizer(OpenAIProvider().get_token_counter("gpt-4o"), "gpt-4o") + proxy = _build(mode="token", savings_profile="coding") + router = _router(proxy) + assert router.config.protect_recent_reads_fraction == 0.3 + + class FakeKompress: + def is_ready(self) -> bool: + return True + + def ensure_background_load(self) -> None: + pass + + def compress(self, content, **kwargs): + compressed = " ".join(content.split()[:20]) + " Retrieve more: hash=deadbeef" + return SimpleNamespace(compressed=compressed, compressed_tokens=len(compressed.split())) + + monkeypatch.setattr(router, "_get_kompress", lambda: FakeKompress()) + + read_output = "\n".join( + f"{i:>6}\tThe release notes describe change number {i} and its rollout plan." + for i in range(1, 121) + ) + messages: list[dict[str, object]] = [ + {"role": "user", "content": "read the notes"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_read_1", + "name": "Read", + "input": {"file_path": "/repo/NOTES.md"}, + } + ], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_read_1", "content": read_output} + ], + }, + ] + for i in range(trailing_turns): + messages.append({"role": "assistant", "content": f"reply {i}"}) + messages.append({"role": "user", "content": f"follow-up {i}"}) + + kwargs = proxy_pipeline_kwargs(proxy.config) + assert kwargs["read_protection_window"] == 0 + if window is not None: + kwargs["read_protection_window"] = window + result = router.apply(messages, tokenizer, **kwargs) + + content = result.messages[2]["content"][0]["content"] + if protected: + assert content == read_output, "a recent Read must stay byte-exact in token mode" + assert "router:excluded:tool" in result.transforms_applied + else: + assert content != read_output, "a Read outside the window still ages out" + + # --------------------------------------------------------------------------- # Baseline: Bash NOT in DEFAULT_EXCLUDE_TOOLS (unchanged by this PR) # --------------------------------------------------------------------------- From b7156714acf1f06cd37f642044455e23e2f89d47 Mon Sep 17 00:00:00 2001 From: Shubham Srivastava Date: Wed, 30 Sep 2026 15:27:09 +0100 Subject: [PATCH 4/4] fix(sdk): keep tool names on Vercel tool-result parts through the OpenAI round trip (#3883) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `openAIToVercel()` emits every `tool-result` part with `toolName: "unknown"`, because an OpenAI `tool` message only carries the `tool_call_id`. So after any successful compression through `headroomMiddleware()` / `compressVercelMessages()`, every tool result in the prompt reaches the AI SDK provider with the wrong tool name. That is not cosmetic for providers that put the name on the wire. In the AI SDK source, `packages/google/src/convert-to-google-messages.ts` builds `functionResponse: { name: toolName, … }` and `packages/mistral/src/convert-to-mistral-chat-messages.ts` sends `name: part.toolName` — both receive `"unknown"` today, which no longer matches the preceding function call. The name is recoverable: the assistant message that precedes the tool result carries `tool_calls[].function.name` for the same id. This records the `id → name` mapping while converting and uses it for the matching `tool-result`, keeping `"unknown"` only as the fallback when no call matches (the existing lone-tool-message test stays valid unchanged). ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - **`sdk/typescript/src/utils/format.ts`** — `openAIToVercel`: a `toolNames` map filled from each assistant `tool_calls` entry as it is converted; the `tool-result` part uses `toolNames.get(msg.tool_call_id) ?? "unknown"`. Three lines net. - **`sdk/typescript/test/utils/format.test.ts`** — `openAIToVercel` carries the name from the matching tool-call onto the tool-result; a Vercel → OpenAI → Vercel round trip preserves `toolName`. - **`sdk/typescript/test/adapters/vercel-ai.test.ts`** — `compressVercelMessages` keeps tool names on tool-result parts after a (mocked) compression that returns the assistant tool call and the tool message. ## Testing - [x] Unit tests pass (`npm test` — vitest, incl. typecheck) - [x] Type checking passes (`npm run typecheck` — tsc --noEmit) - [x] Build passes (`npm run build` — tsup) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ cd sdk/typescript && npx vitest run test/utils/format.test.ts test/adapters/vercel-ai.test.ts Test Files 2 passed (2) Tests 47 passed (47) $ npm test # full suite Test Files 19 passed | 5 skipped (24) Tests 337 passed | 33 skipped (370) Type Errors no errors $ npm run typecheck # tsc --noEmit (clean, no errors) $ npm run build # tsup CJS / ESM / DTS Build success ``` ## Real Behavior Proof - Environment: macOS (Apple silicon), Node v22.16.0, npm 10.9.2; `cd sdk/typescript && npm ci`. - Exact command / steps: a `node --input-type=module` script calling the built `toOpenAI` / `fromOpenAI` from `dist/index.js` on `main` with a user → assistant `tool-call` (`getWeather`) → `tool-result` conversation; then the three new tests on `main` and on this branch; then `npm test`, `npm run typecheck`, `npm run build`. - Observed result: On `main` the round-tripped tool-result part is `{ "type": "tool-result", "toolCallId": "call_1", "toolName": "unknown", … }` while the tool-call part keeps `"getWeather"`; the three new tests fail with `Expected: "getWeather" / Received: "unknown"`. With the change they pass and the round trip yields `toolName: "getWeather"`; full suite 337 passed | 33 skipped; typecheck and build clean. - Not tested: a live provider call. The conversion is pure; the AI SDK provider code cited above is what consumes the name. ## Runtime Rollout Safety - Rollout-managed feature(s): None. Client-side TypeScript SDK conversion fix; no server/runtime feature flag or rollout channel involved. - Minimum rollout channel: N/A. - Stable/default behavior changed: Only the buggy path — tool-result parts now carry the real tool name when the matching tool call is present; everything else is byte-identical, and the `"unknown"` fallback is unchanged for unmatched results. - Kill switch / disable path: N/A — correctness fix, not a gated feature. - Unsafe override required: No. - Qualification impact: None — existing conversion tests unchanged and green. - Rollback path: Revert this PR; no persisted state, migrations, or config to unwind. ## 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` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) N/A ## Additional Notes - A tool result whose matching assistant call was compressed away (or never present) still gets `"unknown"`; recovering a name in that case is impossible from the OpenAI shape and out of scope. - CHANGELOG left untouched — the `no-manual-changelog` CI guard owns it; the Conventional Commit title drives the entry. --- sdk/typescript/src/utils/format.ts | 5 ++- .../test/adapters/vercel-ai.test.ts | 33 +++++++++++++++ sdk/typescript/test/utils/format.test.ts | 41 +++++++++++++++++++ 3 files changed, 78 insertions(+), 1 deletion(-) diff --git a/sdk/typescript/src/utils/format.ts b/sdk/typescript/src/utils/format.ts index c2ac954b68..6a12061ae0 100644 --- a/sdk/typescript/src/utils/format.ts +++ b/sdk/typescript/src/utils/format.ts @@ -328,6 +328,8 @@ export function vercelToOpenAI(messages: any[]): OpenAIMessage[] { export function openAIToVercel(messages: OpenAIMessage[]): any[] { const result: any[] = []; + // tool_call id -> function name, so tool-result parts carry the name the AI SDK expects. + const toolNames = new Map(); for (const msg of messages) { if (msg.role === "system") { @@ -354,6 +356,7 @@ export function openAIToVercel(messages: OpenAIMessage[]): any[] { if (msg.content) parts.push({ type: "text", text: msg.content }); if (msg.tool_calls) { for (const tc of msg.tool_calls) { + toolNames.set(tc.id, tc.function.name); let input: any; try { input = JSON.parse(tc.function.arguments); } catch { input = tc.function.arguments ?? {}; } parts.push({ @@ -380,7 +383,7 @@ export function openAIToVercel(messages: OpenAIMessage[]): any[] { content: [{ type: "tool-result", toolCallId: msg.tool_call_id, - toolName: "unknown", + toolName: toolNames.get(msg.tool_call_id) ?? "unknown", output, }], }); diff --git a/sdk/typescript/test/adapters/vercel-ai.test.ts b/sdk/typescript/test/adapters/vercel-ai.test.ts index ac8634af98..cc4ab6e988 100644 --- a/sdk/typescript/test/adapters/vercel-ai.test.ts +++ b/sdk/typescript/test/adapters/vercel-ai.test.ts @@ -171,6 +171,39 @@ describe("compressVercelMessages", () => { expect(result.messages[1].role).toBe("user"); expect(result.messages[1].content[0].type).toBe("text"); }); + + it("keeps tool names on tool-result parts after compression", async () => { + mockFetch.mockResolvedValueOnce( + mockCompressResponse([ + { + role: "assistant", + content: null, + tool_calls: [ + { id: "call_1", type: "function", function: { name: "getWeather", arguments: '{"city":"Paris"}' } }, + ], + }, + { role: "tool", content: '{"tempC":21}', tool_call_id: "call_1" }, + ]), + ); + + const result = await compressVercelMessages( + [ + { + role: "assistant", + content: [{ type: "tool-call", toolCallId: "call_1", toolName: "getWeather", input: { city: "Paris" } }], + }, + { + role: "tool", + content: [ + { type: "tool-result", toolCallId: "call_1", toolName: "getWeather", output: { type: "json", value: { tempC: 21 } } }, + ], + }, + ], + { model: "gpt-4o", baseUrl: "http://localhost:8787" }, + ); + + expect(result.messages[1].content[0].toolName).toBe("getWeather"); + }); }); describe("withHeadroom", () => { diff --git a/sdk/typescript/test/utils/format.test.ts b/sdk/typescript/test/utils/format.test.ts index f6763156ad..b2a1676b10 100644 --- a/sdk/typescript/test/utils/format.test.ts +++ b/sdk/typescript/test/utils/format.test.ts @@ -638,3 +638,44 @@ describe("round-trip conversion", () => { }); }); }); + +describe("openAIToVercel tool names", () => { + it("carries the tool name from the matching tool-call onto the tool-result part", () => { + const msgs: OpenAIMessage[] = [ + { + role: "assistant", + content: null, + tool_calls: [ + { id: "tc_1", type: "function", function: { name: "getWeather", arguments: '{"city":"Paris"}' } }, + ], + }, + { role: "tool", content: '{"tempC":21}', tool_call_id: "tc_1" }, + ]; + const result = openAIToVercel(msgs); + expect(result[1].content[0]).toEqual({ + type: "tool-result", + toolCallId: "tc_1", + toolName: "getWeather", + output: { type: "json", value: { tempC: 21 } }, + }); + }); + + it("preserves tool names through a Vercel -> OpenAI -> Vercel round trip", () => { + const original = [ + { role: "user", content: [{ type: "text", text: "weather in Paris?" }] }, + { + role: "assistant", + content: [{ type: "tool-call", toolCallId: "call_1", toolName: "getWeather", input: { city: "Paris" } }], + }, + { + role: "tool", + content: [ + { type: "tool-result", toolCallId: "call_1", toolName: "getWeather", output: { type: "json", value: { tempC: 21 } } }, + ], + }, + ]; + const back = openAIToVercel(vercelToOpenAI(original)); + expect(back[2].content[0].toolName).toBe("getWeather"); + }); +}); +