fix(guardrails/headroom): stop compressing the turn the model must act on

[tin-berri]

TLDR

Problem this solves:

• Headroom compressed the user's current instruction into a hash marker
• Model answered the retrieval result instead of the request
• Compressed system prompt was discarded but still billed as savings
• Write-back fused a tool_result turn with the user turn after it

How it solves it:

• Reuse litellm's own "never compress this" policy in the guardrail
• Withhold those rows from the payload instead of pinning them back
• Group the write-back conversion by tool_call_id ownership
• Keep the model's own text in the retrieval follow-up

Relevant issues

• The Headroom guardrail no longer sends the system prompt, the live user turn, or the trailing tool exchange to /v1/compress; everything older still compresses
• The Anthropic guardrail write-back no longer merges consecutive turns, so the request the model sees keeps the boundaries the client sent
• The CCR follow-up keeps any text the model wrote alongside its tool call, and stops echoing tool calls it has no results for

Linear ticket

Resolves LIT-5018

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Screenshots / Proof of Fix

Before

After

Live proxy, a real Claude Code shaped /v1/messages request: top-level system with cache_control, a tools array whose last tool carries cache_control, an assistant tool_use, a tool_result-only user turn, then the live turn holding a slash-command expansion plus the typed line "Reply with exactly one word: ACKNOWLEDGED". A stand-in compression service on :8099 replaces any string content over 400 chars with [Headroom compressed N chars; hash=<24hex>] and serves /v1/retrieve/<hash>.

The provider account on this machine is out of credit, so the upstream is a litellm gateway serving a Bedrock-hosted Anthropic model; the request and response are Anthropic /v1/messages end to end and the calls cost real money. A direct api.anthropic.com re-run is owed once the account has credit.

Before, at abb0e36ca5:

curl -s -X POST http://127.0.0.1:4517/v1/messages \
  -H 'content-type: application/json' -H 'x-litellm-api-key: sk-lit5018' \
  -H 'x-headroom-bypass: true' -d @request.json

MODEL SAID: 'ACKNOWLEDGED'

Same request with compression on, same commit:

curl -s -X POST http://127.0.0.1:4517/v1/messages \
  -H 'content-type: application/json' -H 'x-litellm-api-key: sk-lit5018' -d @request.json

MODEL SAID: "I'm running into an issue retrieving the file contents. Both the `handler.py`
file content (hash `1846b5bf852e04b3932fa9d2`) and your follow-up message (hash
`5db0497f47e8ecd3d18b61a2`) appear to have been compressed by Headroom, but when I try
to retrieve them, I get an error saying they \"were not produced by the current request.\"
This means I can't currently see: 1. The actual contents of `hand...

--- rows the guardrail sent to /v1/compress ---
  [0] system      4118 chars  "You are Claude Code, Anthropic's official CLI for software engineering"
  [1] user          22 chars  'read handler.py for me'
  [2] assistant      4 chars  'null'
  [3] tool        6179 chars  'def handler():\n    return 42\n# padding line of the file that was read\n'
  [4] user        4846 chars  '/team expansion: End-to-end playbook for fixing a litellm bug from a L'

After, at 50b29ce2eb, byte-identical request:

MODEL SAID: 'ACKNOWLEDGED'
usage: cache_read_input_tokens 1754

--- rows the guardrail sent to /v1/compress ---
  [0] user          22 chars  'read handler.py for me'

The payload that actually leaves the gateway, captured by pointing the deployment at a recording upstream and replaying the same request through both commits:

	before abb0e36ca5	after 50b29ce2eb
messages out	3	4
[2] user	tool_result(124b), text(91b)	tool_result(6392b)
[3] user	absent, merged into [2]	text(4797b), text(103b)

Before, the 6179-char file body and the live instruction are both 90-byte markers sharing one turn. After, both are intact and the live turn is its own message.

Compression still runs on history. Same rig, a longer session with an older tool exchange ahead of the current one:

=== before: 7 messages out | tools ['Read', 'Write', 'headroom_retrieve'] ===
  [2] user      tool_result(124b)      <-- MARKER
  [3] assistant text(91b)              <-- MARKER
  [6] user      tool_result(124b), text(91b)   <-- MARKER

=== after: 8 messages out | tools ['Read', 'Write', 'headroom_retrieve'] ===
  [2] user      tool_result(124b)      <-- MARKER
  [3] assistant text(91b)              <-- MARKER
  [6] user      tool_result(2372b)
  [7] user      text(4797b), text(103b)

The old exchange is still compressed and headroom_retrieve is still injected, so CCR is untouched; only the current exchange and the live turn are held back.

Type

🐛 Bug Fix

Changes

• HeadroomGuardrail.apply_guardrail consults get_protected_indices and withholds those rows from /v1/compress instead of sending everything
• Protection is expanded over tool exchanges, so a protected assistant tool call cannot be answered by a marker
• A compress response whose row count differs from the request goes through the fail policy; fail-open returns the caller's own inputs object
• The Anthropic write-back converts tool exchanges as units rather than the whole array, so consecutive turns stop merging
• The CCR follow-up keeps assistant text on all three surfaces and echoes only the retrieve calls it answers
• Grouping yields from a generator consumed once, so it stays linear in the message count

Things a reviewer will ask about

Why does a single-turn request stop compressing? [system, user] is entirely protected, so there is nothing to send and the guardrail returns early. That is the same policy compress() has always applied, and compresr defaults to it too (compress_last_user=False, compress_system=False). If a deployment wants its current prompt compressed, that wants an explicit opt-in field rather than the old behavior of compressing whatever happened to be largest.

Why withhold the rows instead of restoring them after the service replies? Both keep the instruction intact, but sending rows whose result we discard reports savings that never land; that is exactly the system-prompt over-count in this ticket. The cost is that a query-aware compressor no longer sees the newest user message, which is called out in a comment at the withholding site.

Why not convert the write-back one message at a time? It breaks tool pairing. With modify_params on, converting an assistant row apart from its results makes sanitize_messages_for_tool_calling read an orphaned tool call, answer it with a synthetic "tool execution skipped" result, and drop the real one; the file body disappears. Grouping by tool_call_id ownership keeps turns separate without splitting an exchange, and tests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py::test_write_back_keeps_real_tool_results_under_modify_params fails if anyone regresses to per-row.

Why a wall-clock assertion in the grouping test? The first version of that helper accumulated groups by rebuilding a tuple each iteration, which is O(n^2) over an array the caller controls: 20k messages took 312ms, 100k would take minutes. It is linear now (16ms at 100k, 2.06us per message flat from 10k to 80k), and the test fails at 8.97s against a 3s ceiling if anyone reintroduces the accumulator, so the ceiling carries ~200x headroom in the passing direction. The write-back path as a whole is linear on both sides of this change; converting per exchange rather than once costs a 1.55x constant, which is the extra messages the fix deliberately stops merging.

What is deliberately not in here. The ticket also reported that appending headroom_retrieve after the client's cache_control breakpoint breaks prompt caching. Measured with 24 tools and a breakpoint on the last one, a request creates the cache and the next one reads it back at the same token count, with the injected tool sitting after the breakpoint; steady-state caching still hits, so there is no mechanism to fix and the append order is unchanged. Separately, /v1/responses builds structured_messages, hands them to the guardrail and never reads the result back, so compression is computed and discarded there while savings are still recorded; that is a different handler and gets its own ticket.

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

---

Note

Medium Risk
Changes proxy guardrail request shaping and Anthropic message conversion on every guarded /v1/messages path; behavior is heavily tested but mistakes could alter what the model sees or break tool-call sequences.

Overview
Fixes LIT-5018: Headroom was compressing the live user instruction, system prompt, and trailing tool exchange, and Anthropic guardrail write-back was merging separate turns.

Headroom now shares get_protected_indices with built-in compress()—system rows, last user, and last assistant are withheld from /v1/compress, with protection expanded over whole tool exchanges via group_tool_exchanges. Compressed rows are merged back with _restore_protected_messages; mismatched row counts fail (or fail-open returns the original inputs by identity). Retrieval follow-ups reuse shared assistant_text_from_response and only echo retrieve tool calls.

Anthropic write-back converts group_tool_exchanges units through anthropic_messages_pt so tool-result turns stay separate from the following user turn and tool pairing survives modify_params.

group_tool_exchanges is a new linear-time helper in the prompt factory; assistant_text_from_response moves to content_text.py for Compresr/Headroom reuse. Tests and CI include the compression test path.

^{Reviewed by Cursor Bugbot for commit 50b29ce. Bugbot is set up for automated code reviews on this repo. Configure here.}

[greptile-apps]

Greptile Summary

The PR prevents Headroom from compressing protected prompt turns and preserves message boundaries and assistant text during compression write-back and retrieval follow-ups

• Reuses the shared protected-message policy for system, live user, assistant, and associated tool-exchange rows
• Rejects compression responses whose row count cannot be safely interleaved with withheld messages
• Converts Anthropic write-back in tool-call-owned groups to preserve client turn boundaries
• Shares assistant-text extraction across Headroom and Compresr retrieval follow-ups
• Adds focused regression coverage and includes compression tests in the miscellaneous unit-test workflow

Confidence Score: 5/5

The PR appears safe to merge

No blocking failure remains

Important Files Changed

Filename	Overview
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py	Withholds protected turns from compression, restores compressed history positionally, validates response cardinality, and preserves assistant text in retrieval follow-ups
litellm/llms/anthropic/chat/guardrail_translation/handler.py	Converts structured-message write-back by tool-exchange ownership instead of merging the complete message array
litellm/litellm_core_utils/prompt_templates/factory.py	Adds ordered grouping of assistant tool calls with their contiguous owned tool-result rows
litellm/compression/compress.py	Exposes the existing protected-index policy for reuse by compression guardrails
litellm/proxy/guardrails/guardrail_hooks/content_text.py	Centralizes assistant-text extraction for chat completions, Anthropic Messages, and Responses API shapes
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py	Adds regression coverage for protected turns, history compression, fail policy, and retrieval follow-up reconstruction
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py	Covers Anthropic write-back boundaries and preservation of real tool results under parameter sanitization

_{Reviews (2): Last reviewed commit: "fix(guardrails/headroom): stop compressi..." | Re-trigger Greptile}

[veria-ai]

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 1 · PR risk: 0/10

[codecov]

Codecov Report

❌ Patch coverage is 69.23077% with 36 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xy/guardrails/guardrail_hooks/headroom/headroom.py	46.26%	36 Missing ⚠️

📢 Thoughts on this report? Let us know!

[tin-berri]

@greptileai

[tin-berri]

bugbot run

[cursor]

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

^{Reviewed by Cursor Bugbot for commit 50b29ce. Configure here.}

[mateo-berri]

LGTM. Thanks!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_lit5018_headroom_live_turn (50b29ce) with litellm_internal_staging (6e26087)¹}

Footnotes

1. No successful run was found on litellm_internal_staging (c3da121) during the generation of this report, so 2593168 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩