feat: allow externally supplied tool schemas

crates/rmcp/src/transport/streamable_http_client.rs

@@ -839,6 +839,17 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
         let (sse_worker_tx, mut sse_worker_rx) =
             tokio::sync::mpsc::channel::<ServerJsonRpcMessage>(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,15 +862,14 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                 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 {
             request_version_headers(
                 &config.custom_headers,
                 &startup_request,
                 &ProtocolVersion::default(),
-                &empty_tool_cache,
+                &tool_header_cache,
             )
         };
         let (message, session_id) = match self
@@ -908,10 +918,6 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
         } 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<String, Arc<JsonObject>> = 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<String>,
     /// Custom HTTP headers to include with every request
     pub custom_headers: HashMap<HeaderName, HeaderValue>,
+    /// 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<String, Arc<JsonObject>>,
     /// 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<String, Arc<JsonObject>>) -> 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(

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<Arc<JsonObject>> + Send + Sync;
+
 #[non_exhaustive]
 #[derive(Debug, Clone)]
 pub struct StreamableHttpServerConfig {
@@ -1011,6 +1014,9 @@ pub struct StreamableHttpService<S, M> {
     /// 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>>>>>,
+    /// Optional request-aware source for tool schemas. When configured, this is
+    /// authoritative and replaces the service-wide `get_tool` cache.
+    tool_schema_resolver: Option<Arc<ToolSchemaResolver>>,
 }
 
 impl<S, M> Clone for StreamableHttpService<S, M> {
@@ -1021,6 +1027,7 @@ impl<S, M> Clone for StreamableHttpService<S, M> {
             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<Arc<JsonObject>>
+        + Send
+        + Sync
+        + 'static,
+    ) -> Self {
+        self.tool_schema_resolver = Some(Arc::new(resolver));
+        self
+    }
+
     fn get_service(&self) -> Result<S, std::io::Error> {
         (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<Arc<JsonObject>> {
+    /// 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<Arc<JsonObject>> {
+        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

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<HeaderValidationServer, LocalSessionManager> =
+        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>()? {
+                    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,
+        &region_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(())
+}