From 7de6325d595f026e5815e37a70b4edc8b931ee8d Mon Sep 17 00:00:00 2001 From: AlgoVoi Date: Sun, 9 Aug 2026 08:20:55 +0000 Subject: [PATCH 1/5] fix: owner-scope SubscribeToTask and CancelTask for live tasks (#1159) DefaultRequestHandlerV2.on_cancel_task and on_subscribe_to_task resolve the live task through ActiveTaskRegistry.get_or_create(), which keys on task_id alone and returns the cached ActiveTask before the owner-aware TaskStore is consulted. A non-owner who knows the uuid4 task_id could subscribe to another user's live task stream and cancel the task (CWE-639 tenancy bypass). Consult the owner-scoped task_store.get(task_id, context) first and raise TaskNotFoundError when the task is not in the caller's partition, mirroring on_get_task and _setup_active_task and masking a partition mismatch as not-found so existence is not leaked. This also closes the cross-tenant exposure that the #1170 fix (cancel now writes a terminal state for parked tasks) would otherwise surface on the un-guarded cached-active path: a non-owner cancel of a parked (input-required) task would silently write CANCELED. Signed-off-by: AlgoVoi --- .../default_request_handler_v2.py | 16 ++ .../test_default_request_handler_v2.py | 181 ++++++++++++++++++ 2 files changed, 197 insertions(+) diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 872a3bfa2..930dea936 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/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, diff --git a/tests/server/request_handlers/test_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py index b276fb77a..2eafdbaeb 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -1746,3 +1746,184 @@ async def test_aclose_is_idempotent_and_handles_empty(): await handler.aclose() await handler.aclose() + + +# --- Issue #1159: SubscribeToTask / CancelTask must be owner-scoped even when +# the task is LIVE in the ActiveTaskRegistry (the cached-active path that +# previously skipped the owner-aware TaskStore). ------------------------------ + + +class _HangingAgent(AgentExecutor): + """Stays in ``working`` so the task remains live in the registry. + + ``cancel`` writes a terminal CANCELED state, so a legitimate owner cancel + resolves the same way regardless of the #1170 fix. + """ + + def __init__(self) -> None: + self.working = asyncio.Event() + self.release = asyncio.Event() + self.cancel_called = asyncio.Event() + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + task = new_task_from_user_message(context.message) + await event_queue.enqueue_event(task) + updater = TaskUpdater(event_queue, task.id, task.context_id) + await updater.update_status(TaskState.TASK_STATE_WORKING) + self.working.set() + await self.release.wait() + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + self.cancel_called.set() + updater = TaskUpdater(event_queue, context.task_id, context.context_id) + await updater.cancel() + + +class _ParkThenHangAgent(AgentExecutor): + """Parks the task in the non-terminal ``input-required`` state and keeps + the producer alive, so the ActiveTask stays in the registry. ``cancel`` is + cleanup-only (writes no terminal state).""" + + def __init__(self) -> None: + self.parked = asyncio.Event() + self.release = asyncio.Event() + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + task = new_task_from_user_message(context.message) + await event_queue.enqueue_event(task) + updater = TaskUpdater(event_queue, task.id, task.context_id) + await updater.update_status(TaskState.TASK_STATE_INPUT_REQUIRED) + self.parked.set() + await self.release.wait() + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + pass + + +async def _start_live_task(handler, ctx, text): + task = await handler.on_message_send( + SendMessageRequest( + message=Message( + message_id=f'msg-{text}', + role=Role.ROLE_USER, + parts=[Part(text=text)], + ), + configuration=SendMessageConfiguration(return_immediately=True), + ), + ctx, + ) + return task.id + + +@pytest.mark.timeout(10) +@pytest.mark.asyncio +async def test_on_cancel_task_is_owner_scoped_for_live_task(): + """Issue #1159: a non-owner must not cancel another user's LIVE task.""" + agent = _HangingAgent() + handler = DefaultRequestHandlerV2( + agent_executor=agent, + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + alice = _ctx('alice') + bob = _ctx('bob') + + task_id = await _start_live_task(handler, alice, 'work') + await asyncio.wait_for(agent.working.wait(), timeout=5) + # The task is live in the registry -> exercises the cached-active path. + assert await handler._active_task_registry.get(task_id) is not None + + # Bob (non-owner) is rejected, masked as not-found, and never reaches the + # executor's cancel(). + with pytest.raises(TaskNotFoundError): + await handler.on_cancel_task(CancelTaskRequest(id=task_id), bob) + assert not agent.cancel_called.is_set() + + # Alice (owner) can still cancel her own task. + result = await handler.on_cancel_task(CancelTaskRequest(id=task_id), alice) + assert result.status.state == TaskState.TASK_STATE_CANCELED + assert agent.cancel_called.is_set() + + agent.release.set() + await handler.aclose() + + +@pytest.mark.timeout(10) +@pytest.mark.asyncio +async def test_on_subscribe_to_task_is_owner_scoped_for_live_task(): + """Issue #1159: a non-owner must not subscribe to another user's LIVE task.""" + agent = _HangingAgent() + handler = DefaultRequestHandlerV2( + agent_executor=agent, + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + alice = _ctx('alice') + bob = _ctx('bob') + + task_id = await _start_live_task(handler, alice, 'work') + await asyncio.wait_for(agent.working.wait(), timeout=5) + assert await handler._active_task_registry.get(task_id) is not None + + # Bob (non-owner) is rejected on the first iteration of the stream. + with pytest.raises(TaskNotFoundError): + async for _ in handler.on_subscribe_to_task( + SubscribeToTaskRequest(id=task_id), bob + ): + break + + # Alice (owner) can subscribe and receives her own task. + received = None + async for event in handler.on_subscribe_to_task( + SubscribeToTaskRequest(id=task_id), alice + ): + received = event + break + assert received is not None + assert getattr(received, 'id', None) == task_id + + agent.release.set() + await handler.aclose() + + +@pytest.mark.timeout(10) +@pytest.mark.asyncio +async def test_on_cancel_of_parked_task_is_owner_scoped(): + """Issue #1159 x #1170: a non-owner cancel of another user's PARKED + (input-required) live task must be rejected and must NOT write a terminal + CANCELED state. This is the exact cross-tenant regression the #1170 fix + (cancel writes a terminal state) would otherwise expose on the un-guarded + cached-active path.""" + agent = _ParkThenHangAgent() + store = InMemoryTaskStore() + handler = DefaultRequestHandlerV2( + agent_executor=agent, + task_store=store, + agent_card=create_default_agent_card(), + ) + alice = _ctx('alice') + bob = _ctx('bob') + + task_id = await _start_live_task(handler, alice, 'park') + await asyncio.wait_for(agent.parked.wait(), timeout=5) + assert await handler._active_task_registry.get(task_id) is not None + + # Bob (non-owner) is rejected... + with pytest.raises(TaskNotFoundError): + await handler.on_cancel_task(CancelTaskRequest(id=task_id), bob) + + # ...and Alice's task is untouched: still input-required, not CANCELED. + alice_view = await store.get(task_id, alice) + assert alice_view is not None + assert alice_view.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + + agent.release.set() + await handler.aclose() From e15567df3f9a447b3ebf1a431d4bd0e8293d5cf0 Mon Sep 17 00:00:00 2001 From: AlgoVoi Date: Sun, 9 Aug 2026 07:58:32 +0000 Subject: [PATCH 2/5] fix: cancel writes a terminal state instead of reporting a false success (#1170) DefaultRequestHandlerV2 could report a successful cancel while leaving the task with no terminal state, so a client that cancelled had no way to know it could stop polling. Two paths reached that: - Mid-run, ActiveTask.cancel cancelled the producer before awaiting the executor's cancel. The producer is the only component that writes the task's terminal state, so killing it first (and a cleanup-only executor.cancel() that writes nothing) left the task WORKING forever. Await the executor's cancel first, mirroring the V1 handler ordering, so a terminal event can still be written to the still-open event queue. - A task parked in a non-terminal state (e.g. input-required) has no running producer, so cancel returned it untouched. After the task settles, close it out as CANCELED when it is still non-terminal, so the outcome is always visible to the caller, mirroring V1's guard that made a non-cancelled outcome visible instead of a silent no-op. Adds two regression scenarios (converted from the xfail repros in #1171) covering both paths with an empty executor cancel(). Signed-off-by: AlgoVoi --- src/a2a/server/agent_execution/active_task.py | 17 ++- tests/integration/test_scenarios.py | 102 ++++++++++++++++++ 2 files changed, 118 insertions(+), 1 deletion(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index ea1955fdd..7611ee2f6 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -730,7 +730,11 @@ 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. try: await self._agent_executor.cancel( request_context, self._event_queue_agent @@ -739,8 +743,10 @@ async def cancel(self, call_context: ServerCallContext) -> Task: logger.exception( 'Cancel[%s]: Agent cancel failed', self._task_id ) + self._producer_task.cancel() await self._mark_task_as_failed(e) raise + self._producer_task.cancel() else: logger.debug( 'Cancel[%s]: Task already finished [%s] or producer not started [%s], not cancelling', @@ -753,6 +759,15 @@ 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: + task.status.state = TaskState.TASK_STATE_CANCELED + await self._task_manager.save_task_event(task) return task async def aclose(self) -> None: diff --git a/tests/integration/test_scenarios.py b/tests/integration/test_scenarios.py index 762270a9a..a08ffd298 100644 --- a/tests/integration/test_scenarios.py +++ b/tests/integration/test_scenarios.py @@ -2253,3 +2253,105 @@ 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 + + +# Scenario 19 (issue #1170): cancel must leave the task in a state a caller can +# act on. 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 From a382ad5d90319b1e3a815f8e821efb73ecbc9aa5 Mon Sep 17 00:00:00 2001 From: AlgoVoi Date: Sun, 9 Aug 2026 10:43:16 +0100 Subject: [PATCH 3/5] test: wait for persisted parked state before cross-tenant cancel (#1159) update_status enqueues the input-required transition; the owner-scope parked test could read the store before it persisted and see SUBMITTED, so poll for the persisted state before the cancel. Deterministic 12/12 locally. Signed-off-by: AlgoVoi --- .../test_default_request_handler_v2.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/server/request_handlers/test_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py index 2eafdbaeb..8a62e5a17 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -1916,6 +1916,21 @@ async def test_on_cancel_of_parked_task_is_owner_scoped(): await asyncio.wait_for(agent.parked.wait(), timeout=5) assert await handler._active_task_registry.get(task_id) is not None + # update_status enqueues the input-required transition; wait until it is + # persisted to the store before the cross-tenant cancel so the assertions + # below do not race the event pipeline. + alice_view = None + for _ in range(500): + alice_view = await store.get(task_id, alice) + if ( + alice_view is not None + and alice_view.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + ): + break + await asyncio.sleep(0.01) + assert alice_view is not None + assert alice_view.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + # Bob (non-owner) is rejected... with pytest.raises(TaskNotFoundError): await handler.on_cancel_task(CancelTaskRequest(id=task_id), bob) From 864828dcc12ad35b09af1dcc70036998a77168bf Mon Sep 17 00:00:00 2001 From: AlgoVoi Date: Mon, 10 Aug 2026 05:00:51 +0100 Subject: [PATCH 4/5] fix: cancel producer in a finally so CancelledError does not leak it asyncio.CancelledError is a BaseException, so the reordered cancel() path cancelled the producer on neither the `except Exception` arm nor the trailing statement when agent_executor.cancel() raised it, leaving the producer as a pending task. Move the producer cancel into a `finally` so it runs on every exit path, including BaseException. This also keeps the FAILED write ahead of the producer/queue teardown on the except arm (previously the producer was cancelled before _mark_task_as_failed, risking the terminal write being dropped as QueueShutDown). Add a regression test: a hanging execute() with a cancel() that raises CancelledError now cancels the producer instead of leaking it. Thanks to @astrogilda for catching the regression on #1172. Signed-off-by: AlgoVoi --- src/a2a/server/agent_execution/active_task.py | 12 +++- .../agent_execution/test_active_task.py | 61 +++++++++++++++++++ 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 7611ee2f6..6c25d86ce 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -734,7 +734,13 @@ async def cancel(self, call_context: ServerCallContext) -> Task: # 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. + # 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 @@ -743,10 +749,10 @@ async def cancel(self, call_context: ServerCallContext) -> Task: logger.exception( 'Cancel[%s]: Agent cancel failed', self._task_id ) - self._producer_task.cancel() await self._mark_task_as_failed(e) raise - self._producer_task.cancel() + finally: + self._producer_task.cancel() else: logger.debug( 'Cancel[%s]: Task already finished [%s] or producer not started [%s], not cancelling', diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 1be233ee1..5072dd361 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -129,6 +129,67 @@ 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_interrupted_auth( self, From 6e26e2e4d07b6f10c58eeda4c62c25f005a128a7 Mon Sep 17 00:00:00 2001 From: AlgoVoi Date: Mon, 10 Aug 2026 05:14:31 +0100 Subject: [PATCH 5/5] fix: cancel() writes terminal state onto a copy, not the shared task get_task() returns the shared _task_manager._current_task, which may already have been yielded to a subscriber. Writing the CANCELED state onto that object in place changes an already-yielded reference under the reader. Copy the task, set the terminal state on the copy, and persist the copy instead. Also clarify in the cancel() docstring that the terminal state reaches the task store but is not guaranteed to reach an active subscriber stream (the event queues may already be closed), so a streaming client may have to re-read the task. Add a regression test asserting the shared task object is left unchanged while the returned task carries the terminal state. Raised by @astrogilda in review of #1172. Signed-off-by: AlgoVoi --- src/a2a/server/agent_execution/active_task.py | 17 +++++++- .../agent_execution/test_active_task.py | 42 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 6c25d86ce..b73afbe51 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/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 @@ -772,8 +778,15 @@ async def cancel(self, call_context: ServerCallContext) -> Task: # 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: - task.status.state = TaskState.TASK_STATE_CANCELED - await self._task_manager.save_task_event(task) + # 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: diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 5072dd361..dcc024866 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -190,6 +190,48 @@ async def execute_mock(req, q): 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,