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

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.