Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/a2a/server/events/event_consumer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
34 changes: 33 additions & 1 deletion tests/server/events/test_event_consumer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down
Loading