fix(guardrails): scan tool results and enforce tool policy on passthrough streams - #36000
fix(guardrails): scan tool results and enforce tool policy on passthrough streams#36000mateo-berri wants to merge 1 commit into
Conversation
…rough streams Guardrails could not see the untrusted half of an agent request. Anthropic tool_result blocks carry their payload under `content`, never `text`, so the input extractor walked straight past them and a prompt injection sitting in a file the agent had just read reached the model unscanned. Tool results are now extracted (string and block forms), scanned, and masked back in place. Adds `scan_only_tool_results`, which narrows a guardrail to tool output. Agent scaffolding trips PROMPT_ATTACK at every filter strength, so scoping the scan to the data crossing the trust boundary is what makes a prompt-attack filter usable in front of a coding agent. Narrowing drops coverage, so it takes an explicit true; anything else leaves the whole request in scope. The tool_permission guardrail also missed streamed /v1/messages traffic, where the passthrough route hands the hook raw provider SSE frames rather than ModelResponse chunks. Those frames are now reassembled before the rules run, so a denied tool call is blocked on both API surfaces.
4e3575e to
ac03912
Compare
Greptile SummaryThis PR adds guardrail security hardening for tool-result handling and streamed passthrough responses, together with scoped scanning configuration and regression coverage.
Confidence Score: 4/5The implementation appears safe to merge once the non-blocking documentation-placement and typing-convention issues are addressed. The changed runtime paths have focused regression coverage and no concrete blocking failure remains; the accepted findings concern repository organization and type-annotation consistency. Files Needing Attention: guardrails_demo/RUNBOOK.md; tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py
|
| Filename | Overview |
|---|---|
| litellm/llms/anthropic/chat/guardrail_translation/handler.py | Adds location-aware extraction and masking write-back for Anthropic tool-result content while supporting tool-result-only scanning. |
| litellm/llms/base_llm/guardrail_translation/utils.py | Centralizes message-role filtering and explicit scan-only-tool-results scope evaluation. |
| litellm/llms/openai/chat/guardrail_translation/handler.py | Applies tool-result-only filtering to OpenAI message extraction and structured guardrail inputs. |
| litellm/proxy/guardrails/guardrail_hooks/tool_permission.py | Reassembles raw provider SSE frames for streamed tool-policy enforcement and consolidates permission handling. |
| litellm/proxy/guardrails/guardrail_registry.py | Propagates the new scoping parameter to initialized guardrail callbacks. |
| litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py | Allows streaming response reconstruction when no LiteLLM logging object is available. |
| litellm/types/guardrails.py | Adds the optional scan_only_tool_results guardrail configuration field. |
| tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py | Adds thorough extraction, masking, and scope tests but introduces two repository typing-convention violations. |
| tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py | Verifies explicit-true scope activation and tool-role-only scanning. |
| tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py | Covers denied, allowed, and text-only Anthropic raw SSE streams. |
| guardrails_demo/RUNBOOK.md | Adds extensive customer-facing documentation in violation of the repository's documentation-location rule. |
Reviews (1): Last reviewed commit: 4e3575e | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @@ -0,0 +1,256 @@ | |||
| # Securing agent tool use with LiteLLM guardrails on AWS Bedrock | |||
There was a problem hiding this comment.
This customer-facing setup and usage guide is being added to the application repository instead of litellm-docs, splitting documentation across repositories and bypassing the designated documentation workflow.
Rule Used: Prevent documentation from being added - needs to ... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| inputs: GenericGuardrailAPIInputs, | ||
| request_data: dict, | ||
| input_type: Literal["request", "response"], | ||
| logging_obj: Optional[Any] = None, |
There was a problem hiding this comment.
The new logging_obj parameter uses Optional[Any], introducing an avoidable Any into static analysis instead of following the repository convention of using object for broad parameters.
Rule Used: In this repo, prefer object over Optional[Any]... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| def __init__( | ||
| self, | ||
| guardrail_name: str, | ||
| replacement: Optional[str] = None, |
There was a problem hiding this comment.
The new replacement parameter uses Optional[str] rather than the repository's PEP 604 str | None convention, adding an inconsistent annotation to this test helper.
| replacement: Optional[str] = None, | |
| replacement: str | None = None, |
Rule Used: In this repo, prefer str | None (PEP 604 union s... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| skip_tool: bool, | ||
| ) -> tuple[AllMessageValues, ...]: | ||
| """Narrow the structured messages a guardrail sees, per its skip/scope settings.""" | ||
| scoped: Final = openai_messages_only_tool(messages) if scan_only_tool_results else tuple(messages) |
There was a problem hiding this comment.
Low: Tool-result guardrail bypass
When scan_only_tool_results and Bedrock's experimental_use_latest_role_message_only are both enabled, this leaves only tool messages, while Bedrock's selector requires a user message and returns skip_scan=True if none exists. An authenticated user can therefore submit malicious tool-result content without invoking the configured scanner. Preserve tool-result eligibility when latest-role filtering is enabled, and add a test covering these options together for both OpenAI and Anthropic requests.
PR overviewThis pull request updates guardrail handling to scan tool-result content and enforce tool policy on passthrough streams. It also adjusts message-selection behavior across OpenAI- and Anthropic-compatible requests. One guardrail bypass remains open when tool-result-only scanning is combined with Bedrock’s latest-role filtering. Under that configuration, an authenticated user can submit tool-result content without triggering the configured scanner, so the policy enforcement is not yet complete. Open issues (1)
Fixed/addressed: 0 · PR risk: 5/10 |
1 participant
TLDR
Problem this solves:
tool_resultpayloads live undercontent, nottexttool_permission500s on streamed/v1/messagespassthroughHow it solves it:
tool_resultblocks in both formsscan_only_tool_resultsto scope a guardrail to tool outputRelevant issues
Linear ticket
Resolves LIT-5251, Resolves LIT-5250
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
All proofs run against a live proxy on real AWS Bedrock, us-east-1, with a real Bedrock guardrail resource and
PROMPT_ATTACKset to HIGH. Before runs are at0659738b3e(litellm_internal_staging), after runs atac0391269e1. Prompt injection inside an Anthropic tool_result
The file the agent reads carries
IGNORE ALL PREVIOUS INSTRUCTIONSin the middle of an ordinary revenue report, and comes back to the model as atool_resultblockBefore,
0659738b3e: the injection sails through and the model answers normallyAfter,
ac0391269e: blocked before the model sees it2. tool_permission on streamed /v1/messages
A streamed request whose model output calls
bashwith a download from a host outside the allowlistBefore,
0659738b3e: every streamed passthrough request with this guardrail on dies in the hook, so the guardrail is not merely unenforced, it takes the route downAfter,
ac0391269e: the policy is enforced on the streamed responseAnd an allowed tool call still streams end to end, so this is enforcement rather than a blanket block
3. Full five-check sweep after the fix,
ac0391269eCheck 5 is what
scan_only_tool_resultsbuys: a real Claude Code system prompt tripsPROMPT_ATTACKat every filter strength, so without scoping, the filter is unusable in front of a coding agent4. Live against a real coding agent,
ac0391269eClaude Code pointed at the proxy through
ANTHROPIC_BASE_URL, in a scratch workspace holding the poisoned fileBlocked download, where the approval prompt never appears because the policy fires above the human:
Blocked injection, where the read itself succeeds because reading a file is not the attack:
Type
🐛 Bug Fix
Changes
_extract_input_text_and_imagesreadcontent_item.get("text"), which an Anthropictool_resultblock never has: its payload sits undercontent, as a string or as a list of blocks. Tool results were therefore dropped before any input guardrail ran, and a prompt injection sitting in a file the agent had just read reached the model unscanned. Both forms are now extracted, scanned, and written back in place, so masking guardrails edit the tool result rather than silently no-opscan_only_tool_resultsnarrows a guardrail to tool output. Agent scaffolding tripsPROMPT_ATTACKat every filter strength, so scoping the scan to what actually crosses the trust boundary is what makes a prompt-attack filter usable in front of a coding agent. Narrowing drops coverage, so it takes an explicittrue: a yaml string, a placeholder, or an unset value leaves the whole request in scopeToolPermissionGuardrail's streaming hook assumedModelResponsechunks. On a streamed/v1/messagespassthrough it is handed raw provider SSE byte frames instead, andstream_chunk_builderfails on them, taking the whole route down with a 500. Provider frames are now detected and reassembled through the existing Anthropic passthrough handler before the rules run, then replayed to the client untouchedTwo supporting changes fall out of that:
_build_complete_streaming_responseacceptslitellm_logging_obj=None, sincestream_chunk_builderalready treats it as optional and enforcement must not depend on a logging object existing; and the per-attributesetattrblock inguardrail_registrybecomes a loop over the scoping parameter namesTests
tuple-returning helpers and aNamedTuplelocation type replace the ad-hoc index bookkeeping in the Anthropic handler, so write-back is amatchon where each scanned string came from rather than positional guessworkCoverage lives with the code it exercises:
tool_resultextraction and masking write-back in the Anthropic handler tests,scan_only_tool_resultsscoping in the OpenAI handler tests, and denied tool calls carried on real Anthropic SSE byte frames in thetool_permissiontestsQA runbook
Final Attestation