Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 6 additions & 8 deletions crates/cli/src/gateway/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) => {
Expand Down
39 changes: 39 additions & 0 deletions crates/cli/tests/coverage/shared/gateway_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
28 changes: 23 additions & 5 deletions crates/core/src/codec/streaming.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Vec<SseEvent>> {
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<Result<SseEvent>> {
// 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);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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.
Expand Down
36 changes: 36 additions & 0 deletions crates/core/tests/unit/codec/streaming_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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());
}
Loading