diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index a218451b8..cead75a6a 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -399,6 +399,7 @@ class AgentLoop: self._runtime_context_providers: list[RuntimeContextProvider] = [] self._active_tasks: dict[str, set[asyncio.Task[Any]]] = {} self._background_tasks: set[asyncio.Task[Any]] = set() + self._close_mcp_lock = asyncio.Lock() self._session_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = ( weakref.WeakValueDictionary() ) @@ -1338,11 +1339,42 @@ class AgentLoop: await self._publish_next_deferred_automation_turn(session_key) async def close_mcp(self) -> None: - """Drain background work, stop exec sessions, then close MCP connections.""" - if self._background_tasks: - await asyncio.gather(*self._background_tasks, return_exceptions=True) - self._background_tasks.clear() + """Stop active work, then close exec, subagent, and MCP resources. + + Resource teardown must still run if cancellation interrupts task draining. + Gateway shutdown deliberately bounds this coroutine, so keeping the cleanup + phase in ``finally`` prevents a timed-out background task from leaving + subprocess transports alive after the event loop closes. + """ + # The agent loop closes itself from ``run()`` while gateway shutdown also + # performs a guaranteed final close. Serialize those owners so they cannot + # tear down the same subprocess transports concurrently. + close_lock = getattr(self, "_close_mcp_lock", None) + if close_lock is None: + close_lock = self._close_mcp_lock = asyncio.Lock() + async with close_lock: + await self._close_mcp_unlocked() + + async def _close_mcp_unlocked(self) -> None: errors: list[BaseException] = [] + active_task_groups = getattr(self, "_active_tasks", {}) + active_tasks = tuple({task for tasks in active_task_groups.values() for task in tasks}) + active_task_groups.clear() + current_task = asyncio.current_task() + active_tasks = tuple(task for task in active_tasks if task is not current_task) + for task in active_tasks: + if not task.done(): + task.cancel() + try: + if active_tasks: + await asyncio.gather(*active_tasks, return_exceptions=True) + if self._background_tasks: + await asyncio.gather(*self._background_tasks, return_exceptions=True) + except BaseException as exc: + errors.append(exc) + finally: + self._background_tasks.clear() + cleanup_steps = ( self.subagents.close, self._exec_session_manager.close_all, diff --git a/nanobot/cli/gateway_runtime.py b/nanobot/cli/gateway_runtime.py index 8686cd74e..5648f9192 100644 --- a/nanobot/cli/gateway_runtime.py +++ b/nanobot/cli/gateway_runtime.py @@ -206,7 +206,6 @@ async def _close_gateway_runtime( channels: Any, tasks: list[asyncio.Task[Any]], runtime_tasks: asyncio.Future[list[Any]] | None, - runtime_tasks_drained: bool, *, task_wait_timeout: float = 15.0, close_timeout: float = 15.0, @@ -230,16 +229,26 @@ async def _close_gateway_runtime( for task in tasks: if not task.done(): task.cancel() + pending: set[asyncio.Task[Any]] = set() if tasks: # Bounded: a coroutine that swallows cancellation (e.g. an SDK reconnect # loop) must not hold the stop open until systemd's timeout kills the # cgroup. Anything still pending is abandoned and closed underneath. - await asyncio.wait(tasks, timeout=task_wait_timeout) + _done, pending = await asyncio.wait(tasks, timeout=task_wait_timeout) + # A task can swallow the first cancellation while unwinding. Re-cancel + # timed-out tasks so an agent loop stuck draining background work reaches + # its resource-cleanup phase before the explicit final close below. + for task in pending: + task.cancel() + if runtime_tasks is not None and not runtime_tasks.done(): + runtime_tasks.cancel() try: await asyncio.wait_for(agent.close_mcp(), timeout=close_timeout) except BaseException as exc: # noqa: BLE001 - shutdown must proceed logger.warning("Gateway shutdown: agent resource cleanup incomplete: {}", exc) - if runtime_tasks is not None and not runtime_tasks_drained: + # Retrieving an already-finished gather prevents noisy unhandled exceptions, + # but never wait for it here: its children were bounded individually above. + if runtime_tasks is not None and runtime_tasks.done(): with suppress(asyncio.CancelledError, Exception): await runtime_tasks @@ -777,7 +786,6 @@ def _run_gateway( tasks: list[asyncio.Task[Any]] = [] shutdown_task: asyncio.Task[Any] | None = None runtime_tasks: asyncio.Future[list[Any]] | None = None - runtime_tasks_drained = False shutdown_event = asyncio.Event() cli_terminal._ensure_interactive_tty_mode() restore_shutdown_handlers = _install_gateway_shutdown_handlers( @@ -829,7 +837,6 @@ def _run_gateway( return_when=asyncio.FIRST_COMPLETED, ) if runtime_tasks in done: - runtime_tasks_drained = True await runtime_tasks else: runtime_tasks.cancel() @@ -850,9 +857,7 @@ def _run_gateway( agent.stop() # Cancel runtime tasks first, then deterministically close # exec/MCP resources while the event loop is still alive. - await _close_gateway_runtime( - agent, channels, tasks, runtime_tasks, runtime_tasks_drained - ) + await _close_gateway_runtime(agent, channels, tasks, runtime_tasks) # Flush all cached sessions to durable storage before exit. # This prevents data loss on filesystems with write-back # caching (rclone VFS, NFS, FUSE mounts, etc.). diff --git a/tests/agent/test_task_cancel.py b/tests/agent/test_task_cancel.py index 6acb41ff7..bb68183d9 100644 --- a/tests/agent/test_task_cancel.py +++ b/tests/agent/test_task_cancel.py @@ -55,6 +55,62 @@ class TestHandleStop: out = await cmd_stop(ctx) assert "No active task" in out.content + @pytest.mark.asyncio + async def test_close_mcp_cancels_active_turn_before_resources(self): + loop, _bus = _make_loop() + events: list[str] = [] + + async def active_turn(): + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + events.append("turn_cancelled") + raise + + task = asyncio.create_task(active_turn()) + await asyncio.sleep(0) + loop._active_tasks["test:c1"] = {task} + + async def close_subagents(): + events.append("resources_closed") + + loop.subagents.close = close_subagents + loop._exec_session_manager.close_all = AsyncMock() + with patch("nanobot.agent.loop.agent_context.close_mcp", AsyncMock()): + await loop.close_mcp() + + assert events == ["turn_cancelled", "resources_closed"] + assert task.cancelled() + + @pytest.mark.asyncio + async def test_close_mcp_serializes_duplicate_cleanup(self): + loop, _bus = _make_loop() + entered = asyncio.Event() + release = asyncio.Event() + concurrent = 0 + max_concurrent = 0 + + async def close_subagents(): + nonlocal concurrent, max_concurrent + concurrent += 1 + max_concurrent = max(max_concurrent, concurrent) + entered.set() + await release.wait() + concurrent -= 1 + + loop.subagents.close = close_subagents + loop._exec_session_manager.close_all = AsyncMock() + with patch("nanobot.agent.loop.agent_context.close_mcp", AsyncMock()): + first = asyncio.create_task(loop.close_mcp()) + await entered.wait() + second = asyncio.create_task(loop.close_mcp()) + await asyncio.sleep(0) + assert not second.done() + release.set() + await asyncio.gather(first, second) + + assert max_concurrent == 1 + @pytest.mark.asyncio async def test_stop_cancels_active_task(self): from nanobot.bus.events import InboundMessage diff --git a/tests/cli/test_gateway_runtime.py b/tests/cli/test_gateway_runtime.py index 6748e2a79..14c20f973 100644 --- a/tests/cli/test_gateway_runtime.py +++ b/tests/cli/test_gateway_runtime.py @@ -67,7 +67,7 @@ async def test_runtime_tasks_cancelled_before_resources_closed() -> None: task = asyncio.create_task(_cancellable_task(events)) await asyncio.sleep(0) # let the task start (cancellation pre-start skips its body) - await _close_gateway_runtime(agent, channels, [task], None, False) + await _close_gateway_runtime(agent, channels, [task], None) assert events == ["cancelled", "close_mcp"] # cancel happens before close assert channels.stopped == 1 @@ -86,7 +86,7 @@ async def test_pending_background_work_is_drained_before_close_returns() -> None agent.background = asyncio.create_task(background_work()) - await _close_gateway_runtime(agent, channels, [], None, False) + await _close_gateway_runtime(agent, channels, [], None) assert done["done"] is True assert agent.close_calls == 1 @@ -98,14 +98,14 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None: events: list[str] = [] task = asyncio.create_task(_stubborn_task(events)) await asyncio.sleep(0) # let the task start (cancellation pre-start skips its body) + runtime_tasks = asyncio.gather(task) start = time.monotonic() await _close_gateway_runtime( agent, channels, [task], - None, - False, + runtime_tasks, task_wait_timeout=0.05, ) elapsed = time.monotonic() - start @@ -113,7 +113,8 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None: await asyncio.sleep(0) # let the swallowed cancellation handler run assert "swallowed" in events # task was cancelled, then refused to die - assert not task.done() # still running despite cancellation + assert task.done() # the timed-out task received a second cancellation + assert runtime_tasks.done() assert agent.close_calls == 1 # resources still closed underneath it assert elapsed < 1.0 # bounded, not held open by the stubborn task @@ -124,7 +125,7 @@ async def test_hanging_close_is_bounded_and_does_not_raise() -> None: channels = _FakeChannels() start = time.monotonic() - await _close_gateway_runtime(agent, channels, [], None, False, close_timeout=0.05) + await _close_gateway_runtime(agent, channels, [], None, close_timeout=0.05) elapsed = time.monotonic() - start assert agent.close_calls == 1 @@ -137,7 +138,7 @@ async def test_failing_close_is_logged_but_shutdown_proceeds() -> None: agent.raise_on_close = True channels = _FakeChannels() - await _close_gateway_runtime(agent, channels, [], None, False) + await _close_gateway_runtime(agent, channels, [], None) assert agent.close_calls == 1 assert channels.stopped == 1 # teardown continued past the failure @@ -148,20 +149,20 @@ async def test_duplicate_cleanup_is_idempotent() -> None: channels = _FakeChannels() task = asyncio.create_task(_cancellable_task([])) - await _close_gateway_runtime(agent, channels, [task], None, False) - await _close_gateway_runtime(agent, channels, [task], None, False) + await _close_gateway_runtime(agent, channels, [task], None) + await _close_gateway_runtime(agent, channels, [task], None) assert agent.close_calls == 2 # second pass is a clean no-op assert channels.stopped == 2 assert task.cancelled() -async def test_runtime_tasks_gather_is_awaited_when_not_drained() -> None: +async def test_finished_runtime_tasks_gather_is_retrieved() -> None: agent = _FakeAgent() channels = _FakeChannels() runtime_tasks = asyncio.gather(asyncio.sleep(0)) - await _close_gateway_runtime(agent, channels, [], runtime_tasks, False) + await _close_gateway_runtime(agent, channels, [], runtime_tasks) assert runtime_tasks.done() assert agent.close_calls == 1 @@ -173,7 +174,7 @@ async def test_cancelled_runtime_tasks_gather_does_not_raise() -> None: runtime_tasks = asyncio.gather(asyncio.sleep(3600)) runtime_tasks.cancel() - await _close_gateway_runtime(agent, channels, [], runtime_tasks, False) + await _close_gateway_runtime(agent, channels, [], runtime_tasks) assert runtime_tasks.done() # the cancelled gather was awaited without raising assert agent.close_calls == 1