fix(client): wait for slow HTTP discovery in Auto mode

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<ProtocolVersion>,
     },
-    /// 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<ProtocolVersion>,
         legacy_version: Option<ProtocolVersion>,
@@ -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::<Arc<Self>>() {
+            probe.submitted_to_http.store(true, Ordering::Relaxed);
+        }
+    }
+}
+
 /// Client-specific lifecycle entry points.
 pub trait ClientServiceExt: Service<RoleClient> + Sized {
     fn serve_with_lifecycle<T, E, A>(
@@ -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<S, T>(
     peer: &Peer<RoleClient>,
     client_info: &InitializeRequestParams,
     preferred_versions: Vec<ProtocolVersion>,
+    probe: Option<&Arc<DiscoveryProbe>>,
 ) -> Result<DiscoverOutcome, ClientInitializeError>
 where
     S: Service<RoleClient>,
@@ -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");
                     };

crates/rmcp/src/transport/streamable_http_client.rs

@@ -1145,6 +1145,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
     }
     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<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
             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<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                 &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::<C::Error>().await.map_err(
+                WorkerQuitReason::fatal_context("process initialize response"),
             )
-            .await
-        {
-            Ok(res) => {
-                let _ = responder.send(Ok(()));
-                res.expect_initialized::<C::Error>().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<Arc<str>> = if uses_modern_http {

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()))]

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<W: Worker> WorkerSendRequest<W> {
 }
 
 pub struct WorkerTransport<W: Worker> {
+    on_request_submitted: Option<fn(&Extensions)>,
     rx: tokio::sync::mpsc::Receiver<RxJsonRpcMessage<W::Role>>,
     send_service: tokio::sync::mpsc::Sender<WorkerSendRequest<W>>,
     control_send_service: tokio::sync::mpsc::Sender<WorkerSendRequest<W>>,
@@ -179,6 +180,8 @@ pub struct WorkerTransport<W: Worker> {
 
 #[non_exhaustive]
 pub struct WorkerConfig {
+    /// Runs synchronously in `Transport::send`, before its future is polled.
+    pub(crate) on_request_submitted: Option<fn(&Extensions)>,
     pub name: Option<String>,
     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<W: Worker> WorkerTransport<W> {
                 })
         });
         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<W: Worker> Transport<W::Role> for WorkerTransport<W> {
         &mut self,
         item: TxJsonRpcMessage<W::Role>,
     ) -> impl Future<Output = Result<(), Self::Error>> + 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",