Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion src/a2a/server/agent_execution/active_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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',
Expand All @@ -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:
Expand Down
16 changes: 16 additions & 0 deletions src/a2a/server/request_handlers/default_request_handler_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
102 changes: 102 additions & 0 deletions tests/integration/test_scenarios.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
103 changes: 103 additions & 0 deletions tests/server/agent_execution/test_active_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading
Loading