Fixes for types and random issues

Change-Id: I5acd71e8903825bd31b9d483471b4a3cf0f7e340

Author: kthota-g

examples/helloworld/agent_executor.py

@@ -24,14 +24,14 @@ def __init__(self):
     @override
     async def execute(
         self,
-        request: RequestContext,
+        context: RequestContext,
         event_queue: EventQueue,
     ) -> None:
         result = await self.agent.invoke()
         event_queue.enqueue_event(new_agent_text_message(result))
 
     @override
     async def cancel(
-        self, request: RequestContext, event_queue: EventQueue
-    ) -> Task | None:
+        self, context: RequestContext, event_queue: EventQueue
+    ) -> None:
         raise Exception('cancel not supported')

examples/langgraph/agent_executor.py

@@ -4,7 +4,6 @@
 from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.events.event_queue import EventQueue
 from a2a.types import (
-    Task,
     TaskArtifactUpdateEvent,
     TaskState,
     TaskStatus,
@@ -93,6 +92,6 @@ async def execute(
 
     @override
     async def cancel(
-        self, request: RequestContext, event_queue: EventQueue
-    ) -> Task | None:
+        self, context: RequestContext, event_queue: EventQueue
+    ) -> None:
         raise Exception('cancel not supported')

src/a2a/server/events/event_queue.py

@@ -68,4 +68,4 @@ def close(self):
         """Closes the queue for future push events."""
         self.queue.shutdown()
         for child in self._children:
-            child.shutdown()
+            child.close()

src/a2a/server/request_handlers/default_request_handler.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import logging
 
 from collections.abc import AsyncGenerator
@@ -39,11 +40,11 @@ def __init__(
         self,
         agent_executor: AgentExecutor,
         task_store: TaskStore,
-        queue_manager: QueueManager = InMemoryQueueManager(),
+        queue_manager: QueueManager | None = None,
     ) -> None:
         self.agent_executor = agent_executor
         self.task_store = task_store
-        self._queue_manager = queue_manager
+        self._queue_manager = queue_manager or InMemoryQueueManager()
 
     async def on_get_task(self, params: TaskQueryParams) -> Task | None:
         """Default handler for 'tasks/get'."""
@@ -65,10 +66,11 @@ async def on_cancel_task(self, params: TaskIdParams) -> Task | None:
             initial_message=None,
         )
         result_aggregator = ResultAggregator(task_manager)
-        try:
-            queue = await self._queue_manager.tap(task.id)
-        except:
+
+        queue = await self._queue_manager.tap(task.id)
+        if not queue:
             queue = EventQueue()
+
         await self.agent_executor.cancel(
             RequestContext(
                 None,
@@ -134,11 +136,9 @@ async def on_message_send(
             return result
         finally:
             await producer_task
-        if task:
-            try:
-                await self._queue_manager.close(task.id)
-            except NoTaskQueue:
-                pass
+            if task:
+                with contextlib.suppress(NoTaskQueue):
+                    await self._queue_manager.close(task.id)
 
     async def on_message_send_stream(
         self, params: MessageSendParams
@@ -182,13 +182,12 @@ async def on_message_send_stream(
                         logging.info(
                             'Multiple Task objects created in event stream.'
                         )
-                        yield event
+                yield event
         finally:
             await producer_task
-        try:
-            await self._queue_manager.close(task_id)
-        except NoTaskQueue:
-            pass
+            if task_id:
+                with contextlib.suppress(NoTaskQueue):
+                    await self._queue_manager.close(task_id)
 
     async def on_set_task_push_notification_config(
         self, params: TaskPushNotificationConfig
@@ -209,7 +208,6 @@ async def on_resubscribe_to_task(
         task: Task | None = await self.task_store.get(params.id)
         if not task:
             raise ServerError(error=TaskNotFoundError())
-            return
 
         task_manager = TaskManager(
             task_id=task.id,
@@ -220,10 +218,9 @@ async def on_resubscribe_to_task(
 
         result_aggregator = ResultAggregator(task_manager)
 
-        try:
-            queue = await self._queue_manager.tap(task.id)
-        except NoTaskQueue as e:
-            raise ServerError(error=TaskNotFoundError()) from e
+        queue = await self._queue_manager.tap(task.id)
+        if not queue:
+            raise ServerError(error=TaskNotFoundError())
 
         consumer = EventConsumer(queue)
         async for event in result_aggregator.consume_and_emit(consumer):

tests/server/request_handlers/test_jsonrpc_handler.py

@@ -8,6 +8,9 @@
     DefaultRequestHandler,
     JSONRPCHandler,
 )
+from a2a.server.events import (
+    QueueManager,
+)
 from a2a.server.tasks import TaskStore
 from a2a.types import (
     GetTaskRequest,
@@ -375,13 +378,12 @@ async def test_on_resubscribe_existing_task_success(
     ) -> None:
         mock_agent_executor = AsyncMock(spec=AgentExecutor)
         mock_task_store = AsyncMock(spec=TaskStore)
-        mock_queue = AsyncMock(spec=EventQueue)
+        mock_queue_manager = AsyncMock(spec=QueueManager)
         request_handler = DefaultRequestHandler(
-            mock_agent_executor, mock_task_store
+            mock_agent_executor, mock_task_store, mock_queue_manager
         )
         handler = JSONRPCHandler(None, request_handler)
         mock_task = Task(**MINIMAL_TASK, history=[])
-        request_handler._task_queue[mock_task.id] = mock_queue
         events: list[Any] = [
             TaskArtifactUpdateEvent(
                 taskId='task_123',
@@ -407,7 +409,7 @@ async def streaming_coro():
             return_value=streaming_coro(),
         ):
             mock_task_store.get.return_value = mock_task
-            mock_queue.tap.return_value = EventQueue()
+            mock_queue_manager.tap.return_value = EventQueue()
             request = TaskResubscriptionRequest(
                 id='1', params=TaskIdParams(id=mock_task.id)
             )