refactor: move MCP lifecycle out of AgentLoop (#5343)

tests/cli/test_agent_interactive.py

@@ -64,7 +64,7 @@ def test_interactive_agent_routes_a_complete_user_turn(
         def __init__(self, bus) -> None:
             self.bus = bus
             self.stopped = asyncio.Event()
-            self.close_mcp_calls = 0
+            self.aclose_calls = 0
 
         async def run(self) -> None:
             message = await self.bus.consume_inbound()
@@ -97,8 +97,8 @@ def test_interactive_agent_routes_a_complete_user_turn(
         def stop(self) -> None:
             self.stopped.set()
 
-        async def close_mcp(self) -> None:
-            self.close_mcp_calls += 1
+        async def aclose(self) -> None:
+            self.aclose_calls += 1
 
     read_input = AsyncMock(side_effect=["hello nanobot", "exit"])
     print_response = MagicMock()
@@ -136,7 +136,7 @@ def test_interactive_agent_routes_a_complete_user_turn(
     assert inbound.metadata == {"_wants_stream": True}
     loop = seen["loop"]
     assert isinstance(loop, _AgentLoop)
-    assert loop.close_mcp_calls == 1
+    assert loop.aclose_calls == 1
     assert len(renderers) == 1
     renderer = renderers[0]
     assert isinstance(renderer, _Renderer)

tests/cli/test_commands.py

@@ -1543,7 +1543,7 @@ def mock_agent_runtime(tmp_path):
         agent_loop.process_direct = AsyncMock(
             return_value=OutboundMessage(channel="cli", chat_id="direct", content="mock-response"),
         )
-        agent_loop.close_mcp = AsyncMock(return_value=None)
+        agent_loop.aclose = AsyncMock(return_value=None)
         mock_from_config.return_value = agent_loop
 
         yield {
@@ -1621,7 +1621,7 @@ def test_agent_config_sets_active_path(monkeypatch, tmp_path: Path) -> None:
         async def process_direct(self, *_args, **_kwargs):
             return OutboundMessage(channel="cli", chat_id="direct", content="ok")
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
     monkeypatch.setattr("nanobot.cli.agent.AgentLoop", _FakeAgentLoop)
@@ -1662,7 +1662,7 @@ def test_agent_uses_workspace_directory_for_cron_store(monkeypatch, tmp_path: Pa
         async def process_direct(self, *_args, **_kwargs):
             return OutboundMessage(channel="cli", chat_id="direct", content="ok")
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
     monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
@@ -1712,7 +1712,7 @@ def test_agent_workspace_override_does_not_migrate_legacy_cron(
         async def process_direct(self, *_args, **_kwargs):
             return OutboundMessage(channel="cli", chat_id="direct", content="ok")
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
     monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
@@ -1768,7 +1768,7 @@ def test_agent_custom_config_workspace_does_not_migrate_legacy_cron(
         async def process_direct(self, *_args, **_kwargs):
             return OutboundMessage(channel="cli", chat_id="direct", content="ok")
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
     monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
@@ -2062,7 +2062,7 @@ def test_heartbeat_empty_response_still_retains_recent_messages(
         async def process_direct(self, *_args, **_kwargs):
             return SimpleNamespace(content="")
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
         async def run(self) -> None:
@@ -2738,10 +2738,7 @@ def _patch_serve_runtime(monkeypatch, config: Config, seen: dict[str, object]) -
         def __init__(self, **kwargs) -> None:
             seen["workspace"] = kwargs["workspace"]
 
-        async def _connect_mcp(self) -> None:
-            return None
-
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
     def _fake_create_app(
@@ -2749,11 +2746,13 @@ def _patch_serve_runtime(monkeypatch, config: Config, seen: dict[str, object]) -
         model_name: str,
         request_timeout: float,
         api_key: str = "",
+        prepare_agent=None,
     ):
         seen["agent_loop"] = agent_loop
         seen["model_name"] = model_name
         seen["request_timeout"] = request_timeout
         seen["api_key"] = api_key
+        seen["prepare_agent"] = prepare_agent
         return _FakeApiApp()
 
     def _fake_run_app(api_app, host: str, port: int, print):
@@ -2914,7 +2913,7 @@ def test_gateway_unbound_agent_cron_is_skipped(
         async def submit_cron_turn(self, _msg: InboundMessage):
             raise AssertionError("unbound cron job must not run as a bound cron turn")
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
         async def run(self) -> None:
@@ -3033,7 +3032,7 @@ def test_gateway_bound_cron_runs_as_session_turn(
                 content="Checked the repo.",
             )
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
         async def run(self) -> None:
@@ -3253,7 +3252,7 @@ def test_gateway_local_trigger_queue_submits_agent_turns(
             self.runtime_resolver.invalidate.assert_called_once_with()
             await asyncio.Event().wait()
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
         def stop(self) -> None:
@@ -3499,7 +3498,7 @@ def test_gateway_health_endpoint_binds_and_serves_expected_responses(
         async def run(self) -> None:
             await asyncio.Event().wait()
 
-        async def close_mcp(self) -> None:
+        async def aclose(self) -> None:
             return None
 
         def stop(self) -> None:
@@ -3668,7 +3667,7 @@ def test_gateway_health_endpoint_binds_and_serves_expected_responses(
         assert timed_out_writer.output == b""
 
 
-def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
+def test_gateway_agent_task_owns_initial_mcp_provider_close(
     monkeypatch,
     tmp_path: Path,
 ) -> None:
@@ -3696,17 +3695,41 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
             return None
 
         async def run(self) -> None:
+            seen["agent_task"] = asyncio.current_task()
             try:
                 await asyncio.Event().wait()
             finally:
                 seen["agent_task_cleaned_up"] = True
 
-        async def close_mcp(self) -> None:
-            raise AssertionError("gateway must not close MCP from the outer task")
+        async def aclose(self) -> None:
+            seen["agent_closed"] = True
 
         def stop(self) -> None:
             seen["agent_stopped"] = True
 
+    class _FakeMCPProvider:
+        def __init__(self) -> None:
+            self.connect_task: asyncio.Task | None = None
+            self.close_tasks: list[asyncio.Task | None] = []
+
+        @classmethod
+        def from_config(cls, _config, _registry):
+            provider = cls()
+            seen["mcp_provider"] = provider
+            return provider
+
+        async def connect(self) -> None:
+            self.connect_task = asyncio.current_task()
+
+        async def aclose(self) -> None:
+            self.close_tasks.append(asyncio.current_task())
+
+        def runtime_status(self) -> dict[str, str]:
+            return {}
+
+        async def reload(self) -> dict[str, object]:
+            return {"ok": True}
+
     class _FakeChannelManager:
         def __init__(self, _config, _bus, **_kwargs) -> None:
             self.enabled_channels = ["telegram"]
@@ -3753,6 +3776,7 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
         session_manager=lambda _workspace: object(),
     )
     monkeypatch.setattr("nanobot.cli.gateway_runtime.AgentLoop", _FakeAgentLoop)
+    monkeypatch.setattr("nanobot.cli.gateway_runtime.MCPProvider", _FakeMCPProvider)
     monkeypatch.setattr("nanobot.channels.manager.ChannelManager", _FakeChannelManager)
     monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCronService)
     monkeypatch.setattr("asyncio.start_server", _fake_start_server)
@@ -3761,9 +3785,15 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
 
     assert result.exit_code == 0
     assert seen["agent_stopped"] is True
+    assert seen["agent_closed"] is True
     assert seen["agent_task_cleaned_up"] is True
     assert seen["channels_stopped"] is True
     assert seen["cron_stopped"] is True
+    mcp_provider = seen["mcp_provider"]
+    assert isinstance(mcp_provider, _FakeMCPProvider)
+    assert mcp_provider.connect_task is seen["agent_task"]
+    assert mcp_provider.close_tasks[0] is mcp_provider.connect_task
+    assert len(mcp_provider.close_tasks) == 2
 
 
 def test_gateway_shutdown_event_exits_forever_runtime_tasks(
@@ -3800,8 +3830,8 @@ def test_gateway_shutdown_event_exits_forever_runtime_tasks(
             finally:
                 seen["agent_task_cleaned_up"] = True
 
-        async def close_mcp(self) -> None:
-            raise AssertionError("gateway must not close MCP from the outer task")
+        async def aclose(self) -> None:
+            seen["agent_closed"] = True
 
         def stop(self) -> None:
             seen["agent_stopped"] = True
@@ -3881,6 +3911,7 @@ def test_gateway_shutdown_event_exits_forever_runtime_tasks(
 
     assert result.exit_code == 0
     assert seen["agent_stopped"] is True
+    assert seen["agent_closed"] is True
     assert seen["agent_task_cleaned_up"] is True
     assert seen["channel_task_cleaned_up"] is True
     assert seen["channels_stopped"] is True

tests/cli/test_gateway_runtime.py

@@ -22,7 +22,7 @@ class _FakeAgent:
         self.raise_on_close = False
         self.background: asyncio.Task[None] | None = None
 
-    async def close_mcp(self) -> None:
+    async def aclose(self) -> None:
         self.close_calls += 1
         if self.hang_on_close:
             await asyncio.sleep(3600)
@@ -30,7 +30,7 @@ class _FakeAgent:
             raise RuntimeError("cleanup exploded")
         if self.background is not None:
             await self.background
-        self.events.append("close_mcp")
+        self.events.append("aclose")
 
 
 class _FakeChannels:
@@ -43,6 +43,16 @@ class _FakeChannels:
         self.events.append("channels_stopped")
 
 
+class _FakeMCPProvider:
+    def __init__(self, events: list[str] | None = None) -> None:
+        self.close_calls = 0
+        self.events = events if events is not None else []
+
+    async def aclose(self) -> None:
+        self.close_calls += 1
+        self.events.append("mcp_closed")
+
+
 async def _cancellable_task(events: list[str]) -> None:
     try:
         await asyncio.sleep(3600)
@@ -64,13 +74,14 @@ async def _stubborn_task(events: list[str]) -> None:
 async def test_runtime_tasks_cancelled_before_resources_closed() -> None:
     events: list[str] = []
     agent = _FakeAgent(events)
+    provider = _FakeMCPProvider(events)
     channels = _FakeChannels()
     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)
+    await _close_gateway_runtime(agent, provider, channels, [task], None)
 
-    assert events == ["cancelled", "close_mcp"]  # cancel happens before close
+    assert events == ["cancelled", "aclose", "mcp_closed"]
     assert channels.stopped == 1
     assert agent.close_calls == 1
     assert task.cancelled()
@@ -78,6 +89,7 @@ async def test_runtime_tasks_cancelled_before_resources_closed() -> None:
 
 async def test_pending_background_work_is_drained_before_close_returns() -> None:
     agent = _FakeAgent()
+    provider = _FakeMCPProvider()
     channels = _FakeChannels()
     done: dict[str, bool] = {"done": False}
 
@@ -87,7 +99,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)
+    await _close_gateway_runtime(agent, provider, channels, [], None)
 
     assert done["done"] is True
     assert agent.close_calls == 1
@@ -95,6 +107,7 @@ async def test_pending_background_work_is_drained_before_close_returns() -> None
 
 async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
     agent = _FakeAgent()
+    provider = _FakeMCPProvider()
     channels = _FakeChannels()
     events: list[str] = []
     task = asyncio.create_task(_stubborn_task(events))
@@ -104,6 +117,7 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
     start = time.monotonic()
     await _close_gateway_runtime(
         agent,
+        provider,
         channels,
         [task],
         runtime_tasks,
@@ -117,70 +131,88 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
     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 provider.close_calls == 1
     assert elapsed < 1.0  # bounded, not held open by the stubborn task
 
 
 async def test_hanging_close_is_bounded_and_does_not_raise() -> None:
     agent = _FakeAgent()
+    provider = _FakeMCPProvider()
     agent.hang_on_close = True
     channels = _FakeChannels()
 
     start = time.monotonic()
-    await _close_gateway_runtime(agent, channels, [], None, close_timeout=0.05)
+    await _close_gateway_runtime(
+        agent,
+        provider,
+        channels,
+        [],
+        None,
+        close_timeout=0.05,
+    )
     elapsed = time.monotonic() - start
 
     assert agent.close_calls == 1
+    assert provider.close_calls == 1
     assert channels.stopped == 1
     assert elapsed < 1.0
 
 
 async def test_failing_close_is_logged_but_shutdown_proceeds() -> None:
     agent = _FakeAgent()
+    provider = _FakeMCPProvider()
     agent.raise_on_close = True
     channels = _FakeChannels()
 
-    await _close_gateway_runtime(agent, channels, [], None)
+    await _close_gateway_runtime(agent, provider, channels, [], None)
 
     assert agent.close_calls == 1
+    assert provider.close_calls == 1
     assert channels.stopped == 1  # teardown continued past the failure
 
 
 async def test_duplicate_cleanup_is_idempotent() -> None:
     agent = _FakeAgent()
+    provider = _FakeMCPProvider()
     channels = _FakeChannels()
     task = asyncio.create_task(_cancellable_task([]))
 
-    await _close_gateway_runtime(agent, channels, [task], None)
-    await _close_gateway_runtime(agent, channels, [task], None)
+    await _close_gateway_runtime(agent, provider, channels, [task], None)
+    await _close_gateway_runtime(agent, provider, channels, [task], None)
 
     assert agent.close_calls == 2  # second pass is a clean no-op
+    assert provider.close_calls == 2
     assert channels.stopped == 2
     assert task.cancelled()
 
 
 async def test_finished_runtime_tasks_gather_is_retrieved() -> None:
     agent = _FakeAgent()
+    provider = _FakeMCPProvider()
     channels = _FakeChannels()
     finished = asyncio.get_running_loop().create_future()
     finished.set_result(None)
     runtime_tasks = asyncio.gather(finished)
     await asyncio.sleep(0)  # let the gather observe the finished child
 
-    await _close_gateway_runtime(agent, channels, [], runtime_tasks)
+    await _close_gateway_runtime(agent, provider, channels, [], runtime_tasks)
 
     assert runtime_tasks.done()
     assert agent.close_calls == 1
+    assert provider.close_calls == 1
 
 
 async def test_cancelled_runtime_tasks_gather_does_not_raise() -> None:
     agent = _FakeAgent()
+    provider = _FakeMCPProvider()
     channels = _FakeChannels()
     runtime_tasks = asyncio.gather(asyncio.sleep(3600))
     runtime_tasks.cancel()
 
-    await _close_gateway_runtime(agent, channels, [], runtime_tasks)
+    await _close_gateway_runtime(agent, provider, channels, [], runtime_tasks)
     with suppress(asyncio.CancelledError):
         await runtime_tasks  # settle the cancelled gather without raising
 
     assert runtime_tasks.done()  # the cancelled gather was awaited without raising
     assert agent.close_calls == 1
+    assert provider.close_calls == 1