From d77fa74eca88eb8d417cd30fb77d5a63936f939e Mon Sep 17 00:00:00 2001 From: Bailey Hayes Date: Thu, 15 Oct 2026 14:13:09 -0400 Subject: [PATCH] feat(transport): support local streamable HTTP handler futures Signed-off-by: Bailey Hayes --- crates/rmcp/Cargo.toml | 5 + crates/rmcp/README.md | 7 + crates/rmcp/src/service.rs | 4 +- crates/rmcp/src/transport.rs | 2 +- .../src/transport/streamable_http_server.rs | 4 +- .../transport/streamable_http_server/tower.rs | 25 +- .../rmcp/tests/test_local_streamable_http.rs | 217 ++++++++++++++++++ 7 files changed, 248 insertions(+), 16 deletions(-) create mode 100644 crates/rmcp/tests/test_local_streamable_http.rs diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 032aa795c..4692dac41 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -251,6 +251,11 @@ required-features = [ ] path = "tests/test_with_js.rs" +[[test]] +name = "test_local_streamable_http" +required-features = ["server", "local", "transport-streamable-http-server"] +path = "tests/test_local_streamable_http.rs" + [[test]] name = "test_local_client_handler" required-features = ["client", "local"] diff --git a/crates/rmcp/README.md b/crates/rmcp/README.md index f34469ef4..22f4aed98 100644 --- a/crates/rmcp/README.md +++ b/crates/rmcp/README.md @@ -26,6 +26,7 @@ For **getting started**, **usage guides**, and **full MCP feature documentation* | `auth` | OAuth 2.0 authentication support | | | `auth-enterprise-managed` | EMA/XAA refresh-token and ID-JAG exchanges for registered public and confidential clients (includes `auth`) | | | `elicitation` | Elicitation support | | +| `local` | Allow handler futures without `Send`; run services on a Tokio `LocalSet` | | ### Transport features @@ -38,6 +39,12 @@ For **getting started**, **usage guides**, and **full MCP feature documentation* | `transport-streamable-http-client-reqwest` | Streamable HTTP client with default `reqwest` backend | | `transport-streamable-http-server` | Streamable HTTP server transport | +With `local`, the Streamable HTTP server accepts handler futures that hold +`!Send` values across an await. Drive `StreamableHttpService::handle` or its +Tower service inside a `tokio::task::LocalSet`. The handler and service factory +retain their existing thread-safety bounds. The returned Tower future is local +and cannot be mounted directly on routers that require `Send` futures. + ### TLS backend options (for HTTP transports) | Feature | Description | diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index 430c5b77e..cc7cd70ec 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -1322,7 +1322,7 @@ where /// Without the `local` feature this is `tokio::spawn` (requires `Future: Send + 'static`). /// With `local` it uses `tokio::task::spawn_local` (requires only `Future: 'static`). #[cfg(not(feature = "local"))] -fn spawn_service_task(future: F) -> tokio::task::JoinHandle +pub(crate) fn spawn_service_task(future: F) -> tokio::task::JoinHandle where F: Future + Send + 'static, F::Output: Send + 'static, @@ -1331,7 +1331,7 @@ where } #[cfg(feature = "local")] -fn spawn_service_task(future: F) -> tokio::task::JoinHandle +pub(crate) fn spawn_service_task(future: F) -> tokio::task::JoinHandle where F: Future + 'static, F::Output: 'static, diff --git a/crates/rmcp/src/transport.rs b/crates/rmcp/src/transport.rs index 1e195a331..958c9bf37 100644 --- a/crates/rmcp/src/transport.rs +++ b/crates/rmcp/src/transport.rs @@ -111,7 +111,7 @@ pub use auth::{ // pub mod ws; #[cfg(feature = "transport-streamable-http-server-session")] pub mod streamable_http_server; -#[cfg(all(feature = "transport-streamable-http-server", not(feature = "local")))] +#[cfg(feature = "transport-streamable-http-server")] pub use streamable_http_server::tower::{StreamableHttpServerConfig, StreamableHttpService}; #[cfg(feature = "transport-streamable-http-client")] diff --git a/crates/rmcp/src/transport/streamable_http_server.rs b/crates/rmcp/src/transport/streamable_http_server.rs index df1945ab2..14da33fab 100644 --- a/crates/rmcp/src/transport/streamable_http_server.rs +++ b/crates/rmcp/src/transport/streamable_http_server.rs @@ -1,6 +1,6 @@ pub mod session; -#[cfg(all(feature = "transport-streamable-http-server", not(feature = "local")))] +#[cfg(feature = "transport-streamable-http-server")] pub mod tower; pub use session::{RestoreOutcome, SessionId, SessionManager, SessionRestoreMarker}; -#[cfg(all(feature = "transport-streamable-http-server", not(feature = "local")))] +#[cfg(feature = "transport-streamable-http-server")] pub use tower::{StreamableHttpServerConfig, StreamableHttpService}; diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index ab541ff31..543578cfd 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -10,7 +10,7 @@ use std::{ }; use bytes::Bytes; -use futures::{Stream, StreamExt, future::BoxFuture}; +use futures::{Stream, StreamExt}; use http::{HeaderMap, Method, Request, Response, header::ALLOW}; use http_body::Body; use http_body_util::{BodyExt, Full, combinators::BoxBody}; @@ -32,8 +32,8 @@ use crate::{ }, serve_server, service::{ - NotificationContext, RequestContext, Service, negotiate_protocol_version, - serve_directly_with_ct, uses_legacy_lifecycle, + MaybeBoxFuture, NotificationContext, RequestContext, Service, negotiate_protocol_version, + serve_directly_with_ct, spawn_service_task, uses_legacy_lifecycle, }, transport::{ OneshotTransport, TransportAdapterIdentity, @@ -126,8 +126,11 @@ pub struct StreamableHttpServerConfig { /// /// Call [`StreamableHttpServerConfig::enforce_origin_validation`] to enable /// validation with an empty list, rejecting every present Origin value. - /// examples: - /// allowed_origins = ["https://app.example.com:443", "http://localhost:*"] + /// + /// # Example + /// ```text + /// allowed_origins = ["https://app.example.com:443", "http://localhost:*"] + /// ``` pub allowed_origins: Vec, validate_empty_origin_allowlist: bool, /// Optional external session store for cross-instance recovery. @@ -1366,7 +1369,7 @@ where { type Response = BoxResponse; type Error = Infallible; - type Future = BoxFuture<'static, Result>; + type Future = MaybeBoxFuture<'static, Result>; fn call(&mut self, req: http::Request) -> Self::Future { let service = self.clone(); Box::pin(async move { @@ -1404,7 +1407,7 @@ impl Drop for PendingRestoreGuard { // Remove the map entry asynchronously (requires the async write lock). let pending_restores = self.pending_restores.clone(); let session_id = self.session_id.clone(); - tokio::spawn(async move { + spawn_service_task(async move { pending_restores.write().await.remove(&session_id); }); } @@ -1449,7 +1452,7 @@ where let retry = self.config.sse_retry; let server_ct = self.config.cancellation_token.child_token(); - tokio::spawn(async move { + spawn_service_task(async move { let mut sender = Some(sender); if let Some(retry) = retry && let Err(error) = persist_and_forward_event( @@ -1554,7 +1557,7 @@ where peer_info, request_ct.clone(), ); - tokio::spawn(async move { + spawn_service_task(async move { let _ = service.waiting().await; }); @@ -1636,7 +1639,7 @@ where S: crate::ServerHandler + Send + 'static, M: SessionManager, { - tokio::spawn(async move { + spawn_service_task(async move { let svc = serve_server::(service, transport) .await; @@ -2305,7 +2308,7 @@ where peer_info, request_ct.clone(), ); - tokio::spawn(async move { + spawn_service_task(async move { // on service created let _ = service.waiting().await; }); diff --git a/crates/rmcp/tests/test_local_streamable_http.rs b/crates/rmcp/tests/test_local_streamable_http.rs new file mode 100644 index 000000000..087f49d3e --- /dev/null +++ b/crates/rmcp/tests/test_local_streamable_http.rs @@ -0,0 +1,217 @@ +#![cfg(all( + feature = "local", + feature = "server", + feature = "transport-streamable-http-server" +))] + +use std::{rc::Rc, sync::Arc, time::Duration}; + +use bytes::Bytes; +use http::{Request, Response, StatusCode}; +use http_body_util::{BodyExt, Full}; +use rmcp::{ + ServerHandler, + model::{ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerConfig}, + service::{RequestContext, RoleServer}, + transport::{ + StreamableHttpServerConfig, StreamableHttpService, + streamable_http_server::session::{local::LocalSessionManager, never::NeverSessionManager}, + }, +}; +use serde_json::{Value, json}; +use tower_service::Service as _; + +struct LocalServer; + +impl ServerHandler for LocalServer { + async fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + let local = Rc::new(42); + tokio::task::yield_now().await; + Ok(ListToolsResult { + next_cursor: Some(local.to_string()), + ..Default::default() + }) + } + + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } +} + +fn request( + method: &str, + version: &str, + params: Value, + session: Option<&str>, +) -> Request> { + let mut request = Request::post("http://localhost/mcp") + .header("Host", "localhost") + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", version) + .header("Mcp-Method", method); + if let Some(session) = session { + request = request.header("Mcp-Session-Id", session); + } + request + .body(Full::new(Bytes::from( + serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": 1, + "method": method, + "params": params, + })) + .unwrap(), + ))) + .unwrap() +} + +fn list_tools_request(version: &str, session: Option<&str>) -> Request> { + let params = if version == "2026-07-28" { + json!({ + "_meta": { + "io.modelcontextprotocol/protocolVersion": version, + "io.modelcontextprotocol/clientCapabilities": {}, + } + }) + } else { + json!({}) + }; + request("tools/list", version, params, session) +} + +async fn assert_tools(response: Response, json_response: bool) +where + B: http_body::Body, + B::Error: std::fmt::Debug, +{ + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.headers()["Content-Type"], + if json_response { + "application/json" + } else { + "text/event-stream" + }, + ); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let body = std::str::from_utf8(&body).unwrap(); + let message: Value = if json_response { + serde_json::from_str(body).unwrap() + } else { + body.lines() + .filter_map(|line| line.strip_prefix("data: ")) + .find_map(|data| serde_json::from_str(data).ok()) + .expect("SSE stream contains a JSON-RPC response") + }; + assert_eq!(message["id"], 1); + assert_eq!(message["result"]["tools"], json!([])); + assert_eq!(message["result"]["nextCursor"], "42"); +} + +async fn dispatch(json_response: bool, tower: bool) { + tokio::task::LocalSet::new() + .run_until(async { + tokio::time::timeout(Duration::from_secs(5), async { + for version in ["2026-07-28", "2025-11-25"] { + let config = StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(json_response); + let mut service = StreamableHttpService::new( + || Ok(LocalServer), + Arc::new(NeverSessionManager::default()), + config, + ); + let request = list_tools_request(version, None); + let response = if tower { + service.call(request).await.unwrap() + } else { + service.handle(request).await + }; + assert_tools(response, json_response).await; + } + }) + .await + .expect("local HTTP dispatch completes"); + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn json_handler_accepts_non_send_future() { + dispatch(true, false).await; +} + +#[tokio::test(flavor = "current_thread")] +async fn sse_handler_accepts_non_send_future() { + dispatch(false, false).await; +} + +#[tokio::test(flavor = "current_thread")] +async fn json_tower_service_accepts_non_send_future() { + dispatch(true, true).await; +} + +#[tokio::test(flavor = "current_thread")] +async fn sse_tower_service_accepts_non_send_future() { + dispatch(false, true).await; +} + +#[tokio::test(flavor = "current_thread")] +async fn legacy_session_accepts_non_send_future() { + tokio::task::LocalSet::new() + .run_until(async { + tokio::time::timeout(Duration::from_secs(5), async { + let config = StreamableHttpServerConfig::default(); + let service = StreamableHttpService::new( + || Ok(LocalServer), + Arc::new(LocalSessionManager::default()), + config, + ); + let initialize = request( + "initialize", + "2025-11-25", + json!({ + "protocolVersion": "2025-11-25", + "capabilities": {}, + "clientInfo": {"name": "local-test", "version": "1"}, + }), + None, + ); + let response = service.handle(initialize).await; + assert_eq!(response.status(), StatusCode::OK); + let session = response.headers()["Mcp-Session-Id"] + .to_str() + .unwrap() + .to_owned(); + response.into_body().collect().await.unwrap(); + let notification = Request::post("http://localhost/mcp") + .header("Host", "localhost") + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-11-25") + .header("Mcp-Method", "notifications/initialized") + .header("Mcp-Session-Id", &session) + .body(Full::new(Bytes::from_static( + br#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#, + ))) + .unwrap(); + assert_eq!( + service.handle(notification).await.status(), + StatusCode::ACCEPTED + ); + let response = service + .handle(list_tools_request("2025-11-25", Some(&session))) + .await; + assert_tools(response, false).await; + service.config.cancellation_token.cancel(); + }) + .await + .expect("local session dispatch completes"); + }) + .await; +}