From 7625bec652ecde92fa8926e4b1f6cb93edc6b9d5 Mon Sep 17 00:00:00 2001 From: mat973252-coder <234863146+mat973252-coder@users.noreply.github.com> Date: Wed, 9 Sep 2026 07:34:24 +0800 Subject: [PATCH] fix(a2a): parse multi-line SSE events --- crates/aisix-a2a/src/bridge.rs | 112 ++++++++++++------- crates/aisix-a2a/tests/upstream_roundtrip.rs | 48 +++++++- 2 files changed, 117 insertions(+), 43 deletions(-) diff --git a/crates/aisix-a2a/src/bridge.rs b/crates/aisix-a2a/src/bridge.rs index b20f50ef..e3ef59b3 100644 --- a/crates/aisix-a2a/src/bridge.rs +++ b/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 + Send { async_stream::stream! { let mut bytes = resp.bytes_stream(); let mut pending: Vec = Vec::new(); + let mut line_start = 0; loop { let chunk = match bytes.next().await { Some(Ok(chunk)) => chunk, @@ -610,34 +610,37 @@ fn sse_events(resp: reqwest::Response) -> impl futures::Stream } 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 = 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), @@ -645,14 +648,22 @@ fn sse_events(resp: reqwest::Response) -> impl futures::Stream } } -/// 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, 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, 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::>() + .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] diff --git a/crates/aisix-a2a/tests/upstream_roundtrip.rs b/crates/aisix-a2a/tests/upstream_roundtrip.rs index 50698cfe..88eb66c7 100644 --- a/crates/aisix-a2a/tests/upstream_roundtrip.rs +++ b/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,