Skip to content

fix(mcp): scan and mask MCP tool call arguments in unified guardrails - #35142

Open
yassin-berriai wants to merge 1 commit into
litellm_internal_stagingfrom
litellm_mcp_guardrail_tool_args
Open

fix(mcp): scan and mask MCP tool call arguments in unified guardrails#35142
yassin-berriai wants to merge 1 commit into
litellm_internal_stagingfrom
litellm_mcp_guardrail_tool_args

Conversation

@yassin-berriai

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

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • mode: pre_mcp_call guardrails never saw MCP tool arguments
  • presidio et al. could not detect or mask MCP payloads
  • proxy still reported the guardrail as applied

How it solves it:

  • pass argument string leaves to the guardrail as texts
  • fold rewritten leaves back into modified_arguments
  • nested dicts, lists, and non-string values keep their shape

Relevant issues

Linear ticket

Resolves LIT-4944

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 on port 4944 with a real presidio analyzer + anonymizer pair (no mocks, no stubs) and a stdio MCP server whose single tool echoes back verbatim the arguments it received, so the tool result is the evidence of what actually left the gateway.

Guardrail config used for both runs:

guardrails:
  - guardrail_name: "presidio-mcp"
    litellm_params:
      guardrail: presidio
      mode: "pre_mcp_call"
      default_on: true
      presidio_analyzer_api_base: "http://localhost:5943"
      presidio_anonymizer_api_base: "http://localhost:5942"
      pii_entities_config:
        EMAIL_ADDRESS: "MASK"
        PHONE_NUMBER: "MASK"

The call, identical in both runs:

curl -sS -X POST "http://localhost:4944/mcp-rest/tools/call" \
  -H "Authorization: Bearer sk-lit4944" \
  -H "Content-Type: application/json" \
  -d '{
        "server_id": "3b140a205f437a806cc9a46ecce10b67",
        "name": "echo_payload",
        "arguments": {
          "note": "please email jane.doe@example.com and call 415-555-0132",
          "envelope": {
            "reply_to": "ops.lead@example.net",
            "cc": ["auditor@example.org", "no-pii-here"],
            "priority": 2,
            "urgent": true,
            "trace": null
          }
        }
      }'

Before, at 440b1bcf65 (staging, unfixed)

Every value reached the MCP server in the clear:

{"received_by_mcp_server": {
  "envelope": {"cc": ["auditor@example.org", "no-pii-here"], "priority": 2, "reply_to": "ops.lead@example.net", "trace": null, "urgent": true},
  "note": "please email jane.doe@example.com and call 415-555-0132"
}}

and the proxy nonetheless recorded the guardrail as having run, which is what makes this quiet rather than obvious:

'applied_guardrails': ['presidio-mcp', 'presidio-mcp']

After, at f155da12f9

{"received_by_mcp_server": {
  "envelope": {"cc": ["<EMAIL_ADDRESS>", "no-pii-here"], "priority": 2, "reply_to": "<EMAIL_ADDRESS>", "trace": null, "urgent": true},
  "note": "please email <EMAIL_ADDRESS> and call <PHONE_NUMBER>"
}}

Both emails inside the nested object are masked, including the one inside a list; the non-PII string no-pii-here is untouched; and priority: 2, urgent: true, trace: null come through with their original types rather than being stringified.

Type

🐛 Bug Fix

Changes

MCPGuardrailTranslationHandler.process_input_messages built a GenericGuardrailAPIInputs holding only a synthetic tool definition (the tool name plus an empty parameters schema), passed it to apply_guardrail, and returned data untouched. That is the seam every guardrail implementing apply_guardrail goes through for MCP, which is presidio, model_armor, noma, pillar, and bedrock among others. The argument values were never handed over, so detection could not fire; and the one channel the MCP call path reads back, data["modified_arguments"] (consumed by ProxyLogging._convert_mcp_hook_response_to_kwargs), was never written, so a mask could not take effect either.

litellm_content_filter and cisco_ai_defense were unaffected throughout because each reads mcp_arguments and writes modified_arguments itself rather than relying on the shared seam; that bespoke handling is exactly what the other guardrails lacked.

The handler now walks the argument tree for its string leaves in a deterministic depth-first order, hands them over as texts alongside the existing tool definition, and pairs the guardrail's returned texts back to the leaves they came from, rebuilding the arguments with only the leaves the guardrail actually rewrote. Three properties are deliberate. A guardrail that changes nothing writes no modified_arguments, so a clean call goes upstream byte-identical. A guardrail that returns the wrong number of texts is refused rather than trusted, since the pairing is positional and a mismatch could otherwise scramble fields. And arguments nested deeper than DEFAULT_MAX_RECURSE_DEPTH are blocked instead of quietly passing unscanned, matching what litellm_content_filter already does on its own MCP path.

A guardrail declared run_in_parallel shares one payload snapshot and its return value is discarded by design, so the masked arguments are written onto that payload rather than onto a copy of it; returning a copy meant the mask was silently dropped in that mode, which Greptile caught. The end-to-end test is parametrized over both settings and the parallel case fails against the copy-returning version.

Two recursive tree walkers are added to tests/code_coverage_tests/recursive_detector.py's ignore list, alongside the three existing entries that do the same job (content_filter._filter_mcp_argument_value, model_armor._redact_scanned_content, tool_permission._collect_argument_paths). _collect_argument_texts carries the depth cap and fails closed by blocking at it; _replace_argument_texts is transitively bounded, since it only ever runs on a tree the collector already walked under that cap.

Concurrent rewrites are handled explicitly rather than left to whichever guardrail finishes last. Two guardrails opted into run_in_parallel scan the same payload snapshot and each returns a full replacement string derived from the original leaf, so rewrites of different leaves compose, while rewrites of the same leaf cannot be merged at all: applying either result discards the other redaction. Each guardrail therefore compares the leaf as it currently stands against the text it was handed, and an unmergeable collision blocks with a 400 naming the argument instead of silently shipping one redaction and leaking the other. Sequential guardrails see each other output and compose normally.

Worth flagging for reviewers: guardrails previously handed nothing now receive real content, so a deployment running one of them on pre_mcp_call will start detecting, masking, or blocking payloads it silently let through before. That is the point of the fix, but it is a behavior change on upgrade rather than a pure no-op.

Fourteen tests in the mapped file, all of which fail on the pre-fix handler or on a targeted mutation of the fix: the guardrail receiving argument values at all, the mask landing in modified_arguments, nested and list shapes surviving the rewrite, the clean-call and length-mismatch guards, the depth block, and one that drives the whole real path (_convert_mcp_to_llm_format -> pre_call_hook -> _convert_mcp_hook_response_to_kwargs) rather than only the handler in isolation. 452 tests across the content_filter, presidio, unified-guardrail, cisco MCP, and MCP-bridging suites pass unchanged.

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/guardrail_translation/handler.py Outdated
@greptile-apps

greptile-apps Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR makes unified pre-call guardrails inspect and rewrite MCP argument string leaves while preserving nested argument shapes

  • Adds deterministic recursive argument collection and replacement with depth, result-length, and concurrent-rewrite safeguards
  • Mutates the shared MCP payload so parallel guardrail rewrites reach outbound tool calls
  • Adds regression coverage for nested arguments, masking, recursion limits, and sequential and parallel guardrail execution

Confidence Score: 5/5

The PR appears safe to merge

No blocking failure remains; the shared-payload mutation fixes discarded parallel masks, and atomic stale-check/write-back handling prevents parallel rewrites from silently overwriting each other

Important Files Changed

Filename Overview
litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py Collects MCP argument text for unified guardrails, writes masked values back into the outbound payload, and safely handles concurrent rewrites
tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py Adds focused unit and call-path coverage for argument scanning, shape preservation, masking, depth limits, and parallel composition
tests/code_coverage_tests/recursive_detector.py Records the two bounded recursive MCP argument walkers in the recursive-function allowlist

Reviews (4): Last reviewed commit: "fix(mcp): scan and mask MCP tool call ar..." | Re-trigger Greptile

@codecov

codecov Bot commented Jul 29, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.05882% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...mental/mcp_server/guardrail_translation/handler.py 97.05% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_args branch from f55dcab to 9bfbaf8 Compare July 29, 2026 21:32
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

Good catch on the parallel case; you were right, and for a sharper reason than I first assumed.

I went to check whether this was just the documented run_in_parallel constraint ("Only safe for block-only guardrails that do not mutate the request or response", custom_guardrail.py:153-155) before replying. It is not. Driving a masking guardrail through the real pre_call_hook on a plain chat completion shows the mask surviving in parallel mode, because _run_parallel_pre_call_guardrails discards the returned dict but shares the payload object, so an in-place mutator still lands. My handler returned a new dict, which is exactly what that path throws away:

MCP TOOL CALL  run_in_parallel=False -> args sent upstream: {'query': 'mail <EMAIL>'}
MCP TOOL CALL  run_in_parallel=True  -> args sent upstream: {'query': 'mail jane.doe@example.com'}

Fixed by writing the masked arguments onto the caller payload instead of a copy of it, and the end-to-end test is now parametrized over run_in_parallel both ways. Reverting just that line to the copy-returning shape fails the [True] case and passes [False], so the test pins the behavior rather than merely covering it.

Also in this push: the two recursive argument walkers are added to the ignore list in recursive_detector.py, which is what the code-quality failure was. They sit next to the three existing entries doing the same job (content_filter._filter_mcp_argument_value, model_armor._redact_scanned_content, tool_permission._collect_argument_paths); _collect_argument_texts carries the depth cap and fails closed at it, and _replace_argument_texts is transitively bounded because it only runs on a tree the collector already walked under that cap.

Live re-verified on the same proxy plus real presidio rig after the change; masking output is unchanged.

@greptileai please review the current head 9bfbaf8803

strict=fn.get("strict", False) or False, # Default to False if None
),
}
argument_texts = _collect_argument_texts(mcp_arguments)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Low: Unbounded guardrail work from MCP arguments

mcp_arguments is request-controlled, and this collects every string leaf without a cardinality or aggregate-size limit. Guardrails such as Presidio process each text separately and perform analyzer/anonymizer requests, so an authenticated user can send a shallow array with thousands of strings and tie up proxy and guardrail capacity. Enforce a fail-closed maximum leaf count and total text size during collection, or scan the arguments in bounded batches.

@veria-ai

veria-ai Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

PR overview

This PR extends unified guardrails to scan and mask string values within MCP tool-call arguments.

The argument traversal currently has no limit on the number or aggregate size of string values processed. An authenticated user could submit thousands of values to consume proxy and guardrail capacity, so bounded collection or batching is still needed.

Open issues (1)

Fixed/addressed: 0 · PR risk: 5/10

@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_args branch from 9bfbaf8 to cffa067 Compare July 29, 2026 22:01
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

proxy-infra was a real bug in my test, not a flake, and it is fixed in cffa067654.

The failure was TestUnifiedGuardrailCallTypeResolution::test_pass_through_call_type_resolved_from_logging_obj asserting process_output_response was awaited once and getting zero, in a file my diff does not touch. It reproduces deterministically by running my file first in the same process:

pytest .../guardrail_translation/test_mcp_guardrail_handler.py \
       .../pass_through_endpoints/test_passthrough_post_call_guardrails.py::TestUnifiedGuardrailCallTypeResolution
-> 1 failed, 11 passed

Cause: unified_guardrail memoizes its translation mappings in a module-level global, initialized lazily on first use. That test patches load_guardrail_translation_mappings, which only takes effect while the global is still None. My new end-to-end test drives the real pre_call_hook, so it was the first thing in the run to populate it:

global BEFORE any MCP guardrail run: None
global AFTER  MCP guardrail run    : POPULATED (24 entries)

so the patch was ignored, the real mappings had no entry for the mocked call type, and the hook returned early.

My test leaked process-wide state, so the fix is in my fixture: it now saves and restores endpoint_guardrail_translation_mappings alongside litellm.callbacks, and clears ProxyLogging._callback_capabilities_cache (that cache is keyed on id()s of the callback list, so a restored-but-different list can otherwise collide after GC). Passes in both orders and under -n 2.

For the record on the rest of that job: the same 20 Vertex pass-through credential tests fail identically on clean staging 440b1bcf65 locally (670 passed there vs 678 here, the delta being my new tests), so those are environment-dependent and pre-existing rather than anything this PR introduced. codecov/patch is the only other non-green signal and it settles late.

@greptileai please review the current head cffa067654

Comment thread litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py Outdated
@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_args branch from cffa067 to f155da1 Compare July 29, 2026 22:18
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

Right again, and this one does not have a merge-based fix. Fixed fail-closed in f155da12f9.

I reproduced it before deciding what to do: two run_in_parallel guardrails with staggered latency, one masking an email and one a phone number, both rewriting the same argument:

TWO PARALLEL MASKERS -> mail jane.doe@example.com or call <PHONE>
email masked: False | phone masked: True

My first attempt was to compose by writing onto the current payload instead of the pre-await snapshot. That fixes it only when the guardrails touch different leaves. On the same leaf it cannot work in principle: each guardrail returns a whole replacement string derived from the original text, so applying either result discards the other's redaction. Swapping which one finishes first just swaps which redaction is lost:

a-first  -> mail jane.doe@example.com or call <PHONE>
b-first  -> mail <EMAIL> or call 415-555-0132

So the write now composes where composition is well defined, and refuses where it is not. Before writing, each guardrail compares the leaf as it currently stands against the text it was actually handed; if another guardrail already rewrote that leaf, this guardrail's result is stale and the request is blocked with a 400 naming the argument and telling the operator to drop run_in_parallel from one of them. Silently shipping one redaction is the worse outcome here, since dropping a mask leaks exactly the value the guardrail was configured to protect, which is the bug class this PR exists to fix.

Behavior across the three configurations:

parallel, SAME leaf         -> BLOCKED (400, names the argument)
parallel, DIFFERENT leaves  -> {'email': '<EMAIL>', 'phone': '<PHONE>'}   both masks survive
sequential, SAME leaf       -> 'mail <EMAIL> or call <PHONE>'             both masks survive

Three tests cover those, and removing the conflict check fails the blocking one while the other two stay green, so it is pinned rather than merely covered. Sequential remains the correct configuration for stacked rewriting guardrails and now demonstrably composes both redactions.

Live re-verified on the proxy plus real presidio rig; single-guardrail masking output is unchanged.

@greptileai please review the current head f155da12f9

@codspeed-hq

codspeed-hq Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_mcp_guardrail_tool_args (26b6606) with litellm_internal_staging (0e9a624)1

Open in CodSpeed

Footnotes

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

A guardrail configured with mode pre_mcp_call was handed only a synthetic
tool definition (name plus an empty parameters schema), so it never saw the
argument values it was configured to inspect, and any rewrite it returned was
discarded. Detection could not fire and masking could not take effect, while
the applied-guardrails metadata still reported the guardrail as having run.

Pass every string leaf of the tool call arguments as texts, and fold the
guardrail's rewritten leaves back into modified_arguments, which is the channel
the MCP call path reads to decide what to send upstream. The leaf walk reuses
the json_string_leaves / with_json_string_leaves helpers the tool result path
already uses, so both directions share one bounded traversal.

Two guardrails running concurrently under run_in_parallel scan the same payload
snapshot, so each returns a full replacement derived from the original leaf.
Rewrites of the same leaf to different values are rejected rather than silently
losing one redaction; a leaf that already holds this guardrail's own replacement
is convergent and still masks, which is what the bundled content filter does
when it rewrites the arguments itself as well as through texts.
@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_guardrail_tool_args branch from f155da1 to 26b6606 Compare July 31, 2026 18:24
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.

2 participants