Skip to content
Open
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
21 changes: 20 additions & 1 deletion src/a2a/server/agent_execution/active_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,22 @@ async def _process_event_inner(self, event: Event) -> None:
| PushNotificationEvent
| Task,
):
if isinstance(
event, TaskStatusUpdateEvent | TaskArtifactUpdateEvent
):
current = await self.active_task._task_manager.get_task()
if (
current is not None
and current.status.state in TERMINAL_TASK_STATES
):
# A terminal task's status is final. New status updates should be dropped
logger.debug(
'Consumer[%s]: dropping %s for an already-terminal task',
self.active_task._task_id,
type(event).__name__,
)
self.active_task._event_queue_agent.task_done()
return
updated_task = await self._handle_task_event(event)
handled_event = updated_task if isinstance(event, Task) else event

Expand Down Expand Up @@ -1079,7 +1095,10 @@ async def _mark_task_as_failed(self, exception: Exception) -> Task | None:
if self._task_created.is_set():
try:
task = await self._task_manager.get_task()
if task is not None:
if (
task is not None
and task.status.state not in TERMINAL_TASK_STATES
):
await self._event_queue_agent.enqueue_event(
TaskStatusUpdateEvent(
task_id=task.id,
Expand Down
111 changes: 110 additions & 1 deletion tests/server/agent_execution/test_active_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import pytest
import pytest_asyncio

from a2a.server.agent_execution.active_task import ActiveTask
from a2a.server.agent_execution.active_task import ActiveTask, EventConsumer
from a2a.server.agent_execution.agent_executor import AgentExecutor
from a2a.server.agent_execution.context import RequestContext
from a2a.server.context import ServerCallContext
Expand Down Expand Up @@ -234,6 +234,115 @@ async def execute_mock(req, q):
assert shared.status.state == TaskState.TASK_STATE_WORKING
stop_event.set()

@pytest.mark.asyncio
@pytest.mark.parametrize(
('final_state', 'cancel_error_message', 'expected_failed_count'),
[
(TaskState.TASK_STATE_COMPLETED, 'cancel lost the race', 0),
(TaskState.TASK_STATE_WORKING, 'agent cancel blew up', 1),
],
ids=['terminal-task-not-overwritten', 'active-task-still-marks-failed'],
)
async def test_active_task_cancel_exception_marks_failed_only_when_not_terminal(
self,
active_task: ActiveTask,
agent_executor: Mock,
request_context: Mock,
task_manager: Mock,
final_state: TaskState,
cancel_error_message: str,
expected_failed_count: int,
) -> None:
"""Regression test: FAILED is enqueued only if the task was not already terminal when cancel() failed."""
stop_event = asyncio.Event()

async def execute_mock(req: RequestContext, q: EventQueue) -> None:
await stop_event.wait()

get_task_calls = 0

async def get_task_side_effect() -> Task | None:
nonlocal get_task_calls
get_task_calls += 1
# start()'s existence check needs a non-terminal task so it sets _task_created.
if get_task_calls == 1:
return Task(
id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
return Task(id='test-task-id', status=TaskStatus(state=final_state))

agent_executor.execute = AsyncMock(side_effect=execute_mock)
agent_executor.cancel = AsyncMock(
side_effect=RuntimeError(cancel_error_message)
)
task_manager.get_task = AsyncMock(side_effect=get_task_side_effect)

await active_task.enqueue_request(request_context)
await active_task.start(
call_context=ServerCallContext(), create_task_if_missing=True
)
await asyncio.sleep(0.1)
# Precondition: the guard under test only applies once this is set.
assert active_task._task_created.is_set()

with patch.object(
active_task._event_queue_agent,
'enqueue_event',
wraps=active_task._event_queue_agent.enqueue_event,
) as enqueue_spy:
with pytest.raises(RuntimeError, match=cancel_error_message):
await active_task.cancel(request_context)

failed_writes = [
call.args[0]
for call in enqueue_spy.await_args_list
if isinstance(call.args[0], TaskStatusUpdateEvent)
and call.args[0].status.state == TaskState.TASK_STATE_FAILED
]
assert len(failed_writes) == expected_failed_count
assert all(write.task_id == 'test-task-id' for write in failed_writes)

stop_event.set()
await active_task.aclose()

@pytest.mark.asyncio
async def test_event_consumer_drops_status_update_for_terminal_task(
self,
active_task: ActiveTask,
task_manager: Mock,
) -> None:
"""Regression test: FAILED update enqueued before the task turned terminal must be dropped at apply-time too."""
task_manager.get_task = AsyncMock(
return_value=Task(
id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
)
)

stale_failed_event = TaskStatusUpdateEvent(
task_id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_FAILED),
)
await active_task._event_queue_agent.enqueue_event(stale_failed_event)
# Round-trip through the queue so task_done() below has a matching get().
dequeued = await active_task._event_queue_agent.dequeue_event()

# Not started: exercising the consumer's guard directly needs no
# producer/consumer loop.
consumer = EventConsumer(active_task)
with patch.object(
active_task._event_queue_subscribers,
'enqueue_event',
wraps=active_task._event_queue_subscribers.enqueue_event,
) as subscriber_spy:
await consumer._process_event(dequeued)

# Dropped event must reach neither TaskManager nor subscribers.
task_manager.ensure_task_id.assert_not_awaited()
task_manager.process.assert_not_awaited()
subscriber_spy.assert_not_awaited()

@pytest.mark.asyncio
async def test_active_task_interrupted_auth(
self,
Expand Down
Loading