fix(streamable-http): preserve progress in JSON mode (#990)

Author: jstar0

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

@@ -19,6 +19,7 @@ use crate::{
         ClientCapabilities, ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData,
         GetExtensions, Implementation, InitializeRequest, InitializeRequestParams,
         InitializedNotification, JsonObject, JsonRpcError, ProtocolVersion, RequestId,
+        ServerJsonRpcMessage,
     },
     serve_server,
     service::serve_directly,
@@ -48,10 +49,10 @@ pub struct StreamableHttpServerConfig {
     /// If true, the server will create a session for each request and keep it alive.
     /// When enabled, SSE priming events are sent to enable client reconnection.
     pub stateful_mode: bool,
-    /// When true and `stateful_mode` is false, the server returns
-    /// `Content-Type: application/json` directly instead of `text/event-stream`.
-    /// This eliminates SSE framing overhead for simple request-response tools,
-    /// allowed by the MCP Streamable HTTP spec (2025-06-18).
+    /// When true and `stateful_mode` is false, the server prefers
+    /// `Content-Type: application/json` for simple request-response tools.
+    /// If the handler emits a notification or request before the final response,
+    /// the server falls back to `text/event-stream` so no message is lost.
     pub json_response: bool,
     /// Cancellation token for the Streamable HTTP server.
     ///
@@ -1352,31 +1353,47 @@ where
                         let _ = service.waiting().await;
                     });
                     if self.config.json_response {
-                        // JSON-direct mode: await the single response and return as
-                        // application/json, eliminating SSE framing overhead.
-                        // Allowed by MCP Streamable HTTP spec (2025-06-18).
+                        // Prefer JSON for a terminal first message. If the handler
+                        // emits an intermediate notification or request, preserve
+                        // the complete message sequence by falling back to SSE.
                         let cancel = self.config.cancellation_token.child_token();
-                        match tokio::select! {
+                        let Some(message) = (tokio::select! {
                             res = receiver.recv() => res,
                             _ = cancel.cancelled() => None,
-                        } {
-                            Some(message) => {
-                                tracing::trace!(?message);
-                                let body = serde_json::to_vec(&message).map_err(|e| {
-                                    internal_error_response("serialize json response")(e)
-                                })?;
-                                Ok(Response::builder()
-                                    .status(http::StatusCode::OK)
-                                    .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE)
-                                    .body(Full::new(Bytes::from(body)).boxed())
-                                    .expect("valid response"))
-                            }
-                            None => Err(internal_error_response("empty response")(
+                        }) else {
+                            return Err(internal_error_response("empty response")(
                                 std::io::Error::new(
                                     std::io::ErrorKind::UnexpectedEof,
                                     "no response message received from handler",
                                 ),
-                            )),
+                            ));
+                        };
+                        tracing::trace!(?message);
+                        if matches!(
+                            message,
+                            ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_)
+                        ) {
+                            let body = serde_json::to_vec(&message).map_err(|e| {
+                                internal_error_response("serialize json response")(e)
+                            })?;
+                            Ok(Response::builder()
+                                .status(http::StatusCode::OK)
+                                .header(http::header::CONTENT_TYPE, JSON_MIME_TYPE)
+                                .body(Full::new(Bytes::from(body)).boxed())
+                                .expect("valid response"))
+                        } else {
+                            let first = futures::stream::once(async move {
+                                ServerSseMessage::from_message(message)
+                            });
+                            let remaining = ReceiverStream::new(receiver).map(|message| {
+                                tracing::trace!(?message);
+                                ServerSseMessage::from_message(message)
+                            });
+                            Ok(sse_stream_response(
+                                first.chain(remaining),
+                                self.config.sse_keep_alive,
+                                self.config.cancellation_token.child_token(),
+                            ))
                         }
                     } else {
                         // SSE mode (default): original behaviour preserved unchanged

crates/rmcp/tests/test_streamable_http_json_response.rs

@@ -1,13 +1,52 @@
 #![cfg(not(feature = "local"))]
-use rmcp::transport::streamable_http_server::{
-    StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
+use rmcp::{
+    ErrorData, ServerHandler,
+    model::{
+        CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock,
+        ProgressNotificationParam, ServerCapabilities, ServerInfo,
+    },
+    service::RequestContext,
+    transport::streamable_http_server::{
+        StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
+    },
 };
 use tokio_util::sync::CancellationToken;
 
 mod common;
 use common::calculator::Calculator;
 
 const INIT_BODY: &str = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#;
+const CALL_WITH_PROGRESS_BODY: &str = r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"progress","arguments":{},"_meta":{"progressToken":"progress-test-1"}}}"#;
+
+#[derive(Clone)]
+struct ProgressServer;
+
+impl ServerHandler for ProgressServer {
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
+    }
+
+    async fn call_tool(
+        &self,
+        _request: CallToolRequestParams,
+        context: RequestContext<rmcp::RoleServer>,
+    ) -> Result<CallToolResponse, ErrorData> {
+        let progress_token = context
+            .meta
+            .get_progress_token()
+            .expect("request includes progressToken");
+        context
+            .peer
+            .notify_progress(
+                ProgressNotificationParam::new(progress_token, 50.0)
+                    .with_total(100.0)
+                    .with_message("working"),
+            )
+            .await
+            .expect("progress notification is delivered");
+        Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into())
+    }
+}
 
 async fn spawn_server(
     config: StreamableHttpServerConfig,
@@ -34,6 +73,31 @@ async fn spawn_server(
     (client, base_url, ct)
 }
 
+async fn spawn_progress_server(
+    config: StreamableHttpServerConfig,
+) -> (reqwest::Client, String, CancellationToken) {
+    let ct = config.cancellation_token.clone();
+    let service: StreamableHttpService<ProgressServer, LocalSessionManager> =
+        StreamableHttpService::new(|| Ok(ProgressServer), Default::default(), config);
+
+    let router = axum::Router::new().nest_service("/mcp", service);
+    let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = tcp_listener.local_addr().unwrap();
+
+    tokio::spawn({
+        let ct = ct.clone();
+        async move {
+            let _ = axum::serve(tcp_listener, router)
+                .with_graceful_shutdown(async move { ct.cancelled_owned().await })
+                .await;
+        }
+    });
+
+    let client = reqwest::Client::new();
+    let base_url = format!("http://{addr}/mcp");
+    (client, base_url, ct)
+}
+
 #[tokio::test]
 async fn stateless_json_response_returns_application_json() -> anyhow::Result<()> {
     let ct = CancellationToken::new();
@@ -76,6 +140,58 @@ async fn stateless_json_response_returns_application_json() -> anyhow::Result<()
     Ok(())
 }
 
+#[tokio::test]
+async fn stateless_json_response_falls_back_to_sse_for_progress() -> anyhow::Result<()> {
+    let ct = CancellationToken::new();
+    let (client, url, ct) = spawn_progress_server(
+        StreamableHttpServerConfig::default()
+            .with_stateful_mode(false)
+            .with_json_response(true)
+            .with_sse_keep_alive(None)
+            .with_cancellation_token(ct.child_token()),
+    )
+    .await;
+
+    let response = client
+        .post(&url)
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .body(CALL_WITH_PROGRESS_BODY)
+        .send()
+        .await?;
+
+    assert_eq!(response.status(), 200);
+
+    let content_type = response
+        .headers()
+        .get("content-type")
+        .and_then(|value| value.to_str().ok())
+        .unwrap_or("");
+    assert!(
+        content_type.contains("text/event-stream"),
+        "Expected SSE fallback, got: {content_type}"
+    );
+
+    let body = response.text().await?;
+    let messages: Vec<serde_json::Value> = body
+        .lines()
+        .filter_map(|line| line.strip_prefix("data:"))
+        .map(str::trim)
+        .filter(|data| !data.is_empty())
+        .map(serde_json::from_str)
+        .collect::<Result<_, _>>()?;
+    assert_eq!(messages.len(), 2, "Expected progress and result: {body}");
+    assert_eq!(messages[0]["method"], "notifications/progress");
+    assert_eq!(messages[1]["id"], 2);
+    assert!(
+        messages[1]["result"].is_object(),
+        "Expected result object: {body}"
+    );
+
+    ct.cancel();
+    Ok(())
+}
+
 #[tokio::test]
 async fn stateless_sse_mode_default_unchanged() -> anyhow::Result<()> {
     let ct = CancellationToken::new();