fix(streamable-http-server): keep slow per-request-negotiated SSE calls alive

crates/rmcp/src/transport/streamable_http_server/tower.rs

@@ -219,6 +219,17 @@ type PendingRestores = Arc<tokio::sync::RwLock<HashMap<SessionId, RestoreResultS
 #[derive(Debug, Clone)]
 pub struct StreamableHttpServerConfig {
     /// The ping message duration for SSE connections.
+    ///
+    /// A stateless request with per-request protocol metadata is answered with
+    /// an HTTP status taken from its first message. When `json_response` is
+    /// false and the handler has sent nothing for this long, the server commits
+    /// the `text/event-stream` head (HTTP 200) instead, so these pings reach the
+    /// client while the handler runs; its result, or an error, then arrives as
+    /// an SSE event. With the default 15 s, a `-32601` answered after 14 s is a
+    /// `404` with an `application/json` body, and the same answer after 16 s is a
+    /// `200` `text/event-stream` carrying the error as an event. Set this below
+    /// the shortest idle timeout between server and client: the stream is silent
+    /// for at most this long, before and after the head.
     pub sse_keep_alive: Option<Duration>,
     /// The retry interval for SSE priming events.
     pub sse_retry: Option<Duration>,
@@ -1664,7 +1675,8 @@ where
         }
     }
 
-    // The HTTP status must be known before opening an SSE stream.
+    // The first message decides the HTTP status, unless the handler is still
+    // silent after one keep-alive interval: then the SSE head is committed.
     async fn serve_negotiated_request_directly(
         &self,
         service: S,
@@ -1695,19 +1707,43 @@ where
         });
 
         let cancel = self.config.cancellation_token.child_token();
+        // Without a head, keep-alive comments cannot flow, and a long-running
+        // handler leaves the connection idle until its first message: long
+        // enough for clients and proxies to drop it. Committing after one
+        // keep-alive interval keeps the status mapping for prompt answers.
+        let commit_after = self
+            .config
+            .sse_keep_alive
+            .filter(|_| !self.config.json_response);
+        let commit_deadline = async move {
+            match commit_after {
+                Some(after) => tokio::time::sleep(after).await,
+                None => std::future::pending::<()>().await,
+            }
+        };
         // Cancel the handler if the client disconnects while it is still
         // producing its first message (this future is dropped before
         // `receiver.recv()` completes). Disarmed once the handler emits
-        // anything, so a normal response is never cancelled.
+        // anything or the stream is committed, so a normal response is never
+        // cancelled.
         let mut disconnect_guard = Some(request_ct.clone().drop_guard());
         let first = tokio::select! {
+            // A message that is ready when the deadline fires still decides
+            // the status: poll it first.
+            biased;
             message = receiver.recv() => {
                 if let Some(guard) = disconnect_guard.take() {
                     guard.disarm();
                 }
                 message
             }
             _ = cancel.cancelled() => None,
+            () = commit_deadline => {
+                if let Some(guard) = disconnect_guard.take() {
+                    guard.disarm();
+                }
+                return Ok(self.stateless_sse_response(None, receiver, request_ct));
+            }
         }
         .ok_or_else(|| {
             internal_error_response("empty response")(std::io::Error::new(

crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs

@@ -18,8 +18,8 @@ use std::{sync::Arc, time::Duration};
 use rmcp::{
     ErrorData as McpError, RoleServer, ServerHandler,
     model::{
-        CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ServerCapabilities,
-        ServerConfig,
+        CallToolRequestMethod, CallToolRequestParams, CallToolResponse, CallToolResult,
+        ContentBlock, ServerCapabilities, ServerConfig,
     },
     service::RequestContext,
     transport::streamable_http_server::{
@@ -29,6 +29,11 @@ use rmcp::{
 use tokio::sync::Notify;
 use tokio_util::sync::CancellationToken;
 
+/// A short keep-alive lets the SSE server notice a dropped connection quickly
+/// (hyper only observes the disconnect on its next write), and is the deadline
+/// after which a silent per-request-negotiated handler gets its SSE head.
+const KEEP_ALIVE: Duration = Duration::from_millis(200);
+
 #[derive(Clone)]
 struct CancelProbe {
     started: Arc<Notify>,
@@ -43,9 +48,13 @@ impl ServerHandler for CancelProbe {
 
     async fn call_tool(
         &self,
-        _request: CallToolRequestParams,
+        request: CallToolRequestParams,
         context: RequestContext<RoleServer>,
     ) -> Result<CallToolResponse, McpError> {
+        if request.name == "answer_soon" {
+            tokio::time::sleep(KEEP_ALIVE / 2).await;
+            return Err(McpError::method_not_found::<CallToolRequestMethod>());
+        }
         self.started.notify_one();
         // Wait until the per-request cancellation token fires, or give up after a
         // generous timeout so a buggy build fails via the outer assertion rather
@@ -64,6 +73,24 @@ impl ServerHandler for CancelProbe {
 
 const CALL_BODY: &str = r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"wait_for_cancel","arguments":{}}}"#;
 
+fn negotiated_call(id: u32, tool: &str) -> String {
+    serde_json::json!({
+        "jsonrpc": "2.0",
+        "id": id,
+        "method": "tools/call",
+        "params": {
+            "name": tool,
+            "arguments": {},
+            "_meta": {
+                "io.modelcontextprotocol/protocolVersion": "2026-07-28",
+                "io.modelcontextprotocol/clientInfo": {"name": "test", "version": "1.0"},
+                "io.modelcontextprotocol/clientCapabilities": {}
+            }
+        }
+    })
+    .to_string()
+}
+
 struct TestServer {
     url: String,
     server_ct: CancellationToken,
@@ -83,9 +110,7 @@ async fn spawn_stateless_server(json_response: bool) -> anyhow::Result<TestServe
     let config = StreamableHttpServerConfig::default()
         .with_legacy_session_mode(false)
         .with_json_response(json_response)
-        // A short keep-alive lets the SSE server notice a dropped connection
-        // quickly (hyper only observes the disconnect on its next write).
-        .with_sse_keep_alive(Some(Duration::from_millis(100)))
+        .with_sse_keep_alive(Some(KEEP_ALIVE))
         .with_cancellation_token(server_ct.child_token());
 
     let service: StreamableHttpService<CancelProbe, LocalSessionManager> =
@@ -194,3 +219,95 @@ async fn stateless_json_client_disconnect_cancels_request() -> anyhow::Result<()
     server.server_ct.cancel();
     Ok(())
 }
+
+/// Per-request protocol metadata: a handler still running after one keep-alive
+/// interval gets its SSE head, so keep-alive comments reach the client instead
+/// of the connection idling until the result. Dropping that stream still
+/// cancels the handler.
+#[tokio::test]
+async fn stateless_negotiated_sse_commits_head_while_handler_runs() -> anyhow::Result<()> {
+    let server = spawn_stateless_server(false).await?;
+    let client = reqwest::Client::builder()
+        .pool_max_idle_per_host(0)
+        .build()?;
+
+    let mut call = tokio::time::timeout(
+        Duration::from_secs(5),
+        client
+            .post(&server.url)
+            .header("Content-Type", "application/json")
+            .header("Accept", "application/json, text/event-stream")
+            .header("MCP-Protocol-Version", "2026-07-28")
+            .header("Mcp-Method", "tools/call")
+            .header("Mcp-Name", "wait_for_cancel")
+            .body(negotiated_call(2, "wait_for_cancel"))
+            .send(),
+    )
+    .await
+    .expect("the response head should not wait for the handler to finish")?;
+    assert_eq!(call.status(), reqwest::StatusCode::OK);
+    assert_eq!(
+        call.headers()
+            .get(reqwest::header::CONTENT_TYPE)
+            .and_then(|value| value.to_str().ok()),
+        Some("text/event-stream")
+    );
+
+    tokio::time::timeout(Duration::from_secs(5), server.started.notified())
+        .await
+        .expect("tool handler should start");
+    let chunk = tokio::time::timeout(Duration::from_secs(5), call.chunk())
+        .await
+        .expect("a keep-alive comment should arrive while the handler runs")?
+        .expect("the stream should stay open while the handler runs");
+    assert!(
+        chunk.starts_with(b":"),
+        "expected an SSE comment before any message, got {chunk:?}"
+    );
+
+    drop(call);
+    drop(client);
+
+    tokio::time::timeout(Duration::from_secs(10), server.cancelled.notified())
+        .await
+        .expect("RequestContext::ct should fire after client disconnect (negotiated SSE)");
+
+    server.server_ct.cancel();
+    Ok(())
+}
+
+/// Per-request protocol metadata: a handler that answers within one keep-alive
+/// interval keeps the HTTP status its answer maps to, with a JSON body, rather
+/// than a committed event stream carrying the error.
+#[tokio::test]
+async fn stateless_negotiated_answer_within_interval_keeps_its_status() -> anyhow::Result<()> {
+    let server = spawn_stateless_server(false).await?;
+    let client = reqwest::Client::new();
+
+    for id in 0..3 {
+        let answer = client
+            .post(&server.url)
+            .header("Content-Type", "application/json")
+            .header("Accept", "application/json, text/event-stream")
+            .header("MCP-Protocol-Version", "2026-07-28")
+            .header("Mcp-Method", "tools/call")
+            .header("Mcp-Name", "answer_soon")
+            .body(negotiated_call(id, "answer_soon"))
+            .send()
+            .await?;
+        assert_eq!(answer.status(), reqwest::StatusCode::NOT_FOUND);
+        assert_eq!(
+            answer
+                .headers()
+                .get(reqwest::header::CONTENT_TYPE)
+                .and_then(|value| value.to_str().ok()),
+            Some("application/json")
+        );
+        let body: serde_json::Value = serde_json::from_slice(&answer.bytes().await?)?;
+        assert_eq!(body["id"], id);
+        assert_eq!(body["error"]["code"], -32601);
+    }
+
+    server.server_ct.cancel();
+    Ok(())
+}