fix:Handle propagating agent exceptions

Author: kthota-g

examples/helloworld/agent_executor.py

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

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,10 @@ 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()
+                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 +78,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):