Skip to content

fix: publish terminal task state to subscriber streams (#1175) #1191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
anthonyrs06 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: publish terminal task state to subscriber streams (#1175) #1191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
139 changes: 104 additions & 35 deletions src/a2a/server/agent_execution/active_task.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
from __future__ import annotations

import asyncio
import inspect
import logging
import uuid

Expand Down Expand Up @@ -560,17 +561,26 @@ async def _run_producer(self) -> None:
'Producer[%s]: Execution failed',
self._task_id,
)
# Persist the failure directly instead of relying on the closing
# event queue to carry a final status update.
# Persist FAILED (store + push) before finally closes the
# queues (issue #1175). Do not emit a FAILED status event:
# blocking on_message_send would treat that Task as success.
# The producer exception is the stream signal.
if request_context:
task = await self._task_manager.ensure_task_id(
self._task_id,
request_context.context_id or '',
)
if task.status.state not in TERMINAL_TASK_STATES:
task.status.state = TaskState.TASK_STATE_FAILED
await self._task_manager.save_task_event(task)
self._task_created.set()
try:
await self._task_manager.ensure_task_id(
self._task_id,
request_context.context_id or '',
)
await self._persist_and_publish_terminal(
TaskState.TASK_STATE_FAILED,
publish_to_stream=False,
)
self._task_created.set()
except Exception:
logger.exception(
'Producer[%s]: Failed to persist FAILED state',
self._task_id,
)
await self._event_queue_agent.enqueue_event(cast('Event', e))

finally:
Expand Down Expand Up @@ -711,11 +721,11 @@ 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.
The returned task carries a terminal state written to the store.
When a producer is still running, that state is also published to
the live subscriber stream and push *before* the producer is
cancelled (issue #1175). If the producer has already finished, the
subscriber queue is already closed, so the write is store plus push.

Concurrency Guarantee:
Uses `_lock` to ensure we don't attempt to cancel a producer that is
Expand Down Expand Up @@ -755,7 +765,7 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
# Residual: if executor.cancel() raises a BaseException (e.g.
# asyncio.CancelledError), it propagates out of cancel() from
# the finally before the _is_finished.wait() and the CANCELED
# write below, leaving the task non-terminal with no producer
# fallback below, leaving the task non-terminal with no producer
# behind it. This is not the #1170 silent-success shape (the
# caller receives the exception), and it is recoverable: a
# later cancel() takes the else branch and reaches the
Expand All @@ -771,6 +781,21 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
await self._mark_task_as_failed(e)
raise
finally:
try:
# Cleanup-only executor.cancel() or a parked
# input-required task leaves no terminal event.
# Publish CANCELED while the subscriber queue is
# still open — producer teardown lives in
# _run_producer.finally and cannot run until
# _producer_task.cancel() below.
await self._persist_and_publish_terminal(
TaskState.TASK_STATE_CANCELED
)
except Exception:
logger.exception(
'Cancel[%s]: Failed to persist CANCELED state',
self._task_id,
)
self._producer_task.cancel()
else:
logger.debug(
Expand All @@ -784,26 +809,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.
# Fallback when the running-producer path did not write (producer
# already finished or never started). The subscriber queue is
# typically already closed here: store + push only.
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. Copying keeps THIS terminal
# write off the yielded reference. It is not a file-wide guarantee:
# ordinary status writes (save_task_event -> ensure_task -> CopyFrom,
# and the producer-failure write above) still update the shared
# object in place. Making every write copy-on-yield is the general
# property tracked in #1175 and #1191.
updated = Task()
updated.CopyFrom(task)
updated.status.state = TaskState.TASK_STATE_CANCELED
await self._task_manager.save_task_event(updated)
task = updated
persisted = await self._persist_and_publish_terminal(
TaskState.TASK_STATE_CANCELED
)

@mykytanetipa mykytanetipa Sep 7, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_persist_and_publish_terminal is called after await self._is_finished.wait() -> the producer's finally closes the subscriber queue -> await self._event_queue_subscribers.enqueue_event (line 932) -> logs a WARNING each time

Should we have here publish_to_stream=False?

@gomission gomission Sep 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I independently checked this fallback with a real TaskManager/InMemoryTaskStore, the real event queues, and a fake push sender. Your suggested publish_to_stream=False removes the warning while preserving the cancellation outcome.

At PR head 6e9144edf885ae1f8ec1cc6afa5cbfc119187ca0, after the fake executor exits with CancelledError and the subscriber finishes, calling ActiveTask.cancel() gives:

  • returned task and store: CANCELED;
  • fake push sender: exactly one CANCELED notification;
  • exactly one Queue was closed during enqueuing. Event dropped. warning.

Adding the flag only to the post-wait fallback removes that warning. A live-producer control still receives one CANCELED stream event and one fake push, so the live publication path remains exercised.

Revision Live-producer control Finished-producer fallback
PR head Pass Fails only the no-warning assertion
PR head + your one-line flag Pass Pass

This confirms a dropped stream publication attempt; the stored cancellation and fake push call are retained. I did not use HTTP, a model, or real push delivery. The test reaches the finished state through the executor/consumer lifecycle; it does not set private flags or replace queue behavior.

With only that one-line source change, repository lint exits 0 and the full suite reports 1988 passed, 90 skipped, 3 xfailed, 1 xpassed (Python 3.12.11, locked development environment). The separate two-case fixture below also passes.

Two-case local regression

Save as test_cancel_fallback.py and run python -m pytest -q -s test_cancel_fallback.py in an SDK environment (Python 3.12.11 used here).

"""Local lifecycle regression for A2A PR1191, without HTTP or model calls.

The finished-producer case raises CancelledError from the fake executor and
lets ActiveTask close its own queues before cancel() is called. It does not
set private lifecycle flags or replace the store/queue implementation.
"""

import asyncio
import json
import logging
from typing import Any

import pytest
from a2a.server.agent_execution.active_task import ActiveTask
from a2a.server.agent_execution.agent_executor import AgentExecutor
from a2a.server.agent_execution.context import RequestContext
from a2a.server.context import ServerCallContext
from a2a.server.events.event_queue_v2 import EventQueueSource
from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
from a2a.server.tasks.push_notification_sender import (
    PushNotificationEvent,
    PushNotificationSender,
)
from a2a.server.tasks.task_manager import TaskManager
from a2a.types.a2a_pb2 import Task, TaskState, TaskStatus, TaskStatusUpdateEvent


class LocalExecutor(AgentExecutor):
    def __init__(self) -> None:
        self.entered = asyncio.Event()
        self.finish_cancelled = asyncio.Event()

    async def execute(
        self, context: RequestContext, event_queue: EventQueueSource
    ) -> None:
        self.entered.set()
        await self.finish_cancelled.wait()
        raise asyncio.CancelledError

    async def cancel(
        self, context: RequestContext, event_queue: EventQueueSource
    ) -> None:
        # A cleanup-only cancel is the scenario addressed by PR1191.
        return None


class LocalPush(PushNotificationSender):
    def __init__(self) -> None:
        self.states: list[int] = []

    async def send_notification(
        self, task_id: str, event: PushNotificationEvent
    ) -> None:
        assert isinstance(event, TaskStatusUpdateEvent)
        self.states.append(event.status.state)


@pytest.mark.asyncio
@pytest.mark.parametrize("finished_first", [False, True], ids=["live", "finished"])
async def test_cancel_outcome_and_closed_queue(
    finished_first: bool, caplog: pytest.LogCaptureFixture
) -> None:
    context = ServerCallContext()
    store = InMemoryTaskStore()
    await store.save(
        Task(
            id="local-task",
            context_id="local-context",
            status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
        ),
        context,
    )
    manager = TaskManager(store, context, "local-task", "local-context", None)
    executor, push = LocalExecutor(), LocalPush()
    active = ActiveTask(executor, "local-task", manager, push)
    request = RequestContext(
        call_context=context, task_id="local-task", context_id="local-context"
    )
    subscribed = asyncio.Event()
    stream_states: list[int] = []

    async def collect() -> None:
        async for event in active.subscribe(include_initial_task=True):
            if isinstance(event, Task | TaskStatusUpdateEvent):
                stream_states.append(event.status.state)
            subscribed.set()

    collector: asyncio.Task[None] | None = None
    try:
        await active.enqueue_request(request)
        await active.start(call_context=context, create_task_if_missing=False)
        await asyncio.wait_for(executor.entered.wait(), 5)
        collector = asyncio.create_task(collect())
        await asyncio.wait_for(subscribed.wait(), 5)
        if finished_first:
            executor.finish_cancelled.set()
            await asyncio.wait_for(collector, 5)
        caplog.clear()
        with caplog.at_level(logging.WARNING):
            result = await asyncio.wait_for(active.cancel(context), 5)
            await asyncio.wait_for(collector, 5)
        stored = await store.get("local-task", context)
        warnings = [
            record.getMessage()
            for record in caplog.records
            if "Queue was closed during enqueuing" in record.getMessage()
        ]
        observation: dict[str, Any] = {
            "case": "finished" if finished_first else "live",
            "returned": TaskState.Name(result.status.state),
            "stored": TaskState.Name(stored.status.state) if stored else None,
            "fake_push_states": [TaskState.Name(state) for state in push.states],
            "stream_states": [TaskState.Name(state) for state in stream_states],
            "closed_queue_warnings": len(warnings),
        }
        print(json.dumps(observation, sort_keys=True))
        assert stored is not None
        assert (
            result.status.state == stored.status.state == TaskState.TASK_STATE_CANCELED
        )
        assert push.states == [TaskState.TASK_STATE_CANCELED]
        if finished_first:
            assert stream_states == [TaskState.TASK_STATE_WORKING]
        else:
            assert stream_states == [
                TaskState.TASK_STATE_WORKING,
                TaskState.TASK_STATE_CANCELED,
            ]
        assert warnings == []
    finally:
        await active.aclose()
        if collector is not None:
            await asyncio.gather(collector, return_exceptions=True)

Disclosure: Mission, with Codex assistance. This adds regression evidence for @mykytanetipa’s suggestion, also confirmed by @astrogilda, on @anthonyrs06’s existing patch.

if persisted is not None:
task = persisted
return task

async def aclose(self) -> None:
Expand Down Expand Up @@ -871,6 +885,61 @@ async def _maybe_cleanup(self) -> None:
logger.debug('Cleanup[%s]: Triggering cleanup', self._task_id)
self._on_cleanup(self)

async def _persist_and_publish_terminal(
self, state: TaskState, *, publish_to_stream: bool = True
) -> Task | None:
"""Write a terminal state to the store and notify live observers.

Direct ``save_task_event`` after the subscriber queue is closed is
invisible to ``SubscribeToTask`` / ``message/stream`` and to push
(issue #1175). This helper persists a *copy* of the current task
(the shared ``get_task()`` object must not mutate under a reader)
and, when ``publish_to_stream`` is true, emits a
``TaskStatusUpdateEvent`` to subscribers *before* teardown.

Returns the current task unchanged (no store write, no stream
publish, no push) when it is already terminal, so a later
cancel/failure cannot emit a duplicate notification carrying the
earlier state.

Producer-failure keeps ``publish_to_stream=False``: blocking
``on_message_send`` treats a FAILED ``Task`` as a successful
terminal result, so the crash must still surface as the
producer exception on the stream. Store and push still get
FAILED before the queues close.
"""
task = await self._task_manager.get_task()
if task is None or task.status.state in TERMINAL_TASK_STATES:
return task

updated = Task()
updated.CopyFrom(task)
updated.status.state = state
await self._task_manager.save_task_event(updated)
task = updated

event = TaskStatusUpdateEvent(
task_id=task.id,
context_id=task.context_id,
status=TaskStatus(state=task.status.state),
)
if publish_to_stream:
# Subscriber taps dequeue (event, updated_task), matching
# EventConsumer._enqueue_to_subscribers. Event is the
# agent-queue type; widening it would lie about that queue.
updated_task_copy = Task()
updated_task_copy.CopyFrom(task)
await self._event_queue_subscribers.enqueue_event(
cast('Any', (event, updated_task_copy))
)
if self._push_sender and self._task_id:
notification = self._push_sender.send_notification(
self._task_id, event
)
if inspect.isawaitable(notification):
await notification
return task

async def _mark_task_as_failed(self, exception: Exception) -> Task | None:
logger.debug('Marking task %s as failed: %s', self._task_id, exception)
task = None
Expand Down
165 changes: 165 additions & 0 deletions tests/server/agent_execution/test_active_task.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,171 @@ async def execute_mock(req, q):
agent_executor.cancel.assert_called_once()
stop_event.set()

@staticmethod
def _wire_current_task(task_manager: Mock, task: Task) -> None:
"""Make get_task / save_task_event share one current Task."""

async def get_task() -> Task:
return task_manager._current_task

async def ensure_task_id(task_id: str, context_id: str) -> Task:
return task_manager._current_task

async def save_task_event(event: Task) -> None:
if isinstance(event, Task):
task_manager._current_task = event

task_manager._current_task = task
task_manager.get_task = AsyncMock(side_effect=get_task)
task_manager.save_task_event = AsyncMock(side_effect=save_task_event)
task_manager.ensure_task_id = AsyncMock(side_effect=ensure_task_id)

@pytest.mark.asyncio
async def test_cancel_publishes_canceled_to_subscriber_stream(
self,
active_task: ActiveTask,
agent_executor: Mock,
request_context: Mock,
task_manager: Mock,
push_sender: Mock,
) -> None:
"""Issue #1175: cancel() must emit CANCELED on a live subscriber stream.

A cleanup-only executor.cancel() writes nothing. The helper persists a
copy and enqueues TaskStatusUpdateEvent before the producer tears the
subscriber queue down.
"""
stop_event = asyncio.Event()

async def execute_mock(req, q):
await stop_event.wait()

shared = Task(
id='test-task-id',
context_id='test-context-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
self._wire_current_task(task_manager, shared)
agent_executor.execute = AsyncMock(side_effect=execute_mock)
agent_executor.cancel = AsyncMock()

await active_task.enqueue_request(request_context)
await active_task.start(
call_context=ServerCallContext(), create_task_if_missing=True
)
await asyncio.sleep(0.05)

events: list[object] = []

async def collect() -> None:
try:
async for event in active_task.subscribe():
events.append(event)
except Exception: # noqa: BLE001
pass

collector = asyncio.create_task(collect())
await asyncio.sleep(0.05)

result = await active_task.cancel(request_context)
stop_event.set()
await asyncio.wait_for(collector, timeout=2)

assert result.status.state == TaskState.TASK_STATE_CANCELED
assert result is not shared
assert shared.status.state == TaskState.TASK_STATE_WORKING
status_events = [
e
for e in events
if isinstance(e, TaskStatusUpdateEvent)
and e.status.state == TaskState.TASK_STATE_CANCELED
]
assert status_events, f'subscriber saw {events!r}, expected CANCELED'
push_sender.send_notification.assert_awaited()

@pytest.mark.asyncio
async def test_producer_failure_persists_failed_and_notifies_push(
self,
active_task: ActiveTask,
agent_executor: Mock,
request_context: Mock,
task_manager: Mock,
push_sender: Mock,
) -> None:
"""Issue #1175: producer-failure persists FAILED and notifies push.

The crash still surfaces as ValueError on the stream: blocking
on_message_send would treat a FAILED Task as a successful result.
"""
shared = Task(
id='test-task-id',
context_id='test-context-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
self._wire_current_task(task_manager, shared)
request_context.context_id = 'test-context-id'
crash = asyncio.Event()

async def execute_mock(req, q):
await crash.wait()
raise ValueError('Producer crashed')

agent_executor.execute = AsyncMock(side_effect=execute_mock)

await active_task.enqueue_request(request_context)
await active_task.start(
call_context=ServerCallContext(), create_task_if_missing=True
)
# Let the consumer flush _RequestStarted so this tap sees only the
# terminal publish (same timing as the cancel-stream test).
await asyncio.sleep(0.05)

collector_error: list[BaseException] = []

async def collect() -> None:
try:
async for _event in active_task.subscribe():
pass
except ValueError as exc:
collector_error.append(exc)
return

collector = asyncio.create_task(collect())
await asyncio.sleep(0.05)
crash.set()
await asyncio.wait_for(collector, timeout=2)

assert collector_error, 'subscriber hung up instead of seeing the crash'
assert (
task_manager._current_task.status.state
== TaskState.TASK_STATE_FAILED
)
assert shared.status.state == TaskState.TASK_STATE_WORKING
push_sender.send_notification.assert_awaited()

@pytest.mark.asyncio
async def test_persist_terminal_skips_when_already_terminal(
self,
active_task: ActiveTask,
task_manager: Mock,
push_sender: Mock,
) -> None:
"""Already-terminal tasks must not get a duplicate store/push write."""
shared = Task(
id='test-task-id',
context_id='test-context-id',
status=TaskStatus(state=TaskState.TASK_STATE_CANCELED),
)
self._wire_current_task(task_manager, shared)

result = await active_task._persist_and_publish_terminal(
TaskState.TASK_STATE_FAILED
)

assert result is shared
task_manager.save_task_event.assert_not_called()
push_sender.send_notification.assert_not_called()

@pytest.mark.asyncio
async def test_active_task_cancel_producer_cancelled_on_cancellederror(
self,
Expand Down
Loading