From f4b7349068a410da79b9ee805c5845e28762e187 Mon Sep 17 00:00:00 2001 From: meraklbz Date: Mon, 10 Aug 2026 20:14:17 +0800 Subject: [PATCH 1/2] fix: serialize concurrent cancel requests with a per-task lock --- .../default_request_handler.py | 38 ++++++++++ .../test_default_request_handler.py | 76 +++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py index ef61dcca7..52506fde9 100644 --- a/src/a2a/server/request_handlers/default_request_handler.py +++ b/src/a2a/server/request_handlers/default_request_handler.py @@ -134,6 +134,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, list[asyncio.Lock | int]] = {} + 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 +191,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 = [asyncio.Lock(), 0] + self._cancel_locks[task_id] = entry + entry[1] += 1 + + try: + async with entry[0]: + return await self._cancel_task_locked(task_id, context) + finally: + async with self._cancel_locks_guard: + entry[1] -= 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[1] == 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 diff --git a/tests/server/request_handlers/test_default_request_handler.py b/tests/server/request_handlers/test_default_request_handler.py index 727679e7c..7f5f43ee1 100644 --- a/tests/server/request_handlers/test_default_request_handler.py +++ b/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 From ff349094b9c2bd65bb123e611410816d2f5ba098 Mon Sep 17 00:00:00 2001 From: meraklbz Date: Tue, 11 Aug 2026 00:39:21 +0800 Subject: [PATCH 2/2] fix: type the per-task cancel lock entry as a dataclass The ref-counted per-task cancel lock was typed as a heterogeneous list (`list[asyncio.Lock | int]`), which the type checker cannot narrow by index (entry[0] vs entry[1]). Replace it with a small dataclass carrying the lock and the in-flight user count, and add the TOCTOU acronym to the spellcheck allow list. --- .github/actions/spelling/allow.txt | 1 + .../default_request_handler.py | 21 +++++++++++++------ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/.github/actions/spelling/allow.txt b/.github/actions/spelling/allow.txt index 1701f14ef..659a5aa0c 100644 --- a/.github/actions/spelling/allow.txt +++ b/.github/actions/spelling/allow.txt @@ -151,3 +151,4 @@ typeerror UIDs vulnz whl +TOCTOU diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py index 52506fde9..71a838aed 100644 --- a/src/a2a/server/request_handlers/default_request_handler.py +++ b/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. @@ -138,7 +147,7 @@ def __init__( # noqa: PLR0913 # 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, list[asyncio.Lock | int]] = {} + 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. @@ -205,20 +214,20 @@ async def on_cancel_task( async with self._cancel_locks_guard: entry = self._cancel_locks.get(task_id) if entry is None: - entry = [asyncio.Lock(), 0] + entry = _CancelLockEntry(lock=asyncio.Lock()) self._cancel_locks[task_id] = entry - entry[1] += 1 + entry.refs += 1 try: - async with entry[0]: + async with entry.lock: return await self._cancel_task_locked(task_id, context) finally: async with self._cancel_locks_guard: - entry[1] -= 1 + 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[1] == 0 and self._cancel_locks.get(task_id) is entry: + if entry.refs == 0 and self._cancel_locks.get(task_id) is entry: del self._cancel_locks[task_id] async def _cancel_task_locked(