diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 0b0be4e85..3f9fbd47c 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. @@ -296,6 +309,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, @@ -478,6 +492,8 @@ def __init__( # noqa: PLR0913 on_cleanup: Callable[[ActiveTask], None] | None = None, event_stream: TaskEventStream | None = None, interrupted_task_timeout: float | None = None, + *, + max_event_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, ) -> None: """Initializes the ActiveTask. @@ -491,15 +507,23 @@ 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. interrupted_task_timeout: Optional seconds after which an idle interrupted task is finished; ``None`` disables it. """ # --- 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 @@ -519,6 +543,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. @@ -529,6 +554,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`. @@ -782,7 +809,7 @@ async def subscribe( # noqa: PLR0915 # 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 @@ -970,6 +997,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 @@ -1091,6 +1125,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 e5daaaa2e..f0478cf7a 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 @@ -28,19 +29,26 @@ class ActiveTaskRegistry: """A registry for active ActiveTask instances.""" - def __init__( + def __init__( # noqa: PLR0913 self, agent_executor: AgentExecutor, task_store: TaskStore | VersionedTaskStore, push_sender: PushNotificationSender | None = None, event_stream: TaskEventStream | None = None, interrupted_task_timeout: float | 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._interrupted_task_timeout = interrupted_task_timeout 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() @@ -82,6 +90,7 @@ async def get_or_create( on_cleanup=self._on_active_task_cleanup, event_stream=self._event_stream, interrupted_task_timeout=self._interrupted_task_timeout, + 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 ee32880e7..a69e155f8 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, @@ -97,6 +98,22 @@ class DefaultRequestHandlerV2(RequestHandler): 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. + ``interrupted_task_timeout`` releases the `ActiveTask` of a task idle in ``input-required`` or ``auth-required`` after that many seconds; a follow-up recreates it from the store. ``None`` (default) never releases; @@ -124,6 +141,8 @@ def __init__( # noqa: PLR0913 validate_input_modes: bool = False, event_stream: TaskEventStream | None = None, interrupted_task_timeout: float | None = None, + *, + max_event_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, ) -> None: if ( interrupted_task_timeout is not None @@ -184,6 +203,7 @@ def __init__( # noqa: PLR0913 push_sender=self._push_sender, event_stream=self._event_stream, interrupted_task_timeout=interrupted_task_timeout, + max_event_queue_size=max_event_queue_size, ) self._background_tasks = set() @@ -461,7 +481,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_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py index a57be1fa0..2f19637f2 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -3163,6 +3163,43 @@ async def test_follow_up_after_release_resumes_from_store() -> None: await handler.aclose() +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'state', + [TaskState.TASK_STATE_INPUT_REQUIRED, TaskState.TASK_STATE_AUTH_REQUIRED], +) +async def test_event_queue_capacity_survives_interrupted_task_release( + state: TaskState, +) -> None: + """Recreated interrupted tasks retain the configured event capacity.""" + handler = DefaultRequestHandlerV2( + agent_executor=_InterruptingAgent(state, complete_on_follow_up=False), + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + interrupted_task_timeout=0.1, + max_event_queue_size=8, + ) + try: + task = await _send(handler, 'book a flight') + first = await handler._active_task_registry.get(task.id) + assert first is not None + assert first._event_queue_agent.queue.maxsize == 8 + assert first._event_queue_subscribers._incoming_queue.maxsize == 8 + assert await _wait_for_release(handler, task.id) + + result = await _send(handler, 'Friday', task) + recreated = await handler._active_task_registry.get(task.id) + + assert result.status.state == state + assert recreated is not None + assert recreated is not first + assert recreated._event_queue_agent.queue.maxsize == 8 + assert recreated._event_queue_subscribers._incoming_queue.maxsize == 8 + assert await _wait_for_release(handler, task.id) + finally: + await handler.aclose() + + @pytest.mark.asyncio async def test_follow_up_resets_interrupted_task_timeout() -> None: """A follow-up that interrupts again restarts the timer.""" 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) + )