diff --git a/crates/context_server/src/transport/http.rs b/crates/context_server/src/transport/http.rs index aa47a73fc3..6e2ef73475 100644 --- a/crates/context_server/src/transport/http.rs +++ b/crates/context_server/src/transport/http.rs @@ -300,7 +300,7 @@ impl HttpTransport { data_buffer.clear(); } in_message = false; - } else if let Some(data) = line.strip_prefix("data: ") { + } else if let Some(data) = line.strip_prefix("data:") { // Handle data lines let data = data.trim(); if !data.is_empty() { @@ -412,9 +412,13 @@ impl Drop for HttpTransport { mod tests { use super::*; use async_trait::async_trait; + use futures::FutureExt as _; use gpui::TestAppContext; use parking_lot::Mutex as SyncMutex; - use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::{ + sync::atomic::{AtomicBool, AtomicUsize, Ordering}, + time::Duration, + }; /// A mock token provider that returns a configurable token and tracks /// refresh attempts. @@ -492,6 +496,52 @@ mod tests { .unwrap()) } + #[gpui::test] + async fn test_sse_data_field(cx: &mut TestAppContext) { + for data_prefix in ["data:", "data: "] { + let body = format!( + "id:1\nevent:message\n{data_prefix}{{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{{}}}}\n\n" + ); + let client = make_fake_http_client(move |_req| { + let body = body.clone(); + Box::pin(async move { + Ok(Response::builder() + .status(200) + .header("Content-Type", "text/event-stream;charset=UTF-8") + .body(AsyncBody::from(body.into_bytes())) + .unwrap()) + }) + }); + + let transport = HttpTransport::new( + client, + "http://mcp.example.com/mcp".to_string(), + HashMap::default(), + cx.background_executor.clone(), + ); + + transport + .send(r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#.to_string()) + .await + .expect("send should succeed"); + + let mut responses = transport.receive(); + let next_response = responses.next().fuse(); + let timeout = cx.background_executor.timer(Duration::from_secs(1)).fuse(); + futures::pin_mut!(next_response, timeout); + + let response = futures::select_biased! { + response = next_response => response.expect("expected SSE response"), + _ = timeout => panic!("timed out waiting for SSE response with {data_prefix:?}"), + }; + + assert_eq!( + response, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#, + "unexpected SSE response for {data_prefix:?}", + ); + } + } + #[gpui::test] async fn test_bearer_token_attached_to_requests(cx: &mut TestAppContext) { // Capture the Authorization header from the request.