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 ef61dcca7..b680423a5 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. @@ -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). 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 diff --git a/tests/server/request_handlers/test_default_request_handler.py b/tests/server/request_handlers/test_default_request_handler.py index 727679e7c..2fca9fb41 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. + + 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