a2a_client/components/
streaming.rs

1//! Server-Sent Events (SSE) streaming components
2
3use a2a_rs::services::{AsyncA2AClient, StreamItem};
4use axum::response::sse::{Event, KeepAlive, Sse};
5use futures::StreamExt;
6use std::{convert::Infallible, sync::Arc, time::Duration};
7use tracing::{error, info, warn};
8
9use crate::WebA2AClient;
10
11/// Create an SSE stream for task updates
12///
13/// This function handles:
14/// - WebSocket streaming if available
15/// - Fallback to HTTP polling
16/// - Automatic retry logic
17/// - Serialization to JSON events
18pub fn create_sse_stream(
19    client: Arc<WebA2AClient>,
20    task_id: String,
21) -> Sse<impl futures::Stream<Item = Result<Event, Infallible>>> {
22    let stream = async_stream::stream! {
23        info!("Attempting to subscribe to task {} via HTTP stream", task_id);
24
25        let mut retry_count: u32 = 0;
26        let max_retries = 15; // Covers ~2 minutes using exponential backoff
27        let base_delay = Duration::from_millis(500);
28        let max_delay = Duration::from_secs(10);
29        let mut is_terminal = false;
30
31        loop {
32            match client.http.subscribe_to_task(&task_id, Some(50)).await {
33                Ok(mut event_stream) => {
34                    info!("Successfully subscribed to task {} via HTTP stream", task_id);
35                    retry_count = 0; // Reset retries on successful connection
36
37                    while let Some(result) = event_stream.next().await {
38                        match result {
39                            Ok(stream_item) => {
40                                use a2a_rs::domain::TaskStateExt;
41                                let (event_type, event_data) = match &stream_item {
42                                    StreamItem::Task(task) => {
43                                        if let Some(status) = task.status.as_option() {
44                                            if status.state.is_terminal() {
45                                                is_terminal = true;
46                                            }
47                                        }
48                                        match serde_json::to_string(task) {
49                                            Ok(json) => ("task-update", json),
50                                            Err(e) => {
51                                                error!("Failed to serialize task: {}", e);
52                                                continue;
53                                            }
54                                        }
55                                    }
56                                    StreamItem::StatusUpdate(status) => {
57                                        if status.status.state.is_terminal() {
58                                            is_terminal = true;
59                                        }
60                                        match serde_json::to_string(status) {
61                                            Ok(json) => ("task-status", json),
62                                            Err(e) => {
63                                                error!("Failed to serialize status: {}", e);
64                                                continue;
65                                            }
66                                        }
67                                    }
68                                    StreamItem::ArtifactUpdate(artifact) => {
69                                        match serde_json::to_string(artifact) {
70                                            Ok(json) => ("artifact", json),
71                                            Err(e) => {
72                                                error!("Failed to serialize artifact: {}", e);
73                                                continue;
74                                            }
75                                        }
76                                    }
77                                };
78
79                                yield Ok(Event::default()
80                                    .event(event_type)
81                                    .data(event_data));
82                            }
83                            Err(e) => {
84                                warn!("Stream error (continuing): {}", e);
85                                continue;
86                            }
87                        }
88                    }
89
90                    if is_terminal {
91                        info!("Task {} reached terminal state. Ending stream gracefully.", task_id);
92                        break;
93                    } else {
94                        warn!("Stream ended prematurely for task {}. Retrying...", task_id);
95                        // Do not break; it will loop around and reconnect.
96                    }
97                }
98                Err(e) => {
99                    retry_count += 1;
100
101                    if retry_count <= max_retries {
102                        // Exponential delay: base_delay * 2^(retry_count - 1)
103                        let factor = 2u64.saturating_pow(retry_count.saturating_sub(1).min(6));
104                        let delay = base_delay.saturating_mul(factor as u32);
105
106                        // Jitter calculation (0-200ms) based on task_id bytes and system time
107                        let jitter_ms = {
108                            let mut state = 0u64;
109                            for &b in task_id.as_bytes() {
110                                state = state.wrapping_mul(6364136223846793005).wrapping_add(b as u64);
111                            }
112                            let time_ms = std::time::SystemTime::now()
113                                .duration_since(std::time::UNIX_EPOCH)
114                                .unwrap_or_default()
115                                .as_millis() as u64;
116                            state = state.wrapping_mul(6364136223846793005).wrapping_add(time_ms);
117                            state % 200
118                        };
119                        let final_delay = delay.saturating_add(Duration::from_millis(jitter_ms)).min(max_delay);
120
121                        warn!(
122                            "Failed to subscribe to task {} (attempt {}/{}): {}. Retrying in {:?}...",
123                            task_id, retry_count, max_retries, e, final_delay
124                        );
125
126                        tokio::time::sleep(final_delay).await;
127                        continue;
128                    } else {
129                        warn!("Failed to subscribe after {} retries: {}, aborting stream", max_retries, e);
130                        break;
131                    }
132                }
133            }
134        }
135    };
136
137    Sse::new(stream).keep_alive(KeepAlive::default())
138}
a2a_client/components/streaming.rs

a2a_client/components/
streaming.rs
