From 4c047764d5702787e9d391676eec5ed0b5ab351a Mon Sep 17 00:00:00 2001 From: Sankalp Gilda Date: Wed, 15 Jul 2026 18:26:28 -0400 Subject: [PATCH 1/2] fix(server): let subscriber taps evict on full instead of wedging dispatch One undrained sink stalls EventQueueSource._dispatch_loop for every sink: the dispatcher's gather awaits a blocking put on each sink, so a single full queue whose consumer was abandoned blocks the gather forever. The incoming queue fills behind it, producers wedge in enqueue_event, and under sustained load the task's event flow never recovers (see the task dumps and repro on the tracking issue). Split the sink semantics. The default sink is flow control: its consumer drives task state, so a blocking put remains correct backpressure. Tapped subscriber sinks are broadcast observers whose remote consumers can be abandoned, so tap() now accepts evict_on_full: when such a sink is full at delivery time it is force-closed and detached, dispatch to the remaining sinks continues, and blocked producers recover. The full() pre-check is race-free because the dispatcher is the only writer to a sink queue. tap() defaults to evict_on_full=False, preserving the documented blocking contract for existing callers; ActiveTask's subscriber tap opts in explicitly, which fixes the production wedge without changing the public primitive's default behavior. Regression tests cover eviction with continued dispatch, producer recovery after eviction, preserved backpressure on default taps, and the subscriber tap opting in; the eviction tests fail without the fix. --- src/a2a/server/agent_execution/active_task.py | 8 +- src/a2a/server/events/event_queue_v2.py | 72 ++++++++++++++-- .../agent_execution/test_active_task.py | 52 ++++++++++++ tests/server/events/test_event_queue_v2.py | 84 +++++++++++++++++++ 4 files changed, 209 insertions(+), 7 deletions(-) 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..001ad610f 100644 --- a/src/a2a/server/events/event_queue_v2.py +++ b/src/a2a/server/events/event_queue_v2.py @@ -67,6 +67,37 @@ 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. + """ + 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 +109,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 +195,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 +334,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 +383,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 +393,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..8974791a9 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,86 @@ 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) From 9830251685f905fac5ff555a6828193795543113 Mon Sep 17 00:00:00 2001 From: Sankalp Gilda Date: Thu, 16 Jul 2026 12:01:58 -0400 Subject: [PATCH 2/2] fix(server): skip closed sinks in dispatch instead of evicting them A sink mid-graceful-close (close(immediate=False), consumer still draining) can remain in the dispatcher's active_sinks snapshot. With evict_on_full=True and a full queue, delivering to it would upgrade the close to immediate and discard the events the consumer was draining. Skip closed sinks before the evict-on-full check; a close landing after the check is harmless because the put already tolerates a shut-down queue. --- src/a2a/server/events/event_queue_v2.py | 10 ++++++ tests/server/events/test_event_queue_v2.py | 40 ++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/src/a2a/server/events/event_queue_v2.py b/src/a2a/server/events/event_queue_v2.py index 001ad610f..fd2c7be48 100644 --- a/src/a2a/server/events/event_queue_v2.py +++ b/src/a2a/server/events/event_queue_v2.py @@ -86,7 +86,17 @@ async def _deliver_to_sink( 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 ' diff --git a/tests/server/events/test_event_queue_v2.py b/tests/server/events/test_event_queue_v2.py index 8974791a9..828ee1277 100644 --- a/tests/server/events/test_event_queue_v2.py +++ b/tests/server/events/test_event_queue_v2.py @@ -900,3 +900,43 @@ async def test_default_tap_keeps_backpressure() -> None: 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)