From 783c4668e7bc6a4566e11c5f27ba0af7a2605d95 Mon Sep 17 00:00:00 2001
From: Brian Fox <878612+onematchfox@users.noreply.github.com>
Date: Tue, 5 Aug 2025 15:09:36 +0200
Subject: [PATCH] fix: gracefully handle task exceptions in event consumer

---
 src/a2a/server/events/event_consumer.py    |  2 +-
 tests/server/events/test_event_consumer.py | 34 +++++++++++++++++++++-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/a2a/server/events/event_consumer.py b/src/a2a/server/events/event_consumer.py
index 6745847c..607c3c9c 100644
--- a/src/a2a/server/events/event_consumer.py
+++ b/src/a2a/server/events/event_consumer.py
@@ -160,5 +160,5 @@ def agent_task_callback(self, agent_task: asyncio.Task[None]) -> None:
             agent_task: The asyncio.Task that completed.
         """
         logger.debug('Agent task callback triggered.')
-        if agent_task.exception() is not None:
+        if not agent_task.cancelled() and agent_task.done():
             self._exception = agent_task.exception()
diff --git a/tests/server/events/test_event_consumer.py b/tests/server/events/test_event_consumer.py
index 9f061a8b..4116fabd 100644
--- a/tests/server/events/test_event_consumer.py
+++ b/tests/server/events/test_event_consumer.py
@@ -327,18 +327,22 @@ async def test_consume_all_continues_on_queue_empty_if_not_really_closed(
 def test_agent_task_callback_sets_exception(event_consumer: EventConsumer):
     """Test that agent_task_callback sets _exception if the task had one."""
     mock_task = MagicMock(spec=asyncio.Task)
+    mock_task.cancelled.return_value = False
+    mock_task.done.return_value = True
     sample_exception = ValueError('Task failed')
     mock_task.exception.return_value = sample_exception
 
     event_consumer.agent_task_callback(mock_task)
 
     assert event_consumer._exception == sample_exception
-    # mock_task.exception.assert_called_once() # Removing this, as exception() might be called internally by the check
+    mock_task.exception.assert_called_once()
 
 
 def test_agent_task_callback_no_exception(event_consumer: EventConsumer):
     """Test that agent_task_callback does nothing if the task has no exception."""
     mock_task = MagicMock(spec=asyncio.Task)
+    mock_task.cancelled.return_value = False
+    mock_task.done.return_value = True
     mock_task.exception.return_value = None  # No exception
 
     event_consumer.agent_task_callback(mock_task)
@@ -347,6 +351,34 @@ def test_agent_task_callback_no_exception(event_consumer: EventConsumer):
     mock_task.exception.assert_called_once()
 
 
+def test_agent_task_callback_cancelled_task(event_consumer: EventConsumer):
+    """Test that agent_task_callback does nothing if the task has no exception."""
+    mock_task = MagicMock(spec=asyncio.Task)
+    mock_task.cancelled.return_value = True
+    mock_task.done.return_value = True
+    sample_exception = ValueError('Task still running')
+    mock_task.exception.return_value = sample_exception
+
+    event_consumer.agent_task_callback(mock_task)
+
+    assert event_consumer._exception is None  # Should remain None
+    mock_task.exception.assert_not_called()
+
+
+def test_agent_task_callback_not_done_task(event_consumer: EventConsumer):
+    """Test that agent_task_callback does nothing if the task has no exception."""
+    mock_task = MagicMock(spec=asyncio.Task)
+    mock_task.cancelled.return_value = False
+    mock_task.done.return_value = False
+    sample_exception = ValueError('Task is cancelled')
+    mock_task.exception.return_value = sample_exception
+
+    event_consumer.agent_task_callback(mock_task)
+
+    assert event_consumer._exception is None  # Should remain None
+    mock_task.exception.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_consume_all_handles_validation_error(
     event_consumer: EventConsumer, mock_event_queue: AsyncMock