Add extra EventQueue between AgentExecutor and subscribers.

Author: bartek-w

run_test_debug.py

@@ -15,15 +15,15 @@
 
     from a2a.server.agent_execution.agent_executor import AgentExecutor
     from a2a.server.context import ServerCallContext
-    from a2a.server.events.event_queue_v2 import Event, EventQueue
     from a2a.server.tasks.push_notification_sender import (
         PushNotificationSender,
     )
     from a2a.server.tasks.task_manager import TaskManager
 
-
 from a2a.server.events.event_queue_v2 import (
     AsyncQueue,
+    Event,
+    EventQueueSource,
     QueueShutDown,
     _create_async_queue,
 )
@@ -48,6 +48,10 @@
     TaskState.TASK_STATE_FAILED,
     TaskState.TASK_STATE_REJECTED,
 }
+INTERRUPTED_TASK_STATES = {
+    TaskState.TASK_STATE_AUTH_REQUIRED,
+    TaskState.TASK_STATE_INPUT_REQUIRED,
+}
 
 
 class _RequestCompleted:
@@ -76,11 +80,10 @@ class ActiveTask:
       permanently ceased execution and closed its queues.
     """
 
-    def __init__(  # noqa: PLR0913
+    def __init__(
         self,
         agent_executor: AgentExecutor,
         task_id: str,
-        event_queue: EventQueue,
         task_manager: TaskManager,
         push_sender: PushNotificationSender | None = None,
         on_cleanup: Callable[[ActiveTask], None] | None = None,
@@ -90,8 +93,6 @@ def __init__(  # noqa: PLR0913
         Args:
             agent_executor: The executor to run the agent logic (producer).
             task_id: The unique identifier of the task being managed.
-            event_queue: The queue for events produced by the agent. Acts as the pipe
-                         between the producer and consumer tasks.
             task_manager: The manager for task state and database persistence.
             push_sender: Optional sender for out-of-band push notifications.
             on_cleanup: Optional callback triggered when the task is fully finished
@@ -101,7 +102,10 @@ def __init__(  # noqa: PLR0913
         # --- Core Dependencies ---
         self._agent_executor = agent_executor
         self._task_id = task_id
-        self._event_queue = event_queue
+        self._event_queue_agent = EventQueueSource()
+        self._event_queue_subscribers = EventQueueSource(
+            create_default_sink=False
+        )
         self._task_manager = task_manager
         self._push_sender = push_sender
         self._on_cleanup = on_cleanup
@@ -284,7 +288,7 @@ async def _run_producer(self) -> None:
 
                         try:
                             await self._agent_executor.execute(
-                                request_context, self._event_queue
+                                request_context, self._event_queue_agent
                             )
                             logger.debug(
                                 'Producer[%s]: Execution finished successfully',
@@ -301,7 +305,7 @@ async def _run_producer(self) -> None:
                                 self._task_id,
                             )
                             # TODO: Hide from external consumers
-                            await self._event_queue.enqueue_event(
+                            await self._event_queue_agent.enqueue_event(
                                 cast('Event', _RequestCompleted(request_id))
                             )
                             self._request_queue.task_done()
@@ -319,7 +323,8 @@ async def _run_producer(self) -> None:
                         self._exception = e
             finally:
                 self._request_queue.shutdown(immediate=True)
-                await self._event_queue.close(immediate=False)
+                await self._event_queue_agent.close(immediate=False)
+                await self._event_queue_subscribers.close(immediate=False)
         finally:
             logger.debug('Producer[%s]: Completed', self._task_id)
 
@@ -345,7 +350,7 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                             'Consumer[%s]: Waiting for event',
                             self._task_id,
                         )
-                        event = await self._event_queue.dequeue_event()
+                        event = await self._event_queue_agent.dequeue_event()
                         logger.debug(
                             'Consumer[%s]: Dequeued event %s',
                             self._task_id,
@@ -373,15 +378,14 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
 
                                 # Check for AUTH_REQUIRED or INPUT_REQUIRED or TERMINAL states
                                 res = await self._task_manager.get_task()
-                                is_interrupted = res and res.status.state in (
-                                    TaskState.TASK_STATE_AUTH_REQUIRED,
-                                    TaskState.TASK_STATE_INPUT_REQUIRED,
+                                is_interrupted = (
+                                    res
+                                    and res.status.state
+                                    in INTERRUPTED_TASK_STATES
                                 )
-                                is_terminal = res and res.status.state in (
-                                    TaskState.TASK_STATE_COMPLETED,
-                                    TaskState.TASK_STATE_CANCELED,
-                                    TaskState.TASK_STATE_FAILED,
-                                    TaskState.TASK_STATE_REJECTED,
+                                is_terminal = (
+                                    res
+                                    and res.status.state in TERMINAL_TASK_STATES
                                 )
 
                                 # If we hit a breakpoint or terminal state, lock in the result.
@@ -427,9 +431,11 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                     await self._push_sender.send_notification(
                                         self._task_id, event
                                     )
-
                         finally:
-                            self._event_queue.task_done()
+                            await self._event_queue_subscribers.enqueue_event(
+                                event
+                            )
+                            self._event_queue_agent.task_done()
                 except QueueShutDown:
                     logger.debug(
                         'Consumer[%s]: Event queue shut down', self._task_id
@@ -482,7 +488,7 @@ async def subscribe(  # noqa: PLR0912, PLR0915
                 self._reference_count,
             )
 
-        tapped_queue = await self._event_queue.tap()
+        tapped_queue = await self._event_queue_subscribers.tap()
         request_id = await self.enqueue_request(request) if request else None
 
         try:
@@ -562,17 +568,14 @@ async def cancel(self, call_context: ServerCallContext) -> Task | Message:
         )
 
         async with self._lock:
-            # if self._producer_task and not self._producer_task.done():
             if not self._is_finished.is_set() and self._producer_task:
                 logger.debug(
                     'Cancel[%s]: Cancelling producer task', self._task_id
                 )
-                # We do NOT await self._agent_executor.cancel here
-                # because it might take a while and we want to block on `_is_finished.wait()`
                 self._producer_task.cancel()
                 try:
                     await self._agent_executor.cancel(
-                        request_context, self._event_queue
+                        request_context, self._event_queue_agent
                     )
                 except Exception as e:
                     logger.exception(

src/a2a/server/agent_execution/active_task.py

@@ -13,7 +13,6 @@
     from a2a.server.tasks.task_store import TaskStore
 
 from a2a.server.agent_execution.active_task import ActiveTask
-from a2a.server.events.event_queue_v2 import EventQueueSource
 from a2a.server.tasks.task_manager import TaskManager
 
 
@@ -48,7 +47,6 @@ async def get_or_create(
             if task_id in self._active_tasks:
                 return self._active_tasks[task_id]
 
-            event_queue = EventQueueSource()
             task_manager = TaskManager(
                 task_id=task_id,
                 context_id=context_id,
@@ -60,7 +58,6 @@ async def get_or_create(
             active_task = ActiveTask(
                 agent_executor=self._agent_executor,
                 task_id=task_id,
-                event_queue=event_queue,
                 task_manager=task_manager,
                 push_sender=self._push_sender,
                 on_cleanup=self._on_active_task_cleanup,
@@ -75,6 +72,7 @@ async def get_or_create(
 
     def _on_active_task_cleanup(self, active_task: ActiveTask) -> None:
         """Called by ActiveTask when it's finished and has no subscribers."""
+        logger.debug('Active task %s cleanup scheduled', active_task.task_id)
         task = asyncio.create_task(self._remove_task(active_task.task_id))
         self._cleanup_tasks.add(task)
         task.add_done_callback(self._cleanup_tasks.discard)

src/a2a/server/agent_execution/active_task_registry.py

@@ -28,7 +28,11 @@ class EventQueueSource(EventQueue):
     in `_incoming_queue` and distributed to all child Sinks by a background dispatcher task.
     """
 
-    def __init__(self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE) -> None:
+    def __init__(
+        self,
+        max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE,
+        create_default_sink: bool = True,
+    ) -> None:
         """Initializes the EventQueueSource."""
         if max_queue_size <= 0:
             raise ValueError('max_queue_size must be greater than 0')
@@ -41,10 +45,15 @@ def __init__(self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE) -> None:
         self._is_closed = False
 
         # Internal sink for backward compatibility
-        self._default_sink = EventQueueSink(
-            parent=self, max_queue_size=max_queue_size
-        )
-        self._sinks.add(self._default_sink)
+        self._default_sink: EventQueueSink | None
+        if create_default_sink:
+            self._default_sink = EventQueueSink(
+                parent=self, max_queue_size=max_queue_size
+            )
+            self._sinks.add(self._default_sink)
+        else:
+            self._default_sink = None
+
         self._dispatcher_task = asyncio.create_task(self._dispatch_loop())
 
         self._dispatcher_task_expected_to_cancel = False
@@ -54,6 +63,8 @@ def __init__(self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE) -> None:
     @property
     def queue(self) -> AsyncQueue[Event]:
         """Returns the underlying asyncio.Queue of the default sink."""
+        if self._default_sink is None:
+            raise ValueError('No default sink available.')
         return self._default_sink.queue
 
     async def _dispatch_loop(self) -> None:
@@ -183,10 +194,14 @@ async def enqueue_event(self, event: Event) -> None:
 
     async def dequeue_event(self) -> Event:
         """Dequeues an event from the default internal sink queue."""
+        if self._default_sink is None:
+            raise ValueError('No default sink available.')
         return await self._default_sink.dequeue_event()
 
     def task_done(self) -> None:
         """Signals that a formerly enqueued task is complete via the default internal sink queue."""
+        if self._default_sink is None:
+            raise ValueError('No default sink available.')
         self._default_sink.task_done()
 
     async def close(self, immediate: bool = False) -> None:

src/a2a/server/events/event_queue_v2.py

@@ -11,6 +11,10 @@
     RequestContextBuilder,
     SimpleRequestContextBuilder,
 )
+from a2a.server.agent_execution.active_task import (
+    INTERRUPTED_TASK_STATES,
+    TERMINAL_TASK_STATES,
+)
 from a2a.server.agent_execution.active_task_registry import ActiveTaskRegistry
 from a2a.server.request_handlers.request_handler import (
     RequestHandler,
@@ -30,7 +34,6 @@
     SubscribeToTaskRequest,
     Task,
     TaskPushNotificationConfig,
-    TaskState,
     TaskStatusUpdateEvent,
 )
 from a2a.utils.errors import (
@@ -63,12 +66,6 @@
 
 logger = logging.getLogger(__name__)
 
-TERMINAL_TASK_STATES = {
-    TaskState.TASK_STATE_COMPLETED,
-    TaskState.TASK_STATE_CANCELED,
-    TaskState.TASK_STATE_FAILED,
-    TaskState.TASK_STATE_REJECTED,
-}
 
 # TODO: cleanup context_id management
 
@@ -241,14 +238,7 @@ async def on_message_send(  # noqa: D102
             return task
 
         try:
-            result_states = {
-                TaskState.TASK_STATE_COMPLETED,
-                TaskState.TASK_STATE_FAILED,
-                TaskState.TASK_STATE_CANCELED,
-                TaskState.TASK_STATE_REJECTED,
-                TaskState.TASK_STATE_INPUT_REQUIRED,
-                TaskState.TASK_STATE_AUTH_REQUIRED,
-            }
+            result_states = TERMINAL_TASK_STATES | INTERRUPTED_TASK_STATES
 
             result = None
             async for event in active_task.subscribe(request=request_context):

src/a2a/server/request_handlers/default_request_handler_v2.py

@@ -20,7 +20,7 @@
 from a2a.server.request_handlers.default_request_handler import (
     LegacyRequestHandler,
 )
-from a2a.server.routes import CallContextBuilder
+from a2a.server.request_handlers import GrpcServerCallContextBuilder
 from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
 from a2a.types import a2a_pb2_grpc
 from a2a.types.a2a_pb2 import (
@@ -89,7 +89,7 @@ def user_name(self) -> str:
         return 'test-user'
 
 
-class MockCallContextBuilder(CallContextBuilder):
+class MockCallContextBuilder(GrpcServerCallContextBuilder):
     def build(self, request: Any) -> ServerCallContext:
         return ServerCallContext(
             user=MockUser(), state={'headers': {'a2a-version': '1.0'}}
@@ -996,6 +996,7 @@ async def cancel(
             await task1
     else:
         await task1
+
     await task2
 
     # Consume remaining events if any