fix(streamable-http-server): look up tool schemas per request

crates/rmcp/src/transport/streamable_http_server/tower.rs

@@ -1338,10 +1338,6 @@ pub struct StreamableHttpService<S, M> {
     /// than racing to replay the initialize handshake. `None` when no external
     /// session store is configured (avoids allocating the map).
     pending_restores: Option<PendingRestores>,
-    /// 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<std::sync::RwLock<HashMap<String, Option<Arc<JsonObject>>>>>,
 }
 
 impl<S, M> Clone for StreamableHttpService<S, M> {
@@ -1351,7 +1347,6 @@ impl<S, M> Clone for StreamableHttpService<S, M> {
             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<S, std::io::Error> {
@@ -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<Arc<JsonObject>> {
-        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

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<S>(
+    service_factory: impl Fn() -> Result<S, std::io::Error> + 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<HeaderValidationServer, LocalSessionManager> =
-        StreamableHttpService::new(|| Ok(HeaderValidationServer), Default::default(), config);
+    let service: StreamableHttpService<S, LocalSessionManager> =
+        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,90 @@ async fn rejects_missing_param_header_with_32020() -> anyhow::Result<()> {
     Ok(())
 }
 
+/// Server whose `deploy` tool gains or loses the `x-mcp-header` annotation at runtime.
+#[derive(Clone, Default)]
+struct ChangingToolServer {
+    annotated: Arc<AtomicBool>,
+}
+
+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())
+    }
+
+    fn get_tool(&self, name: &str) -> Option<Tool> {
+        if !self.annotated.load(Ordering::SeqCst) {
+            return None;
+        }
+        HeaderValidationServer.get_tool(name)
+    }
+}
+
+#[tokio::test]
+async fn validates_param_header_once_tool_gains_annotation() -> anyhow::Result<()> {
+    let server = ChangingToolServer::default();
+    let (client, url, ct) = spawn_changing_tool_server(&server).await;
+
+    // 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 = 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);
+
+    ct.cancel();
+    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) {