diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index 8f0d9eeab..92e9b07cb 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -219,6 +219,17 @@ type PendingRestores = Arc, /// The retry interval for SSE priming events. pub sse_retry: Option, @@ -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,12 +1707,30 @@ 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(); @@ -1708,6 +1738,12 @@ where 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( diff --git a/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs index 05efe2110..67c73ae57 100644 --- a/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs +++ b/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, @@ -43,9 +48,13 @@ impl ServerHandler for CancelProbe { async fn call_tool( &self, - _request: CallToolRequestParams, + request: CallToolRequestParams, context: RequestContext, ) -> Result { + if request.name == "answer_soon" { + tokio::time::sleep(KEEP_ALIVE / 2).await; + return Err(McpError::method_not_found::()); + } 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 = @@ -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(()) +}