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 new_in_memory_queue_with_options, new_in_memory_queue_with_persistence, InMemoryQueueReader,
19 InMemoryQueueWriter, DEFAULT_MAX_EVENT_SIZE, DEFAULT_QUEUE_CAPACITY, DEFAULT_WRITE_TIMEOUT,
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 #[must_use]
154 pub const fn with_max_concurrent_queues(mut self, max: usize) -> Self {
155 self.max_concurrent_queues = Some(max);
156 self
157 }
158
159 pub async fn get_or_create(
170 &self,
171 task_id: &TaskId,
172 ) -> (Arc<InMemoryQueueWriter>, Option<InMemoryQueueReader>) {
173 let mut map = self.writers.write().await;
174 #[allow(clippy::option_if_let_else)]
175 let result = if let Some(existing) = map.get(task_id) {
176 (Arc::clone(existing), None)
177 } else if self
178 .max_concurrent_queues
179 .is_some_and(|max| map.len() >= max)
180 {
181 let (writer, _reader) = new_in_memory_queue_with_options(
184 self.capacity,
185 self.max_event_size,
186 self.write_timeout,
187 );
188 (Arc::new(writer), None)
189 } else {
190 let (writer, reader) = new_in_memory_queue_with_options(
191 self.capacity,
192 self.max_event_size,
193 self.write_timeout,
194 );
195 let writer = Arc::new(writer);
196 map.insert(task_id.clone(), Arc::clone(&writer));
197 (writer, Some(reader))
198 };
199 let queue_count = map.len();
200 drop(map);
201 if let Some(ref metrics) = self.metrics {
202 metrics.on_queue_depth_change(queue_count);
203 }
204 result
205 }
206
207 pub async fn get_or_create_with_persistence(
215 &self,
216 task_id: &TaskId,
217 ) -> (
218 Arc<InMemoryQueueWriter>,
219 Option<InMemoryQueueReader>,
220 Option<tokio::sync::mpsc::Receiver<A2aResult<StreamResponse>>>,
221 ) {
222 let mut map = self.writers.write().await;
223 #[allow(clippy::option_if_let_else)]
224 let result = if let Some(existing) = map.get(task_id) {
225 (Arc::clone(existing), None, None)
226 } else if self
227 .max_concurrent_queues
228 .is_some_and(|max| map.len() >= max)
229 {
230 let (writer, _reader) = new_in_memory_queue_with_options(
231 self.capacity,
232 self.max_event_size,
233 self.write_timeout,
234 );
235 (Arc::new(writer), None, None)
236 } else {
237 let (writer, reader, persistence_rx) = new_in_memory_queue_with_persistence(
238 self.capacity,
239 self.max_event_size,
240 self.write_timeout,
241 );
242 let writer = Arc::new(writer);
243 map.insert(task_id.clone(), Arc::clone(&writer));
244 (writer, Some(reader), Some(persistence_rx))
245 };
246 let queue_count = map.len();
247 drop(map);
248 if let Some(ref metrics) = self.metrics {
249 metrics.on_queue_depth_change(queue_count);
250 }
251 result
252 }
253
254 #[allow(clippy::option_if_let_else)]
272 pub(crate) async fn lease(
273 &self,
274 task_id: &TaskId,
275 with_persistence: bool,
276 capacity: Option<usize>,
277 ) -> QueueLease {
278 let capacity = capacity.unwrap_or(self.capacity);
279 let mut map = self.writers.write().await;
280 let lease = if map.contains_key(task_id) {
281 QueueLease::Existing
282 } else if self
283 .max_concurrent_queues
284 .is_some_and(|max| map.len() >= max)
285 {
286 QueueLease::CapacityExhausted
287 } else if with_persistence {
288 let (writer, reader, persistence_rx) = new_in_memory_queue_with_persistence(
289 capacity,
290 self.max_event_size,
291 self.write_timeout,
292 );
293 let writer = Arc::new(writer);
294 map.insert(task_id.clone(), Arc::clone(&writer));
295 QueueLease::Created {
296 writer,
297 reader,
298 persistence_rx: Some(persistence_rx),
299 }
300 } else {
301 let (writer, reader) =
302 new_in_memory_queue_with_options(capacity, self.max_event_size, self.write_timeout);
303 let writer = Arc::new(writer);
304 map.insert(task_id.clone(), Arc::clone(&writer));
305 QueueLease::Created {
306 writer,
307 reader,
308 persistence_rx: None,
309 }
310 };
311 let queue_count = map.len();
312 drop(map);
313 if let Some(ref metrics) = self.metrics {
314 metrics.on_queue_depth_change(queue_count);
315 }
316 lease
317 }
318
319 pub(crate) async fn writer_for_cancel(&self, task_id: &TaskId) -> Arc<InMemoryQueueWriter> {
329 {
330 let map = self.writers.read().await;
331 if let Some(writer) = map.get(task_id) {
332 return Arc::clone(writer);
333 }
334 }
335 let (writer, _reader) = new_in_memory_queue_with_options(
336 self.capacity,
337 self.max_event_size,
338 self.write_timeout,
339 );
340 Arc::new(writer)
341 }
342
343 pub async fn subscribe(&self, task_id: &TaskId) -> Option<InMemoryQueueReader> {
351 let map = self.writers.read().await;
352 map.get(task_id).map(|writer| writer.subscribe())
353 }
354
355 pub(crate) async fn raw_subscribe(
362 &self,
363 task_id: &TaskId,
364 ) -> Option<tokio::sync::broadcast::Receiver<A2aResult<StreamResponse>>> {
365 let map = self.writers.read().await;
366 map.get(task_id).map(|writer| writer.raw_subscribe())
367 }
368
369 pub async fn subscribe_with_snapshot(
378 &self,
379 task_id: &TaskId,
380 snapshot: StreamResponse,
381 ) -> Option<InMemoryQueueReader> {
382 let map = self.writers.read().await;
383 let writer = map.get(task_id)?;
384 let rx = writer.raw_subscribe();
388 drop(map);
389 Some(InMemoryQueueReader::with_first_event(rx, snapshot))
390 }
391
392 pub async fn destroy(&self, task_id: &TaskId) {
394 let mut map = self.writers.write().await;
395 map.remove(task_id);
396 let queue_count = map.len();
397 drop(map);
398 if let Some(ref metrics) = self.metrics {
399 metrics.on_queue_depth_change(queue_count);
400 }
401 }
402
403 pub async fn active_count(&self) -> usize {
405 let map = self.writers.read().await;
406 map.len()
407 }
408
409 pub(crate) async fn has_queue(&self, task_id: &TaskId) -> bool {
415 self.writers.read().await.contains_key(task_id)
416 }
417
418 #[must_use]
421 pub(crate) const fn max_concurrent_queues(&self) -> Option<usize> {
422 self.max_concurrent_queues
423 }
424
425 pub async fn destroy_all(&self) {
427 let mut map = self.writers.write().await;
428 map.clear();
429 }
430}
431
432#[cfg(test)]
433mod tests {
434 use super::*;
435 use crate::streaming::event_queue::{EventQueueReader, EventQueueWriter};
436 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
437 use a2a_protocol_types::task::{ContextId, TaskState, TaskStatus};
438
439 fn make_status_event(task_id: &str, state: TaskState) -> StreamResponse {
441 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
442 task_id: TaskId::new(task_id),
443 context_id: ContextId::new("ctx-test"),
444 status: TaskStatus {
445 state,
446 message: None,
447 timestamp: None,
448 },
449 metadata: None,
450 })
451 }
452
453 #[test]
456 fn max_concurrent_queues_reports_configured_limit() {
457 assert_eq!(EventQueueManager::new().max_concurrent_queues(), None);
459 assert_eq!(
461 EventQueueManager::new()
462 .with_max_concurrent_queues(42)
463 .max_concurrent_queues(),
464 Some(42)
465 );
466 }
467
468 #[tokio::test]
469 async fn manager_get_or_create_new_task() {
470 let manager = EventQueueManager::new();
471 let task_id = TaskId::new("task-1");
472
473 let (writer, reader) = manager.get_or_create(&task_id).await;
474 assert!(
475 reader.is_some(),
476 "first get_or_create should return a reader"
477 );
478
479 writer
481 .write(make_status_event("task-1", TaskState::Working))
482 .await
483 .expect("write through manager writer should succeed");
484
485 assert_eq!(
486 manager.active_count().await,
487 1,
488 "should have 1 active queue"
489 );
490 }
491
492 #[tokio::test]
493 async fn manager_get_or_create_existing_task_returns_no_reader() {
494 let manager = EventQueueManager::new();
495 let task_id = TaskId::new("task-1");
496
497 let (_w1, r1) = manager.get_or_create(&task_id).await;
498 assert!(r1.is_some(), "first call should return a reader");
499
500 let (_w2, r2) = manager.get_or_create(&task_id).await;
501 assert!(
502 r2.is_none(),
503 "second call for same task should return None reader"
504 );
505
506 assert_eq!(
507 manager.active_count().await,
508 1,
509 "should still have only 1 active queue"
510 );
511 }
512
513 #[tokio::test]
514 async fn manager_subscribe_existing_task() {
515 use crate::streaming::event_queue::EventQueueReader;
516
517 let manager = EventQueueManager::new();
518 let task_id = TaskId::new("task-1");
519
520 let (writer, _reader) = manager.get_or_create(&task_id).await;
521
522 let sub = manager.subscribe(&task_id).await;
523 assert!(
524 sub.is_some(),
525 "subscribe should return a reader for existing task"
526 );
527
528 let mut sub_reader = sub.unwrap();
529 writer
530 .write(make_status_event("task-1", TaskState::Working))
531 .await
532 .expect("write should succeed");
533 drop(writer);
534
535 let r = sub_reader.read().await;
536 assert!(r.is_some(), "subscriber should receive the event");
537 }
538
539 #[tokio::test]
540 async fn manager_subscribe_nonexistent_task_returns_none() {
541 let manager = EventQueueManager::new();
542 let task_id = TaskId::new("no-such-task");
543
544 let sub = manager.subscribe(&task_id).await;
545 assert!(
546 sub.is_none(),
547 "subscribe should return None for nonexistent task"
548 );
549 }
550
551 #[tokio::test]
552 async fn manager_destroy_removes_queue() {
553 let manager = EventQueueManager::new();
554 let task_id = TaskId::new("task-1");
555
556 let (_writer, _reader) = manager.get_or_create(&task_id).await;
557 assert_eq!(manager.active_count().await, 1);
558
559 manager.destroy(&task_id).await;
560 assert_eq!(
561 manager.active_count().await,
562 0,
563 "destroy should remove the queue"
564 );
565 }
566
567 #[tokio::test]
568 async fn manager_destroy_all_clears_queues() {
569 let manager = EventQueueManager::new();
570
571 let _q1 = manager.get_or_create(&TaskId::new("t1")).await;
572 let _q2 = manager.get_or_create(&TaskId::new("t2")).await;
573 assert_eq!(manager.active_count().await, 2);
574
575 manager.destroy_all().await;
576 assert_eq!(
577 manager.active_count().await,
578 0,
579 "destroy_all should clear all queues"
580 );
581 }
582
583 #[tokio::test]
584 async fn lease_reports_existing_and_has_queue() {
585 let manager = EventQueueManager::new();
586 let task = TaskId::new("t-lease");
587
588 assert!(matches!(
590 manager.lease(&task, true, None).await,
591 QueueLease::Created { .. }
592 ));
593 assert!(manager.has_queue(&task).await, "queue should now be live");
594
595 assert!(matches!(
598 manager.lease(&task, true, None).await,
599 QueueLease::Existing
600 ));
601
602 assert!(!manager.has_queue(&TaskId::new("other")).await);
604 }
605
606 #[tokio::test]
607 async fn manager_max_concurrent_queues_enforced() {
608 let manager = EventQueueManager::new().with_max_concurrent_queues(1);
609
610 let (_w1, r1) = manager.get_or_create(&TaskId::new("t1")).await;
611 assert!(r1.is_some(), "first queue should be created successfully");
612
613 let (_w2, r2) = manager.get_or_create(&TaskId::new("t2")).await;
615 assert!(
616 r2.is_none(),
617 "second queue should return None reader when limit is reached"
618 );
619 assert_eq!(
620 manager.active_count().await,
621 1,
622 "should still have only 1 queue (second was not stored)"
623 );
624 }
625
626 #[tokio::test]
627 async fn manager_with_capacity_and_max_event_size() {
628 let manager = EventQueueManager::with_capacity(4).with_max_event_size(10); let task_id = TaskId::new("t1");
631 let (writer, _reader) = manager.get_or_create(&task_id).await;
632
633 let event = make_status_event("t1", TaskState::Working);
634 let result = writer.write(event).await;
635 assert!(
636 result.is_err(),
637 "event should be rejected by the size limit configured on the manager"
638 );
639 }
640
641 #[tokio::test]
657 async fn with_capacity_uses_the_given_capacity_not_the_default() {
658 let manager = EventQueueManager::with_capacity(1);
659 let task_id = TaskId::new("cap");
660 let (writer, reader) = manager.get_or_create(&task_id).await;
661 let mut reader = reader.expect("first get_or_create yields a reader");
662
663 writer
665 .write(make_status_event("cap", TaskState::Working))
666 .await
667 .expect("first write");
668 writer
669 .write(make_status_event("cap", TaskState::Completed))
670 .await
671 .expect("second write");
672
673 let first = reader.read().await.expect("reader is still open");
674 assert!(
675 first.is_err(),
676 "a capacity-1 queue must surface an overrun to the reader; \
677 got Ok, which is what DEFAULT_QUEUE_CAPACITY (256) would give"
678 );
679 }
680
681 #[tokio::test]
690 async fn lease_capacity_override_beats_the_managers_own() {
691 let manager = EventQueueManager::new();
692 let task_id = TaskId::new("override");
693 let crate::streaming::QueueLease::Created { writer, reader, .. } =
694 manager.lease(&task_id, false, Some(1)).await
695 else {
696 panic!("first lease must create a queue");
697 };
698 let mut reader = reader;
699
700 writer
701 .write(make_status_event("override", TaskState::Working))
702 .await
703 .expect("first write");
704 writer
705 .write(make_status_event("override", TaskState::Completed))
706 .await
707 .expect("second write");
708
709 assert!(
710 reader.read().await.expect("reader is still open").is_err(),
711 "the override asked for capacity 1; buffering both events is what \
712 the manager's own default (256) would do"
713 );
714 }
715
716 #[tokio::test]
719 async fn lease_without_an_override_uses_the_managers_capacity() {
720 let manager = EventQueueManager::with_capacity(1);
721 let task_id = TaskId::new("no-override");
722 let crate::streaming::QueueLease::Created { writer, reader, .. } =
723 manager.lease(&task_id, false, None).await
724 else {
725 panic!("first lease must create a queue");
726 };
727 let mut reader = reader;
728
729 writer
730 .write(make_status_event("no-override", TaskState::Working))
731 .await
732 .expect("first write");
733 writer
734 .write(make_status_event("no-override", TaskState::Completed))
735 .await
736 .expect("second write");
737
738 assert!(
739 reader.read().await.expect("reader is still open").is_err(),
740 "None must fall back to the manager's capacity of 1, not to the default"
741 );
742 }
743
744 #[tokio::test]
754 async fn first_queue_is_tracked_when_a_concurrency_limit_is_set() {
755 let manager = EventQueueManager::new().with_max_concurrent_queues(1);
756
757 let first = TaskId::new("q1");
758 let (_w, reader, persistence) = manager.get_or_create_with_persistence(&first).await;
759 assert!(
760 reader.is_some(),
761 "the first queue is below the limit and must be tracked, \
762 with a reader; None means the at-capacity branch was taken"
763 );
764 assert!(
765 persistence.is_some(),
766 "a tracked queue gets a persistence rx"
767 );
768 assert_eq!(
769 manager.active_count().await,
770 1,
771 "first queue must be stored"
772 );
773
774 let second = TaskId::new("q2");
776 let (_w2, reader2, _p2) = manager.get_or_create_with_persistence(&second).await;
777 assert!(
778 reader2.is_none(),
779 "the second queue exceeds the limit and must not be tracked"
780 );
781 assert_eq!(
782 manager.active_count().await,
783 1,
784 "an over-limit queue must not be stored"
785 );
786 }
787
788 #[tokio::test]
791 async fn raw_subscribe_returns_a_receiver_for_a_live_queue() {
792 let manager = EventQueueManager::new();
793 let task_id = TaskId::new("raw");
794 let (_writer, _reader) = manager.get_or_create(&task_id).await;
795
796 assert!(
797 manager.raw_subscribe(&task_id).await.is_some(),
798 "a live queue must yield a receiver"
799 );
800 assert!(
801 manager
802 .raw_subscribe(&TaskId::new("absent"))
803 .await
804 .is_none(),
805 "an unknown task must yield None — pins that Some is not blanket"
806 );
807 }
808
809 #[tokio::test]
816 async fn subscribe_with_snapshot_returns_a_reader_that_yields_the_snapshot() {
817 let manager = EventQueueManager::new();
818 let task_id = TaskId::new("snap");
819 let (_writer, _reader) = manager.get_or_create(&task_id).await;
820
821 let snapshot = make_status_event("snap", TaskState::Working);
822 let reader = manager.subscribe_with_snapshot(&task_id, snapshot).await;
823 let mut reader = reader.expect("a live queue must yield a reader");
824
825 let first = reader
826 .read()
827 .await
828 .expect("reader is open")
829 .expect("snapshot is delivered as Ok");
830 match first {
831 StreamResponse::StatusUpdate(ev) => {
832 assert_eq!(
833 ev.status.state,
834 TaskState::Working,
835 "snapshot arrives first"
836 );
837 }
838 other => panic!("expected the snapshot StatusUpdate first, got {other:?}"),
839 }
840
841 assert!(
842 manager
843 .subscribe_with_snapshot(
844 &TaskId::new("absent"),
845 make_status_event("absent", TaskState::Working)
846 )
847 .await
848 .is_none(),
849 "an unknown task must yield None — pins that Some is not blanket"
850 );
851 }
852}