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..670aa0e5 --- /dev/null +++ b/tests/e2e/src/cases/mcp-guardrail-e2e.test.ts @@ -0,0 +1,218 @@ +import { createHash, randomUUID } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + ProxyClient, + 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 } }, + }); + }; + + 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, + }); + // 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, + }); + } + + // 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 () => { + 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("")