Type fixes and message stream fixes

Change-Id: I6b16e04e3ba0e9c01ba2268ab144e943017a8f2c

Author: kthota-g

examples/langgraph/agent_executor.py

@@ -1,19 +1,16 @@
-from typing import Any
-
 from agent import CurrencyAgent
 from typing_extensions import override
 
 from a2a.server.agent_execution import AgentExecutor, RequestContext
-from a2a.server.events.event_queue import EventQueue, Event
-from a2a.utils import new_agent_text_message, new_text_artifact, new_task
+from a2a.server.events.event_queue import EventQueue
 from a2a.types import (
     Task,
-    TextPart,
-    TaskStatusUpdateEvent,
     TaskArtifactUpdateEvent,
-    TaskStatus,
     TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
 )
+from a2a.utils import new_agent_text_message, new_task, new_text_artifact
 
 
 class CurrencyAgentExecutor(AgentExecutor):
@@ -31,6 +28,9 @@ async def execute(
         query = context.get_user_input()
         task = context.current_task
 
+        if not context.message:
+            raise Exception('No message provided')
+
         if not task:
             task = new_task(context.message)
             event_queue.enqueue_event(task)
@@ -92,6 +92,7 @@ async def execute(
                 )
 
     @override
-    async def cancel(self, request: RequestContext, event_queue: EventQueue) -> Task | None:
-        raise Exception("cancel not supported")
-
+    async def cancel(
+        self, request: RequestContext, event_queue: EventQueue
+    ) -> Task | None:
+        raise Exception('cancel not supported')

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.13"
 dependencies = [
     "httpx>=0.28.1",
     "httpx-sse>=0.4.0",
+    "mypy>=1.15.0",
     "pydantic>=2.11.3",
     "sse-starlette>=2.3.3",
     "starlette>=0.46.2",

src/a2a/server/events/event_consumer.py

@@ -29,11 +29,11 @@ async def consume_one(self) -> Event:
         logger.debug('Attempting to consume one event.')
         try:
             event = await self.queue.dequeue_event(no_wait=True)
-        except asyncio.QueueEmpty:
+        except asyncio.QueueEmpty as e:
             logger.warning('Event queue was empty in consume_one.')
             raise ServerError(
                 InternalError(message='Agent did not return any response')
-            ) from None
+            ) from e
 
         logger.debug(f'Dequeued event of type: {type(event)} in consume_one.')
 

src/a2a/server/tasks/result_aggregator.py

@@ -28,22 +28,19 @@ class ResultAggregator:
 
     def __init__(self, task_manager: TaskManager):
         self.task_manager = task_manager
-        self._message = None
+        self._message: Message | None = None
 
     @property
-    def current_result(self) -> Task | Message | None:
+    async def current_result(self) -> Task | Message | None:
         if self._message:
             return self._message
-        return self.task_manager.get_task()
+        return await self.task_manager.get_task()
 
     async def consume_and_emit(
         self, consumer: EventConsumer
     ) -> AsyncGenerator[Event, None]:
         """Processes the event stream and emits the same event stream out."""
         async for event in consumer.consume_all():
-            if isinstance(event, Message):
-                self._current_task_or_message = event
-                break
             await self.task_manager.process(event)
             yield event
 
@@ -53,17 +50,18 @@ async def consume_all(
         """Processes the entire event stream and returns the final result."""
         async for event in consumer.consume_all():
             if isinstance(event, Message):
-                self._current_task_or_message = event
+                self._message = event
                 return event
             await self.task_manager.process(event)
         return await self.task_manager.get_task()
 
-    async def consume_and_emit_task(
-        self, consumer: EventConsumer
-    ) -> AsyncGenerator[Event, None]:
-        """Processes the event stream and emits the current state of the task."""
-        async for event in consumer.consume_all():
-            if isinstance(event, Message):
-                self._current_task_or_message = event
-                break
-            yield await self.task_manager.process(event)
+    # async def consume_and_emit_task(
+    #     self, consumer: EventConsumer
+    # ) -> AsyncGenerator[Event, None]:
+    #     """Processes the event stream and emits the current state of the task."""
+    #     async for event in consumer.consume_all():
+    #         if isinstance(event, Message):
+    #             self._current_task_or_message = event
+    #             break
+    #         await self.task_manager.process(event)
+    #     return await self.task_manager.get_task()

src/a2a/server/tasks/task_manager.py

@@ -127,14 +127,17 @@ async def ensure_task(
 
         return task
 
-    async def process(self, event: Event) -> Task | Message:
+    async def process(self, event: Event) -> Event:
         """Processes an event, store the task state and return the task or message.
 
         The returned Task or Message represent the current status of the result.
         """
-        if isinstance(event, Message):
-            return Message
-        return await self.save_task_event(event)
+        if isinstance(
+            event, Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
+        ):
+            await self.save_task_event(event)
+
+        return event
 
     def _init_task_obj(self, task_id: str, context_id: str) -> Task:
         """Initializes a new task object."""