fix(gateway): own completion batch task lifecycle

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:

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()

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()