diff --git a/contributors/emails/yuzilong.leif@gmail.com b/contributors/emails/yuzilong.leif@gmail.com new file mode 100644 index 0000000000000..2e0dc96163c74 --- /dev/null +++ b/contributors/emails/yuzilong.leif@gmail.com @@ -0,0 +1,2 @@ +yuzilongleif-collab +# PR #71898 diff --git a/gateway/run.py b/gateway/run.py index 4870a187cfc9e..4aaecc9d6fd4a 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -5704,6 +5704,14 @@ def __init__(self, config: Optional[GatewayConfig] = None): self._completion_deliveries_inflight: set[tuple[str, str, object]] = set() self._completion_deliveries_delivered: "OrderedDict[tuple[str, str, object], None]" = OrderedDict() self._completion_delivery_retention = 2048 + # Agent-triggered terminal completions from one conversation often land + # in the same scheduler tick. Hold them briefly so the agent receives + # 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 @@ -12220,6 +12228,16 @@ def _phase_elapsed() -> float: _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) @@ -21170,6 +21188,197 @@ async def _deliver_completion_notification( except Exception: logger.debug("Could not release durable completion claim", exc_info=True) + @staticmethod + def _completion_notification_batch_key(evt: dict) -> tuple[str, ...]: + """Return a routing-complete key for short-window process fan-in.""" + return tuple(str(evt.get(field) or "") for field in ( + "session_key", + "platform", + "chat_type", + "chat_id", + "thread_id", + "user_id", + )) + + @staticmethod + def _format_coalesced_process_completions(entries: list[tuple[str, dict, asyncio.Future]]) -> str: + """Build one bounded synthetic event from several redacted completions.""" + lines = [ + f"[IMPORTANT: {len(entries)} background processes completed for this session.", + "Treat these results as one completion batch and send at most one " + "consolidated user-facing response.", + ] + shown = entries[:10] + for _text, evt, _future in shown: + session_id = str(evt.get("session_id") or "unknown") + exit_code = evt.get("exit_code") + reason = str(evt.get("completion_reason") or "exited") + # Completion-event output is normally passed through the terminal + # redactor at the producer seam, but that redactor is deliberately + # configurable. This synthetic turn is gateway user-facing input, + # so keep the unconditional gateway floor here as defence in depth. + # Redact before slicing: truncating first can leave a credential + # fragment that no longer matches the authoritative patterns. + output = _redact_gateway_user_facing_secrets( + str(evt.get("output") or "") + ).strip() + if len(output) > 800: + output = f"[… truncated …]\n{output[-800:]}" + lines.append( + f"\n- {session_id}: exit_code={exit_code}, reason={reason}" + ) + if output: + lines.append(output) + omitted = len(entries) - len(shown) + if omitted: + lines.append( + f"\n- … and {omitted} more completion(s); inspect them with " + "the process tool if they affect the conclusion." + ) + lines.append( + "If a result does not change the current conclusion, absorb it silently.]" + ) + return "\n".join(lines) + + def _record_coalesced_completion_siblings(self, events: list[dict]) -> None: + """Extend a successful primary delivery claim to its batched siblings.""" + with self._completion_delivery_lock: + for evt in events: + identity = self._completion_delivery_identity(evt) + if identity is None: + continue + self._completion_deliveries_inflight.discard(identity) + self._completion_deliveries_delivered[identity] = None + while ( + len(self._completion_deliveries_delivered) + > self._completion_delivery_retention + ): + self._completion_deliveries_delivered.popitem(last=False) + + async def _flush_process_completion_batch(self, key: tuple[str, ...]) -> None: + """Deliver one short-window completion batch and resolve its waiters.""" + current_task = asyncio.current_task() + entries: list[tuple[str, dict, asyncio.Future]] = [] + delivered: Optional[bool] = False + try: + await asyncio.sleep(self._completion_notification_batch_window) + entries = self._completion_notification_batches.pop(key, []) + # Detach before adapter delivery. A completion that arrives while + # this batch is in flight must be able to schedule the next flush. + if self._completion_notification_batch_tasks.get(key) is current_task: + self._completion_notification_batch_tasks.pop(key, None) + if not entries: + return + 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 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]: + """Fan in concurrent process completions that share one conversation.""" + # Some unit tests construct GatewayRunner with object.__new__. Keep the + # 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() + self._completion_notification_batches.setdefault(key, []).append( + (synth_text, evt, future) + ) + if key not in self._completion_notification_batch_tasks: + 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: """Fill platform/chat_id/thread_id/chat_type on an async-delegation event. @@ -21340,7 +21549,7 @@ async def _run_process_watcher(self, watcher: dict) -> None: synth_text = format_process_notification(completion_evt) if not synth_text: break - delivered = await self._deliver_completion_notification( + delivered = await self._enqueue_process_completion_notification( synth_text, completion_evt, ) if delivered is False: diff --git a/tests/gateway/test_completion_delivery.py b/tests/gateway/test_completion_delivery.py index 3c3e09967f415..6cc37d532d8be 100644 --- a/tests/gateway/test_completion_delivery.py +++ b/tests/gateway/test_completion_delivery.py @@ -11,7 +11,7 @@ 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 @@ -316,3 +317,419 @@ async def _instant_sleep(*_a, **_kw): delivered = adapter.handle_message.await_args.args[0] assert secret not in delivered.text assert "HOME=/home/user" in delivered.text + + +def test_concurrent_process_watchers_coalesce_one_session_completion_turn(monkeypatch): + """Concurrent terminal watchers for one session must re-enter the agent once.""" + import tools.process_registry as pr_module + + registry = ProcessRegistry() + watchers = [] + for index in range(3): + session = ProcessSession( + id=f"proc_batch_{index}", + command=f"printf batch-{index}", + task_id=f"task-{index}", + started_at=1000.0 + index, + output_buffer=f"batch-{index}\n", + exited=True, + exit_code=0, + notify_on_complete=True, + ) + registry._finished[session.id] = session + watchers.append({ + "session_id": session.id, + "check_interval": 0, + "session_key": "agent:main:telegram:dm:123", + "platform": "telegram", + "chat_type": "dm", + "chat_id": "123", + "notify_on_complete": True, + }) + monkeypatch.setattr(pr_module, "process_registry", registry) + + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + + async def _exercise(): + await asyncio.gather(*( + runner._run_process_watcher(watcher) + for watcher in watchers + )) + + asyncio.run(_exercise()) + + adapter.handle_message.assert_awaited_once() + delivered = adapter.handle_message.await_args.args[0] + assert "3 background processes completed" in delivered.text + for index in range(3): + assert f"proc_batch_{index}" in delivered.text + + +def test_completion_arriving_during_batch_delivery_schedules_next_flush(): + """A new event cannot be stranded behind an in-flight batch for its route.""" + first_delivery_entered = asyncio.Event() + release_first_delivery = asyncio.Event() + delivery_count = 0 + + async def _deliver(_event): + nonlocal delivery_count + delivery_count += 1 + if delivery_count == 1: + first_delivery_entered.set() + await release_first_delivery.wait() + + adapter = SimpleNamespace(handle_message=AsyncMock(side_effect=_deliver)) + runner = _runner(adapter) + + async def _exercise(): + first = asyncio.create_task(runner._enqueue_process_completion_notification( + "first completion", + _completion_event(started_at=1.0, session_id="proc_first"), + )) + await first_delivery_entered.wait() + second = asyncio.create_task(runner._enqueue_process_completion_notification( + "second completion", + _completion_event(started_at=2.0, session_id="proc_second"), + )) + release_first_delivery.set() + assert await first is True + assert await asyncio.wait_for(second, timeout=1.0) is True + + asyncio.run(_exercise()) + + assert adapter.handle_message.await_count == 2 + + +def test_completion_batches_do_not_cross_conversation_routes(): + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + + first = _completion_event(started_at=1.0, session_id="proc_route_a") + second = _completion_event(started_at=2.0, session_id="proc_route_b") + second["session_key"] = "agent:main:telegram:dm:456" + second["chat_id"] = "456" + + async def _exercise(): + return await asyncio.gather( + runner._enqueue_process_completion_notification("first", first), + runner._enqueue_process_completion_notification("second", second), + ) + + assert asyncio.run(_exercise()) == [True, True] + assert adapter.handle_message.await_count == 2 + + +def test_failed_coalesced_delivery_retries_all_entries(): + attempts = 0 + + async def _deliver(_event): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("temporary adapter failure") + + adapter = SimpleNamespace(handle_message=AsyncMock(side_effect=_deliver)) + runner = _runner(adapter) + events = [ + _completion_event(started_at=float(index), session_id=f"proc_retry_{index}") + for index in range(2) + ] + + async def _enqueue_all(): + return await asyncio.gather(*( + runner._enqueue_process_completion_notification(f"event-{index}", event) + for index, event in enumerate(events) + )) + + async def _exercise(): + assert await _enqueue_all() == [False, False] + assert await _enqueue_all() == [True, True] + + asyncio.run(_exercise()) + assert adapter.handle_message.await_count == 2 + + +def test_coalesced_success_records_every_completion_identity(): + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + events = [ + _completion_event(started_at=float(index), session_id=f"proc_ledger_{index}") + for index in range(3) + ] + + async def _exercise(): + return await asyncio.gather(*( + runner._enqueue_process_completion_notification(f"event-{index}", event) + for index, event in enumerate(events) + )) + + assert asyncio.run(_exercise()) == [True, True, True] + for event in events: + identity = runner._completion_delivery_identity(event) + 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_coalesced_format_force_redacts_output_when_redaction_disabled(monkeypatch): + """A user setting cannot disable the gateway's outbound secret floor.""" + import agent.redact as redact_module + + secret = "abc123randomopaquetokenvalue999" + monkeypatch.setattr(redact_module, "_REDACT_ENABLED", False) + + async def _format(): + loop = asyncio.get_running_loop() + first = _completion_event(started_at=1.0, session_id="proc_secret") + first["output"] = ( + f"MY_SERVICE_TOKEN={secret}\n" + "HOME=/home/user\n" + ) + second = _completion_event(started_at=2.0, session_id="proc_control") + return GatewayRunner._format_coalesced_process_completions([ + ("first", first, loop.create_future()), + ("second", second, loop.create_future()), + ]) + + text = asyncio.run(_format()) + + assert secret not in text + assert "HOME=/home/user" in text + + +def test_coalesced_format_redacts_before_truncating_output(monkeypatch): + """Truncation cannot remove the prefix needed to recognize a secret.""" + import agent.redact as redact_module + + marker = "SHOULD_NOT_SURVIVE" + monkeypatch.setattr(redact_module, "_REDACT_ENABLED", False) + + async def _format(): + loop = asyncio.get_running_loop() + first = _completion_event(started_at=1.0, session_id="proc_long_secret") + first["output"] = f"MY_SERVICE_TOKEN={'x' * 900}{marker}\n" + second = _completion_event(started_at=2.0, session_id="proc_control") + return GatewayRunner._format_coalesced_process_completions([ + ("first", first, loop.create_future()), + ("second", second, loop.create_future()), + ]) + + text = asyncio.run(_format()) + + assert marker not in text + + +def test_duplicate_primary_does_not_discard_fresh_batch_sibling(): + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + duplicate = _completion_event(started_at=1.0, session_id="proc_duplicate") + fresh = _completion_event(started_at=2.0, session_id="proc_fresh") + duplicate_identity = runner._completion_delivery_identity(duplicate) + runner._completion_deliveries_delivered[duplicate_identity] = None + + async def _exercise(): + return await asyncio.gather( + runner._enqueue_process_completion_notification("duplicate", duplicate), + runner._enqueue_process_completion_notification("fresh", fresh), + ) + + assert asyncio.run(_exercise()) == [True, True] + adapter.handle_message.assert_awaited_once() + fresh_identity = runner._completion_delivery_identity(fresh) + assert fresh_identity in runner._completion_deliveries_delivered + + +def test_batch_format_failure_resolves_waiters_for_retry(monkeypatch): + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + monkeypatch.setattr( + runner, + "_format_coalesced_process_completions", + MagicMock(side_effect=ValueError("bad batch")), + ) + events = [ + _completion_event(started_at=float(index), session_id=f"proc_format_{index}") + for index in range(2) + ] + + async def _exercise(): + pending = asyncio.gather(*( + runner._enqueue_process_completion_notification(f"event-{index}", event) + for index, event in enumerate(events) + )) + return await asyncio.wait_for(pending, timeout=1.0) + + 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 47ca88af2d2ab..5e41786d6ab2c 100644 --- a/tests/gateway/test_gateway_shutdown.py +++ b/tests/gateway/test_gateway_shutdown.py @@ -94,6 +94,49 @@ async def block_forever(_event): 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()