Skip to content

feat(mcp): scan and mask MCP tool results via post_mcp_call guardrails - #35155

Merged
yassin-berriai merged 1 commit into
litellm_internal_stagingfrom
litellm_mcp_guardrail_tool_result
Jul 30, 2026
Merged

feat(mcp): scan and mask MCP tool results via post_mcp_call guardrails#35155
yassin-berriai merged 1 commit into
litellm_internal_stagingfrom
litellm_mcp_guardrail_tool_result

Conversation

@yassin-berriai

@yassin-berriai yassin-berriai commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • guardrails could not see MCP tool results, only requests
  • a tool returning sensitive data bypassed every guardrail
  • no mode existed to scan the outbound direction

How it solves it:

  • new post_mcp_call mode, selectable in the dashboard
  • tool result text flows through the unified apply_guardrail seam
  • a guardrail can mask values in the result or reject it

Relevant issues

Linear ticket

Resolves LIT-4935

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)

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_customer tool deliberately returns sensitive data, so what the client receives is the evidence.

curl -sS -X POST "http://localhost:4946/mcp-rest/tools/call" \
  -H "Authorization: Bearer sk-lit4935" -H "Content-Type: application/json" \
  -d '{"server_id":"<id>","name":"lookup_customer","arguments":{"customer_id":"C-1001"}}'

Before, on staging 440b1bcf65. There is no mode that scans a tool result, so the strongest available configuration is mode: 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:

ValueError: 'post_mcp_call' is not a valid GuardrailEventHooks
ERROR:    Application startup failed. Exiting.

After, at 268e884eed, with mode: 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: BLOCK instead of MASK. 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 structuredContent while its text says nothing sensitive, proves the structured path:

text:              lookup complete
structuredContent: {'record': {'email': '<EMAIL_ADDRESS>', 'phone': '<PHONE_NUMBER>'}, 'balance': 42.0}

For the UI half, post_mcp_call appears 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_call and during_mcp_call. The result went back to the client unscanned: MCPGuardrailTranslationHandler.process_output_response was 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_call event hook and an MCP-shaped dispatcher. ProxyLogging.post_mcp_call_hook gates on should_run_guardrail(post_mcp_call) plus the presence of apply_guardrail, then routes the result through MCPGuardrailTranslationHandler.process_output_response with input_type="response", resolving the handler at runtime via load_guardrail_translation_mappings()[CallTypes.call_mcp_tool] so proxy/utils.py never imports mcp at 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_hook was considered and rejected: its other_callbacks loop (proxy/utils.py:2404) is ungated, so every registered non-guardrail CustomLogger would be handed a CallToolResult where it expects an LLM response. unified_guardrail.py is 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.py beside the existing extract_mcp_tool_result_error_message, covering CallToolResult models, 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 at custom_logger.py:519 now says so instead of advertising a contract that does not hold. Honoring it naively would be actively wrong, because the dispatcher unwraps to response.mcp_tool_call_response, and cisco's blocking object carries a bare [TextContent(...)] list:

type the dispatcher would hand back  : list
has .content attribute               : False
extract_mcp_tool_result_error_message(unwrapped)   : None
extract_mcp_tool_result_error_message(CallToolResult): x

Assigning that would break .content consumers and, worse, make a cisco block log as a success, since error extraction reports None for the list. Cisco already takes effect by mutating in place.

structuredContent is scanned and masked too, not just content. 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 one apply_guardrail call and the returned texts are split back to their origins, so a single pass covers the whole client-visible result. The JSON walk lives in mcp_server/utils.py beside 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_filter takes 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's input_type="response" path takes an unmask branch when output_parse_pii populated pii_tokens, so it will not scan a result for newly-leaked values in that configuration; that is pre-existing presidio behavior shared with the LLM post_call path, 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_guardrail gate and the apply_guardrail filter 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 across tests/test_litellm/proxy/guardrails/. The 23 failures in a full mcp_server directory 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

  • 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

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

Comment thread litellm/proxy/_experimental/mcp_server/server.py Outdated
Comment thread litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py Outdated
@greptile-apps

greptile-apps Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Adds outbound MCP result guardrails and addresses the previously reported gaps:

  • Runs post_mcp_call enforcement before MCP results reach logging callbacks, including calls without a logging object
  • Scans and rewrites text blocks, structured string values, dictionary keys, and numeric values with fail-closed handling
  • Propagates rewritten results and guardrail rejections through protocol, REST, tool-search, Responses API, and chat-completions paths
  • Adds the new guardrail mode to the dashboard and expands regression coverage

Confidence Score: 5/5

The PR appears safe to merge

No blocking failure remains

Important Files Changed

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

codecov Bot commented Jul 29, 2026

Copy link
Copy Markdown

@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_result branch from eab8e79 to 86a749f Compare July 29, 2026 23:43
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

Real bypass, and fixing it properly closed a second gap I had listed as a follow-up. Fixed in 86a749f347.

The guardrail was called from inside _fire_mcp_tool_call_logging, which both call paths reach only behind a logging check (if litellm_logging_obj: in call_mcp_tool, and an early if logging_obj is None: return result on the REST path). So enforcement depended on logging being configured, and litellm_logging_obj comes straight off kwargs.get("litellm_logging_obj", None).

Rather than widen that condition, the guardrail moved out of the logging function entirely and onto execute_mcp_tool's single return path, unconditional on logging. That is the choke point all three dispatch routes funnel through, so it also closes the tool_search.py:147 gap the PR body had listed as a known follow-up: tool search calls execute_mcp_tool directly and never reached _fire_mcp_tool_call_logging at all, meaning tool-search-dispatched calls were previously unguarded even when logging was configured. Both holes are one fix now, and the guardrail still runs before the success/failure logging so the log records the masked text.

The two tests that asserted the old location were retargeted to the new one, and a third pins the bypass directly: a call with litellm_logging_obj=None must still be masked. Re-gating on a logging object fails exactly that test and leaves the other two green:

mutant (re-add `or litellm_logging_obj is None`) -> 1 failed, 2 passed
fixed                                            -> 3 passed

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: tool_search is no longer a gap, structuredContent remains one.

@greptileai please review the current head 86a749f347

Comment thread litellm/responses/mcp/litellm_proxy_mcp_handler.py Outdated
@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_result branch from 86a749f to 226020c Compare July 29, 2026 23:58
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

Both findings addressed in 226020c966.

Responses-path ordering. You were right, and it was two problems on that path, not one. litellm_logging_obj.post_call(...) and async_post_mcp_tool_call_hook ran before the guardrail, so logging callbacks saw the unmasked result; and the guardrail was nested inside if litellm_logging_obj:, which is the same bypass shape as the one just fixed on the shared path. The guardrail now runs immediately after call_tool returns, unconditional on logging, with logging after it. Callbacks and the spend log see the masked result.

structuredContent. I had documented this as a follow-up; you are right that it should not merge that way. structuredContent is serialized to the client alongside content:

structuredContent is part of the wire result: {'email': 'jane@example.com', 'nested': {'phone': '415-555-0132'}}

so a masked content beside an unmasked structuredContent leaks exactly what was redacted, behind a response that looks protected — the same failure mode as dropping a mask outright.

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 structuredContent. Different fields compose normally; only a genuine repeat of a redacted value blocks.

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 "email jane@example.com" is not a substring of {"email": "jane@example.com"}. Both shapes are pinned:

mutant: drop the guard entirely            -> 1 failed, 11 passed
mutant: compare whole texts, not spans     -> 1 failed, 11 passed
fixed                                      -> 12 passed

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: misc / Run tests fails at test_openapi_compliance.py::TestRequestCompliance::test_content_schema_uses_discriminator, which reproduces identically on clean staging 440b1bcf65 locally, so it is pre-existing and not from this PR.

@greptileai please review the current head 226020c966

@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_result branch from 226020c to f85765d Compare July 30, 2026 00:09
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

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 content, so a value living only in structuredContent was never handed to the guardrail at all. structuredContent is now scanned and masked properly in f85765d40e.

Both content text blocks and every string leaf of structuredContent go into a single apply_guardrail call, and the returned texts are split back to their sources, so one guardrail pass covers the whole client-visible result. The JSON walk lives in mcp_server/utils.py next to the existing content helpers rather than in the handler, so it is shared rather than a second private copy.

Live proof against a tool whose PII exists only in structuredContent (text says nothing sensitive):

text:              lookup complete
structuredContent: {'record': {'email': '<EMAIL_ADDRESS>', 'phone': '<PHONE_NUMBER>'}, 'balance': 42.0}

Masked in place, shape preserved, and the float left alone. Four tests cover it, including the exact gap you named (test_value_present_only_in_structured_content_is_masked) and depth-cap behavior; both plausible regressions die:

mutant: scan content only            -> 3 failed, 10 passed
mutant: never write structured back   -> 2 failed, 11 passed
fixed                                -> 13 passed

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 ValueError subclass rather than an HTTPException, because mcp_server/utils.py deliberately imports neither fastapi nor the litellm package; the handler converts it to a 400.

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.

misc / Run tests remains red at test_openapi_compliance.py::TestRequestCompliance::test_content_schema_uses_discriminator, which reproduces identically on clean staging 440b1bcf65 locally and is unrelated to this PR.

@greptileai please review the current head f85765d40e

Comment thread litellm/proxy/_experimental/mcp_server/utils.py Outdated
@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_result branch from f85765d to 01ebb91 Compare July 30, 2026 00:42
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

CI caught a real bug in my own change, now fixed in 01ebb919e1. Worth writing up because the failure mode is not obvious.

proxy-infra failed on my new depth test with my own JSONTooDeeplyNestedError escaping the except that exists to catch it. It passed locally and failed under xdist. Cause: mcp_server/utils.py is deliberately reloaded by tests that override its environment-backed constants (there is a note to that effect above those constants, and test_mcp_server_identity_env.py does it). A reload gives every class defined there a fresh identity, so the handler's except JSONTooDeeplyNestedError no longer matches the class the reloaded module raises. Reproduced deterministically by running that reload test first:

pytest .../test_mcp_server_identity_env.py .../test_mcp_guardrail_handler.py -p no:randomly
-> 1 failed, 15 passed

Defining a custom exception in a module that is reloaded by design was the wrong shape. The walk now returns None for too-deep instead of raising, and the handler converts that to the 400. A sentinel has no identity to lose, so the whole class of bug is gone rather than patched.

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 importlib.reload itself, and that leaked: it made two test_mcp_env_vars.py tests fail. I was about to attribute those to the pre-existing baseline, because staging does fail them too. Deselecting only my test proved otherwise:

my dir + env_vars, with my reload test     -> 2 failed, 75 passed
my dir + env_vars, that test deselected    -> 76 passed

So the test is now a direct contract assertion on the helper (json_string_leaves(deep) is None) with no global mutation. It kills the same mutant:

mutant: raise instead of returning the sentinel -> 2 failed, 12 passed
fixed                                            -> 14 passed

The residual two test_mcp_env_vars.py failures are staging's own: test_mcp_server_identity_env.py reloads the module and they fail identically on clean staging 440b1bcf65 with the same two files and nothing of mine involved:

identity_env + env_vars on clean staging -> 2 failed, 64 passed

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 01ebb919e1

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
@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_result branch from 01ebb91 to 268e884 Compare July 30, 2026 01:10
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

Fixed in 268e884eed, following the precedent already in the codebase rather than inventing a rule.

Dictionary keys and non-string scalars in structuredContent were used purely as reconstruction paths and passed through verbatim, so a map keyed by an identifier ({"jane@example.com": {...}}, a common API shape) reached the client with the key never shown to the guardrail.

litellm_content_filter already answers this question for MCP tool call arguments: it scans dict keys and stringified numerics and blocks with "matched a masking rule on a non-rewritable field", because renaming a key or rewriting a number changes the payload's contract rather than redacting a value (content_filter.py:1695-1720). The result path now does the same. Keys and numeric leaves go into the same guardrail call as the text, and a match on any of them raises a 400 naming the reason instead of being silently preserved or silently rewritten.

Three tests cover it, and the mutant that stops collecting them dies:

mutant: structured_labels = ()  -> 2 failed, 15 passed
fixed                           -> 17 passed

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 json_string_leaves (returns None past the cap so the caller blocks rather than skipping deeper values) and carries its recursion-detector entry with that reason.

3138 tests pass across the MCP server directory, proxy-utils, cisco MCP, and Responses suites, with only the two test_mcp_env_vars.py tests deselected that clean staging also fails from its own module reload. Both lint budget gates within ceiling. Structured masking re-verified live and unchanged.

@greptileai please review the current head 268e884eed

@codspeed-hq

codspeed-hq Bot commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_mcp_guardrail_tool_result (268e884) with litellm_internal_staging (551e5d0)1

Open in CodSpeed

Footnotes

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

@yassin-berriai
yassin-berriai enabled auto-merge (squash) July 30, 2026 21:08
@yassin-berriai
yassin-berriai merged commit 5c16132 into litellm_internal_staging Jul 30, 2026
80 checks passed
@yassin-berriai
yassin-berriai deleted the litellm_mcp_guardrail_tool_result branch July 30, 2026 21:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants