fix(http): use raw SSE stream for POST responses

The SseAutoReconnectStream wrapper introduces a 3-second reconnect
delay (from the priming event retry field) when the inner stream
ends, preventing the 50ms drain from completing. Use a thin
raw_sse_to_jsonrpc adapter for per-request POST SSE streams instead,
so the stream ends immediately when the server closes the channel.

Author: DaleSeo

crates/rmcp/src/transport/streamable_http_client.rs

@@ -281,6 +281,39 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
 }
 
 impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
+    /// Convert a raw SSE stream into a JSON-RPC message stream without
+    /// reconnection logic.  Used for per-request POST SSE responses where
+    /// we close the stream after the first response and want the underlying
+    /// HTTP connection to be returned to the pool promptly.
+    fn raw_sse_to_jsonrpc(
+        stream: BoxedSseStream,
+    ) -> impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> + Send + 'static
+    {
+        stream.filter_map(|event| async {
+            match event {
+                Err(e) => Some(Err(StreamableHttpError::Sse(e))),
+                Ok(sse) => {
+                    let is_message =
+                        matches!(sse.event.as_deref(), None | Some("") | Some("message"));
+                    if !is_message {
+                        return None;
+                    }
+                    let data = sse.data?;
+                    if data.trim().is_empty() {
+                        return None;
+                    }
+                    match serde_json::from_str::<ServerJsonRpcMessage>(&data) {
+                        Ok(msg) => Some(Ok(msg)),
+                        Err(e) => {
+                            tracing::debug!("failed to deserialize server message: {e}");
+                            None
+                        }
+                    }
+                }
+            }
+        })
+    }
+
     async fn execute_sse_stream(
         sse_stream: impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>>
         + Send
@@ -730,38 +763,12 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                                                 Ok(())
                                             }
                                             Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
-                                                if let Some(sid) = &session_id {
-                                                    let sse_stream = SseAutoReconnectStream::new(
-                                                        stream,
-                                                        StreamableHttpClientReconnect {
-                                                            client: self.client.clone(),
-                                                            session_id: sid.clone(),
-                                                            uri: config.uri.clone(),
-                                                            auth_header: config.auth_header.clone(),
-                                                            custom_headers: protocol_headers
-                                                                .clone(),
-                                                        },
-                                                        self.config.retry_config.clone(),
-                                                    );
-                                                    streams.spawn(Self::execute_sse_stream(
-                                                        sse_stream,
-                                                        sse_worker_tx.clone(),
-                                                        true,
-                                                        transport_task_ct.child_token(),
-                                                    ));
-                                                } else {
-                                                    let sse_stream =
-                                                    SseAutoReconnectStream::never_reconnect(
-                                                        stream,
-                                                        StreamableHttpError::<C::Error>::UnexpectedEndOfStream,
-                                                    );
-                                                    streams.spawn(Self::execute_sse_stream(
-                                                        sse_stream,
-                                                        sse_worker_tx.clone(),
-                                                        true,
-                                                        transport_task_ct.child_token(),
-                                                    ));
-                                                }
+                                                streams.spawn(Self::execute_sse_stream(
+                                                    Self::raw_sse_to_jsonrpc(stream),
+                                                    sse_worker_tx.clone(),
+                                                    true,
+                                                    transport_task_ct.child_token(),
+                                                ));
                                                 tracing::trace!("got new sse stream after re-init");
                                                 Ok(())
                                             }
@@ -781,36 +788,16 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                             Ok(())
                         }
                         Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
-                            if let Some(session_id) = &session_id {
-                                let sse_stream = SseAutoReconnectStream::new(
-                                    stream,
-                                    StreamableHttpClientReconnect {
-                                        client: self.client.clone(),
-                                        session_id: session_id.clone(),
-                                        uri: config.uri.clone(),
-                                        auth_header: config.auth_header.clone(),
-                                        custom_headers: protocol_headers.clone(),
-                                    },
-                                    self.config.retry_config.clone(),
-                                );
-                                streams.spawn(Self::execute_sse_stream(
-                                    sse_stream,
-                                    sse_worker_tx.clone(),
-                                    true,
-                                    transport_task_ct.child_token(),
-                                ));
-                            } else {
-                                let sse_stream = SseAutoReconnectStream::never_reconnect(
-                                    stream,
-                                    StreamableHttpError::<C::Error>::UnexpectedEndOfStream,
-                                );
-                                streams.spawn(Self::execute_sse_stream(
-                                    sse_stream,
-                                    sse_worker_tx.clone(),
-                                    true,
-                                    transport_task_ct.child_token(),
-                                ));
-                            }
+                            // Per-request POST SSE streams use a thin
+                            // adapter instead of SseAutoReconnectStream so
+                            // the stream ends immediately when the server
+                            // closes the channel, enabling connection reuse.
+                            streams.spawn(Self::execute_sse_stream(
+                                Self::raw_sse_to_jsonrpc(stream),
+                                sse_worker_tx.clone(),
+                                true,
+                                transport_task_ct.child_token(),
+                            ));
                             tracing::trace!("got new sse stream");
                             Ok(())
                         }