feat: Implementation of DefaultRequestHandlerV2 (#933)

This pull request introduces a significant refactoring of the agent execution layer, implementing an ActiveTask system and a DefaultRequestHandlerV2 to better manage task lifecycles, concurrency, and event streaming.

Fixes #869

Author: bartek-w

src/a2a/server/agent_execution/active_task.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from a2a.server.agent_execution.agent_executor import AgentExecutor
+    from a2a.server.context import ServerCallContext
+    from a2a.server.tasks.push_notification_sender import PushNotificationSender
+    from a2a.server.tasks.task_store import TaskStore
+
+from a2a.server.agent_execution.active_task import ActiveTask
+from a2a.server.tasks.task_manager import TaskManager
+
+
+logger = logging.getLogger(__name__)
+
+
+class ActiveTaskRegistry:
+    """A registry for active ActiveTask instances."""
+
+    def __init__(
+        self,
+        agent_executor: AgentExecutor,
+        task_store: TaskStore,
+        push_sender: PushNotificationSender | None = None,
+    ):
+        self._agent_executor = agent_executor
+        self._task_store = task_store
+        self._push_sender = push_sender
+        self._active_tasks: dict[str, ActiveTask] = {}
+        self._lock = asyncio.Lock()
+        self._cleanup_tasks: set[asyncio.Task[None]] = set()
+
+    async def get_or_create(
+        self,
+        task_id: str,
+        call_context: ServerCallContext,
+        context_id: str | None = None,
+        create_task_if_missing: bool = False,
+    ) -> ActiveTask:
+        """Retrieves an existing ActiveTask or creates a new one."""
+        async with self._lock:
+            if task_id in self._active_tasks:
+                return self._active_tasks[task_id]
+
+            task_manager = TaskManager(
+                task_id=task_id,
+                context_id=context_id,
+                task_store=self._task_store,
+                initial_message=None,
+                context=call_context,
+            )
+
+            active_task = ActiveTask(
+                agent_executor=self._agent_executor,
+                task_id=task_id,
+                task_manager=task_manager,
+                push_sender=self._push_sender,
+                on_cleanup=self._on_active_task_cleanup,
+            )
+            self._active_tasks[task_id] = active_task
+
+        await active_task.start(
+            call_context=call_context,
+            create_task_if_missing=create_task_if_missing,
+        )
+        return active_task
+
+    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)
+
+    async def _remove_task(self, task_id: str) -> None:
+        async with self._lock:
+            self._active_tasks.pop(task_id, None)
+            logger.debug('Removed active task for %s from registry', task_id)
+
+    async def get(self, task_id: str) -> ActiveTask | None:
+        """Retrieves an existing task."""
+        async with self._lock:
+            return self._active_tasks.get(task_id)

src/a2a/server/agent_execution/active_task_registry.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 
 from a2a.server.agent_execution.context import RequestContext
-from a2a.server.events.event_queue import EventQueue
+from a2a.server.events.event_queue_v2 import EventQueue
 
 
 class AgentExecutor(ABC):
@@ -23,6 +23,18 @@ async def execute(
         return once the agent's execution for this request is complete or
         yields control (e.g., enters an input-required state).
 
+        TODO: Document request lifecycle and AgentExecutor responsibilities:
+        - Should not close the event_queue.
+        - Guarantee single execution per request (no concurrent execution).
+        - Throwing exception will result in TaskState.TASK_STATE_ERROR (CHECK!)
+        - Once call is completed it should not access context or event_queue
+        - Before completing the call it SHOULD update task status to terminal or interrupted state.
+        - Explain AUTH_REQUIRED workflow.
+        - Explain INPUT_REQUIRED workflow.
+        - Explain how cancelation work (executor task will be canceled, cancel() is called, order of calls, etc)
+        - Explain if execute can wait for cancel and if cancel can wait for execute.
+        - Explain behaviour of streaming / not-immediate when execute() returns in active state.
+
         Args:
             context: The request context containing the message, task ID, etc.
             event_queue: The queue to publish events to.
@@ -38,6 +50,10 @@ async def cancel(
         in the context and publish a `TaskStatusUpdateEvent` with state
         `TaskState.TASK_STATE_CANCELED` to the `event_queue`.
 
+        TODO: Document cancelation workflow.
+        - What if TaskState.TASK_STATE_CANCELED is not set by cancel() ?
+        - How it can interact with execute() ?
+
         Args:
             context: The request context containing the task ID to cancel.
             event_queue: The queue to publish the cancellation status update to.

src/a2a/server/agent_execution/agent_executor.py

@@ -120,7 +120,7 @@ def current_task(self) -> Task | None:
         return self._current_task
 
     @current_task.setter
-    def current_task(self, task: Task) -> None:
+    def current_task(self, task: Task | None) -> None:
         """Sets the current task object."""
         self._current_task = task
 

src/a2a/server/agent_execution/context.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

@@ -3,7 +3,10 @@
 import logging
 
 from a2a.server.request_handlers.default_request_handler import (
-    DefaultRequestHandler,
+    LegacyRequestHandler,
+)
+from a2a.server.request_handlers.default_request_handler_v2 import (
+    DefaultRequestHandlerV2,
 )
 from a2a.server.request_handlers.request_handler import (
     RequestHandler,
@@ -40,11 +43,15 @@ def __init__(self, *args, **kwargs):
             ) from _original_error
 
 
+DefaultRequestHandler = DefaultRequestHandlerV2
+
 __all__ = [
     'DefaultGrpcServerCallContextBuilder',
     'DefaultRequestHandler',
+    'DefaultRequestHandlerV2',
     'GrpcHandler',
     'GrpcServerCallContextBuilder',
+    'LegacyRequestHandler',
     'RequestHandler',
     'build_error_response',
     'prepare_response_object',

src/a2a/server/request_handlers/__init__.py

@@ -74,7 +74,7 @@
 
 
 @trace_class(kind=SpanKind.SERVER)
-class DefaultRequestHandler(RequestHandler):
+class LegacyRequestHandler(RequestHandler):
     """Default request handler for all incoming requests.
 
     This handler provides default implementations for all A2A JSON-RPC methods,