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)]
266 pub(crate) async fn lease(&self, task_id: &TaskId, with_persistence: bool) -> QueueLease {
267 let mut map = self.writers.write().await;
268 let lease = if map.contains_key(task_id) {
269 QueueLease::Existing
270 } else if self
271 .max_concurrent_queues
272 .is_some_and(|max| map.len() >= max)
273 {
274 QueueLease::CapacityExhausted
275 } else if with_persistence {
276 let (writer, reader, persistence_rx) = new_in_memory_queue_with_persistence(
277 self.capacity,
278 self.max_event_size,
279 self.write_timeout,
280 );
281 let writer = Arc::new(writer);
282 map.insert(task_id.clone(), Arc::clone(&writer));
283 QueueLease::Created {
284 writer,
285 reader,
286 persistence_rx: Some(persistence_rx),
287 }
288 } else {
289 let (writer, reader) = new_in_memory_queue_with_options(
290 self.capacity,
291 self.max_event_size,
292 self.write_timeout,
293 );
294 let writer = Arc::new(writer);
295 map.insert(task_id.clone(), Arc::clone(&writer));
296 QueueLease::Created {
297 writer,
298 reader,
299 persistence_rx: None,
300 }
301 };
302 let queue_count = map.len();
303 drop(map);
304 if let Some(ref metrics) = self.metrics {
305 metrics.on_queue_depth_change(queue_count);
306 }
307 lease
308 }
309
310 pub(crate) async fn writer_for_cancel(&self, task_id: &TaskId) -> Arc<InMemoryQueueWriter> {
320 {
321 let map = self.writers.read().await;
322 if let Some(writer) = map.get(task_id) {
323 return Arc::clone(writer);
324 }
325 }
326 let (writer, _reader) = new_in_memory_queue_with_options(
327 self.capacity,
328 self.max_event_size,
329 self.write_timeout,
330 );
331 Arc::new(writer)
332 }
333
334 pub async fn subscribe(&self, task_id: &TaskId) -> Option<InMemoryQueueReader> {
342 let map = self.writers.read().await;
343 map.get(task_id).map(|writer| writer.subscribe())
344 }
345
346 pub(crate) async fn raw_subscribe(
353 &self,
354 task_id: &TaskId,
355 ) -> Option<tokio::sync::broadcast::Receiver<A2aResult<StreamResponse>>> {
356 let map = self.writers.read().await;
357 map.get(task_id).map(|writer| writer.raw_subscribe())
358 }
359
360 pub async fn subscribe_with_snapshot(
369 &self,
370 task_id: &TaskId,
371 snapshot: StreamResponse,
372 ) -> Option<InMemoryQueueReader> {
373 let map = self.writers.read().await;
374 let writer = map.get(task_id)?;
375 let rx = writer.raw_subscribe();
379 drop(map);
380 Some(InMemoryQueueReader::with_first_event(rx, snapshot))
381 }
382
383 pub async fn destroy(&self, task_id: &TaskId) {
385 let mut map = self.writers.write().await;
386 map.remove(task_id);
387 let queue_count = map.len();
388 drop(map);
389 if let Some(ref metrics) = self.metrics {
390 metrics.on_queue_depth_change(queue_count);
391 }
392 }
393
394 pub async fn active_count(&self) -> usize {
396 let map = self.writers.read().await;
397 map.len()
398 }
399
400 pub(crate) async fn has_queue(&self, task_id: &TaskId) -> bool {
406 self.writers.read().await.contains_key(task_id)
407 }
408
409 #[must_use]
412 pub(crate) const fn max_concurrent_queues(&self) -> Option<usize> {
413 self.max_concurrent_queues
414 }
415
416 pub async fn destroy_all(&self) {
418 let mut map = self.writers.write().await;
419 map.clear();
420 }
421}
422
423#[cfg(test)]
424mod tests {
425 use super::*;
426 use crate::streaming::event_queue::{EventQueueReader, EventQueueWriter};
427 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
428 use a2a_protocol_types::task::{ContextId, TaskState, TaskStatus};
429
430 fn make_status_event(task_id: &str, state: TaskState) -> StreamResponse {
432 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
433 task_id: TaskId::new(task_id),
434 context_id: ContextId::new("ctx-test"),
435 status: TaskStatus {
436 state,
437 message: None,
438 timestamp: None,
439 },
440 metadata: None,
441 })
442 }
443
444 #[test]
447 fn max_concurrent_queues_reports_configured_limit() {
448 assert_eq!(EventQueueManager::new().max_concurrent_queues(), None);
450 assert_eq!(
452 EventQueueManager::new()
453 .with_max_concurrent_queues(42)
454 .max_concurrent_queues(),
455 Some(42)
456 );
457 }
458
459 #[tokio::test]
460 async fn manager_get_or_create_new_task() {
461 let manager = EventQueueManager::new();
462 let task_id = TaskId::new("task-1");
463
464 let (writer, reader) = manager.get_or_create(&task_id).await;
465 assert!(
466 reader.is_some(),
467 "first get_or_create should return a reader"
468 );
469
470 writer
472 .write(make_status_event("task-1", TaskState::Working))
473 .await
474 .expect("write through manager writer should succeed");
475
476 assert_eq!(
477 manager.active_count().await,
478 1,
479 "should have 1 active queue"
480 );
481 }
482
483 #[tokio::test]
484 async fn manager_get_or_create_existing_task_returns_no_reader() {
485 let manager = EventQueueManager::new();
486 let task_id = TaskId::new("task-1");
487
488 let (_w1, r1) = manager.get_or_create(&task_id).await;
489 assert!(r1.is_some(), "first call should return a reader");
490
491 let (_w2, r2) = manager.get_or_create(&task_id).await;
492 assert!(
493 r2.is_none(),
494 "second call for same task should return None reader"
495 );
496
497 assert_eq!(
498 manager.active_count().await,
499 1,
500 "should still have only 1 active queue"
501 );
502 }
503
504 #[tokio::test]
505 async fn manager_subscribe_existing_task() {
506 use crate::streaming::event_queue::EventQueueReader;
507
508 let manager = EventQueueManager::new();
509 let task_id = TaskId::new("task-1");
510
511 let (writer, _reader) = manager.get_or_create(&task_id).await;
512
513 let sub = manager.subscribe(&task_id).await;
514 assert!(
515 sub.is_some(),
516 "subscribe should return a reader for existing task"
517 );
518
519 let mut sub_reader = sub.unwrap();
520 writer
521 .write(make_status_event("task-1", TaskState::Working))
522 .await
523 .expect("write should succeed");
524 drop(writer);
525
526 let r = sub_reader.read().await;
527 assert!(r.is_some(), "subscriber should receive the event");
528 }
529
530 #[tokio::test]
531 async fn manager_subscribe_nonexistent_task_returns_none() {
532 let manager = EventQueueManager::new();
533 let task_id = TaskId::new("no-such-task");
534
535 let sub = manager.subscribe(&task_id).await;
536 assert!(
537 sub.is_none(),
538 "subscribe should return None for nonexistent task"
539 );
540 }
541
542 #[tokio::test]
543 async fn manager_destroy_removes_queue() {
544 let manager = EventQueueManager::new();
545 let task_id = TaskId::new("task-1");
546
547 let (_writer, _reader) = manager.get_or_create(&task_id).await;
548 assert_eq!(manager.active_count().await, 1);
549
550 manager.destroy(&task_id).await;
551 assert_eq!(
552 manager.active_count().await,
553 0,
554 "destroy should remove the queue"
555 );
556 }
557
558 #[tokio::test]
559 async fn manager_destroy_all_clears_queues() {
560 let manager = EventQueueManager::new();
561
562 let _q1 = manager.get_or_create(&TaskId::new("t1")).await;
563 let _q2 = manager.get_or_create(&TaskId::new("t2")).await;
564 assert_eq!(manager.active_count().await, 2);
565
566 manager.destroy_all().await;
567 assert_eq!(
568 manager.active_count().await,
569 0,
570 "destroy_all should clear all queues"
571 );
572 }
573
574 #[tokio::test]
575 async fn lease_reports_existing_and_has_queue() {
576 let manager = EventQueueManager::new();
577 let task = TaskId::new("t-lease");
578
579 assert!(matches!(
581 manager.lease(&task, true).await,
582 QueueLease::Created { .. }
583 ));
584 assert!(manager.has_queue(&task).await, "queue should now be live");
585
586 assert!(matches!(
589 manager.lease(&task, true).await,
590 QueueLease::Existing
591 ));
592
593 assert!(!manager.has_queue(&TaskId::new("other")).await);
595 }
596
597 #[tokio::test]
598 async fn manager_max_concurrent_queues_enforced() {
599 let manager = EventQueueManager::new().with_max_concurrent_queues(1);
600
601 let (_w1, r1) = manager.get_or_create(&TaskId::new("t1")).await;
602 assert!(r1.is_some(), "first queue should be created successfully");
603
604 let (_w2, r2) = manager.get_or_create(&TaskId::new("t2")).await;
606 assert!(
607 r2.is_none(),
608 "second queue should return None reader when limit is reached"
609 );
610 assert_eq!(
611 manager.active_count().await,
612 1,
613 "should still have only 1 queue (second was not stored)"
614 );
615 }
616
617 #[tokio::test]
618 async fn manager_with_capacity_and_max_event_size() {
619 let manager = EventQueueManager::with_capacity(4).with_max_event_size(10); let task_id = TaskId::new("t1");
622 let (writer, _reader) = manager.get_or_create(&task_id).await;
623
624 let event = make_status_event("t1", TaskState::Working);
625 let result = writer.write(event).await;
626 assert!(
627 result.is_err(),
628 "event should be rejected by the size limit configured on the manager"
629 );
630 }
631
632 #[tokio::test]
648 async fn with_capacity_uses_the_given_capacity_not_the_default() {
649 let manager = EventQueueManager::with_capacity(1);
650 let task_id = TaskId::new("cap");
651 let (writer, reader) = manager.get_or_create(&task_id).await;
652 let mut reader = reader.expect("first get_or_create yields a reader");
653
654 writer
656 .write(make_status_event("cap", TaskState::Working))
657 .await
658 .expect("first write");
659 writer
660 .write(make_status_event("cap", TaskState::Completed))
661 .await
662 .expect("second write");
663
664 let first = reader.read().await.expect("reader is still open");
665 assert!(
666 first.is_err(),
667 "a capacity-1 queue must surface an overrun to the reader; \
668 got Ok, which is what DEFAULT_QUEUE_CAPACITY (256) would give"
669 );
670 }
671
672 #[tokio::test]
682 async fn first_queue_is_tracked_when_a_concurrency_limit_is_set() {
683 let manager = EventQueueManager::new().with_max_concurrent_queues(1);
684
685 let first = TaskId::new("q1");
686 let (_w, reader, persistence) = manager.get_or_create_with_persistence(&first).await;
687 assert!(
688 reader.is_some(),
689 "the first queue is below the limit and must be tracked, \
690 with a reader; None means the at-capacity branch was taken"
691 );
692 assert!(
693 persistence.is_some(),
694 "a tracked queue gets a persistence rx"
695 );
696 assert_eq!(
697 manager.active_count().await,
698 1,
699 "first queue must be stored"
700 );
701
702 let second = TaskId::new("q2");
704 let (_w2, reader2, _p2) = manager.get_or_create_with_persistence(&second).await;
705 assert!(
706 reader2.is_none(),
707 "the second queue exceeds the limit and must not be tracked"
708 );
709 assert_eq!(
710 manager.active_count().await,
711 1,
712 "an over-limit queue must not be stored"
713 );
714 }
715
716 #[tokio::test]
719 async fn raw_subscribe_returns_a_receiver_for_a_live_queue() {
720 let manager = EventQueueManager::new();
721 let task_id = TaskId::new("raw");
722 let (_writer, _reader) = manager.get_or_create(&task_id).await;
723
724 assert!(
725 manager.raw_subscribe(&task_id).await.is_some(),
726 "a live queue must yield a receiver"
727 );
728 assert!(
729 manager
730 .raw_subscribe(&TaskId::new("absent"))
731 .await
732 .is_none(),
733 "an unknown task must yield None — pins that Some is not blanket"
734 );
735 }
736
737 #[tokio::test]
744 async fn subscribe_with_snapshot_returns_a_reader_that_yields_the_snapshot() {
745 let manager = EventQueueManager::new();
746 let task_id = TaskId::new("snap");
747 let (_writer, _reader) = manager.get_or_create(&task_id).await;
748
749 let snapshot = make_status_event("snap", TaskState::Working);
750 let reader = manager.subscribe_with_snapshot(&task_id, snapshot).await;
751 let mut reader = reader.expect("a live queue must yield a reader");
752
753 let first = reader
754 .read()
755 .await
756 .expect("reader is open")
757 .expect("snapshot is delivered as Ok");
758 match first {
759 StreamResponse::StatusUpdate(ev) => {
760 assert_eq!(
761 ev.status.state,
762 TaskState::Working,
763 "snapshot arrives first"
764 );
765 }
766 other => panic!("expected the snapshot StatusUpdate first, got {other:?}"),
767 }
768
769 assert!(
770 manager
771 .subscribe_with_snapshot(
772 &TaskId::new("absent"),
773 make_status_event("absent", TaskState::Working)
774 )
775 .await
776 .is_none(),
777 "an unknown task must yield None — pins that Some is not blanket"
778 );
779 }
780}