From 4dd5126120b85c6903a9eb40fd4c2ec4e6b4dcbc Mon Sep 17 00:00:00 2001 From: Mega Mind <146339422+M3gA-Mind@users.noreply.github.com> Date: Mon, 29 Jun 2026 09:57:52 +0530 Subject: [PATCH] perf(mcp): return MCP tool results as soon as the reply frame arrives (#4220) --- src/openhuman/mcp_client/client.rs | 35 +++++++++++++--- src/openhuman/mcp_client/client_helpers.rs | 28 +++++++++++++ src/openhuman/mcp_client/client_tests.rs | 49 ++++++++++++++++++++++ 3 files changed, 107 insertions(+), 5 deletions(-) diff --git a/src/openhuman/mcp_client/client.rs b/src/openhuman/mcp_client/client.rs index 4319bac85..de04612ac 100644 --- a/src/openhuman/mcp_client/client.rs +++ b/src/openhuman/mcp_client/client.rs @@ -2,6 +2,7 @@ use crate::openhuman::config::{McpAuthConfig, McpClientIdentityConfig}; use crate::openhuman::workflows::types::ToolResult; use anyhow::Context; use base64::Engine; +use futures_util::StreamExt; use parking_lot::Mutex; #[cfg(test)] use reqwest::header::HeaderMap; @@ -746,8 +747,8 @@ pub fn redact_endpoint(raw: &str) -> String { #[path = "client_helpers.rs"] mod client_helpers; use client_helpers::{ - header_to_string, parse_sse_events, parse_sse_message, parse_www_authenticate_challenge, - x_mcp_headers_from_schema, + first_complete_sse_data, header_to_string, parse_sse_events, parse_sse_message, + parse_www_authenticate_challenge, x_mcp_headers_from_schema, }; impl McpHttpClient { @@ -759,8 +760,6 @@ impl McpHttpClient { .and_then(|v| v.to_str().ok()) .unwrap_or("") .to_string(); - let text = response.text().await?; - if status == reqwest::StatusCode::UNAUTHORIZED { let resource_metadata = parse_www_authenticate_challenge(&headers) .and_then(|challenge| challenge.resource_metadata); @@ -774,12 +773,38 @@ impl McpHttpClient { })); } if !status.is_success() { + let text = response.text().await.unwrap_or_default(); anyhow::bail!("MCP HTTP {} — {}", status.as_u16(), text); } let payload: Value = if content_type.starts_with("text/event-stream") { - parse_sse_message(&text)? + // Read the SSE body incrementally and return as soon as the single + // JSON-RPC reply frame arrives, instead of buffering to stream-close + // (`response.text().await`). A server that holds the stream open + // after replying would otherwise stall this call until the request + // timeout, and those stalls compound across a skill's tool chain + // (#4195). The request-level reqwest timeout still bounds the worst + // case (a server that never replies). + let mut raw: Vec = Vec::new(); + let mut stream = response.bytes_stream(); + let mut frame: Option = None; + while let Some(chunk) = stream.next().await { + raw.extend_from_slice(&chunk?); + // Decode the whole buffer each pass so a multi-byte UTF-8 + // sequence split across chunk boundaries is never corrupted. + if let Some(data) = first_complete_sse_data(&String::from_utf8_lossy(&raw))? { + frame = Some(data); + break; + } + } + match frame { + Some(data) => data, + // Stream ended without a data frame — fall back to the whole-body + // parser for a clear "no data frame" error that includes the body. + None => parse_sse_message(&String::from_utf8_lossy(&raw))?, + } } else { + let text = response.text().await?; serde_json::from_str(&text).map_err(|e| { anyhow::anyhow!("Failed to parse MCP JSON response: {e} — body: {text}") })? diff --git a/src/openhuman/mcp_client/client_helpers.rs b/src/openhuman/mcp_client/client_helpers.rs index baa83558c..4c4ffae02 100644 --- a/src/openhuman/mcp_client/client_helpers.rs +++ b/src/openhuman/mcp_client/client_helpers.rs @@ -13,6 +13,34 @@ pub(super) fn parse_sse_message(body: &str) -> anyhow::Result { Ok(event) } +/// Return the first JSON `data:` frame from the **fully-terminated** prefix of a +/// partially-received SSE buffer, or `None` if no complete event carrying a data +/// frame has arrived yet. +/// +/// MCP Streamable HTTP lets a server keep the POST response's SSE stream open +/// after it has already emitted the single JSON-RPC reply, so reading the whole +/// body to stream-close (`response.text().await`) stalls every tool call until +/// the server closes or the request times out — the dominant transport cause of +/// the multi-minute skill latency in #4195. Reading incrementally and stopping +/// at the first data frame lets the call return the instant the reply lands. +/// +/// Only the portion up to the last blank-line event boundary is parsed, so a +/// half-received final `data:` line is never decoded as truncated JSON. +/// Keepalive comments and dataless events are skipped (returns `None`, keep +/// reading). An SSE event terminates on a blank line; `\r\n` is normalised to +/// `\n` first so CRLF streams split on the same boundary. +pub(super) fn first_complete_sse_data(buffer: &str) -> anyhow::Result> { + let normalized = buffer.replace("\r\n", "\n"); + // The terminator of the last complete event. Without one, nothing is + // fully received yet. + let Some(boundary) = normalized.rfind("\n\n") else { + return Ok(None); + }; + let complete = &normalized[..=boundary]; + let events = parse_sse_events(complete)?; + Ok(events.into_iter().find_map(|event| event.data)) +} + pub(super) fn parse_sse_events(body: &str) -> anyhow::Result> { let mut events = Vec::new(); let mut event_type: Option = None; diff --git a/src/openhuman/mcp_client/client_tests.rs b/src/openhuman/mcp_client/client_tests.rs index 9edfa0a36..47024df7c 100644 --- a/src/openhuman/mcp_client/client_tests.rs +++ b/src/openhuman/mcp_client/client_tests.rs @@ -394,6 +394,55 @@ fn parse_sse_events_handles_multiple_frames() { assert_eq!(events[1].data.as_ref().unwrap()["b"], 2); } +// #4195 — the incremental SSE reader must surface the JSON-RPC reply the moment +// a complete `data:` frame is buffered, so a server that holds the stream open +// after replying no longer stalls the tool call until the request timeout. + +#[test] +fn first_complete_sse_data_returns_none_until_event_terminated() { + // The data line has arrived but the terminating blank line has not — a + // half-received frame must NOT be parsed (it could be truncated JSON). + assert!(first_complete_sse_data("event: message\ndata: {\"a\":1}\n") + .expect("ok") + .is_none()); + // Nothing complete at all. + assert!(first_complete_sse_data("event: mess") + .expect("ok") + .is_none()); +} + +#[test] +fn first_complete_sse_data_returns_first_complete_frame() { + // A fully-terminated event yields its data immediately, even though more + // bytes (here, the start of a second frame) trail behind it. + let buffer = "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"ok\":true}}\n\ndata: {\"b\":2}"; + let data = first_complete_sse_data(buffer) + .expect("ok") + .expect("first complete frame"); + assert_eq!(data["result"]["ok"], true); +} + +#[test] +fn first_complete_sse_data_skips_keepalive_and_dataless_events() { + // Leading SSE comment + a dataless event must be skipped, returning the + // first event that actually carries a data frame. + let buffer = ": keepalive\n\nevent: ping\n\ndata: {\"id\":7}\n\n"; + let data = first_complete_sse_data(buffer) + .expect("ok") + .expect("data frame after keepalive"); + assert_eq!(data["id"], 7); +} + +#[test] +fn first_complete_sse_data_handles_crlf_boundaries() { + // CRLF streams must split on the same blank-line boundary. + let buffer = "event: message\r\ndata: {\"id\":9}\r\n\r\n"; + let data = first_complete_sse_data(buffer) + .expect("ok") + .expect("crlf data frame"); + assert_eq!(data["id"], 9); +} + #[test] fn parse_www_authenticate_extracts_resource_metadata() { let mut headers = HeaderMap::new();