Skip to main content

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}