From 6f217d21c41ad759e3800757706bd8ef68f9b48e Mon Sep 17 00:00:00 2001 From: Jarvis Date: Sun, 16 Aug 2026 09:58:41 +0800 Subject: [PATCH 1/2] feat(mcp): scan structured tool output, block as a tool error, scope guardrails per server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three gaps in how guardrails govern MCP tool calls. **`structuredContent` was never scanned.** The output hook fed the guardrail chain the decoded `text` content blocks and only fell back to the serialized result when there were none. A tool result carries `structuredContent` alongside `content`, the gateway relays it to the client verbatim, and the spec only RECOMMENDS mirroring it into a text block — so a tool returning clean prose plus a sensitive structured payload passed inspection. The hook now also walks that field's string leaves. Object keys stay out of the scan: they are the tool's declared output schema rather than its data, the same reason the content blocks are decoded instead of scanned as raw JSON. **A block surfaced as a JSON-RPC protocol error.** MCP separates "this request was not valid" (a protocol error, which clients treat as a transport-level failure) from "the tool call did not succeed" (`isError` on the result, which the calling agent reads as tool output). A policy rejection is the second kind: the request was well-formed, and the caller should learn in-band that content policy stopped it. Both hooks now answer HTTP 200 with `result.content[0].text` + `isError: true` instead of `error.code: -32600`. **Guardrails could not be scoped to an MCP server.** An MCP tool call resolves no model, so the only attachment scopes that reached it were env, api_key and team — an operator could not guard one registered server without guarding all MCP traffic. `guardrail_attachments` accepts a new `scope_type: mcp_server` whose `scope_id` is the registered server's id; `RequestContext` carries the called server so the index selects on it. Model and MCP-server scopes are mutually exclusive per request and share a specificity tier. An attachment whose `scope_id` is empty no longer matches a request that lacks the dimension, so the "absent" sentinel cannot read as a wildcard. Behavior change: an MCP client that matched on `error.code == -32600` to detect a policy rejection must read `result.isError` instead. --- crates/aisix-core/src/models/guardrail.rs | 18 +- crates/aisix-core/src/models/snapshot.rs | 2 +- crates/aisix-guardrails/src/build.rs | 64 ++++ crates/aisix-guardrails/src/index.rs | 103 ++++++- crates/aisix-proxy/src/audio.rs | 2 + crates/aisix-proxy/src/chat.rs | 1 + crates/aisix-proxy/src/completions.rs | 1 + crates/aisix-proxy/src/embeddings.rs | 1 + crates/aisix-proxy/src/error.rs | 8 +- crates/aisix-proxy/src/images.rs | 1 + crates/aisix-proxy/src/jobs.rs | 2 + crates/aisix-proxy/src/mcp.rs | 283 ++++++++++++++++-- crates/aisix-proxy/src/messages.rs | 1 + crates/aisix-proxy/src/passthrough.rs | 1 + crates/aisix-proxy/src/realtime.rs | 1 + crates/aisix-proxy/src/rerank.rs | 1 + crates/aisix-proxy/src/responses.rs | 1 + crates/aisix-proxy/src/videos.rs | 1 + .../guardrail_attachment.schema.json | 7 +- tests/e2e/src/cases/mcp-guardrail-e2e.test.ts | 221 ++++++++++++++ tests/e2e/src/harness/index.ts | 6 +- tests/e2e/src/harness/upstream-mcp.ts | 39 ++- 22 files changed, 711 insertions(+), 54 deletions(-) create mode 100644 tests/e2e/src/cases/mcp-guardrail-e2e.test.ts diff --git a/crates/aisix-core/src/models/guardrail.rs b/crates/aisix-core/src/models/guardrail.rs index 3a2ca612..2489d4d8 100644 --- a/crates/aisix-core/src/models/guardrail.rs +++ b/crates/aisix-core/src/models/guardrail.rs @@ -962,20 +962,26 @@ impl Resource for Guardrail { /// Which dimension of the request a guardrail attachment is scoped to. /// /// `Env` applies to every request in the environment. The narrower scopes let -/// operators attach a guardrail to only the models, API keys, or teams that -/// need it. +/// operators attach a guardrail to only the models, MCP servers, API keys, or +/// teams that need it. +/// +/// `Model` and `McpServer` select dimensions a request carries only one of: an +/// MCP tool call resolves no model, and an LLM request routes to no MCP server. +/// A `Model`-scoped guardrail therefore never inspects MCP traffic, and an +/// `McpServer`-scoped one never inspects model traffic. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] #[serde(rename_all = "snake_case")] pub enum GuardrailScopeType { Env, Model, + McpServer, ApiKey, Team, } /// Guardrail attachment that scopes one guardrail to an environment, model, -/// caller API key, or team. AISIX loads attachments with the guardrail -/// definitions and uses `scope_type` plus `scope_id` to decide which +/// MCP server, caller API key, or team. AISIX loads attachments with the +/// guardrail definitions and uses `scope_type` plus `scope_id` to decide which /// guardrails apply to each request. #[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema, PartialEq, Eq)] pub struct GuardrailAttachment { @@ -986,8 +992,8 @@ pub struct GuardrailAttachment { /// What dimension of the request this attachment is scoped to. pub scope_type: GuardrailScopeType, - /// The UUID of the specific resource (model / api_key / team). - /// `None` when `scope_type` is `Env` (applies to all requests). + /// The UUID of the specific resource (model / mcp_server / api_key / + /// team). `None` when `scope_type` is `Env` (applies to all requests). pub scope_id: Option, /// Higher number = higher precedence. When the same guardrail appears diff --git a/crates/aisix-core/src/models/snapshot.rs b/crates/aisix-core/src/models/snapshot.rs index 8db2bce1..839d25fa 100644 --- a/crates/aisix-core/src/models/snapshot.rs +++ b/crates/aisix-core/src/models/snapshot.rs @@ -28,7 +28,7 @@ pub struct AisixSnapshot { pub guardrails: ResourceTable, /// Attachment rows: `/aisix//guardrail_attachments/`. /// Each row binds a guardrail definition to a scope (env / model / - /// api_key / team). `GuardrailIndex::build_from_snapshot` consumes + /// mcp_server / api_key / team). `GuardrailIndex::build_from_snapshot` consumes /// both this table and `guardrails` to build the per-request resolver. pub guardrail_attachments: ResourceTable, /// Per-env cache policies. Stage 2 honors only the existence of an diff --git a/crates/aisix-guardrails/src/build.rs b/crates/aisix-guardrails/src/build.rs index 248c4941..dbd9b7eb 100644 --- a/crates/aisix-guardrails/src/build.rs +++ b/crates/aisix-guardrails/src/build.rs @@ -995,6 +995,7 @@ pub fn build_index_from_snapshot( let scope_kind = match attachment.scope_type { GuardrailScopeType::Env => ScopeKind::Env, GuardrailScopeType::Model => ScopeKind::Model, + GuardrailScopeType::McpServer => ScopeKind::McpServer, GuardrailScopeType::ApiKey => ScopeKind::ApiKey, GuardrailScopeType::Team => ScopeKind::Team, }; @@ -1887,6 +1888,7 @@ mod tests { let index = build_index_from_snapshot(&shuffled_table(), &attachments, None); let chain = index.resolve(&RequestContext { model_id: "m", + mcp_server_id: "", api_key_id: "k", team_id: None, }); @@ -1937,6 +1939,7 @@ mod tests { let ctx = RequestContext { model_id: "m1", + mcp_server_id: "", api_key_id: "k1", team_id: None, }; @@ -1977,6 +1980,7 @@ mod tests { // Verify the guardrail does not fire (not just that the index is empty). let ctx = RequestContext { model_id: "m", + mcp_server_id: "", api_key_id: "k", team_id: None, }; @@ -2025,6 +2029,7 @@ mod tests { ); let ctx = RequestContext { model_id: "any", + mcp_server_id: "", api_key_id: "any", team_id: None, }; @@ -2066,6 +2071,7 @@ mod tests { let ctx = RequestContext { model_id: "any-model", + mcp_server_id: "", api_key_id: "any-key", team_id: None, }; @@ -2140,6 +2146,7 @@ mod tests { let ctx = RequestContext { model_id: "m1", + mcp_server_id: "", api_key_id: "k1", team_id: None, }; @@ -2193,6 +2200,7 @@ mod tests { let live = LiveGuardrailIndex::new(SnapshotHandle::new(AisixSnapshot::new()), None); let chain = live.resolve(&RequestContext { model_id: "m", + mcp_server_id: "", api_key_id: "k", team_id: None, }); @@ -2259,6 +2267,7 @@ mod tests { LiveGuardrailIndex::new_with_sink(SnapshotHandle::new(snap), None, Some(sink.clone())); let ctx = RequestContext { model_id: "m1", + mcp_server_id: "", api_key_id: "k1", team_id: None, }; @@ -2481,6 +2490,7 @@ mod tests { let index = build_index_from_snapshot(&guardrails, &attachments, None); let chain = index.resolve(&RequestContext { model_id: "m-A", + mcp_server_id: "", api_key_id: "k", team_id: None, }); @@ -2495,6 +2505,59 @@ mod tests { ); } + #[tokio::test] + async fn mcp_server_attachment_builds_into_the_index() { + // The wire `scope_type: "mcp_server"` survives the snapshot build and + // selects on the called server, leaving model traffic alone. + let guardrails: ResourceTable = ResourceTable::default(); + guardrails.insert(entry( + "kw", + "g-1", + parse( + r#"{ + "name": "kw", + "kind": "keyword", + "hook_point": "input", + "patterns": [{ "kind": "literal", "value": "AKIA" }] + }"#, + ), + )); + let attachments: ResourceTable = ResourceTable::default(); + attachments.insert(attachment_entry( + "a-mcp", + parse_attachment( + r#"{ "guardrail_id": "g-1", "scope_type": "mcp_server", "scope_id": "mcp-A", "priority": 50 }"#, + ), + )); + + let index = build_index_from_snapshot(&guardrails, &attachments, None); + assert_eq!(index.len(), 1, "the attachment must not be skipped"); + + let matched = index.resolve(&RequestContext { + model_id: "", + mcp_server_id: "mcp-A", + api_key_id: "k", + team_id: None, + }); + assert_eq!(matched.len(), 1); + + let other_server = index.resolve(&RequestContext { + model_id: "", + mcp_server_id: "mcp-B", + api_key_id: "k", + team_id: None, + }); + assert!(other_server.is_empty()); + + let llm = index.resolve(&RequestContext { + model_id: "m-A", + mcp_server_id: "", + api_key_id: "k", + team_id: None, + }); + assert!(llm.is_empty(), "model traffic carries no MCP server"); + } + #[tokio::test] async fn resolved_chain_applied_empty_when_no_attachment_matches() { // A model-scoped attachment that doesn't match the request resolves to @@ -2524,6 +2587,7 @@ mod tests { let index = build_index_from_snapshot(&guardrails, &attachments, None); let chain = index.resolve(&RequestContext { model_id: "m-OTHER", + mcp_server_id: "", api_key_id: "k", team_id: None, }); diff --git a/crates/aisix-guardrails/src/index.rs b/crates/aisix-guardrails/src/index.rs index 68ee652a..a82ca1eb 100644 --- a/crates/aisix-guardrails/src/index.rs +++ b/crates/aisix-guardrails/src/index.rs @@ -8,10 +8,11 @@ //! //! | `scope_type` | meaning | //! |---|---| -//! | `env` | applies to every request in the environment | -//! | `model` | applies only when the request targets this model UUID | -//! | `api_key` | applies only when authenticated with this API-key UUID | -//! | `team` | applies only when the API key belongs to this team UUID | +//! | `env` | applies to every request in the environment | +//! | `model` | applies only when the request targets this model UUID | +//! | `mcp_server` | applies only to MCP tool calls routed to this server UUID | +//! | `api_key` | applies only when authenticated with this API-key UUID | +//! | `team` | applies only when the API key belongs to this team UUID | //! //! `GuardrailIndex` holds the pre-built runtime guardrails for a snapshot. //! `resolve(ctx)` filters + deduplicates the entries and returns the chain @@ -37,6 +38,7 @@ use crate::{Guardrail, GuardrailChain}; pub enum ScopeKind { Env, Model, + McpServer, ApiKey, Team, } @@ -79,7 +81,13 @@ impl IndexEntry { fn applies_to(&self, ctx: &RequestContext<'_>) -> bool { match self.scope_kind { ScopeKind::Env => true, - ScopeKind::Model => self.scope_id.as_deref() == Some(ctx.model_id), + // `model_id` / `mcp_server_id` are empty on the requests that have + // no such identity (an MCP tool call resolves no model; an LLM + // request routes to no MCP server), so an attachment whose + // `scope_id` is itself empty must NOT match everything that lacks + // the dimension — compare only non-empty ids. + ScopeKind::Model => matches_id(self.scope_id.as_deref(), ctx.model_id), + ScopeKind::McpServer => matches_id(self.scope_id.as_deref(), ctx.mcp_server_id), ScopeKind::ApiKey => self.scope_id.as_deref() == Some(ctx.api_key_id), ScopeKind::Team => ctx .team_id @@ -90,12 +98,22 @@ impl IndexEntry { } } +/// Equality for a scope dimension a request may not carry at all: an empty +/// id on either side never matches. +fn matches_id(scope_id: Option<&str>, ctx_id: &str) -> bool { + !ctx_id.is_empty() && scope_id == Some(ctx_id) +} + /// Per-request context used by [`GuardrailIndex::resolve`] to select and /// deduplicate the applicable guardrails. #[derive(Debug, Clone, Copy)] pub struct RequestContext<'a> { /// UUID of the model the request targets (virtual or concrete). + /// Empty for a request that resolves no model, such as an MCP tool call. pub model_id: &'a str, + /// UUID of the registered MCP server an MCP tool call is routed to. + /// Empty for every request that is not an MCP tool call. + pub mcp_server_id: &'a str, /// UUID of the API key used to authenticate the request. pub api_key_id: &'a str, /// UUID of the team the API key belongs to. `None` if the key is not @@ -115,12 +133,15 @@ pub struct GuardrailIndex { } /// Scope specificity rank: higher = more specific → wins dedup on equal priority. -/// ApiKey > Team > Model > Env, matching the P0c spec in #379. +/// ApiKey > Team > Model > Env, matching the P0c spec in #379. `McpServer` +/// shares `Model`'s rank: the two dimensions are mutually exclusive within a +/// request (an MCP tool call resolves no model and an LLM request routes to no +/// MCP server), so their relative order can never decide a deduplication. fn scope_specificity(k: &ScopeKind) -> u8 { match k { ScopeKind::ApiKey => 3, ScopeKind::Team => 2, - ScopeKind::Model => 1, + ScopeKind::Model | ScopeKind::McpServer => 1, ScopeKind::Env => 0, } } @@ -231,6 +252,7 @@ mod tests { fn ctx<'a>(model: &'a str, apikey: &'a str, team: Option<&'a str>) -> RequestContext<'a> { RequestContext { model_id: model, + mcp_server_id: "", api_key_id: apikey, team_id: team, } @@ -309,6 +331,72 @@ mod tests { ); } + // 3b. McpServer-scope attachment applies only to tool calls routed to + // that server — and never to model traffic, which carries no server id. + #[tokio::test] + async fn mcp_server_scope_only_matching_server() { + let g = kw("g1", "secret"); + let idx = GuardrailIndex::from_entries(vec![entry( + "g1", + ScopeKind::McpServer, + Some("mcp-A"), + 50, + g, + )]); + + let mcp_ctx = |server: &'static str| RequestContext { + model_id: "", + mcp_server_id: server, + api_key_id: "k1", + team_id: None, + }; + + let chain_a = idx.resolve(&mcp_ctx("mcp-A")); + assert!(chain_a.check_input(&req("secret")).await.is_block()); + + let chain_b = idx.resolve(&mcp_ctx("mcp-B")); + assert_eq!( + chain_b.check_input(&req("secret")).await, + GuardrailVerdict::Allow + ); + + // An LLM request carries no MCP server, so the attachment is inert. + let chain_llm = idx.resolve(&ctx("model-A", "k1", None)); + assert_eq!( + chain_llm.check_input(&req("secret")).await, + GuardrailVerdict::Allow + ); + } + + // 3c. A dimension a request does not carry never matches, even against an + // attachment whose own scope_id is empty — the sentinel for "absent" must + // not read as a wildcard. + #[tokio::test] + async fn empty_scope_id_never_matches_an_absent_dimension() { + let idx = GuardrailIndex::from_entries(vec![ + entry("g1", ScopeKind::McpServer, Some(""), 50, kw("g1", "secret")), + entry("g2", ScopeKind::Model, Some(""), 50, kw("g2", "secret")), + ]); + + // An MCP tool call has no model; an LLM request has no MCP server. + let mcp_chain = idx.resolve(&RequestContext { + model_id: "", + mcp_server_id: "mcp-A", + api_key_id: "k1", + team_id: None, + }); + assert_eq!( + mcp_chain.check_input(&req("secret")).await, + GuardrailVerdict::Allow + ); + + let llm_chain = idx.resolve(&ctx("model-A", "k1", None)); + assert_eq!( + llm_chain.check_input(&req("secret")).await, + GuardrailVerdict::Allow + ); + } + // 4. ApiKey-scope attachment applies only to the matching key. #[tokio::test] async fn api_key_scope_only_matching_key() { @@ -569,6 +657,7 @@ mod tests { let team = format!("scope-{}", (i + 7) % 10); let _ = idx.resolve(&RequestContext { model_id: &model, + mcp_server_id: "", api_key_id: &key, team_id: Some(&team), }); diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 769a0964..077d0571 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -577,6 +577,7 @@ async fn multipart_dispatch( // The transcript RESPONSE is scanned/masked after the upstream call. let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; @@ -1082,6 +1083,7 @@ async fn speech_dispatch( // RPM slot. (Output is binary audio, not scannable text — no output hook.) let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 596c0da5..5ac6c397 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1232,6 +1232,7 @@ async fn dispatch( // guardrail context. let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 73996245..ec1c7578 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -294,6 +294,7 @@ async fn dispatch( // (matching /v1/chat/completions). let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index f4bf2458..482b22e9 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -319,6 +319,7 @@ async fn dispatch( // block doesn't burn an RPM slot (matching /v1/chat/completions). let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/error.rs b/crates/aisix-proxy/src/error.rs index d1d954ca..f65f9a72 100644 --- a/crates/aisix-proxy/src/error.rs +++ b/crates/aisix-proxy/src/error.rs @@ -306,9 +306,11 @@ pub enum ProxyError { /// metadata, safe to surface (#519 B.4b) — but never the matched-pattern /// detail (per #153 that detail stays in `tracing` only; echoing it lets /// callers enumerate the blocklist or extract the blocked output). -/// `side` is `"request"` (input hook) or `"response"` (output hook). -/// Every proxy endpoint family builds its 422 / SSE-error message through -/// this helper so the envelope shape can't drift between siblings. +/// `side` is `"request"` (input hook) or `"response"` (output hook) — the MCP +/// endpoints pass `"tool call"` / `"tool result"` instead. Every endpoint +/// family builds its rejection text through this helper so the wording can't +/// drift between siblings, even where the envelope differs (422 or an SSE +/// error event on the LLM routes; an `isError` tool result on `/mcp`). pub(crate) fn guardrail_block_message(side: &str, guardrail_name: Option<&str>) -> String { match guardrail_name { Some(name) => format!("{side} blocked by content policy (guardrail '{name}')"), diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index 0ae43e1a..a9d26e1a 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -259,6 +259,7 @@ async fn dispatch( // text, so there is no output hook.) let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/jobs.rs b/crates/aisix-proxy/src/jobs.rs index da012acb..74dbb55d 100644 --- a/crates/aisix-proxy/src/jobs.rs +++ b/crates/aisix-proxy/src/jobs.rs @@ -445,6 +445,7 @@ async fn scan_input_blob( ) -> Result<(), ProxyError> { let ctx = aisix_guardrails::RequestContext { model_id: &target.model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; @@ -488,6 +489,7 @@ async fn scan_output_blob( ) -> Result<(), ProxyError> { let ctx = aisix_guardrails::RequestContext { model_id: &target.model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index 72b6db66..14886c8c 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -270,14 +270,23 @@ async fn dispatch( // Resolve the guardrail chain once and run BOTH directions through the SAME // chain as LLM traffic: the tool arguments (input) before the call, and the // tool result (output) after. MCP has no model, so an empty `model_id` - // matches env / api-key / team-scoped guardrails, never a Model-scoped one. + // never matches a Model-scoped guardrail; the called server's id carries + // the MCP-side dimension instead, so env / mcp-server / api-key / team + // scopes all apply. An unregistered server name (a malformed namespaced + // tool) leaves the id empty and simply matches no MCP-server scope. // An empty chain short-circuits, keeping the no-guardrail path cheap (and // skipping the response buffering the output check needs). let rpc_id = peek.as_ref().and_then(|p| p.id.clone()); let guardrail_chain = is_tool_call .then(|| { + let mcp_server_id = snapshot + .mcp_servers + .get_by_name(&mcp_server) + .map(|entry| entry.id.clone()) + .unwrap_or_default(); let ctx = aisix_guardrails::RequestContext { model_id: "", + mcp_server_id: &mcp_server_id, api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; @@ -437,9 +446,7 @@ async fn output_guardrail_block( // and LLM output on the same representation: a keyword guardrail sees the // decoded prose, so envelope field names (`content`, `type`, `text`) can't // trip a false positive, and escaped characters can't hide blocked content. - // Fall back to the whole serialized result for non-standard shapes so - // nothing escapes inspection. - let result_text = result + let mut scanned: Vec = result .get("content") .and_then(|c| c.as_array()) .map(|blocks| { @@ -447,11 +454,26 @@ async fn output_guardrail_block( .iter() .filter(|b| b.get("type").and_then(|t| t.as_str()) == Some("text")) .filter_map(|b| b.get("text").and_then(|t| t.as_str())) - .collect::>() - .join("\n") + .filter(|text| !text.is_empty()) + .map(str::to_owned) + .collect() }) - .filter(|text| !text.is_empty()) - .unwrap_or_else(|| result.to_string()); + .unwrap_or_default(); + // `structuredContent` is serialized to the client ALONGSIDE `content`, and + // the spec only RECOMMENDS mirroring it into a text block — so a tool can + // return clean prose and carry the sensitive value here. Scan its string + // leaves, for the same reason the content blocks are decoded first: the + // values are the tool's data, while the object keys are its output schema. + if let Some(structured) = result.get("structuredContent") { + collect_string_leaves(structured, &mut scanned); + } + // Fall back to the whole serialized result for non-standard shapes so + // nothing escapes inspection. + let result_text = if scanned.is_empty() { + result.to_string() + } else { + scanned.join("\n") + }; let resp = aisix_gateway::ChatResponse { id: String::new(), model: String::new(), @@ -478,6 +500,24 @@ async fn output_guardrail_block( } } +/// Push every non-empty string leaf of `value` — walking objects and arrays — +/// onto `out`, in document order. Object KEYS are skipped: they are the tool's +/// declared output-schema field names rather than its data, so scanning them +/// would reintroduce the field-name false positives that decoding the content +/// blocks avoids. The walk is iterative, so a deeply nested result cannot +/// recurse the handler's stack. +fn collect_string_leaves(value: &serde_json::Value, out: &mut Vec) { + let mut stack = vec![value]; + while let Some(node) = stack.pop() { + match node { + serde_json::Value::String(text) if !text.is_empty() => out.push(text.clone()), + serde_json::Value::Array(items) => stack.extend(items.iter().rev()), + serde_json::Value::Object(map) => stack.extend(map.values().rev()), + _ => {} + } + } +} + /// Emit a usage event for a single MCP tool call into the same sink as LLM /// usage. MCP calls carry no token cost yet, so token/cost fields stay zero; /// the event records who called which tool, the outcome, and the latency. @@ -523,12 +563,18 @@ fn emit_tool_call_usage( .fan_out(&event, None, exporters.iter().map(|e| &e.value)); } -/// Build the MCP-native response for a guardrail block: a JSON-RPC error -/// echoing the request id, served as HTTP 200 with a JSON body (the MCP -/// Streamable HTTP shape). Both the input and output hooks funnel through here; -/// `side` (`"tool call"` for input arguments, `"tool result"` for output) -/// selects the caller-visible wording. Unlike the LLM path's 422, an MCP client -/// expects a JSON-RPC envelope, so the block surfaces as an error it can handle. +/// Build the MCP-native response for a guardrail block: a `tools/call` result +/// flagged `isError`, echoing the request id, served as HTTP 200 with a JSON +/// body (the MCP Streamable HTTP shape). Both the input and output hooks funnel +/// through here; `side` (`"tool call"` for input arguments, `"tool result"` for +/// output) selects the caller-visible wording. +/// +/// A tool-execution error rather than a JSON-RPC protocol error: MCP separates +/// "this request was not valid" (a protocol error, which a client surfaces as a +/// transport-level failure) from "the tool call did not succeed" (`isError` on +/// the result, which the calling agent reads as tool output and can adapt to). +/// A policy rejection is the second kind — the request was well-formed and the +/// caller should learn, in-band, that content policy stopped it. fn jsonrpc_guardrail_block( id: Option, side: &str, @@ -538,7 +584,10 @@ fn jsonrpc_guardrail_block( let body = serde_json::json!({ "jsonrpc": "2.0", "id": id.unwrap_or(serde_json::Value::Null), - "error": { "code": -32600, "message": message } + "result": { + "content": [{ "type": "text", "text": message }], + "isError": true, + } }); ( StatusCode::OK, @@ -1147,11 +1196,42 @@ mod tests { /// Seed an env-scoped guardrail (from its JSON) by RCU-inserting it + an /// attachment into the live snapshot handle. fn seed_guardrail(handle: &SnapshotHandle, guardrail_json: &str) { + seed_guardrail_with_attachment( + handle, + guardrail_json, + r#"{"guardrail_id":"g1","scope_type":"env","priority":50}"#, + ); + } + + /// Register an MCP server under `name` with the resource id `id`, so a + /// tool call naming it resolves the id an `mcp_server`-scoped attachment + /// matches on. + fn seed_mcp_server(handle: &SnapshotHandle, id: &'static str, name: &str) { + use aisix_core::models::McpServer; + let server: McpServer = serde_json::from_value(serde_json::json!({ + "name": name, + // Never dialled: the guardrail verdicts under test are decided + // before the gateway is built (input) or on a synthesized body + // (output). + "url": "http://127.0.0.1:1/mcp", + })) + .expect("valid mcp server"); + handle.rcu(|snap| { + let new = snap.clone(); + new.mcp_servers + .insert(ResourceEntry::new(id, server.clone(), 1)); + new + }); + } + + fn seed_guardrail_with_attachment( + handle: &SnapshotHandle, + guardrail_json: &str, + attachment_json: &str, + ) { use aisix_core::models::{Guardrail, GuardrailAttachment}; let guardrail: Guardrail = serde_json::from_str(guardrail_json).unwrap(); - let attachment: GuardrailAttachment = - serde_json::from_str(r#"{"guardrail_id":"g1","scope_type":"env","priority":50}"#) - .unwrap(); + let attachment: GuardrailAttachment = serde_json::from_str(attachment_json).unwrap(); handle.rcu(|snap| { let new = snap.clone(); new.guardrails @@ -1181,10 +1261,10 @@ mod tests { seed_guardrail(&handle, INPUT_GUARD); // Arguments carrying the forbidden token are blocked by the same - // guardrail chain LLM input uses — surfaced as an MCP-native JSON-RPC - // error (HTTP 200) before the gateway/upstream is reached. A distinctive - // request id (7) proves the handler echoes the caller's id (not a - // constant) through the block envelope both hooks funnel through. + // guardrail chain LLM input uses — surfaced as an MCP-native tool-error + // result (HTTP 200) before the gateway/upstream is reached. A + // distinctive request id (7) proves the handler echoes the caller's id + // (not a constant) through the block envelope both hooks funnel through. let blocked = router .clone() .oneshot(mcp_request_with_id( @@ -1206,13 +1286,17 @@ mod tests { serde_json::json!(7), "the block must echo the request id" ); - assert_eq!(envelope["error"]["code"], -32600); + assert_eq!( + envelope["result"]["isError"], + serde_json::json!(true), + "a policy rejection is a tool-execution error, not a protocol error" + ); assert!( - envelope.get("result").is_none(), - "a guardrail block carries no result" + envelope.get("error").is_none(), + "a guardrail block must not surface as a JSON-RPC protocol error" ); assert!( - envelope["error"]["message"] + envelope["result"]["content"][0]["text"] .as_str() .unwrap_or_default() .contains("content policy"), @@ -1244,6 +1328,7 @@ mod tests { let index = LiveGuardrailIndex::new(handle, None); let chain = index.resolve(&RequestContext { model_id: "", + mcp_server_id: "", api_key_id: "ak-1", team_id: None, }); @@ -1301,6 +1386,7 @@ mod tests { seed_guardrail(&handle, FIELD_NAME_GUARD); let chain = LiveGuardrailIndex::new(handle, None).resolve(&RequestContext { model_id: "", + mcp_server_id: "", api_key_id: "ak-1", team_id: None, }); @@ -1330,7 +1416,8 @@ mod tests { async fn output_block_envelope_echoes_id_and_shape() { // Both hooks funnel the block through `jsonrpc_guardrail_block`; assert // the wire envelope directly so a regression that nulls the id or shifts - // the code/status/content-type is caught without an rmcp upstream. + // the result shape/status/content-type is caught without an rmcp + // upstream. let resp = jsonrpc_guardrail_block( Some(serde_json::json!(42)), "tool result", @@ -1353,13 +1440,17 @@ mod tests { serde_json::json!(42), "the original JSON-RPC id must be echoed, not nulled" ); - assert_eq!(v["error"]["code"], -32600); + // A tool-execution error (`isError` on the result), not a JSON-RPC + // protocol error: the calling agent reads the rejection as tool output + // it can react to instead of a transport-level failure. + assert_eq!(v["result"]["isError"], serde_json::json!(true)); assert!( - v.get("result").is_none(), - "a block envelope carries no result" + v.get("error").is_none(), + "a block envelope carries no protocol error" ); + assert_eq!(v["result"]["content"][0]["type"], "text"); assert!( - v["error"]["message"] + v["result"]["content"][0]["text"] .as_str() .unwrap_or_default() .contains("tool result blocked by content policy"), @@ -1367,6 +1458,136 @@ mod tests { ); } + /// A tool result can carry data in `structuredContent` that is NOT mirrored + /// into a text content block (the spec only recommends mirroring), and the + /// gateway relays that field to the client verbatim. Scanning only the + /// content blocks would therefore let it through unread. + #[tokio::test] + async fn output_guardrail_scans_structured_content() { + use aisix_guardrails::{LiveGuardrailIndex, RequestContext}; + + let handle = SnapshotHandle::new(snapshot_with_key()); + seed_guardrail(&handle, OUTPUT_GUARD); + let chain = LiveGuardrailIndex::new(handle, None).resolve(&RequestContext { + model_id: "", + mcp_server_id: "", + api_key_id: "ak-1", + team_id: None, + }); + + // Clean prose, sensitive structured payload — the client sees both. + let structured_only = br#"{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"lookup ok"}],"structuredContent":{"record":{"notes":["forbidden-token"]}}}}"#; + assert!( + output_guardrail_block(&chain, structured_only, "lookup", &mut Vec::new()) + .await + .is_some(), + "structuredContent reaches the client, so it must be scanned" + ); + + // A result with no content blocks at all — a tool may return only + // structured output — is scanned through the same path. + let no_content = br#"{"jsonrpc":"2.0","id":1,"result":{"content":[],"structuredContent":{"note":"forbidden-token"}}}"#; + assert!( + output_guardrail_block(&chain, no_content, "lookup", &mut Vec::new()) + .await + .is_some(), + "a structured-only result must still be scanned" + ); + + // Clean on both sides passes. + let clean = br#"{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"lookup ok"}],"structuredContent":{"record":{"notes":["all good"]}}}}"#; + assert!( + output_guardrail_block(&chain, clean, "lookup", &mut Vec::new()) + .await + .is_none(), + "a clean structured result must not be blocked" + ); + } + + /// The structured walk carries the same "scan data, not field names" rule + /// the content blocks follow: an object KEY matching the pattern is the + /// tool's output schema, not its data, and must not fire. + #[tokio::test] + async fn structured_content_keys_do_not_trip_the_guardrail() { + use aisix_guardrails::{LiveGuardrailIndex, RequestContext}; + + const KEY_NAME_GUARD: &str = r#"{"name":"key-name-guard","kind":"keyword","hook_point":"output","patterns":[{"kind":"literal","value":"ssn"}]}"#; + let handle = SnapshotHandle::new(snapshot_with_key()); + seed_guardrail(&handle, KEY_NAME_GUARD); + let chain = LiveGuardrailIndex::new(handle, None).resolve(&RequestContext { + model_id: "", + mcp_server_id: "", + api_key_id: "ak-1", + team_id: None, + }); + + let key_only = br#"{"jsonrpc":"2.0","id":1,"result":{"content":[],"structuredContent":{"ssn":"redacted upstream"}}}"#; + assert!( + output_guardrail_block(&chain, key_only, "lookup", &mut Vec::new()) + .await + .is_none(), + "a schema field name must not be treated as tool data" + ); + + // The same pattern in a VALUE fires, proving the guardrail is live here. + let value_hit = br#"{"jsonrpc":"2.0","id":1,"result":{"content":[],"structuredContent":{"field":"ssn 123-45-6789"}}}"#; + assert!( + output_guardrail_block(&chain, value_hit, "lookup", &mut Vec::new()) + .await + .is_some(), + "the pattern in a structured VALUE must block" + ); + } + + /// An `mcp_server`-scoped attachment governs only the tool calls routed to + /// that server — the dimension a Model scope cannot express for MCP. + #[tokio::test] + async fn mcp_server_scoped_guardrail_applies_only_to_that_server() { + let handle = SnapshotHandle::new(snapshot_with_key()); + let hub = Arc::new(aisix_gateway::Hub::new()); + let state = ProxyState::new(handle.clone(), hub, &cfg()).without_cache(); + let router = build_router(state); + seed_mcp_server(&handle, "mcp-ghost", "ghost"); + seed_mcp_server(&handle, "mcp-other", "other"); + seed_guardrail_with_attachment( + &handle, + INPUT_GUARD, + r#"{"guardrail_id":"g1","scope_type":"mcp_server","scope_id":"mcp-ghost","priority":50}"#, + ); + + let blocked_body = |name: &str| { + mcp_request( + "tools/call", + serde_json::json!({ "name": name, "arguments": { "q": "forbidden-token" } }), + ) + }; + let body_text = |resp: Response| async move { + let bytes = axum::body::to_bytes(resp.into_body(), 64 * 1024) + .await + .expect("read body"); + String::from_utf8_lossy(&bytes).into_owned() + }; + + let scoped = router + .clone() + .oneshot(blocked_body("ghost__tool")) + .await + .expect("router responds"); + assert!( + body_text(scoped).await.contains("content policy"), + "the attached server's tool call must be blocked" + ); + + let other = router + .oneshot(blocked_body("other__tool")) + .await + .expect("router responds"); + assert!( + !body_text(other).await.contains("content policy"), + "a server the guardrail is not attached to must be untouched" + ); + } + /// #698: a tool-call usage event must reach the per-env observability /// exporters via the OTLP fan-out — pre-fix MCP usage was emitted only /// into the CP sink, so exporters never saw /mcp traffic. Uses the ghost diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 9565fdd8..e69e6a03 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -545,6 +545,7 @@ async fn dispatch( // block doesn't burn an RPM slot (matching /v1/chat/completions). let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index b38e7892..95c016c7 100644 --- a/crates/aisix-proxy/src/passthrough.rs +++ b/crates/aisix-proxy/src/passthrough.rs @@ -338,6 +338,7 @@ async fn dispatch( // content/DLP policy as the typed surfaces. Empty chain → no scan, no cost. let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index 0a0febd8..b08ba21f 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -481,6 +481,7 @@ async fn run_session( let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index aba22bfd..a4fb37b5 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -273,6 +273,7 @@ async fn dispatch( // indices/scores, not generated text, so there is no output hook.) let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 1c6341cb..a52c29ad 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -525,6 +525,7 @@ async fn dispatch( // block doesn't burn an RPM slot (matching /v1/chat/completions). let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/crates/aisix-proxy/src/videos.rs b/crates/aisix-proxy/src/videos.rs index 902ae51d..e0f638a5 100644 --- a/crates/aisix-proxy/src/videos.rs +++ b/crates/aisix-proxy/src/videos.rs @@ -1641,6 +1641,7 @@ async fn dispatch_create( // block doesn't burn an RPM slot (#542). let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &target.model_entry.id, + mcp_server_id: "", api_key_id: &auth.entry.id, team_id: auth.key().team_id.as_deref(), }; diff --git a/schemas/resources/guardrail_attachment.schema.json b/schemas/resources/guardrail_attachment.schema.json index 60b4f3c4..4e6cfb27 100644 --- a/schemas/resources/guardrail_attachment.schema.json +++ b/schemas/resources/guardrail_attachment.schema.json @@ -2,17 +2,18 @@ "$schema": "http://json-schema.org/draft-07/schema#", "definitions": { "GuardrailScopeType": { - "description": "Which dimension of the request a guardrail attachment is scoped to.\n\n`Env` applies to every request in the environment. The narrower scopes let operators attach a guardrail to only the models, API keys, or teams that need it.", + "description": "Which dimension of the request a guardrail attachment is scoped to.\n\n`Env` applies to every request in the environment. The narrower scopes let operators attach a guardrail to only the models, MCP servers, API keys, or teams that need it.\n\n`Model` and `McpServer` select dimensions a request carries only one of: an MCP tool call resolves no model, and an LLM request routes to no MCP server. A `Model`-scoped guardrail therefore never inspects MCP traffic, and an `McpServer`-scoped one never inspects model traffic.", "enum": [ "env", "model", + "mcp_server", "api_key", "team" ], "type": "string" } }, - "description": "Guardrail attachment that scopes one guardrail to an environment, model, caller API key, or team. AISIX loads attachments with the guardrail definitions and uses `scope_type` plus `scope_id` to decide which guardrails apply to each request.", + "description": "Guardrail attachment that scopes one guardrail to an environment, model, MCP server, caller API key, or team. AISIX loads attachments with the guardrail definitions and uses `scope_type` plus `scope_id` to decide which guardrails apply to each request.", "properties": { "enabled": { "default": true, @@ -37,7 +38,7 @@ "type": "integer" }, "scope_id": { - "description": "The UUID of the specific resource (model / api_key / team). `None` when `scope_type` is `Env` (applies to all requests).", + "description": "The UUID of the specific resource (model / mcp_server / api_key / team). `None` when `scope_type` is `Env` (applies to all requests).", "type": [ "string", "null" diff --git a/tests/e2e/src/cases/mcp-guardrail-e2e.test.ts b/tests/e2e/src/cases/mcp-guardrail-e2e.test.ts new file mode 100644 index 00000000..c08eac62 --- /dev/null +++ b/tests/e2e/src/cases/mcp-guardrail-e2e.test.ts @@ -0,0 +1,221 @@ +import { createHash, randomUUID } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startMcpUpstream, + waitConfigPropagation, + type McpUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: guardrails over the MCP gateway, against a real gateway + etcd + two +// real MCP upstreams (official TypeScript SDK servers). +// +// Pinned contract: +// - a policy rejection is a TOOL-EXECUTION error (`result.isError`), never a +// JSON-RPC protocol error, so the calling agent reads it as tool output; +// - `structuredContent` is scanned, not just the text content blocks — the +// gateway relays that field to the client verbatim, so a tool that returns +// clean prose plus a sensitive structured payload must not slip through; +// - an `mcp_server`-scoped attachment governs only the tool calls routed to +// that server, the dimension no model scope can express for MCP. + +const KEY = "sk-mcp-guardrail-e2e"; +const sha256 = (s: string) => createHash("sha256").update(s).digest("hex"); + +/** Distinct patterns so the three guardrails can coexist in one env. */ +const ALPHA_ONLY_PATTERN = "poison-alpha-only"; +const EVERYWHERE_PATTERN = "blocked-everywhere"; +const STRUCTURED_PATTERN = "classified-payload"; + +interface RpcReply { + status: number; + json?: { + result?: { + content?: Array<{ type: string; text?: string }>; + structuredContent?: Record; + isError?: boolean; + }; + error?: { code: number; message: string }; + }; +} + +describe("mcp guardrails e2e: /mcp", () => { + let app: SpawnedApp | undefined; + let alpha: McpUpstream | undefined; + let beta: McpUpstream | undefined; + let etcdReachable = false; + let seed: SeedClient; + + const post = async (body: unknown): Promise => { + const res = await fetch(`${app!.proxyUrl}/mcp`, { + method: "POST", + headers: { + authorization: `Bearer ${KEY}`, + "content-type": "application/json", + accept: "application/json, text/event-stream", + }, + body: JSON.stringify(body), + }); + const text = await res.text(); + let json: RpcReply["json"]; + try { + json = text ? (JSON.parse(text) as RpcReply["json"]) : undefined; + } catch { + json = undefined; + } + return { status: res.status, json }; + }; + + /** Spec-faithful per-operation handshake (the endpoint is stateless). */ + const callTool = async (name: string, text: string): Promise => { + await post({ + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "mcp-guardrail-e2e", version: "0.1" }, + }, + }); + return post({ + jsonrpc: "2.0", + id: 3, + method: "tools/call", + params: { name, arguments: { text } }, + }); + }; + + const blockedByPolicy = (reply: RpcReply): boolean => + reply.json?.result?.isError === true && + (reply.json.result.content?.[0]?.text ?? "").includes("content policy"); + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + alpha = await startMcpUpstream("alpha", { structuredTool: true }); + beta = await startMcpUpstream("beta"); + app = await spawnApp(); + seed = new SeedClient(etcd, app.etcdPrefix); + + const alphaId = randomUUID(); + await seed.update("mcp_servers", alphaId, { + display_name: "alpha", + url: alpha.url, + enabled: true, + }); + await seed.update("mcp_servers", randomUUID(), { + display_name: "beta", + url: beta.url, + enabled: true, + }); + await seed.createApiKey({ + key_hash: sha256(KEY), + allowed_models: [], + allowed_tools: ["*"], + }); + + // One guardrail per behaviour under test, each with its own pattern, so + // all three can be attached at once without interfering. + const alphaOnly = await seed.createGuardrail({ + name: "mcp-alpha-only", + kind: "keyword", + patterns: [{ kind: "literal", value: ALPHA_ONLY_PATTERN }], + }); + const everywhere = await seed.createGuardrail({ + name: "mcp-everywhere", + kind: "keyword", + patterns: [{ kind: "literal", value: EVERYWHERE_PATTERN }], + }); + const structured = await seed.createGuardrail({ + name: "mcp-structured", + kind: "keyword", + hook_point: "output", + patterns: [{ kind: "literal", value: STRUCTURED_PATTERN }], + }); + + await seed.update("guardrail_attachments", randomUUID(), { + guardrail_id: alphaOnly.id, + scope_type: "mcp_server", + scope_id: alphaId, + priority: 100, + }); + for (const id of [everywhere.id, structured.id]) { + await seed.update("guardrail_attachments", randomUUID(), { + guardrail_id: id, + scope_type: "env", + priority: 50, + }); + } + + // Probe the narrowest condition — the server-scoped attachment matching + // alpha — so no assertion races a row that has not landed yet. + await waitConfigPropagation(async () => + blockedByPolicy(await callTool("alpha__echo", ALPHA_ONLY_PATTERN)), + ); + }, 60_000); + + afterAll(async () => { + await app?.exit(); + await alpha?.close(); + await beta?.close(); + }); + + test("a policy rejection is a tool error, not a protocol error", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + const reply = await callTool("beta__echo", `say ${EVERYWHERE_PATTERN}`); + + expect(reply.status).toBe(200); + expect(reply.json?.error).toBeUndefined(); + expect(reply.json?.result?.isError).toBe(true); + expect(reply.json?.result?.content?.[0]?.text).toContain("content policy"); + // The firing rule is named so an operator can find it; the matched + // content never is. + expect(reply.json?.result?.content?.[0]?.text).toContain("mcp-everywhere"); + expect(reply.json?.result?.content?.[0]?.text).not.toContain( + EVERYWHERE_PATTERN, + ); + }); + + test("structuredContent reaches the client, so it is scanned", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + // First: the field really is relayed verbatim. This is the reason it must + // be scanned — without this leg, a passing block test could just mean the + // gateway had dropped the field. + const clean = await callTool("alpha__lookup", "ordinary value"); + expect(clean.json?.result?.isError).toBeFalsy(); + expect(clean.json?.result?.structuredContent).toEqual({ + record: { note: "ordinary value" }, + }); + + // The text block `lookup` returns is constant and clean, so only the + // structured payload carries the pattern. + const blocked = await callTool("alpha__lookup", STRUCTURED_PATTERN); + expect(blocked.json?.result?.isError).toBe(true); + expect(blocked.json?.result?.content?.[0]?.text).toContain("tool result"); + expect(blocked.json?.result?.structuredContent).toBeUndefined(); + }); + + test("an mcp_server scope governs only its own server", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + const onAlpha = await callTool("alpha__echo", `carrying ${ALPHA_ONLY_PATTERN}`); + expect(onAlpha.json?.result?.isError).toBe(true); + expect(onAlpha.json?.result?.content?.[0]?.text).toContain("mcp-alpha-only"); + + // Same content, different server: the attachment does not reach it, and + // the call runs through to the upstream. + const onBeta = await callTool("beta__echo", `carrying ${ALPHA_ONLY_PATTERN}`); + expect(onBeta.json?.result?.isError).toBeFalsy(); + expect(onBeta.json?.result?.content?.[0]?.text).toBe( + `beta:carrying ${ALPHA_ONLY_PATTERN}`, + ); + }); +}); diff --git a/tests/e2e/src/harness/index.ts b/tests/e2e/src/harness/index.ts index 79c08e3e..92c31dc8 100644 --- a/tests/e2e/src/harness/index.ts +++ b/tests/e2e/src/harness/index.ts @@ -4,7 +4,11 @@ export { ProxyClient } from "./proxy.js"; export { EtcdClient } from "./etcd.js"; export { SeedClient } from "./seed.js"; export { startOpenAiUpstream, type OpenAiUpstream, type ReceivedRequest } from "./upstream-openai.js"; -export { startMcpUpstream, type McpUpstream } from "./upstream-mcp.js"; +export { + startMcpUpstream, + type McpUpstream, + type McpUpstreamOptions, +} from "./upstream-mcp.js"; export { startA2aUpstream, type A2aUpstream, diff --git a/tests/e2e/src/harness/upstream-mcp.ts b/tests/e2e/src/harness/upstream-mcp.ts index 980aba20..2378a3fc 100644 --- a/tests/e2e/src/harness/upstream-mcp.ts +++ b/tests/e2e/src/harness/upstream-mcp.ts @@ -18,6 +18,16 @@ export interface McpUpstream { close(): Promise; } +export interface McpUpstreamOptions { + /** + * Also expose a `lookup` tool whose result carries `structuredContent` + * alongside a fixed, always-clean text block — the shape a tool uses to + * return machine-readable output. Off by default so the tool inventory + * every other suite asserts on stays `echo` + `reverse`. + */ + structuredTool?: boolean; +} + /** * A real MCP upstream server built on the official TypeScript SDK, speaking * the stateless Streamable HTTP transport with JSON responses — the exact @@ -31,9 +41,12 @@ export interface McpUpstream { * A fresh SDK `Server` + transport is built per request (the SDK's stateless * pattern); the gateway reconnects per operation, so nothing is shared. */ -export async function startMcpUpstream(label: string): Promise { +export async function startMcpUpstream( + label: string, + options: McpUpstreamOptions = {}, +): Promise { const httpServer: HttpServer = createServer((req, res) => { - void handle(label, req, res); + void handle(label, req, res, options); }); await new Promise((resolve) => httpServer.listen(0, "127.0.0.1", resolve), @@ -53,6 +66,7 @@ async function handle( label: string, req: IncomingMessage, res: ServerResponse, + options: McpUpstreamOptions, ): Promise { try { if (req.method !== "POST") { @@ -85,10 +99,31 @@ async function handle( properties: { text: { type: "string" } }, }, }, + ...(options.structuredTool + ? [ + { + name: "lookup", + description: "return the text as structured output", + inputSchema: { + type: "object" as const, + properties: { text: { type: "string" } }, + }, + }, + ] + : []), ], })); server.setRequestHandler(CallToolRequestSchema, async (request) => { const text = String(request.params.arguments?.text ?? ""); + if (request.params.name === "lookup") { + // The text block is deliberately constant and clean: only + // `structuredContent` carries the caller's value, which is exactly + // the case a content-blocks-only scan would miss. + return { + content: [{ type: "text", text: "lookup ok" }], + structuredContent: { record: { note: text } }, + }; + } const out = request.params.name === "reverse" ? [...text].reverse().join("") From 8931c977a5fb7cf89109a4f91d0810748266bc9c Mon Sep 17 00:00:00 2001 From: Jarvis Date: Sun, 16 Aug 2026 10:31:03 +0800 Subject: [PATCH 2/2] test(e2e): gate the mcp guardrail spec on the caller key, seeded last MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The readiness gate exercised the very behaviour the spec asserts (an mcp_server-scoped block), so a real assertion failure would have surfaced as a beforeAll timeout, and its JSON-parse fallback turned an upstream or transport fault into "not ready". It also probed only the alpha attachment while later tests assert on the env-scoped guardrails and the beta server — separate etcd keys with independent propagation. Follow the harness convention instead: seed the caller API key after every other resource and gate on it authenticating through ProxyClient.listModels, which cannot throw. Etcd applies writes in revision order, so that one condition implies the whole seed set. --- tests/e2e/src/cases/mcp-guardrail-e2e.test.ts | 27 +++++++++---------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/tests/e2e/src/cases/mcp-guardrail-e2e.test.ts b/tests/e2e/src/cases/mcp-guardrail-e2e.test.ts index c08eac62..670aa0e5 100644 --- a/tests/e2e/src/cases/mcp-guardrail-e2e.test.ts +++ b/tests/e2e/src/cases/mcp-guardrail-e2e.test.ts @@ -2,6 +2,7 @@ import { createHash, randomUUID } from "node:crypto"; import { afterAll, beforeAll, describe, expect, test } from "vitest"; import { EtcdClient, + ProxyClient, SeedClient, spawnApp, startMcpUpstream, @@ -89,10 +90,6 @@ describe("mcp guardrails e2e: /mcp", () => { }); }; - const blockedByPolicy = (reply: RpcReply): boolean => - reply.json?.result?.isError === true && - (reply.json.result.content?.[0]?.text ?? "").includes("content policy"); - beforeAll(async () => { const etcd = new EtcdClient(); etcdReachable = await etcd.ping(); @@ -114,12 +111,6 @@ describe("mcp guardrails e2e: /mcp", () => { url: beta.url, enabled: true, }); - await seed.createApiKey({ - key_hash: sha256(KEY), - allowed_models: [], - allowed_tools: ["*"], - }); - // One guardrail per behaviour under test, each with its own pattern, so // all three can be attached at once without interfering. const alphaOnly = await seed.createGuardrail({ @@ -153,11 +144,17 @@ describe("mcp guardrails e2e: /mcp", () => { }); } - // Probe the narrowest condition — the server-scoped attachment matching - // alpha — so no assertion races a row that has not landed yet. - await waitConfigPropagation(async () => - blockedByPolicy(await callTool("alpha__echo", ALPHA_ONLY_PATTERN)), - ); + // The caller key is written LAST, so the key authenticating implies every + // row above it is already in the snapshot (etcd applies in revision + // order). The gate deliberately touches neither `/mcp` nor a guardrail: + // a broken assertion must fail as an assertion, not as a gate timeout. + await seed.createApiKey({ + key_hash: sha256(KEY), + allowed_models: [], + allowed_tools: ["*"], + }); + const proxy = new ProxyClient(app.proxyUrl, KEY); + await waitConfigPropagation(async () => (await proxy.listModels()).status === 200); }, 60_000); afterAll(async () => {