fix: Handle propagating agent exceptions (a2aproject#20)

* fix: Fix mypy errors and enable streaming for hw example

* fix:Handle propagating agent exceptions

* fix:Handle propagating agent exceptions

* Add comment for queue read timeout

Author: kthota-g

src/a2a/server/agent_execution/base_agent_executor.py

@@ -22,6 +22,8 @@ class EventConsumer:
 
     def __init__(self, queue: EventQueue):
         self.queue = queue
+        self._timeout = 0.5
+        self._exception: BaseException | None = None
         logger.debug('EventConsumer initialized')
 
     async def consume_one(self) -> Event:
@@ -45,8 +47,15 @@ async def consume_all(self) -> AsyncGenerator[Event]:
         """Consume all the generated streaming events from the agent."""
         logger.debug('Starting to consume all events from the queue.')
         while True:
+            if self._exception:
+                raise self._exception
             try:
-                event = await self.queue.dequeue_event()
+                # We use a timeout when waiting for an event from the queue.
+                # This is required because it allows the loop to check if 
+                # `self._exception` has been set by the `agent_task_callback`. 
+                # Without the timeout, loop might hang indefinitely if no events are
+                # enqueued by the agent and the agent simply threw an exception
+                event = await asyncio.wait_for(self.queue.dequeue_event(), timeout=self._timeout)
                 logger.debug(
                     f'Dequeued event of type: {type(event)} in consume_all.'
                 )
@@ -74,5 +83,16 @@ async def consume_all(self) -> AsyncGenerator[Event]:
                     logger.debug('Stopping event consumption in consume_all.')
                     self.queue.close()
                     break
+            except asyncio.TimeoutError:
+                # continue polling until there is a final event
+                continue
             except asyncio.QueueShutDown:
                 break
+        
+                
+
+
+
+    def agent_task_callback(self, agent_task: asyncio.Task[None]):
+            if agent_task.exception() is not None:
+                self._exception = agent_task.exception()

src/a2a/server/events/event_consumer.py

@@ -138,6 +138,7 @@ async def on_message_send(
         await self._register_producer(task_id, producer_task)
 
         consumer = EventConsumer(queue)
+        producer_task.add_done_callback(consumer.agent_task_callback)
 
         interrupted = False
         try:
@@ -192,6 +193,7 @@ async def on_message_send_stream(
 
         try:
             consumer = EventConsumer(queue)
+            producer_task.add_done_callback(consumer.agent_task_callback)
             async for event in result_aggregator.consume_and_emit(consumer):
                 # Now we know we have a Task, register the queue
                 if isinstance(event, Task):