fix(gateway): make resource teardown cancellation-safe

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,

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

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

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