feat(mcp): scan and mask MCP tool results via post_mcp_call guardrails - #35155
Conversation
|
|
Greptile SummaryAdds outbound MCP result guardrails and addresses the previously reported gaps:
Confidence Score: 5/5The PR appears safe to merge No blocking failure remains
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/server.py | Moves outbound guardrail enforcement onto the shared tool-execution return path before logging and forwards the rewritten result |
| litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py | Implements unified scanning, masking, and rejection for MCP text and structured output |
| litellm/proxy/_experimental/mcp_server/utils.py | Adds shared MCP result mutation and bounded structured-content traversal helpers |
| litellm/proxy/_experimental/mcp_server/rest_endpoints.py | Returns guardrailed results and preserves guardrail rejections instead of treating them as ignorable logging failures |
| litellm/proxy/utils.py | Adds dispatch for apply_guardrail-based callbacks configured for post_mcp_call |
| litellm/responses/mcp/litellm_proxy_mcp_handler.py | Runs outbound result guardrails before post-call callbacks on Responses and chat-completions MCP paths |
| ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx | Exposes post_mcp_call in the guardrail mode selector |
Reviews (6): Last reviewed commit: "feat(guardrails): scan and mask MCP tool..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
eab8e79 to
86a749f
Compare
|
Real bypass, and fixing it properly closed a second gap I had listed as a follow-up. Fixed in The guardrail was called from inside Rather than widen that condition, the guardrail moved out of the logging function entirely and onto The two tests that asserted the old location were retargeted to the new one, and a third pins the bypass directly: a call with 356 tests pass across the MCP server, REST endpoint, guardrail-translation, proxy-utils, and cisco MCP suites; both lint budget gates are within ceiling. Live re-verified on the proxy plus real presidio rig, masking output unchanged. The PR body's follow-up list is updated: @greptileai please review the current head |
86a749f to
226020c
Compare
|
Both findings addressed in Responses-path ordering. You were right, and it was two problems on that path, not one. structuredContent. I had documented this as a follow-up; you are right that it should not merge that way. so a masked Masking it in place needs the nested-JSON rewrite from #35142, which is not in this branch. Rather than duplicate that walk and guarantee a merge conflict, this now fails closed: the handler diffs each original text against the masked one to recover the spans the guardrail actually removed, and rejects with a 400 if any of those spans still appear in the serialized The diff step matters and my first attempt got it wrong: comparing whole text blocks never fires, because the guardrail replaces an entity inside a sentence, so Three tests cover it: a repeat of a masked value blocks, unrelated structured data does not, and a clean result is untouched. Upgrading block to mask is now a small follow-up once #35142 lands, and the PR body says so. 744 tests pass across the MCP server, REST endpoint, guardrail-translation, proxy-utils, cisco MCP, and Responses suites; both lint budget gates within ceiling; live re-verified on the proxy plus real presidio rig with masking output unchanged. On the one red check: @greptileai please review the current head |
226020c to
f85765d
Compare
|
You were right that the reject-only guard did not close this, and the distinction you drew is the one that mattered: it only caught values already redacted from Both Live proof against a tool whose PII exists only in structuredContent (text says nothing sensitive): Masked in place, shape preserved, and the float left alone. Four tests cover it, including the exact gap you named ( Two details worth noting. The walk fails closed at a depth cap rather than truncating, so nothing passes unscanned, and it carries recursion-detector ignore entries with that reason. And it raises a plain 745 tests pass across the MCP server, REST endpoint, guardrail-translation, proxy-utils, cisco MCP, and Responses suites; both lint budget gates within ceiling; masking of the plain-text case re-verified live and unchanged.
@greptileai please review the current head |
f85765d to
01ebb91
Compare
|
CI caught a real bug in my own change, now fixed in
Defining a custom exception in a module that is reloaded by design was the wrong shape. The walk now returns While pinning that I introduced and then removed a second problem, which is worth flagging since I nearly misreported it. My first regression test called So the test is now a direct contract assertion on the helper ( The residual two 3135 tests pass across the MCP server directory, proxy-utils, cisco MCP, and Responses suites; both lint budget gates within ceiling; structuredContent masking re-verified live and unchanged. @greptileai please review the current head |
Guardrails could only see the MCP tool call request (pre_mcp_call / during_mcp_call); the tool result went back to the client unscanned, so a tool that returns sensitive data bypassed every configured guardrail. Adds a `post_mcp_call` event hook that runs after the tool executes and routes the result through the unified apply_guardrail seam, so a text guardrail (e.g. presidio) can mask sensitive values in the tool output or reject the result without any MCP-specific code of its own. - MCPGuardrailTranslationHandler.process_output_response now extracts the tool result's text content into GenericGuardrailAPIInputs["texts"], calls apply_guardrail with input_type="response", and writes the returned text back into the content list in place (the logging payload already references that object, so a copy would leave the unmasked text in the spend log) - ProxyLogging.post_mcp_call_hook dispatches guardrails that implement apply_guardrail, gated on should_run_guardrail(post_mcp_call); guardrails implementing async_post_mcp_tool_call_hook keep their existing dispatch and are not run twice - both MCP tool-call paths (mcp_server and the Responses API handler) now honor the rewritten result, and the REST path no longer swallows a guardrail rejection as a logging failure - shared, duck-typed MCP content helpers live in mcp_server/utils.py next to extract_mcp_tool_result_error_message - documents that async_post_mcp_tool_call_hook's return value is discarded by every call site, so that hook only takes effect by mutating in place
01ebb91 to
268e884
Compare
|
Fixed in Dictionary keys and non-string scalars in
Three tests cover it, and the mutant that stops collecting them dies: A sensitive key blocks, a sensitive numeric value blocks, and ordinary keys and numbers pass through untouched alongside normal text masking. The collector shares the depth-cap sentinel contract with 3138 tests pass across the MCP server directory, proxy-utils, cisco MCP, and Responses suites, with only the two @greptileai please review the current head |
TLDR
Problem this solves:
How it solves it:
post_mcp_callmode, selectable in the dashboardapply_guardrailseamRelevant issues
Linear ticket
Resolves LIT-4935
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
Live proxy with a real presidio analyzer + anonymizer pair (no mocks) and a stdio MCP server whose
lookup_customertool deliberately returns sensitive data, so what the client receives is the evidence.Before, on staging
440b1bcf65. There is no mode that scans a tool result, so the strongest available configuration ismode: pre_mcp_call, and the value reaches the client untouched:{"content":[{"type":"text","text":"customer C-1001: email jane.doe@example.com, phone 415-555-0132, account balance 42.00"}],"isError":false}Configuring the mode this PR adds is not a partial fix on staging, it is a boot failure, which is the cleanest statement of what was missing:
After, at
268e884eed, withmode: post_mcp_call. Email and phone masked, the non-sensitive balance untouched:{"content":[{"type":"text","text":"customer C-1001: email <EMAIL_ADDRESS>, phone <PHONE_NUMBER>, account balance 42.00"}],"isError":false}After, with
EMAIL_ADDRESS: BLOCKinstead ofMASK. The rejection propagates rather than degrading to an unguarded result:{"detail":{"error":"blocked_pii_entity","message":"Blocked entity detected: EMAIL_ADDRESS by Guardrail: presidio-mcp-output"}}A second tool, returning PII only in
structuredContentwhile its text says nothing sensitive, proves the structured path:For the UI half,
post_mcp_callappears in the guardrail mode dropdown as "After MCP Tool Call - Runs after MCP tool execution and checks the tool result" at http://localhost:4000/ui/?page=guardrails -> Add Guardrail -> Mode.Type
🆕 New Feature
Changes
Guardrails already ran on the MCP tool-call request through
pre_mcp_callandduring_mcp_call. The result went back to the client unscanned:MCPGuardrailTranslationHandler.process_output_responsewas a stub logging "Output processing not implemented for MCP tools", and the only outbound seam,async_post_mcp_tool_call_hook, was implemented by exactly one guardrail (cisco).This adds a
post_mcp_callevent hook and an MCP-shaped dispatcher.ProxyLogging.post_mcp_call_hookgates onshould_run_guardrail(post_mcp_call)plus the presence ofapply_guardrail, then routes the result throughMCPGuardrailTranslationHandler.process_output_responsewithinput_type="response", resolving the handler at runtime viaload_guardrail_translation_mappings()[CallTypes.call_mcp_tool]soproxy/utils.pynever importsmcpat module scope. The upshot is that a text guardrail such as presidio can mask an MCP tool result with no MCP-specific code of its own.Reusing
post_call_success_hookwas considered and rejected: itsother_callbacksloop (proxy/utils.py:2404) is ungated, so every registered non-guardrailCustomLoggerwould be handed aCallToolResultwhere it expects an LLM response.unified_guardrail.pyis untouched.The result's content list is rewritten in place rather than replaced with a copy, because the logging payload captured before the hook runs references that same list, so a copy would leave the unmasked text in the spend log and the OTel span. The shape handling is duck-typed and lives in
mcp_server/utils.pybeside the existingextract_mcp_tool_result_error_message, coveringCallToolResultmodels, dicts, and non-text blocks; images and embedded resources report no text and are never handed to the guardrail.The guardrail runs on
execute_mcp_tool's return path rather than inside the logging helper, so enforcement never depends on logging being configured and every dispatch route gets it: the MCP protocol handler, the REST endpoint, and tool search all funnel through that one function. Both tool-call paths forward the rewritten result, and the REST path re-raises guardrail rejections. That second part matters: the existing "logging failed (continuing)" swallow would otherwise have converted a block into a silent leak on that route.One thing deliberately not done:
async_post_mcp_tool_call_hook's return value is still discarded, and the docstring atcustom_logger.py:519now says so instead of advertising a contract that does not hold. Honoring it naively would be actively wrong, because the dispatcher unwraps toresponse.mcp_tool_call_response, and cisco's blocking object carries a bare[TextContent(...)]list:Assigning that would break
.contentconsumers and, worse, make a cisco block log as a success, since error extraction reportsNonefor the list. Cisco already takes effect by mutating in place.structuredContentis scanned and masked too, not justcontent. It is serialized to the client alongside the text blocks, so a sensitive value living only there would otherwise reach the client having never been shown to the guardrail. Both sources go into oneapply_guardrailcall and the returned texts are split back to their origins, so a single pass covers the whole client-visible result. The JSON walk lives inmcp_server/utils.pybeside the content helpers so it is shared rather than private to this handler, and it fails closed at a depth cap instead of truncating, so nothing passes unscanned. Dictionary keys and non-string scalars are scanned too, but a match on one blocks with a 400 rather than being rewritten, because renaming a key or editing a number changes the payload contract instead of redacting a value;litellm_content_filtertakes the same position on MCP tool call arguments. One follow-up remains: the new mode needs a docs PR in the docs repo. Separately, presidio'sinput_type="response"path takes an unmask branch whenoutput_parse_piipopulatedpii_tokens, so it will not scan a result for newly-leaked values in that configuration; that is pre-existing presidio behavior shared with the LLMpost_callpath, not introduced here, but it is worth knowing when configuring this mode.Twenty-one new behavior tests, each failing on the pre-fix tree, across the handler, the dispatcher, and both call sites; plus targeted mutations confirming the
should_run_guardrailgate and theapply_guardrailfilter are each independently load-bearing. 745 tests pass across the proxy-utils, MCP server, REST endpoint, guardrail-translation, cisco MCP, and Responses suites, and 2685 acrosstests/test_litellm/proxy/guardrails/. The 23 failures in a fullmcp_serverdirectory run (test_semantic_tool_filter.py,test_mcp_env_vars.py) reproduce identically on clean staging with the same command, 23 there and 23 here, so they are pre-existing and unrelated.Final Attestation