fix: add init_timeout for streamable-http sessions (#811)

Author: DaleSeo

crates/rmcp/src/transport/streamable_http_server/session/local.rs

@@ -930,6 +930,8 @@ pub enum LocalSessionWorkerError {
     FailToHandleMessage(SessionError),
     #[error("keep alive timeout after {}ms", _0.as_millis())]
     KeepAliveTimeout(Duration),
+    #[error("init timeout after {}ms", _0.as_millis())]
+    InitTimeout(Duration),
     #[error("Transport closed")]
     TransportClosed,
     #[error("Tokio join error {0}")]
@@ -959,13 +961,24 @@ impl Worker for LocalSessionWorker {
             FromHttpService(SessionEvent),
             FromHandler(WorkerSendRequest<LocalSessionWorker>),
         }
-        // waiting for initialize request
-        let evt = self.event_rx.recv().await.ok_or_else(|| {
-            WorkerQuitReason::fatal(
-                LocalSessionWorkerError::TransportTerminated,
-                "get initialize request",
-            )
-        })?;
+        let init_timeout = self.session_config.init_timeout.unwrap_or(Duration::MAX);
+        let evt = tokio::select! {
+            evt = self.event_rx.recv() => evt.ok_or_else(|| {
+                WorkerQuitReason::fatal(
+                    LocalSessionWorkerError::TransportTerminated,
+                    "get initialize request",
+                )
+            })?,
+            _ = context.cancellation_token.cancelled() => {
+                return Err(WorkerQuitReason::Cancelled);
+            }
+            _ = tokio::time::sleep(init_timeout) => {
+                return Err(WorkerQuitReason::fatal(
+                    LocalSessionWorkerError::InitTimeout(init_timeout),
+                    "waiting for initialize request",
+                ));
+            }
+        };
         let SessionEvent::InitializeRequest { request, responder } = evt else {
             return Err(WorkerQuitReason::fatal(
                 LocalSessionWorkerError::UnexpectedEvent(evt),
@@ -1122,13 +1135,18 @@ pub struct SessionConfig {
     /// resume requests. After this duration, completed entries are evicted
     /// and resume will return an error. Default is 60 seconds.
     pub completed_cache_ttl: Duration,
+    /// Maximum duration to wait for the `initialize` request after session
+    /// creation. If not received within this window, the session is
+    /// terminated. Default is 60 seconds. Set to `None` to disable.
+    pub init_timeout: Option<Duration>,
 }
 
 impl SessionConfig {
     pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
     pub const DEFAULT_KEEP_ALIVE: Duration = Duration::from_secs(300);
     pub const DEFAULT_SSE_RETRY: Duration = Duration::from_secs(3);
     pub const DEFAULT_COMPLETED_CACHE_TTL: Duration = Duration::from_secs(60);
+    pub const DEFAULT_INIT_TIMEOUT: Duration = Duration::from_secs(60);
 }
 
 impl Default for SessionConfig {
@@ -1138,6 +1156,7 @@ impl Default for SessionConfig {
             keep_alive: Some(Self::DEFAULT_KEEP_ALIVE),
             sse_retry: Some(Self::DEFAULT_SSE_RETRY),
             completed_cache_ttl: Self::DEFAULT_COMPLETED_CACHE_TTL,
+            init_timeout: Some(Self::DEFAULT_INIT_TIMEOUT),
         }
     }
 }

crates/rmcp/tests/test_streamable_http_init_timeout.rs

@@ -0,0 +1,63 @@
+#![cfg(all(feature = "transport-streamable-http-server", not(feature = "local")))]
+
+use std::time::Duration;
+
+use rmcp::{
+    model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId},
+    transport::streamable_http_server::session::{SessionManager, local::LocalSessionManager},
+};
+
+#[tokio::test]
+async fn test_init_timeout_terminates_pre_init_session() -> anyhow::Result<()> {
+    let mut manager = LocalSessionManager::default();
+    manager.session_config.init_timeout = Some(Duration::from_millis(200));
+
+    // Bind the transport so its drop-guard doesn't cancel the worker — we
+    // want termination via init_timeout, not via cancellation.
+    let (session_id, _transport) = manager.create_session().await?;
+
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    let message = ClientJsonRpcMessage::request(
+        ClientRequest::PingRequest(PingRequest::default()),
+        RequestId::Number(1),
+    );
+    let result = manager.initialize_session(&session_id, message).await;
+
+    assert!(
+        result.is_err(),
+        "expected worker to be dead; got: {result:?}"
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_init_timeout_none_keeps_worker_alive() -> anyhow::Result<()> {
+    let mut manager = LocalSessionManager::default();
+    manager.session_config.init_timeout = None;
+
+    let (session_id, _transport) = manager.create_session().await?;
+
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    let message = ClientJsonRpcMessage::request(
+        ClientRequest::PingRequest(PingRequest::default()),
+        RequestId::Number(1),
+    );
+    // Liveness probe: a live worker accepts the send then stalls waiting for
+    // a handler response (none is wired up), tripping the outer timeout. A
+    // dead worker would fail the send and return immediately.
+    let probe = tokio::time::timeout(
+        Duration::from_millis(200),
+        manager.initialize_session(&session_id, message),
+    )
+    .await;
+
+    assert!(
+        probe.is_err(),
+        "expected worker to be alive; got: {probe:?}"
+    );
+
+    Ok(())
+}