DefaultRequestHandlerV2

Author: bartek-w

src/a2a/server/agent_execution/active_task.py

@@ -0,0 +1,90 @@
+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.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.context import ServerCallContext
+from a2a.server.events.event_queue_v2 import EventQueueSource
+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,
+        context_id: str | None = None,
+        call_context: ServerCallContext | 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]
+
+            event_queue = EventQueueSource()
+            task_manager = TaskManager(
+                task_id=task_id,
+                context_id=context_id,
+                task_store=self._task_store,
+                initial_message=None,
+                context=ServerCallContext(),
+            )
+
+            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,
+            )
+            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."""
+        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,19 @@ 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:
+        - 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-immediat when task is completed in active state.
+
+
         Args:
             context: The request context containing the message, task ID, etc.
             event_queue: The queue to publish events to.
@@ -38,6 +51,9 @@ async def cancel(
         in the context and publish a `TaskStatusUpdateEvent` with state
         `TaskState.TASK_STATE_CANCELED` to the `event_queue`.
 
+        TODO:
+        - Should not close the event_queue.
+
         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

@@ -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,
@@ -38,9 +41,13 @@ def __init__(self, *args, **kwargs):
             ) from _original_error
 
 
+DefaultRequestHandler = DefaultRequestHandlerV2
+
 __all__ = [
     'DefaultRequestHandler',
+    'DefaultRequestHandlerV2',
     'GrpcHandler',
+    '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,