feat(mcp): output guardrails over MCP tool results#673
Merged
Conversation
Completes guardrails-over-MCP: the tool *result* now runs through the same guardrail chain as the LLM response. The /mcp handler resolves the chain once and runs both directions — input on the arguments (before the call), output on the result (after) — so a guardrail configured for the output hook fires on MCP tool results too. When a guardrail chain is attached the handler buffers the gateway response, extracts the JSON-RPC `result`, feeds it to `check_output` as assistant text, and on a Block replaces the result with an MCP-native JSON-RPC error (recorded as guardrail_blocked). An error response (no `result`) has nothing to scan and passes through. The response body is buffered ONLY when a chain is attached, so the no-guardrail path is unchanged. Tested: `output_guardrail_block` blocks a result whose content carries a guardrail-blocked token, allows a clean result, and skips an error response with no result. Closes #671. Refs AISIX-Cloud#894
📝 WalkthroughWalkthroughThe MCP Input + Output Guardrail Pipeline
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
crates/aisix-proxy/src/mcp.rs (1)
711-756: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd an endpoint-level output-block regression.
This only tests the helper. Please also cover
/mcpreturning the synthesized JSON-RPC error and emittingguardrail_blocked=trueafter a blocked tool result, so the buffering/replacement/usage path is locked down.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aisix-proxy/src/mcp.rs` around lines 711 - 756, The current test only covers output_guardrail_block directly, but it does not verify the /mcp endpoint behavior when a tool result is blocked. Add an endpoint-level regression around the MCP handler in mcp.rs that exercises a blocked tool response end-to-end, asserts the synthesized JSON-RPC error is returned, and checks that the response metadata/logging includes guardrail_blocked=true. Reuse the existing output_guardrail_block and /mcp request flow so the buffering, replacement, and usage path are covered together.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/aisix-proxy/src/mcp.rs`:
- Around line 182-186: The early BAD_GATEWAY return in the MCP upstream
buffering path skips the usage event for a tool call that already reached
upstream. Update the error branch in mcp.rs around the to_bytes handling to emit
the same MCP usage event before returning, using the existing tool-call/usage
reporting path in this flow so the failure is still recorded. Keep the fix
localized to the resp_body buffering error handling in the mcp response path.
- Around line 180-188: Gate the buffering logic in mcp.rs on output-capable
guardrails only: the current response path under guardrail_chain always calls
to_bytes with state.request_body_limit_bytes and can fail large valid MCP
responses even when no output hook will run. Update the conditional around the
response handling to check runs_on_output() before entering the
resp_parts/resp_body buffering branch, and keep input-only guardrail chains on
the non-buffering path so output_guardrail_block is only reached when it can
actually inspect the response.
---
Nitpick comments:
In `@crates/aisix-proxy/src/mcp.rs`:
- Around line 711-756: The current test only covers output_guardrail_block
directly, but it does not verify the /mcp endpoint behavior when a tool result
is blocked. Add an endpoint-level regression around the MCP handler in mcp.rs
that exercises a blocked tool response end-to-end, asserts the synthesized
JSON-RPC error is returned, and checks that the response metadata/logging
includes guardrail_blocked=true. Reuse the existing output_guardrail_block and
/mcp request flow so the buffering, replacement, and usage path are covered
together.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3c4065cb-aae5-4bf6-aef0-8368f7ea698c
📒 Files selected for processing (1)
crates/aisix-proxy/src/mcp.rs
Comment on lines
+180
to
+188
| let response = if let Some(chain) = &guardrail_chain { | ||
| let (resp_parts, resp_body) = response.into_parts(); | ||
| let resp_bytes = match to_bytes(resp_body, state.request_body_limit_bytes).await { | ||
| Ok(bytes) => bytes, | ||
| Err(_) => { | ||
| return (StatusCode::BAD_GATEWAY, "invalid upstream response").into_response() | ||
| } | ||
| }; | ||
| if let Some(guardrail_name) = output_guardrail_block(chain, &resp_bytes, &mcp_tool).await { |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the target file and related symbols.
ast-grep outline crates/aisix-proxy/src/mcp.rs --view expanded || true
printf '\n--- search runs_on_output ---\n'
rg -n "runs_on_output|guardrail_chain|output_guardrail_block|request_body_limit_bytes" crates -S
printf '\n--- read relevant slice of mcp.rs ---\n'
sed -n '150,230p' crates/aisix-proxy/src/mcp.rs
printf '\n--- locate Guardrail definition ---\n'
rg -n "trait Guardrail|fn runs_on_output|runs_on_output\\(" crates -SRepository: api7/aisix
Length of output: 14701
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- mcp.rs around guardrail_chain creation ---'
sed -n '100,145p' crates/aisix-proxy/src/mcp.rs
printf '\n%s\n' '--- GuardrailChain API ---'
sed -n '250,330p' crates/aisix-guardrails/src/lib.rs
printf '\n%s\n' '--- GuardrailChain implementation ---'
sed -n '120,180p' crates/aisix-guardrails/src/chain.rs
printf '\n%s\n' '--- relevant mcp tests around input/output guardrails ---'
sed -n '660,770p' crates/aisix-proxy/src/mcp.rsRepository: api7/aisix
Length of output: 11520
Gate response buffering on output-capable guardrails.
guardrail_chain can be input-only, but this branch still buffers every MCP tool response and applies state.request_body_limit_bytes, so a large valid result can fail with 502 invalid upstream response even though no output hook will inspect it. Filter on runs_on_output() before taking the buffering path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/aisix-proxy/src/mcp.rs` around lines 180 - 188, Gate the buffering
logic in mcp.rs on output-capable guardrails only: the current response path
under guardrail_chain always calls to_bytes with state.request_body_limit_bytes
and can fail large valid MCP responses even when no output hook will run. Update
the conditional around the response handling to check runs_on_output() before
entering the resp_parts/resp_body buffering branch, and keep input-only
guardrail chains on the non-buffering path so output_guardrail_block is only
reached when it can actually inspect the response.
Comment on lines
+182
to
+186
| let resp_bytes = match to_bytes(resp_body, state.request_body_limit_bytes).await { | ||
| Ok(bytes) => bytes, | ||
| Err(_) => { | ||
| return (StatusCode::BAD_GATEWAY, "invalid upstream response").into_response() | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Emit usage before returning the buffering error.
This early BAD_GATEWAY return skips the MCP usage event for a tool call that reached upstream.
Proposed fix
let resp_bytes = match to_bytes(resp_body, state.request_body_limit_bytes).await {
Ok(bytes) => bytes,
Err(_) => {
+ emit_tool_call_usage(
+ &state,
+ &auth,
+ &mcp_server,
+ &mcp_tool,
+ StatusCode::BAD_GATEWAY.as_u16(),
+ latency,
+ false,
+ );
return (StatusCode::BAD_GATEWAY, "invalid upstream response").into_response()
}
};📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| let resp_bytes = match to_bytes(resp_body, state.request_body_limit_bytes).await { | |
| Ok(bytes) => bytes, | |
| Err(_) => { | |
| return (StatusCode::BAD_GATEWAY, "invalid upstream response").into_response() | |
| } | |
| let resp_bytes = match to_bytes(resp_body, state.request_body_limit_bytes).await { | |
| Ok(bytes) => bytes, | |
| Err(_) => { | |
| emit_tool_call_usage( | |
| &state, | |
| &auth, | |
| &mcp_server, | |
| &mcp_tool, | |
| StatusCode::BAD_GATEWAY.as_u16(), | |
| latency, | |
| false, | |
| ); | |
| return (StatusCode::BAD_GATEWAY, "invalid upstream response").into_response() | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/aisix-proxy/src/mcp.rs` around lines 182 - 186, The early BAD_GATEWAY
return in the MCP upstream buffering path skips the usage event for a tool call
that already reached upstream. Update the error branch in mcp.rs around the
to_bytes handling to emit the same MCP usage event before returning, using the
existing tool-call/usage reporting path in this flow so the failure is still
recorded. Keep the fix localized to the resp_body buffering error handling in
the mcp response path.
Independent audit of the output-guardrail change surfaced two gaps: - The output-block wire contract (HTTP 200, JSON-RPC error code -32600, echoed request id) was only exercised through the isolated helper — a regression that nulled the id or shifted the code would pass CI. Add a direct envelope test on `jsonrpc_guardrail_block` (id echo + shape) and strengthen the input router test to parse the envelope and assert the echoed id end-to-end (both hooks funnel through the same block seam). - Output guardrails scanned the serialized JSON envelope, not the decoded tool text — diverging from the LLM path: envelope field names (`content`/`type`/`text`) could false-positive and escaped characters could hide blocked content. Scan the `text`-type content blocks instead, falling back to the full serialized result for non-standard shapes so nothing escapes inspection. Add a regression test proving an envelope field name no longer fires while decoded text still does. Also distinguishes the caller-visible wording: an output block now reads "tool result blocked by content policy", an input block "tool call". cargo test -p aisix-proxy: 514 green; fmt + clippy --workspace --all-targets -D warnings clean.
Output-guardrail scanning rested on the `/mcp` gateway's `json_response = true` config: a `tools/call` returns a single `application/json` object, which `output_guardrail_block` parses to reach the tool result. The parse previously used `.ok()?`, which fails *open* — an unparseable body (e.g. if that config ever regressed to SSE framing) returned `None` (allow), slipping an unscanned tool result past the guardrail. Fail closed instead: an unparseable body blocks. A protocol-level error envelope (valid JSON, no `result`) still passes through — it has no tool output to scan. Add a regression test that an SSE-framed body blocks. Surfaced by two independent audits of the output-guardrail change. cargo test -p aisix-proxy --lib mcp 13 green; fmt + clippy clean.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Completes guardrails over MCP: the tool result now runs through the same guardrail chain as an LLM response. The
/mcphandler resolves the chain once and runs it in both directions:tools/callarguments before dispatching (shipped previously, feat(mcp): input guardrails over MCP tool calls #672)A guardrail configured for the
outputhook now fires on MCP tool results, exactly as it does for chat completions.How
When a guardrail chain is attached to the request, the handler:
result, feeds it tocheck_outputas assistant text,Block, replaces the result with an MCP-native JSON-RPC error (code -32600, the guardrail block message) and recordsguardrail_blockedon the usage event,An error response (no
resultmember) has nothing to scan and passes straight through. The response body is buffered only when a chain is attached, so the common no-guardrail path is byte-for-byte unchanged.Why this matters
One API key governs both the LLM call and the MCP tool call through the same DP pipeline — input + output guardrails included. Surveying the field as of 2026-06-30 (primary sources): guardrails applied to MCP tool traffic are either unsupported, marked "coming soon", or shipped only behind an experimental flag elsewhere. Same-key, same-pipeline governance over MCP is the differentiator.
Test
output_guardrail_blockunit-tested three ways: blocks a result whose content carries a blocked token, allows a clean result, and skips an error response that has noresult. Existing input-guardrail and usage-event tests still pass (cargo test -p aisix-proxy: 512 tests green; fmt +clippy --workspace --all-targets -D warningsclean).Closes #671.
Refs AISIX-Cloud#894.
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Independent audit response
A cold independent audit (no shared context) returned FIX-FIRST with one HIGH and one MEDIUM. Both are addressed in code:
code -32600, echoed request id) was only exercised via the isolated helper, so a regression that nulled the id or shifted the code would have passed CI. Fixed: added a direct envelope test onjsonrpc_guardrail_block(id-echo + shape) and strengthened the input router test to parse the envelope and assert the echoed id end-to-end — both hooks funnel through the samerpc_id+ block seam, so this exercises the handler wiring the output path shares.content/type/text) could false-positive and escaped characters could hide blocked content. Fixed: scan thetext-type content blocks, falling back to the full serialized result for non-standard shapes so nothing escapes inspection. Added a regression test proving an envelope field name no longer fires while decoded text still does."tool result blocked by content policy"(was"tool call"); the stale doc on the block builder is updated. Adopted.Residual gap (justified, not blocking): a full router-level output-path test (real MCP upstream returns a result → block / clean re-wrap) needs an in-process MCP upstream stub the proxy harness doesn't have yet. The output decision (
output_guardrail_block) and the wire builder (jsonrpc_guardrail_block) are unit-tested, and the input path proves the handler's id-capture/echo seam end-to-end. Building the stub harness is tracked as #674.Verification:
cargo test -p aisix-proxy514 green;cargo fmt --all -- --checkandcargo clippy --workspace --all-targets -- -D warningsclean.Follow-up hardening (both audits, same finding)
A second independent audit and a separate protocol-shape verification independently flagged that output scanning rested on one load-bearing flag (
json_response = true): if the/mcpresponse ever became SSE-framed,output_guardrail_blockwould fail to parse the body and — under the original.ok()?— fail open, passing an unscanned tool result through. Fixed inb11635d: an unparseable body now fails closed (blocks); a protocol-level error envelope (valid JSON, noresult) still passes through since it has no tool output to scan. Regression test added (SSE-framed body blocks). Both audits' final recommendation was MERGE; this closes their one residual note.