fix(mcp): scan and mask MCP tool call arguments in unified guardrails - #35142
fix(mcp): scan and mask MCP tool call arguments in unified guardrails#35142yassin-berriai wants to merge 1 commit into
Conversation
|
|
Greptile SummaryThis PR makes unified pre-call guardrails inspect and rewrite MCP argument string leaves while preserving nested argument shapes
Confidence Score: 5/5The 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
|
| 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 Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
f55dcab to
9bfbaf8
Compare
|
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 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 Also in this push: the two recursive argument walkers are added to the ignore list in Live re-verified on the same proxy plus real presidio rig after the change; masking output is unchanged. @greptileai please review the current head |
| strict=fn.get("strict", False) or False, # Default to False if None | ||
| ), | ||
| } | ||
| argument_texts = _collect_argument_texts(mcp_arguments) |
There was a problem hiding this comment.
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.
PR overviewThis 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 |
9bfbaf8 to
cffa067
Compare
|
The failure was Cause: 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 For the record on the rest of that job: the same 20 Vertex pass-through credential tests fail identically on clean staging @greptileai please review the current head |
cffa067 to
f155da1
Compare
|
Right again, and this one does not have a merge-based fix. Fixed fail-closed in I reproduced it before deciding what to do: two 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: 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 Behavior across the three configurations: 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 |
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.
f155da1 to
26b6606
Compare
TLDR
Problem this solves:
mode: pre_mcp_callguardrails never saw MCP tool argumentsHow it solves it:
textsmodified_argumentsRelevant issues
Linear ticket
Resolves LIT-4944
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 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:
The call, identical in both runs:
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:
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-hereis untouched; andpriority: 2,urgent: true,trace: nullcome through with their original types rather than being stringified.Type
🐛 Bug Fix
Changes
MCPGuardrailTranslationHandler.process_input_messagesbuilt aGenericGuardrailAPIInputsholding only a synthetic tool definition (the tool name plus an emptyparametersschema), passed it toapply_guardrail, and returneddatauntouched. That is the seam every guardrail implementingapply_guardrailgoes 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 byProxyLogging._convert_mcp_hook_response_to_kwargs), was never written, so a mask could not take effect either.litellm_content_filterandcisco_ai_defensewere unaffected throughout because each readsmcp_argumentsand writesmodified_argumentsitself 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
textsalongside 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 nomodified_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 thanDEFAULT_MAX_RECURSE_DEPTHare blocked instead of quietly passing unscanned, matching whatlitellm_content_filteralready does on its own MCP path.A guardrail declared
run_in_parallelshares 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_textscarries the depth cap and fails closed by blocking at it;_replace_argument_textsis 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_parallelscan 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_callwill 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