diff --git a/README.md b/README.md index 4d97de14b..5349cb75f 100644 --- a/README.md +++ b/README.md @@ -1503,8 +1503,8 @@ untouched. - `Mcp-Param-*` — selected `tools/call` arguments, promoted from the tool's input schema. -To promote a tool argument into a routing header, annotate the top-level schema -property with `x-mcp-header`: +To promote a tool argument into a routing header, annotate its schema property +with `x-mcp-header`: ```rust, ignore // A `region` argument surfaces as the `Mcp-Param-Region` request header. @@ -1517,7 +1517,29 @@ let schema = serde_json::json!({ ``` Annotations must be non-empty RFC 9110 tokens, case-insensitively unique, and -applied only to top-level primitive (`string`/`integer`/`boolean`) properties. +applied only to primitive (`string`/`integer`/`boolean`) properties. Nested +properties are allowed when every step from the root is a `properties` key; an +annotation under `items`, `oneOf`/`anyOf`/`allOf`/`not`, `if`/`then`/`else`, or +`$ref` makes the tool invalid, and clients drop it from `tools/list`. + +With `schemars`, a nested struct is emitted behind a `$ref` by default. Mark it +`#[schemars(inline)]` so its annotated fields stay on the `properties` path: + +```rust, ignore +#[derive(serde::Deserialize, schemars::JsonSchema)] +#[schemars(inline)] +struct Location { + #[schemars(extend("x-mcp-header" = "Region"))] + region: String, +} + +#[derive(serde::Deserialize, schemars::JsonSchema)] +struct DeployRequest { + location: Location, + image: String, +} +``` + Values that cannot travel as a bare header (leading/trailing whitespace, control/non-ASCII characters) are transparently Base64-wrapped as `=?base64??=`. diff --git a/crates/rmcp/src/transport/common/mcp_headers.rs b/crates/rmcp/src/transport/common/mcp_headers.rs index 54e99e535..e638c87e2 100644 --- a/crates/rmcp/src/transport/common/mcp_headers.rs +++ b/crates/rmcp/src/transport/common/mcp_headers.rs @@ -99,38 +99,78 @@ fn is_tchar(c: char) -> bool { ) } -/// Top-level properties carrying an `x-mcp-header` annotation, as `(property, header)` pairs. -fn param_header_annotations(input_schema: &JsonObject) -> Vec<(String, String)> { +/// An `x-mcp-header` annotation on a property reached from the schema root +/// through `properties` keys alone. +struct PropertyAnnotation<'a> { + /// Property names from the root to the annotated property. + path: Vec<&'a str>, + schema: &'a Value, + header: &'a Value, +} + +/// Every `x-mcp-header` annotation along `properties` chains, nested ones included. +fn property_annotations(input_schema: &JsonObject) -> Vec> { let mut out = Vec::new(); - if let Some(Value::Object(props)) = input_schema.get("properties") { - for (prop, schema) in props { - if let Some(Value::String(header)) = schema.get("x-mcp-header") - && !header.is_empty() - { - out.push((prop.clone(), header.clone())); - } + collect_property_annotations(input_schema.get("properties"), &mut Vec::new(), &mut out); + out +} + +fn collect_property_annotations<'a>( + properties: Option<&'a Value>, + path: &mut Vec<&'a str>, + out: &mut Vec>, +) { + let Some(Value::Object(properties)) = properties else { + return; + }; + for (name, schema) in properties { + path.push(name); + if let Some(header) = schema.get("x-mcp-header") { + out.push(PropertyAnnotation { + path: path.clone(), + schema, + header, + }); } + collect_property_annotations(schema.get("properties"), path, out); + path.pop(); } - out +} + +/// Well-formed annotations as `(path, header)` pairs, for header emission and checks. +fn param_header_annotations(input_schema: &JsonObject) -> Vec<(Vec<&str>, &str)> { + property_annotations(input_schema) + .into_iter() + .filter_map(|annotation| match annotation.header { + Value::String(header) if !header.is_empty() => Some((annotation.path, header.as_str())), + _ => None, + }) + .collect() +} + +/// The instance value at `path` in `arguments`, if every step is present. +fn value_at<'a>(arguments: &'a Value, path: &[&str]) -> Option<&'a Value> { + path.iter().try_fold(arguments, |value, key| value.get(key)) } /// Validates the `x-mcp-header` annotations in a tool input schema. /// /// Annotations must be non-empty RFC 9110 tokens, case-insensitively unique, -/// applied only to top-level primitive (`string`/`integer`/`boolean`) properties. -/// Returns the offending reason on the first violation. +/// applied only to primitive (`string`/`integer`/`boolean`) properties reached +/// from the root through `properties` keys alone. Returns the offending reason +/// on the first violation. #[cfg(feature = "client-side-sse")] pub(crate) fn validate_param_header_annotations(input_schema: &JsonObject) -> Result<(), String> { - let Some(Value::Object(props)) = input_schema.get("properties") else { - return Ok(()); - }; + if input_schema.contains_key("x-mcp-header") { + return Err("x-mcp-header must annotate a property, not the schema root".to_owned()); + } + if has_off_chain_annotation(input_schema, true) { + return Err("x-mcp-header must sit on a `properties`-only path from the root".to_owned()); + } let mut seen = std::collections::HashSet::new(); - for (prop, schema) in props { - reject_nested_annotations(schema, prop)?; - let Some(raw) = schema.get("x-mcp-header") else { - continue; - }; - let Value::String(header) = raw else { + for annotation in property_annotations(input_schema) { + let prop = annotation.path.join("."); + let Value::String(header) = annotation.header else { return Err(format!("property `{prop}`: x-mcp-header must be a string")); }; if header.is_empty() { @@ -146,7 +186,7 @@ pub(crate) fn validate_param_header_annotations(input_schema: &JsonObject) -> Re "property `{prop}`: duplicate x-mcp-header `{header}` (case-insensitive)" )); } - match schema.get("type").and_then(Value::as_str) { + match annotation.schema.get("type").and_then(Value::as_str) { Some("string" | "integer" | "boolean") => {} other => { return Err(format!( @@ -159,20 +199,42 @@ pub(crate) fn validate_param_header_annotations(input_schema: &JsonObject) -> Re Ok(()) } -/// Rejects `x-mcp-header` on nested properties (only top-level promotion is supported). +/// True if `schema` holds an `x-mcp-header` off the `properties` chain. +/// +/// `on_chain` says whether `schema` is the root or a property reached from it +/// through `properties` keys alone. #[cfg(feature = "client-side-sse")] -fn reject_nested_annotations(schema: &Value, path: &str) -> Result<(), String> { - if let Some(Value::Object(nested)) = schema.get("properties") { - for (key, value) in nested { - if value.get("x-mcp-header").is_some() { - return Err(format!( - "property `{path}.{key}`: x-mcp-header is not supported on nested properties" - )); - } - reject_nested_annotations(value, &format!("{path}.{key}"))?; +fn has_off_chain_annotation(schema: &JsonObject, on_chain: bool) -> bool { + schema.iter().any(|(key, value)| match key.as_str() { + "x-mcp-header" => !on_chain, + "properties" => schema_map_has_off_chain_annotation(value, on_chain), + "$defs" | "definitions" | "dependentSchemas" | "patternProperties" => { + schema_map_has_off_chain_annotation(value, false) } + // Instance data, not subschemas. + "default" | "const" | "enum" | "examples" => false, + _ => subschema_has_off_chain_annotation(value, false), + }) +} + +/// Like [`has_off_chain_annotation`] for each value of a name-to-schema map. +#[cfg(feature = "client-side-sse")] +fn schema_map_has_off_chain_annotation(map: &Value, on_chain: bool) -> bool { + map.as_object().is_some_and(|map| { + map.values() + .any(|schema| subschema_has_off_chain_annotation(schema, on_chain)) + }) +} + +#[cfg(feature = "client-side-sse")] +fn subschema_has_off_chain_annotation(value: &Value, on_chain: bool) -> bool { + match value { + Value::Object(schema) => has_off_chain_annotation(schema, on_chain), + Value::Array(items) => items + .iter() + .any(|item| subschema_has_off_chain_annotation(item, false)), + _ => false, } - Ok(()) } /// Wraps a value as `=?base64??=` when it cannot travel as a bare header value. @@ -240,8 +302,8 @@ pub(crate) fn standard_request_headers( && let (Some(schema), Some(arguments)) = (tool_schema, params.and_then(|p| p.get("arguments"))) { - for (prop, header) in param_header_annotations(schema) { - let Some(arg) = arguments.get(&prop) else { + for (path, header) in param_header_annotations(schema) { + let Some(arg) = value_at(arguments, &path) else { continue; }; let Some(encoded) = primitive_to_string(arg).map(|s| encode_header_value(&s)) else { @@ -297,10 +359,11 @@ pub(crate) fn validate_request_headers( && let Some(schema) = tool_schema { let arguments = params.and_then(|p| p.get("arguments")); - for (prop, header) in param_header_annotations(schema) { + for (path, header) in param_header_annotations(schema) { + let prop = path.join("."); let full = format!("{HEADER_MCP_PARAM_PREFIX}{header}"); let header_value = header_str(headers, &full)?; - let arg = arguments.and_then(|a| a.get(&prop)); + let arg = arguments.and_then(|a| value_at(a, &path)); let body_value = arg.filter(|v| !v.is_null()).and_then(primitive_to_string); match (header_value, body_value) { @@ -537,15 +600,176 @@ mod tests { } #[test] - fn rejects_nested_annotation() { + fn accepts_nested_annotation() { let schema = schema_with(json!({ "outer": { "type": "object", "properties": { "inner": { "type": "string", "x-mcp-header": "Inner" } } } })); + assert!(validate_param_header_annotations(&schema).is_ok()); + } + + #[test] + fn rejects_nested_number_type() { + let schema = schema_with(json!({ + "outer": { + "type": "object", + "properties": { "inner": { "type": "number", "x-mcp-header": "Inner" } } + } + })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_duplicate_across_nesting_levels() { + let schema = schema_with(json!({ + "region": { "type": "string", "x-mcp-header": "Region" }, + "outer": { + "type": "object", + "properties": { "inner": { "type": "string", "x-mcp-header": "region" } } + } + })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_annotation_under_items() { + let schema = schema_with(json!({ + "list": { + "type": "array", + "items": { "type": "string", "x-mcp-header": "Item" } + } + })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_annotation_under_properties_of_items() { + let schema = schema_with(json!({ + "list": { + "type": "array", + "items": { + "type": "object", + "properties": { "id": { "type": "string", "x-mcp-header": "Id" } } + } + } + })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_annotation_under_composition() { + let schema = schema_with(json!({ + "choice": { + "oneOf": [{ "type": "string", "x-mcp-header": "Choice" }] + } + })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_annotation_under_conditional() { + let schema = schema_with(json!({ + "value": { + "type": "object", + "then": { + "properties": { "x": { "type": "string", "x-mcp-header": "X" } } + } + } + })); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_annotation_in_ref_target() { + let schema = json!({ + "type": "object", + "properties": { "region": { "$ref": "#/$defs/region" } }, + "$defs": { "region": { "type": "string", "x-mcp-header": "Region" } } + }) + .as_object() + .unwrap() + .clone(); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_annotation_in_def_named_like_a_keyword() { + let schema = json!({ + "type": "object", + "$defs": { "default": { "type": "string", "x-mcp-header": "Region" } } + }) + .as_object() + .unwrap() + .clone(); + assert!(validate_param_header_annotations(&schema).is_err()); + } + + #[test] + fn rejects_annotation_on_root() { + let schema = json!({ "type": "object", "x-mcp-header": "Root" }) + .as_object() + .unwrap() + .clone(); assert!(validate_param_header_annotations(&schema).is_err()); } + + #[test] + fn ignores_annotation_key_in_default_value() { + let schema = schema_with(json!({ + "config": { + "type": "object", + "default": { "x-mcp-header": "not-a-schema" } + } + })); + assert!(validate_param_header_annotations(&schema).is_ok()); + } + } + + #[cfg(feature = "server")] + mod derived_schema { + use schemars::JsonSchema; + + use super::*; + use crate::handler::server::common::schema_for_type; + + #[derive(JsonSchema)] + #[expect(dead_code)] + struct RefLocation { + #[schemars(extend("x-mcp-header" = "Region"))] + region: String, + } + + #[derive(JsonSchema)] + #[expect(dead_code)] + struct RefDeploy { + location: RefLocation, + } + + #[derive(JsonSchema)] + #[schemars(inline)] + #[expect(dead_code)] + struct InlineLocation { + #[schemars(extend("x-mcp-header" = "Region"))] + region: String, + } + + #[derive(JsonSchema)] + #[expect(dead_code)] + struct InlineDeploy { + location: InlineLocation, + } + + #[test] + fn rejects_nested_struct_behind_ref() { + assert!(validate_param_header_annotations(&schema_for_type::()).is_err()); + } + + #[test] + fn accepts_inlined_nested_struct() { + assert!(validate_param_header_annotations(&schema_for_type::()).is_ok()); + } } mod standard_request_headers { @@ -590,6 +814,70 @@ mod tests { Some("us-west1") ); } + + fn nested_call_headers(arguments: Value) -> HashMap { + let schema = schema_with(json!({ + "location": { + "type": "object", + "properties": { "region": { "type": "string", "x-mcp-header": "Region" } } + } + })); + let request = json!({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { "name": "deploy", "arguments": arguments } + }); + super::super::standard_request_headers(&request, Some(&schema)) + .into_iter() + .map(|(name, value)| (name.as_str().to_owned(), value.to_str().unwrap().to_owned())) + .collect() + } + + #[test] + fn sets_nested_param_header() { + let headers = nested_call_headers(json!({ "location": { "region": "eu" } })); + assert_eq!( + headers.get("mcp-param-region").map(String::as_str), + Some("eu") + ); + } + + #[test] + fn sets_nested_param_header_after_sibling_property() { + let schema = schema_with(json!({ + "image": { "type": "string" }, + "location": { + "type": "object", + "properties": { "region": { "type": "string", "x-mcp-header": "Region" } } + } + })); + let request = json!({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { + "name": "deploy", + "arguments": { "image": "app:1", "location": { "region": "eu" } } + } + }); + let headers = super::super::standard_request_headers(&request, Some(&schema)); + assert_eq!( + headers + .iter() + .find(|(name, _)| name == "mcp-param-region") + .map(|(_, value)| value.to_str().unwrap()), + Some("eu") + ); + } + + #[test] + fn omits_nested_param_header_when_parent_absent() { + let headers = nested_call_headers(json!({})); + assert_eq!(headers.get("mcp-param-region"), None); + } + + #[test] + fn omits_param_header_for_same_name_at_wrong_depth() { + let headers = nested_call_headers(json!({ "region": "eu" })); + assert_eq!(headers.get("mcp-param-region"), None); + } } mod validate_request_headers { @@ -704,6 +992,44 @@ mod tests { assert!(validate_request_headers(&headers, &request, Some(&schema)).is_err()); } + fn nested_schema() -> JsonObject { + schema_with(json!({ + "location": { + "type": "object", + "properties": { "region": { "type": "string", "x-mcp-header": "Region" } } + } + })) + } + + fn nested_request() -> Value { + json!({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { "name": "deploy", "arguments": { "location": { "region": "eu" } } } + }) + } + + #[test] + fn accepts_matching_nested_param() { + let headers = header_map(&[ + ("Mcp-Method", "tools/call"), + ("Mcp-Name", "deploy"), + ("Mcp-Param-Region", "eu"), + ]); + assert!( + validate_request_headers(&headers, &nested_request(), Some(&nested_schema())) + .is_ok() + ); + } + + #[test] + fn rejects_missing_nested_param() { + let headers = header_map(&[("Mcp-Method", "tools/call"), ("Mcp-Name", "deploy")]); + assert_eq!( + validate_request_headers(&headers, &nested_request(), Some(&nested_schema())), + Err("missing Mcp-Param-Region header for `location.region`".to_owned()) + ); + } + #[test] fn rejects_duplicate_param_for_absent_argument() { let schema = schema_with(json!({ diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 9f1a4bcdb..ae946a33f 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -2554,6 +2554,46 @@ mod tests { ); } + #[test] + fn cache_tools_keeps_nested_annotations_and_drops_array_annotations() { + let nested = tool( + "nested", + json!({ + "type": "object", + "properties": { "region": { "type": "string", "x-mcp-header": "Region" } }, + }), + ); + let in_items = tool( + "in_items", + json!({ "type": "array", "items": { "type": "string", "x-mcp-header": "Item" } }), + ); + let mut message = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::with_all_items(vec![nested, in_items])), + NumberOrString::Number(1), + ); + let mut cache = HashMap::new(); + + cache_tools_from_response(&mut cache, &mut message, &ProtocolVersion::V_2026_07_28); + + let ServerJsonRpcMessage::Response(response) = &mut message else { + panic!("expected tools/list response"); + }; + let ServerResult::ListToolsResult(result) = &mut response.result else { + panic!("expected tools/list result"); + }; + assert_eq!( + ( + result + .tools + .iter() + .map(|tool| tool.name.as_ref()) + .collect::>(), + cache.keys().map(String::as_str).collect::>(), + ), + (vec!["nested"], vec!["nested"]) + ); + } + #[test] fn cache_tools_preserves_pre_standard_header_results() { let invalid = tool("legacy", json!({ "type": "string", "x-mcp-header": "" }));