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
112 changes: 101 additions & 11 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 @@ -730,7 +740,10 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
logger.debug(
'Cancel[%s]: Cancelling producer task', self._task_id
)
self._producer_task.cancel()
# Await executor.cancel before cancelling the producer so a
# terminal write can still reach the open subscriber queue
# (#1172 / #1175). Producer cancel stays in finally so a
# BaseException from executor.cancel cannot leak the producer.
try:
await self._agent_executor.cancel(
request_context, self._event_queue_agent
Expand All @@ -741,13 +754,42 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
)
await self._mark_task_as_failed(e)
raise
finally:
try:
task = await self._task_manager.get_task()
if (
task is not None
and task.status.state not in TERMINAL_TASK_STATES
):
# Cleanup-only executor.cancel() or a parked
# input-required task leaves no terminal event.
# Write CANCELED and publish it while queues
# are still open so a live subscriber does not
# have to poll.
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(
'Cancel[%s]: Task already finished [%s] or producer not started [%s], not cancelling',
self._task_id,
self._is_finished.is_set(),
self._producer_task,
)
task = await self._task_manager.get_task()
if (
task is not None
and task.status.state not in TERMINAL_TASK_STATES
):
await self._persist_and_publish_terminal(
TaskState.TASK_STATE_CANCELED
)

await self._is_finished.wait()
task = await self._task_manager.get_task()
Expand Down Expand Up @@ -820,6 +862,54 @@ 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.

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:
return None

if task.status.state not in TERMINAL_TASK_STATES:
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:
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
142 changes: 142 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,148 @@ 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_active_task_interrupted_auth(
self,
Expand Down
Loading