1use std::collections::HashMap;
9use std::sync::Arc;
10
11use a2a_protocol_types::task::TaskId;
12use tokio::sync::RwLock;
13
14use a2a_protocol_types::error::A2aResult;
15use a2a_protocol_types::events::StreamResponse;
16
17use super::{
18 DEFAULT_MAX_EVENT_SIZE, DEFAULT_QUEUE_CAPACITY, DEFAULT_WRITE_TIMEOUT, InMemoryQueueReader,
19 InMemoryQueueWriter, StreamEvent, new_in_memory_queue_with_options,
20 new_in_memory_queue_with_persistence,
21};
22use crate::metrics::Metrics;
23
24#[allow(clippy::large_enum_variant)]
39pub enum QueueLease {
40 Created {
43 writer: Arc<InMemoryQueueWriter>,
44 reader: InMemoryQueueReader,
45 persistence_rx: Option<tokio::sync::mpsc::Receiver<A2aResult<StreamEvent>>>,
46 },
47 Existing,
52 CapacityExhausted,
55}
56
57#[derive(Clone)]
65pub struct EventQueueManager {
66 writers: Arc<RwLock<HashMap<TaskId, Arc<InMemoryQueueWriter>>>>,
67 capacity: usize,
69 max_event_size: usize,
71 write_timeout: std::time::Duration,
73 max_concurrent_queues: Option<usize>,
75 metrics: Option<Arc<dyn Metrics>>,
77}
78
79impl std::fmt::Debug for EventQueueManager {
80 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81 f.debug_struct("EventQueueManager")
82 .field("writers", &"<RwLock<HashMap<...>>>")
83 .field("capacity", &self.capacity)
84 .field("max_event_size", &self.max_event_size)
85 .field("write_timeout", &self.write_timeout)
86 .field("max_concurrent_queues", &self.max_concurrent_queues)
87 .field("metrics", &self.metrics.is_some())
88 .finish()
89 }
90}
91
92impl Default for EventQueueManager {
93 fn default() -> Self {
94 Self {
95 writers: Arc::default(),
96 capacity: DEFAULT_QUEUE_CAPACITY,
97 max_event_size: DEFAULT_MAX_EVENT_SIZE,
98 write_timeout: DEFAULT_WRITE_TIMEOUT,
99 max_concurrent_queues: None,
100 metrics: None,
101 }
102 }
103}
104
105impl EventQueueManager {
106 #[must_use]
116 pub fn new() -> Self {
117 Self::default()
118 }
119
120 #[must_use]
122 pub fn with_capacity(capacity: usize) -> Self {
123 Self {
124 writers: Arc::default(),
125 capacity,
126 max_event_size: DEFAULT_MAX_EVENT_SIZE,
127 write_timeout: DEFAULT_WRITE_TIMEOUT,
128 max_concurrent_queues: None,
129 metrics: None,
130 }
131 }
132
133 #[must_use]
138 pub const fn with_max_event_size(mut self, max_event_size: usize) -> Self {
139 self.max_event_size = max_event_size;
140 self
141 }
142
143 #[must_use]
145 pub fn with_metrics(mut self, metrics: Arc<dyn Metrics>) -> Self {
146 self.metrics = Some(metrics);
147 self
148 }
149
150 fn observed(&self, writer: InMemoryQueueWriter) -> InMemoryQueueWriter {
153 match &self.metrics {
154 Some(metrics) => writer.with_metrics(Arc::clone(metrics)),
155 None => writer,
156 }
157 }
158
159 #[must_use]
164 pub const fn with_max_concurrent_queues(mut self, max: usize) -> Self {
165 self.max_concurrent_queues = Some(max);
166 self
167 }
168
169 pub async fn get_or_create(
180 &self,
181 task_id: &TaskId,
182 ) -> (Arc<InMemoryQueueWriter>, Option<InMemoryQueueReader>) {
183 let mut map = self.writers.write().await;
184 #[allow(clippy::option_if_let_else)]
185 let result = if let Some(existing) = map.get(task_id) {
186 (Arc::clone(existing), None)
187 } else if self
188 .max_concurrent_queues
189 .is_some_and(|max| map.len() >= max)
190 {
191 let (writer, _reader) = new_in_memory_queue_with_options(
194 self.capacity,
195 self.max_event_size,
196 self.write_timeout,
197 );
198 (Arc::new(writer), None)
199 } else {
200 let (writer, reader) = new_in_memory_queue_with_options(
201 self.capacity,
202 self.max_event_size,
203 self.write_timeout,
204 );
205 let writer = Arc::new(writer);
206 map.insert(task_id.clone(), Arc::clone(&writer));
207 (writer, Some(reader))
208 };
209 let queue_count = map.len();
210 drop(map);
211 if let Some(ref metrics) = self.metrics {
212 metrics.on_queue_depth_change(queue_count);
213 }
214 result
215 }
216
217 pub async fn get_or_create_with_persistence(
225 &self,
226 task_id: &TaskId,
227 ) -> (
228 Arc<InMemoryQueueWriter>,
229 Option<InMemoryQueueReader>,
230 Option<tokio::sync::mpsc::Receiver<A2aResult<StreamEvent>>>,
231 ) {
232 let mut map = self.writers.write().await;
233 #[allow(clippy::option_if_let_else)]
234 let result = if let Some(existing) = map.get(task_id) {
235 (Arc::clone(existing), None, None)
236 } else if self
237 .max_concurrent_queues
238 .is_some_and(|max| map.len() >= max)
239 {
240 let (writer, _reader) = new_in_memory_queue_with_options(
241 self.capacity,
242 self.max_event_size,
243 self.write_timeout,
244 );
245 (Arc::new(writer), None, None)
246 } else {
247 let (writer, reader, persistence_rx) = new_in_memory_queue_with_persistence(
248 self.capacity,
249 self.max_event_size,
250 self.write_timeout,
251 );
252 let writer = Arc::new(self.observed(writer));
253 map.insert(task_id.clone(), Arc::clone(&writer));
254 (writer, Some(reader), Some(persistence_rx))
255 };
256 let queue_count = map.len();
257 drop(map);
258 if let Some(ref metrics) = self.metrics {
259 metrics.on_queue_depth_change(queue_count);
260 }
261 result
262 }
263
264 #[allow(clippy::option_if_let_else)]
282 pub(crate) async fn lease(
283 &self,
284 task_id: &TaskId,
285 with_persistence: bool,
286 capacity: Option<usize>,
287 ) -> QueueLease {
288 let capacity = capacity.unwrap_or(self.capacity);
289 let mut map = self.writers.write().await;
290 let lease = if map.contains_key(task_id) {
291 QueueLease::Existing
292 } else if self
293 .max_concurrent_queues
294 .is_some_and(|max| map.len() >= max)
295 {
296 QueueLease::CapacityExhausted
297 } else if with_persistence {
298 let (writer, reader, persistence_rx) = new_in_memory_queue_with_persistence(
299 capacity,
300 self.max_event_size,
301 self.write_timeout,
302 );
303 let writer = Arc::new(self.observed(writer).with_terminal_gate());
306 map.insert(task_id.clone(), Arc::clone(&writer));
307 QueueLease::Created {
308 writer,
309 reader,
310 persistence_rx: Some(persistence_rx),
311 }
312 } else {
313 let (writer, reader) =
314 new_in_memory_queue_with_options(capacity, self.max_event_size, self.write_timeout);
315 let writer = Arc::new(writer);
316 map.insert(task_id.clone(), Arc::clone(&writer));
317 QueueLease::Created {
318 writer,
319 reader,
320 persistence_rx: None,
321 }
322 };
323 let queue_count = map.len();
324 drop(map);
325 if let Some(ref metrics) = self.metrics {
326 metrics.on_queue_depth_change(queue_count);
327 }
328 lease
329 }
330
331 pub(crate) async fn current_seq(&self, task_id: &TaskId) -> Option<u64> {
350 let map = self.writers.read().await;
351 map.get(task_id).map(|w| w.current_seq())
352 }
353
354 pub(crate) async fn writer_for_cancel(&self, task_id: &TaskId) -> Arc<InMemoryQueueWriter> {
355 {
356 let map = self.writers.read().await;
357 if let Some(writer) = map.get(task_id) {
358 return Arc::clone(writer);
359 }
360 }
361 let (writer, _reader) = new_in_memory_queue_with_options(
362 self.capacity,
363 self.max_event_size,
364 self.write_timeout,
365 );
366 Arc::new(writer)
367 }
368
369 pub async fn subscribe(&self, task_id: &TaskId) -> Option<InMemoryQueueReader> {
377 let map = self.writers.read().await;
378 map.get(task_id).map(|writer| writer.subscribe())
379 }
380
381 pub(crate) async fn raw_subscribe(
388 &self,
389 task_id: &TaskId,
390 ) -> Option<tokio::sync::broadcast::Receiver<A2aResult<StreamEvent>>> {
391 let map = self.writers.read().await;
392 map.get(task_id).map(|writer| writer.raw_subscribe())
393 }
394
395 pub async fn subscribe_with_snapshot(
404 &self,
405 task_id: &TaskId,
406 snapshot: StreamResponse,
407 ) -> Option<InMemoryQueueReader> {
408 let map = self.writers.read().await;
409 let writer = map.get(task_id)?;
410 let rx = writer.raw_subscribe();
414 drop(map);
415 Some(InMemoryQueueReader::with_first_event(rx, snapshot))
416 }
417
418 pub async fn destroy(&self, task_id: &TaskId) {
420 let mut map = self.writers.write().await;
421 map.remove(task_id);
422 let queue_count = map.len();
423 drop(map);
424 if let Some(ref metrics) = self.metrics {
425 metrics.on_queue_depth_change(queue_count);
426 }
427 }
428
429 pub async fn active_count(&self) -> usize {
431 let map = self.writers.read().await;
432 map.len()
433 }
434
435 pub(crate) async fn has_queue(&self, task_id: &TaskId) -> bool {
441 self.writers.read().await.contains_key(task_id)
442 }
443
444 #[must_use]
447 pub(crate) const fn max_concurrent_queues(&self) -> Option<usize> {
448 self.max_concurrent_queues
449 }
450
451 pub async fn destroy_all(&self) {
453 let mut map = self.writers.write().await;
454 map.clear();
455 }
456}
457
458#[cfg(test)]
459mod tests {
460 use super::*;
461 use crate::streaming::event_queue::{EventQueueReader, EventQueueWriter};
462 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
463 use a2a_protocol_types::task::{ContextId, TaskState, TaskStatus};
464
465 fn make_status_event(task_id: &str, state: TaskState) -> StreamResponse {
467 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
468 task_id: TaskId::new(task_id),
469 context_id: ContextId::new("ctx-test"),
470 status: TaskStatus {
471 state,
472 message: None,
473 timestamp: None,
474 },
475 metadata: None,
476 })
477 }
478
479 #[test]
482 fn max_concurrent_queues_reports_configured_limit() {
483 assert_eq!(EventQueueManager::new().max_concurrent_queues(), None);
485 assert_eq!(
487 EventQueueManager::new()
488 .with_max_concurrent_queues(42)
489 .max_concurrent_queues(),
490 Some(42)
491 );
492 }
493
494 #[tokio::test]
495 async fn manager_get_or_create_new_task() {
496 let manager = EventQueueManager::new();
497 let task_id = TaskId::new("task-1");
498
499 let (writer, reader) = manager.get_or_create(&task_id).await;
500 assert!(
501 reader.is_some(),
502 "first get_or_create should return a reader"
503 );
504
505 writer
507 .write(make_status_event("task-1", TaskState::Working))
508 .await
509 .expect("write through manager writer should succeed");
510
511 assert_eq!(
512 manager.active_count().await,
513 1,
514 "should have 1 active queue"
515 );
516 }
517
518 #[tokio::test]
519 async fn manager_get_or_create_existing_task_returns_no_reader() {
520 let manager = EventQueueManager::new();
521 let task_id = TaskId::new("task-1");
522
523 let (_w1, r1) = manager.get_or_create(&task_id).await;
524 assert!(r1.is_some(), "first call should return a reader");
525
526 let (_w2, r2) = manager.get_or_create(&task_id).await;
527 assert!(
528 r2.is_none(),
529 "second call for same task should return None reader"
530 );
531
532 assert_eq!(
533 manager.active_count().await,
534 1,
535 "should still have only 1 active queue"
536 );
537 }
538
539 #[tokio::test]
540 async fn manager_subscribe_existing_task() {
541 use crate::streaming::event_queue::EventQueueReader;
542
543 let manager = EventQueueManager::new();
544 let task_id = TaskId::new("task-1");
545
546 let (writer, _reader) = manager.get_or_create(&task_id).await;
547
548 let sub = manager.subscribe(&task_id).await;
549 assert!(
550 sub.is_some(),
551 "subscribe should return a reader for existing task"
552 );
553
554 let mut sub_reader = sub.unwrap();
555 writer
556 .write(make_status_event("task-1", TaskState::Working))
557 .await
558 .expect("write should succeed");
559 drop(writer);
560
561 let r = tokio::time::timeout(std::time::Duration::from_secs(5), sub_reader.read())
565 .await
566 .expect("subscriber read must not hang");
567 assert!(r.is_some(), "subscriber should receive the event");
568 }
569
570 #[tokio::test]
571 async fn manager_subscribe_nonexistent_task_returns_none() {
572 let manager = EventQueueManager::new();
573 let task_id = TaskId::new("no-such-task");
574
575 let sub = manager.subscribe(&task_id).await;
576 assert!(
577 sub.is_none(),
578 "subscribe should return None for nonexistent task"
579 );
580 }
581
582 #[tokio::test]
583 async fn manager_destroy_removes_queue() {
584 let manager = EventQueueManager::new();
585 let task_id = TaskId::new("task-1");
586
587 let (_writer, _reader) = manager.get_or_create(&task_id).await;
588 assert_eq!(manager.active_count().await, 1);
589
590 manager.destroy(&task_id).await;
591 assert_eq!(
592 manager.active_count().await,
593 0,
594 "destroy should remove the queue"
595 );
596 }
597
598 #[tokio::test]
599 async fn manager_destroy_all_clears_queues() {
600 let manager = EventQueueManager::new();
601
602 let _q1 = manager.get_or_create(&TaskId::new("t1")).await;
603 let _q2 = manager.get_or_create(&TaskId::new("t2")).await;
604 assert_eq!(manager.active_count().await, 2);
605
606 manager.destroy_all().await;
607 assert_eq!(
608 manager.active_count().await,
609 0,
610 "destroy_all should clear all queues"
611 );
612 }
613
614 #[tokio::test]
615 async fn lease_reports_existing_and_has_queue() {
616 let manager = EventQueueManager::new();
617 let task = TaskId::new("t-lease");
618
619 assert!(matches!(
621 manager.lease(&task, true, None).await,
622 QueueLease::Created { .. }
623 ));
624 assert!(manager.has_queue(&task).await, "queue should now be live");
625
626 assert!(matches!(
629 manager.lease(&task, true, None).await,
630 QueueLease::Existing
631 ));
632
633 assert!(!manager.has_queue(&TaskId::new("other")).await);
635 }
636
637 #[tokio::test]
638 async fn manager_max_concurrent_queues_enforced() {
639 let manager = EventQueueManager::new().with_max_concurrent_queues(1);
640
641 let (_w1, r1) = manager.get_or_create(&TaskId::new("t1")).await;
642 assert!(r1.is_some(), "first queue should be created successfully");
643
644 let (_w2, r2) = manager.get_or_create(&TaskId::new("t2")).await;
646 assert!(
647 r2.is_none(),
648 "second queue should return None reader when limit is reached"
649 );
650 assert_eq!(
651 manager.active_count().await,
652 1,
653 "should still have only 1 queue (second was not stored)"
654 );
655 }
656
657 #[tokio::test]
658 async fn manager_with_capacity_and_max_event_size() {
659 let manager = EventQueueManager::with_capacity(4).with_max_event_size(10); let task_id = TaskId::new("t1");
662 let (writer, _reader) = manager.get_or_create(&task_id).await;
663
664 let event = make_status_event("t1", TaskState::Working);
665 let result = writer.write(event).await;
666 assert!(
667 result.is_err(),
668 "event should be rejected by the size limit configured on the manager"
669 );
670 }
671
672 #[tokio::test]
688 async fn with_capacity_uses_the_given_capacity_not_the_default() {
689 let manager = EventQueueManager::with_capacity(1);
690 let task_id = TaskId::new("cap");
691 let (writer, reader) = manager.get_or_create(&task_id).await;
692 let mut reader = reader.expect("first get_or_create yields a reader");
693
694 writer
696 .write(make_status_event("cap", TaskState::Working))
697 .await
698 .expect("first write");
699 writer
700 .write(make_status_event("cap", TaskState::Completed))
701 .await
702 .expect("second write");
703
704 let first = reader.read().await.expect("reader is still open");
705 assert!(
706 first.is_err(),
707 "a capacity-1 queue must surface an overrun to the reader; \
708 got Ok, which is what DEFAULT_QUEUE_CAPACITY (256) would give"
709 );
710 }
711
712 #[tokio::test]
721 async fn lease_capacity_override_beats_the_managers_own() {
722 let manager = EventQueueManager::new();
723 let task_id = TaskId::new("override");
724 let crate::streaming::QueueLease::Created { writer, reader, .. } =
725 manager.lease(&task_id, false, Some(1)).await
726 else {
727 panic!("first lease must create a queue");
728 };
729 let mut reader = reader;
730
731 writer
732 .write(make_status_event("override", TaskState::Working))
733 .await
734 .expect("first write");
735 writer
736 .write(make_status_event("override", TaskState::Completed))
737 .await
738 .expect("second write");
739
740 assert!(
741 reader.read().await.expect("reader is still open").is_err(),
742 "the override asked for capacity 1; buffering both events is what \
743 the manager's own default (256) would do"
744 );
745 }
746
747 #[tokio::test]
750 async fn lease_without_an_override_uses_the_managers_capacity() {
751 let manager = EventQueueManager::with_capacity(1);
752 let task_id = TaskId::new("no-override");
753 let crate::streaming::QueueLease::Created { writer, reader, .. } =
754 manager.lease(&task_id, false, None).await
755 else {
756 panic!("first lease must create a queue");
757 };
758 let mut reader = reader;
759
760 writer
761 .write(make_status_event("no-override", TaskState::Working))
762 .await
763 .expect("first write");
764 writer
765 .write(make_status_event("no-override", TaskState::Completed))
766 .await
767 .expect("second write");
768
769 assert!(
770 reader.read().await.expect("reader is still open").is_err(),
771 "None must fall back to the manager's capacity of 1, not to the default"
772 );
773 }
774
775 #[tokio::test]
785 async fn first_queue_is_tracked_when_a_concurrency_limit_is_set() {
786 let manager = EventQueueManager::new().with_max_concurrent_queues(1);
787
788 let first = TaskId::new("q1");
789 let (_w, reader, persistence) = manager.get_or_create_with_persistence(&first).await;
790 assert!(
791 reader.is_some(),
792 "the first queue is below the limit and must be tracked, \
793 with a reader; None means the at-capacity branch was taken"
794 );
795 assert!(
796 persistence.is_some(),
797 "a tracked queue gets a persistence rx"
798 );
799 assert_eq!(
800 manager.active_count().await,
801 1,
802 "first queue must be stored"
803 );
804
805 let second = TaskId::new("q2");
807 let (_w2, reader2, _p2) = manager.get_or_create_with_persistence(&second).await;
808 assert!(
809 reader2.is_none(),
810 "the second queue exceeds the limit and must not be tracked"
811 );
812 assert_eq!(
813 manager.active_count().await,
814 1,
815 "an over-limit queue must not be stored"
816 );
817 }
818
819 #[tokio::test]
822 async fn raw_subscribe_returns_a_receiver_for_a_live_queue() {
823 let manager = EventQueueManager::new();
824 let task_id = TaskId::new("raw");
825 let (_writer, _reader) = manager.get_or_create(&task_id).await;
826
827 assert!(
828 manager.raw_subscribe(&task_id).await.is_some(),
829 "a live queue must yield a receiver"
830 );
831 assert!(
832 manager
833 .raw_subscribe(&TaskId::new("absent"))
834 .await
835 .is_none(),
836 "an unknown task must yield None — pins that Some is not blanket"
837 );
838 }
839
840 #[tokio::test]
847 async fn subscribe_with_snapshot_returns_a_reader_that_yields_the_snapshot() {
848 let manager = EventQueueManager::new();
849 let task_id = TaskId::new("snap");
850 let (_writer, _reader) = manager.get_or_create(&task_id).await;
851
852 let snapshot = make_status_event("snap", TaskState::Working);
853 let reader = manager.subscribe_with_snapshot(&task_id, snapshot).await;
854 let mut reader = reader.expect("a live queue must yield a reader");
855
856 let first = reader
857 .read()
858 .await
859 .expect("reader is open")
860 .expect("snapshot is delivered as Ok");
861 match first.event {
862 StreamResponse::StatusUpdate(ev) => {
863 assert_eq!(
864 ev.status.state,
865 TaskState::Working,
866 "snapshot arrives first"
867 );
868 }
869 other => panic!("expected the snapshot StatusUpdate first, got {other:?}"),
870 }
871
872 assert!(
873 manager
874 .subscribe_with_snapshot(
875 &TaskId::new("absent"),
876 make_status_event("absent", TaskState::Working)
877 )
878 .await
879 .is_none(),
880 "an unknown task must yield None — pins that Some is not blanket"
881 );
882 }
883}