diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index 84d1dc4f6..19a1fa5d1 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -519,16 +519,14 @@ fn sse_json_stream(response: reqwest::Response) -> LlmJsonStream { while let Some(chunk) = bytes.next().await { match chunk { Ok(buffer) => { - match decoder.push_bytes(&buffer) { - Ok(events) => { - for event in events { - yield Ok(event.data); + for result in decoder.push_bytes_results(&buffer) { + match result { + Ok(event) => yield Ok(event.data), + Err(error) => { + yield Err(error); + return; } } - Err(error) => { - yield Err(error); - return; - } } } Err(error) => { diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index 62535ab2e..d6ad7872c 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -874,6 +874,45 @@ fn structured_upstream_failure_classification_matches_retry_policy() { ); } +#[tokio::test] +async fn sse_json_stream_yields_valid_event_before_later_batch_error() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let sse_body = concat!( + "data: {\"chunk\":\"first\"}\n\n", + "data: {not valid json}\n\n" + ); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + sse_body.len(), + sse_body + ); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 1024]; + let _ = socket.read(&mut request).await.unwrap(); + socket.write_all(response.as_bytes()).await.unwrap(); + }); + + let response = test_http_client() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let mut stream = sse_json_stream(response); + + assert_eq!( + stream.next().await.unwrap().unwrap(), + json!({"chunk": "first"}) + ); + let error = stream.next().await.unwrap().unwrap_err().to_string(); + assert!(error.contains("SSE data payload"), "{error}"); + assert!(error.contains("not valid json"), "{error}"); + assert!(stream.next().await.is_none()); + + server.await.unwrap(); +} + #[tokio::test] async fn retry_aware_buffered_body_read_failure_stays_structured() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/crates/core/src/codec/streaming.rs b/crates/core/src/codec/streaming.rs index 92bf2a08d..e1bbbd777 100644 --- a/crates/core/src/codec/streaming.rs +++ b/crates/core/src/codec/streaming.rs @@ -128,21 +128,39 @@ impl SseEventDecoder { /// or skip the frame; frames with no `data:` line at all (e.g. SSE heartbeats) are silently /// dropped. pub fn push_bytes(&mut self, bytes: &[u8]) -> Result> { + self.push_bytes_results(bytes).into_iter().collect() + } + + /// Appends `bytes` and returns each completed frame's result in wire order. + /// + /// Unlike [`Self::push_bytes`], this preserves successful events that precede a malformed + /// frame in the same byte batch. Decoding stops after the first error so callers can emit the + /// preceding events, surface the error, and terminate the stream. + pub fn push_bytes_results(&mut self, bytes: &[u8]) -> Vec> { // Normalize CRLF to LF on append so the framing search only needs to find `\n\n`. Some // providers emit mixed line endings on the wire; normalizing once here keeps the inner - // loop cheap. + // loop cheap. If CRLF is split across chunks, retain the trailing CR until the next append + // and remove it only when the next byte completes the sequence. + if self.buffer.ends_with('\r') && bytes.first() == Some(&b'\n') { + self.buffer.pop(); + } let chunk = String::from_utf8_lossy(bytes).replace("\r\n", "\n"); self.buffer.push_str(&chunk); - let mut events = Vec::new(); + let mut results = Vec::new(); while let Some(cut) = self.buffer.find("\n\n") { let frame: String = self.buffer.drain(..cut).collect(); // Drop the `\n\n` terminator itself. self.buffer.drain(..2); - if let Some(event) = parse_sse_frame(&frame)? { - events.push(event); + match parse_sse_frame(&frame) { + Ok(Some(event)) => results.push(Ok(event)), + Ok(None) => {} + Err(error) => { + results.push(Err(error)); + break; + } } } - Ok(events) + results } /// Drains any remaining buffered frame at end of stream. diff --git a/crates/core/tests/unit/codec/streaming_tests.rs b/crates/core/tests/unit/codec/streaming_tests.rs index 07742c899..3a9856429 100644 --- a/crates/core/tests/unit/codec/streaming_tests.rs +++ b/crates/core/tests/unit/codec/streaming_tests.rs @@ -31,6 +31,23 @@ fn buffers_partial_frames_across_pushes() { assert_eq!(events[0].data, json!({"a": 1})); } +#[test] +fn normalizes_crlf_terminator_split_across_pushes() { + let mut decoder = SseEventDecoder::new(); + assert!( + decoder + .push_bytes_results(b"data: {\"a\":1}\r\n\r") + .is_empty() + ); + + let results = decoder.push_bytes_results(b"\n"); + assert_eq!(results.len(), 1); + assert_eq!( + results.into_iter().next().unwrap().unwrap().data, + json!({"a": 1}) + ); +} + #[test] fn drops_frames_without_data_lines() { let mut decoder = SseEventDecoder::new(); @@ -76,3 +93,22 @@ fn surfaces_parse_errors_with_payload_context() { assert!(message.contains("SSE data payload"), "{message}"); assert!(message.contains("not valid json"), "{message}"); } + +#[test] +fn preserves_successes_before_a_later_parse_error() { + let mut decoder = SseEventDecoder::new(); + let mut results = decoder + .push_bytes_results( + b"event: good\ndata: {\"chunk\":\"first\"}\n\nevent: bad\ndata: {not valid json}\n\n", + ) + .into_iter(); + + let first = results.next().unwrap().unwrap(); + assert_eq!(first.event.as_deref(), Some("good")); + assert_eq!(first.data, json!({"chunk": "first"})); + + let error = results.next().unwrap().unwrap_err().to_string(); + assert!(error.contains("SSE data payload"), "{error}"); + assert!(error.contains("not valid json"), "{error}"); + assert!(results.next().is_none()); +}