Merge "Add QueueManager abstraction and implement in memory version." into main

Author: pstephengoogle

src/a2a/server/events/__init__.py

@@ -1,5 +1,14 @@
 from a2a.server.events.event_consumer import EventConsumer
 from a2a.server.events.event_queue import Event, EventQueue
+from a2a.server.events.queue_manager import QueueManager, TaskQueueExists, NoTaskQueue
+from a2a.server.events.in_memory_queue_manager import InMemoryQueueManager
 
-
-__all__ = ['Event', 'EventConsumer', 'EventQueue']
+__all__ = [
+    'Event',
+    'EventConsumer',
+    'EventQueue',
+    'QueueManager',
+    'TaskQueueExists',
+    'NoTaskQueue',
+    'InMemoryQueueManager',
+]

src/a2a/server/events/in_memory_queue_manager.py

@@ -0,0 +1,49 @@
+import asyncio
+from a2a.server.events.event_queue import EventQueue
+from a2a.server.events.queue_manager import QueueManager, TaskQueueExists, NoTaskQueue
+
+class InMemoryQueueManager(QueueManager):
+    """InMemoryQueueManager is used for a single binary management.
+
+    This implements the QueueManager but requires all incoming interactions
+    to hit the same binary that manages the queues.
+
+    This works for single binary solution. Needs a distributed approach for
+    true scalable deployment.
+    """
+
+    def __init__(self):
+        self._task_queue: dict[str, EventQueue] = {}
+        self._lock = asyncio.Lock()
+
+    async def add(self, task_id: str, queue: EventQueue):
+        async with self._lock:
+            if task_id in self._task_queue:
+                raise TaskQueueExists()
+            self._task_queue[task_id] = queue
+
+    async def get(self, task_id: str) -> EventQueue | None:
+        async with self._lock:
+            if task_id not in self._task_queue:
+                return None
+            return self._task_queue[task_id]
+
+    async def tap(self, task_id: str) -> EventQueue | None:
+        async with self._lock:
+            if task_id not in self._task_queue:
+                return None
+            return self._task_queue[task_id].tap()
+
+    async def close(self, task_id: str):
+        async with self._lock:
+            if task_id not in self._task_queue:
+                raise NoTaskQueue()
+            del self._task_queue[task_id]
+
+    async def create_or_tap(self, task_id: str) -> EventQueue:
+        async with self._lock:
+            if task_id not in self._task_queue:
+                queue = EventQueue()
+                self._task_queue[task_id] = queue
+                return queue
+            return self._task_queue[task_id].tap()

src/a2a/server/events/queue_manager.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from a2a.server.events.event_queue import EventQueue
+
+
+class QueueManager(ABC):
+    """Interface for managing the event queue lifecycles."""
+
+    @abstractmethod
+    async def add(self, task_id: str, queue: EventQueue):
+        pass
+
+    @abstractmethod
+    async def get(self, task_id: str) -> EventQueue | None:
+        pass
+
+    @abstractmethod
+    async def tap(self, task_id: str) -> EventQueue | None:
+        pass
+
+    @abstractmethod
+    async def close(self, task_id: str):
+        pass
+
+    @abstractmethod
+    async def create_or_tap(self, task_id: str) -> EventQueue:
+        pass
+
+
+class TaskQueueExists(Exception):
+    pass
+
+class NoTaskQueue(Exception):
+    pass

src/a2a/server/request_handlers/default_request_handler.py

@@ -4,7 +4,15 @@
 from collections.abc import AsyncGenerator
 
 from a2a.server.agent_execution import AgentExecutor, RequestContext
-from a2a.server.events import Event, EventConsumer, EventQueue
+from a2a.server.events import (
+    EventConsumer,
+    EventQueue,
+    Event,
+    QueueManager,
+    TaskQueueExists,
+    NoTaskQueue,
+    InMemoryQueueManager,
+)
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.server.tasks import ResultAggregator, TaskManager, TaskStore
 from a2a.types import (
@@ -28,13 +36,14 @@ class DefaultRequestHandler(RequestHandler):
     """Default request handler for all incoming requests."""
 
     def __init__(
-        self, agent_executor: AgentExecutor, task_store: TaskStore
+        self,
+        agent_executor: AgentExecutor,
+        task_store: TaskStore,
+        queue_manager: QueueManager = InMemoryQueueManager(),
     ) -> None:
         self.agent_executor = agent_executor
         self.task_store = task_store
-        # This works for single binary solution. Needs a distributed approach for
-        # true scalable deployment.
-        self._task_queue: dict[str, EventQueue] = {}
+        self._queue_manager = queue_manager
 
     async def on_get_task(self, params: TaskQueryParams) -> Task | None:
         """Default handler for 'tasks/get'."""
@@ -56,8 +65,10 @@ async def on_cancel_task(self, params: TaskIdParams) -> Task | None:
             initial_message=None,
         )
         result_aggregator = ResultAggregator(task_manager)
-
-        queue = EventQueue()
+        try:
+            queue = await self._queue_manager.tap(task.id)
+        except:
+            queue = EventQueue()
         await self.agent_executor.cancel(
             RequestContext(
                 None,
@@ -95,15 +106,11 @@ async def on_message_send(
         task: Task | None = await task_manager.get_task()
         if task:
             task = task_manager.update_with_message(params.message, task)
+            queue = await self._queue_manager.create_or_tap(task.id)
+        else:
+            queue = EventQueue()
         result_aggregator = ResultAggregator(task_manager)
         # TODO to manage the non-blocking flows.
-
-        queue = EventQueue()
-        # If this is a follow up on an existing task, register the queue now
-        task_id: str | None = task.id if task else None
-        if task_id:
-            self._task_queue[task_id] = queue
-
         producer_task = asyncio.create_task(
             self._run_event_stream(
                 RequestContext(
@@ -127,6 +134,11 @@ async def on_message_send(
             return result
         finally:
             await producer_task
+        if task:
+            try:
+                await self._queue_manager.close(task.id)
+            except NoTaskQueue:
+                pass
 
     async def on_message_send_stream(
         self, params: MessageSendParams
@@ -142,14 +154,11 @@ async def on_message_send_stream(
 
         if task:
             task = task_manager.update_with_message(params.message, task)
-
+            queue = await self._queue_manager.create_or_tap(task.id)
+        else:
+            queue = EventQueue()
         result_aggregator = ResultAggregator(task_manager)
-        queue = EventQueue()
-
         task_id: str | None = task.id if task else None
-        if task_id:
-            self._task_queue[task_id] = queue
-
         producer_task = asyncio.create_task(
             self._run_event_stream(
                 RequestContext(
@@ -165,13 +174,20 @@ async def on_message_send_stream(
             consumer = EventConsumer(queue)
             async for event in result_aggregator.consume_and_emit(consumer):
                 # Now we know we have a Task, register the queue
-                if isinstance(event, Task) and event.id not in self._task_queue:
-                    self._task_queue[event.id] = queue
-                    task_id = event.id
-                yield event
-
+                if isinstance(event, Task):
+                    try:
+                        await self._queue_manager.add(event.id, queue)
+                        task_id = event.id
+                    except TaskQueueExists:
+                        logging.info(
+                            'Multiple Task objects created in event stream.')
+                        yield event
         finally:
             await producer_task
+        try:
+            await self._queue_manager.close(task_id)
+        except NoTaskQueue:
+            pass
 
     async def on_set_task_push_notification_config(
         self, request: TaskPushNotificationConfig
@@ -203,12 +219,11 @@ async def on_resubscribe_to_task(
 
         result_aggregator = ResultAggregator(task_manager)
 
-        # Need to tap the existing queue.
-        if not task.id in self._task_queue:
+        try:
+            queue = await self._queue_manager.tap(task.id)
+        except NoTaskQueue:
             raise ServerError(error=TaskNotFoundError())
-            return
 
-        queue = self._task_queue[task.id].tap()
         consumer = EventConsumer(queue)
         async for event in result_aggregator.consume_and_emit(consumer):
             yield event