fix: serialize concurrent cancel requests with a per-task lock

.github/actions/spelling/allow.txt

@@ -1,4 +1,4 @@
 a2a
 A2A
 A2AFastAPI
 AAgent
@@ -28,7 +28,7 @@
 AUser
 autouse
 backticks
 base64url
 buf
 bufbuild
 cla
@@ -43,7 +43,7 @@
 drivername
 DSNs
 dunders
 ES256
 euo
 EUR
 evt
@@ -58,8 +58,8 @@
 gle
 GVsb
 hazmat
 HS256
 HS384
 ietf
 importlib
 initdb
@@ -100,10 +100,10 @@
 Oneof
 OpenAPI
 openapiv
 openapiv2
 opensource
 otherurl
 pb2
 podman
 Podman
 poolclass
@@ -126,7 +126,7 @@
 respx
 resub
 rmi
 RS256
 RUF
 SECP256R1
 SFIXED
@@ -151,3 +151,4 @@
 UIDs
 vulnz
 whl
+TOCTOU

src/a2a/server/request_handlers/default_request_handler.py

@@ -2,6 +2,7 @@
 import logging
 
 from collections.abc import AsyncGenerator, Awaitable, Callable
+from dataclasses import dataclass
 from typing import cast
 
 from a2a.server.agent_execution import (
@@ -76,6 +77,14 @@
 }
 
 
+@dataclass
+class _CancelLockEntry:
+    """Per-task cancel lock plus its in-flight user count."""
+
+    lock: asyncio.Lock
+    refs: int = 0
+
+
 @trace_class(kind=SpanKind.SERVER)
 class LegacyRequestHandler(RequestHandler):
     """Default request handler for all incoming requests.
@@ -134,6 +143,12 @@ def __init__(  # noqa: PLR0913
         # TODO: Likely want an interface for managing this, like AgentExecutionManager.
         self._running_agents = {}
         self._running_agents_lock = asyncio.Lock()
+        # Per-task cancellation locks. Each value is [lock, user_count].
+        # user_count tracks how many concurrent on_cancel_task calls are
+        # using the lock so entries can be garbage-collected safely once
+        # the last caller finishes. Both are protected by _cancel_locks_guard.
+        self._cancel_locks: dict[str, _CancelLockEntry] = {}
+        self._cancel_locks_guard = asyncio.Lock()
         # Tracks background tasks (e.g., deferred cleanups) to avoid orphaning
         # asyncio tasks and to surface unexpected exceptions.
         self._background_tasks = set()
@@ -185,8 +200,40 @@ async def on_cancel_task(
         """Default handler for 'tasks/cancel'.
 
         Attempts to cancel the task managed by the `AgentExecutor`.
+
+        The terminal-state check and the cancel call are serialized with a
+        per-task lock so concurrent cancels cannot both pass the check
+        (TOCTOU, BUG-44). V2's ``ActiveTaskRegistry`` already serializes
+        cancels per ``ActiveTask`` and is not affected.
         """
         task_id = params.id
+
+        # Acquire (or create) the per-task lock and bump its user count.
+        # The guard section is synchronous (no await), so the count is
+        # updated atomically with respect to other coroutines.
+        async with self._cancel_locks_guard:
+            entry = self._cancel_locks.get(task_id)
+            if entry is None:
+                entry = _CancelLockEntry(lock=asyncio.Lock())
+                self._cancel_locks[task_id] = entry
+            entry.refs += 1
+
+        try:
+            async with entry.lock:
+                return await self._cancel_task_locked(task_id, context)
+        finally:
+            async with self._cancel_locks_guard:
+                entry.refs -= 1
+                # Only drop the entry when no other caller is waiting on it;
+                # otherwise a new caller could grab a fresh lock and bypass
+                # the serialization.
+                if entry.refs == 0 and self._cancel_locks.get(task_id) is entry:
+                    del self._cancel_locks[task_id]
+
+    async def _cancel_task_locked(
+        self, task_id: str, context: ServerCallContext
+    ) -> Task | None:
+        """Cancels a task; callers must hold the per-task cancel lock."""
         task: Task | None = await self.task_store.get(task_id, context)
         if not task:
             raise TaskNotFoundError

tests/server/request_handlers/test_default_request_handler.py

@@ -3143,3 +3143,79 @@ async def test_on_get_task_push_notification_config_is_owner_scoped(
             ),
             _ctx('bob'),
         )
+
+
+@pytest.mark.asyncio
+async def test_on_cancel_task_serializes_concurrent_cancels(agent_card):
+    """Concurrent cancels of the same task must be serialized (BUG-44).
+
+    The second cancel's terminal-state check must observe the outcome of the
+    first cancel instead of racing past it.
+    """
+    task_id = 'toctou_task'
+    mock_task_store = AsyncMock(spec=TaskStore)
+    mock_task_store.get.return_value = create_sample_task(task_id=task_id)
+
+    mock_queue_manager = AsyncMock(spec=QueueManager)
+    mock_event_queue = AsyncMock(spec=EventQueueLegacy)
+    mock_queue_manager.tap.return_value = mock_event_queue
+
+    # The first cancel blocks inside agent_executor.cancel until released;
+    # once released, the task is considered canceled in the store.
+    release_first = asyncio.Event()
+    calls = 0
+    mock_agent_executor = AsyncMock(spec=AgentExecutor)
+
+    async def blocking_cancel(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            await release_first.wait()
+        mock_task_store.get.return_value = create_sample_task(
+            task_id=task_id, status_state=TaskState.TASK_STATE_CANCELED
+        )
+
+    mock_agent_executor.cancel.side_effect = blocking_cancel
+
+    mock_result_aggregator_instance = AsyncMock(spec=ResultAggregator)
+    mock_result_aggregator_instance.consume_all.return_value = (
+        create_sample_task(
+            task_id=task_id, status_state=TaskState.TASK_STATE_CANCELED
+        )
+    )
+
+    request_handler = DefaultRequestHandler(
+        agent_executor=mock_agent_executor,
+        task_store=mock_task_store,
+        queue_manager=mock_queue_manager,
+        agent_card=agent_card,
+    )
+
+    context = create_server_call_context()
+    params = CancelTaskRequest(id=task_id)
+
+    with patch(
+        'a2a.server.request_handlers.default_request_handler.ResultAggregator',
+        return_value=mock_result_aggregator_instance,
+    ):
+        first = asyncio.create_task(
+            request_handler.on_cancel_task(params, context)
+        )
+        await asyncio.sleep(0)  # let the first cancel reach the executor
+        second = asyncio.create_task(
+            request_handler.on_cancel_task(params, context)
+        )
+        await asyncio.sleep(0)  # let the second cancel block on the lock
+        release_first.set()
+        first_result = await first
+
+        # The second cancel re-checks state under the per-task lock and must
+        # observe the terminal state left by the first cancel.
+        with pytest.raises(TaskNotCancelableError):
+            await second
+
+    assert first_result.status.state == TaskState.TASK_STATE_CANCELED
+    # The second cancel must never reach the executor.
+    assert calls == 1
+    # The per-task lock entry is cleaned up once no caller is using it.
+    assert task_id not in request_handler._cancel_locks