diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 21d68be89..647406851 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -2,7 +2,15 @@ #![expect(deprecated)] pub(super) mod cache; -use std::{borrow::Cow, num::NonZeroUsize, sync::Arc, time::Duration}; +use std::{ + borrow::Cow, + num::NonZeroUsize, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; use cache::CacheGeneration; pub use cache::{ClientCacheConfig, MAX_CLIENT_CACHE_TTL}; @@ -711,8 +719,9 @@ pub enum ClientLifecycleMode { Discover { preferred_versions: Vec, }, - /// Probe with `server/discover`, falling back when the peer reports that it is legacy or does - /// not respond within 10 seconds. + /// Probe with `server/discover` and fall back if the peer reports that it is legacy. + /// Streamable HTTP waits for the discovery POST within the caller's startup deadline. + /// Other transports also fall back if discovery does not finish within 10 seconds. Auto { preferred_versions: Vec, legacy_version: Option, @@ -721,6 +730,21 @@ pub enum ClientLifecycleMode { const DEFAULT_AUTO_DISCOVER_TIMEOUT: Duration = Duration::from_secs(10); +/// Shared with the transport so Auto can tell an outstanding HTTP POST from a silent peer. +#[derive(Default)] +pub(crate) struct DiscoveryProbe { + submitted_to_http: AtomicBool, +} + +impl DiscoveryProbe { + #[cfg(feature = "transport-streamable-http-client")] + pub(crate) fn observe_http_request(extensions: &Extensions) { + if let Some(probe) = extensions.get::>() { + probe.submitted_to_http.store(true, Ordering::Relaxed); + } + } +} + /// Client-specific lifecycle entry points. pub trait ClientServiceExt: Service + Sized { fn serve_with_lifecycle( @@ -853,6 +877,7 @@ where &peer, &client_info, preferred_versions, + None, ) .await? { @@ -865,18 +890,25 @@ where preferred_versions, legacy_version, } => { - let discover_result = tokio::time::timeout( - auto_discover_timeout, - discover_startup( + let probe = Arc::new(DiscoveryProbe::default()); + let discover_result = { + let discovery = discover_startup( &service, &mut transport, &id_provider, &peer, &client_info, preferred_versions, - ), - ) - .await; + Some(&probe), + ); + tokio::pin!(discovery); + match tokio::time::timeout(auto_discover_timeout, &mut discovery).await { + Err(_) if probe.submitted_to_http.load(Ordering::Relaxed) => { + Ok(discovery.await) + } + result => result, + } + }; match discover_result { Ok(Ok(DiscoverOutcome::Modern)) => {} Ok(Ok(DiscoverOutcome::Legacy(discover_error))) => { @@ -1026,6 +1058,7 @@ async fn discover_startup( peer: &Peer, client_info: &InitializeRequestParams, preferred_versions: Vec, + probe: Option<&Arc>, ) -> Result where S: Service, @@ -1047,6 +1080,9 @@ where ); let mut discover = DiscoverRequest::new(DiscoverRequestParams {}); discover.extensions.insert(meta); + if let Some(probe) = probe { + discover.extensions.insert(Arc::clone(probe)); + } let id = id_provider.next_request_id(); transport .send(ClientJsonRpcMessage::request( @@ -2656,9 +2692,12 @@ mod tests { client.cancel().await.unwrap(); } - #[tokio::test] - async fn auto_startup_falls_back_when_discover_is_ignored() { - use crate::model::{InitializeResult, ServerCapabilities}; + #[rstest::rstest] + #[case::initial_probe(false)] + #[case::version_retry(true)] + #[tokio::test(start_paused = true)] + async fn auto_startup_falls_back_when_discover_is_ignored(#[case] retry: bool) { + use crate::model::{ErrorCode, InitializeResult, ServerCapabilities}; tokio::task::LocalSet::new() .run_until(async { @@ -2678,8 +2717,27 @@ mod tests { ClientRequest::DiscoverRequest(_) )); + if retry { + let error = ErrorData::new( + ErrorCode::UNSUPPORTED_PROTOCOL_VERSION, + "Unsupported version", + Some(serde_json::json!({"supported": ["2026-07-28"]})), + ); + server + .send(ServerJsonRpcMessage::error(error, Some(discover.id))) + .await + .unwrap(); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Request(request)) + if matches!(request.request, ClientRequest::DiscoverRequest(_)) + )); + } let ClientJsonRpcMessage::Request(initialize) = - server.receive().await.expect("expected initialize request") + tokio::time::timeout(Duration::from_secs(1), server.receive()) + .await + .expect("fallback deadline") + .expect("expected initialize request") else { panic!("expected initialize request"); }; diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 7b1981ff0..3379ea121 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -1145,6 +1145,10 @@ impl Worker for StreamableHttpClientWorker { } fn config(&self) -> super::worker::WorkerConfig { super::worker::WorkerConfig { + #[cfg(feature = "client")] + on_request_submitted: Some(crate::service::DiscoveryProbe::observe_http_request), + #[cfg(not(feature = "client"))] + on_request_submitted: None, name: Some("StreamableHttpClientWorker".into()), channel_buffer_capacity: self.config.channel_buffer_capacity, } @@ -1169,6 +1173,10 @@ impl Worker for StreamableHttpClientWorker { ClientJsonRpcMessage::Request(request) if matches!(&request.request, ClientRequest::InitializeRequest(_)) ); + let is_discover = matches!( + &startup_request, ClientJsonRpcMessage::Request(request) + if matches!(request.request, ClientRequest::DiscoverRequest(_)) + ); 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 { @@ -1181,32 +1189,34 @@ impl Worker for StreamableHttpClientWorker { &empty_tool_cache, ) }; - let (message, session_id) = match self - .client - .post_message_with_max_sse_event_size( - config.uri.clone(), - startup_request, - None, - config.auth_header.clone(), - bootstrap_headers.clone(), - config.max_sse_event_size, + let post = self.client.post_message_with_max_sse_event_size( + config.uri.clone(), + startup_request, + None, + config.auth_header.clone(), + bootstrap_headers.clone(), + config.max_sse_event_size, + ); + let startup = async { + let response = match post.await { + Ok(response) => response, + Err(err) => { + let msg = format!("{err:?}"); + let _ = responder.send(Err(err)); + return Err(WorkerQuitReason::fatal( + StreamableHttpError::TransportChannelClosed, + msg, + )); + } + }; + let _ = responder.send(Ok(())); + response.expect_initialized::().await.map_err( + WorkerQuitReason::fatal_context("process initialize response"), ) - .await - { - Ok(res) => { - let _ = responder.send(Ok(())); - res.expect_initialized::().await.map_err( - WorkerQuitReason::fatal_context("process initialize response"), - )? - } - Err(err) => { - let msg = format!("{:?}", err); - let _ = responder.send(Err(err)); - return Err(WorkerQuitReason::fatal( - StreamableHttpError::TransportChannelClosed, - msg, - )); - } + }; + let (message, session_id) = tokio::select! { + _ = transport_task_ct.cancelled(), if is_discover => return Err(WorkerQuitReason::Cancelled), + response = startup => response?, }; let mut uses_modern_http = !is_legacy_startup; let mut session_id: Option> = if uses_modern_http { diff --git a/crates/rmcp/src/transport/streamable_http_server/session/local.rs b/crates/rmcp/src/transport/streamable_http_server/session/local.rs index a3e6d19f4..d65358011 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -1051,6 +1051,7 @@ impl Worker for LocalSessionWorker { crate::transport::worker::WorkerConfig { name: Some(format!("streamable-http-session-{}", self.id)), channel_buffer_capacity: self.session_config.channel_capacity, + ..Default::default() } } #[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))] diff --git a/crates/rmcp/src/transport/worker.rs b/crates/rmcp/src/transport/worker.rs index e32da7d7d..1a2107c67 100644 --- a/crates/rmcp/src/transport/worker.rs +++ b/crates/rmcp/src/transport/worker.rs @@ -13,7 +13,7 @@ use tracing::{Instrument, Level}; use super::{IntoTransport, Transport}; use crate::{ - model::{CancelledNotification, JsonRpcMessage, RequestId}, + model::{CancelledNotification, Extensions, GetExtensions, JsonRpcMessage, RequestId}, service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage}, }; @@ -167,6 +167,7 @@ impl WorkerSendRequest { } pub struct WorkerTransport { + on_request_submitted: Option, rx: tokio::sync::mpsc::Receiver>, send_service: tokio::sync::mpsc::Sender>, control_send_service: tokio::sync::mpsc::Sender>, @@ -179,6 +180,8 @@ pub struct WorkerTransport { #[non_exhaustive] pub struct WorkerConfig { + /// Runs synchronously in `Transport::send`, before its future is polled. + pub(crate) on_request_submitted: Option, pub name: Option, pub channel_buffer_capacity: usize, } @@ -188,6 +191,7 @@ impl Default for WorkerConfig { Self { name: None, channel_buffer_capacity: 16, + on_request_submitted: None, } } } @@ -254,6 +258,7 @@ impl WorkerTransport { }) }); Self { + on_request_submitted: config.on_request_submitted, rx: from_transport_rx, send_service: to_transport_tx, control_send_service: control_to_transport_tx, @@ -347,6 +352,11 @@ impl Transport for WorkerTransport { &mut self, item: TxJsonRpcMessage, ) -> impl Future> + Send + 'static { + if let Some(on_request_submitted) = self.on_request_submitted + && let JsonRpcMessage::Request(request) = &item + { + on_request_submitted(request.request.extensions()); + } let control_generation = self.control_generation.load(Ordering::SeqCst); let mut cancellation_target = None; let registration = if W::supports_request_cancellation() { @@ -452,6 +462,39 @@ mod tests { } } + #[tokio::test(start_paused = true)] + async fn auto_falls_back_while_discovery_send_awaits_acknowledgment() { + use crate::{ + ClientLifecycleMode, ClientServiceExt, + model::{ClientConfig, ClientRequest, ProtocolVersion}, + }; + let (worker, context) = tokio::sync::oneshot::channel(); + let transport = WorkerTransport::spawn(TestWorker(worker)); + let startup = ClientConfig::default().serve_with_lifecycle( + transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: None, + }, + ); + let (_, ()) = tokio::join!(startup, async { + let mut context = context.await.unwrap(); + // Keep the discovery send pending while the bridge accepts the next request. + let probe = context.recv_from_handler().await.unwrap(); + assert!(matches!(&probe.message, JsonRpcMessage::Request(request) + if matches!(request.request, ClientRequest::DiscoverRequest(_)))); + let initialize = + tokio::time::timeout(Duration::from_secs(15), context.recv_from_handler()) + .await + .expect("Auto should fall back after 10 seconds") + .unwrap(); + assert!( + matches!(initialize.message, JsonRpcMessage::Request(request) + if matches!(request.request, ClientRequest::InitializeRequest(_))) + ); + }); + } + fn cancellation_message(id: RequestId) -> ClientJsonRpcMessage { serde_json::from_value(serde_json::json!({ "jsonrpc": "2.0", diff --git a/crates/rmcp/tests/test_discover_http_client_startup.rs b/crates/rmcp/tests/test_discover_http_client_startup.rs index 70aa219dc..bf9bee4c0 100644 --- a/crates/rmcp/tests/test_discover_http_client_startup.rs +++ b/crates/rmcp/tests/test_discover_http_client_startup.rs @@ -5,7 +5,7 @@ feature = "transport-streamable-http-server" ))] -use std::{borrow::Cow, sync::Arc}; +use std::{borrow::Cow, sync::Arc, time::Duration}; use axum::{ Router, @@ -14,12 +14,16 @@ use axum::{ http::{Response, StatusCode}, routing::post, }; +use futures::StreamExt; use rmcp::{ - ClientLifecycleMode, ClientServiceExt, ServerHandler, - model::{ClientConfig, DiscoverResult, ErrorCode, ErrorData, ProtocolVersion}, + ClientLifecycleMode, ClientServiceExt, RoleClient, ServerHandler, + model::{ + ClientConfig, ClientJsonRpcMessage, ClientRequest, DiscoverResult, ErrorCode, ErrorData, + ProtocolVersion, ServerJsonRpcMessage, + }, service::{MaybeSendFuture, RequestContext, RoleServer}, transport::{ - StreamableHttpClientTransport, + StreamableHttpClientTransport, Transport, streamable_http_client::StreamableHttpClientTransportConfig, streamable_http_server::{ StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, @@ -103,8 +107,51 @@ async fn plain_text_legacy_http_handler( .expect("build notification response") } -#[tokio::test] -async fn discover_http_client_bootstraps_headers_without_initialize() { +struct DelayedTransport { + inner: T, + defer_discovery_send: bool, +} + +impl> Transport for DelayedTransport { + type Error = T::Error; + + fn send( + &mut self, + message: ClientJsonRpcMessage, + ) -> impl Future> + Send + 'static { + let defer = self.defer_discovery_send + && matches!(&message, ClientJsonRpcMessage::Request(r) + if matches!(r.request, ClientRequest::DiscoverRequest(_))); + let send = self.inner.send(message); + async move { + if defer { + tokio::time::sleep(Duration::from_secs(12)).await; + } + send.await + } + } + + async fn receive(&mut self) -> Option { + tokio::task::yield_now().await; + self.inner.receive().await + } + + async fn close(&mut self) -> Result<(), Self::Error> { + self.inner.close().await + } +} + +#[rstest::rstest] +#[case::discover_json(false, true, false)] +#[case::auto_json(true, true, false)] +#[case::auto_sse(true, false, false)] +#[case::queued_http_send(true, true, true)] +#[tokio::test(start_paused = true)] +async fn discover_http_client_bootstraps_headers_without_initialize( + #[case] auto: bool, + #[case] json_response: bool, + #[case] defer_discovery_send: bool, +) { let ct = CancellationToken::new(); let service: StreamableHttpService = StreamableHttpService::new( @@ -112,10 +159,22 @@ async fn discover_http_client_bootstraps_headers_without_initialize() { Default::default(), StreamableHttpServerConfig::default() .with_legacy_session_mode(false) - .with_json_response(true) + .with_json_response(json_response) .with_cancellation_token(ct.child_token()), ); - let router = axum::Router::new().nest_service("/mcp", service); + // Send headers immediately but delay body chunks past the 10-second fallback. + let router = Router::new() + .nest_service("/mcp", service) + .layer(axum::middleware::map_response( + |response: Response| async { + response.map(|body| { + Body::from_stream(body.into_data_stream().then(|chunk| async move { + tokio::time::sleep(Duration::from_secs(12)).await; + chunk + })) + }) + }, + )); let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); @@ -132,12 +191,22 @@ async fn discover_http_client_bootstraps_headers_without_initialize() { let transport = StreamableHttpClientTransport::from_config( StreamableHttpClientTransportConfig::with_uri(format!("http://{address}/mcp")), ); + let preferred_versions = vec![ProtocolVersion::V_2026_07_28]; + let lifecycle = if auto { + ClientLifecycleMode::Auto { + preferred_versions, + legacy_version: None, + } + } else { + ClientLifecycleMode::Discover { preferred_versions } + }; let client = ClientConfig::default() .serve_with_lifecycle( - transport, - ClientLifecycleMode::Discover { - preferred_versions: vec![ProtocolVersion::V_2026_07_28], + DelayedTransport { + inner: transport, + defer_discovery_send, }, + lifecycle, ) .await .expect("discover HTTP client should start"); @@ -275,6 +344,7 @@ async fn version_negotiating_http_handler( .expect("build version rejection response"); } + tokio::time::sleep(Duration::from_secs(12)).await; let result = DiscoverResult::new(vec![ProtocolVersion::V_2025_11_25], Default::default()); return Response::builder() .status(StatusCode::OK) @@ -296,7 +366,7 @@ async fn version_negotiating_http_handler( /// A recognized modern error in a 4xx discover rejection identifies a modern /// server, so the client retries at a supported version instead of falling back /// to legacy `initialize`. -#[tokio::test] +#[tokio::test(start_paused = true)] async fn auto_http_client_retries_version_after_modern_4xx_rejection() { let ct = CancellationToken::new(); let state = VersionNegotiatingHttpState::default();