test/a2a/test_cancel_race.py
Rama main · réplica de lectura del repositorio original.
# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
"""``AgentExecutor.cancel`` racing the turn's own terminal transition.
The SDK calls ``cancel`` without synchronising with ``execute``: it only checks that its consumer
has not yet processed a terminal event, so a cancel can land after the turn has already completed
(or a turn can try to complete after a cancel ended it). These tests drive the executor directly
on a recording queue to put each side into that window deterministically.
"""
import asyncio
import logging
import threading
from collections.abc import Awaitable, Callable, Coroutine
from contextlib import AsyncExitStack, ExitStack
from typing import Any
import pytest
from a2a.server.agent_execution import RequestContext
from a2a.server.context import ServerCallContext
from a2a.server.events import Event, EventQueue, EventQueueLegacy
from a2a.types import (
Artifact,
Message,
Part,
Role,
SendMessageRequest,
Task,
TaskArtifactUpdateEvent,
TaskState,
TaskStatusUpdateEvent,
)
from ag2 import Agent, Context
from ag2.a2a.events import A2ATaskArtifactUpdate, A2ATaskStatusUpdate
from ag2.a2a.executor import AgentExecutor
from ag2.a2a.extension import MIME_TOOL_CALL
from ag2.a2a.mappers.events import chunk_to_text_artifact
from ag2.a2a.mappers.messages import build_user_message
from ag2.a2a.mappers.parts import is_data_part_with_mime
from ag2.a2a.testing import ObservedAgentExecutor, TaskStatusRecorder, TurnGate
from ag2.context import ConversationContext
from ag2.events import ClientToolCallEvent, ModelMessageChunk, TextInput, ToolCallEvent
from ag2.observers import observer
from ag2.testing import TestConfig
from ag2.tools.final.function_tool import FunctionDefinition, FunctionToolSchema
_TASK_ID = "task-1"
_CONTEXT_ID = "ctx-1"
_TEXT_CLOSED = "text closed"
class _RecordingQueue(EventQueue):
def __init__(self) -> None:
self.events: list[Event] = []
self.on_event: Callable[[Event], Awaitable[None]] | None = None
async def enqueue_event(self, event: Event) -> None:
# The hook runs before the event lands, as a slow enqueue would.
if self.on_event is not None:
await self.on_event(event)
self.events.append(event)
def states(self) -> list[TaskState]:
return [e.status.state for e in self.events if isinstance(e, TaskStatusUpdateEvent)]
def lifecycle(self) -> list[TaskState | str]:
"""Status transitions plus each close of streamed text, in wire order."""
return [
e.status.state if isinstance(e, TaskStatusUpdateEvent) else _TEXT_CLOSED
for e in self.events
if isinstance(e, TaskStatusUpdateEvent) or (isinstance(e, TaskArtifactUpdateEvent) and e.append)
]
def _request() -> RequestContext:
message = Message(
role=Role.ROLE_USER,
parts=[Part(text="go")],
message_id="msg-1",
task_id=_TASK_ID,
context_id=_CONTEXT_ID,
)
return RequestContext(ServerCallContext(), request=SendMessageRequest(message=message))
def _client_tool_request() -> RequestContext:
message = build_user_message(
[TextInput("go")],
tool_schemas=[FunctionToolSchema(function=FunctionDefinition(name="lookup"))],
task_id=_TASK_ID,
context_id=_CONTEXT_ID,
message_id="msg-1",
)
return RequestContext(ServerCallContext(), request=SendMessageRequest(message=message))
def _cancel_request() -> RequestContext:
return RequestContext(ServerCallContext(), task_id=_TASK_ID, context_id=_CONTEXT_ID)
class _ScopedStatuses(TaskStatusRecorder):
def __init__(self) -> None:
super().__init__()
self.closed = asyncio.Event()
self.states_at_close: list[TaskState] = []
def register(self, stack: ExitStack | AsyncExitStack, context: ConversationContext) -> None:
super().register(stack, context)
stack.callback(self.close)
def close(self) -> None:
self.states_at_close = list(self.states)
self.closed.set()
class _CancelOnStatus:
"""Lifecycle observer that cancels the task once it reaches ``state``, then raises ``error`` if given."""
def __init__(
self,
state: TaskState,
queue: _RecordingQueue,
*,
spawn_cancel: bool = False,
error: Exception | None = None,
) -> None:
self.state = state
self.queue = queue
self.spawn_cancel = spawn_cancel
self.error = error
self.executor: AgentExecutor | None = None
self.observed: list[TaskState] = []
async def on_status(self, event: A2ATaskStatusUpdate) -> None:
self.observed.append(event.state)
if event.state != self.state:
return
assert self.executor is not None
cancellation = self.executor.cancel(_cancel_request(), self.queue)
if self.spawn_cancel:
await asyncio.create_task(cancellation)
else:
await cancellation
if self.error is not None:
raise self.error
class _PauseOnStatus:
"""Lifecycle observer that blocks delivery of ``state`` until resumed."""
def __init__(self, state: TaskState) -> None:
self.state = state
self.paused = asyncio.Event()
self._resumed = asyncio.Event()
async def on_status(self, event: A2ATaskStatusUpdate) -> None:
if event.state == self.state:
self.paused.set()
await self._resumed.wait()
def resume(self) -> None:
self._resumed.set()
class _PausingArtifactQueue(_RecordingQueue):
def __init__(self) -> None:
super().__init__()
self.started = asyncio.Event()
self.release = asyncio.Event()
async def enqueue_event(self, event: Event) -> None:
if isinstance(event, TaskArtifactUpdateEvent) and not self.started.is_set():
self.started.set()
await self.release.wait()
await super().enqueue_event(event)
class _DropAndCancel:
"""Interrupting lifecycle observer that drops the first ``state`` update and cancels the task."""
def __init__(self, state: TaskState, queue: _RecordingQueue) -> None:
self.state = state
self.queue = queue
self.executor: AgentExecutor | None = None
self.pending: list[asyncio.Task[None]] = []
async def on_status(self, event: A2ATaskStatusUpdate) -> A2ATaskStatusUpdate | None:
if event.state != self.state or self.pending:
return event
assert self.executor is not None
self.pending.append(asyncio.create_task(self.executor.cancel(_cancel_request(), self.queue)))
return None
class _SendClientCall:
def __init__(self) -> None:
self.started = asyncio.Event()
self.release = asyncio.Event()
self.sender: asyncio.Task[None] | None = None
self.recovered = False
async def ask_remote(self, context: Context) -> str:
self.started.set()
await self.release.wait()
self.sender = asyncio.create_task(context.send(ClientToolCallEvent(name="lookup", arguments="{}")))
try:
await self.sender
except asyncio.CancelledError:
self.recovered = True
return "fallback"
async def _drain_queue(queue: EventQueueLegacy) -> None:
while True:
await queue.dequeue_event()
queue.task_done()
@pytest.mark.asyncio
class TestCancelRacingTheTurn:
@pytest.mark.parametrize("cancel_sender", [False, True])
async def test_published_client_call_is_kept_when_sender_is_canceled(self, cancel_sender: bool) -> None:
tool = _SendClientCall()
config = TestConfig(ToolCallEvent(name="ask_remote", arguments="{}"), "done")
executor = AgentExecutor(Agent("server", config=config, tools=[tool.ask_remote]))
queue = EventQueueLegacy()
running = asyncio.create_task(executor.execute(_client_tool_request(), queue))
draining: asyncio.Task[None] | None = None
events: list[Event] = []
try:
await asyncio.wait_for(tool.started.wait(), timeout=1.0)
child = await queue.tap(max_queue_size=1)
await child.enqueue_event(Message(message_id="prefill"))
tool.release.set()
while True:
event = await asyncio.wait_for(queue.dequeue_event(), timeout=1.0)
queue.task_done()
events.append(event)
if isinstance(event, TaskArtifactUpdateEvent):
break
assert tool.sender is not None
if cancel_sender:
tool.sender.cancel()
draining = asyncio.create_task(_drain_queue(child))
await asyncio.wait_for(running, timeout=1.0)
while not queue.queue.empty():
events.append(await queue.dequeue_event())
queue.task_done()
finally:
tool.release.set()
await queue.close(immediate=True)
tasks = [task for task in (running, draining, tool.sender) if task is not None]
for task in tasks:
if not task.done():
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
assert tool.recovered is cancel_sender
assert [e.status.state for e in events if isinstance(e, TaskStatusUpdateEvent)] == [
TaskState.TASK_STATE_WORKING,
TaskState.TASK_STATE_INPUT_REQUIRED,
]
assert len([e for e in events if isinstance(e, TaskArtifactUpdateEvent)]) == 1
@pytest.mark.parametrize("kind", ["text", "status"])
@pytest.mark.parametrize("cancel_count", [1, 2])
async def test_cancel_after_partial_enqueue_finishes_wire_commit(self, kind: str, cancel_count: int) -> None:
gate = TurnGate()
statuses = TaskStatusRecorder()
config = TestConfig(ToolCallEvent(name="hold", arguments="{}"), ModelMessageChunk("thinking"), "thinking")
executor = ObservedAgentExecutor(Agent("server", config=config, tools=[gate.hold]), statuses)
queue = EventQueueLegacy()
running = asyncio.create_task(executor.execute(_request(), queue))
canceling: asyncio.Task[None] | None = None
draining: asyncio.Task[None] | None = None
events: list[Event] = []
try:
await asyncio.wait_for(gate.entered.wait(), timeout=1.0)
child = await queue.tap(max_queue_size=1)
if kind == "text":
await child.enqueue_event(Message(message_id="prefill"))
gate.release()
while True:
event = await asyncio.wait_for(queue.dequeue_event(), timeout=1.0)
queue.task_done()
events.append(event)
if isinstance(event, TaskArtifactUpdateEvent):
if kind == "text":
break
if event.last_chunk:
await child.dequeue_event()
child.task_done()
if isinstance(event, TaskStatusUpdateEvent) and event.status.state == TaskState.TASK_STATE_COMPLETED:
break
if kind == "text":
canceling = asyncio.create_task(executor.cancel(_cancel_request(), queue))
await asyncio.sleep(0)
for _ in range(cancel_count):
running.cancel()
await asyncio.sleep(0)
draining = asyncio.create_task(_drain_queue(child))
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(running, timeout=1.0)
if canceling is not None:
await asyncio.wait_for(canceling, timeout=1.0)
expected = TaskState.TASK_STATE_CANCELED if kind == "text" else TaskState.TASK_STATE_COMPLETED
await asyncio.wait_for(statuses.wait_for(expected), timeout=1.0)
while not queue.queue.empty():
events.append(await queue.dequeue_event())
queue.task_done()
finally:
gate.release()
await queue.close(immediate=True)
tasks = [task for task in (running, canceling, draining) if task is not None]
for task in tasks:
if not task.done():
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
assert statuses.states == [TaskState.TASK_STATE_WORKING, expected]
assert [e.last_chunk for e in events if isinstance(e, TaskArtifactUpdateEvent)] == [False, True]
assert [e.status.state for e in events if isinstance(e, TaskStatusUpdateEvent)] == [
TaskState.TASK_STATE_WORKING,
expected,
]
@pytest.mark.parametrize("eager", [False, True], ids=["default-factory", "eager-factory"])
async def test_filtered_notification_does_not_stall_later_notifications(self, eager: bool) -> None:
eager_factory = getattr(asyncio, "eager_task_factory", None)
if eager and eager_factory is None:
pytest.skip("eager_task_factory requires Python 3.12+")
loop = asyncio.get_running_loop()
previous_factory = loop.get_task_factory()
queue = _RecordingQueue()
dropper = _DropAndCancel(TaskState.TASK_STATE_INPUT_REQUIRED, queue)
statuses = TaskStatusRecorder()
executor = ObservedAgentExecutor(
Agent("server", config=TestConfig(ToolCallEvent(name="lookup", arguments="{}"))),
observer(A2ATaskStatusUpdate, dropper.on_status, interrupt=True),
statuses,
)
dropper.executor = executor
try:
loop.set_task_factory(eager_factory if eager else None)
await asyncio.wait_for(executor.execute(_client_tool_request(), queue), timeout=1.0)
await asyncio.wait_for(asyncio.gather(*dropper.pending), timeout=1.0)
finally:
loop.set_task_factory(previous_factory)
assert queue.states() == [
TaskState.TASK_STATE_WORKING,
TaskState.TASK_STATE_INPUT_REQUIRED,
TaskState.TASK_STATE_CANCELED,
]
assert statuses.states == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_CANCELED]
@pytest.mark.parametrize("spawn_cancel", [False, True], ids=["inline", "child-task"])
@pytest.mark.parametrize("state", [TaskState.TASK_STATE_INPUT_REQUIRED, TaskState.TASK_STATE_COMPLETED])
async def test_lifecycle_observer_can_await_cancellation(self, state: TaskState, spawn_cancel: bool) -> None:
queue = _RecordingQueue()
canceler = _CancelOnStatus(state, queue, spawn_cancel=spawn_cancel)
statuses = TaskStatusRecorder()
answer = (
ToolCallEvent(name="lookup", arguments="{}") if state == TaskState.TASK_STATE_INPUT_REQUIRED else "done"
)
executor = ObservedAgentExecutor(
Agent("server", config=TestConfig(ModelMessageChunk("thinking"), answer)),
observer(A2ATaskStatusUpdate, canceler.on_status),
statuses,
)
canceler.executor = executor
await asyncio.wait_for(executor.execute(_client_tool_request(), queue), timeout=1.0)
expected = [TaskState.TASK_STATE_WORKING, state]
if state == TaskState.TASK_STATE_INPUT_REQUIRED:
expected.append(TaskState.TASK_STATE_CANCELED)
assert queue.states() == expected
assert canceler.observed == expected
assert statuses.states == expected
async def test_producer_cancellation_does_not_drop_queued_terminal_notification(self) -> None:
queue = _RecordingQueue()
pause = _PauseOnStatus(TaskState.TASK_STATE_INPUT_REQUIRED)
statuses = _ScopedStatuses()
executor = ObservedAgentExecutor(
Agent("server", config=TestConfig(ToolCallEvent(name="lookup", arguments="{}"))),
observer(A2ATaskStatusUpdate, pause.on_status),
statuses,
)
running = asyncio.create_task(executor.execute(_client_tool_request(), queue))
try:
await asyncio.wait_for(pause.paused.wait(), timeout=1.0)
await asyncio.wait_for(executor.cancel(_cancel_request(), queue), timeout=1.0)
running.cancel()
with pytest.raises(asyncio.CancelledError):
await running
assert not statuses.closed.is_set()
finally:
pause.resume()
if not running.done():
running.cancel()
await asyncio.gather(running, return_exceptions=True)
await asyncio.wait_for(statuses.wait_for(TaskState.TASK_STATE_CANCELED), timeout=1.0)
await asyncio.wait_for(statuses.closed.wait(), timeout=1.0)
assert statuses.states_at_close == statuses.states
assert statuses.states == [
TaskState.TASK_STATE_WORKING,
TaskState.TASK_STATE_INPUT_REQUIRED,
TaskState.TASK_STATE_CANCELED,
]
async def test_observer_failure_does_not_drop_queued_terminal_notification(
self, caplog: pytest.LogCaptureFixture
) -> None:
queue = _RecordingQueue()
canceler = _CancelOnStatus(
TaskState.TASK_STATE_INPUT_REQUIRED, queue, error=ValueError("observer failed after canceling")
)
statuses = TaskStatusRecorder()
executor = ObservedAgentExecutor(
Agent("server", config=TestConfig(ToolCallEvent(name="lookup", arguments="{}"))),
statuses,
observer(A2ATaskStatusUpdate, canceler.on_status),
)
canceler.executor = executor
# The transitions are on the wire, so the observer's failure is logged rather than raised.
await asyncio.wait_for(executor.execute(_client_tool_request(), queue), timeout=1.0)
assert "observer failed after canceling" in caplog.text
assert queue.states() == [
TaskState.TASK_STATE_WORKING,
TaskState.TASK_STATE_INPUT_REQUIRED,
TaskState.TASK_STATE_CANCELED,
]
assert statuses.states == canceler.observed == queue.states()
@pytest.mark.parametrize("kind", ["text", "tool-call", "custom-artifact"])
async def test_cancel_waits_for_an_inflight_artifact(self, kind: str) -> None:
if kind == "text":
event = ModelMessageChunk("thinking")
elif kind == "tool-call":
event = ToolCallEvent(name="lookup", arguments="{}")
else:
event = chunk_to_text_artifact(
ModelMessageChunk("custom"),
artifact_id="custom",
task_id=_TASK_ID,
context_id=_CONTEXT_ID,
append=False,
)
executor = AgentExecutor(Agent("server", config=TestConfig(event, "done")))
queue = _PausingArtifactQueue()
running = asyncio.create_task(executor.execute(_client_tool_request(), queue))
canceling: asyncio.Task[None] | None = None
try:
await asyncio.wait_for(queue.started.wait(), timeout=1.0)
canceling = asyncio.create_task(executor.cancel(_cancel_request(), queue))
await asyncio.sleep(0)
queue.release.set()
await asyncio.wait_for(asyncio.gather(running, canceling), timeout=1.0)
finally:
queue.release.set()
pending = [task for task in (running, canceling) if task is not None and not task.done()]
for task in pending:
task.cancel()
await asyncio.gather(*pending, return_exceptions=True)
assert queue.states() == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_CANCELED]
assert isinstance(queue.events[-1], TaskStatusUpdateEvent)
async def test_cancel_landing_after_the_turn_completed_leaves_it_completed(self) -> None:
executor = AgentExecutor(Agent("server", config=TestConfig("done")))
queue = _RecordingQueue()
cancels: list[asyncio.Task[None]] = []
async def cancel_on_completion(event: Event) -> None:
if isinstance(event, TaskStatusUpdateEvent) and event.status.state == TaskState.TASK_STATE_COMPLETED:
cancels.append(asyncio.create_task(executor.cancel(_cancel_request(), queue)))
await asyncio.sleep(0)
queue.on_event = cancel_on_completion
await executor.execute(_request(), queue)
await asyncio.gather(*cancels)
assert queue.states() == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_COMPLETED]
async def test_cancel_after_the_turn_finished_writes_nothing(self) -> None:
# No turn is running here; the SDK closes a still-open task as CANCELED itself.
executor = AgentExecutor(Agent("server", config=TestConfig("done")))
queue = _RecordingQueue()
await executor.execute(_request(), queue)
finished = len(queue.events)
await executor.cancel(_cancel_request(), queue)
assert len(queue.events) == finished
async def test_turn_finishing_after_a_cancel_ends_quietly(self) -> None:
gate = TurnGate()
statuses = TaskStatusRecorder()
config = TestConfig(ToolCallEvent(name="hold", arguments="{}"), "done")
executor = ObservedAgentExecutor(Agent("server", config=config, tools=[gate.hold]), statuses)
queue = _RecordingQueue()
running = asyncio.create_task(executor.execute(_request(), queue))
await gate.entered.wait()
await executor.cancel(_cancel_request(), queue)
gate.release()
await running
assert queue.states() == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_CANCELED]
# Server-side observers see the same lifecycle as the wire, not the completion that lost.
assert statuses.states == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_CANCELED]
async def test_cancel_landing_during_the_text_close_follows_input_required(self) -> None:
# Asking for a client tool closes the streamed text before the task leaves ``working``; a
# cancel arriving during that close waits for the input request, then cancels the task.
config = TestConfig(ModelMessageChunk("thinking"), ToolCallEvent(name="lookup", arguments="{}"))
executor = AgentExecutor(Agent("server", config=config))
queue = _RecordingQueue()
cancels: list[asyncio.Task[None]] = []
async def cancel_during_close(event: Event) -> None:
if isinstance(event, TaskArtifactUpdateEvent) and event.append and event.last_chunk:
cancels.append(asyncio.create_task(executor.cancel(_cancel_request(), queue)))
await asyncio.sleep(0)
queue.on_event = cancel_during_close
await executor.execute(_client_tool_request(), queue)
await asyncio.gather(*cancels)
# The text is closed once, before the task leaves ``working``.
assert queue.lifecycle() == [
TaskState.TASK_STATE_WORKING,
_TEXT_CLOSED,
TaskState.TASK_STATE_INPUT_REQUIRED,
TaskState.TASK_STATE_CANCELED,
]
@pytest.mark.asyncio
@pytest.mark.parametrize("client_tool", [False, True], ids=["completed", "input-required"])
async def test_request_observers_close_after_the_last_notification(client_tool: bool) -> None:
statuses = _ScopedStatuses()
answer = ToolCallEvent(name="lookup", arguments="{}") if client_tool else "done"
executor = ObservedAgentExecutor(Agent("server", config=TestConfig(answer)), statuses)
queue = _RecordingQueue()
await executor.execute(_client_tool_request(), queue)
assert statuses.closed.is_set()
assert statuses.states_at_close == queue.states()
@pytest.mark.asyncio
async def test_observers_close_when_a_later_registration_fails() -> None:
statuses = _ScopedStatuses()
executor = ObservedAgentExecutor(Agent("server", config=TestConfig("done")), statuses, _FailingObserver())
with pytest.raises(ValueError, match="registration failed"):
await executor.execute(_request(), _RecordingQueue())
assert statuses.closed.is_set()
class _FailingObserver:
def register(self, stack: ExitStack | AsyncExitStack, context: ConversationContext) -> None:
raise ValueError("registration failed")
@pytest.mark.asyncio
async def test_observers_close_after_a_turn_failure() -> None:
statuses = _ScopedStatuses()
executor = ObservedAgentExecutor(Agent("unconfigured"), statuses)
queue = _RecordingQueue()
with pytest.raises(RuntimeError, match="Agent.config is not set"):
await executor.execute(_request(), queue)
assert statuses.closed.is_set()
assert statuses.states_at_close == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_FAILED]
@pytest.mark.asyncio
async def test_cancelled_startup_keeps_observers_until_working_is_delivered() -> None:
statuses = _ScopedStatuses()
pause = _PauseOnStatus(TaskState.TASK_STATE_WORKING)
executor = ObservedAgentExecutor(
Agent("server", config=TestConfig("done")), observer(A2ATaskStatusUpdate, pause.on_status), statuses
)
running = asyncio.create_task(executor.execute(_request(), _RecordingQueue()))
try:
await asyncio.wait_for(pause.paused.wait(), timeout=1.0)
running.cancel()
with pytest.raises(asyncio.CancelledError):
await running
assert not statuses.closed.is_set()
finally:
pause.resume()
if not running.done():
running.cancel()
await asyncio.gather(running, return_exceptions=True)
await asyncio.wait_for(statuses.closed.wait(), timeout=1.0)
assert statuses.states_at_close == [TaskState.TASK_STATE_WORKING]
class _FailingCleanupObserver:
def __init__(self, error: BaseException) -> None:
self.error = error
def register(self, stack: ExitStack | AsyncExitStack, context: ConversationContext) -> None:
stack.callback(self.close)
def close(self) -> None:
raise self.error
@pytest.mark.asyncio
@pytest.mark.parametrize("error", [ValueError("cleanup failed"), asyncio.CancelledError("cleanup cancelled")])
async def test_observer_cleanup_failure_does_not_fail_the_completed_turn(
error: BaseException, caplog: pytest.LogCaptureFixture
) -> None:
statuses = _ScopedStatuses()
executor = ObservedAgentExecutor(
Agent("server", config=TestConfig("done")), statuses, _FailingCleanupObserver(error)
)
queue = _RecordingQueue()
with caplog.at_level(logging.ERROR, logger="ag2.a2a.executor"):
await executor.execute(_request(), queue)
assert statuses.closed.is_set()
assert statuses.states_at_close == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_COMPLETED]
assert queue.states() == statuses.states_at_close
assert "A2A lifecycle subscriber cleanup failed" in caplog.text
async def leave_note(context: Context) -> str:
"""Server tool that puts its own artifact on the wire, passed through by the executor."""
note = Artifact(artifact_id="note", parts=[Part(text="N")])
await context.send(
A2ATaskArtifactUpdate(update=TaskArtifactUpdateEvent(task_id=_TASK_ID, context_id=_CONTEXT_ID, artifact=note))
)
return "noted"
def _describe(event: Event) -> str:
if isinstance(event, Task):
return "task"
if isinstance(event, TaskStatusUpdateEvent):
return TaskState.Name(event.status.state)
assert isinstance(event, TaskArtifactUpdateEvent)
if event.artifact.artifact_id == "note":
return "note:" + "".join(part.text for part in event.artifact.parts)
if any(is_data_part_with_mime(part, MIME_TOOL_CALL) for part in event.artifact.parts):
return "client-call"
text = "".join(part.text for part in event.artifact.parts)
return f"text:{text}{':last' if event.last_chunk else ''}"
@pytest.mark.asyncio
async def test_wire_order_is_the_order_the_turn_produced_its_output() -> None:
config = TestConfig(
ModelMessageChunk("a"),
ToolCallEvent(name="leave_note", arguments="{}"),
ModelMessageChunk("b"),
ToolCallEvent(name="lookup", arguments="{}"),
)
executor = AgentExecutor(Agent("server", config=config, tools=[leave_note]))
queue = _RecordingQueue()
await asyncio.wait_for(executor.execute(_client_tool_request(), queue), timeout=1.0)
assert [_describe(event) for event in queue.events] == [
"task",
"TASK_STATE_WORKING",
"text:a",
"note:N",
"text:b",
"client-call",
"text::last",
"TASK_STATE_INPUT_REQUIRED",
]
@pytest.mark.asyncio
async def test_transitions_do_not_wait_for_a_slow_observer() -> None:
gate = TurnGate()
pause = _PauseOnStatus(TaskState.TASK_STATE_WORKING)
statuses = TaskStatusRecorder()
config = TestConfig(ModelMessageChunk("thinking"), ToolCallEvent(name="hold", arguments="{}"), "done")
executor = ObservedAgentExecutor(
Agent("server", config=config, tools=[gate.hold]), observer(A2ATaskStatusUpdate, pause.on_status), statuses
)
queue = _RecordingQueue()
running = asyncio.create_task(executor.execute(_request(), queue))
try:
# The observer is still holding ``working`` while the turn streams and reaches its tool.
await asyncio.wait_for(gate.entered.wait(), timeout=1.0)
assert [_describe(event) for event in queue.events] == ["task", "TASK_STATE_WORKING", "text:thinking"]
await asyncio.wait_for(executor.cancel(_cancel_request(), queue), timeout=1.0)
assert queue.states() == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_CANCELED]
finally:
pause.resume()
gate.release()
await asyncio.wait_for(running, timeout=1.0)
# Delivery caught up, in wire order.
await asyncio.wait_for(statuses.wait_for(TaskState.TASK_STATE_CANCELED), timeout=1.0)
assert statuses.states == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_CANCELED]
async def _close_the_loop_while_an_observer_holds_working() -> None:
pause = _PauseOnStatus(TaskState.TASK_STATE_WORKING)
executor = ObservedAgentExecutor(
Agent("server", config=TestConfig("done")), observer(A2ATaskStatusUpdate, pause.on_status)
)
running = asyncio.create_task(executor.execute(_request(), _RecordingQueue()))
await pause.paused.wait()
running.cancel()
with pytest.raises(asyncio.CancelledError):
await running
# Returning leaves delivery pending; ``asyncio.run`` cancels it and waits for it to stop.
async def _close_the_loop_while_a_cancel_waits_behind_a_blocked_write() -> None:
queue = _PausingArtifactQueue()
executor = AgentExecutor(Agent("server", config=TestConfig(ModelMessageChunk("a"), "done")))
running = asyncio.create_task(executor.execute(_request(), queue))
await queue.started.wait()
cancelling = asyncio.create_task(executor.cancel(_cancel_request(), queue))
# One step queues the cancel behind the blocked chunk; returning then stops the writer under it.
await asyncio.sleep(0)
assert not running.done()
assert not cancelling.done()
class _LoopRun:
"""Runs a scenario under its own ``asyncio.run``, which returns only once every task it leaves has stopped."""
def __init__(self, scenario: Callable[[], Coroutine[Any, Any, None]]) -> None:
self.scenario = scenario
self.finished = threading.Event()
self.error: BaseException | None = None
def run(self) -> None:
try:
asyncio.run(self.scenario())
except BaseException as exc:
self.error = exc
self.finished.set()
@pytest.mark.parametrize(
"scenario",
[_close_the_loop_while_an_observer_holds_working, _close_the_loop_while_a_cancel_waits_behind_a_blocked_write],
ids=["observer-holds-delivery", "cancel-queued-behind-blocked-write"],
)
def test_closing_the_event_loop_stops_the_request(scenario: Callable[[], Coroutine[Any, Any, None]]) -> None:
loop_run = _LoopRun(scenario)
thread = threading.Thread(target=loop_run.run, daemon=True)
thread.start()
thread.join(timeout=5.0)
assert loop_run.finished.is_set()
assert loop_run.error is None
async def _refuse_the_closing_chunk(event: Event) -> None:
if isinstance(event, TaskArtifactUpdateEvent) and event.last_chunk:
raise ConnectionError("queue refused the closing chunk")
@pytest.mark.asyncio
async def test_task_whose_closing_chunk_cannot_be_written_ends_failed() -> None:
statuses = TaskStatusRecorder()
executor = ObservedAgentExecutor(Agent("server", config=TestConfig(ModelMessageChunk("a"), "a")), statuses)
queue = _RecordingQueue()
queue.on_event = _refuse_the_closing_chunk
with pytest.raises(ConnectionError, match="closing chunk"):
await executor.execute(_request(), queue)
assert queue.states() == [TaskState.TASK_STATE_WORKING, TaskState.TASK_STATE_FAILED]
assert statuses.states == queue.states()