a2a_client/components/
streaming.rs

1//! Server-Sent Events (SSE) streaming components
2
3use a2a_rs::{RetryPolicy, StreamItem};
4use axum::response::sse::{Event, KeepAlive, Sse};
5use futures::StreamExt;
6use std::{convert::Infallible, sync::Arc};
7use tracing::{error, warn};
8
9use crate::WebA2AClient;
10
11/// Create an SSE stream of task updates for an Axum endpoint.
12///
13/// This is a thin serialization adapter over
14/// [`WebA2AClient::subscribe_resilient`]: the reusable core owns reconnect +
15/// exponential backoff and `Last-Event-ID` resumption, so this function only
16/// maps each [`StreamItem`] to a typed Axum [`Event`] (`task-update` /
17/// `task-status` / `artifact`), tagging it with the server event id so a browser
18/// `EventSource` resumes automatically. The stream ends when the task reaches a
19/// terminal state (or retries are exhausted).
20pub fn create_sse_stream(
21    client: Arc<WebA2AClient>,
22    task_id: String,
23) -> Sse<impl futures::Stream<Item = Result<Event, Infallible>>> {
24    let updates = client.subscribe_resilient(&task_id, RetryPolicy::default());
25
26    let stream = updates.filter_map(|result| async move {
27        let event = match result {
28            Ok(event) => event,
29            Err(e) => {
30                warn!("Stream error (ending): {}", e);
31                return None;
32            }
33        };
34
35        let (event_type, data) = match &event.item {
36            StreamItem::Task(task) => ("task-update", serde_json::to_string(task)),
37            StreamItem::StatusUpdate(status) => ("task-status", serde_json::to_string(status)),
38            StreamItem::ArtifactUpdate(artifact) => ("artifact", serde_json::to_string(artifact)),
39        };
40
41        match data {
42            Ok(json) => {
43                let mut sse = Event::default().event(event_type).data(json);
44                if let Some(id) = event.event_id {
45                    sse = sse.id(id.to_string());
46                }
47                Some(Ok(sse))
48            }
49            Err(e) => {
50                error!("Failed to serialize {event_type}: {e}");
51                None
52            }
53        }
54    });
55
56    Sse::new(stream).keep_alive(KeepAlive::default())
57}
a2a_client/components/streaming.rs

a2a_client/components/
streaming.rs
