fix: owner-scope cancel/subscribe and write terminal state on cancel (#1159, #1170)

src/a2a/server/agent_execution/active_task.py

@@ -707,6 +707,12 @@ async def subscribe(
     async def cancel(self, call_context: ServerCallContext) -> Task:
         """Cancels the running active task.
 
+        The returned task carries a terminal state, which is written to the
+        task store. That write is not guaranteed to reach an active subscriber
+        stream: by the time cancel writes it the event queues may already be
+        closed, so a client that was streaming may have to re-read the task to
+        observe the terminal state.
+
         Concurrency Guarantee:
         Uses `_lock` to ensure we don't attempt to cancel a producer that is
         already winding down or hasn't started. It fires the cancellation signal
@@ -730,7 +736,17 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
                 logger.debug(
                     'Cancel[%s]: Cancelling producer task', self._task_id
                 )
-                self._producer_task.cancel()
+                # Await the executor's cancel before cancelling the producer,
+                # so the component that owns the task's terminal state can
+                # still write to the still-open event queue. Cancelling the
+                # producer first can tear the queue down first and drop that
+                # write. Mirrors the V1 handler ordering. The producer cancel
+                # is in a finally so it still runs when executor.cancel() raises
+                # a BaseException such as asyncio.CancelledError, which
+                # `except Exception` does not catch; otherwise the producer
+                # would leak as a pending task. Cancelling after
+                # _mark_task_as_failed also keeps the FAILED write ahead of the
+                # queue teardown, so it is not dropped as QueueShutDown.
                 try:
                     await self._agent_executor.cancel(
                         request_context, self._event_queue_agent
@@ -741,6 +757,8 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
                     )
                     await self._mark_task_as_failed(e)
                     raise
+                finally:
+                    self._producer_task.cancel()
             else:
                 logger.debug(
                     'Cancel[%s]: Task already finished [%s] or producer not started [%s], not cancelling',
@@ -753,6 +771,22 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
         task = await self._task_manager.get_task()
         if not task:
             raise RuntimeError('Task should have been created')
+        # A cleanup-only executor.cancel() may not write a terminal state, and
+        # a task parked in a non-terminal state (e.g. input-required) has no
+        # running producer to write one either. Close it out as CANCELED so a
+        # caller that cancelled is never left polling a live task. Mirrors the
+        # V1 handler, which made a non-cancelled outcome visible instead of
+        # reporting success and changing nothing.
+        if task.status.state not in TERMINAL_TASK_STATES:
+            # Write a copy rather than mutating the task in place: get_task()
+            # returns the shared _task_manager._current_task, which may already
+            # have been yielded to a subscriber, and that reference must not
+            # change under the reader.
+            updated = Task()
+            updated.CopyFrom(task)
+            updated.status.state = TaskState.TASK_STATE_CANCELED
+            await self._task_manager.save_task_event(updated)
+            task = updated
         return task
 
     async def aclose(self) -> None:

src/a2a/server/agent_execution/active_task_registry.py

@@ -15,6 +15,7 @@
 
 from a2a.server.agent_execution.active_task import ActiveTask
 from a2a.server.tasks.task_manager import TaskManager
+from a2a.utils.errors import TaskNotFoundError
 
 
 logger = logging.getLogger(__name__)
@@ -49,25 +50,44 @@ async def get_or_create(
         async with self._lock:
             if self._closed:
                 raise RuntimeError('ActiveTaskRegistry is closed')
-            if task_id in self._active_tasks:
-                return self._active_tasks[task_id]
-
-            task_manager = TaskManager(
-                task_id=task_id,
-                context_id=context_id,
-                task_store=self._task_store,
-                initial_message=initial_message,
-                context=call_context,
-            )
-
-            active_task = ActiveTask(
-                agent_executor=self._agent_executor,
-                task_id=task_id,
-                task_manager=task_manager,
-                push_sender=self._push_sender,
-                on_cleanup=self._on_active_task_cleanup,
-            )
-            self._active_tasks[task_id] = active_task
+            existing = self._active_tasks.get(task_id)
+            if existing is None:
+                task_manager = TaskManager(
+                    task_id=task_id,
+                    context_id=context_id,
+                    task_store=self._task_store,
+                    initial_message=initial_message,
+                    context=call_context,
+                )
+
+                active_task = ActiveTask(
+                    agent_executor=self._agent_executor,
+                    task_id=task_id,
+                    task_manager=task_manager,
+                    push_sender=self._push_sender,
+                    on_cleanup=self._on_active_task_cleanup,
+                )
+                self._active_tasks[task_id] = active_task
+
+        if existing is not None:
+            # Owner-aware guard on the cache-hit path. The miss path below is
+            # owner-scoped by ActiveTask.start(), which reads through the task
+            # store with call_context and raises TaskNotFoundError when the
+            # task is not owned and create_task_if_missing is false. A cache
+            # hit returned before that check, so resolving a live task by id
+            # was an unauthenticated lookup (issue #1159, CWE-639): every
+            # call site had to guard first. Enforce it here instead, so the
+            # hit path is symmetric with the miss path and no future caller
+            # can reintroduce the gap. Skipped when create_task_if_missing is
+            # set, which is the on_message_send create path establishing
+            # ownership. Masked as not-found so existence is not leaked. Done
+            # outside _lock because the store read is I/O and the miss-path
+            # check runs outside the lock too.
+            if not create_task_if_missing and not await self._task_store.get(
+                task_id, call_context
+            ):
+                raise TaskNotFoundError
+            return existing
 
         await active_task.start(
             call_context=call_context,

src/a2a/server/request_handlers/default_request_handler_v2.py

@@ -165,6 +165,15 @@ async def on_cancel_task(  # noqa: D102
     ) -> Task | None:
         task_id = params.id
 
+        # Owner-aware isolation: mirror on_get_task / _setup_active_task and
+        # consult the owner-scoped store before resolving the live task from
+        # the registry, which keys on task_id alone. Without this a non-owner
+        # who knows the uuid4 task_id could cancel another user's live task
+        # (issue #1159, CWE-639). Mask a partition mismatch as not-found so
+        # existence is not leaked.
+        if not await self.task_store.get(task_id, context):
+            raise TaskNotFoundError
+
         try:
             active_task = await self._active_task_registry.get_or_create(
                 task_id, call_context=context, create_task_if_missing=False
@@ -395,6 +404,13 @@ async def on_subscribe_to_task(  # noqa: D102
     ) -> AsyncGenerator[Event, None]:
         task_id = params.id
 
+        # Owner-aware isolation before resolving the live task from the
+        # registry (which is keyed on task_id alone); see on_cancel_task.
+        # Prevents a non-owner from subscribing to another user's live
+        # task stream (issue #1159, CWE-639).
+        if not await self.task_store.get(task_id, context):
+            raise TaskNotFoundError
+
         active_task = await self._active_task_registry.get_or_create(
             task_id,
             call_context=context,

tests/integration/test_scenarios.py

@@ -2253,3 +2253,108 @@ async def cancel(
     # Verify that agent can see context_id the same as in task_1 and different task id
     assert agent_context2_id == agent_context1_id
     assert agent_task2_id != agent_task1_id
+
+
+# Cancel-terminal-state (issue #1170): cancel must leave the task in a state a
+# caller can act on. These two scenarios are @technicalpickles' repro from
+# https://github.com/a2aproject/a2a-python/pull/1171, folded in here verbatim
+# with the strict-xfail markers dropped since this PR fixes what they pin. Both
+# executors below define cancel() as an empty teardown, which is what
+# InputRequiredAgent, SlowAgent and DummyAgentExecutor use in this file and what
+# a real cleanup-only executor looks like.
+_TERMINAL_STATES = {
+    TaskState.TASK_STATE_CANCELED,
+    TaskState.TASK_STATE_COMPLETED,
+    TaskState.TASK_STATE_FAILED,
+    TaskState.TASK_STATE_REJECTED,
+}
+
+
+async def _start_task(client, text):
+    it = client.send_message(
+        SendMessageRequest(
+            message=Message(
+                message_id='test-msg',
+                role=Role.ROLE_USER,
+                parts=[Part(text=text)],
+            ),
+            configuration=SendMessageConfiguration(return_immediately=True),
+        )
+    )
+    res = await it.__anext__()
+    return res.task.id if res.HasField('task') else res.status_update.task_id
+
+
+@pytest.mark.timeout(5.0)
+@pytest.mark.asyncio
+async def test_scenario_19_mid_run_cancel_reaches_a_terminal_state():
+    """A caller who cancels should at least be able to stop polling."""
+    started = asyncio.Event()
+    hang = asyncio.Event()
+
+    class SilentCancelAgent(AgentExecutor):
+        async def execute(
+            self, context: RequestContext, event_queue: EventQueue
+        ):
+            task = new_task_from_user_message(context.message)
+            task.status.state = TaskState.TASK_STATE_WORKING
+            await event_queue.enqueue_event(task)
+            started.set()
+            await hang.wait()
+
+        async def cancel(
+            self, context: RequestContext, event_queue: EventQueue
+        ):
+            pass
+
+    client = await create_client(
+        create_handler(SilentCancelAgent(), use_legacy=False),
+        agent_card=agent_card(),
+        streaming=True,
+    )
+    task_id = await _start_task(client, 'hello')
+    await asyncio.wait_for(started.wait(), timeout=1.0)
+
+    await client.cancel_task(CancelTaskRequest(id=task_id))
+    task_after = await client.get_task(GetTaskRequest(id=task_id))
+
+    assert task_after.status.state in _TERMINAL_STATES
+
+
+@pytest.mark.timeout(5.0)
+@pytest.mark.asyncio
+async def test_scenario_19_cancel_of_parked_task_does_not_silently_succeed():
+    """input-required is non-terminal, so cancel should either work or raise
+    TaskNotCancelableError. Reporting success and changing nothing is the one
+    outcome a caller cannot act on."""
+
+    class ParkingAgent(AgentExecutor):
+        async def execute(
+            self, context: RequestContext, event_queue: EventQueue
+        ):
+            task = new_task_from_user_message(context.message)
+            task.status.state = TaskState.TASK_STATE_INPUT_REQUIRED
+            await event_queue.enqueue_event(task)
+
+        async def cancel(
+            self, context: RequestContext, event_queue: EventQueue
+        ):
+            pass
+
+    client = await create_client(
+        create_handler(ParkingAgent(), use_legacy=False),
+        agent_card=agent_card(),
+        streaming=True,
+    )
+    task_id = await _start_task(client, 'start')
+
+    for _ in range(50):
+        parked = await client.get_task(GetTaskRequest(id=task_id))
+        if parked.status.state == TaskState.TASK_STATE_INPUT_REQUIRED:
+            break
+        await asyncio.sleep(0.02)
+    assert parked.status.state == TaskState.TASK_STATE_INPUT_REQUIRED
+
+    result = await client.cancel_task(CancelTaskRequest(id=task_id))
+
+    assert result.status.state != TaskState.TASK_STATE_INPUT_REQUIRED

tests/server/agent_execution/test_active_task.py

@@ -129,6 +129,109 @@ async def execute_mock(req, q):
         agent_executor.cancel.assert_called_once()
         stop_event.set()
 
+    @pytest.mark.asyncio
+    async def test_active_task_cancel_producer_cancelled_on_cancellederror(
+        self,
+        active_task: ActiveTask,
+        agent_executor: Mock,
+        request_context: Mock,
+        task_manager: Mock,
+    ) -> None:
+        """Regression: executor.cancel() raising asyncio.CancelledError must
+        still cancel the producer task.
+
+        asyncio.CancelledError is a BaseException, so `except Exception` does
+        not catch it. The producer cancel therefore lives in a `finally`;
+        without it the producer would leak as a pending task on this path.
+        """
+        hang = asyncio.Event()
+        producer_cancelled = asyncio.Event()
+
+        async def execute_mock(req, q):
+            try:
+                await hang.wait()
+            except asyncio.CancelledError:
+                producer_cancelled.set()
+                raise
+
+        agent_executor.execute = AsyncMock(side_effect=execute_mock)
+        agent_executor.cancel = AsyncMock(side_effect=asyncio.CancelledError)
+        task_manager.get_task = AsyncMock(
+            return_value=Task(
+                id='test-task-id',
+                status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
+            )
+        )
+
+        await active_task.enqueue_request(request_context)
+        await active_task.start(
+            call_context=ServerCallContext(), create_task_if_missing=True
+        )
+        await asyncio.sleep(0.1)  # let the producer reach `await hang.wait()`
+        assert active_task._producer_task is not None
+
+        # The CancelledError from executor.cancel propagates out of cancel()...
+        with pytest.raises(asyncio.CancelledError):
+            await active_task.cancel(request_context)
+        agent_executor.cancel.assert_awaited_once()
+
+        # ...but the producer must have received the cancellation, not been
+        # left blocked on `hang.wait()`. `_run_producer` swallows the
+        # CancelledError and returns, so the executor's coroutine seeing it is
+        # the signal that `_producer_task.cancel()` actually fired. Without the
+        # `finally` this event never sets and the producer leaks pending.
+        for _ in range(50):
+            if producer_cancelled.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert producer_cancelled.is_set(), (
+            'producer was never cancelled -> it leaked as a pending task'
+        )
+        hang.set()
+        await active_task.aclose()
+
+    @pytest.mark.asyncio
+    async def test_active_task_cancel_does_not_mutate_shared_task(
+        self,
+        active_task: ActiveTask,
+        agent_executor: Mock,
+        request_context: Mock,
+        task_manager: Mock,
+    ) -> None:
+        """cancel() must not write the terminal state onto the shared task.
+
+        get_task() returns _task_manager._current_task, which may already have
+        been yielded to a subscriber. The terminal write goes onto a copy so
+        that reference does not change under the reader.
+        """
+        stop_event = asyncio.Event()
+
+        async def execute_mock(req, q):
+            await stop_event.wait()
+
+        shared = Task(
+            id='test-task-id',
+            status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
+        )
+        agent_executor.execute = AsyncMock(side_effect=execute_mock)
+        agent_executor.cancel = AsyncMock()
+        task_manager.get_task = AsyncMock(return_value=shared)
+
+        await active_task.enqueue_request(request_context)
+        await active_task.start(
+            call_context=ServerCallContext(), create_task_if_missing=True
+        )
+        await asyncio.sleep(0.1)
+
+        result = await active_task.cancel(request_context)
+
+        # The returned task is CANCELED...
+        assert result.status.state == TaskState.TASK_STATE_CANCELED
+        # ...but the shared object get_task handed out is untouched.
+        assert result is not shared
+        assert shared.status.state == TaskState.TASK_STATE_WORKING
+        stop_event.set()
+
     @pytest.mark.asyncio
     async def test_active_task_interrupted_auth(
         self,