[Bug]: CancelTask racing the task's own completion overwrites COMPLETED with FAILED #1397

What happened?

A task can finish on its own while a CancelTask request for it is in flight. When it does, DefaultRequestHandlerV2 (via ActiveTask.cancel) can overwrite the task's terminal COMPLETED state with FAILED. Reproduced on 1.2.1; the same code is on main (43d1eff) and in 1.3.0.

The window is real. ActiveTask.cancel only checks _is_finished, which the consumer sets once it has processed the terminal event. A complete() that the executor has already enqueued, but the consumer has not yet handled, passes that check. So AgentExecutor.cancel() runs against a task that has already finished.

A completed task is overwritten with FAILED

If executor.cancel() raises, ActiveTask.cancel calls _mark_task_as_failed. That method enqueues TASK_STATE_FAILED without checking whether the task is already terminal. Its sibling _write_failed_status does check.

An executor that shares one TaskUpdater between execute and cancel hits this naturally: when cancel() loses to complete(), TaskUpdater.update_status raises RuntimeError("... already in a terminal state"). The connected client has already received COMPLETED and the result over the stream, yet GetTask afterwards reports FAILED.

The spec defines TASK_STATE_COMPLETED as "a terminal state" (§4.1.3 TaskState; a2a.proto: "Indicates that a task has finished successfully. This is a terminal state."). A streaming client is told it is final, because the stream "MUST close when the task reaches a terminal state" (§3.1.2). After that, the stored state should not change, and in particular should not contradict what was streamed.

Expected: a terminal state is never replaced; a failing cancel() on an already-terminal task leaves it COMPLETED.

Related question: what should CancelTask answer in this race?

Even when cancel() does not raise, ActiveTask.cancel re-reads the task and returns it in whatever terminal state it reached. So CancelTask responds successfully with a TASK_STATE_COMPLETED task. The spec doesn't settle this case clearly. §3.1.5 Cancel Task says:

• "success is not guaranteed (e.g., the task might have already completed or failed…)" and "returns its updated state", which supports the current behaviour;
• Outputs: "Updated Task with cancellation status", and the error TaskNotCancelableError: "The task is not in a cancelable state (e.g., already completed, failed, or canceled)", which supports an error.

The SDK itself answers both ways. The cross-replica path, _cancel_remote, raises TaskNotCancelableError when the reloaded task is not CANCELED; the local path returns the task. Is that difference intended? If not, it would help to pick one.

Repro (self-contained, a2a-sdk 1.2.1, no transport). The executor's cancel() lets the in-flight work finish and complete the task, which deterministically puts the cancel in the window above:

import asyncio

from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.context import ServerCallContext
from a2a.server.events import EventQueue
from a2a.server.request_handlers import DefaultRequestHandlerV2
from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
from a2a.types import (
    AgentCapabilities,
    AgentCard,
    CancelTaskRequest,
    GetTaskRequest,
    Message,
    Part,
    Role,
    SendMessageRequest,
    Task,
    TaskState,
    TaskStatus,
)


class FinishesDuringCancel(AgentExecutor):
    """The task's own work completes it while cancel() is in progress."""

    def __init__(self, cancel_raises: bool) -> None:
        self.cancel_raises = cancel_raises
        self.started = asyncio.Event()
        self.release = asyncio.Event()
        self.completed = asyncio.Event()

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        task_id, context_id = context.task_id, context.context_id
        await event_queue.enqueue_event(
            Task(id=task_id, context_id=context_id, status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED))
        )
        updater = TaskUpdater(event_queue, task_id, context_id)
        await updater.start_work()
        self.started.set()
        await self.release.wait()
        await updater.complete()
        self.completed.set()

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        self.release.set()
        await self.completed.wait()
        if self.cancel_raises:
            # What a shared TaskUpdater raises when cancel() loses the race to complete().
            raise RuntimeError("Task is already in a terminal state.")


async def drain(stream) -> None:
    async for _ in stream:
        pass


async def run(cancel_raises: bool) -> None:
    executor = FinishesDuringCancel(cancel_raises)
    handler = DefaultRequestHandlerV2(
        agent_executor=executor,
        task_store=InMemoryTaskStore(),
        agent_card=AgentCard(name="repro", capabilities=AgentCapabilities(streaming=True)),
    )
    ctx = ServerCallContext()
    message = Message(role=Role.ROLE_USER, parts=[Part(text="go")], message_id="m1")
    stream = handler.on_message_send_stream(SendMessageRequest(message=message), ctx)
    first = await anext(stream)
    task_id = first.id if isinstance(first, Task) else first.task_id
    # Keep reading the stream, as a connected client does.
    draining = asyncio.create_task(drain(stream))
    await executor.started.wait()

    try:
        result = await handler.on_cancel_task(CancelTaskRequest(id=task_id), ctx)
        outcome = f"returned task in {TaskState.Name(result.status.state)}"
    except Exception as e:
        outcome = f"raised {type(e).__name__}: {e}"
    await draining
    stored = await handler.on_get_task(GetTaskRequest(id=task_id), ctx)
    print(f"cancel_raises={cancel_raises}: CancelTask {outcome}; stored state {TaskState.Name(stored.status.state)}")


asyncio.run(run(cancel_raises=False))
asyncio.run(run(cancel_raises=True))

Output (3 of 3 runs):

cancel_raises=False: CancelTask returned task in TASK_STATE_COMPLETED; stored state TASK_STATE_COMPLETED
cancel_raises=True: CancelTask raised RuntimeError: Task is already in a terminal state.; stored state TASK_STATE_FAILED

The second line is the bug: the stored state should stay TASK_STATE_COMPLETED. The first line shows the behaviour the related question is about.

Possible fix: in _mark_task_as_failed (or before calling it from cancel), skip the FAILED write when the task is already terminal, as _write_failed_status does.

We hit this in AG2's A2A server, whose executor shares one TaskUpdater between execute and cancel.

Relevant log output

Cancel[<task id>]: Agent cancel failed
Traceback (most recent call last):
  File ".../a2a/server/agent_execution/active_task.py", line 834, in cancel
    await self._agent_executor.cancel(
  File "repro.py", line 54, in cancel
    raise RuntimeError("Task is already in a terminal state.")
RuntimeError: Task is already in a terminal state.

Code of Conduct

• I agree to follow this project's Code of Conduct