a2a_rs/adapter/streaming/
fanout.rs

1//! Streaming fan-out adapter over an [`AsyncEventLog`].
2//!
3//! [`StreamingFanout`] is the [`AsyncStreamingHandler`] adapter. It owns the
4//! in-process side of streaming — a broadcast channel per task, plus an optional
5//! set of synchronous callback subscribers — and delegates ids and retention to
6//! the event log it was built over. It deliberately does *not*:
7//!
8//! - touch the task store (so it cannot replay current task state on subscribe —
9//!   the initial `Task` snapshot is delivered by the application service before
10//!   stream items, which is spec-compliant), nor
11//! - fire push-webhook notifications (that is the [`AsyncPushNotifier`] port's
12//!   job, orchestrated by the
13//!   [`TaskStatusBroadcast`](crate::application::TaskStatusBroadcast) mixin).
14//!
15//! The split with the log is what durability turns on. Who is listening right
16//! now is this process's business and a restart may forget it; what the stream
17//! already said is the log's, and a durable log still has it after the restart.
18//! [`InMemoryStreamingHandler`] pairs the fan-out with an in-memory log, which
19//! is the zero-configuration default; `SqlxTaskStorage` implements the same port
20//! for a log that outlives the process.
21//!
22//! [`AsyncPushNotifier`]: crate::port::AsyncPushNotifier
23
24use std::collections::HashMap;
25use std::pin::Pin;
26use std::sync::Arc;
27
28use async_trait::async_trait;
29use futures::{Stream, StreamExt};
30use tokio::sync::Mutex;
31use tokio::sync::broadcast;
32
33use crate::adapter::storage::event_log::InMemoryEventLog;
34use crate::domain::{A2AError, TaskArtifactUpdateEvent, TaskStatusUpdateEvent};
35use crate::port::AsyncStreamingHandler;
36use crate::port::event_log::AsyncEventLog;
37use crate::port::streaming_handler::{SeqEvent, Subscriber, UpdateEvent};
38
39type StatusSubscribers = Vec<Box<dyn Subscriber<TaskStatusUpdateEvent> + Send + Sync>>;
40type ArtifactSubscribers = Vec<Box<dyn Subscriber<TaskArtifactUpdateEvent> + Send + Sync>>;
41
42/// Capacity of the per-task broadcast channel. A reader that falls this far
43/// behind is told it lagged and reconnects; the log, not this channel, is what
44/// it resumes from.
45const CHANNEL_CAPACITY: usize = 256;
46
47/// Per-task in-process state: a broadcast channel for live readers and any
48/// synchronous callback subscribers.
49struct TaskChannel {
50    sender: broadcast::Sender<SeqEvent>,
51    status: StatusSubscribers,
52    artifacts: ArtifactSubscribers,
53}
54
55impl TaskChannel {
56    fn new() -> Self {
57        let (sender, _) = broadcast::channel(CHANNEL_CAPACITY);
58        Self {
59            sender,
60            status: Vec::new(),
61            artifacts: Vec::new(),
62        }
63    }
64}
65
66/// Fan-out of task updates to live readers and callback subscribers, over an
67/// [`AsyncEventLog`] that assigns the ids and keeps the events for replay.
68///
69/// Cloning shares the underlying per-task state (an `Arc<Mutex<…>>`) and the
70/// log, so a clone observes the same channels and subscribers.
71#[derive(Clone)]
72pub struct StreamingFanout<L> {
73    /// One lock per task rather than one over the map. The map lock is held
74    /// only long enough to hand out a task's channel; everything that awaits —
75    /// the log write, a callback subscriber — happens under that task's own
76    /// lock, so two tasks streaming at once do not queue behind each other.
77    /// Within a task the lock still orders the log write against the send, which
78    /// is what keeps ids and delivery in the same order.
79    tasks: Arc<Mutex<HashMap<String, Arc<Mutex<TaskChannel>>>>>,
80    log: L,
81}
82
83/// The zero-configuration streaming handler: fan-out over an in-process event
84/// log. Resumption works within one run of a server and no further — see
85/// [`InMemoryEventLog`].
86pub type InMemoryStreamingHandler = StreamingFanout<InMemoryEventLog>;
87
88impl<L> StreamingFanout<L> {
89    /// Build a fan-out over `log`.
90    ///
91    /// The log is what a client resumes from, so pass a durable one where
92    /// resumption has to survive a restart.
93    pub fn over(log: L) -> Self {
94        Self {
95            tasks: Arc::new(Mutex::new(HashMap::new())),
96            log,
97        }
98    }
99
100    /// This task's channel, created if it is the first anyone has asked.
101    async fn channel(&self, task_id: &str) -> Arc<Mutex<TaskChannel>> {
102        self.tasks
103            .lock()
104            .await
105            .entry(task_id.to_string())
106            .or_insert_with(|| Arc::new(Mutex::new(TaskChannel::new())))
107            .clone()
108    }
109}
110
111impl StreamingFanout<InMemoryEventLog> {
112    /// Create a handler over a fresh in-memory log.
113    pub fn new() -> Self {
114        Self::over(InMemoryEventLog::new())
115    }
116}
117
118impl Default for StreamingFanout<InMemoryEventLog> {
119    fn default() -> Self {
120        Self::new()
121    }
122}
123
124impl<L: AsyncEventLog> StreamingFanout<L> {
125    /// Log `event`, then publish it to live readers.
126    ///
127    /// Called with the task's channel locked, which is what keeps the id order
128    /// and the delivery order the same: two concurrent broadcasts cannot take
129    /// their ids in one order and reach the channel in the other.
130    async fn publish(
131        &self,
132        channel: &TaskChannel,
133        task_id: &str,
134        event: UpdateEvent,
135    ) -> Result<(), A2AError> {
136        let seq = self.log.append(task_id, event).await?;
137        // A send error just means there are no live readers; the log still has
138        // the event for a later resume, so it is ignored.
139        let _ = channel.sender.send(seq);
140        Ok(())
141    }
142}
143
144#[async_trait]
145impl<L: AsyncEventLog + Clone + 'static> AsyncStreamingHandler for StreamingFanout<L> {
146    async fn add_status_subscriber(
147        &self,
148        task_id: &str,
149        subscriber: Box<dyn Subscriber<TaskStatusUpdateEvent> + Send + Sync>,
150    ) -> Result<String, A2AError> {
151        #[cfg(feature = "tracing")]
152        tracing::info!(
153            task_id = %task_id,
154            "✅ Adding subscriber for status updates"
155        );
156
157        self.channel(task_id)
158            .await
159            .lock()
160            .await
161            .status
162            .push(subscriber);
163
164        Ok(format!("status-{}-{}", task_id, uuid::Uuid::new_v4()))
165    }
166
167    async fn add_artifact_subscriber(
168        &self,
169        task_id: &str,
170        subscriber: Box<dyn Subscriber<TaskArtifactUpdateEvent> + Send + Sync>,
171    ) -> Result<String, A2AError> {
172        self.channel(task_id)
173            .await
174            .lock()
175            .await
176            .artifacts
177            .push(subscriber);
178
179        Ok(format!("artifact-{}-{}", task_id, uuid::Uuid::new_v4()))
180    }
181
182    async fn remove_subscription(&self, _subscription_id: &str) -> Result<(), A2AError> {
183        Err(A2AError::UnsupportedOperation(
184            "Subscription removal by ID is not supported by the in-memory streaming handler"
185                .to_string(),
186        ))
187    }
188
189    /// Drop the task's live readers and callback subscribers.
190    ///
191    /// The event log is left alone: a task with nobody listening is exactly the
192    /// case a client is about to resume from. Deleting the events is
193    /// [`AsyncEventLog::discard`], which an operator's retention sweep calls.
194    async fn remove_task_subscribers(&self, task_id: &str) -> Result<(), A2AError> {
195        let mut guard = self.tasks.lock().await;
196        guard.remove(task_id);
197        Ok(())
198    }
199
200    async fn get_subscriber_count(&self, task_id: &str) -> Result<usize, A2AError> {
201        let Some(channel) = self.tasks.lock().await.get(task_id).cloned() else {
202            return Ok(0);
203        };
204        let channel = channel.lock().await;
205        Ok(channel.status.len() + channel.artifacts.len() + channel.sender.receiver_count())
206    }
207
208    async fn broadcast_status_update(
209        &self,
210        task_id: &str,
211        update: TaskStatusUpdateEvent,
212    ) -> Result<(), A2AError> {
213        #[cfg(feature = "tracing")]
214        tracing::debug!(
215            task_id = %task_id,
216            state = ?update.status.state,
217            "📡 Broadcasting status update to subscribers"
218        );
219
220        let channel = self.channel(task_id).await;
221        let channel = channel.lock().await;
222        self.publish(&channel, task_id, UpdateEvent::StatusUpdate(update.clone()))
223            .await?;
224        for subscriber in channel.status.iter() {
225            if let Err(e) = subscriber.on_update(update.clone()).await {
226                #[cfg(feature = "tracing")]
227                tracing::error!(task_id = %task_id, error = %e, "❌ Failed to notify subscriber");
228                #[cfg(not(feature = "tracing"))]
229                let _ = e;
230            }
231        }
232        Ok(())
233    }
234
235    async fn broadcast_artifact_update(
236        &self,
237        task_id: &str,
238        update: TaskArtifactUpdateEvent,
239    ) -> Result<(), A2AError> {
240        let channel = self.channel(task_id).await;
241        let channel = channel.lock().await;
242        self.publish(
243            &channel,
244            task_id,
245            UpdateEvent::ArtifactUpdate(update.clone()),
246        )
247        .await?;
248        for subscriber in channel.artifacts.iter() {
249            if let Err(e) = subscriber.on_update(update.clone()).await {
250                #[cfg(feature = "tracing")]
251                tracing::error!(task_id = %task_id, error = %e, "❌ Failed to notify subscriber");
252                #[cfg(not(feature = "tracing"))]
253                let _ = e;
254            }
255        }
256        Ok(())
257    }
258
259    async fn status_update_stream(
260        &self,
261        _task_id: &str,
262    ) -> Result<Pin<Box<dyn Stream<Item = Result<TaskStatusUpdateEvent, A2AError>> + Send>>, A2AError>
263    {
264        Err(A2AError::UnsupportedOperation(
265            "Status-only update stream is not supported; use combined_update_stream".to_string(),
266        ))
267    }
268
269    async fn artifact_update_stream(
270        &self,
271        _task_id: &str,
272    ) -> Result<
273        Pin<Box<dyn Stream<Item = Result<TaskArtifactUpdateEvent, A2AError>> + Send>>,
274        A2AError,
275    > {
276        Err(A2AError::UnsupportedOperation(
277            "Artifact-only update stream is not supported; use combined_update_stream".to_string(),
278        ))
279    }
280
281    async fn combined_update_stream(
282        &self,
283        task_id: &str,
284        from_event_id: Option<u64>,
285    ) -> Result<Pin<Box<dyn Stream<Item = Result<SeqEvent, A2AError>> + Send>>, A2AError> {
286        // Subscribing and reading the log under the task's lock is what makes
287        // the two halves meet exactly: a broadcast cannot land between them, so
288        // the reader neither misses an event nor sees one twice.
289        let channel = self.channel(task_id).await;
290        let guard = channel.lock().await;
291        let receiver = guard.sender.subscribe();
292        let replay = match from_event_id {
293            Some(from) => self.log.replay(task_id, from).await?,
294            None => Default::default(),
295        };
296        drop(guard);
297
298        // A replay that starts partway through the gap is a fragment of what the
299        // client missed, not the remainder of it — and replaying it would re-apply
300        // updates older than the task snapshot the service sends first, which for
301        // an appending artifact means duplicated content. Stream live instead and
302        // let the snapshot be the client's state.
303        let replay = if replay.complete {
304            replay.events
305        } else {
306            #[cfg(feature = "tracing")]
307            tracing::warn!(
308                task_id = %task_id,
309                from_event_id = ?from_event_id,
310                dropped = replay.events.len(),
311                "event log no longer covers the requested resume point; streaming live from the task snapshot"
312            );
313            Vec::new()
314        };
315
316        let live = futures::stream::unfold(receiver, |mut rx| async move {
317            match rx.recv().await {
318                Ok(event) => Some((Ok(event), rx)),
319                // Reader fell behind the broadcast channel: surface an error so a
320                // resilient client reconnects and resumes from its last id.
321                Err(broadcast::error::RecvError::Lagged(n)) => Some((
322                    Err(A2AError::Internal(format!(
323                        "streaming reader lagged, dropped {n} events"
324                    ))),
325                    rx,
326                )),
327                Err(broadcast::error::RecvError::Closed) => None,
328            }
329        });
330
331        let stream = futures::stream::iter(replay.into_iter().map(Ok)).chain(live);
332        Ok(Box::pin(stream))
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339    use crate::domain::{TaskState, TaskStatus, TaskStatusUpdateEvent};
340
341    fn status_event(task_id: &str, state: TaskState) -> TaskStatusUpdateEvent {
342        TaskStatusUpdateEvent {
343            task_id: task_id.to_string(),
344            context_id: "ctx".to_string(),
345            kind: "status-update".to_string(),
346            status: TaskStatus::new(state, None),
347            metadata: None,
348        }
349    }
350
351    fn seq_state(seq: &SeqEvent) -> ::buffa::EnumValue<TaskState> {
352        match &seq.event {
353            UpdateEvent::StatusUpdate(e) => e.status.state,
354            UpdateEvent::ArtifactUpdate(_) => panic!("expected status update"),
355        }
356    }
357
358    /// A live `combined_update_stream` reader receives broadcasts in order, each
359    /// tagged with a monotonic id starting at 1.
360    #[tokio::test]
361    async fn live_stream_delivers_in_order_with_ids() {
362        let handler = InMemoryStreamingHandler::new();
363        let mut stream = handler.combined_update_stream("t1", None).await.unwrap();
364
365        handler
366            .broadcast_status_update("t1", status_event("t1", TaskState::Working))
367            .await
368            .unwrap();
369        handler
370            .broadcast_status_update("t1", status_event("t1", TaskState::Completed))
371            .await
372            .unwrap();
373
374        let first = stream.next().await.unwrap().unwrap();
375        let second = stream.next().await.unwrap().unwrap();
376        assert_eq!(first.id, 1);
377        assert_eq!(
378            seq_state(&first),
379            ::buffa::EnumValue::from(TaskState::Working)
380        );
381        assert_eq!(second.id, 2);
382        assert_eq!(
383            seq_state(&second),
384            ::buffa::EnumValue::from(TaskState::Completed)
385        );
386    }
387
388    /// Subscribing with `from_event_id` replays the logged tail with a greater
389    /// id before any live updates.
390    #[tokio::test]
391    async fn resume_replays_buffered_tail() {
392        let handler = InMemoryStreamingHandler::new();
393        // Emit two events with no live reader; they are retained in the log.
394        handler
395            .broadcast_status_update("t1", status_event("t1", TaskState::Working))
396            .await
397            .unwrap();
398        handler
399            .broadcast_status_update("t1", status_event("t1", TaskState::Completed))
400            .await
401            .unwrap();
402
403        // Resume from id 1: only event 2 should replay.
404        let mut stream = handler.combined_update_stream("t1", Some(1)).await.unwrap();
405        let replayed = stream.next().await.unwrap().unwrap();
406        assert_eq!(replayed.id, 2);
407        assert_eq!(
408            seq_state(&replayed),
409            ::buffa::EnumValue::from(TaskState::Completed)
410        );
411    }
412
413    /// Past the log's capacity the tail is a fragment of the gap rather than its
414    /// remainder, so it is dropped and the client resumes from the snapshot the
415    /// service sends ahead of the stream.
416    #[tokio::test]
417    async fn an_uncoverable_resume_streams_live_instead_of_a_partial_tail() {
418        let handler = StreamingFanout::over(InMemoryEventLog::with_capacity(2));
419        for _ in 0..5 {
420            handler
421                .broadcast_status_update("t1", status_event("t1", TaskState::Working))
422                .await
423                .unwrap();
424        }
425
426        let mut stream = handler.combined_update_stream("t1", Some(1)).await.unwrap();
427        handler
428            .broadcast_status_update("t1", status_event("t1", TaskState::Completed))
429            .await
430            .unwrap();
431
432        let next = stream.next().await.unwrap().unwrap();
433        assert_eq!(
434            next.id, 6,
435            "events 4 and 5 are a fragment of the gap, so the stream starts live at 6"
436        );
437    }
438
439    /// Dropping a task's subscribers leaves its events replayable — a task with
440    /// nobody listening is what a resume attaches to.
441    #[tokio::test]
442    async fn removing_subscribers_keeps_the_log() {
443        let handler = InMemoryStreamingHandler::new();
444        handler
445            .broadcast_status_update("t1", status_event("t1", TaskState::Working))
446            .await
447            .unwrap();
448        handler.remove_task_subscribers("t1").await.unwrap();
449
450        let mut stream = handler.combined_update_stream("t1", Some(0)).await.unwrap();
451        let replayed = stream.next().await.unwrap().unwrap();
452        assert_eq!(replayed.id, 1);
453    }
454
455    /// A synchronous callback subscriber still receives broadcasts (the push API
456    /// rides alongside the broadcast channel).
457    #[tokio::test]
458    async fn callback_subscriber_still_notified() {
459        use std::sync::Mutex as StdMutex;
460
461        #[derive(Default, Clone)]
462        struct Recorder {
463            seen: Arc<StdMutex<Vec<::buffa::EnumValue<TaskState>>>>,
464        }
465        #[async_trait]
466        impl Subscriber<TaskStatusUpdateEvent> for Recorder {
467            async fn on_update(&self, update: TaskStatusUpdateEvent) -> Result<(), A2AError> {
468                self.seen.lock().unwrap().push(update.status.state);
469                Ok(())
470            }
471        }
472
473        let handler = InMemoryStreamingHandler::new();
474        let recorder = Recorder::default();
475        handler
476            .add_status_subscriber("t1", Box::new(recorder.clone()))
477            .await
478            .unwrap();
479        handler
480            .broadcast_status_update("t1", status_event("t1", TaskState::Working))
481            .await
482            .unwrap();
483
484        assert_eq!(
485            *recorder.seen.lock().unwrap(),
486            vec![::buffa::EnumValue::from(TaskState::Working)]
487        );
488    }
489}
a2a_rs/adapter/streaming/fanout.rs

a2a_rs/adapter/streaming/
fanout.rs
