diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 9b744e04..1d1de4ab 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -873,89 +873,93 @@ class AgentLoop: async def run(self) -> None: """Run the agent loop, dispatching messages as tasks to stay responsive to /stop.""" self._running = True - await self._connect_mcp() - logger.info("Agent loop started") + try: + await self._connect_mcp() + logger.info("Agent loop started") - while self._running: - try: - msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0) - except asyncio.TimeoutError: - self.auto_compact.check_expired( - self._schedule_background, - active_session_keys=self._pending_queues.keys(), - ) - continue - except asyncio.CancelledError: - # Preserve real task cancellation so shutdown can complete cleanly. - # Only ignore non-task CancelledError signals that may leak from integrations. - if not self._running or asyncio.current_task().cancelling(): - raise - continue - except Exception as e: - logger.warning("Error consuming inbound message: {}, continuing...", e) - continue + while self._running: + try: + msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0) + except asyncio.TimeoutError: + self.auto_compact.check_expired( + self._schedule_background, + active_session_keys=self._pending_queues.keys(), + ) + continue + except asyncio.CancelledError: + # Preserve real task cancellation so shutdown can complete cleanly. + # Only ignore non-task CancelledError signals that may leak from integrations. + if not self._running or asyncio.current_task().cancelling(): + raise + continue + except Exception as e: + logger.warning("Error consuming inbound message: {}, continuing...", e) + continue - raw = msg.content.strip() - effective_key = self._effective_session_key(msg) - if await agent_context.handle_runtime_control(self, msg, self.tools): - continue - if self.commands.is_priority(raw): - await self._dispatch_command_inline( - msg, effective_key, raw, - self.commands.dispatch_priority, - ) - continue - if self._cron_turns.defer_if_active( - msg, - session_key=effective_key, - active_session_keys=self._pending_queues.keys(), - ): - logger.info( - "Deferred cron turn for active session {}", - effective_key, - ) - continue - # If this session already has an active pending queue (i.e. a task - # is processing this session), route the message there for mid-turn - # injection instead of creating a competing task. - if effective_key in self._pending_queues: - # Non-priority commands must not be queued for injection; - # dispatch them directly (same pattern as priority commands). - if self.commands.is_dispatchable_command(raw): + raw = msg.content.strip() + effective_key = self._effective_session_key(msg) + if await agent_context.handle_runtime_control(self, msg, self.tools): + continue + if self.commands.is_priority(raw): await self._dispatch_command_inline( msg, effective_key, raw, - self.commands.dispatch, + self.commands.dispatch_priority, ) continue - pending_msg = msg - if effective_key != msg.session_key: - pending_msg = dataclasses.replace( - msg, - session_key_override=effective_key, - ) - try: - self._pending_queues[effective_key].put_nowait(pending_msg) - except asyncio.QueueFull: - logger.warning( - "Pending queue full for session {}, falling back to queued task", - effective_key, - ) - else: + if self._cron_turns.defer_if_active( + msg, + session_key=effective_key, + active_session_keys=self._pending_queues.keys(), + ): logger.info( - "Routed follow-up message to pending queue for session {}", + "Deferred cron turn for active session {}", effective_key, ) continue - # Compute the effective session key before dispatching - # This ensures /stop command can find tasks correctly when unified session is enabled - task = asyncio.create_task(self._dispatch(msg)) - self._active_tasks.setdefault(effective_key, []).append(task) - task.add_done_callback( - lambda t, k=effective_key: self._active_tasks.get(k, []) - and self._active_tasks[k].remove(t) - if t in self._active_tasks.get(k, []) - else None - ) + # If this session already has an active pending queue (i.e. a task + # is processing this session), route the message there for mid-turn + # injection instead of creating a competing task. + if effective_key in self._pending_queues: + # Non-priority commands must not be queued for injection; + # dispatch them directly (same pattern as priority commands). + if self.commands.is_dispatchable_command(raw): + await self._dispatch_command_inline( + msg, effective_key, raw, + self.commands.dispatch, + ) + continue + pending_msg = msg + if effective_key != msg.session_key: + pending_msg = dataclasses.replace( + msg, + session_key_override=effective_key, + ) + try: + self._pending_queues[effective_key].put_nowait(pending_msg) + except asyncio.QueueFull: + logger.warning( + "Pending queue full for session {}, falling back to queued task", + effective_key, + ) + else: + logger.info( + "Routed follow-up message to pending queue for session {}", + effective_key, + ) + continue + # Compute the effective session key before dispatching + # This ensures /stop command can find tasks correctly when unified session is enabled + task = asyncio.create_task(self._dispatch(msg)) + self._active_tasks.setdefault(effective_key, []).append(task) + task.add_done_callback( + lambda t, k=effective_key: self._active_tasks.get(k, []) + and self._active_tasks[k].remove(t) + if t in self._active_tasks.get(k, []) + else None + ) + finally: + # MCP stdio transports use AnyIO cancel scopes; close them from the task that opened them. + await self.close_mcp() async def _dispatch(self, msg: InboundMessage) -> None: """Process a message: per-session serial, cross-session concurrent.""" diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index 3388b63b..8caa714c 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -1105,16 +1105,23 @@ def _run_gateway( console.print(f"[yellow]Could not open browser ({e}); visit {open_browser_url}[/yellow]") async def run(): + tasks: list[asyncio.Task] = [] try: await cron.start() tasks = [ - agent.run(), - channels.start_all(), + asyncio.create_task(agent.run(), name="nanobot-agent-loop"), + asyncio.create_task(channels.start_all(), name="nanobot-channels"), ] if health_server_enabled: - tasks.append(_health_server(config.gateway.host, port)) + tasks.append(asyncio.create_task( + _health_server(config.gateway.host, port), + name="nanobot-health-server", + )) if open_browser_url: - tasks.append(_open_browser_when_ready()) + tasks.append(asyncio.create_task( + _open_browser_when_ready(), + name="nanobot-open-browser", + )) await asyncio.gather(*tasks) except KeyboardInterrupt: console.print("\nShutting down...") @@ -1124,9 +1131,13 @@ def _run_gateway( console.print("\n[red]Error: Gateway crashed unexpectedly[/red]") console.print(traceback.format_exc()) finally: - await agent.close_mcp() cron.stop() agent.stop() + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) await channels.stop_all() # Flush all cached sessions to durable storage before exit. # This prevents data loss on filesystems with write-back diff --git a/tests/agent/test_mcp_connection.py b/tests/agent/test_mcp_connection.py index d70c6b37..45f4bd47 100644 --- a/tests/agent/test_mcp_connection.py +++ b/tests/agent/test_mcp_connection.py @@ -135,6 +135,43 @@ async def test_connect_mcp_retries_when_no_servers_connect(tmp_path, monkeypatch assert loop._mcp_stacks == {} +@pytest.mark.asyncio +async def test_agent_loop_run_closes_mcp_from_connection_owner_task( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +): + loop = _make_loop(tmp_path, mcp_servers={"playwright": object()}) + connected = asyncio.Event() + owner_tasks: list[asyncio.Task | None] = [] + closed_tasks: list[asyncio.Task | None] = [] + + class _OwnerCheckedStack: + def __init__(self) -> None: + self.owner = asyncio.current_task() + owner_tasks.append(self.owner) + + async def aclose(self) -> None: + closed_tasks.append(asyncio.current_task()) + assert asyncio.current_task() is self.owner + + async def _fake_connect(servers, _registry): + stacks = {name: _OwnerCheckedStack() for name in servers} + connected.set() + return stacks + + monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) + + task = asyncio.create_task(loop.run()) + await asyncio.wait_for(connected.wait(), timeout=1) + loop.stop() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert owner_tasks + assert closed_tasks == owner_tasks + assert loop._mcp_stacks == {} + + @pytest.mark.asyncio async def test_reload_mcp_servers_adds_and_removes_tools_without_restart( tmp_path, diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index ef462348..7825bc5d 100644 --- a/tests/cli/test_commands.py +++ b/tests/cli/test_commands.py @@ -1960,6 +1960,103 @@ def test_gateway_health_endpoint_binds_and_serves_expected_responses( assert missing_response.endswith("\r\n\r\nNot Found") +def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup( + monkeypatch, + tmp_path: Path, +) -> None: + config_file = _write_instance_config(tmp_path) + config = Config() + config.gateway.port = 18791 + seen: dict[str, object] = {} + + class _FakeSessionManager: + def flush_all(self) -> int: + return 0 + + class _FakeAgentLoop: + @classmethod + def from_config(cls, config, bus=None, **extra): + return cls(**extra) + + def __init__(self, **_kwargs) -> None: + self.model = "test-model" + self.provider = object() + self.sessions = _FakeSessionManager() + + def llm_runtime(self) -> None: + return None + + async def run(self) -> None: + 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") + + def stop(self) -> None: + seen["agent_stopped"] = True + + class _FakeChannelManager: + def __init__(self, _config, _bus, **_kwargs) -> None: + self.enabled_channels = ["telegram"] + + async def start_all(self) -> None: + await asyncio.Event().wait() + + async def stop_all(self) -> None: + seen["channels_stopped"] = True + + class _FakeCronService: + def __init__(self, _store_path: Path) -> None: + self.on_job = None + + async def start(self) -> None: + return None + + def stop(self) -> None: + seen["cron_stopped"] = True + + def status(self) -> dict[str, int]: + return {"jobs": 0} + + def register_system_job(self, _job) -> None: + return None + + class _FakeServer: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> bool: + return False + + async def serve_forever(self) -> None: + raise _StopGatewayError("stop") + + async def _fake_start_server(_handler, _host: str, _port: int): + return _FakeServer() + + _patch_cli_command_runtime( + monkeypatch, + config, + message_bus=lambda: object(), + session_manager=lambda _workspace: object(), + ) + monkeypatch.setattr("nanobot.cli.commands.AgentLoop", _FakeAgentLoop) + monkeypatch.setattr("nanobot.channels.manager.ChannelManager", _FakeChannelManager) + monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCronService) + monkeypatch.setattr("asyncio.start_server", _fake_start_server) + + result = runner.invoke(app, ["gateway", "--config", str(config_file)]) + + assert result.exit_code == 0 + assert seen["agent_stopped"] is True + assert seen["agent_task_cleaned_up"] is True + assert seen["channels_stopped"] is True + assert seen["cron_stopped"] is True + + def test_serve_uses_api_config_defaults_and_workspace_override( monkeypatch, tmp_path: Path ) -> None: