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

[chopmob-cloud]

Summary

Two related fixes to DefaultRequestHandlerV2, shipped together because the first is a prerequisite for the second to be safe.

Fixes #1159 (owner-scoping, the security base). on_cancel_task and on_subscribe_to_task resolved a live task via ActiveTaskRegistry.get_or_create(task_id) by task_id alone, skipping the owner-aware task_store.get(task_id, context) that on_get_task and the send path already perform. A caller who knows another tenant's task_id could subscribe to its live stream and cancel it (CWE-639). Both handlers now consult the owner-aware store first and fail closed with TaskNotFoundError (masking existence, matching on_get_task).

Fixes #1170 (cancel writes a terminal state). ActiveTask.cancel cancelled the producer before AgentExecutor.cancel, so the component that owns the terminal state could not write it; and a task parked non-terminal (e.g. input-required) reported cancel-success without transitioning. The executor cancel now runs before the producer cancel (producer still cancelled on the error path), and cancellation that leaves the task non-terminal is closed out as CANCELED.

Why one PR

#1170 alone makes #1159 worse: for a parked task, a cross-tenant cancel via the #1159 bypass changes from a harmless no-op into an actual cross-tenant CANCELED write. Landing #1159's owner check together (as the base commit) removes that window. Merging them together, or #1159 before #1170, is safe; #1170 must not land first.

Tests

• Owner-scope regressions: non-owner cancel and non-owner subscribe are rejected (TaskNotFoundError), including a parked task; the owner still succeeds.
• Cancel-terminal regressions: mid-run cancel reaches a terminal state; cancel of a parked task does not silently succeed (the maintainer repros from test(server): add xfail repros for cancel leaving no terminal state #1171, converted from xfail to passing).

Validation

• Owner-scope tests: reverting the [Bug]: SubscribeToTask and CancelTask bypass owner-aware TaskStore for live tasks #1159 handler change fails all three; restore passes. Full request-handler module: 62 passed.
• Cancel-terminal tests: reverting the Cancel reports success but writes no terminal state (DefaultRequestHandlerV2) #1170 change fails both scenario_19 tests; restore passes. Full integration module: no regression.
• Combined: with both applied, a cross-tenant cancel of a parked task is rejected rather than writing CANCELED.
• ruff clean at the pinned version.

[github-actions]

🧪 Code Coverage (vs main)

⬇️ Download Full Report

	Base	PR	Delta
src/a2a/server/agent_execution/active_task.py	95.09%	94.70%	🔴 -0.39%
src/a2a/server/events/event_queue_v2.py	91.79%	91.28%	🔴 -0.51%
src/a2a/server/request_handlers/default_request_handler_v2.py	94.17%	94.37%	🟢 +0.20%
src/a2a/utils/telemetry.py	91.47%	90.70%	🔴 -0.78%
Total	93.00%	92.96%	🔴 -0.03%

Generated by coverage-comment.yml

[chopmob-cloud]

Independently validated the #1170 half of this against the repros in #1171, which adds two strict-xfail scenarios (test_scenario_19_mid_run_cancel_reaches_a_terminal_state and test_scenario_19_cancel_of_parked_task_does_not_silently_succeed).

On a clean checkout of current main both xfail as expected (the bug is present); run against this branch they XPASS, and forced with --runxfail they pass outright: 2 passed on this branch versus 2 failed on main. This PR's own suite for the two changed files is green (174 passed, 0 failed).

So the repros in #1171 are effectively an independent check on the cancel-terminal-state fix here. If useful they could be folded into this PR with the xfail markers dropped, giving direct regression coverage alongside the owner-scoping (#1159) tests.

[astrogilda]

Outside contributor here, not a maintainer. I have one merged fix in active_task.py (#1105, the aclose() drain, shipped in 1.1.2) and one open PR in event_queue_v2.py (#1137), so this is a neighbourhood I have had my hands in rather than any kind of authority. I ran this branch rather than only reading it, and I found one thing I think is a real regression plus two things worth a maintainer's opinion. Everything below is reproducible from a382ad5 and cff6727.

First, on whether the bugs are real. I took your three new handler tests and ran them against unmodified main with your source changes reverted: all three fail with Failed: DID NOT RAISE <class 'a2a.utils.errors.TaskNotFoundError'>, meaning Bob genuinely does cancel and subscribe to Alice's live task today. That is #1159 reproducing under its own tests, and it is worth saying plainly in a thread that has been quiet for six days: this is a live cross-tenant authorization gap, not a hypothetical. On this branch all three pass.

Second, the #1170 half, cross-checked against #1171 rather than against your own tests. I checked out main and this branch into two worktrees, dropped in #1171's tests/integration/test_scenarios.py verbatim in both, and ran only its two scenarios. On main both report XFAIL (exit 0). On this branch both report XPASS(strict), which pytest surfaces as a failure precisely because the bug they pin is gone, and with --runxfail both pass outright (2 passed). So @technicalpickles' deliberately fix-free repro is an independent confirmation of this branch, and I would suggest folding those two tests in here with the markers dropped, since they cover the mid-run path at the client-and-transport level rather than at the handler level.

Third, no regressions in the existing suite: 1768 passed on main versus 1773 passed on this branch, same 90 skipped, same 3 xfailed, --ignore=tests/integration/cross_version on both because that directory needs network I do not have. So the five new tests are the entire delta.

Now the thing I think needs changing. The reordering leaves a path where the producer task is never cancelled at all. The new shape cancels the producer in the except Exception arm and again after the try, but asyncio.CancelledError is a BaseException, so if agent_executor.cancel() raises it, neither statement runs and the producer survives as a pending asyncio.Task. On main this could not happen, because the producer was cancelled unconditionally before the executor was awaited, so this is a regression rather than a pre-existing gap. I built a probe with a hanging execute() and a cancel() that raises asyncio.CancelledError, and asserted on active._producer_task.done() after the call: on main it is True, on this branch it is False. That is the same failure class @mykytanetipa flagged on my #1105, where a dispatcher cancelled but not awaited surfaced as Task was destroyed but it is pending!. A try/finally closes it and removes the duplicated call:

try:
    await self._agent_executor.cancel(
        request_context, self._event_queue_agent
    )
except Exception as e:
    logger.exception('Cancel[%s]: Agent cancel failed', self._task_id)
    await self._mark_task_as_failed(e)
    raise
finally:
    self._producer_task.cancel()

The second point is about where the owner check belongs, and I raise it as a question because I may be missing a reason for the current placement. The registry's miss path is already owner-scoped: get_or_create on a miss calls ActiveTask.start(), which reads through TaskManager.get_task() with the caller's call_context and raises TaskNotFoundError when create_task_if_missing is false. I verified that by evicting a task from the registry while leaving it in the store and then having a non-owner cancel it on unmodified main: rejected with TaskNotFoundError. So the gap is exactly the cache-hit early return at active_task_registry.py:52-53, which returns before call_context is ever consulted. Guarding it there, conditioned on not create_task_if_missing so on_message_send is unaffected, would make the hit path symmetric with the miss path and would close the class for every present and future caller of the registry, rather than for the two handlers that happen to be known today. Two call-site checks fix the two known holes and leave ActiveTaskRegistry itself as an unauthenticated lookup. If there is a reason the check has to live in the handler, I would find that useful to know.

The third point is scope rather than correctness. The terminal state now reaches the store but never reaches the wire. I subscribed to a live task, cancelled it, and collected the events the subscriber actually received: the stream yields the initial Task in WORKING and then ends, while task_store.get afterwards returns CANCELED. A client that was streaming therefore still cannot tell that the task was cancelled without polling, and push_sender does not fire either, since _update_task_state only notifies on PushNotificationEvent. I do not think that blocks this PR, because at that point both queues are already closed and a direct store write is the only option left, which is also why the producer-failure path at active_task.py:570-572 does exactly the same thing, under a comment that says so. But it is worth a sentence in the docstring or a follow-up issue, because "cancel now writes a terminal state" reads as a stronger promise than what a streaming client receives. Relatedly, mutating task.status.state in place mutates the shared _task_manager._current_task, which is the same object already handed to a subscriber, so an already-yielded reference changes underneath the reader. I noticed that only because it confounded my first measurement, and it may well be fine given the same pattern exists elsewhere in the file.

On splitting: I would keep these together. The repo has split bundled PRs before, including one of mine, but here your third test is the argument against it, because it shows the two changes interact. Landing #1170's unconditional terminal write without #1159's owner check would turn a read-only cross-tenant subscribe into a cross-tenant write, which is strictly worse than either bug alone. If a maintainer does want them separated, the safe order is #1159 first and #1170 second, never the reverse.

Reproduction details if any of the above is worth re-checking: two git worktree checkouts at cff6727 and a382ad5, uv sync --all-extras --group dev, Python 3.10, pytest -p no:randomly.

[chopmob-cloud]

Thanks, this is a good catch and you are right on the first point. It is a
regression we introduced and it should not ship as-is.

The producer leak (point 1): confirmed, adopting your try/finally.

The reorder moved the producer cancel to after the executor await, so the only
statements that cancel it are the except Exception arm and the trailing line.
asyncio.CancelledError has been a BaseException (not an Exception) since
3.8, so when agent_executor.cancel() raises it (a cancelled cancel, a
shutdown, or an executor that raises it directly, as in your probe) neither runs
and the producer is left pending. On main the unconditional pre-await
self._producer_task.cancel() made this impossible, so this is a regression, not
a pre-existing gap. I confirmed the control flow: with the branch shape the
producer cancel fires on neither path once cancel() raises CancelledError;
your finally restores it.

Taking your shape as written:

try:
    await self._agent_executor.cancel(
        request_context, self._event_queue_agent
    )
except Exception as e:
    logger.exception('Cancel[%s]: Agent cancel failed', self._task_id)
    await self._mark_task_as_failed(e)
    raise
finally:
    self._producer_task.cancel()

One extra reason it is the right shape, beyond closing the BaseException hole:
it also fixes an ordering inconsistency the current except arm has with this
PR's own stated intent. Today that arm runs self._producer_task.cancel() and
then await self._mark_task_as_failed(e), but _mark_task_as_failed enqueues
the FAILED TaskStatusUpdateEvent onto _event_queue_agent and swallows
QueueShutDown. Cancelling the producer first, then awaiting, lets the producer
tear the queue down at that await point and the FAILED write gets silently
dropped, which is the same "cancel producer first drops the terminal write"
hazard the reorder was meant to avoid. The finally writes FAILED to the
still-open queue first, then cancels, so both the Exception and
CancelledError paths are correct. I will add a regression test on your probe
(hanging execute, cancel raising CancelledError, assert the producer is
cancelled).

Owner check placement (point 2).

You are right that the miss path is already owner-scoped through
ActiveTask.start() -> TaskManager.get_task(call_context), and that the only
hole is the cache-hit early return. We guarded at the two call sites to mirror
the existing handler-level checks (on_get_task, _setup_active_task) and keep
ActiveTaskRegistry auth-unaware, but I agree the registry-level guard, gated on
not create_task_if_missing as you note so on_message_send is unaffected, is
the stronger invariant and would let the two call-site checks go away. I would
rather a maintainer weigh in on whether auth belongs in the registry before I
move it there, since it shifts which layer owns the check. Happy to do it either
way.

Wire visibility and shared mutation (point 3).

Agreed it is out of scope for this PR: at that point both queues are closed, a
direct store write is the only option left, and it mirrors the producer-failure
path at active_task.py:570-572 for exactly that reason. I will tighten the
docstring so "cancel writes a terminal state" reads as "reaches the store," not
"reaches an active stream," and open a follow-up for wire-level cancel
visibility (a synthetic terminal event to live subscribers plus the
push_sender path). On the in-place task.status.state mutation of the shared
_current_task: good catch, I will stop mutating the yielded object in place and
write a copy so an already-yielded reference does not change under a reader.

On splitting: agreed, keep them together, and thanks for spelling out why (the
#1170 write on the un-guarded path turns a cross-tenant read into a cross-tenant
write). If they are ever split, #1159 first.

I have folded in the try/finally plus a regression test that reproduces your
probe (a hanging execute, a cancel that raises asyncio.CancelledError,
asserting the producer receives the cancellation rather than leaking pending). It
fails on the current branch and passes with the fix, and the existing cancel
tests, the owner-scope handler tests, and the #1170 scenario tests stay green.

[chopmob-cloud]

Opened #1175 to track the stream and push visibility follow-up from this review: the terminal state reaches the store but not an active subscriber stream. It covers both the cancel path and the parallel producer-failure FAILED write.