fix(a2a): parse multi-line SSE events

crates/aisix-a2a/src/bridge.rs

@@ -590,17 +590,17 @@ impl A2aBridge for HttpBridge {
     }
 }
 
-/// Parse an upstream SSE body into the JSON-RPC envelope of each `data:` field.
+/// Parse an upstream SSE body into one JSON-RPC envelope per event.
 ///
-/// Deliberately minimal: A2A carries one JSON-RPC envelope per `data:` line, so
-/// `event:` / `id:` / `retry:` fields and comments are metadata this gateway has
-/// no use for and passes over. A `data:` line that is not JSON ends the stream
-/// with an error rather than being skipped — a caller that silently dropped
-/// events would report a truncated task as a complete one.
+/// Join an event's `data:` fields with newlines before parsing, as specified by
+/// https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation.
+/// Metadata and comments are ignored. Malformed JSON fails the stream once the
+/// event ends, rather than silently reporting a truncated task as complete.
 fn sse_events(resp: reqwest::Response) -> impl futures::Stream<Item = A2aEvent> + Send {
     async_stream::stream! {
         let mut bytes = resp.bytes_stream();
         let mut pending: Vec<u8> = Vec::new();
+        let mut line_start = 0;
         loop {
             let chunk = match bytes.next().await {
                 Some(Ok(chunk)) => chunk,
@@ -610,49 +610,60 @@ fn sse_events(resp: reqwest::Response) -> impl futures::Stream<Item = A2aEvent>
                 }
                 None => break,
             };
-            pending.extend_from_slice(&chunk);
-            // A single event is bounded even though the stream is not: an
-            // upstream that never emits a newline must not grow this buffer
-            // without limit.
-            if pending.len() > MAX_SSE_EVENT_BYTES {
-                yield Err(A2aError::Request(
-                    "upstream SSE event exceeded size cap".to_string(),
-                ));
-                return;
-            }
-            while let Some(newline) = pending.iter().position(|b| *b == b'\n') {
-                let line: Vec<u8> = pending.drain(..=newline).collect();
-                match parse_sse_data_line(&line) {
-                    Ok(Some(event)) => yield Ok(event),
-                    Ok(None) => {}
-                    Err(e) => {
-                        yield Err(e);
-                        return;
+            for byte in chunk {
+                pending.push(byte);
+                // Bound the whole event, including multiple data lines, rather
+                // than a network chunk that may contain many small events.
+                if pending.len() > MAX_SSE_EVENT_BYTES {
+                    yield Err(A2aError::Request(
+                        "upstream SSE event exceeded size cap".to_string(),
+                    ));
+                    return;
+                }
+                if byte == b'\n' {
+                    let line = &pending[line_start..];
+                    if line == b"\n" || line == b"\r\n" {
+                        match parse_sse_frame(&pending) {
+                            Ok(Some(event)) => yield Ok(event),
+                            Ok(None) => {}
+                            Err(e) => {
+                                yield Err(e);
+                                return;
+                            }
+                        }
+                        pending.clear();
                     }
+                    line_start = pending.len();
                 }
             }
         }
         // A body that ends without its final newline still carries an event —
-        // and if that last line is malformed it fails the stream like any
-        // other. Swallowing the error here would make a truncated task read as
-        // a clean end, which is the exact failure the per-line rule exists to
-        // prevent.
-        match parse_sse_data_line(&pending) {
+        // preserve that compatibility, but parse all of its data fields together.
+        // A malformed trailing event must still fail rather than end quietly.
+        match parse_sse_frame(&pending) {
             Ok(Some(event)) => yield Ok(event),
             Ok(None) => {}
             Err(e) => yield Err(e),
         }
     }
 }
 
-/// Extract the JSON-RPC envelope from one SSE line, or `None` when the line
-/// carries no `data:` field.
-fn parse_sse_data_line(line: &[u8]) -> Result<Option<serde_json::Value>, A2aError> {
-    let text = std::str::from_utf8(line)
+/// Extract the JSON-RPC envelope from an SSE event's joined data fields.
+fn parse_sse_frame(frame: &[u8]) -> Result<Option<serde_json::Value>, A2aError> {
+    let text = std::str::from_utf8(frame)
         .map_err(|_| A2aError::Request("upstream SSE event was not valid UTF-8".to_string()))?;
-    let Some(payload) = text.trim_end_matches(['\r', '\n']).strip_prefix("data:") else {
-        return Ok(None);
-    };
+    let payload = text
+        .lines()
+        .filter_map(|line| {
+            let data = if line == "data" {
+                ""
+            } else {
+                line.strip_prefix("data:")?
+            };
+            Some(data.strip_prefix(' ').unwrap_or(data))
+        })
+        .collect::<Vec<_>>()
+        .join("\n");
     let payload = payload.trim();
     if payload.is_empty() {
         return Ok(None);
@@ -784,7 +795,7 @@ mod tests {
 
     #[test]
     fn sse_lines_yield_only_data_payloads() {
-        let event = |line: &str| parse_sse_data_line(line.as_bytes()).unwrap();
+        let event = |line: &str| parse_sse_frame(line.as_bytes()).unwrap();
 
         assert_eq!(
             event("data: {\"jsonrpc\":\"2.0\",\"id\":1}\n").unwrap()["id"],
@@ -803,16 +814,39 @@ mod tests {
         assert!(event("data:\n").is_none());
     }
 
+    #[test]
+    fn sse_frame_joins_data_fields_before_parsing() {
+        for newline in ["\n", "\r\n"] {
+            let frame = [
+                "data: {\"jsonrpc\":\"2.0\",",
+                ": keep-alive",
+                "event: status-update",
+                "data",
+                "data:\"result\":{\"final\":true}}",
+                "",
+                "",
+            ]
+            .join(newline);
+            assert_eq!(
+                parse_sse_frame(frame.as_bytes()).unwrap().unwrap(),
+                serde_json::json!({"jsonrpc": "2.0", "result": {"final": true}})
+            );
+        }
+        // Joining without the required newline would silently turn invalid
+        // JSON into a different, valid string.
+        assert!(parse_sse_frame(b"data: {\"text\":\"hel\ndata: lo\"}\n\n").is_err());
+    }
+
     #[test]
     fn a_data_line_that_is_not_json_is_an_error_not_a_skip() {
         // Silently dropping it would let a truncated task read as a complete
         // one, which is worse than failing the stream.
-        let err = parse_sse_data_line(b"data: not-json\n").unwrap_err();
+        let err = parse_sse_frame(b"data: not-json\n").unwrap_err();
         assert!(
             matches!(err, A2aError::Request(ref m) if m.contains("malformed JSON-RPC event")),
             "got {err:?}"
         );
-        assert!(parse_sse_data_line(b"data: \xff\xfe\n").is_err());
+        assert!(parse_sse_frame(b"data: \xff\xfe\n").is_err());
     }
 
     #[test]

crates/aisix-a2a/tests/upstream_roundtrip.rs

@@ -371,7 +371,7 @@ async fn the_card_fetch_deadline_covers_the_whole_candidate_walk() {
 }
 
 /// An upstream that answers `message/stream` with a real SSE body, written in
-/// awkward chunks: two events in one write, an event split across two writes,
+/// awkward chunks: two events in one write, an event split across several writes,
 /// and comment / `event:` framing in between. Proves the reader reassembles
 /// across chunk boundaries rather than assuming one chunk is one event.
 async fn spawn_streaming_agent() -> SocketAddr {
@@ -388,8 +388,9 @@ async fn spawn_streaming_agent() -> SocketAddr {
                 ": open\ndata: {{\"jsonrpc\":\"2.0\",\"id\":\"s\",\"result\":{{\"seq\":1,\"version\":{seen_version},\"accept\":\"{accept}\",\"headers\":{seen_headers}}}}}\n\n\
                  event: status-update\ndata: {{\"jsonrpc\":\"2.0\",\"id\":\"s\",\"result\":{{\"seq\":2}}}}\n\n"
             )),
-            Ok("data: {\"jsonrpc\":\"2.0\",\"id\":\"s\",\"resu".to_string()),
-            Ok("lt\":{\"seq\":3,\"final\":true}}\n\n".to_string()),
+            Ok("data: {\"jsonrpc\":\"2.0\",\r\ndata: \"id\":\"s\",\r".to_string()),
+            Ok("\n: keep-alive\r\nevent: status-update\r\ndata: \"resu".to_string()),
+            Ok("lt\":{\"seq\":3,\"final\":true}}\r\n\r\n".to_string()),
         ];
         (
             [(axum::http::header::CONTENT_TYPE, "text/event-stream")],
@@ -426,7 +427,7 @@ async fn streams_events_as_they_arrive_across_chunk_boundaries() {
     assert_eq!(events.len(), 3, "got {events:#?}");
     assert_eq!(events[0]["result"]["seq"], 1);
     assert_eq!(events[1]["result"]["seq"], 2);
-    // Reassembled from two writes that split mid-JSON.
+    // One envelope spans multiple data fields and writes, including a split CRLF.
     assert_eq!(events[2]["result"]["seq"], 3);
     assert_eq!(events[2]["result"]["final"], true);
     // A streaming call is still an A2A call: it announces its version and asks
@@ -563,6 +564,45 @@ async fn a_malformed_final_line_fails_the_stream() {
     );
 }
 
+#[tokio::test]
+async fn an_oversized_multiline_event_fails_the_stream() {
+    use futures::StreamExt;
+
+    async fn oversized() -> impl IntoResponse {
+        // Each line fits below the cap, but without a blank line they belong
+        // to one event. Discarding lines as they arrive would evade the cap.
+        let line = format!(": {}\n", "x".repeat(1024));
+        let chunks = (0..16 * 1024).map(move |_| Ok::<_, std::convert::Infallible>(line.clone()));
+        (
+            [(axum::http::header::CONTENT_TYPE, "text/event-stream")],
+            axum::body::Body::from_stream(futures::stream::iter(chunks)),
+        )
+    }
+    let app = Router::new().route("/a2a", post(oversized));
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app.into_make_service())
+            .await
+            .unwrap();
+    });
+
+    let bridge = HttpBridge::new(upstream(format!("http://{addr}/a2a"), A2aAuth::None));
+    let events: Vec<_> = bridge
+        .send_stream(&json!({"jsonrpc":"2.0","id":"s","method":"message/stream"}))
+        .await
+        .expect("stream opens")
+        .collect()
+        .await;
+    server.abort();
+
+    assert_eq!(events.len(), 1, "got {events:#?}");
+    assert!(
+        matches!(&events[0], Err(A2aError::Request(message)) if message.contains("size cap")),
+        "got {events:#?}"
+    );
+}
+
 // ---------------------------------------------------------------------------
 // `forward_client_headers` — the operator names inbound client headers that
 // must reach this agent. The gateway rebuilds the outbound JSON-RPC message,