From daf717c129758b5dc2c7e86bdab7ea6419d57f26 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Fri, 21 Aug 2026 11:20:43 +0100 Subject: [PATCH] feat: allow externally supplied tool schemas Signed-off-by: lucarlig --- .../src/transport/streamable_http_client.rs | 70 +++++++++++-- .../transport/streamable_http_server/tower.rs | 47 +++++++-- .../test_streamable_http_standard_headers.rs | 99 +++++++++++++++++++ 3 files changed, 203 insertions(+), 13 deletions(-) diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index fe563ee3f..1b9557a50 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -839,6 +839,17 @@ impl Worker for StreamableHttpClientWorker { let (sse_worker_tx, mut sse_worker_rx) = tokio::sync::mpsc::channel::(channel_buffer_capacity); let config = self.config.clone(); + // SEP-2243: tool input schemas (name -> schema) used to promote annotated + // tools/call arguments to Mcp-Param-* headers. Callers may seed the cache + // when another component owns tool discovery; tools/list responses refresh it. + let mut tool_header_cache = config.tool_schemas.clone(); + tool_header_cache.retain(|name, schema| { + let Err(reason) = mcp_headers::validate_param_header_annotations(schema) else { + return true; + }; + tracing::warn!(tool = %name, "ignoring configured schema with invalid x-mcp-header annotations: {reason}"); + false + }); let transport_task_ct = context.cancellation_token.clone(); let _drop_guard = transport_task_ct.clone().drop_guard(); let WorkerSendRequest { @@ -851,7 +862,6 @@ impl Worker for StreamableHttpClientWorker { if matches!(&request.request, ClientRequest::InitializeRequest(_)) ); let mut saved_init_request = is_legacy_startup.then(|| startup_request.clone()); - let empty_tool_cache = HashMap::new(); let (bootstrap_version, bootstrap_headers) = if is_legacy_startup { (ProtocolVersion::default(), config.custom_headers.clone()) } else { @@ -859,7 +869,7 @@ impl Worker for StreamableHttpClientWorker { &config.custom_headers, &startup_request, &ProtocolVersion::default(), - &empty_tool_cache, + &tool_header_cache, ) }; let (message, session_id) = match self @@ -908,10 +918,6 @@ impl Worker for StreamableHttpClientWorker { } else { (bootstrap_version, bootstrap_headers) }; - // SEP-2243: tool input schemas (name -> schema) cached from tools/list responses, - // used to promote annotated tools/call arguments to Mcp-Param-* headers. - let mut tool_header_cache: HashMap> = HashMap::new(); - // Store session info for cleanup when run() exits (not spawned, so cleanup completes before close() returns) let mut session_cleanup_info = session_id.as_ref().map(|sid| SessionCleanupInfo { client: self.client.clone(), @@ -1684,6 +1690,12 @@ pub struct StreamableHttpClientTransportConfig { pub auth_header: Option, /// Custom HTTP headers to include with every request pub custom_headers: HashMap, + /// Tool input schemas used to construct SEP-2243 `Mcp-Param-*` headers. + /// + /// This lets callers seed schemas learned outside this transport, so the + /// client does not need to issue `tools/list` before `tools/call`. + /// Schemas returned by later `tools/list` responses replace matching entries. + pub tool_schemas: HashMap>, /// Maximum raw size of one SSE event accepted from the server. /// /// The built-in reqwest and Unix socket clients enforce this value. Custom @@ -1748,6 +1760,12 @@ impl StreamableHttpClientTransportConfig { self } + /// Set tool input schemas used to construct SEP-2243 `Mcp-Param-*` headers. + pub fn tool_schemas(mut self, tool_schemas: HashMap>) -> Self { + self.tool_schemas = tool_schemas; + self + } + /// Set the maximum raw size of one SSE event accepted from the server. pub fn max_sse_event_size(mut self, bytes: usize) -> Self { self.max_sse_event_size = bytes; @@ -1777,6 +1795,7 @@ impl Default for StreamableHttpClientTransportConfig { allow_stateless: true, auth_header: None, custom_headers: HashMap::new(), + tool_schemas: HashMap::new(), max_sse_event_size: DEFAULT_MAX_SSE_EVENT_SIZE, reinit_on_expired_session: true, } @@ -2088,6 +2107,45 @@ mod tests { ) } + #[test] + fn configured_tool_schema_builds_param_headers_without_tools_list() { + let schema = json!({ + "type": "object", + "properties": { + "region": { "type": "string", "x-mcp-header": "Region" }, + }, + }); + let config = StreamableHttpClientTransportConfig::with_uri("http://localhost/mcp") + .tool_schemas(HashMap::from([( + "deploy".to_owned(), + Arc::new(schema.as_object().expect("object schema").clone()), + )])); + let message: ClientJsonRpcMessage = serde_json::from_value(json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "deploy", + "arguments": { "region": "us-west1" }, + }, + })) + .expect("tools/call message"); + + let headers = build_request_headers( + &HashMap::new(), + &message, + &config.tool_schemas, + &ProtocolVersion::V_2026_07_28, + ); + + assert_eq!( + headers + .get(&HeaderName::from_static("mcp-param-region")) + .expect("configured schema creates parameter header"), + &HeaderValue::from_static("us-west1") + ); + } + #[test] fn cache_tools_removes_invalid_header_annotations() { let valid = tool( diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index f1fef585f..62ecea558 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -55,6 +55,9 @@ use crate::{ pub(crate) const DEFAULT_MAX_REQUEST_BODY_BYTES: usize = 4 * 1024 * 1024; const STATELESS_STREAM_CHANNEL_CAPACITY: usize = 16; +type ToolSchemaResolver = + dyn Fn(&http::request::Parts, &str) -> Option> + Send + Sync; + #[non_exhaustive] #[derive(Debug, Clone)] pub struct StreamableHttpServerConfig { @@ -1011,6 +1014,9 @@ pub struct StreamableHttpService { /// Populated lazily via `get_tool` so the service factory runs at most once /// per tool name. `None` value means the tool exposes no schema. tool_schemas: Arc>>>>, + /// Optional request-aware source for tool schemas. When configured, this is + /// authoritative and replaces the service-wide `get_tool` cache. + tool_schema_resolver: Option>, } impl Clone for StreamableHttpService { @@ -1021,6 +1027,7 @@ impl Clone for StreamableHttpService { service_factory: self.service_factory.clone(), pending_restores: self.pending_restores.clone(), tool_schemas: self.tool_schemas.clone(), + tool_schema_resolver: self.tool_schema_resolver.clone(), } } } @@ -1102,8 +1109,27 @@ where service_factory: Arc::new(service_factory), pending_restores, tool_schemas: Arc::new(std::sync::RwLock::new(HashMap::new())), + tool_schema_resolver: None, } } + + /// Use a request-aware tool schema source for SEP-2243 validation. + /// + /// The resolver receives the HTTP request parts and the `tools/call` name. + /// It is useful when routing or authorization state in request extensions + /// determines which tool definition applies. Once configured, the resolver + /// is authoritative; `None` does not fall back to the service-wide cache. + pub fn with_tool_schema_resolver( + mut self, + resolver: impl Fn(&http::request::Parts, &str) -> Option> + + Send + + Sync + + 'static, + ) -> Self { + self.tool_schema_resolver = Some(Arc::new(resolver)); + self + } + fn get_service(&self) -> Result { (self.service_factory)() } @@ -1267,10 +1293,13 @@ where Ok(self.stateless_sse_response(Some(first), receiver, request_ct)) } - /// Returns the cached input schema for `name`, constructing a service once - /// per name to read its `ServerHandler::get_tool` definition. Used to - /// validate SEP-2243 `Mcp-Param-*` headers against the request body. - fn tool_schema(&self, name: &str) -> Option> { + /// Returns the input schema for `name` in this request. A configured + /// request-aware resolver is authoritative; otherwise the schema is cached + /// by name from `ServerHandler::get_tool`. + fn tool_schema(&self, parts: &http::request::Parts, name: &str) -> Option> { + if let Some(resolver) = &self.tool_schema_resolver { + return resolver(parts, name); + } if let Ok(cache) = self.tool_schemas.read() { if let Some(schema) = cache.get(name) { return schema.clone(); @@ -1759,7 +1788,9 @@ where validate_protocol_version_header(&part.headers, has_per_request_version)?; validate_request_protocol_version_meta(&part.headers, &message)?; // Validate SEP-2243 standard headers against the body - validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?; + validate_standard_headers(&part.headers, &message, |name| { + self.tool_schema(&part, name) + })?; // inject request part to extensions match &mut message { @@ -1812,7 +1843,7 @@ where message_has_per_request_protocol_version(&message), )?; validate_standard_headers(&part.headers, &message, |name| { - self.tool_schema(name) + self.tool_schema(&part, name) })?; validate_request_protocol_version_meta(&part.headers, &message)?; let ClientJsonRpcMessage::Request(request) = message else { @@ -1941,7 +1972,9 @@ where } } // Validate SEP-2243 standard headers against the body - validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?; + validate_standard_headers(&part.headers, &message, |name| { + self.tool_schema(&part, name) + })?; validate_request_protocol_version_meta(&part.headers, &message)?; validate_required_protocol_meta(&self.config, &message)?; let service = self diff --git a/crates/rmcp/tests/test_streamable_http_standard_headers.rs b/crates/rmcp/tests/test_streamable_http_standard_headers.rs index 510433e70..8954b6a2d 100644 --- a/crates/rmcp/tests/test_streamable_http_standard_headers.rs +++ b/crates/rmcp/tests/test_streamable_http_standard_headers.rs @@ -17,6 +17,12 @@ const SEP_VERSION: &str = "2026-07-28"; #[derive(Clone, Default)] struct HeaderValidationServer; +#[derive(Clone, Copy)] +enum TenantSchema { + Region, + Location, +} + impl ServerHandler for HeaderValidationServer { fn get_info(&self) -> ServerInfo { ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) @@ -59,6 +65,59 @@ async fn spawn_server() -> (reqwest::Client, String, CancellationToken) { (reqwest::Client::new(), format!("http://{addr}/mcp"), ct) } +async fn spawn_request_aware_server() -> (reqwest::Client, [String; 2], CancellationToken) { + let config = StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(true) + .with_sse_keep_alive(None) + .with_cancellation_token(CancellationToken::new()); + let ct = config.cancellation_token.clone(); + let service: StreamableHttpService = + StreamableHttpService::new(|| Ok(HeaderValidationServer), Default::default(), config) + .with_tool_schema_resolver(|parts, name| { + if name != "deploy" { + return None; + } + let header = match parts.extensions.get::()? { + TenantSchema::Region => "Region", + TenantSchema::Location => "Location", + }; + let schema = serde_json::json!({ + "type": "object", + "properties": { + "region": { "type": "string", "x-mcp-header": header }, + }, + }); + Some(Arc::new(schema.as_object().expect("object schema").clone())) + }); + + let region_router = axum::Router::new() + .nest_service("/region/mcp", service.clone()) + .layer(axum::Extension(TenantSchema::Region)); + let location_router = axum::Router::new() + .nest_service("/location/mcp", service) + .layer(axum::Extension(TenantSchema::Location)); + let router = region_router.merge(location_router); + let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = tcp_listener.local_addr().unwrap(); + tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(tcp_listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + ( + reqwest::Client::new(), + [ + format!("http://{addr}/region/mcp"), + format!("http://{addr}/location/mcp"), + ], + ct, + ) +} + /// POSTs a `tools/call` with the given protocol-version and optional SEP-2243 headers. async fn post_tool_call( client: &reqwest::Client, @@ -313,3 +372,43 @@ async fn rejects_missing_param_header_with_32020() -> anyhow::Result<()> { ct.cancel(); Ok(()) } + +#[tokio::test] +async fn request_aware_schema_resolver_isolates_same_named_tools() -> anyhow::Result<()> { + let (client, [region_url, location_url], ct) = spawn_request_aware_server().await; + + let region_response = post_tool_call( + &client, + ®ion_url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + Some("us-west1"), + ) + .await; + let region_body: serde_json::Value = region_response.json().await?; + assert_ne!( + region_body["error"]["code"], -32020, + "the region tenant must use its request-scoped schema: {region_body}" + ); + + let location_response = post_tool_call( + &client, + &location_url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + Some("us-west1"), + ) + .await; + assert_eq!(location_response.status(), 400); + let location_body: serde_json::Value = location_response.json().await?; + assert_eq!(location_body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +}