From 0a434e57d15892e1a67e4a8764086ae072703e7b Mon Sep 17 00:00:00 2001 From: Dale Seo <5466341+DaleSeo@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:04:30 -0400 Subject: [PATCH 1/2] fix(streamable-http-server): look up tool schemas per request --- .../transport/streamable_http_server/tower.rs | 38 +++------- .../test_streamable_http_standard_headers.rs | 74 ++++++++++++++++++- 2 files changed, 82 insertions(+), 30 deletions(-) diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index ab541ff31..b44960581 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -1338,10 +1338,6 @@ pub struct StreamableHttpService { /// than racing to replay the initialize handshake. `None` when no external /// session store is configured (avoids allocating the map). pending_restores: Option, - /// Caches tool input schemas by name for SEP-2243 `Mcp-Param-*` validation. - /// 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>>>>, } impl Clone for StreamableHttpService { @@ -1351,7 +1347,6 @@ impl Clone for StreamableHttpService { session_manager: self.session_manager.clone(), service_factory: self.service_factory.clone(), pending_restores: self.pending_restores.clone(), - tool_schemas: self.tool_schemas.clone(), } } } @@ -1430,7 +1425,6 @@ where session_manager, service_factory: Arc::new(service_factory), pending_restores, - tool_schemas: Arc::new(std::sync::RwLock::new(HashMap::new())), } } fn get_service(&self) -> Result { @@ -1599,24 +1593,12 @@ 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. + /// Looks up `name`'s input schema for SEP-2243 `Mcp-Param-*` validation. fn tool_schema(&self, name: &str) -> Option> { - if let Ok(cache) = self.tool_schemas.read() - && let Some(schema) = cache.get(name) - { - return schema.clone(); - } - let schema = self - .get_service() - .ok() - .and_then(|service| service.get_tool(name)) - .map(|tool| tool.input_schema); - if let Ok(mut cache) = self.tool_schemas.write() { - cache.insert(name.to_owned(), schema.clone()); - } - schema + self.get_service() + .ok()? + .get_tool(name) + .map(|tool| tool.input_schema) } /// Spawn a task that runs `serve_server` for the given session, waits for @@ -2269,13 +2251,15 @@ where validate_protocol_version_header(&part.headers, has_per_request_version)?; } } - // Validate SEP-2243 standard headers against the body - validate_standard_headers(&part.headers, &message, |name| self.tool_schema(name))?; - validate_request_protocol_version_meta(&part.headers, &message)?; - validate_required_protocol_meta(&self.config, &message)?; let service = self .get_service() .map_err(internal_error_response("get service"))?; + // Validate SEP-2243 standard headers against the body + validate_standard_headers(&part.headers, &message, |name| { + service.get_tool(name).map(|tool| tool.input_schema) + })?; + validate_request_protocol_version_meta(&part.headers, &message)?; + validate_required_protocol_meta(&self.config, &message)?; match message { ClientJsonRpcMessage::Request(mut request) => { let negotiates_per_request = has_per_request_version diff --git a/crates/rmcp/tests/test_streamable_http_standard_headers.rs b/crates/rmcp/tests/test_streamable_http_standard_headers.rs index 5bbfedd64..d5344588e 100644 --- a/crates/rmcp/tests/test_streamable_http_standard_headers.rs +++ b/crates/rmcp/tests/test_streamable_http_standard_headers.rs @@ -1,6 +1,9 @@ #![cfg(not(feature = "local"))] //! SEP-2243 server-side validation of `Mcp-Method` / `Mcp-Name` / `Mcp-Param-*` headers. -use std::sync::Arc; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; use rmcp::{ ServerHandler, @@ -36,14 +39,23 @@ impl ServerHandler for HeaderValidationServer { } async fn spawn_server() -> (reqwest::Client, String, CancellationToken) { + spawn_server_with(|| Ok(HeaderValidationServer)).await +} + +async fn spawn_server_with( + service_factory: impl Fn() -> Result + Send + Sync + 'static, +) -> (reqwest::Client, String, CancellationToken) +where + S: ServerHandler + Send + 'static, +{ 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); + let service: StreamableHttpService = + StreamableHttpService::new(service_factory, Default::default(), config); let router = axum::Router::new().nest_service("/mcp", service); let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -314,6 +326,62 @@ async fn rejects_missing_param_header_with_32020() -> anyhow::Result<()> { Ok(()) } +/// Server whose `deploy` tool gains the `x-mcp-header` annotation at runtime. +#[derive(Clone, Default)] +struct ChangingToolServer { + annotated: Arc, +} + +impl ServerHandler for ChangingToolServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + fn get_tool(&self, name: &str) -> Option { + if !self.annotated.load(Ordering::SeqCst) { + return None; + } + HeaderValidationServer.get_tool(name) + } +} + +#[tokio::test] +async fn validates_param_header_against_current_tool_schema() -> anyhow::Result<()> { + let server = ChangingToolServer::default(); + let (client, url, ct) = spawn_server_with({ + let server = server.clone(); + move || Ok(server.clone()) + }) + .await; + let send_mismatch = || { + post_tool_call( + &client, + &url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + Some("eu-central1"), + ) + }; + + let before: serde_json::Value = send_mismatch().await.json().await?; + assert_ne!( + before["error"]["code"], -32020, + "unannotated tool must not validate Mcp-Param-*, got: {before}" + ); + + server.annotated.store(true, Ordering::SeqCst); + let response = send_mismatch().await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + ct.cancel(); + Ok(()) +} + /// Spawns the server in legacy session mode: the initialize handshake creates /// a real session, exercising the path where the validator used to be bypassed. async fn spawn_legacy_server() -> (reqwest::Client, String, CancellationToken) { From 46108493c3b061a50bfa7c81a0966f8dbc155c50 Mon Sep 17 00:00:00 2001 From: Dale Seo <5466341+DaleSeo@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:33:00 -0400 Subject: [PATCH 2/2] test: cover stale tool schemas in header tests --- .../test_streamable_http_standard_headers.rs | 78 +++++++++++++------ 1 file changed, 53 insertions(+), 25 deletions(-) diff --git a/crates/rmcp/tests/test_streamable_http_standard_headers.rs b/crates/rmcp/tests/test_streamable_http_standard_headers.rs index d5344588e..e531ac00f 100644 --- a/crates/rmcp/tests/test_streamable_http_standard_headers.rs +++ b/crates/rmcp/tests/test_streamable_http_standard_headers.rs @@ -326,12 +326,34 @@ async fn rejects_missing_param_header_with_32020() -> anyhow::Result<()> { Ok(()) } -/// Server whose `deploy` tool gains the `x-mcp-header` annotation at runtime. +/// Server whose `deploy` tool gains or loses the `x-mcp-header` annotation at runtime. #[derive(Clone, Default)] struct ChangingToolServer { annotated: Arc, } +async fn spawn_changing_tool_server( + server: &ChangingToolServer, +) -> (reqwest::Client, String, CancellationToken) { + let server = server.clone(); + spawn_server_with(move || Ok(server.clone())).await +} + +/// POSTs a `deploy` call whose `Mcp-Param-Region` header contradicts the body. +async fn post_mismatched_region(client: &reqwest::Client, url: &str) -> reqwest::Response { + post_tool_call( + client, + url, + SEP_VERSION, + "deploy", + serde_json::json!({ "region": "us-west1" }), + Some("tools/call"), + Some("deploy"), + Some("eu-central1"), + ) + .await +} + impl ServerHandler for ChangingToolServer { fn get_info(&self) -> ServerConfig { ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) @@ -346,34 +368,18 @@ impl ServerHandler for ChangingToolServer { } #[tokio::test] -async fn validates_param_header_against_current_tool_schema() -> anyhow::Result<()> { +async fn validates_param_header_once_tool_gains_annotation() -> anyhow::Result<()> { let server = ChangingToolServer::default(); - let (client, url, ct) = spawn_server_with({ - let server = server.clone(); - move || Ok(server.clone()) - }) - .await; - let send_mismatch = || { - post_tool_call( - &client, - &url, - SEP_VERSION, - "deploy", - serde_json::json!({ "region": "us-west1" }), - Some("tools/call"), - Some("deploy"), - Some("eu-central1"), - ) - }; + let (client, url, ct) = spawn_changing_tool_server(&server).await; - let before: serde_json::Value = send_mismatch().await.json().await?; - assert_ne!( - before["error"]["code"], -32020, - "unannotated tool must not validate Mcp-Param-*, got: {before}" - ); + // Header validation passes, so the request fails later on its missing `_meta`. + let response = post_mismatched_region(&client, &url).await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32602); server.annotated.store(true, Ordering::SeqCst); - let response = send_mismatch().await; + let response = post_mismatched_region(&client, &url).await; assert_eq!(response.status(), 400); let body: serde_json::Value = response.json().await?; assert_eq!(body["error"]["code"], -32020); @@ -382,6 +388,28 @@ async fn validates_param_header_against_current_tool_schema() -> anyhow::Result< Ok(()) } +#[tokio::test] +async fn skips_param_header_validation_once_tool_loses_annotation() -> anyhow::Result<()> { + let server = ChangingToolServer { + annotated: Arc::new(AtomicBool::new(true)), + }; + let (client, url, ct) = spawn_changing_tool_server(&server).await; + + let response = post_mismatched_region(&client, &url).await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32020); + + server.annotated.store(false, Ordering::SeqCst); + let response = post_mismatched_region(&client, &url).await; + assert_eq!(response.status(), 400); + let body: serde_json::Value = response.json().await?; + assert_eq!(body["error"]["code"], -32602); + + ct.cancel(); + Ok(()) +} + /// Spawns the server in legacy session mode: the initialize handshake creates /// a real session, exercising the path where the validator used to be bypassed. async fn spawn_legacy_server() -> (reqwest::Client, String, CancellationToken) {