lc_a2a/client/sse.rs
1use std::collections::VecDeque;
2
3use crate::protocol::TaskPushNotification;
4
5use super::A2AError;
6
7/// A live SSE connection to an A2A server (P2-1).
8///
9/// Consume it with [`A2ASseStream::next`], which yields one
10/// [`TaskPushNotification`] per complete SSE event until the server closes the
11/// stream.
12pub struct A2ASseStream {
13 response: reqwest::Response,
14 parser: A2aSseParser,
15 pending: VecDeque<TaskPushNotification>,
16}
17
18impl A2ASseStream {
19 pub(crate) fn new(response: reqwest::Response) -> Self {
20 Self {
21 response,
22 parser: A2aSseParser::new(),
23 pending: VecDeque::new(),
24 }
25 }
26
27 /// Wait for the next event, or `None` when the stream ends.
28 ///
29 /// Returns an error if the connection breaks; a malformed event is logged
30 /// and skipped so one bad payload does not kill the whole stream
31 /// (0.22.0 audit fix).
32 pub async fn next(&mut self) -> Option<Result<TaskPushNotification, A2AError>> {
33 loop {
34 if let Some(event) = self.pending.pop_front() {
35 return Some(Ok(event));
36 }
37 match self.response.chunk().await {
38 Ok(Some(chunk)) => {
39 let text = String::from_utf8_lossy(&chunk);
40 self.pending.extend(self.parser.feed(&text));
41 // Loop so queued events are returned even if a chunk
42 // carried none (or we keep reading on an empty chunk).
43 }
44 Ok(None) => return None,
45 Err(e) => return Some(Err(A2AError::from(e))),
46 }
47 }
48 }
49}
50
51/// Incremental parser for A2A SSE event frames.
52///
53/// Mirrors the SSE parsing in `lc-providers/src/openai/sse.rs`: events are
54/// terminated by a blank line (`\n\n`), and the payload is the `data:` field
55/// (multi-line `data:` fields are joined with newlines per the SSE spec).
56struct A2aSseParser {
57 buffer: String,
58}
59
60impl A2aSseParser {
61 fn new() -> Self {
62 Self {
63 buffer: String::new(),
64 }
65 }
66
67 /// Feed a chunk of the response body; returns any complete notifications.
68 ///
69 /// 0.22.0 audit fix: an event whose JSON payload fails to parse is logged
70 /// and skipped; only transport-level errors surface to the caller.
71 fn feed(&mut self, chunk: &str) -> Vec<TaskPushNotification> {
72 self.buffer.push_str(chunk);
73 let mut out = Vec::new();
74 while let Some(pos) = self.buffer.find("\n\n") {
75 let event_text: String = self.buffer[..pos].to_string();
76 self.buffer.drain(..=pos + 1);
77
78 let data = event_text
79 .lines()
80 .filter_map(|line| line.strip_prefix("data:"))
81 .map(|line| line.trim_start().trim_end())
82 .collect::<Vec<_>>()
83 .join("\n");
84 if data.is_empty() || data == "[DONE]" {
85 continue;
86 }
87 match serde_json::from_str::<TaskPushNotification>(&data) {
88 Ok(notification) => out.push(notification),
89 Err(e) => {
90 log::warn!("skipping malformed SSE event `{data}`: {e}");
91 }
92 }
93 }
94 out
95 }
96}