From 7536655b8fea565fe32fa5f5aa61567a86f1f004 Mon Sep 17 00:00:00 2001 From: yuzilongleif-collab <235949691+yuzilongleif-collab@users.noreply.github.com> Date: Fri, 31 Jul 2026 21:33:16 +0000 Subject: [PATCH] fix(gateway): own completion batch task lifecycle --- gateway/run.py | 124 +++++++++++---- tests/gateway/test_completion_delivery.py | 176 +++++++++++++++++++++- tests/gateway/test_gateway_shutdown.py | 43 ++++++ 3 files changed, 313 insertions(+), 30 deletions(-) diff --git a/gateway/run.py b/gateway/run.py index 3f1718e077..190a2b2dc1 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -6578,7 +6578,9 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew # one synthetic turn instead of one turn per process (#70300). self._completion_notification_batches: dict[tuple[str, ...], list[tuple[str, dict, asyncio.Future]]] = {} self._completion_notification_batch_tasks: dict[tuple[str, ...], asyncio.Task] = {} + self._completion_notification_batch_flush_tasks: set[asyncio.Task] = set() self._completion_notification_batch_window = 0.1 + self._completion_notification_batches_stopping = False # Cache AIAgent instances per session to preserve prompt caching. # Without this, a new AIAgent is created per message, rebuilding the @@ -13898,6 +13900,16 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew _agent, context="shutdown idle-cache" ) + # Completion flush tasks can be sleeping in their fan-in window or + # blocked in adapter delivery. Cancel and await them while adapters + # are still alive so every watcher receives a retryable result + # before platform teardown begins. + cancel_completion_batches = getattr( + self, "_cancel_process_completion_batch_tasks", None + ) + if cancel_completion_batches is not None: + await cancel_completion_batches() + for platform, adapter in list(self.adapters.items()): await self._bounded_adapter_teardown(adapter, platform) @@ -24051,40 +24063,73 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew self._completion_notification_batch_tasks.pop(key, None) if not entries: return - try: - if len(entries) == 1: - synth_text = entries[0][0] - else: - synth_text = self._format_coalesced_process_completions(entries) + if len(entries) == 1: + synth_text = entries[0][0] + else: + synth_text = self._format_coalesced_process_completions(entries) - # A duplicate primary can legitimately return None from the - # lifecycle dedupe seam. Try the next batch identity so a - # fresh sibling is never discarded with that duplicate. - delivered = None - for _text, candidate_evt, _future in entries: - delivered = await self._deliver_completion_notification( - synth_text, candidate_evt, - ) - if delivered is not None: - break - if delivered is True and len(entries) > 1: - self._record_coalesced_completion_siblings( - [evt for _text, evt, _future in entries] - ) - except Exception: - logger.exception("Coalesced process completion delivery failed") - delivered = False - finally: - # Never strand watcher futures if formatting or delivery fails. - # False follows the existing watcher retry path. - for _text, _evt, future in entries: - if not future.done(): - future.set_result(delivered) + # A duplicate primary can legitimately return None from the + # lifecycle dedupe seam. Try the next batch identity so a + # fresh sibling is never discarded with that duplicate. + delivered = None + for _text, candidate_evt, _future in entries: + delivered = await self._deliver_completion_notification( + synth_text, candidate_evt, + ) + if delivered is not None: + break + if delivered is True and len(entries) > 1: + self._record_coalesced_completion_siblings( + [evt for _text, evt, _future in entries] + ) + except asyncio.CancelledError: + # Shutdown may cancel us either during the fan-in window or while + # adapter delivery is blocked. Recover entries that have not yet + # detached and resolve every waiter as retryable before adapters + # are torn down. + delivered = False + if not entries: + entries = self._completion_notification_batches.pop(key, []) + raise + except Exception: + logger.exception("Coalesced process completion delivery failed") + delivered = False finally: + # Never strand watcher futures if formatting, delivery, or task + # cancellation interrupts a batch. False follows the existing + # watcher retry path; None remains the ordinary dedupe result. + for _text, _evt, future in entries: + if not future.done(): + future.set_result(delivered) # Do not remove a newer flush task that reused the same route key. if self._completion_notification_batch_tasks.get(key) is current_task: self._completion_notification_batch_tasks.pop(key, None) + async def _cancel_process_completion_batch_tasks(self) -> None: + """Settle pending completion batches before adapter teardown.""" + self._completion_notification_batches_stopping = True + tasks = { + task + for task in getattr( + self, "_completion_notification_batch_flush_tasks", set() + ) + if not task.done() + } + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + # Defensive cleanup for an orphaned queue with no live flush task. + batches = getattr(self, "_completion_notification_batches", {}) + for entries in batches.values(): + for _text, _evt, future in entries: + if not future.done(): + future.set_result(False) + batches.clear() + getattr(self, "_completion_notification_batch_tasks", {}).clear() + getattr(self, "_completion_notification_batch_flush_tasks", set()).clear() + async def _enqueue_process_completion_notification( self, synth_text: str, evt: dict, ) -> Optional[bool]: @@ -24093,8 +24138,17 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew # batching seam lazy so those focused lifecycle tests remain valid. if not hasattr(self, "_completion_notification_batches"): self._completion_notification_batches = {} + if not hasattr(self, "_completion_notification_batch_tasks"): self._completion_notification_batch_tasks = {} + if not hasattr(self, "_completion_notification_batch_flush_tasks"): + self._completion_notification_batch_flush_tasks = set() + if not hasattr(self, "_completion_notification_batch_window"): self._completion_notification_batch_window = 0.1 + if not hasattr(self, "_completion_notification_batches_stopping"): + self._completion_notification_batches_stopping = False + + if self._completion_notification_batches_stopping: + return False key = self._completion_notification_batch_key(evt) future = asyncio.get_running_loop().create_future() @@ -24102,9 +24156,21 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew (synth_text, evt, future) ) if key not in self._completion_notification_batch_tasks: - self._completion_notification_batch_tasks[key] = asyncio.create_task( + task = asyncio.create_task( self._flush_process_completion_batch(key) ) + self._completion_notification_batch_tasks[key] = task + # Keep the flush alive and include it in the gateway's normal + # lifecycle accounting. Focused tests that construct a runner via + # object.__new__ lazily receive the same ownership set. + if not hasattr(self, "_background_tasks"): + self._background_tasks = set() + self._background_tasks.add(task) + self._completion_notification_batch_flush_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + task.add_done_callback( + self._completion_notification_batch_flush_tasks.discard + ) return await future def _enrich_async_delegation_routing(self, evt: dict) -> None: diff --git a/tests/gateway/test_completion_delivery.py b/tests/gateway/test_completion_delivery.py index 4aba541b86..dfc4bc48ba 100644 --- a/tests/gateway/test_completion_delivery.py +++ b/tests/gateway/test_completion_delivery.py @@ -11,7 +11,7 @@ import json import queue from collections import OrderedDict from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -46,6 +46,7 @@ def _runner(adapter, *, origins=None): runner._completion_deliveries_inflight = set() runner._completion_deliveries_delivered = OrderedDict() runner._completion_delivery_retention = 2048 + runner._background_tasks = set() return runner @@ -469,6 +470,30 @@ def test_coalesced_success_records_every_completion_identity(): assert identity in runner._completion_deliveries_delivered +def test_coalesced_format_bounds_details_and_reports_omitted_count(): + async def _format(): + loop = asyncio.get_running_loop() + entries = [ + ( + f"event-{index}", + _completion_event( + started_at=float(index), session_id=f"proc_bound_{index}" + ), + loop.create_future(), + ) + for index in range(12) + ] + return GatewayRunner._format_coalesced_process_completions(entries) + + text = asyncio.run(_format()) + + for index in range(10): + assert f"proc_bound_{index}" in text + assert "proc_bound_10" not in text + assert "proc_bound_11" not in text + assert "and 2 more completion(s)" in text + + def test_duplicate_primary_does_not_discard_fresh_batch_sibling(): adapter = SimpleNamespace(handle_message=AsyncMock()) runner = _runner(adapter) @@ -511,3 +536,152 @@ def test_batch_format_failure_resolves_waiters_for_retry(monkeypatch): assert asyncio.run(_exercise()) == [False, False] adapter.handle_message.assert_not_awaited() + + +def test_shutdown_cancels_batch_during_window_and_settles_waiter_for_retry(): + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + sleep_entered = asyncio.Event() + release_sleep = asyncio.Event() + real_sleep = asyncio.sleep + event = _completion_event(started_at=1.0, session_id="proc_cancel_window") + + async def _controlled_sleep(delay): + if delay == runner._completion_notification_batch_window: + sleep_entered.set() + await release_sleep.wait() + return + await real_sleep(delay) + + async def _exercise(): + pending = asyncio.create_task( + runner._enqueue_process_completion_notification("completion", event) + ) + await sleep_entered.wait() + flush_task = next(iter(runner._completion_notification_batch_tasks.values())) + assert flush_task in runner._background_tasks + + await runner._cancel_process_completion_batch_tasks() + + assert await asyncio.wait_for(pending, timeout=1.0) is False + assert flush_task.cancelled() + assert flush_task not in runner._background_tasks + assert runner._completion_notification_batches == {} + assert runner._completion_notification_batch_tasks == {} + + with patch("gateway.run.asyncio.sleep", new=_controlled_sleep): + asyncio.run(_exercise()) + adapter.handle_message.assert_not_awaited() + + +def test_shutdown_cancels_blocked_batch_delivery_and_keeps_it_retryable(): + delivery_entered = asyncio.Event() + + async def _blocked_delivery(_event): + delivery_entered.set() + await asyncio.Event().wait() + + adapter = SimpleNamespace(handle_message=AsyncMock(side_effect=_blocked_delivery)) + runner = _runner(adapter) + runner._completion_notification_batch_window = 0 + event = _completion_event(started_at=1.0, session_id="proc_cancel_delivery") + + async def _exercise(): + pending = asyncio.create_task( + runner._enqueue_process_completion_notification("completion", event) + ) + await delivery_entered.wait() + flush_task = next(iter(runner._completion_notification_batch_flush_tasks)) + + await runner._cancel_process_completion_batch_tasks() + + assert await asyncio.wait_for(pending, timeout=1.0) is False + assert flush_task.cancelled() + assert runner._completion_delivery_identity(event) not in runner._completion_deliveries_inflight + assert runner._completion_delivery_identity(event) not in runner._completion_deliveries_delivered + assert runner._completion_notification_batches == {} + assert runner._completion_notification_batch_tasks == {} + + asyncio.run(_exercise()) + adapter.handle_message.assert_awaited_once() + + +def test_completion_enqueue_stays_retryable_after_shutdown_starts(): + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + + async def _exercise(): + await runner._cancel_process_completion_batch_tasks() + return await runner._enqueue_process_completion_notification( + "completion", + _completion_event(started_at=1.0, session_id="proc_after_shutdown"), + ) + + assert asyncio.run(_exercise()) is False + assert runner._completion_notification_batches == {} + assert runner._completion_notification_batch_tasks == {} + adapter.handle_message.assert_not_awaited() + + +def test_successful_batch_releases_all_lifecycle_task_references(): + adapter = SimpleNamespace(handle_message=AsyncMock(return_value=None)) + runner = _runner(adapter) + runner._completion_notification_batch_window = 0 + + async def _exercise(): + result = await runner._enqueue_process_completion_notification( + "completion", + _completion_event(started_at=1.0, session_id="proc_success_cleanup"), + ) + await asyncio.sleep(0) + return result + + assert asyncio.run(_exercise()) is True + assert runner._completion_notification_batch_tasks == {} + assert runner._completion_notification_batch_flush_tasks == set() + assert runner._background_tasks == set() + + +def test_shutdown_cancels_overlapping_flushes_for_same_route(): + delivery_entered = asyncio.Event() + + async def _blocked_delivery(_event): + delivery_entered.set() + await asyncio.Event().wait() + + adapter = SimpleNamespace(handle_message=AsyncMock(side_effect=_blocked_delivery)) + runner = _runner(adapter) + runner._completion_notification_batch_window = 0 + first_event = _completion_event(started_at=1.0, session_id="proc_old_flush") + second_event = _completion_event(started_at=2.0, session_id="proc_new_flush") + + async def _exercise(): + first = asyncio.create_task( + runner._enqueue_process_completion_notification("first", first_event) + ) + await delivery_entered.wait() + + # The first task has detached from the route index while blocked in + # adapter delivery. A new completion for the same route must create a + # second flush, and shutdown must still own and cancel both tasks. + assert runner._completion_notification_batch_tasks == {} + runner._completion_notification_batch_window = 3600 + second = asyncio.create_task( + runner._enqueue_process_completion_notification("second", second_event) + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + flush_tasks = set(runner._completion_notification_batch_flush_tasks) + assert len(flush_tasks) == 2 + + await runner._cancel_process_completion_batch_tasks() + + assert await asyncio.gather(first, second) == [False, False] + assert all(task.cancelled() for task in flush_tasks) + assert runner._completion_notification_batches == {} + assert runner._completion_notification_batch_tasks == {} + assert runner._completion_notification_batch_flush_tasks == set() + assert runner._background_tasks == set() + + asyncio.run(_exercise()) + adapter.handle_message.assert_awaited_once() diff --git a/tests/gateway/test_gateway_shutdown.py b/tests/gateway/test_gateway_shutdown.py index 367c550e40..3aaeb51696 100644 --- a/tests/gateway/test_gateway_shutdown.py +++ b/tests/gateway/test_gateway_shutdown.py @@ -104,6 +104,49 @@ async def test_gateway_stop_interrupts_running_agents_and_cancels_adapter_tasks( assert runner._shutdown_event.is_set() is True +@pytest.mark.asyncio +async def test_gateway_stop_settles_completion_batch_before_adapter_disconnect(): + runner, adapter = make_restart_runner() + runner._completion_notification_batch_window = 3600 + event = { + "session_id": "shutdown-batch", + "started_at": 1.0, + "session_key": "telegram:dm:123456:u1", + "platform": "telegram", + "chat_type": "dm", + "chat_id": "123456", + "user_id": "u1", + "exit_code": 0, + "output": "done", + } + call_order: list[str] = [] + original_cancel = runner._cancel_process_completion_batch_tasks + + async def _tracked_cancel(): + call_order.append("batch_cancel_start") + await original_cancel() + call_order.append("batch_cancel_done") + + async def _disconnect(): + call_order.append("disconnect") + + runner._cancel_process_completion_batch_tasks = _tracked_cancel + adapter.disconnect = _disconnect + pending = asyncio.create_task( + runner._enqueue_process_completion_notification("completion", event) + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert runner._completion_notification_batch_flush_tasks + + with patch("gateway.status.remove_pid_file"), patch("gateway.status.write_runtime_status"): + await runner.stop() + + assert await asyncio.wait_for(pending, timeout=1.0) is False + assert call_order == ["batch_cancel_start", "batch_cancel_done", "disconnect"] + assert runner._completion_notification_batch_flush_tasks == set() + + @pytest.mark.asyncio async def test_in_chat_restart_skips_home_shutdown_even_with_active_session(): runner, adapter = make_restart_runner()