diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index b0154c8d6..0d80676a9 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -607,7 +607,13 @@ async def subscribe( self._reference_count, ) - tapped_queue = await self._event_queue_subscribers.tap() + # Subscriber sinks belong to remote consumers that can be abandoned + # (e.g. a non-blocking send whose HTTP response already returned). + # evict_on_full keeps one undrained subscriber from wedging dispatch + # for every other subscriber and, transitively, the producer. + tapped_queue = await self._event_queue_subscribers.tap( + evict_on_full=True + ) request_id = await self.enqueue_request(request) if request else None try: diff --git a/src/a2a/server/events/event_queue_v2.py b/src/a2a/server/events/event_queue_v2.py index 3386a732e..fd2c7be48 100644 --- a/src/a2a/server/events/event_queue_v2.py +++ b/src/a2a/server/events/event_queue_v2.py @@ -67,6 +67,47 @@ def queue(self) -> AsyncQueue[Event]: raise ValueError('No default sink available.') return self._default_sink.queue + async def _deliver_to_sink( + self, sink: 'EventQueueSink', event: Event + ) -> None: + """Puts one event into one sink, evicting an evict-on-full sink if full. + + A sink whose consumer stopped draining (an abandoned subscriber) fills + up and, with a blocking put, would stall this dispatcher's ``gather`` + forever: the incoming queue then fills, producers wedge in + ``enqueue_event``, and the task's event flow never recovers. Sinks + tapped with ``evict_on_full=True`` are therefore treated as broadcast + observers: when full, the sink is force-closed and detached so + dispatch to the remaining sinks continues and blocked producers + recover. Sinks with ``evict_on_full=False`` (including the default + sink) keep the blocking put, which is the documented flow-control + contract. + + The ``full()`` pre-check is race-free: this dispatcher is the only + writer to a sink queue and consumers only read, so the queue cannot + grow between the check and the put. + + A sink that closed after the dispatch snapshot was taken (e.g. a + graceful ``close(immediate=False)`` whose consumer is still + draining) is skipped: treating its full queue as an abandoned + consumer would upgrade the close to immediate and discard the + events the consumer is draining. The check is best-effort — a + close landing after it is harmless, because the put below + tolerates a shut-down queue. + """ + if sink.is_closed(): + return + if sink._evict_on_full and sink.queue.full(): # noqa: SLF001 + logger.warning( + 'Evicting event queue sink %r: its queue is full and the ' + 'consumer is not draining it. Closing the sink so dispatch ' + 'to the remaining sinks continues.', + sink, + ) + await sink.close(immediate=True) + return + await sink._put_internal(event) # noqa: SLF001 + async def _dispatch_loop(self) -> None: try: while True: @@ -78,7 +119,7 @@ async def _dispatch_loop(self) -> None: if active_sinks: results = await asyncio.gather( *( - sink._put_internal(event) # noqa: SLF001 + self._deliver_to_sink(sink, event) for sink in active_sinks ), return_exceptions=True, @@ -164,17 +205,33 @@ async def _join_incoming_queue(self) -> None: ) async def tap( - self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE + self, + max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, + evict_on_full: bool = False, ) -> 'EventQueueSink': """Taps the event queue to create a new child queue that receives future events. Note: The tapped queue may receive some old events if the incoming event queue is lagging behind and hasn't dispatched them yet. + + Args: + max_queue_size: Bound for the new sink's queue. + evict_on_full: When True, the dispatcher treats this sink as a + broadcast observer: if its queue is full at delivery time, the + sink is force-closed and detached instead of blocking dispatch + (and, transitively, every producer) behind it. When False (the + default), a full sink back-pressures dispatch, preserving the + documented flow-control behavior. Use True for sinks whose + consumer may be abandoned, such as remote subscribers. """ async with self._lock: if self._is_closed: raise QueueShutDown('Cannot tap a closed EventQueueSource.') - sink = EventQueueSink(parent=self, max_queue_size=max_queue_size) + sink = EventQueueSink( + parent=self, + max_queue_size=max_queue_size, + evict_on_full=evict_on_full, + ) self._sinks.add(sink) return sink @@ -287,12 +344,21 @@ def __init__( self, parent: EventQueueSource, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, + evict_on_full: bool = False, ) -> None: - """Initializes the EventQueueSink.""" + """Initializes the EventQueueSink. + + Args: + parent: The EventQueueSource this sink belongs to. + max_queue_size: Bound for this sink's queue. + evict_on_full: Dispatch policy for this sink; see + ``EventQueueSource.tap``. + """ if max_queue_size <= 0: raise ValueError('max_queue_size must be greater than 0') self._parent = parent + self._evict_on_full = evict_on_full self._queue: AsyncQueue[Event] = create_async_queue( maxsize=max_queue_size ) @@ -327,7 +393,9 @@ def task_done(self) -> None: self._queue.task_done() async def tap( - self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE + self, + max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE, + evict_on_full: bool = False, ) -> 'EventQueueSink': """Creates a child queue that receives future events. @@ -335,7 +403,9 @@ async def tap( queue is lagging behind and hasn't dispatched them yet. """ # Delegate tap to the parent source so all sinks are flat under the source - return await self._parent.tap(max_queue_size=max_queue_size) + return await self._parent.tap( + max_queue_size=max_queue_size, evict_on_full=evict_on_full + ) async def close(self, immediate: bool = False) -> None: """Closes the child sink queue. diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index ce9e2c068..3de018cb2 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -895,3 +895,55 @@ async def execute_mock(req, q): assert len(events) == 0 await active_task.cancel(request_context) + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_subscriber_taps_are_evict_on_full(): + """Subscriber sinks opt into evict-on-full dispatch. + + A subscriber can be abandoned (non-blocking send whose response already + returned), so its sink must never wedge dispatch for other subscribers. + """ + agent_executor = Mock() + task_manager = Mock() + + active_task = ActiveTask( + agent_executor=agent_executor, + task_id='test-task-id', + task_manager=task_manager, + push_sender=Mock(), + ) + + async def slow_execute(req, q): + await asyncio.sleep(10) + + agent_executor.execute = AsyncMock(side_effect=slow_execute) + task_manager.get_task = AsyncMock( + return_value=Task( + id='test-task-id', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + await active_task.enqueue_request(Mock(spec=RequestContext)) + await active_task.start( + call_context=ServerCallContext(), create_task_if_missing=True + ) + + subscriber = active_task.subscribe(include_initial_task=True) + task = await anext(subscriber) + assert task.id == 'test-task-id' + + sinks = active_task._event_queue_subscribers._sinks + assert sinks, 'subscribe should have tapped a sink' + assert all(sink._evict_on_full for sink in sinks) + + await subscriber.aclose() + background = [ + t + for t in (active_task._producer_task, active_task._consumer_task) + if t is not None + ] + for t in background: + t.cancel() + await asyncio.gather(*background, return_exceptions=True) diff --git a/tests/server/events/test_event_queue_v2.py b/tests/server/events/test_event_queue_v2.py index 27bceea4c..828ee1277 100644 --- a/tests/server/events/test_event_queue_v2.py +++ b/tests/server/events/test_event_queue_v2.py @@ -8,6 +8,7 @@ from a2a.server.events.event_queue import ( DEFAULT_MAX_QUEUE_SIZE, + Event, EventQueue, QueueShutDown, ) @@ -816,3 +817,126 @@ async def consumer_second() -> None: # Validate that: after unblocking second consumer everything ends smoothly. assert len(consumed_first) == 50 assert len(consumed_second) == 50 + + +@pytest.mark.asyncio +async def test_evict_on_full_sink_is_evicted_and_dispatch_continues() -> None: + """A full evict_on_full sink is evicted; other sinks keep receiving. + + Regression test for the runtime wedge reported on #1101: one abandoned + subscriber sink fills up, the dispatcher's gather blocks on it forever, + and no other sink receives anything again. + """ + queue = EventQueueSource() + stuck_sink = await queue.tap(max_queue_size=1, evict_on_full=True) + + consumed: list[Event] = [] + + async def consume_default() -> None: + while True: + try: + event = await queue.dequeue_event() + consumed.append(event) + queue.task_done() + except QueueShutDown: + break + + consumer_task = asyncio.create_task(consume_default()) + + # The first event fills the stuck sink (capacity 1, never drained). The + # second and third reach the default sink only if the dispatcher + # survives the stuck sink. + for i in range(3): + await queue.enqueue_event(create_sample_message(str(i))) + + async def default_sink_received_all() -> None: + while len(consumed) < 3: + await asyncio.sleep(0.01) + + await asyncio.wait_for(default_sink_received_all(), timeout=2.0) + + # The stuck sink was evicted: closed and detached from the source. + assert stuck_sink.is_closed() + assert stuck_sink not in queue._sinks # noqa: SLF001 + + await queue.close(immediate=True) + await consumer_task + + +@pytest.mark.asyncio +async def test_producer_unblocked_after_evict_on_full() -> None: + """Producers blocked on a full incoming queue recover after eviction. + + Without eviction, the dispatcher never finishes its gather, the incoming + queue stays full, and every producer wedges in enqueue_event permanently. + """ + queue = EventQueueSource(max_queue_size=1, create_default_sink=False) + await queue.tap(max_queue_size=1, evict_on_full=True) + + # Sink capacity 1 + incoming capacity 1 both fill, so without eviction + # the later enqueues block forever. With eviction the pipeline drains + # and every enqueue completes. + async def produce() -> None: + for i in range(4): + await queue.enqueue_event(create_sample_message(str(i))) + + await asyncio.wait_for(produce(), timeout=2.0) + await queue.close(immediate=True) + + +@pytest.mark.asyncio +async def test_default_tap_keeps_backpressure() -> None: + """A default tap (evict_on_full=False) preserves blocking dispatch.""" + queue = EventQueueSource() + stuck_sink = await queue.tap(max_queue_size=1) + + for i in range(3): + await queue.enqueue_event(create_sample_message(str(i))) + + # Give the dispatcher time to block on the full sink; the sink must + # survive and stay attached (flow control, not eviction). + await asyncio.sleep(0.3) + assert not stuck_sink.is_closed() + assert stuck_sink in queue._sinks # noqa: SLF001 + + await queue.close(immediate=True) + + +@pytest.mark.asyncio +async def test_graceful_close_not_upgraded_to_immediate_by_eviction() -> None: + """A sink mid-graceful-close is skipped, never force-closed by eviction. + + ``close(immediate=False)`` marks the sink closed and waits for the + consumer to drain the remaining events, but the dispatcher may still + hold the sink in an ``active_sinks`` snapshot taken before the close + removed it. Delivering to that sink while its queue is full must not + trip the evict-on-full path (``close(immediate=True)``), which would + discard the events the consumer is still draining. + """ + queue = EventQueueSource() + sink = await queue.tap(max_queue_size=1, evict_on_full=True) + + # Fill the sink so the evict-on-full branch would be reachable. + await sink._put_internal(create_sample_message('0')) # noqa: SLF001 + + # Start a graceful close: it marks the sink closed, then blocks in + # queue.join() until the pending event is consumed. + close_task = asyncio.create_task(sink.close(immediate=False)) + while not sink.is_closed(): + await asyncio.sleep(0) + + # Simulate the dispatcher delivering from a stale snapshot: the sink + # is full and marked closed. It must be skipped, not evicted. + await queue._deliver_to_sink( # noqa: SLF001 + sink, create_sample_message('1') + ) + + # The pending event survived the delivery attempt and the consumer + # finishes draining, which lets the graceful close complete. + event = await sink.dequeue_event() + assert isinstance(event, Message) + assert event.message_id == '0' + sink.task_done() + await asyncio.wait_for(close_task, timeout=2.0) + + await queue.close(immediate=True)