feat(guardrails/headroom): add CCR (compress-cache-retrieve) via agentic loop - #31681
Conversation
|
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d643a0a571
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| function = tool.get("function") | ||
| if tool.get("type") == "function" and isinstance(function, dict): | ||
| if function.get("name") == HEADROOM_RETRIEVE_TOOL_NAME: | ||
| return True | ||
| return False |
There was a problem hiding this comment.
For /v1/messages requests, the Anthropic guardrail translation converts guardrailed_tools back from OpenAI format into Anthropic tools (litellm/llms/anthropic/chat/guardrail_translation/handler.py lines 167-175), which makes the injected retrieve tool look like {"type": "custom", "name": "headroom_retrieve", ...}. This predicate only accepts OpenAI type=function tools, so async_should_run_agentic_loop returns false even when Claude emits the headroom_retrieve tool_use; Headroom CCR never retrieves the compressed content on the Anthropic Messages path.
Useful? React with 👍 / 👎.
| tool_results.append( | ||
| { | ||
| "role": "tool", | ||
| "tool_call_id": tc.get("id"), | ||
| "content": content, | ||
| } | ||
| ) |
There was a problem hiding this comment.
When the Anthropic tool_use branch in _extract_headroom_tool_calls is used, this plan still appends OpenAI chat-completion tool messages (role: "tool" with tool_call_id). _execute_anthropic_agentic_plan sends patch.messages to anthropic_messages.acreate, whose existing interception implementations build an assistant tool_use content block followed by a user tool_result content block; with these OpenAI-shaped messages the follow-up is not a valid Anthropic Messages tool-result turn, so the retrieval loop will fail instead of producing the final answer.
Useful? React with 👍 / 👎.
Greptile SummaryThis PR extends
Confidence Score: 5/5Safe to merge. The three API-surface bugs identified in the previous review are all fixed with correct implementations and regression tests using plain-dict responses rather than MagicMock objects. The core CCR logic is well-structured: hash validation is scoped per litellm_call_id, each API surface gets a distinct and correctly shaped follow-up, and the shared helpers in factory.py centralise format detection so it cannot silently regress per-surface. Both findings are edge-case robustness concerns with no correctness bugs on the happy path. The response-type detection helpers in headroom.py use single-attribute heuristics that could misfire if a future response type carries both output and content as lists.
|
| Filename | Overview |
|---|---|
| litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py | Adds CCR agentic loop to HeadroomGuardrail: injects a retrieve tool on compression, validates hashes per-call-id, and builds correct follow-up messages for chat, Responses API, and Anthropic shapes. One edge-case gap in the UUID fallback propagation. |
| litellm/litellm_core_utils/prompt_templates/factory.py | Adds shared cross-surface helpers: NormalizedToolCall TypedDict, get_tool_calls_from_response (chat/responses/Anthropic), has_tool_with_name (OpenAI and Anthropic native shapes). Clean, well-abstracted additions with correct fallthrough behavior. |
| tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py | Comprehensive new test coverage for CCR: tool injection, agentic loop detection across all three API surfaces, hash validation/rejection, per-call-id security scoping, and correct follow-up message shapes for Anthropic and Responses API. Tests use plain-dict responses (not just MagicMock) to catch TypedDict runtime bugs. |
| tests/llm_translation/test_prompt_factory.py | New unit tests for get_tool_calls_from_response and has_tool_with_name covering chat, responses API, Anthropic (MagicMock and plain-dict), empty, and non-list tool inputs. |
| CLAUDE.md | Adds one coding convention bullet: prefer shared helpers in factory.py when branching on API surface, rather than duplicating format-detection logic per module. |
Reviews (12): Last reviewed commit: "fix(guardrails/headroom): match Anthropi..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
…via agentic loop
When Headroom's /v1/compress returns messages containing hash markers
(hash=[a-f0-9]{24}), inject a headroom_retrieve tool into the request.
When the LLM calls that tool, intercept via async_should_run_agentic_loop
and async_build_agentic_loop_plan, call GET /v1/retrieve/{hash} on the
Headroom sidecar, and replay the LLM with the original content as a tool
result -- all transparent to the caller.
…AI and Anthropic response formats
…t detection in CCR loop
…06 ruff violations
7b4b9f3 to
08d1d75
Compare
…rror to fix BLE001
|
@greptileai review |
…or CCR tool calls
|
@greptileai review |
|
@greptileai review |
…urrent request Previously any LLM-supplied hash in a headroom_retrieve tool call was forwarded to the Headroom retrieve API, letting a crafted tool call fetch arbitrary cached content. Validate the hash against the set produced by compressing the current request's messages before calling retrieve.
|
@greptileai review |
Merging this PR will improve performance by 14.45%
|
| Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|
| ⚡ | test_completion_simple_message |
4.6 ms | 4 ms | +14.45% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing litellm_headroom_ccr (57b37c9) with litellm_internal_staging (be4d0d8)
|
@greptileai review |
…ses API replay shape Hash validation now also checks an in-memory cache of hashes actually returned by /v1/compress, not just whether the hash text appears somewhere in the request's messages. The message-text check alone is forgeable: an attacker can plant a hash-shaped string in their own prompt and have it treated as valid. Responses API follow-up now emits function_call/function_call_output items keyed by call_id instead of chat-style assistant/tool messages, since the Responses API does not accept the latter as input. Also fixes call_id/id field priority when extracting tool calls from Responses API output, since call_id (not id) is what must match between the function_call and its output.
|
@greptileai review |
UP037 flags quotes on annotations that are already lazily evaluated via `from __future__ import annotations`.
|
@greptileai review |
Functional under asyncio_mode=auto, but every other async test in the file has the decorator for consistency.
|
Code review result from my claude agent: Requesting changes for three reasons:
Required fixes:
|
… replay shape Two real gaps found in review: 1. The instance-wide issued-hash cache combined with a message-text check did not actually scope retrieval to the request that produced the hash. A hash issued for request A stays in the shared cache until TTL expiry, and the message-text check is satisfied by any request whose own messages happen to echo that hash string. Request B could plant A's hash in its own prompt and retrieve A's content. Fixed by keying the issued-hash cache by litellm_call_id, matching the pattern already used in compression_interception: a hash is only honored when it was issued under the exact call_id resolving for the current request. 2. The Anthropic Messages replay path fell through to the chat-style assistant/tool-message builder, which Anthropic does not accept. Anthropic requires the tool_use block echoed in an assistant message paired with a tool_result block in a user message, keyed by tool_use_id. Added a dedicated branch for this shape.
|
Confirmed both correctness/security points against the existing compression_interception implementation, and fixed both in 16b0c4a:
|
|
@greptileai review |
Add a bullet to the coding-conventions list: look for or add a shared helper when logic branches on API surface (chat completions vs Anthropic Messages vs Responses API), instead of duplicating format-detection per module.
… shared cross-API tool util Live e2e testing against the real Anthropic API surfaced two bugs the mocked unit tests couldn't catch because they used MagicMock responses instead of realistic response shapes: 1. has_headroom_retrieve_tool only recognized OpenAI-shaped function tools. By the time an Anthropic Messages response reaches the agentic-loop gate, the tool this guardrail injected has already been transformed into Anthropic's native shape (type: "custom", top-level "name"), so the gate never fired for real Anthropic requests. 2. AnthropicMessagesResponse is a TypedDict, so real responses are plain dicts at runtime, not objects with attribute access. The extractors and format detectors used bare getattr(), which silently returns nothing for dict responses instead of reading the actual key. Extracted the cross-API-surface tool-call extraction and tool-presence check into litellm/litellm_core_utils/prompt_templates/factory.py (get_tool_calls_from_response, has_tool_with_name) so this format fragmentation is handled in one place instead of being duplicated per-guardrail, and reused the existing repair-aware parse_tool_call_arguments from common_utils instead of a naive json.loads. headroom.py now delegates to these shared helpers. Confirmed live against the real Anthropic API: the retrieve loop now fires and successfully retrieves the correct hash's content through the full compress -> tool-call -> retrieve -> replay round-trip.
|
@greptileai review |
|
Follow-up: while validating the /v1/messages fix live against the real Anthropic API (not just mocked unit tests), found and fixed two more real bugs in the same area:
Both are confirmed fixed by re-running the same live test against the real API: the retrieve loop now fires and successfully retrieves content through the full round-trip. Extracted the cross-surface tool-call detection into litellm_core_utils/prompt_templates/factory.py (get_tool_calls_from_response, has_tool_with_name) so this doesn't get re-broken per-guardrail, and added regression tests using plain dict responses (not MagicMock) so the shape bug specifically can't regress silently again. |
Use lowercase list/dict generics in the new factory.py tool-call helpers instead of typing.List/Dict, drop the now-unused Tuple import in headroom.py, and reorder the new factory import ahead of the llms.custom_httpx import to satisfy import sorting.
|
@greptileai review |
| if tool.get("type") == "function" and isinstance(function, dict): | ||
| if function.get("name") == tool_name: | ||
| return True | ||
| if tool.get("type") == "custom" and tool.get("name") == tool_name: |
There was a problem hiding this comment.
Non-blocker:
Anthropic official tool-use docs define client tools by name + input_schema; type: "custom" is not required, so has_tool_with_name should treat no-type Anthropic tools as valid matches too.
Anthropic's documented client tool format is just name + input_schema; type: "custom" is only one possible value, not a requirement. Match any non-OpenAI-shaped tool on its top-level name instead of requiring type == "custom".
|
Fixed in 5bc3838. has_tool_with_name now matches any non-OpenAI-shaped tool on its top-level name regardless of the type field (or its absence), instead of requiring type == "custom". Added a regression test for the no-type-field case. |
|
@greptileai review |
50b936c
into
litellm_internal_staging
…tic loop (BerriAI#31681) * feat(guardrails/headroom): add CCR (compress-cache-retrieve) support via agentic loop When Headroom's /v1/compress returns messages containing hash markers (hash=[a-f0-9]{24}), inject a headroom_retrieve tool into the request. When the LLM calls that tool, intercept via async_should_run_agentic_loop and async_build_agentic_loop_plan, call GET /v1/retrieve/{hash} on the Headroom sidecar, and replay the LLM with the original content as a tool result -- all transparent to the caller. * style: run ruff format on headroom guardrail and tests * fix(guardrails/headroom): detect headroom_retrieve calls in both OpenAI and Anthropic response formats * test(guardrails/headroom): add test for Anthropic content block format detection in CCR loop * ci: trigger CI checks * fix(guardrails/headroom): replace List/Dict with list/dict to fix UP006 ruff violations * fix(guardrails/headroom): replace except Exception with except ValueError to fix BLE001 * fix(guardrails/headroom): add Responses API output format detection for CCR tool calls * refactor(guardrails/headroom): extract format-specific helpers to fix C901 complexity * fix(guardrails/headroom): scope CCR retrieval to hashes produced by current request Previously any LLM-supplied hash in a headroom_retrieve tool call was forwarded to the Headroom retrieve API, letting a crafted tool call fetch arbitrary cached content. Validate the hash against the set produced by compressing the current request's messages before calling retrieve. * fix(guardrails/headroom): track issued hashes server-side, fix Responses API replay shape Hash validation now also checks an in-memory cache of hashes actually returned by /v1/compress, not just whether the hash text appears somewhere in the request's messages. The message-text check alone is forgeable: an attacker can plant a hash-shaped string in their own prompt and have it treated as valid. Responses API follow-up now emits function_call/function_call_output items keyed by call_id instead of chat-style assistant/tool messages, since the Responses API does not accept the latter as input. Also fixes call_id/id field priority when extracting tool calls from Responses API output, since call_id (not id) is what must match between the function_call and its output. * fix(guardrails/headroom): drop redundant quoted type annotations UP037 flags quotes on annotations that are already lazily evaluated via `from __future__ import annotations`. * test(guardrails/headroom): add missing pytest.mark.asyncio decorators Functional under asyncio_mode=auto, but every other async test in the file has the decorator for consistency. * fix(guardrails/headroom): scope CCR hashes per call_id, fix Anthropic replay shape Two real gaps found in review: 1. The instance-wide issued-hash cache combined with a message-text check did not actually scope retrieval to the request that produced the hash. A hash issued for request A stays in the shared cache until TTL expiry, and the message-text check is satisfied by any request whose own messages happen to echo that hash string. Request B could plant A's hash in its own prompt and retrieve A's content. Fixed by keying the issued-hash cache by litellm_call_id, matching the pattern already used in compression_interception: a hash is only honored when it was issued under the exact call_id resolving for the current request. 2. The Anthropic Messages replay path fell through to the chat-style assistant/tool-message builder, which Anthropic does not accept. Anthropic requires the tool_use block echoed in an assistant message paired with a tool_result block in a user message, keyed by tool_use_id. Added a dedicated branch for this shape. * docs: note proactive API-fragmentation helper convention Add a bullet to the coding-conventions list: look for or add a shared helper when logic branches on API surface (chat completions vs Anthropic Messages vs Responses API), instead of duplicating format-detection per module. * fix(guardrails/headroom): fix Anthropic tool-shape detection, extract shared cross-API tool util Live e2e testing against the real Anthropic API surfaced two bugs the mocked unit tests couldn't catch because they used MagicMock responses instead of realistic response shapes: 1. has_headroom_retrieve_tool only recognized OpenAI-shaped function tools. By the time an Anthropic Messages response reaches the agentic-loop gate, the tool this guardrail injected has already been transformed into Anthropic's native shape (type: "custom", top-level "name"), so the gate never fired for real Anthropic requests. 2. AnthropicMessagesResponse is a TypedDict, so real responses are plain dicts at runtime, not objects with attribute access. The extractors and format detectors used bare getattr(), which silently returns nothing for dict responses instead of reading the actual key. Extracted the cross-API-surface tool-call extraction and tool-presence check into litellm/litellm_core_utils/prompt_templates/factory.py (get_tool_calls_from_response, has_tool_with_name) so this format fragmentation is handled in one place instead of being duplicated per-guardrail, and reused the existing repair-aware parse_tool_call_arguments from common_utils instead of a naive json.loads. headroom.py now delegates to these shared helpers. Confirmed live against the real Anthropic API: the retrieve loop now fires and successfully retrieves the correct hash's content through the full compress -> tool-call -> retrieve -> replay round-trip. * fix(guardrails/headroom): fix ruff-strict UP006/I001 budget violations Use lowercase list/dict generics in the new factory.py tool-call helpers instead of typing.List/Dict, drop the now-unused Tuple import in headroom.py, and reorder the new factory import ahead of the llms.custom_httpx import to satisfy import sorting. * fix(guardrails/headroom): match Anthropic tools without a type field Anthropic's documented client tool format is just name + input_schema; type: "custom" is only one possible value, not a requirement. Match any non-OpenAI-shaped tool on its top-level name instead of requiring type == "custom".
… /v1/retrieve/{hash} (#702)
Lossy /v1/compress rewrites now advertise their retrieval hash in the
guardrail's regex-locked hash=<24hex> form (first 24 hex of the blake3
content address; the first 16 are the existing tee name). The new
/v1/retrieve/{hash} endpoint resolves it back to the verbatim original
as {"original_content": ...} — the exact contract LiteLLM's headroom
guardrail CCR loop (BerriAI/litellm#31681) speaks. Marker shape is
pinned by contract tests so drift fails CI; hashes are pure functions
of content, so stubs stay byte-stable (#498). Local handles unchanged.
Co-authored-by: Cursor <cursoragent@cursor.com>
… /v1/retrieve/{hash} (#702)
Lossy /v1/compress rewrites now advertise their retrieval hash in the
guardrail's regex-locked hash=<24hex> form (first 24 hex of the blake3
content address; the first 16 are the existing tee name). The new
/v1/retrieve/{hash} endpoint resolves it back to the verbatim original
as {"original_content": ...} — the exact contract LiteLLM's headroom
guardrail CCR loop (BerriAI/litellm#31681) speaks. Marker shape is
pinned by contract tests so drift fails CI; hashes are pure functions
of content, so stubs stay byte-stable (#498). Local handles unchanged.
Co-authored-by: Cursor <cursoragent@cursor.com>
4 participants
Relevant issues
Linear ticket
Pre-Submission checklist
Screenshots / Proof of Fix
Headroom sidecar running at https://headroom-mock.onrender.com.
Run the proxy:
Config (
headroom_ccr_test.yaml):/v1/chat/completions
Proxy logs showing the full CCR loop:
Response (LLM answered correctly with retrieved full context):
{"choices":[{"message":{"content":"2 + 2 equals 4."}}]}The full CCR round-trip: 5078 tokens compressed to 519, LLM retrieved 18313 chars of original content via the agentic loop, answered correctly.
/v1/messages (Anthropic Messages API)
A secret string was placed past the compression truncation point (padded with filler text) so the model can only answer by retrieving the compressed content:
Proxy logs showing the full CCR loop against the real Anthropic API:
This surfaced two real bugs that mocked unit tests (which used
MagicMockresponses) had been masking:has_headroom_retrieve_toolonly recognized OpenAI-shaped function tools; by the time a Messages API response reaches the gate, the injected tool has already been transformed into Anthropic's native{"type": "custom", "name": ...}shape, so the gate never fired for real Anthropic traffic.AnthropicMessagesResponseis aTypedDict, so real responses are plain dicts at runtime, not objects with attribute access — the extractors used baregetattr(), which silently returns nothing for dict responses.Both are fixed, with regression tests using plain-dict responses (not
MagicMock) so this can't regress silently again./v1/responses
Routes through the same shared
get_tool_calls_from_response/has_tool_with_namehelpers and the dedicatedfunction_call/function_call_outputreplay branch; covered by unit tests (test_async_should_run_agentic_loop_detects_responses_api_output_format,test_async_build_agentic_loop_plan_builds_responses_api_function_call_items) sinceResponsesAPIResponseis a real Pydantic model at runtime (not a TypedDict), so it wasn't affected by the dict-vs-object bug above.Type
New Feature
Changes
Extends the existing HeadroomGuardrail with CCR (Compress-Cache-Retrieve) support via LiteLLM's agentic loop infrastructure.
How it works:
The Headroom sidecar must run on the same host as LiteLLM (Headroom's loopback-only restriction on /v1/retrieve in production). Configure it as a guardrail:
API surface coverage (verified live against real OpenAI and Anthropic APIs, plus unit tests for the Responses API replay shape):
New shared helpers (
litellm/litellm_core_utils/prompt_templates/factory.py):get_tool_calls_from_response(response)— extracts normalized tool calls from a response regardless of API surface (chat completions, Responses API, or Anthropic Messages)has_tool_with_name(tools, name)— checks whether a tools list includes a given tool, regardless of shape (OpenAI function tools or Anthropic's native tool shape)These exist so future guardrails/integrations that need to detect tool calls across API surfaces don't have to re-derive this parsing per module (see the new CLAUDE.md convention note).
Security: the CCR retrieval loop validates that a requested hash was actually issued by this exact request's own compression call (scoped by litellm_call_id, matching the pattern used by compression_interception), not just that the hash string appears somewhere in the request's messages, which would be forgeable via prompt injection.