From 6a9a97dd4dc398b7db83528b9836c7e798ba022b Mon Sep 17 00:00:00 2001 From: boulea7 Date: Sat, 24 Oct 2026 05:50:52 +0800 Subject: [PATCH] feat(server): make V2 event queue capacity configurable --- src/a2a/server/agent_execution/active_task.py | 57 ++- .../agent_execution/active_task_registry.py | 11 +- .../default_request_handler_v2.py | 22 +- .../agent_execution/test_active_task.py | 68 ++++ .../test_active_task_registry.py | 12 + tests/server/cluster/conftest.py | 4 +- .../cluster/test_cancel_multireplica.py | 12 +- .../server/cluster/test_send_multireplica.py | 7 +- .../cluster/test_subscribe_multireplica.py | 17 +- .../test_event_queue_capacity.py | 343 ++++++++++++++++++ .../test_evicted_message_response.py | 243 +++++++++++++ 11 files changed, 781 insertions(+), 15 deletions(-) create mode 100644 tests/server/request_handlers/test_event_queue_capacity.py create mode 100644 tests/server/request_handlers/test_evicted_message_response.py diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index e02ff3315..352edc003 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -60,6 +60,7 @@ from a2a.server.cluster.event_stream import VersionedEvent from a2a.server.cluster.task_store import ConcurrentTaskModificationError from a2a.server.events.event_queue_v2 import ( + DEFAULT_MAX_QUEUE_SIZE, AsyncQueue, Event, EventQueueSource, @@ -140,6 +141,9 @@ async def run(self) -> None: await self._process_event(event) + except asyncio.CancelledError as error: + self.active_task._consumer_error = error + raise except QueueShutDown: logger.debug( 'Consumer[%s]: Event queue shut down', @@ -152,16 +156,25 @@ async def run(self) -> None: if isinstance(e, ConcurrentTaskModificationError): # Other writers kept advancing the task; report a protocol error. error = InternalError(message=str(e)) + self.active_task._consumer_error = error updated_task = None try: updated_task = await self._write_failed_status() + except Exception: + logger.exception( + 'Consumer[%s]: Failed to persist failure status', + self.active_task._task_id, + ) finally: # Subscribers must always hear about the failure, or the # request waiting on them never returns. await self._enqueue_to_subscribers( cast('Event', error), updated_task ) + finally: + # Validation ends before _run_consumer waits for producer cleanup. + self.active_task._response_validation_done.set() async def _write_failed_status(self) -> Task | None: """Records FAILED once a failure ends the event stream. @@ -295,6 +308,7 @@ def _handle_message_event(self, event: Message) -> None: 'Multiple Message objects received.' ) self.task_mode = False + self.active_task._message_result = event async def _handle_task_event( self, @@ -476,6 +490,8 @@ def __init__( # noqa: PLR0913 push_sender: PushNotificationSender | None = None, on_cleanup: Callable[[ActiveTask], None] | None = None, event_stream: TaskEventStream | None = None, + *, + max_event_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, ) -> None: """Initializes the ActiveTask. @@ -489,13 +505,21 @@ def __init__( # noqa: PLR0913 the task from the ActiveTaskRegistry. event_stream: Optional cross-replica stream; applied events are published to it so other replicas can observe them. + max_event_queue_size: Positive event capacity for each local event + buffer, including internal lifecycle markers. Applies + to both sources, the agent's default sink and each + subscriber tap, not the pending-request queue or the + cross-replica stream. """ # --- Core Dependencies --- self._agent_executor = agent_executor self._task_id = task_id - self._event_queue_agent = EventQueueSource() + self._max_event_queue_size = max_event_queue_size + self._event_queue_agent = EventQueueSource( + max_queue_size=max_event_queue_size + ) self._event_queue_subscribers = EventQueueSource( - create_default_sink=False + max_queue_size=max_event_queue_size, create_default_sink=False ) self._task_manager = task_manager self._push_sender = push_sender @@ -512,6 +536,7 @@ def __init__( # noqa: PLR0913 # _task_created is set when initial version of task is stored in DB. self._task_created = asyncio.Event() + self._response_validation_done = asyncio.Event() # `_is_finished` is set EXACTLY ONCE when the consumer loop exits, signifying # the absolute end of the task's active lifecycle. @@ -522,6 +547,8 @@ def __init__( # noqa: PLR0913 self._producer_task: asyncio.Task[None] | None = None # The background task reading from _event_queue and updating the DB. self._consumer_task: asyncio.Task[None] | None = None + self._message_result: Message | None = None + self._consumer_error: BaseException | None = None # Tracks how many active SSE/gRPC streams are currently tailing this task. # Protected by `_lock`. @@ -768,7 +795,7 @@ async def subscribe( # evict_on_full=True: a subscriber that falls behind is dropped, so it # cannot wedge dispatch for the others. See _deliver_to_sink(). tapped_queue = await self._event_queue_subscribers.tap( - evict_on_full=True + max_queue_size=self._max_event_queue_size, evict_on_full=True ) request_id = await self.enqueue_request(request) if request else None @@ -955,6 +982,13 @@ async def aclose(self) -> None: background tasks are cancelled, so teardown is bounded even when a subscriber sink was never drained. It is safe to call multiple times. """ + # Immediate queue closure can discard events before validation ends. + if ( + not self._response_validation_done.is_set() + and self._consumer_error is None + ): + self._consumer_error = asyncio.CancelledError() + # Shut down the request queue first, mirroring the producer's # ``finally``. If `start()` was never called, the producer/consumer # ``finally`` blocks never run, and without this a caller parked in @@ -1027,6 +1061,23 @@ async def _mark_task_as_failed(self, exception: Exception) -> Task | None: pass return task + async def get_result(self) -> Task | Message: + """Returns a task snapshot or a fully validated direct Message result. + + An evicted subscriber may miss the Message or a validation error. Wait + for Message validation, independently of producer cleanup, so later + invalid events or executor failures still take precedence. Cancelling + this waiter does not cancel the shared consumer. + """ + message = self._message_result + if message is not None: + await self._response_validation_done.wait() + if self._consumer_error is not None: + raise self._consumer_error + if message is not None: + return message + return await self.get_task() + async def get_task(self) -> Task: """Get task from db.""" await self._task_created.wait() diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index cae2a55ce..7d42b8dbb 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -18,6 +18,7 @@ from a2a.server.agent_execution.active_task import ActiveTask from a2a.server.cluster.task_store import StoredTask +from a2a.server.events.event_queue import DEFAULT_MAX_QUEUE_SIZE from a2a.server.tasks.task_manager import TaskManager from a2a.utils.errors import TaskNotFoundError @@ -34,11 +35,18 @@ def __init__( task_store: TaskStore | VersionedTaskStore, push_sender: PushNotificationSender | None = None, event_stream: TaskEventStream | None = None, - ): + *, + max_event_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, + ) -> None: + """Initializes a registry with a positive per-buffer event capacity.""" + if max_event_queue_size <= 0: + raise ValueError('max_event_queue_size must be greater than 0') + self._agent_executor = agent_executor self._task_store = task_store self._push_sender = push_sender self._event_stream = event_stream + self._max_event_queue_size = max_event_queue_size self._active_tasks: dict[str, ActiveTask] = {} self._lock = threading.RLock() self._cleanup_tasks: set[asyncio.Task[None]] = set() @@ -73,6 +81,7 @@ async def get_or_create( push_sender=self._push_sender, on_cleanup=self._on_active_task_cleanup, event_stream=self._event_stream, + max_event_queue_size=self._max_event_queue_size, ) self._active_tasks[task_id] = active_task diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index b135f529f..b749a98e4 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -24,6 +24,7 @@ LegacyTaskStoreAdapter, VersionedTaskStore, ) +from a2a.server.events.event_queue import DEFAULT_MAX_QUEUE_SIZE from a2a.server.request_handlers.request_handler import ( RequestHandler, validate, @@ -96,6 +97,22 @@ class DefaultRequestHandlerV2(RequestHandler): with `DefaultRequestHandler` but is not used: v2 delegates event streaming to an in-memory `ActiveTaskRegistry`. Passing a non-``None`` value emits a `DeprecationWarning` and logs a warning. + + ``max_event_queue_size`` sets the positive event capacity of each local + event buffer (default 1024): the two source inputs, the agent's default + sink and one sink per local subscription. Internal request-start and + request-completion markers also occupy slots. This is a per-buffer event + count, not a task memory limit; pending requests, shared storage and remote + ``TaskEventStream`` buffers are outside its scope. + + A full default sink backpressures the agent when the internal consumer + falls behind. Full subscriber sinks are evicted, including live but slow + consumers and the subscriptions used by send and streaming requests. + Smaller capacities can end a subscription sooner; a non-streaming send + whose subscription is evicted can return a non-terminal stored task + snapshot rather than waiting for the agent to finish. Direct Message + responses remain available after eviction once the consumer has validated + the full response; later validation and executor errors take precedence. """ _background_tasks: set[asyncio.Task] @@ -118,6 +135,8 @@ def __init__( # noqa: PLR0913 push_url_validator: Callable[[str], Awaitable[bool]] | None = None, validate_input_modes: bool = False, event_stream: TaskEventStream | None = None, + *, + max_event_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, ) -> None: if queue_manager is not None: message = ( @@ -172,6 +191,7 @@ def __init__( # noqa: PLR0913 task_store=self._versioned_store, push_sender=self._push_sender, event_stream=self._event_stream, + max_event_queue_size=max_event_queue_size, ) self._background_tasks = set() @@ -445,7 +465,7 @@ async def on_message_send( # noqa: D102 if result is None: logger.debug('Missing result for task %s', request_context.task_id) - result = await active_task.get_task() + result = await active_task.get_result() if isinstance(result, Task): result = apply_history_length(result, params.configuration) diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 254423e18..d86510cf8 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -2,6 +2,7 @@ import contextlib import logging +from typing import cast from unittest.mock import AsyncMock, Mock, patch import pytest @@ -11,6 +12,7 @@ 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 import Event from a2a.server.events.event_queue_v2 import EventQueueSource as EventQueue from a2a.server.tasks.push_notification_sender import PushNotificationSender from a2a.server.tasks.task_manager import TaskManager @@ -1336,3 +1338,69 @@ async def slow_execute(req, q): for t in background: t.cancel() await asyncio.gather(*background, return_exceptions=True) + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +@pytest.mark.parametrize('max_event_queue_size', [1, 2]) +async def test_small_event_queue_capacity_evicts_slow_subscriber( + max_event_queue_size: int, +) -> None: + task_manager = Mock(spec=TaskManager) + task_manager.get_task = AsyncMock( + return_value=Task( + id='test-task-id', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + active_task = ActiveTask( + agent_executor=Mock(spec=AgentExecutor), + task_id='test-task-id', + task_manager=task_manager, + max_event_queue_size=max_event_queue_size, + ) + active_task._task_created.set() + source = active_task._event_queue_subscribers + slow = active_task.subscribe(include_initial_task=True) + fast = active_task.subscribe(include_initial_task=True) + subscriptions = [slow, fast] + readers: list[asyncio.Task[Event]] = [] + try: + assert isinstance(await anext(slow), Task) + slow_sink = next(iter(source._sinks)) + assert isinstance(await anext(fast), Task) + + for index in range(max_event_queue_size + 2): + reader = asyncio.create_task(anext(fast)) + readers.append(reader) + event = Task(id=f'event-{index}') + await source.enqueue_event(cast('Event', (event, None))) + await source.test_only_join_incoming_queue() + assert await reader == event + assert slow_sink.is_closed() == (index >= max_event_queue_size) + + assert slow_sink not in source._sinks + assert len(source._sinks) == 1 + + late = active_task.subscribe(include_initial_task=True) + subscriptions.append(late) + assert isinstance(await anext(late), Task) + fast_reader = asyncio.create_task(anext(fast)) + late_reader = asyncio.create_task(anext(late)) + readers.extend((fast_reader, late_reader)) + event = Task(id='after-eviction') + await source.enqueue_event(cast('Event', (event, None))) + await source.test_only_join_incoming_queue() + assert await fast_reader == event + assert await late_reader == event + finally: + for reader in readers: + if not reader.done(): + reader.cancel() + await asyncio.gather(*readers, return_exceptions=True) + for subscription in subscriptions: + await subscription.aclose() + await active_task.aclose() + + assert not source._sinks + assert source._dispatcher_task.done() diff --git a/tests/server/agent_execution/test_active_task_registry.py b/tests/server/agent_execution/test_active_task_registry.py index 7f0e3e46b..083a5ae41 100644 --- a/tests/server/agent_execution/test_active_task_registry.py +++ b/tests/server/agent_execution/test_active_task_registry.py @@ -37,6 +37,18 @@ def _make_registry() -> ActiveTaskRegistry: ) +@pytest.mark.parametrize('max_event_queue_size', [0, -1]) +def test_registry_rejects_nonpositive_event_queue_size( + max_event_queue_size: int, +) -> None: + with pytest.raises(ValueError, match='max_event_queue_size'): + ActiveTaskRegistry( + agent_executor=_SlowExecutor(), + task_store=InMemoryTaskStore(), + max_event_queue_size=max_event_queue_size, + ) + + @pytest.mark.timeout(5) @pytest.mark.asyncio async def test_aclose_reaps_active_tasks_and_empties_registry(): diff --git a/tests/server/cluster/conftest.py b/tests/server/cluster/conftest.py index 97da1eaf6..9f9a65e8c 100644 --- a/tests/server/cluster/conftest.py +++ b/tests/server/cluster/conftest.py @@ -18,7 +18,7 @@ VersionedTaskStore, ) from a2a.server.context import ServerCallContext -from a2a.server.events.event_queue import Event +from a2a.server.events.event_queue import DEFAULT_MAX_QUEUE_SIZE, Event from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope from a2a.server.request_handlers.default_request_handler_v2 import ( DefaultRequestHandlerV2, @@ -390,6 +390,7 @@ def make_replica( store: VersionedTaskStore, stream: TaskEventStream, agent: AgentExecutor, + max_event_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, ) -> DefaultRequestHandlerV2: """Builds a handler ('replica') wired to a shared store and stream.""" return DefaultRequestHandlerV2( @@ -397,6 +398,7 @@ def make_replica( task_store=store, agent_card=streaming_agent_card(), event_stream=stream, + max_event_queue_size=max_event_queue_size, ) diff --git a/tests/server/cluster/test_cancel_multireplica.py b/tests/server/cluster/test_cancel_multireplica.py index bc270d50f..4fb29e77c 100644 --- a/tests/server/cluster/test_cancel_multireplica.py +++ b/tests/server/cluster/test_cancel_multireplica.py @@ -24,12 +24,17 @@ @pytest.mark.asyncio @pytest.mark.timeout(15) +@pytest.mark.parametrize('max_event_queue_size', [1024, 1]) async def test_cancel_on_non_owning_replica_stops_remote_agent( - shared_store, shared_stream, context + shared_store, shared_stream, context, max_event_queue_size: int ) -> None: agent = LongRunningAgent() - replica_a = make_replica(shared_store, shared_stream, agent) - replica_b = make_replica(shared_store, shared_stream, agent) + replica_a = make_replica( + shared_store, shared_stream, agent, max_event_queue_size + ) + replica_b = make_replica( + shared_store, shared_stream, agent, max_event_queue_size + ) try: a_task = asyncio.create_task( drain( @@ -52,6 +57,7 @@ async def test_cancel_on_non_owning_replica_stops_remote_agent( CancelTaskRequest(id=task_id), context ) assert result.status.state == TaskState.TASK_STATE_CANCELED + assert await replica_b._active_task_registry.get(task_id) is None # noqa: SLF001 # A's agent observes the CAS conflict and aborts; its stream ends. await asyncio.wait_for(agent.aborted.wait(), timeout=8) diff --git a/tests/server/cluster/test_send_multireplica.py b/tests/server/cluster/test_send_multireplica.py index fabd02dca..bc675e23a 100644 --- a/tests/server/cluster/test_send_multireplica.py +++ b/tests/server/cluster/test_send_multireplica.py @@ -17,10 +17,13 @@ @pytest.mark.asyncio @pytest.mark.timeout(10) +@pytest.mark.parametrize('max_event_queue_size', [1024, 1]) async def test_single_send_completes_on_shared_store( - shared_store, shared_stream, context + shared_store, shared_stream, context, max_event_queue_size: int ) -> None: - handler = make_replica(shared_store, shared_stream, CompletingAgent()) + handler = make_replica( + shared_store, shared_stream, CompletingAgent(), max_event_queue_size + ) try: result = await handler.on_message_send( build_send_request('hi'), context diff --git a/tests/server/cluster/test_subscribe_multireplica.py b/tests/server/cluster/test_subscribe_multireplica.py index 4cc04c8bb..25cae2fbb 100644 --- a/tests/server/cluster/test_subscribe_multireplica.py +++ b/tests/server/cluster/test_subscribe_multireplica.py @@ -16,16 +16,21 @@ @pytest.mark.asyncio @pytest.mark.timeout(15) +@pytest.mark.parametrize('max_event_queue_size', [1024, 1]) async def test_resubscribe_on_non_owning_replica_streams_events( - shared_store, shared_stream, context + shared_store, shared_stream, context, max_event_queue_size: int ) -> None: """Replica A runs the agent; a resubscribe on replica B streams its events. Pre-fix, B would build a local ActiveTask, emit the snapshot, then hang. """ agent = ControlledAgent() - replica_a = make_replica(shared_store, shared_stream, agent) - replica_b = make_replica(shared_store, shared_stream, agent) + replica_a = make_replica( + shared_store, shared_stream, agent, max_event_queue_size + ) + replica_b = make_replica( + shared_store, shared_stream, agent, max_event_queue_size + ) try: a_events: list = [] @@ -47,6 +52,7 @@ async def run_a() -> None: # Resubscribe on B (not running the agent). Collect until COMPLETED. b_states: list = [] + snapshot_received = asyncio.Event() async def run_b() -> None: async for ev in replica_b.on_subscribe_to_task( @@ -54,11 +60,14 @@ async def run_b() -> None: ): if getattr(ev, 'status', None): b_states.append(ev.status.state) + if len(b_states) == 1: + snapshot_received.set() if ev.status.state == TaskState.TASK_STATE_COMPLETED: return b_task = asyncio.create_task(run_b()) - await asyncio.sleep(0.1) # let B read snapshot + start tailing + await asyncio.wait_for(snapshot_received.wait(), timeout=5) + assert await replica_b._active_task_registry.get(task_id) is None # noqa: SLF001 # Release the agent on A -> it completes; B must observe via the stream. agent.release.set() diff --git a/tests/server/request_handlers/test_event_queue_capacity.py b/tests/server/request_handlers/test_event_queue_capacity.py new file mode 100644 index 000000000..08291d393 --- /dev/null +++ b/tests/server/request_handlers/test_event_queue_capacity.py @@ -0,0 +1,343 @@ +import asyncio + +from collections.abc import AsyncGenerator +from typing import cast + +import pytest + +from a2a.server.agent_execution import AgentExecutor, RequestContext +from a2a.server.context import ServerCallContext +from a2a.server.events.event_queue import ( + DEFAULT_MAX_QUEUE_SIZE, + Event, + EventQueue, +) +from a2a.server.events.event_queue_v2 import EventQueueSink, EventQueueSource +from a2a.server.request_handlers import ( + DefaultRequestHandler, + DefaultRequestHandlerV2, +) +from a2a.server.tasks import InMemoryTaskStore +from a2a.types.a2a_pb2 import ( + AgentCapabilities, + AgentCard, + CancelTaskRequest, + Message, + Part, + Role, + SendMessageRequest, + SubscribeToTaskRequest, + Task, + TaskState, + TaskStatus, + TaskStatusUpdateEvent, +) + + +class _GatedExecutor(AgentExecutor): + def __init__(self) -> None: + self.release = asyncio.Event() + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + await event_queue.enqueue_event( + Task( + id=context.task_id or '', + context_id=context.context_id or '', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + await self.release.wait() + await event_queue.enqueue_event( + TaskStatusUpdateEvent( + task_id=context.task_id or '', + context_id=context.context_id or '', + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED), + ) + ) + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + return None + + +def _agent_card() -> AgentCard: + return AgentCard( + name='queue-capacity-test', + version='1.0', + capabilities=AgentCapabilities(streaming=True), + ) + + +def _send_request() -> SendMessageRequest: + return SendMessageRequest( + message=Message( + message_id='message-1', + role=Role.ROLE_USER, + parts=[Part(text='start')], + ) + ) + + +@pytest.mark.parametrize('max_event_queue_size', [0, -1]) +def test_handler_rejects_nonpositive_event_queue_size( + max_event_queue_size: int, +) -> None: + with pytest.raises(ValueError, match='max_event_queue_size'): + DefaultRequestHandlerV2( + agent_executor=_GatedExecutor(), + task_store=InMemoryTaskStore(), + agent_card=_agent_card(), + max_event_queue_size=max_event_queue_size, + ) + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize('max_event_queue_size', [None, 1, 3]) +async def test_handler_event_queue_capacity_reaches_all_local_buffers( + max_event_queue_size: int | None, +) -> None: + kwargs = ( + {} + if max_event_queue_size is None + else {'max_event_queue_size': max_event_queue_size} + ) + expected = ( + DEFAULT_MAX_QUEUE_SIZE + if max_event_queue_size is None + else max_event_queue_size + ) + executor = _GatedExecutor() + handler = DefaultRequestHandler( + agent_executor=executor, + task_store=InMemoryTaskStore(), + agent_card=_agent_card(), + **kwargs, + ) + context = ServerCallContext() + stream = handler.on_message_send_stream(_send_request(), context) + subscriptions: list[AsyncGenerator[Event, None]] = [] + reader: asyncio.Task[Event] | None = None + try: + first = await anext(stream) + assert isinstance(first, Task) + active = await handler._active_task_registry.get(first.id) + assert active is not None + + for _ in range(2): + subscription = handler.on_subscribe_to_task( + SubscribeToTaskRequest(id=first.id), context + ) + subscriptions.append(subscription) + assert isinstance(await anext(subscription), Task) + + agent_source = active._event_queue_agent + subscriber_source = active._event_queue_subscribers + assert DEFAULT_MAX_QUEUE_SIZE == 1024 + assert agent_source._incoming_queue.maxsize == expected + assert agent_source.queue.maxsize == expected + assert subscriber_source._incoming_queue.maxsize == expected + assert subscriber_source._default_sink is None + assert len(subscriber_source._sinks) == 3 + assert all( + sink.queue.maxsize == expected for sink in subscriber_source._sinks + ) + + reader = asyncio.create_task(anext(stream)) + cancelled = await handler.on_cancel_task( + CancelTaskRequest(id=first.id), context + ) + assert cancelled is not None + assert cancelled.status.state == TaskState.TASK_STATE_CANCELED + with pytest.raises(StopAsyncIteration): + await reader + finally: + if reader is not None: + if not reader.done(): + reader.cancel() + await asyncio.gather(reader, return_exceptions=True) + for subscription in subscriptions: + await subscription.aclose() + await stream.aclose() + await handler.aclose() + + assert active._producer_task is not None + assert active._producer_task.done() + assert active._consumer_task is not None + assert active._consumer_task.done() + assert agent_source._dispatcher_task.done() + assert subscriber_source._dispatcher_task.done() + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize('streaming', [False, True]) +async def test_small_event_queue_capacity_completes_requests( + streaming: bool, +) -> None: + executor = _GatedExecutor() + executor.release.set() + handler = DefaultRequestHandlerV2( + agent_executor=executor, + task_store=InMemoryTaskStore(), + agent_card=_agent_card(), + max_event_queue_size=1, + ) + try: + if streaming: + events = [ + event + async for event in handler.on_message_send_stream( + _send_request(), ServerCallContext() + ) + ] + assert isinstance(events[0], Task) + result = events[-1] + assert isinstance(result, TaskStatusUpdateEvent) + else: + result = await handler.on_message_send( + _send_request(), ServerCallContext() + ) + assert isinstance(result, Task) + assert result.status.state == TaskState.TASK_STATE_COMPLETED + finally: + await handler.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_evicted_send_returns_current_task_snapshot( + monkeypatch: pytest.MonkeyPatch, +) -> None: + resume_requester = asyncio.Event() + requester_waiting = asyncio.Event() + requester_evicted = asyncio.Event() + original_dequeue = EventQueueSink.dequeue_event + original_close = EventQueueSink.close + + async def gated_dequeue(sink: EventQueueSink) -> Event: + if sink._evict_on_full: + requester_waiting.set() + await resume_requester.wait() + return await original_dequeue(sink) + + async def observe_close( + sink: EventQueueSink, immediate: bool = False + ) -> None: + await original_close(sink, immediate=immediate) + if sink._evict_on_full and immediate: + requester_evicted.set() + + monkeypatch.setattr(EventQueueSink, 'dequeue_event', gated_dequeue) + monkeypatch.setattr(EventQueueSink, 'close', observe_close) + handler = DefaultRequestHandlerV2( + agent_executor=_GatedExecutor(), + task_store=InMemoryTaskStore(), + agent_card=_agent_card(), + max_event_queue_size=1, + ) + request = asyncio.create_task( + handler.on_message_send(_send_request(), ServerCallContext()) + ) + try: + await requester_waiting.wait() + # The internal request-start marker fills the tap; the Task evicts it. + await requester_evicted.wait() + resume_requester.set() + result = await request + assert isinstance(result, Task) + assert result.status.state == TaskState.TASK_STATE_WORKING + active = await handler._active_task_registry.get(result.id) + assert active is not None + assert active._producer_task is not None + assert not active._producer_task.done() + finally: + resume_requester.set() + if not request.done(): + request.cancel() + await asyncio.gather(request, return_exceptions=True) + await handler.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_small_event_queue_backpressures_producer_while_store_waits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = InMemoryTaskStore() + store_waiting = asyncio.Event() + resume_store = asyncio.Event() + producer_blocked = asyncio.Event() + producer_resumed = asyncio.Event() + original_save = store.save + + async def gated_save(task: Task, context: ServerCallContext) -> None: + store_waiting.set() + await resume_store.wait() + await original_save(task, context) + + executor = _GatedExecutor() + + async def burst(context: RequestContext, event_queue: EventQueue) -> None: + source = cast('EventQueueSource', event_queue) + original_enqueue = source.enqueue_event + + async def observe_enqueue(event: Event) -> None: + blocked = ( + store_waiting.is_set() + and source.queue.full() + and source._incoming_queue.full() + ) + if blocked: + producer_blocked.set() + await original_enqueue(event) + if blocked: + producer_resumed.set() + + monkeypatch.setattr(source, 'enqueue_event', observe_enqueue) + await event_queue.enqueue_event( + Task( + id=context.task_id or '', + context_id=context.context_id or '', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + for _ in range(10): + await event_queue.enqueue_event( + TaskStatusUpdateEvent( + task_id=context.task_id or '', + context_id=context.context_id or '', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + await executor.release.wait() + + monkeypatch.setattr(store, 'save', gated_save) + monkeypatch.setattr(executor, 'execute', burst) + handler = DefaultRequestHandlerV2( + agent_executor=executor, + task_store=store, + agent_card=_agent_card(), + max_event_queue_size=1, + ) + request = asyncio.create_task( + handler.on_message_send(_send_request(), ServerCallContext()) + ) + try: + await store_waiting.wait() + await producer_blocked.wait() + assert not producer_resumed.is_set() + assert not request.done() + resume_store.set() + await producer_resumed.wait() + finally: + resume_store.set() + await handler.aclose() + if not request.done(): + request.cancel() + await asyncio.wait_for( + asyncio.gather(request, return_exceptions=True), timeout=1 + ) diff --git a/tests/server/request_handlers/test_evicted_message_response.py b/tests/server/request_handlers/test_evicted_message_response.py new file mode 100644 index 000000000..f98327fa0 --- /dev/null +++ b/tests/server/request_handlers/test_evicted_message_response.py @@ -0,0 +1,243 @@ +import asyncio + +from collections.abc import AsyncGenerator +from contextlib import aclosing +from typing import Any + +import pytest + +from a2a.server.agent_execution import AgentExecutor, RequestContext +from a2a.server.agent_execution.active_task import ActiveTask, EventConsumer +from a2a.server.context import ServerCallContext +from a2a.server.events.event_queue import Event, EventQueue +from a2a.server.events.event_queue_v2 import EventQueueSink, EventQueueSource +from a2a.server.request_handlers import DefaultRequestHandlerV2 +from a2a.server.tasks import InMemoryTaskStore +from a2a.types.a2a_pb2 import ( + AgentCard, + Message, + Part, + Role, + SendMessageRequest, + Task, + TaskState, + TaskStatus, + TaskStatusUpdateEvent, +) +from a2a.utils.errors import InvalidAgentResponseError + + +class _MessageExecutor(AgentExecutor): + def __init__(self, outcome: str) -> None: + self.outcome = outcome + self.release = asyncio.Event() + self.release_tail = asyncio.Event() + if outcome == 'message': + self.release.set() + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + if self.outcome == 'status': + await event_queue.enqueue_event( + TaskStatusUpdateEvent( + task_id=context.task_id or '', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + await self.release_tail.wait() + return + message = Message( + message_id='reply', + role=Role.ROLE_AGENT, + parts=[Part(text='response')], + ) + await event_queue.enqueue_event(message) + await self.release.wait() + if self.outcome in ('duplicate', 'duplicate-tail'): + await event_queue.enqueue_event(message) + elif self.outcome in ('task', 'task-tail'): + await event_queue.enqueue_event( + Task( + id=context.task_id or '', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + elif self.outcome in ('error', 'save-error'): + raise RuntimeError('executor failed after Message') + if self.outcome.endswith('-tail'): + await self.release_tail.wait() + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + return None + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +@pytest.mark.parametrize( + ('capacity', 'outcome'), + [ + (1, 'message'), + (2, 'message'), + (None, 'message'), + (1, 'duplicate'), + (1, 'task'), + (1, 'error'), + (1, 'save-error'), + (1, 'cancel'), + (1, 'duplicate-tail'), + (None, 'duplicate-tail'), + (1, 'task-tail'), + (None, 'task-tail'), + (1, 'status'), + (None, 'status'), + (1, 'force-close'), + (1, 'consumer-cancel'), + ], +) +async def test_evicted_message_send_finishes_validation( + capacity: int | None, + outcome: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + resume_requester = asyncio.Event() + delivery_ready = asyncio.Event() + subscription_ended = asyncio.Event() + message_received = asyncio.Event() + original_dequeue = EventQueueSink.dequeue_event + original_delivery = EventQueueSource._deliver_to_sink + original_subscribe = ActiveTask.subscribe + deliveries = 0 + active_tasks: list[ActiveTask] = [] + + async def gated_dequeue(sink: EventQueueSink) -> Event: + if sink._evict_on_full: + await resume_requester.wait() + return await original_dequeue(sink) + + async def observed_delivery( + source: EventQueueSource, + sink: EventQueueSink, + event: Event, + ) -> None: + nonlocal deliveries + await original_delivery(source, sink, event) + if sink._evict_on_full: + deliveries += 1 + required = 3 if capacity == 2 else 2 + if sink.is_closed() or deliveries >= required: + delivery_ready.set() + + async def observed_subscribe( + active: ActiveTask, + **kwargs: Any, + ) -> AsyncGenerator[Event, None]: + active_tasks.append(active) + try: + async with aclosing(original_subscribe(active, **kwargs)) as sub: + async for event in sub: + if isinstance(event, Message): + message_received.set() + yield event + finally: + subscription_ended.set() + + async def failing_status_write(consumer: EventConsumer) -> Task | None: + raise ValueError('failure status could not be saved') + + monkeypatch.setattr(EventQueueSink, 'dequeue_event', gated_dequeue) + monkeypatch.setattr(EventQueueSource, '_deliver_to_sink', observed_delivery) + monkeypatch.setattr(ActiveTask, 'subscribe', observed_subscribe) + if outcome == 'save-error': + monkeypatch.setattr( + EventConsumer, '_write_failed_status', failing_status_write + ) + executor = _MessageExecutor(outcome) + kwargs = {} if capacity is None else {'max_event_queue_size': capacity} + handler = DefaultRequestHandlerV2( + agent_executor=executor, + task_store=InMemoryTaskStore(), + agent_card=AgentCard(name='message-test', version='1'), + **kwargs, + ) + request = asyncio.create_task( + handler.on_message_send( + SendMessageRequest( + message=Message( + message_id='request', + role=Role.ROLE_USER, + parts=[Part(text='hello')], + ) + ), + ServerCallContext(), + ) + ) + try: + await delivery_ready.wait() + resume_requester.set() + if capacity is None and outcome != 'message': + if outcome != 'status': + await message_received.wait() + else: + await subscription_ended.wait() + if outcome != 'message': + if outcome != 'status': + assert not request.done() + if outcome == 'cancel': + request.cancel() + with pytest.raises(asyncio.CancelledError): + await request + assert active_tasks[0]._consumer_task is not None + assert not active_tasks[0]._consumer_task.cancelled() + assert not active_tasks[0]._consumer_task.done() + elif outcome in ('force-close', 'consumer-cancel'): + active = active_tasks[0] + if outcome == 'force-close': + await asyncio.wait_for(active.aclose(), timeout=1) + else: + assert active._consumer_task is not None + active._consumer_task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(request), timeout=1) + else: + executor.release.set() + if outcome in ( + 'duplicate', + 'task', + 'duplicate-tail', + 'task-tail', + 'status', + ): + with pytest.raises(InvalidAgentResponseError): + await asyncio.wait_for(asyncio.shield(request), timeout=1) + assert not executor.release_tail.is_set() + elif outcome in ('error', 'save-error'): + with pytest.raises(RuntimeError, match='executor failed'): + await asyncio.wait_for(asyncio.shield(request), timeout=1) + else: + result = await asyncio.wait_for( + asyncio.shield(request), timeout=1 + ) + assert isinstance(result, Message) + assert result.message_id == 'reply' + finally: + resume_requester.set() + executor.release.set() + executor.release_tail.set() + if not request.done(): + request.cancel() + await asyncio.wait_for(handler.aclose(), timeout=2) + await asyncio.wait_for( + asyncio.gather(request, return_exceptions=True), + timeout=1, + ) + + assert request.done() + assert not handler._active_task_registry._active_tasks + assert all( + task is None or task.done() + for active in active_tasks + for task in (active._producer_task, active._consumer_task) + )