litellm/tests/test_litellm/proxy/guardrails
Krrish Dholakia 50b936c75e
feat(guardrails/headroom): add CCR (compress-cache-retrieve) via agentic loop (#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".
2026-06-30 19:19:34 -07:00
..
guardrail_hooks feat(guardrails/headroom): add CCR (compress-cache-retrieve) via agentic loop (#31681) 2026-06-30 19:19:34 -07:00
test_content_filter_path_traversal.py Litellm OSS Staging 010626 (#29422) 2026-06-01 21:42:51 -07:00
test_content_utils.py fix(guardrails): redact `data["input"]` for Responses-API mask paths 2026-05-01 04:41:57 +00:00
test_custom_code_security.py fix(guardrails): return HTTP 400 for litellm content filter blocks (#28418) 2026-05-30 14:28:46 -07:00
test_deferred_guardrail_logging.py feat(proxy): enforce key/team guardrails on bedrock passthrough routes (#30194) 2026-06-12 07:23:29 -07:00
test_guardrail_coverage.py fix(tests): use canonical litellm_enterprise import path (#27699) 2026-05-12 12:32:57 -07:00
test_guardrail_endpoints.py feat(guardrails): wire apply_guardrail into proxy logging callbacks (#28970) 2026-05-28 09:41:02 -07:00
test_guardrail_registry.py fix(guardrails): stop re-initializing DB guardrails on every poll (#30542) 2026-06-16 11:17:49 -07:00
test_init_guardrails.py Litellm fix GitHub action testing (#11163) 2025-05-26 14:41:42 -07:00
test_llm_as_a_judge.py feat(guardrails): LLM-as-a-Judge guardrail (#26360) 2026-04-24 17:15:32 -07:00
test_mcp_jwt_signer.py fix(mcp): JWT on tools/list and REST tools/call server resolution (#28227) 2026-05-20 13:31:44 -07:00
test_pillar_guardrails.py chore(proxy): harden request control fields 2026-04-29 22:35:17 -07:00
test_prompt_security_guardrails.py fixing prompt-security's guardrail implementation (#19374) 2026-01-21 20:09:40 -08:00
test_qostodian_nexus_guardrail.py feat: add Qohash Nexus guardrail hook (#24927) 2026-05-01 17:26:32 +05:30