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, new_in_memory_queue_with_options, new_in_memory_queue_with_persistence,
20};
21use crate::metrics::Metrics;
22
23#[allow(clippy::large_enum_variant)]
38pub enum QueueLease {
39 Created {
42 writer: Arc<InMemoryQueueWriter>,
43 reader: InMemoryQueueReader,
44 persistence_rx: Option<tokio::sync::mpsc::Receiver<A2aResult<StreamResponse>>>,
45 },
46 Existing,
51 CapacityExhausted,
54}
55
56#[derive(Clone)]
64pub struct EventQueueManager {
65 writers: Arc<RwLock<HashMap<TaskId, Arc<InMemoryQueueWriter>>>>,
66 capacity: usize,
68 max_event_size: usize,
70 write_timeout: std::time::Duration,
72 max_concurrent_queues: Option<usize>,
74 metrics: Option<Arc<dyn Metrics>>,
76}
77
78impl std::fmt::Debug for EventQueueManager {
79 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80 f.debug_struct("EventQueueManager")
81 .field("writers", &"<RwLock<HashMap<...>>>")
82 .field("capacity", &self.capacity)
83 .field("max_event_size", &self.max_event_size)
84 .field("write_timeout", &self.write_timeout)
85 .field("max_concurrent_queues", &self.max_concurrent_queues)
86 .field("metrics", &self.metrics.is_some())
87 .finish()
88 }
89}
90
91impl Default for EventQueueManager {
92 fn default() -> Self {
93 Self {
94 writers: Arc::default(),
95 capacity: DEFAULT_QUEUE_CAPACITY,
96 max_event_size: DEFAULT_MAX_EVENT_SIZE,
97 write_timeout: DEFAULT_WRITE_TIMEOUT,
98 max_concurrent_queues: None,
99 metrics: None,
100 }
101 }
102}
103
104impl EventQueueManager {
105 #[must_use]
115 pub fn new() -> Self {
116 Self::default()
117 }
118
119 #[must_use]
121 pub fn with_capacity(capacity: usize) -> Self {
122 Self {
123 writers: Arc::default(),
124 capacity,
125 max_event_size: DEFAULT_MAX_EVENT_SIZE,
126 write_timeout: DEFAULT_WRITE_TIMEOUT,
127 max_concurrent_queues: None,
128 metrics: None,
129 }
130 }
131
132 #[must_use]
137 pub const fn with_max_event_size(mut self, max_event_size: usize) -> Self {
138 self.max_event_size = max_event_size;
139 self
140 }
141
142 #[must_use]
144 pub fn with_metrics(mut self, metrics: Arc<dyn Metrics>) -> Self {
145 self.metrics = Some(metrics);
146 self
147 }
148
149 fn observed(&self, writer: InMemoryQueueWriter) -> InMemoryQueueWriter {
152 match &self.metrics {
153 Some(metrics) => writer.with_metrics(Arc::clone(metrics)),
154 None => writer,
155 }
156 }
157
158 #[must_use]
163 pub const fn with_max_concurrent_queues(mut self, max: usize) -> Self {
164 self.max_concurrent_queues = Some(max);
165 self
166 }
167
168 pub async fn get_or_create(
179 &self,
180 task_id: &TaskId,
181 ) -> (Arc<InMemoryQueueWriter>, Option<InMemoryQueueReader>) {
182 let mut map = self.writers.write().await;
183 #[allow(clippy::option_if_let_else)]
184 let result = if let Some(existing) = map.get(task_id) {
185 (Arc::clone(existing), None)
186 } else if self
187 .max_concurrent_queues
188 .is_some_and(|max| map.len() >= max)
189 {
190 let (writer, _reader) = new_in_memory_queue_with_options(
193 self.capacity,
194 self.max_event_size,
195 self.write_timeout,
196 );
197 (Arc::new(writer), None)
198 } else {
199 let (writer, reader) = new_in_memory_queue_with_options(
200 self.capacity,
201 self.max_event_size,
202 self.write_timeout,
203 );
204 let writer = Arc::new(writer);
205 map.insert(task_id.clone(), Arc::clone(&writer));
206 (writer, Some(reader))
207 };
208 let queue_count = map.len();
209 drop(map);
210 if let Some(ref metrics) = self.metrics {
211 metrics.on_queue_depth_change(queue_count);
212 }
213 result
214 }
215
216 pub async fn get_or_create_with_persistence(
224 &self,
225 task_id: &TaskId,
226 ) -> (
227 Arc<InMemoryQueueWriter>,
228 Option<InMemoryQueueReader>,
229 Option<tokio::sync::mpsc::Receiver<A2aResult<StreamResponse>>>,
230 ) {
231 let mut map = self.writers.write().await;
232 #[allow(clippy::option_if_let_else)]
233 let result = if let Some(existing) = map.get(task_id) {
234 (Arc::clone(existing), None, None)
235 } else if self
236 .max_concurrent_queues
237 .is_some_and(|max| map.len() >= max)
238 {
239 let (writer, _reader) = new_in_memory_queue_with_options(
240 self.capacity,
241 self.max_event_size,
242 self.write_timeout,
243 );
244 (Arc::new(writer), None, None)
245 } else {
246 let (writer, reader, persistence_rx) = new_in_memory_queue_with_persistence(
247 self.capacity,
248 self.max_event_size,
249 self.write_timeout,
250 );
251 let writer = Arc::new(self.observed(writer));
252 map.insert(task_id.clone(), Arc::clone(&writer));
253 (writer, Some(reader), Some(persistence_rx))
254 };
255 let queue_count = map.len();
256 drop(map);
257 if let Some(ref metrics) = self.metrics {
258 metrics.on_queue_depth_change(queue_count);
259 }
260 result
261 }
262
263 #[allow(clippy::option_if_let_else)]
281 pub(crate) async fn lease(
282 &self,
283 task_id: &TaskId,
284 with_persistence: bool,
285 capacity: Option<usize>,
286 ) -> QueueLease {
287 let capacity = capacity.unwrap_or(self.capacity);
288 let mut map = self.writers.write().await;
289 let lease = if map.contains_key(task_id) {
290 QueueLease::Existing
291 } else if self
292 .max_concurrent_queues
293 .is_some_and(|max| map.len() >= max)
294 {
295 QueueLease::CapacityExhausted
296 } else if with_persistence {
297 let (writer, reader, persistence_rx) = new_in_memory_queue_with_persistence(
298 capacity,
299 self.max_event_size,
300 self.write_timeout,
301 );
302 let writer = Arc::new(self.observed(writer));
303 map.insert(task_id.clone(), Arc::clone(&writer));
304 QueueLease::Created {
305 writer,
306 reader,
307 persistence_rx: Some(persistence_rx),
308 }
309 } else {
310 let (writer, reader) =
311 new_in_memory_queue_with_options(capacity, self.max_event_size, self.write_timeout);
312 let writer = Arc::new(writer);
313 map.insert(task_id.clone(), Arc::clone(&writer));
314 QueueLease::Created {
315 writer,
316 reader,
317 persistence_rx: None,
318 }
319 };
320 let queue_count = map.len();
321 drop(map);
322 if let Some(ref metrics) = self.metrics {
323 metrics.on_queue_depth_change(queue_count);
324 }
325 lease
326 }
327
328 pub(crate) async fn writer_for_cancel(&self, task_id: &TaskId) -> Arc<InMemoryQueueWriter> {
338 {
339 let map = self.writers.read().await;
340 if let Some(writer) = map.get(task_id) {
341 return Arc::clone(writer);
342 }
343 }
344 let (writer, _reader) = new_in_memory_queue_with_options(
345 self.capacity,
346 self.max_event_size,
347 self.write_timeout,
348 );
349 Arc::new(writer)
350 }
351
352 pub async fn subscribe(&self, task_id: &TaskId) -> Option<InMemoryQueueReader> {
360 let map = self.writers.read().await;
361 map.get(task_id).map(|writer| writer.subscribe())
362 }
363
364 pub(crate) async fn raw_subscribe(
371 &self,
372 task_id: &TaskId,
373 ) -> Option<tokio::sync::broadcast::Receiver<A2aResult<StreamResponse>>> {
374 let map = self.writers.read().await;
375 map.get(task_id).map(|writer| writer.raw_subscribe())
376 }
377
378 pub async fn subscribe_with_snapshot(
387 &self,
388 task_id: &TaskId,
389 snapshot: StreamResponse,
390 ) -> Option<InMemoryQueueReader> {
391 let map = self.writers.read().await;
392 let writer = map.get(task_id)?;
393 let rx = writer.raw_subscribe();
397 drop(map);
398 Some(InMemoryQueueReader::with_first_event(rx, snapshot))
399 }
400
401 pub async fn destroy(&self, task_id: &TaskId) {
403 let mut map = self.writers.write().await;
404 map.remove(task_id);
405 let queue_count = map.len();
406 drop(map);
407 if let Some(ref metrics) = self.metrics {
408 metrics.on_queue_depth_change(queue_count);
409 }
410 }
411
412 pub async fn active_count(&self) -> usize {
414 let map = self.writers.read().await;
415 map.len()
416 }
417
418 pub(crate) async fn has_queue(&self, task_id: &TaskId) -> bool {
424 self.writers.read().await.contains_key(task_id)
425 }
426
427 #[must_use]
430 pub(crate) const fn max_concurrent_queues(&self) -> Option<usize> {
431 self.max_concurrent_queues
432 }
433
434 pub async fn destroy_all(&self) {
436 let mut map = self.writers.write().await;
437 map.clear();
438 }
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444 use crate::streaming::event_queue::{EventQueueReader, EventQueueWriter};
445 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
446 use a2a_protocol_types::task::{ContextId, TaskState, TaskStatus};
447
448 fn make_status_event(task_id: &str, state: TaskState) -> StreamResponse {
450 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
451 task_id: TaskId::new(task_id),
452 context_id: ContextId::new("ctx-test"),
453 status: TaskStatus {
454 state,
455 message: None,
456 timestamp: None,
457 },
458 metadata: None,
459 })
460 }
461
462 #[test]
465 fn max_concurrent_queues_reports_configured_limit() {
466 assert_eq!(EventQueueManager::new().max_concurrent_queues(), None);
468 assert_eq!(
470 EventQueueManager::new()
471 .with_max_concurrent_queues(42)
472 .max_concurrent_queues(),
473 Some(42)
474 );
475 }
476
477 #[tokio::test]
478 async fn manager_get_or_create_new_task() {
479 let manager = EventQueueManager::new();
480 let task_id = TaskId::new("task-1");
481
482 let (writer, reader) = manager.get_or_create(&task_id).await;
483 assert!(
484 reader.is_some(),
485 "first get_or_create should return a reader"
486 );
487
488 writer
490 .write(make_status_event("task-1", TaskState::Working))
491 .await
492 .expect("write through manager writer should succeed");
493
494 assert_eq!(
495 manager.active_count().await,
496 1,
497 "should have 1 active queue"
498 );
499 }
500
501 #[tokio::test]
502 async fn manager_get_or_create_existing_task_returns_no_reader() {
503 let manager = EventQueueManager::new();
504 let task_id = TaskId::new("task-1");
505
506 let (_w1, r1) = manager.get_or_create(&task_id).await;
507 assert!(r1.is_some(), "first call should return a reader");
508
509 let (_w2, r2) = manager.get_or_create(&task_id).await;
510 assert!(
511 r2.is_none(),
512 "second call for same task should return None reader"
513 );
514
515 assert_eq!(
516 manager.active_count().await,
517 1,
518 "should still have only 1 active queue"
519 );
520 }
521
522 #[tokio::test]
523 async fn manager_subscribe_existing_task() {
524 use crate::streaming::event_queue::EventQueueReader;
525
526 let manager = EventQueueManager::new();
527 let task_id = TaskId::new("task-1");
528
529 let (writer, _reader) = manager.get_or_create(&task_id).await;
530
531 let sub = manager.subscribe(&task_id).await;
532 assert!(
533 sub.is_some(),
534 "subscribe should return a reader for existing task"
535 );
536
537 let mut sub_reader = sub.unwrap();
538 writer
539 .write(make_status_event("task-1", TaskState::Working))
540 .await
541 .expect("write should succeed");
542 drop(writer);
543
544 let r = sub_reader.read().await;
545 assert!(r.is_some(), "subscriber should receive the event");
546 }
547
548 #[tokio::test]
549 async fn manager_subscribe_nonexistent_task_returns_none() {
550 let manager = EventQueueManager::new();
551 let task_id = TaskId::new("no-such-task");
552
553 let sub = manager.subscribe(&task_id).await;
554 assert!(
555 sub.is_none(),
556 "subscribe should return None for nonexistent task"
557 );
558 }
559
560 #[tokio::test]
561 async fn manager_destroy_removes_queue() {
562 let manager = EventQueueManager::new();
563 let task_id = TaskId::new("task-1");
564
565 let (_writer, _reader) = manager.get_or_create(&task_id).await;
566 assert_eq!(manager.active_count().await, 1);
567
568 manager.destroy(&task_id).await;
569 assert_eq!(
570 manager.active_count().await,
571 0,
572 "destroy should remove the queue"
573 );
574 }
575
576 #[tokio::test]
577 async fn manager_destroy_all_clears_queues() {
578 let manager = EventQueueManager::new();
579
580 let _q1 = manager.get_or_create(&TaskId::new("t1")).await;
581 let _q2 = manager.get_or_create(&TaskId::new("t2")).await;
582 assert_eq!(manager.active_count().await, 2);
583
584 manager.destroy_all().await;
585 assert_eq!(
586 manager.active_count().await,
587 0,
588 "destroy_all should clear all queues"
589 );
590 }
591
592 #[tokio::test]
593 async fn lease_reports_existing_and_has_queue() {
594 let manager = EventQueueManager::new();
595 let task = TaskId::new("t-lease");
596
597 assert!(matches!(
599 manager.lease(&task, true, None).await,
600 QueueLease::Created { .. }
601 ));
602 assert!(manager.has_queue(&task).await, "queue should now be live");
603
604 assert!(matches!(
607 manager.lease(&task, true, None).await,
608 QueueLease::Existing
609 ));
610
611 assert!(!manager.has_queue(&TaskId::new("other")).await);
613 }
614
615 #[tokio::test]
616 async fn manager_max_concurrent_queues_enforced() {
617 let manager = EventQueueManager::new().with_max_concurrent_queues(1);
618
619 let (_w1, r1) = manager.get_or_create(&TaskId::new("t1")).await;
620 assert!(r1.is_some(), "first queue should be created successfully");
621
622 let (_w2, r2) = manager.get_or_create(&TaskId::new("t2")).await;
624 assert!(
625 r2.is_none(),
626 "second queue should return None reader when limit is reached"
627 );
628 assert_eq!(
629 manager.active_count().await,
630 1,
631 "should still have only 1 queue (second was not stored)"
632 );
633 }
634
635 #[tokio::test]
636 async fn manager_with_capacity_and_max_event_size() {
637 let manager = EventQueueManager::with_capacity(4).with_max_event_size(10); let task_id = TaskId::new("t1");
640 let (writer, _reader) = manager.get_or_create(&task_id).await;
641
642 let event = make_status_event("t1", TaskState::Working);
643 let result = writer.write(event).await;
644 assert!(
645 result.is_err(),
646 "event should be rejected by the size limit configured on the manager"
647 );
648 }
649
650 #[tokio::test]
666 async fn with_capacity_uses_the_given_capacity_not_the_default() {
667 let manager = EventQueueManager::with_capacity(1);
668 let task_id = TaskId::new("cap");
669 let (writer, reader) = manager.get_or_create(&task_id).await;
670 let mut reader = reader.expect("first get_or_create yields a reader");
671
672 writer
674 .write(make_status_event("cap", TaskState::Working))
675 .await
676 .expect("first write");
677 writer
678 .write(make_status_event("cap", TaskState::Completed))
679 .await
680 .expect("second write");
681
682 let first = reader.read().await.expect("reader is still open");
683 assert!(
684 first.is_err(),
685 "a capacity-1 queue must surface an overrun to the reader; \
686 got Ok, which is what DEFAULT_QUEUE_CAPACITY (256) would give"
687 );
688 }
689
690 #[tokio::test]
699 async fn lease_capacity_override_beats_the_managers_own() {
700 let manager = EventQueueManager::new();
701 let task_id = TaskId::new("override");
702 let crate::streaming::QueueLease::Created { writer, reader, .. } =
703 manager.lease(&task_id, false, Some(1)).await
704 else {
705 panic!("first lease must create a queue");
706 };
707 let mut reader = reader;
708
709 writer
710 .write(make_status_event("override", TaskState::Working))
711 .await
712 .expect("first write");
713 writer
714 .write(make_status_event("override", TaskState::Completed))
715 .await
716 .expect("second write");
717
718 assert!(
719 reader.read().await.expect("reader is still open").is_err(),
720 "the override asked for capacity 1; buffering both events is what \
721 the manager's own default (256) would do"
722 );
723 }
724
725 #[tokio::test]
728 async fn lease_without_an_override_uses_the_managers_capacity() {
729 let manager = EventQueueManager::with_capacity(1);
730 let task_id = TaskId::new("no-override");
731 let crate::streaming::QueueLease::Created { writer, reader, .. } =
732 manager.lease(&task_id, false, None).await
733 else {
734 panic!("first lease must create a queue");
735 };
736 let mut reader = reader;
737
738 writer
739 .write(make_status_event("no-override", TaskState::Working))
740 .await
741 .expect("first write");
742 writer
743 .write(make_status_event("no-override", TaskState::Completed))
744 .await
745 .expect("second write");
746
747 assert!(
748 reader.read().await.expect("reader is still open").is_err(),
749 "None must fall back to the manager's capacity of 1, not to the default"
750 );
751 }
752
753 #[tokio::test]
763 async fn first_queue_is_tracked_when_a_concurrency_limit_is_set() {
764 let manager = EventQueueManager::new().with_max_concurrent_queues(1);
765
766 let first = TaskId::new("q1");
767 let (_w, reader, persistence) = manager.get_or_create_with_persistence(&first).await;
768 assert!(
769 reader.is_some(),
770 "the first queue is below the limit and must be tracked, \
771 with a reader; None means the at-capacity branch was taken"
772 );
773 assert!(
774 persistence.is_some(),
775 "a tracked queue gets a persistence rx"
776 );
777 assert_eq!(
778 manager.active_count().await,
779 1,
780 "first queue must be stored"
781 );
782
783 let second = TaskId::new("q2");
785 let (_w2, reader2, _p2) = manager.get_or_create_with_persistence(&second).await;
786 assert!(
787 reader2.is_none(),
788 "the second queue exceeds the limit and must not be tracked"
789 );
790 assert_eq!(
791 manager.active_count().await,
792 1,
793 "an over-limit queue must not be stored"
794 );
795 }
796
797 #[tokio::test]
800 async fn raw_subscribe_returns_a_receiver_for_a_live_queue() {
801 let manager = EventQueueManager::new();
802 let task_id = TaskId::new("raw");
803 let (_writer, _reader) = manager.get_or_create(&task_id).await;
804
805 assert!(
806 manager.raw_subscribe(&task_id).await.is_some(),
807 "a live queue must yield a receiver"
808 );
809 assert!(
810 manager
811 .raw_subscribe(&TaskId::new("absent"))
812 .await
813 .is_none(),
814 "an unknown task must yield None — pins that Some is not blanket"
815 );
816 }
817
818 #[tokio::test]
825 async fn subscribe_with_snapshot_returns_a_reader_that_yields_the_snapshot() {
826 let manager = EventQueueManager::new();
827 let task_id = TaskId::new("snap");
828 let (_writer, _reader) = manager.get_or_create(&task_id).await;
829
830 let snapshot = make_status_event("snap", TaskState::Working);
831 let reader = manager.subscribe_with_snapshot(&task_id, snapshot).await;
832 let mut reader = reader.expect("a live queue must yield a reader");
833
834 let first = reader
835 .read()
836 .await
837 .expect("reader is open")
838 .expect("snapshot is delivered as Ok");
839 match first {
840 StreamResponse::StatusUpdate(ev) => {
841 assert_eq!(
842 ev.status.state,
843 TaskState::Working,
844 "snapshot arrives first"
845 );
846 }
847 other => panic!("expected the snapshot StatusUpdate first, got {other:?}"),
848 }
849
850 assert!(
851 manager
852 .subscribe_with_snapshot(
853 &TaskId::new("absent"),
854 make_status_event("absent", TaskState::Working)
855 )
856 .await
857 .is_none(),
858 "an unknown task must yield None — pins that Some is not blanket"
859 );
860 }
861}