From 85e803e7c1117dd62df0a433c608f0b36fd94aa7 Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Wed, 12 Aug 2026 01:16:46 +0800 Subject: [PATCH 1/6] refactor: move MCP lifecycle out of AgentLoop --- nanobot/agent/context.py | 20 +- nanobot/agent/loop.py | 45 +- nanobot/agent/tools/mcp.py | 749 +++++++++--------- nanobot/agent/tools/self.py | 2 +- nanobot/api/server.py | 33 +- nanobot/bus/events.py | 1 - nanobot/channels/manager.py | 5 +- .../tests/test_websocket_http_routes.py | 15 +- nanobot/cli/agent.py | 61 +- nanobot/cli/commands.py | 13 +- nanobot/cli/gateway_runtime.py | 56 +- nanobot/nanobot.py | 28 +- nanobot/webui/gateway_services.py | 4 +- nanobot/webui/settings_routes.py | 37 +- nanobot/webui/ws_http.py | 4 +- tests/agent/conftest.py | 3 - tests/agent/test_auto_compact.py | 80 +- tests/agent/test_consolidate_offset.py | 10 +- .../test_loop_direct_websocket_status.py | 1 - tests/agent/test_loop_save_turn.py | 3 - tests/agent/test_loop_session_policy.py | 3 +- tests/agent/test_loop_tool_context.py | 17 + tests/agent/test_mcp_connection.py | 221 +++--- tests/agent/test_mcp_reconnect_crash.py | 41 +- tests/agent/test_task_cancel.py | 25 +- tests/agent/tools/test_self_tool.py | 21 - tests/cli/test_agent_interactive.py | 8 +- tests/cli/test_commands.py | 69 +- tests/cli/test_gateway_runtime.py | 54 +- tests/test_api_attachment.py | 3 +- tests/test_api_stream.py | 24 +- tests/test_nanobot_facade.py | 17 +- tests/test_openai_api.py | 66 +- tests/tools/test_exec_session_tools.py | 24 +- tests/tools/test_mcp_tool.py | 48 +- tests/webui/test_settings_routes.py | 32 +- 36 files changed, 1004 insertions(+), 839 deletions(-) diff --git a/nanobot/agent/context.py b/nanobot/agent/context.py index 44169ba9b3f..d659f2751d4 100644 --- a/nanobot/agent/context.py +++ b/nanobot/agent/context.py @@ -42,29 +42,11 @@ def session_extra(metadata: Mapping[str, Any] | None) -> dict[str, Any]: ) -async def connect_mcp(state: Any, tools: ToolRegistry) -> None: - await mcp_tools.connect_missing_servers(state, tools) - - -def mcp_runtime_status(state: Any) -> dict[str, mcp_tools.MCPRuntimeStatus]: - return mcp_tools.runtime_status(state) - - -async def close_mcp(state: Any) -> None: - await mcp_tools.close_mcp_servers(state) - - async def handle_runtime_control(state: Any, msg: InboundMessage, tools: ToolRegistry) -> bool: if msg.metadata.get(INBOUND_META_RUNTIME_CONTROL) == RUNTIME_CONTROL_SESSION_DISCARD: await state.discard_session(msg.session_key) return True - for handler in ( - image_generation_tools.handle_runtime_control, - mcp_tools.handle_runtime_control, - ): - if await handler(state, msg, tools): - return True - return False + return await image_generation_tools.handle_runtime_control(state, msg, tools) class ContextBuilder: diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 35514c5e7a4..aae9d1a376e 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -95,11 +95,9 @@ ) if TYPE_CHECKING: - from nanobot.agent.tools.mcp import MCPConnection, MCPRuntimeStatus from nanobot.config.schema import ( ChannelsConfig, Config, - MCPServerConfig, ProviderConfig, ToolsConfig, ) @@ -271,7 +269,7 @@ def __init__( cron_service: CronService | None = None, restrict_to_workspace: bool = False, session_manager: SessionManager | None = None, - mcp_servers: dict[str, MCPServerConfig] | None = None, + tool_registry: ToolRegistry | None = None, channels_config: ChannelsConfig | None = None, timezone: str | None = None, session_ttl_minutes: int = 0, @@ -379,7 +377,7 @@ def __init__( self.context = ContextBuilder(workspace, timezone=timezone, disabled_skills=disabled_skills) self.sessions = session_manager or SessionManager(workspace) self.sessions.set_file_cap_archiver(self.context.memory.raw_archive) - self.tools = ToolRegistry() + self.tools = tool_registry if tool_registry is not None else ToolRegistry() # One file-read/write tracker per logical session. The tool registry is # shared by this loop, so tools resolve the active state via contextvars. self._file_state_store = FileStateStore() @@ -399,15 +397,11 @@ def __init__( ) self._unified_session = unified_session self._running = False - self._mcp_servers = mcp_servers or {} - self._mcp_stacks: dict[str, MCPConnection] = {} - self._mcp_runtime_statuses: dict[str, MCPRuntimeStatus] = {} - self._mcp_connecting = False self._runtime_context_providers: list[RuntimeContextProvider] = [] self._active_tasks: dict[str, set[asyncio.Task[Any]]] = {} self._discarding_sessions: set[str] = set() self._background_tasks: set[asyncio.Task[Any]] = set() - self._close_mcp_lock = asyncio.Lock() + self._close_lock = asyncio.Lock() self._session_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = ( weakref.WeakValueDictionary() ) @@ -486,8 +480,6 @@ def from_config( config, provider_snapshot_loader, ) - from nanobot.agent.plugins import agent_plugin_mcp_servers - return cls( bus=bus, provider=provider, @@ -502,7 +494,6 @@ def from_config( provider_retry_mode=defaults.provider_retry_mode, tool_hint_max_length=defaults.tool_hint_max_length, restrict_to_workspace=config.tools.restrict_to_workspace, - mcp_servers=agent_plugin_mcp_servers(config.workspace_path, config.tools.mcp_servers), channels_config=config.channels, timezone=defaults.timezone, unified_session=defaults.unified_session, @@ -643,14 +634,6 @@ def _register_default_tools( logger.info("Registered {} tools: {}", len(registered), registered) - async def _connect_mcp(self) -> None: - """Connect configured MCP servers.""" - await agent_context.connect_mcp(self, self.tools) - - def mcp_runtime_status(self) -> dict[str, MCPRuntimeStatus]: - """Return connection state learned from real MCP runtime attempts.""" - return agent_context.mcp_runtime_status(self) - def register_runtime_context_provider( self, provider: RuntimeContextProvider, @@ -1162,7 +1145,6 @@ async def run(self) -> None: """Run the agent loop, dispatching messages as tasks to stay responsive to /stop.""" self._running = True try: - await self._connect_mcp() logger.info("Agent loop started") while self._running: @@ -1253,8 +1235,7 @@ async def run(self) -> None: active_tasks.add(task) task.add_done_callback(active_tasks.discard) finally: - # MCP stdio transports use AnyIO cancel scopes; close them from the task that opened them. - await self.close_mcp() + await self.aclose() async def _dispatch(self, msg: InboundMessage) -> None: """Process a message: per-session serial, cross-session concurrent.""" @@ -1372,24 +1353,24 @@ async def _dispatch(self, msg: InboundMessage) -> None: await delivery.idle() await self._publish_next_deferred_automation_turn(session_key) - async def close_mcp(self) -> None: - """Stop active work, then close exec, subagent, and MCP resources. + async def aclose(self) -> None: + """Stop active work, then close resources owned by the agent loop. 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 + # The loop closes itself from ``run()`` while application 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) + # tear down the same resources concurrently. + close_lock = getattr(self, "_close_lock", None) if close_lock is None: - close_lock = self._close_mcp_lock = asyncio.Lock() + close_lock = self._close_lock = asyncio.Lock() async with close_lock: - await self._close_mcp_unlocked() + await self._aclose_unlocked() - async def _close_mcp_unlocked(self) -> None: + async def _aclose_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}) @@ -1412,7 +1393,6 @@ async def _close_mcp_unlocked(self) -> None: cleanup_steps = ( self.subagents.close, self._exec_session_manager.close_all, - lambda: agent_context.close_mcp(self), ) for cleanup in cleanup_steps: try: @@ -2301,7 +2281,6 @@ async def process_direct( """Process an external message directly and return the outbound payload.""" if channel == "system": raise ValueError("channel 'system' is reserved for internal messages") - await self._connect_mcp() metadata: dict[str, Any] = {} if not persist_user_message: metadata[turn_continuation.SKIP_USER_PERSIST_META] = True diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index e590d130ce4..5e736c46800 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -1,4 +1,6 @@ -"""MCP client: connects to MCP servers and wraps their tools as native nanobot tools.""" +"""MCP client and dynamic tool-provider lifecycle.""" + +from __future__ import annotations import asyncio import hashlib @@ -7,23 +9,15 @@ import re import shutil import urllib.parse -from collections.abc import AsyncIterator, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping from contextlib import AsyncExitStack, suppress -from typing import TYPE_CHECKING, Any, Literal, Mapping, Protocol, cast -from weakref import WeakKeyDictionary +from typing import TYPE_CHECKING, Any, Literal, Protocol, cast import httpx from loguru import logger from nanobot.agent.tools.base import Tool, ToolResult from nanobot.agent.tools.registry import ToolRegistry -from nanobot.bus.events import ( - INBOUND_META_RUNTIME_CONTROL, - RUNTIME_CONTROL_ACK, - RUNTIME_CONTROL_MCP_RELOAD, - InboundMessage, -) -from nanobot.bus.queue import MessageBus from nanobot.security.network import ( PinnedDNSAsyncTransport, env_proxy_applies_to_url, @@ -39,7 +33,7 @@ from mcp.types import Tool as MCPToolDefinition from nanobot.agent.tools.mcp_oauth import MCPOAuthHandlers - from nanobot.config.schema import MCPServerConfig + from nanobot.config.schema import Config, MCPServerConfig # Transient connection errors that warrant a single retry. # These typically happen when an MCP server restarts or a network @@ -60,12 +54,9 @@ # Characters allowed in tool names by model providers (Anthropic, OpenAI, etc.). # Replace anything outside [a-zA-Z0-9_-] with underscore and collapse runs. _SANITIZE_RE = re.compile(r"_+") -_RELOAD_LOCKS: WeakKeyDictionary[Any, asyncio.Lock] = WeakKeyDictionary() _ReconnectCallback = Callable[[str, str, Tool], Awaitable[Tool | None]] +MCPServerLoader = Callable[[], Mapping[str, "MCPServerConfig"]] MCPRuntimeStatus = Literal["connecting", "connected", "failed"] -_MCP_RUNTIME_STATUSES: frozenset[MCPRuntimeStatus] = frozenset( - ("connecting", "connected", "failed") -) class MCPConnection(Protocol): @@ -492,11 +483,11 @@ class _MCPWrapperBase(Tool): """Common reconnect handling for wrappers bound to one MCP server session.""" _plugin_discoverable = False - _session: "ClientSession" + _session: ClientSession _server_name: str _name: str - def _set_mcp_connection(self, session: "ClientSession", server_name: str) -> None: + def _set_mcp_connection(self, session: ClientSession, server_name: str) -> None: self._session = session self._server_name = server_name self._reconnect: _ReconnectCallback | None = None @@ -586,9 +577,9 @@ class MCPToolWrapper(_MCPWrapperBase): def __init__( self, - session: "ClientSession", + session: ClientSession, server_name: str, - tool_def: "MCPToolDefinition", + tool_def: MCPToolDefinition, tool_timeout: int = 30, ): self._set_mcp_connection(session, server_name) @@ -748,9 +739,9 @@ class MCPResourceWrapper(_MCPWrapperBase): def __init__( self, - session: "ClientSession", + session: ClientSession, server_name: str, - resource_def: "Resource", + resource_def: Resource, resource_timeout: int = 30, ): self._set_mcp_connection(session, server_name) @@ -852,9 +843,9 @@ class MCPPromptWrapper(_MCPWrapperBase): def __init__( self, - session: "ClientSession", + session: ClientSession, server_name: str, - prompt_def: "Prompt", + prompt_def: Prompt, prompt_timeout: int = 30, ): self._set_mcp_connection(session, server_name) @@ -985,10 +976,10 @@ async def execute(self, **kwargs: Any) -> str: async def connect_mcp_servers( - mcp_servers: "dict[str, MCPServerConfig]", + mcp_servers: dict[str, MCPServerConfig], registry: ToolRegistry, *, - oauth_handlers: Mapping[str, "MCPOAuthHandlers"] | None = None, + oauth_handlers: Mapping[str, MCPOAuthHandlers] | None = None, ) -> dict[str, MCPConnection]: """Connect to configured MCP servers and register their tools, resources, prompts. @@ -1002,7 +993,7 @@ async def connect_mcp_servers( from mcp.client.streamable_http import streamable_http_client async def open_single_server( - name: str, cfg: "MCPServerConfig", server_stack: AsyncExitStack + name: str, cfg: MCPServerConfig, server_stack: AsyncExitStack ) -> bool: try: transport_type = cfg.type @@ -1244,7 +1235,7 @@ def httpx_client_factory( return False async def connect_single_server( - name: str, cfg: "MCPServerConfig" + name: str, cfg: MCPServerConfig ) -> tuple[str, MCPConnection | None]: loop = asyncio.get_running_loop() ready: asyncio.Future[bool] = loop.create_future() @@ -1301,369 +1292,375 @@ def session_extra(metadata: Mapping[str, Any] | None) -> dict[str, Any]: return {"mcp_presets": mcp_presets} if isinstance(mcp_presets, list) and mcp_presets else {} -def _runtime_status_store( - state: Any, - *, - create: bool = False, -) -> dict[str, MCPRuntimeStatus] | None: - raw_statuses: object = getattr(state, "_mcp_runtime_statuses", None) - if isinstance(raw_statuses, dict): - return cast(dict[str, MCPRuntimeStatus], raw_statuses) - if not create: - return None - statuses: dict[str, MCPRuntimeStatus] = {} - state._mcp_runtime_statuses = statuses - return statuses - - -def runtime_status(state: Any) -> dict[str, MCPRuntimeStatus]: - """Return the latest connection-attempt result for configured MCP servers.""" - statuses = _runtime_status_store(state) - raw_configured: object = getattr(state, "_mcp_servers", None) - if statuses is None or not isinstance(raw_configured, dict): - return {} - configured = cast(dict[str, Any], raw_configured) - return { - name: status - for name, status in statuses.items() - if name in configured and status in _MCP_RUNTIME_STATUSES - } +def _configured_servers(config: Config) -> dict[str, MCPServerConfig]: + from nanobot.agent.plugins import agent_plugin_mcp_servers + return agent_plugin_mcp_servers( + config.workspace_path, + config.tools.mcp_servers, + ) -def _set_runtime_status( - state: Any, - server_names: Mapping[str, Any] | set[str] | list[str] | tuple[str, ...], - status: MCPRuntimeStatus, -) -> None: - statuses = _runtime_status_store(state, create=True) - assert statuses is not None - for name in server_names: - statuses[name] = status - - -def _record_connection_result( - state: Any, - attempted: Mapping[str, Any] | set[str] | list[str] | tuple[str, ...], - connected: Mapping[str, Any] | set[str] | list[str] | tuple[str, ...], -) -> None: - attempted_names = set(attempted) - connected_names = set(connected) - _set_runtime_status(state, connected_names, "connected") - _set_runtime_status(state, attempted_names - connected_names, "failed") - - -async def connect_missing_servers(state: Any, registry: ToolRegistry) -> None: - """Connect configured MCP servers that are not currently live.""" - async with _reload_lock(state): - if getattr(state, "_mcp_closing", False): - return - configured_missing = { - name: cfg for name, cfg in state._mcp_servers.items() if name not in state._mcp_stacks - } - oauth_servers = { - name: cfg - for name, cfg in configured_missing.items() - if getattr(cfg, "auth", None) == "oauth" - } - authorization_pending: set[str] = set() - if oauth_servers: - from nanobot.agent.tools.mcp_oauth import mcp_oauth_has_credentials - authorization_pending = { - name - for name, cfg in oauth_servers.items() - if not mcp_oauth_has_credentials(name, cfg.url) - } - statuses = _runtime_status_store(state) - if statuses is not None: - for name in authorization_pending: - statuses.pop(name, None) - missing_servers = { - name: cfg - for name, cfg in configured_missing.items() - if name not in authorization_pending +def _load_current_servers() -> dict[str, MCPServerConfig]: + from nanobot.config.loader import load_config, resolve_config_env_vars + + return _configured_servers(resolve_config_env_vars(load_config())) + + +class MCPProvider: + """Own configured MCP connections and their dynamic tool registrations.""" + + def __init__( + self, + servers: Mapping[str, MCPServerConfig], + registry: ToolRegistry, + *, + server_loader: MCPServerLoader | None = None, + ) -> None: + self._servers = dict(servers) + self._registry = registry + self._server_loader = server_loader or _load_current_servers + self._connections: dict[str, MCPConnection] = {} + self._runtime_statuses: dict[str, MCPRuntimeStatus] = {} + self._lock = asyncio.Lock() + self._closing = False + + @classmethod + def from_config( + cls, + config: Config, + registry: ToolRegistry, + *, + server_loader: MCPServerLoader | None = None, + ) -> MCPProvider: + return cls( + _configured_servers(config), + registry, + server_loader=server_loader, + ) + + @property + def configured_server_names(self) -> set[str]: + return set(self._servers) + + @property + def connected_server_names(self) -> set[str]: + return set(self._connections) + + def runtime_status(self) -> dict[str, MCPRuntimeStatus]: + """Return the latest connection-attempt result for configured servers.""" + return { + name: status + for name, status in self._runtime_statuses.items() + if name in self._servers } - if state._mcp_connecting or not missing_servers: - return - state._mcp_connecting = True - _set_runtime_status(state, missing_servers, "connecting") + + def _set_runtime_status( + self, + server_names: Iterable[str], + status: MCPRuntimeStatus, + ) -> None: + for name in server_names: + self._runtime_statuses[name] = status + + def _record_connection_result( + self, + attempted: Iterable[str], + connected: Iterable[str], + ) -> None: + attempted_names = set(attempted) + connected_names = set(connected) + self._set_runtime_status(connected_names, "connected") + self._set_runtime_status(attempted_names - connected_names, "failed") + + @staticmethod + async def _close_connection(name: str, connection: MCPConnection) -> None: try: - connected = await connect_mcp_servers(missing_servers, registry) - if getattr(state, "_mcp_closing", False): - for connection in connected.values(): - await connection.aclose() - return - state._mcp_stacks.update(connected) - _record_connection_result(state, missing_servers, connected) - _attach_reconnect_handlers(state, registry, connected) - if connected: - logger.info("MCP connected servers: {}", sorted(connected)) - else: - logger.warning("No MCP servers connected successfully (will retry next message)") + await connection.aclose() except asyncio.CancelledError: if task_is_cancelling(): raise - _set_runtime_status(state, missing_servers, "failed") - logger.warning("MCP connection cancelled (will retry next message)") - except BaseException as e: - _set_runtime_status(state, missing_servers, "failed") - logger.warning("Failed to connect MCP servers (will retry next message): {}", e) - finally: - state._mcp_connecting = False - - -async def reload_servers(state: Any, registry: ToolRegistry) -> dict[str, Any]: - """Reconcile live MCP connections with the current config file.""" - async with _reload_lock(state): - if getattr(state, "_mcp_closing", False): - return { - "ok": False, - "message": "MCP connections are shutting down.", - "requires_restart": True, + logger.debug("MCP server '{}' cleanup error (can be ignored)", name) + except (RuntimeError, BaseExceptionGroup): + logger.debug("MCP server '{}' cleanup error (can be ignored)", name) + + @classmethod + async def _close_connections(cls, connections: Mapping[str, MCPConnection]) -> None: + cancellation: asyncio.CancelledError | None = None + for name, connection in connections.items(): + try: + await cls._close_connection(name, connection) + except asyncio.CancelledError as exc: + cancellation = cancellation or exc + if cancellation is not None: + raise cancellation + + async def connect(self) -> None: + """Connect configured servers that are not currently live.""" + async with self._lock: + if self._closing: + return + configured_missing = { + name: cfg + for name, cfg in self._servers.items() + if name not in self._connections } - try: - from nanobot.agent.plugins import agent_plugin_mcp_servers - from nanobot.config.loader import load_config, resolve_config_env_vars - - config = resolve_config_env_vars(load_config()) - next_servers = agent_plugin_mcp_servers( - config.workspace_path, - config.tools.mcp_servers, - ) - except Exception as exc: - logger.warning("MCP hot reload could not read config: {}", exc) - return { - "ok": False, - "message": "Could not reload MCP config. Restart nanobot to pick up changes.", - "requires_restart": True, - "error": str(exc), + oauth_servers = { + name: cfg + for name, cfg in configured_missing.items() + if cfg.auth == "oauth" } + authorization_pending: set[str] = set() + if oauth_servers: + from nanobot.agent.tools.mcp_oauth import mcp_oauth_has_credentials + + authorization_pending = { + name + for name, cfg in oauth_servers.items() + if not mcp_oauth_has_credentials(name, cfg.url) + } + for name in authorization_pending: + self._runtime_statuses.pop(name, None) + missing_servers = { + name: cfg + for name, cfg in configured_missing.items() + if name not in authorization_pending + } + if not missing_servers: + return + self._set_runtime_status(missing_servers, "connecting") + try: + connected = await connect_mcp_servers(missing_servers, self._registry) + if self._closing: + await self._close_connections(connected) + return + self._connections.update(connected) + self._record_connection_result(missing_servers, connected) + self._attach_reconnect_handlers(connected) + if connected: + logger.info("MCP connected servers: {}", sorted(connected)) + else: + logger.warning( + "No MCP servers connected successfully " + "(will retry on the next readiness check)" + ) + except asyncio.CancelledError: + if task_is_cancelling(): + raise + self._set_runtime_status(missing_servers, "failed") + logger.warning( + "MCP connection cancelled (will retry on the next readiness check)" + ) + except BaseException as exc: + self._set_runtime_status(missing_servers, "failed") + logger.warning( + "Failed to connect MCP servers " + "(will retry on the next readiness check): {}", + exc, + ) - current_servers = dict(state._mcp_servers) - current_names = set(current_servers) - next_names = set(next_servers) - from nanobot.agent.tools.mcp_oauth import mcp_oauth_has_credentials - - authorization_pending = { - name - for name, cfg in next_servers.items() - if cfg.auth == "oauth" and not mcp_oauth_has_credentials(name, cfg.url) - } - removed = sorted(current_names - next_names) - added = sorted(next_names - current_names) - changed = sorted( - name - for name in current_names & next_names - if _server_signature(current_servers[name]) != _server_signature(next_servers[name]) - ) - - tools_removed = 0 - for name in [*removed, *changed]: - tools_removed += _unregister_server_tools(registry, name) - await _close_server(state, name) - - runtime_statuses = _runtime_status_store(state) - if runtime_statuses is not None: - for name in [*removed, *authorization_pending]: - runtime_statuses.pop(name, None) - - state._mcp_servers = next_servers - retry_missing = sorted( - name - for name in next_names - if name not in state._mcp_stacks - and name not in set(added) | set(changed) - and name not in authorization_pending - ) - to_connect_names = sorted( - (set(added) | set(changed) | set(retry_missing)) - authorization_pending - ) - to_connect = {name: next_servers[name] for name in to_connect_names} - connected: dict[str, MCPConnection] = {} - if to_connect: - _set_runtime_status(state, to_connect, "connecting") - connected = await connect_mcp_servers(to_connect, registry) - if getattr(state, "_mcp_closing", False): - for connection in connected.values(): - await connection.aclose() + async def reload(self) -> dict[str, Any]: + """Reconcile live MCP connections with the current configuration.""" + async with self._lock: + if self._closing: + return self._closing_result() + try: + next_servers = dict(self._server_loader()) + except Exception as exc: + logger.warning("MCP hot reload could not read config: {}", exc) return { "ok": False, - "message": "MCP connections are shutting down.", + "message": "Could not reload MCP config. Restart nanobot to pick up changes.", "requires_restart": True, + "error": str(exc), } - state._mcp_stacks.update(connected) - _record_connection_result(state, to_connect, connected) - _attach_reconnect_handlers(state, registry, connected) - - failed = sorted(set(to_connect) - set(connected)) - unchanged = not removed and not added and not changed and not retry_missing - ok = not failed - if failed: - message = "MCP config reloaded, but some servers did not connect: " + ", ".join(failed) - elif unchanged: - message = "MCP config is already live." - elif retry_missing and not added and not changed and not removed: - message = "MCP connections refreshed without restarting nanobot." - else: - message = "MCP config reloaded without restarting nanobot." - - logger.info( - "MCP hot reload: added={} changed={} removed={} retried={} connected={} failed={} tools_removed={}", - added, - changed, - removed, - retry_missing, - sorted(connected), - failed, - tools_removed, - ) - return { - "ok": ok, - "message": message, - "added": added, - "changed": changed, - "removed": removed, - "retried": retry_missing, - "connected": sorted(state._mcp_stacks), - "configured": sorted(state._mcp_servers), - "failed": failed, - "tools_removed": tools_removed, - "requires_restart": False, - } + current_servers = dict(self._servers) + current_names = set(current_servers) + next_names = set(next_servers) + from nanobot.agent.tools.mcp_oauth import mcp_oauth_has_credentials -async def request_mcp_reload( - bus: MessageBus, - *, - timeout: float = 15.0, -) -> dict[str, Any]: - """Ask the running agent loop to reconcile live MCP connections.""" - loop = asyncio.get_running_loop() - ack: asyncio.Future[dict[str, Any]] = loop.create_future() - await bus.publish_inbound( - InboundMessage( - channel="system", - sender_id="webui-settings", - chat_id="runtime", - content=RUNTIME_CONTROL_MCP_RELOAD, - metadata={ - INBOUND_META_RUNTIME_CONTROL: RUNTIME_CONTROL_MCP_RELOAD, - RUNTIME_CONTROL_ACK: ack, - }, - ) - ) - try: - result = await asyncio.wait_for(ack, timeout=timeout) - except asyncio.TimeoutError: - return { - "ok": False, - "message": "MCP hot reload timed out. Restart nanobot to pick up changes.", - "requires_restart": True, - } - return result if isinstance(cast(object, result), dict) else { - "ok": False, - "message": "MCP hot reload returned an unexpected response.", - "requires_restart": True, - } + authorization_pending = { + name + for name, cfg in next_servers.items() + if cfg.auth == "oauth" and not mcp_oauth_has_credentials(name, cfg.url) + } + removed = sorted(current_names - next_names) + added = sorted(next_names - current_names) + changed = sorted( + name + for name in current_names & next_names + if _server_signature(current_servers[name]) + != _server_signature(next_servers[name]) + ) + tools_removed = 0 + for name in [*removed, *changed]: + tools_removed += _unregister_server_tools(self._registry, name) + await self._close_server(name) -async def handle_runtime_control(state: Any, msg: InboundMessage, registry: ToolRegistry) -> bool: - metadata = msg.metadata if isinstance(cast(object, msg.metadata), dict) else {} - control = metadata.get(INBOUND_META_RUNTIME_CONTROL) - if control != RUNTIME_CONTROL_MCP_RELOAD: - return False + for name in [*removed, *authorization_pending]: + self._runtime_statuses.pop(name, None) - ack = metadata.get(RUNTIME_CONTROL_ACK) - try: - result = await reload_servers(state, registry) - except Exception as exc: - logger.exception("MCP hot reload failed") - result = { + self._servers = next_servers + retry_missing = sorted( + name + for name in next_names + if name not in self._connections + and name not in set(added) | set(changed) + and name not in authorization_pending + ) + to_connect_names = sorted( + (set(added) | set(changed) | set(retry_missing)) + - authorization_pending + ) + to_connect = {name: next_servers[name] for name in to_connect_names} + connected: dict[str, MCPConnection] = {} + if to_connect: + self._set_runtime_status(to_connect, "connecting") + connected = await connect_mcp_servers(to_connect, self._registry) + if self._closing: + await self._close_connections(connected) + return self._closing_result() + self._connections.update(connected) + self._record_connection_result(to_connect, connected) + self._attach_reconnect_handlers(connected) + + failed = sorted(set(to_connect) - set(connected)) + unchanged = not removed and not added and not changed and not retry_missing + ok = not failed + if failed: + message = ( + "MCP config reloaded, but some servers did not connect: " + + ", ".join(failed) + ) + elif unchanged: + message = "MCP config is already live." + elif retry_missing and not added and not changed and not removed: + message = "MCP connections refreshed without restarting nanobot." + else: + message = "MCP config reloaded without restarting nanobot." + + logger.info( + "MCP hot reload: added={} changed={} removed={} retried={} " + "connected={} failed={} tools_removed={}", + added, + changed, + removed, + retry_missing, + sorted(connected), + failed, + tools_removed, + ) + return { + "ok": ok, + "message": message, + "added": added, + "changed": changed, + "removed": removed, + "retried": retry_missing, + "connected": sorted(self._connections), + "configured": sorted(self._servers), + "failed": failed, + "tools_removed": tools_removed, + "requires_restart": False, + } + + @staticmethod + def _closing_result() -> dict[str, Any]: + return { "ok": False, - "message": "MCP hot reload failed. Restart nanobot to pick up changes.", + "message": "MCP connections are shutting down.", "requires_restart": True, - "error": str(exc), } - if isinstance(ack, asyncio.Future) and not ack.done(): - cast(asyncio.Future[dict[str, Any]], ack).set_result(result) - return True - - -def _reload_lock(state: Any) -> asyncio.Lock: - try: - return _RELOAD_LOCKS[state] - except KeyError: - lock = asyncio.Lock() - _RELOAD_LOCKS[state] = lock - return lock + def _attach_reconnect_handlers(self, server_names: Iterable[str]) -> None: + async def reconnect( + server_name: str, + tool_name: str, + stale_tool: Tool, + ) -> Tool | None: + return await self._refresh_terminated_server( + server_name, + tool_name, + stale_tool, + ) -def _attach_reconnect_handlers( - state: Any, - registry: ToolRegistry, - server_names: Mapping[str, Any] | set[str] | list[str] | tuple[str, ...], -) -> None: - async def reconnect(server_name: str, tool_name: str, stale_tool: Tool) -> Tool | None: - return await _refresh_terminated_server( - state, - registry, - server_name, - tool_name, - stale_tool, - ) + for server_name in server_names: + for tool_name in list(self._registry.tool_names): + tool = self._registry.get(tool_name) + if not _tool_belongs_to_server(tool, tool_name, server_name): + continue + if isinstance(tool, _MCPWrapperBase): + tool.set_reconnect_handler(reconnect) - for server_name in server_names: - for tool_name in list(registry.tool_names): - tool = registry.get(tool_name) - if not _tool_belongs_to_server(tool, tool_name, server_name): - continue - if isinstance(tool, _MCPWrapperBase): - tool.set_reconnect_handler(reconnect) + async def _refresh_terminated_server( + self, + server_name: str, + tool_name: str, + stale_tool: Tool, + ) -> Tool | None: + async with self._lock: + if self._closing: + return None + cfg = self._servers.get(server_name) + if cfg is None: + logger.warning( + "MCP server '{}' session terminated but is no longer configured", + server_name, + ) + return None + current_tool = self._registry.get(tool_name) + if ( + current_tool is not None + and current_tool is not stale_tool + and server_name in self._connections + ): + return current_tool -async def _refresh_terminated_server( - state: Any, - registry: ToolRegistry, - server_name: str, - tool_name: str, - stale_tool: Tool, -) -> Tool | None: - async with _reload_lock(state): - if getattr(state, "_mcp_closing", False): - return None - cfg = state._mcp_servers.get(server_name) - if cfg is None: logger.warning( - "MCP server '{}' session terminated but is no longer configured", + "MCP server '{}' session terminated; refreshing connection", server_name, ) - return None + _unregister_server_tools(self._registry, server_name) + await self._close_server(server_name) - current_tool = registry.get(tool_name) - if ( - current_tool is not None - and current_tool is not stale_tool - and server_name in state._mcp_stacks - ): - return current_tool - - logger.warning("MCP server '{}' session terminated; refreshing connection", server_name) - _unregister_server_tools(registry, server_name) - await _close_server(state, server_name) - - _set_runtime_status(state, {server_name}, "connecting") - connected = await connect_mcp_servers({server_name: cfg}, registry) - if getattr(state, "_mcp_closing", False): - for connection in connected.values(): - await connection.aclose() - return None - state._mcp_stacks.update(connected) - _record_connection_result(state, {server_name}, connected) - _attach_reconnect_handlers(state, registry, connected) - if server_name not in connected: - logger.warning("MCP server '{}' reconnect failed after session termination", server_name) - return None - return registry.get(tool_name) + self._set_runtime_status({server_name}, "connecting") + connected = await connect_mcp_servers( + {server_name: cfg}, + self._registry, + ) + if self._closing: + await self._close_connections(connected) + return None + self._connections.update(connected) + self._record_connection_result({server_name}, connected) + self._attach_reconnect_handlers(connected) + if server_name not in connected: + logger.warning( + "MCP server '{}' reconnect failed after session termination", + server_name, + ) + return None + return self._registry.get(tool_name) + + async def _close_server(self, server_name: str) -> None: + connection = self._connections.pop(server_name, None) + if connection is None: + return + await self._close_connection(server_name, connection) + + async def aclose(self) -> None: + """Close every connection while excluding reconnect and hot reload.""" + self._closing = True + async with self._lock: + connections = dict(self._connections) + self._connections.clear() + self._runtime_statuses.clear() + for name in self._servers: + _unregister_server_tools(self._registry, name) + await self._close_connections(connections) def _server_signature(cfg: Any) -> Any: @@ -1690,37 +1687,3 @@ def _unregister_server_tools(registry: ToolRegistry, server_name: str) -> int: registry.unregister(tool_name) removed += 1 return removed - - -async def _close_server(state: Any, server_name: str) -> None: - stack = state._mcp_stacks.pop(server_name, None) - if stack is None: - return - try: - await stack.aclose() - except asyncio.CancelledError: - if task_is_cancelling(): - raise - logger.debug("MCP server '{}' cleanup error (can be ignored)", server_name) - except (RuntimeError, BaseExceptionGroup): - logger.debug("MCP server '{}' cleanup error (can be ignored)", server_name) - - -async def close_mcp_servers(state: Any) -> None: - """Close every MCP connection while excluding reconnect and hot reload.""" - state._mcp_closing = True - async with _reload_lock(state): - connections = list(state._mcp_stacks.items()) - state._mcp_stacks.clear() - statuses = _runtime_status_store(state) - if statuses is not None: - statuses.clear() - for name, connection in connections: - try: - await connection.aclose() - except asyncio.CancelledError: - if task_is_cancelling(): - raise - logger.debug("MCP server '{}' cleanup error (can be ignored)", name) - except (RuntimeError, BaseExceptionGroup): - logger.debug("MCP server '{}' cleanup error (can be ignored)", name) diff --git a/nanobot/agent/tools/self.py b/nanobot/agent/tools/self.py index 075c43806e3..45ed72dd67f 100644 --- a/nanobot/agent/tools/self.py +++ b/nanobot/agent/tools/self.py @@ -78,7 +78,7 @@ def enabled(cls, ctx: ToolContext) -> bool: "runner", "sessions", "consolidator", "dream", "auto_compact", "context", "commands", # Sensitive runtime state (credentials, message routing, task tracking) - "_mcp_servers", "_mcp_stacks", "_pending_queues", + "_pending_queues", "_session_locks", "_active_tasks", "_background_tasks", # Security boundaries (inspect + modify both blocked) "restrict_to_workspace", "channels_config", diff --git a/nanobot/api/server.py b/nanobot/api/server.py index 9ad57deefde..59da093fea7 100644 --- a/nanobot/api/server.py +++ b/nanobot/api/server.py @@ -48,6 +48,7 @@ _MODEL_NAME_KEY = web.AppKey[str]("model_name") _REQUEST_TIMEOUT_KEY = web.AppKey[float]("request_timeout") _SESSION_LOCKS_KEY = web.AppKey[dict[str, asyncio.Lock]]("session_locks") +_PREPARE_AGENT_KEY = web.AppKey[Callable[[], Awaitable[None]] | None]("prepare_agent") _MISSING = object() @@ -66,6 +67,17 @@ def _app_value( return app.get(legacy_key, default) +async def _prepare_agent(app: Any) -> None: + prepare: Callable[[], Awaitable[None]] | None = _app_value( + app, + _PREPARE_AGENT_KEY, + "prepare_agent", + None, + ) + if prepare is not None: + await prepare() + + # --------------------------------------------------------------------------- # Response helpers # --------------------------------------------------------------------------- @@ -346,8 +358,9 @@ async def _run() -> None: nonlocal stream_failed try: async with session_lock: - response = await asyncio.wait_for( - agent_loop.process_direct( + async with asyncio.timeout(timeout_s): + await _prepare_agent(request.app) + response = await agent_loop.process_direct( content=text, media=media_paths if media_paths else None, session_key=session_key, @@ -355,9 +368,7 @@ async def _run() -> None: chat_id=API_CHAT_ID, on_stream=_on_stream, on_stream_end=_on_stream_end, - ), - timeout=timeout_s, - ) + ) if not emitted_content: response_text = _response_text(response) if response_text.strip(): @@ -390,16 +401,15 @@ async def _run() -> None: try: async with session_lock: try: - response = await asyncio.wait_for( - agent_loop.process_direct( + async with asyncio.timeout(timeout_s): + await _prepare_agent(request.app) + response = await agent_loop.process_direct( content=text, media=media_paths if media_paths else None, session_key=session_key, channel="api", chat_id=API_CHAT_ID, - ), - timeout=timeout_s, - ) + ) response_text = _response_text(response) if not response_text or not response_text.strip(): logger.warning("Empty response for session {}, using fallback", session_key) @@ -452,6 +462,7 @@ def create_app( model_name: str = "nanobot", request_timeout: float = 120.0, api_key: str = "", + prepare_agent: Callable[[], Awaitable[None]] | None = None, ) -> web.Application: """Create the aiohttp application. @@ -460,12 +471,14 @@ def create_app( model_name: Model name reported in responses. request_timeout: Per-request timeout in seconds. api_key: Optional API key for Bearer-token authentication on API routes. + prepare_agent: Optional application-owned readiness callback run before each turn. """ app = web.Application(client_max_size=20 * 1024 * 1024) # 20MB for base64 images app[_AGENT_LOOP_KEY] = agent_loop app[_MODEL_NAME_KEY] = model_name app[_REQUEST_TIMEOUT_KEY] = request_timeout app[_SESSION_LOCKS_KEY] = {} # per-user locks, keyed by session_key + app[_PREPARE_AGENT_KEY] = prepare_agent @web.middleware async def auth_middleware( diff --git a/nanobot/bus/events.py b/nanobot/bus/events.py index 2e2cad6ce46..28ada6b784b 100644 --- a/nanobot/bus/events.py +++ b/nanobot/bus/events.py @@ -16,7 +16,6 @@ # loop to update runtime state without going through a user session. INBOUND_META_RUNTIME_CONTROL = "_runtime_control" RUNTIME_CONTROL_ACK = "_ack" -RUNTIME_CONTROL_MCP_RELOAD = "mcp_reload" RUNTIME_CONTROL_IMAGE_GENERATION_RELOAD = "image_generation_reload" RUNTIME_CONTROL_SESSION_DISCARD = "session_discard" diff --git a/nanobot/channels/manager.py b/nanobot/channels/manager.py index 43765181378..0015b712f3b 100644 --- a/nanobot/channels/manager.py +++ b/nanobot/channels/manager.py @@ -5,7 +5,7 @@ import asyncio import hashlib import inspect -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Awaitable, Callable, Iterable, Mapping from contextlib import suppress from pathlib import Path from typing import TYPE_CHECKING, Any, cast @@ -101,6 +101,7 @@ def __init__( webui_runtime_surface: str = "browser", webui_runtime_capabilities: dict[str, Any] | None = None, webui_mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None, + webui_mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None, webui_skill_state_action: Callable[[set[str]], None] | None = None, config_path: Path | None = None, ): @@ -121,6 +122,7 @@ def __init__( self._webui_runtime_surface = webui_runtime_surface self._webui_runtime_capabilities = dict(webui_runtime_capabilities or {}) self._webui_mcp_runtime_status = webui_mcp_runtime_status + self._webui_mcp_reload = webui_mcp_reload self._webui_skill_state_action = webui_skill_state_action self.channels: dict[str, BaseChannel] = {} self._channel_owners: dict[str, str] = {} @@ -190,6 +192,7 @@ def _build_channel( channel_feature_action=self.apply_channel_feature_action, channel_runtime_status=self.get_status, mcp_runtime_status=self._webui_mcp_runtime_status, + mcp_reload=self._webui_mcp_reload, skill_state_action=self._webui_skill_state_action, logger=logger, ) diff --git a/nanobot/channels/websocket/tests/test_websocket_http_routes.py b/nanobot/channels/websocket/tests/test_websocket_http_routes.py index 6740c37d72d..4f77b455f90 100644 --- a/nanobot/channels/websocket/tests/test_websocket_http_routes.py +++ b/nanobot/channels/websocket/tests/test_websocket_http_routes.py @@ -75,6 +75,7 @@ def _make_handler( local_trigger_pending_ids: Any | None = None, channel_feature_action: Any | None = None, channel_runtime_status: Any | None = None, + mcp_reload: Any | None = None, ) -> GatewayServices: config = WebSocketConfig.model_validate(cfg) if isinstance(cfg, dict) else cfg workspace = workspace_path or Path.cwd() @@ -94,6 +95,7 @@ def _make_handler( local_trigger_pending_ids=local_trigger_pending_ids, channel_feature_action=channel_feature_action, channel_runtime_status=channel_runtime_status, + mcp_reload=mcp_reload, ) @@ -111,6 +113,7 @@ def _ch( local_trigger_pending_ids: Any | None = None, channel_feature_action: Any | None = None, channel_runtime_status: Any | None = None, + mcp_reload: Any | None = None, **extra: Any, ) -> WebSocketChannel: cfg: dict[str, Any] = { @@ -134,6 +137,7 @@ def _ch( local_trigger_pending_ids=local_trigger_pending_ids, channel_feature_action=channel_feature_action, channel_runtime_status=channel_runtime_status, + mcp_reload=mcp_reload, ) return InProcessHttpChannel(cfg, bus, gateway=gateway) @@ -2054,14 +2058,15 @@ def _custom_action( _custom_action, ) - async def _hot_reload(_bus): + async def _hot_reload(): return {"ok": True, "message": "MCP config reloaded.", "requires_restart": False} - monkeypatch.setattr( - "nanobot.webui.settings_routes.request_mcp_reload", - _hot_reload, + channel = _ch( + bus, + session_manager=_seed_session(tmp_path), + port=29913, + mcp_reload=_hot_reload, ) - channel = _ch(bus, session_manager=_seed_session(tmp_path), port=29913) server_task = asyncio.create_task(channel.start()) try: deny = await _http_get("http://127.0.0.1:29913/api/settings/mcp-presets") diff --git a/nanobot/cli/agent.py b/nanobot/cli/agent.py index 854383763f1..b8cd7647ba9 100644 --- a/nanobot/cli/agent.py +++ b/nanobot/cli/agent.py @@ -13,6 +13,8 @@ from nanobot import __logo__ from nanobot.agent.hooks import create_file_edit_activity_hook from nanobot.agent.loop import AgentLoop +from nanobot.agent.tools.mcp import MCPProvider +from nanobot.agent.tools.registry import ToolRegistry from nanobot.bus.outbound_events import ( StreamDeltaEvent, StreamedResponseEvent, @@ -84,6 +86,8 @@ def agent( # Create cron service with workspace-scoped store cron_store_path = runtime_config.workspace_path / "cron" / "jobs.json" cron = CronService(cron_store_path) + tools = ToolRegistry() + mcp_provider = MCPProvider.from_config(runtime_config, tools) _set_nanobot_logs(logs) @@ -95,6 +99,7 @@ def agent( cron_service=cron, image_generation_provider_configs=image_gen_provider_configs(runtime_config), hook_factories=[create_file_edit_activity_hook], + tool_registry=tools, ) except ValueError as exc: _print_agent_start_error(exc) @@ -106,6 +111,12 @@ def agent( render_markdown=False, ) + async def _close_runtime() -> None: + try: + await agent_loop.aclose() + finally: + await mcp_provider.aclose() + # Shared reference for progress callbacks _thinking: ThinkingSpinner | None = None @@ -149,30 +160,33 @@ async def _cli_progress( if message: # Single message mode — direct call, no bus needed async def run_once() -> None: - renderer = StreamRenderer( - render_markdown=markdown, - bot_name=runtime_config.agents.defaults.bot_name, - bot_icon=runtime_config.agents.defaults.bot_icon, - ) - response = await agent_loop.process_direct( - message, - session_id, - on_progress=_make_progress(renderer), - on_stream=renderer.on_delta, - on_stream_end=renderer.on_end, - ) - if not renderer.streamed: - await renderer.close() - print_kwargs: dict[str, Any] = {} - if renderer.header_printed: - print_kwargs["show_header"] = False - cli_terminal._print_agent_response( - response.content if response else "", + try: + await mcp_provider.connect() + renderer = StreamRenderer( render_markdown=markdown, - metadata=response.metadata if response else None, - **print_kwargs, + bot_name=runtime_config.agents.defaults.bot_name, + bot_icon=runtime_config.agents.defaults.bot_icon, ) - await agent_loop.close_mcp() + response = await agent_loop.process_direct( + message, + session_id, + on_progress=_make_progress(renderer), + on_stream=renderer.on_delta, + on_stream_end=renderer.on_end, + ) + if not renderer.streamed: + await renderer.close() + print_kwargs: dict[str, Any] = {} + if renderer.header_printed: + print_kwargs["show_header"] = False + cli_terminal._print_agent_response( + response.content if response else "", + render_markdown=markdown, + metadata=response.metadata if response else None, + **print_kwargs, + ) + finally: + await _close_runtime() asyncio.run(run_once()) else: @@ -209,6 +223,7 @@ def _handle_signal(signum: int, _frame: FrameType | None) -> None: signal.signal(signal.SIGPIPE, signal.SIG_IGN) async def run_interactive() -> None: + await mcp_provider.connect() bus_task = asyncio.create_task(agent_loop.run()) turn_done = asyncio.Event() turn_done.set() @@ -347,6 +362,6 @@ async def _consume_outbound() -> None: agent_loop.stop() outbound_task.cancel() await asyncio.gather(bus_task, outbound_task, return_exceptions=True) - await agent_loop.close_mcp() + await _close_runtime() asyncio.run(run_interactive()) diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index efd97c6270b..318e55b5739 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -49,6 +49,8 @@ from nanobot import optional_features as feature_support # noqa: E402 from nanobot.agent.hooks import create_file_edit_activity_hook # noqa: E402 from nanobot.agent.loop import AgentLoop # noqa: E402 +from nanobot.agent.tools.mcp import MCPProvider # noqa: E402 +from nanobot.agent.tools.registry import ToolRegistry # noqa: E402 from nanobot.cli import terminal as cli_terminal # noqa: E402 from nanobot.cli.agent import agent # noqa: E402 from nanobot.cli.gateway import create_gateway_app # noqa: E402 @@ -351,12 +353,15 @@ def serve( sync_workspace_templates(runtime_config.workspace_path) bus = MessageBus() session_manager = SessionManager(runtime_config.workspace_path) + tools = ToolRegistry() + mcp_provider = MCPProvider.from_config(runtime_config, tools) try: agent_loop = AgentLoop.from_config( runtime_config, bus, session_manager=session_manager, image_generation_provider_configs=image_gen_provider_configs(runtime_config), hook_factories=[create_file_edit_activity_hook], + tool_registry=tools, ) except ValueError as exc: console.print(f"[red]Error: {exc}[/red]") @@ -378,13 +383,17 @@ def serve( api_app = create_app( agent_loop, model_name=model_name, request_timeout=timeout, api_key=api_key, + prepare_agent=mcp_provider.connect, ) async def on_startup(_app: Any) -> None: - await agent_loop._connect_mcp() + await mcp_provider.connect() async def on_cleanup(_app: Any) -> None: - await agent_loop.close_mcp() + try: + await agent_loop.aclose() + finally: + await mcp_provider.aclose() api_app.on_startup.append(on_startup) api_app.on_cleanup.append(on_cleanup) diff --git a/nanobot/cli/gateway_runtime.py b/nanobot/cli/gateway_runtime.py index 3f48cff142f..f31cb0e90ee 100644 --- a/nanobot/cli/gateway_runtime.py +++ b/nanobot/cli/gateway_runtime.py @@ -14,6 +14,8 @@ from nanobot import __logo__, __version__ from nanobot.agent.hooks import create_file_edit_activity_hook from nanobot.agent.loop import AgentLoop +from nanobot.agent.tools.mcp import MCPProvider +from nanobot.agent.tools.registry import ToolRegistry from nanobot.cli import terminal as cli_terminal from nanobot.cli.runtime_config import _migrate_cron_store from nanobot.cli.webui_support import ( @@ -233,6 +235,7 @@ def _print_gateway_health_endpoint(host: str, port: int) -> None: async def _close_gateway_runtime( agent: AgentLoop, + mcp_provider: MCPProvider, channels: Any, tasks: list[asyncio.Task[Any]], runtime_tasks: asyncio.Future[list[Any]] | None, @@ -240,18 +243,13 @@ async def _close_gateway_runtime( task_wait_timeout: float = 15.0, close_timeout: float = 15.0, ) -> None: - """Cancel runtime tasks, then deterministically close agent resources. + """Cancel runtime tasks, then deterministically close application resources. Order matters: runtime tasks (including the agent loop and any in-flight - turn) are cancelled and awaited -- bounded -- before exec sessions, - subagents, and MCP servers are torn down, so no active turn is using a - shared resource when it closes. The final close is bounded and idempotent: - the agent loop's own finally also calls ``close_mcp()``, so this runs again - as a no-op when that path already completed, and as the guaranteed final - close when it was skipped or cut short (which previously left asyncio - subprocess transports alive past ``loop.close()``, producing - "RuntimeError: Event loop is closed" noise and potentially orphaned - processes at interpreter exit). + turn) are cancelled and awaited -- bounded -- before the loop-owned resources + and the application-owned MCP provider are torn down. The final close is + bounded and idempotent, so it also covers a cancelled or incomplete loop + cleanup without leaving subprocess transports alive past ``loop.close()``. """ # Some SDKs swallow task cancellation while attempting to reconnect. # Close channel transports before waiting for their runners to exit. @@ -272,10 +270,14 @@ async def _close_gateway_runtime( 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) + for label, close in ( + ("agent", agent.aclose), + ("MCP provider", mcp_provider.aclose), + ): + try: + await asyncio.wait_for(close(), timeout=close_timeout) + except BaseException as exc: # noqa: BLE001 - shutdown must proceed + logger.warning("Gateway shutdown: {} cleanup incomplete: {}", label, exc) # 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(): @@ -414,6 +416,9 @@ def _load_gateway_provider_snapshot( route_policy=WebuiTurnRoutePolicy(session_manager), ) + tools = ToolRegistry() + mcp_provider = MCPProvider.from_config(config, tools) + # Create agent with cron service agent = AgentLoop.from_config( config, bus, @@ -431,6 +436,7 @@ def _load_gateway_provider_snapshot( hooks=[TokenUsageHook(timezone_name=config.agents.defaults.timezone)], local_trigger_store=trigger_store, hook_factories=[create_file_edit_activity_hook], + tool_registry=tools, ) def _schedule_webui_background(awaitable: Awaitable[None]) -> None: agent.schedule_background(cast(Coroutine[Any, Any, None], awaitable)) @@ -512,6 +518,7 @@ async def _silent(*_args: Any, **_kwargs: Any) -> None: prompt, last_cursor = result key = dream_session_key() dream_runtime = agent.dream_runtime() + await mcp_provider.connect() resp = await agent.process_direct( prompt, session_key=key, @@ -589,6 +596,7 @@ async def _silent(*_args: Any, **_kwargs: Any) -> None: if isinstance(message_tool, MessageTool): suppress_token = message_tool.set_suppress_delivery(True) try: + await mcp_provider.connect() resp = await agent.process_direct( prompt, session_key="heartbeat", @@ -668,7 +676,8 @@ def _webui_skill_state_action(disabled_skills: set[str]) -> None: webui_static_dist=webui_static_dist, webui_runtime_surface=webui_runtime_surface, webui_runtime_capabilities=webui_runtime_capabilities, - webui_mcp_runtime_status=agent.mcp_runtime_status, + webui_mcp_runtime_status=mcp_provider.runtime_status, + webui_mcp_reload=mcp_provider.reload, webui_skill_state_action=_webui_skill_state_action, config_path=Path(config_path), ) @@ -844,6 +853,13 @@ async def run() -> None: await cron.start() # Re-read once on first admission to close the watcher subscription window. agent.runtime_resolver.invalidate() + async def _run_agent() -> None: + try: + await mcp_provider.connect() + await agent.run() + finally: + await mcp_provider.aclose() + tasks = [ asyncio.create_task( watch_config_file( @@ -852,7 +868,7 @@ async def run() -> None: ), name="nanobot-config-watcher", ), - asyncio.create_task(agent.run(), name="nanobot-agent-loop"), + asyncio.create_task(_run_agent(), name="nanobot-agent-loop"), asyncio.create_task(channels.start_all(), name="nanobot-channels"), asyncio.create_task( run_local_trigger_queue( @@ -910,7 +926,13 @@ async def run() -> None: 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) + await _close_gateway_runtime( + agent, + mcp_provider, + 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/nanobot/nanobot.py b/nanobot/nanobot.py index c1b3395370e..2720438386c 100644 --- a/nanobot/nanobot.py +++ b/nanobot/nanobot.py @@ -10,6 +10,8 @@ from nanobot.agent.hook import AgentHook, SDKCaptureHook from nanobot.agent.hooks import create_file_edit_activity_hook from nanobot.agent.loop import AgentLoop +from nanobot.agent.tools.mcp import MCPProvider +from nanobot.agent.tools.registry import ToolRegistry from nanobot.config.schema import Config from nanobot.providers.image_generation import image_gen_provider_configs from nanobot.sdk.clients import MemoryClient, RuntimeClient, SessionClient @@ -71,9 +73,16 @@ class Nanobot: print(result.content) """ - def __init__(self, loop: AgentLoop, *, config: Config | None = None) -> None: + def __init__( + self, + loop: AgentLoop, + *, + config: Config | None = None, + mcp_provider: MCPProvider | None = None, + ) -> None: self._loop = loop self._config = config + self._mcp_provider = mcp_provider self.sessions = SessionClient(loop) self.memory = MemoryClient(loop) self.runtime = RuntimeClient(loop) @@ -120,12 +129,15 @@ def from_config( elif model_preset is not None: config.agents.defaults.model_preset = model_preset + tools = ToolRegistry() + mcp_provider = MCPProvider.from_config(config, tools) loop = AgentLoop.from_config( config, image_generation_provider_configs=image_gen_provider_configs(config), hook_factories=[create_file_edit_activity_hook], + tool_registry=tools, ) - return cls(loop, config=config) + return cls(loop, config=config, mcp_provider=mcp_provider) async def run( self, @@ -178,6 +190,8 @@ async def run( ) if runtime is not None: kwargs["runtime"] = runtime + if self._mcp_provider is not None: + await self._mcp_provider.connect() response = await self._loop.process_direct( message, **kwargs, @@ -259,6 +273,8 @@ async def _run() -> RunResult: if override_runtime is not None: kwargs["runtime"] = override_runtime try: + if self._mcp_provider is not None: + await self._mcp_provider.connect() response = await self._loop.process_direct( message, **kwargs, @@ -327,8 +343,12 @@ async def stream( await run.aclose() async def aclose(self) -> None: - """Release resources held by this instance (MCP connections, etc.).""" - await self._loop.close_mcp() + """Release resources held by this instance.""" + try: + await self._loop.aclose() + finally: + if self._mcp_provider is not None: + await self._mcp_provider.aclose() async def __aenter__(self) -> Nanobot: return self diff --git a/nanobot/webui/gateway_services.py b/nanobot/webui/gateway_services.py index 0460f582af8..9820fdbb73b 100644 --- a/nanobot/webui/gateway_services.py +++ b/nanobot/webui/gateway_services.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Awaitable, Mapping from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, Callable @@ -66,6 +66,7 @@ def build_gateway_services( channel_feature_action: Callable[..., Any] | None = None, channel_runtime_status: Callable[[], dict[str, Any]] | None = None, mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None, + mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None, skill_state_action: Callable[[set[str]], None] | None = None, logger: Any = default_logger, ) -> GatewayServices: @@ -119,6 +120,7 @@ def build_gateway_services( channel_feature_action=channel_feature_action, channel_runtime_status=channel_runtime_status, mcp_runtime_status=mcp_runtime_status, + mcp_reload=mcp_reload, skill_state_action=skill_state_action, log=logger, ) diff --git a/nanobot/webui/settings_routes.py b/nanobot/webui/settings_routes.py index 633bff4d564..c938ae10672 100644 --- a/nanobot/webui/settings_routes.py +++ b/nanobot/webui/settings_routes.py @@ -5,14 +5,13 @@ import asyncio import html import json -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from typing import Any, cast from websockets.http11 import Request as WsRequest from websockets.http11 import Response from nanobot.agent.tools.image_generation import request_image_generation_reload -from nanobot.agent.tools.mcp import request_mcp_reload from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH from nanobot.api.runtime import ApiRuntime, api_runtime_paths from nanobot.bus.queue import MessageBus @@ -71,6 +70,7 @@ _WEBUI_MUTATION_REQUEST_ATTR = "_nanobot_webui_mutation_request" _CHANNEL_CONNECT_ACTIONS = frozenset({"start", "poll", "cancel"}) _MCP_OAUTH_CALLBACK_URL_MAX_BYTES = 8 * 1024 +_MCP_RELOAD_TIMEOUT_SECONDS = 15.0 _query_first = contracts.query_first @@ -227,6 +227,7 @@ def __init__( channel_feature_action: Callable[..., Any] | None = None, channel_runtime_status: Callable[[], dict[str, Any]] | None = None, mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None, + mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None, mcp_oauth_redirect_uri: Callable[[WsRequest], str] | None = None, ) -> None: self.settings = settings @@ -241,6 +242,7 @@ def __init__( self._channel_feature_action = channel_feature_action self._channel_runtime_status = channel_runtime_status self._mcp_runtime_status = mcp_runtime_status + self._mcp_reload = mcp_reload self._mcp_oauth_redirect_uri = mcp_oauth_redirect_uri self._mcp_oauth = McpOAuthManager() self._restart_sections: set[str] = set() @@ -472,7 +474,7 @@ def _system_operations(self) -> system_domain.SystemSettingsOperations: approve_code=approve_code, deny_code=deny_code, mcp_presets_action=mcp_presets_settings_action, - reload_mcp=lambda: request_mcp_reload(self.bus), + reload_mcp=self._reload_mcp_runtime, mcp_runtime_status=self._mcp_runtime_status, check_for_update=check_for_update, channel_feature_action=self._channel_feature_action, @@ -499,6 +501,33 @@ async def _apply_image_generation_runtime_change( self._restart_sections.discard("image") return updated + async def _reload_mcp_runtime(self) -> dict[str, Any]: + if self._mcp_reload is None: + return { + "ok": False, + "message": "MCP runtime reload is unavailable. Restart nanobot to apply changes.", + "requires_restart": True, + } + try: + return await asyncio.wait_for( + self._mcp_reload(), + timeout=_MCP_RELOAD_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + return { + "ok": False, + "message": "MCP hot reload timed out. Restart nanobot to pick up changes.", + "requires_restart": True, + } + except Exception as exc: + self.logger.exception("MCP hot reload failed") + return { + "ok": False, + "message": "MCP hot reload failed. Restart nanobot to pick up changes.", + "requires_restart": True, + "error": str(exc), + } + def _parse_mcp_settings_query(self, request: WsRequest) -> QueryParams: return self._query(request) @@ -622,7 +651,7 @@ async def _handle_mcp_oauth_start(self, request: WsRequest) -> Response: name, cfg, redirect_uri, - reload_mcp=lambda: request_mcp_reload(self.bus), + reload_mcp=self._reload_mcp_runtime, reset_credentials=reset, ) except Exception as exc: diff --git a/nanobot/webui/ws_http.py b/nanobot/webui/ws_http.py index 19ef9014bb1..363db6d3f0a 100644 --- a/nanobot/webui/ws_http.py +++ b/nanobot/webui/ws_http.py @@ -14,7 +14,7 @@ import mimetypes import re import time -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from pathlib import Path from typing import TYPE_CHECKING, Any, cast from urllib.parse import quote, unquote, urlsplit, urlunsplit @@ -308,6 +308,7 @@ def __init__( channel_feature_action: Callable[..., Any] | None = None, channel_runtime_status: Callable[[], dict[str, Any]] | None = None, mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None, + mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None, skill_state_action: Callable[[set[str]], None] | None = None, log: Any = logger, ) -> None: @@ -351,6 +352,7 @@ def __init__( channel_feature_action=channel_feature_action, channel_runtime_status=channel_runtime_status, mcp_runtime_status=mcp_runtime_status, + mcp_reload=mcp_reload, mcp_oauth_redirect_uri=self._mcp_oauth_redirect_uri, ) diff --git a/tests/agent/conftest.py b/tests/agent/conftest.py index 9bb4b65b3c8..bb8ec7db7a6 100644 --- a/tests/agent/conftest.py +++ b/tests/agent/conftest.py @@ -46,7 +46,6 @@ def make_loop( context_window_tokens: int = 128_000, session_ttl_minutes: int = 0, unified_session: bool = False, - mcp_servers: dict | None = None, tools_config=None, model_presets: dict | None = None, hooks: list | None = None, @@ -72,8 +71,6 @@ def make_loop( session_ttl_minutes=session_ttl_minutes, unified_session=unified_session, ) - if mcp_servers is not None: - kwargs["mcp_servers"] = mcp_servers if tools_config is not None: kwargs["tools_config"] = tools_config if model_presets is not None: diff --git a/tests/agent/test_auto_compact.py b/tests/agent/test_auto_compact.py index f29af04bf7d..7f579366f9b 100644 --- a/tests/agent/test_auto_compact.py +++ b/tests/agent/test_auto_compact.py @@ -310,7 +310,7 @@ async def test_is_expired_boundary(self, tmp_path): assert loop.auto_compact._is_expired(ts) is True ts2 = datetime.now() - timedelta(minutes=14, seconds=59) assert loop.auto_compact._is_expired(ts2) is False - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_is_expired_string_timestamp(self, tmp_path): @@ -320,7 +320,7 @@ async def test_is_expired_string_timestamp(self, tmp_path): assert loop.auto_compact._is_expired(ts) is True assert loop.auto_compact._is_expired(None) is False assert loop.auto_compact._is_expired("") is False - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_check_expired_only_archives_expired_sessions(self, tmp_path): @@ -343,7 +343,7 @@ async def test_check_expired_only_archives_expired_sessions(self, tmp_path): active_after = loop.sessions.get_or_create("cli:active") assert len(active_after.messages) == 1 assert active_after.messages[0]["content"] == "recent" - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_archives_full_tail_without_deleting_history(self, tmp_path): @@ -367,7 +367,7 @@ async def test_auto_compact_archives_full_tail_without_deleting_history(self, tm assert len(visible) == loop.auto_compact._RECENT_SUFFIX_MESSAGES assert visible[0]["content"] == "msg user 2" assert visible[-1]["content"] == "msg assistant 5" - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_extends_recent_suffix_to_user_turn(self, tmp_path): @@ -398,7 +398,7 @@ async def test_auto_compact_extends_recent_suffix_to_user_turn(self, tmp_path): for m in visible for tc in (m.get("tool_calls") or []) ) - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_stores_summary(self, tmp_path): @@ -422,7 +422,7 @@ async def test_auto_compact_stores_summary(self, tmp_path): assert len(session_after.get_history(max_messages=12)) == ( loop.auto_compact._RECENT_SUFFIX_MESSAGES ) - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_empty_session(self, tmp_path): @@ -436,7 +436,7 @@ async def test_auto_compact_empty_session(self, tmp_path): session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == 0 assert "cli:test" not in loop.auto_compact._summaries - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_respects_last_consolidated(self, tmp_path): @@ -455,7 +455,7 @@ async def test_auto_compact_respects_last_consolidated(self, tmp_path): await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) assert len(archived_messages) == 10 - await loop.close_mcp() + await loop.aclose() class TestAutoCompactIdleDetection: @@ -474,7 +474,7 @@ async def test_no_auto_compact_when_ttl_disabled(self, tmp_path): session_after = loop.sessions.get_or_create("cli:test") assert any(m["content"] == "old message" for m in session_after.messages) - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_triggers_on_idle(self, tmp_path): @@ -503,7 +503,7 @@ async def test_auto_compact_triggers_on_idle(self, tmp_path): for m in session_after.get_history(max_messages=len(session_after.messages)) ) assert any(m["content"] == "new msg" for m in session_after.messages) - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_no_auto_compact_when_active(self, tmp_path): @@ -517,7 +517,7 @@ async def test_no_auto_compact_when_active(self, tmp_path): session_after = loop.sessions.get_or_create("cli:test") assert any(m["content"] == "recent message" for m in session_after.messages) - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_does_not_affect_priority_commands(self, tmp_path): @@ -540,7 +540,7 @@ async def test_auto_compact_does_not_affect_priority_commands(self, tmp_path): # Session should be untouched since priority commands skip _process_message session_after = loop.sessions.get_or_create("cli:test") assert any(m["content"] == "old message" for m in session_after.messages) - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_with_slash_new(self, tmp_path): @@ -562,7 +562,7 @@ async def test_auto_compact_with_slash_new(self, tmp_path): session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == 0 - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_shortcut_command_persisted_with_command_flag(self, tmp_path): @@ -581,7 +581,7 @@ async def test_shortcut_command_persisted_with_command_flag(self, tmp_path): assert session_after.messages[1]["role"] == "assistant" assert session_after.messages[1].get("_command") is True assert AgentLoop._PENDING_USER_TURN_KEY not in session_after.metadata - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_shortcut_command_excluded_from_get_history(self, tmp_path): @@ -597,7 +597,7 @@ async def test_shortcut_command_excluded_from_get_history(self, tmp_path): assert len(history) == 2 assert all(m["content"] != "/help" for m in history) assert all(m["content"] != "help text" for m in history) - await loop.close_mcp() + await loop.aclose() class TestAutoCompactSystemMessages: @@ -628,7 +628,7 @@ async def test_auto_compact_triggers_for_system_messages(self, tmp_path): m["content"] == "old user 0" for m in session_after.get_history(max_messages=len(session_after.messages)) ) - await loop.close_mcp() + await loop.aclose() class TestAutoCompactEdgeCases: @@ -656,7 +656,7 @@ async def test_auto_compact_with_nothing_summary(self, tmp_path): # "(nothing)" summary should not be stored assert "cli:test" not in loop.auto_compact._summaries - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_archive_failure_preserves_raw_history(self, tmp_path): @@ -677,7 +677,7 @@ async def test_auto_compact_archive_failure_preserves_raw_history(self, tmp_path loop.auto_compact._RECENT_SUFFIX_MESSAGES ) - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_auto_compact_preserves_runtime_checkpoint_before_check(self, tmp_path): @@ -709,7 +709,7 @@ async def test_auto_compact_preserves_runtime_checkpoint_before_check(self, tmp_ assert any(m["content"] == "previous message" for m in session_after.messages) assert any(m["content"] == "interrupted response" for m in session_after.messages) - await loop.close_mcp() + await loop.aclose() class TestAutoCompactIntegration: @@ -779,7 +779,7 @@ async def test_full_lifecycle(self, tmp_path): # The new message should be processed (response exists) assert response is not None - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_runtime_context_markers_not_persisted_for_multi_paragraph_turn(self, tmp_path): @@ -807,7 +807,7 @@ async def test_runtime_context_markers_not_persisted_for_multi_paragraph_turn(se content = str(persisted.get("content", "")) assert "[Runtime Context" not in content assert "[/Runtime Context]" not in content - await loop.close_mcp() + await loop.aclose() class TestProactiveAutoCompact: @@ -870,7 +870,7 @@ async def test_no_check_when_ttl_disabled(self, tmp_path): session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == 1 - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_proactive_archive_on_idle_tick(self, tmp_path): @@ -897,7 +897,7 @@ async def test_proactive_archive_on_idle_tick(self, tmp_path): entry = loop.auto_compact._summaries.get("cli:test") assert entry is not None assert entry[0] == "User chatted about old things." - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_proactive_archive_skips_dream_sessions(self, tmp_path): @@ -918,7 +918,7 @@ async def test_proactive_archive_skips_dream_sessions(self, tmp_path): assert _fake_compact.state["count"] == 0 assert "dream:20260602-155256" not in loop.auto_compact._archiving assert "dream:20260602-155256" not in loop.auto_compact._summaries - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_no_proactive_archive_when_active(self, tmp_path): @@ -932,7 +932,7 @@ async def test_no_proactive_archive_when_active(self, tmp_path): session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == 1 - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_no_duplicate_archive(self, tmp_path): @@ -968,7 +968,7 @@ async def _slow_compact(key, *, runtime, max_suffix=8): # Clean up block_forever.set() await _drain_background_tasks(loop) - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_proactive_archive_error_does_not_block(self, tmp_path): @@ -989,7 +989,7 @@ async def _failing_compact(key, *, runtime, max_suffix=8): # Key should be removed from _archiving (finally block) assert "cli:test" not in loop.auto_compact._archiving - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_proactive_archive_skips_empty_sessions(self, tmp_path): @@ -1005,7 +1005,7 @@ async def test_proactive_archive_skips_empty_sessions(self, tmp_path): # Empty session should not produce a summary assert "cli:test" not in loop.auto_compact._summaries - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_skip_expired_session_with_active_agent_task(self, tmp_path): @@ -1026,7 +1026,7 @@ async def test_skip_expired_session_with_active_agent_task(self, tmp_path): session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == 12 # All messages preserved - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_archive_after_active_task_completes(self, tmp_path): @@ -1047,7 +1047,7 @@ async def test_archive_after_active_task_completes(self, tmp_path): # Second tick: task completed, should archive await self._run_check_expired(loop) assert _fake_compact.state["count"] == 1 - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_partial_active_set_only_archives_inactive_expired(self, tmp_path): @@ -1083,7 +1083,7 @@ async def test_partial_active_set_only_archives_inactive_expired(self, tmp_path) assert len(s2_after.messages) == 12 # Preserved s3_after = loop.sessions.get_or_create("cli:recent") assert len(s3_after.messages) == 1 # Preserved - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_no_reschedule_after_successful_archive(self, tmp_path): @@ -1104,7 +1104,7 @@ async def test_no_reschedule_after_successful_archive(self, tmp_path): # Second tick: should NOT re-schedule because the session has no removable tail. await self._run_check_expired(loop) assert _fake_compact.state["count"] == 1 # Still 1, not re-scheduled - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_empty_session_does_not_schedule_idle_compact(self, tmp_path): @@ -1124,7 +1124,7 @@ async def test_empty_session_does_not_schedule_idle_compact(self, tmp_path): await self._run_check_expired(loop) assert _fake_compact.state["count"] == 0 assert "cli:test" not in loop.auto_compact._summaries - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_session_can_be_compacted_again_after_new_messages(self, tmp_path): @@ -1155,7 +1155,7 @@ async def test_session_can_be_compacted_again_after_new_messages(self, tmp_path) # Second compact cycle should succeed await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) assert _fake_compact.state["count"] == 2 - await loop.close_mcp() + await loop.aclose() class TestSummaryPersistence: @@ -1182,7 +1182,7 @@ async def test_summary_persisted_in_session_metadata(self, tmp_path): assert meta is not None assert meta["text"] == "User said hello." assert "last_active" in meta - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_summary_recovered_after_restart(self, tmp_path): @@ -1218,7 +1218,7 @@ async def test_summary_recovered_after_restart(self, tmp_path): assert "Previous conversation summary" in summary # _last_summary persists in metadata for restart survival. assert "_last_summary" in reloaded.metadata - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_metadata_persists_for_restart(self, tmp_path): @@ -1246,7 +1246,7 @@ async def test_metadata_persists_for_restart(self, tmp_path): assert "Summary." in summary2 # _last_summary persists in metadata for restart survival. assert "_last_summary" in reloaded.metadata - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_metadata_cleanup_on_inmemory_path(self, tmp_path): @@ -1272,7 +1272,7 @@ async def test_metadata_cleanup_on_inmemory_path(self, tmp_path): assert summary is not None # _last_summary persists in metadata for restart survival. assert "_last_summary" in reloaded.metadata - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_new_summary_overrides_old(self, tmp_path): @@ -1314,7 +1314,7 @@ async def test_new_summary_overrides_old(self, tmp_path): _, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test") assert summary2 is not None assert "Second summary." in summary2 - await loop.close_mcp() + await loop.aclose() @pytest.mark.asyncio async def test_new_command_clears_last_summary(self, tmp_path): @@ -1342,4 +1342,4 @@ async def test_new_command_clears_last_summary(self, tmp_path): # After /new, metadata should no longer contain _last_summary fresh = loop.sessions.get_or_create("cli:test") assert "_last_summary" not in fresh.metadata - await loop.close_mcp() + await loop.aclose() diff --git a/tests/agent/test_consolidate_offset.py b/tests/agent/test_consolidate_offset.py index 315e3eb61e1..6f34075996f 100644 --- a/tests/agent/test_consolidate_offset.py +++ b/tests/agent/test_consolidate_offset.py @@ -538,7 +538,7 @@ async def _failing_summarize(_messages, *, runtime, session_key=None) -> bool: session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == 0 - await loop.close_mcp() + await loop.aclose() assert call_count == 1 @pytest.mark.asyncio @@ -572,7 +572,7 @@ async def _fake_summarize(messages, *, runtime, session_key=None) -> bool: assert response is not None assert "new session started" in response.content.lower() - await loop.close_mcp() + await loop.aclose() assert archived_count == 3 assert archived_session_key == "cli:test" @@ -603,8 +603,8 @@ async def _ok_summarize(_messages, *, runtime, session_key=None) -> bool: assert loop.sessions.get_or_create("cli:test").messages == [] @pytest.mark.asyncio - async def test_close_mcp_drains_background_tasks(self, tmp_path: Path) -> None: - """close_mcp waits for background tasks to complete.""" + async def test_aclose_drains_background_tasks(self, tmp_path: Path) -> None: + """aclose waits for background tasks to complete.""" from nanobot.bus.events import InboundMessage loop = self._make_loop(tmp_path) @@ -632,5 +632,5 @@ async def _slow_summarize(_messages, *, runtime, session_key=None) -> bool: assert not archived.is_set() release_archive.set() - await loop.close_mcp() + await loop.aclose() assert archived.is_set() diff --git a/tests/agent/test_loop_direct_websocket_status.py b/tests/agent/test_loop_direct_websocket_status.py index 9bbb0d91e9f..19cdac492a4 100644 --- a/tests/agent/test_loop_direct_websocket_status.py +++ b/tests/agent/test_loop_direct_websocket_status.py @@ -93,7 +93,6 @@ async def test_process_direct_websocket_clears_run_status(tmp_path) -> None: @pytest.mark.asyncio async def test_process_direct_reuses_existing_session_lock(tmp_path) -> None: loop = _make_loop(tmp_path) - loop._connect_mcp = AsyncMock() session_key = "api:fixed" lock = loop._session_locks.setdefault(session_key, asyncio.Lock()) await lock.acquire() diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index 4359a142e01..bf0315ad8df 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -1519,13 +1519,11 @@ async def fake_run(spec): @pytest.mark.asyncio async def test_process_direct_rejects_reserved_system_channel(tmp_path: Path) -> None: loop = _make_full_loop(tmp_path) - loop._connect_mcp = AsyncMock() # type: ignore[method-assign] loop._process_message = AsyncMock(return_value=None) # type: ignore[method-assign] with pytest.raises(ValueError, match="reserved for internal messages"): await loop.process_direct("external input", channel="system") - loop._connect_mcp.assert_not_awaited() loop._process_message.assert_not_awaited() @@ -1534,7 +1532,6 @@ async def test_process_direct_skip_user_persist_does_not_save_retry_user( tmp_path: Path, ) -> None: loop = _make_full_loop(tmp_path) - loop._connect_mcp = AsyncMock() session = loop.sessions.get_or_create("api:default") session.add_message("user", "hello") session.add_message("assistant", "previous empty-response attempt") diff --git a/tests/agent/test_loop_session_policy.py b/tests/agent/test_loop_session_policy.py index 88dae71e3e6..3a2a065f505 100644 --- a/tests/agent/test_loop_session_policy.py +++ b/tests/agent/test_loop_session_policy.py @@ -123,8 +123,7 @@ async def wait_for_discard(key: str) -> None: await asyncio.sleep(0) loop.provider.chat_with_retry = AsyncMock(side_effect=block_provider) - monkeypatch.setattr(loop, "_connect_mcp", AsyncMock()) - monkeypatch.setattr(loop, "close_mcp", AsyncMock()) + monkeypatch.setattr(loop, "aclose", AsyncMock()) terminate_exec_sessions = AsyncMock(return_value=1) monkeypatch.setattr( loop._exec_session_manager, diff --git a/tests/agent/test_loop_tool_context.py b/tests/agent/test_loop_tool_context.py index 2b6c15964fe..ddb6a764e97 100644 --- a/tests/agent/test_loop_tool_context.py +++ b/tests/agent/test_loop_tool_context.py @@ -11,6 +11,7 @@ current_request_context, reset_request_context, ) +from nanobot.agent.tools.registry import ToolRegistry from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus from nanobot.providers.base import LLMResponse, ToolCallRequest @@ -56,6 +57,22 @@ def prepare_call(self, name: str, arguments: dict): return (self.tool, arguments, None) if name == "cron" else (None, arguments, None) +def test_loop_registers_default_tools_in_injected_registry(tmp_path: Path) -> None: + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + registry = ToolRegistry() + + loop = AgentLoop( + bus=MessageBus(), + provider=provider, + workspace=tmp_path, + tool_registry=registry, + ) + + assert loop.tools is registry + assert registry.has("read_file") + + @pytest.mark.asyncio async def test_loop_binds_request_context_for_tool_execution(tmp_path: Path) -> None: provider = MagicMock() diff --git a/tests/agent/test_mcp_connection.py b/tests/agent/test_mcp_connection.py index 50476553bae..dde56650e3d 100644 --- a/tests/agent/test_mcp_connection.py +++ b/tests/agent/test_mcp_connection.py @@ -1,4 +1,4 @@ -"""Tests for MCP connection lifecycle in AgentLoop.""" +"""Tests for the application-owned MCP provider lifecycle.""" from __future__ import annotations @@ -6,7 +6,7 @@ from contextlib import AsyncExitStack from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import anyio import pytest @@ -15,11 +15,10 @@ from mcp.shared.message import SessionMessage from mcp.types import ErrorData -from nanobot.agent.loop import AgentLoop from nanobot.agent.tools import mcp as mcp_runtime from nanobot.agent.tools.base import Tool -from nanobot.agent.tools.mcp import MCPResourceWrapper, MCPToolWrapper -from nanobot.bus.queue import MessageBus +from nanobot.agent.tools.mcp import MCPProvider, MCPResourceWrapper, MCPToolWrapper +from nanobot.agent.tools.registry import ToolRegistry from nanobot.config.loader import load_config, save_config from nanobot.config.schema import MCPServerConfig @@ -74,18 +73,20 @@ async def execute(self, **_kwargs: Any) -> str: return "ok" -def _make_loop(tmp_path, *, mcp_servers: dict | None = None) -> AgentLoop: - bus = MessageBus() - provider = MagicMock() - provider.get_default_model.return_value = "test-model" - provider.generation.max_tokens = 4096 - return AgentLoop( - bus=bus, - provider=provider, - workspace=tmp_path, - model="test-model", - mcp_servers=mcp_servers or {"test": object()}, +def _stdio_server(command: str = "test-mcp") -> MCPServerConfig: + return MCPServerConfig(type="stdio", command=command) + + +def _make_provider( + *, + mcp_servers: dict[str, MCPServerConfig] | None = None, +) -> tuple[MCPProvider, ToolRegistry]: + registry = ToolRegistry() + provider = MCPProvider( + mcp_servers if mcp_servers is not None else {"test": _stdio_server()}, + registry, ) + return provider, registry @pytest.mark.asyncio @@ -140,7 +141,7 @@ async def own_connection() -> None: @pytest.mark.asyncio async def test_connect_mcp_retries_when_no_servers_connect(tmp_path, monkeypatch: pytest.MonkeyPatch): - loop = _make_loop(tmp_path) + provider, _registry = _make_provider() attempts = 0 async def _fake_connect(_servers, _registry): @@ -150,12 +151,12 @@ async def _fake_connect(_servers, _registry): monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) - await loop._connect_mcp() - await loop._connect_mcp() + await provider.connect() + await provider.connect() assert attempts == 2 - assert loop._mcp_stacks == {} - assert loop.mcp_runtime_status() == {"test": "failed"} + assert provider.connected_server_names == set() + assert provider.runtime_status() == {"test": "failed"} @pytest.mark.asyncio @@ -168,7 +169,7 @@ async def test_connect_mcp_does_not_report_failure_before_oauth_authorization( auth="oauth", url="https://mcp.example.com/mcp", ) - loop = _make_loop(tmp_path, mcp_servers={"oauth-app": cfg}) + provider, _registry = _make_provider(mcp_servers={"oauth-app": cfg}) connect = AsyncMock() monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", connect) monkeypatch.setattr( @@ -176,19 +177,20 @@ async def test_connect_mcp_does_not_report_failure_before_oauth_authorization( lambda _name, _url: False, ) - await loop._connect_mcp() + await provider.connect() connect.assert_not_awaited() - assert loop.mcp_runtime_status() == {} + assert provider.runtime_status() == {} @pytest.mark.asyncio -async def test_agent_loop_run_closes_mcp_from_connection_owner_task( +async def test_mcp_provider_closes_connections_independently_from_agent_loop( tmp_path, monkeypatch: pytest.MonkeyPatch, ): - loop = _make_loop(tmp_path, mcp_servers={"playwright": object()}) - connected = asyncio.Event() + provider, registry = _make_provider( + mcp_servers={"playwright": _stdio_server("playwright")} + ) owner_tasks: list[asyncio.Task | None] = [] closed_tasks: list[asyncio.Task | None] = [] @@ -203,40 +205,38 @@ async def aclose(self) -> None: 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) + await provider.connect() + registry.register(_FakeMcpTool("mcp_playwright_search")) + await provider.aclose() assert owner_tasks assert closed_tasks == owner_tasks - assert loop._mcp_stacks == {} + assert provider.connected_server_names == set() + assert registry.get("mcp_playwright_search") is None @pytest.mark.asyncio async def test_close_server_ignores_server_cancelled_error(tmp_path): - loop = _make_loop(tmp_path) + provider, _registry = _make_provider() class _ServerCancelledStack: async def aclose(self) -> None: raise asyncio.CancelledError() - loop._mcp_stacks = {"test": _ServerCancelledStack()} + provider._connections = {"test": _ServerCancelledStack()} - await mcp_runtime._close_server(loop, "test") + await provider._close_server("test") - assert loop._mcp_stacks == {} + assert provider.connected_server_names == set() @pytest.mark.asyncio -async def test_close_mcp_servers_continues_after_server_cancelled_error(tmp_path): - loop = _make_loop(tmp_path) +async def test_provider_close_continues_after_server_cancelled_error(tmp_path): + provider, _registry = _make_provider() closed: list[str] = [] class _ServerCancelledStack: @@ -247,21 +247,53 @@ class _TrackedStack: async def aclose(self) -> None: closed.append("second") - loop._mcp_stacks = { + provider._connections = { "first": _ServerCancelledStack(), "second": _TrackedStack(), } - await mcp_runtime.close_mcp_servers(loop) + await provider.aclose() assert closed == ["second"] - assert loop._mcp_stacks == {} + assert provider.connected_server_names == set() + + +@pytest.mark.asyncio +async def test_provider_close_finishes_other_connections_before_propagating_cancellation( + tmp_path, +): + provider, _registry = _make_provider() + started = asyncio.Event() + closed: list[str] = [] + + class _BlockingStack: + async def aclose(self) -> None: + started.set() + await asyncio.Event().wait() + + class _TrackedStack: + async def aclose(self) -> None: + closed.append("second") + + provider._connections = { + "first": _BlockingStack(), + "second": _TrackedStack(), + } + task = asyncio.create_task(provider.aclose()) + await asyncio.wait_for(started.wait(), timeout=1) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert closed == ["second"] + assert provider.connected_server_names == set() @pytest.mark.asyncio @pytest.mark.parametrize("close_all", [False, True], ids=["single", "all"]) async def test_mcp_cleanup_re_raises_external_cancellation(tmp_path, close_all: bool): - loop = _make_loop(tmp_path) + provider, _registry = _make_provider() started = asyncio.Event() class _BlockingStack: @@ -269,12 +301,12 @@ async def aclose(self) -> None: started.set() await asyncio.Event().wait() - loop._mcp_stacks = {"test": _BlockingStack()} + provider._connections = {"test": _BlockingStack()} if close_all: - task = asyncio.create_task(mcp_runtime.close_mcp_servers(loop)) + task = asyncio.create_task(provider.aclose()) else: - task = asyncio.create_task(mcp_runtime._close_server(loop, "test")) + task = asyncio.create_task(provider._close_server("test")) await asyncio.wait_for(started.wait(), timeout=1) task.cancel() @@ -312,41 +344,38 @@ async def _fake_connect(servers, registry): return stacks monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) - loop = _make_loop(tmp_path, mcp_servers={}) + provider, registry = _make_provider(mcp_servers={}) - added = await mcp_runtime.reload_servers(loop, loop.tools) + added = await provider.reload() assert added["ok"] is True assert added["added"] == ["browserbase"] - assert loop.tools.has("mcp_browserbase_navigate") - assert "browserbase" in loop._mcp_stacks + assert registry.has("mcp_browserbase_navigate") + assert provider.connected_server_names == {"browserbase"} config = load_config() del config.tools.mcp_servers["browserbase"] save_config(config) - removed = await mcp_runtime.reload_servers(loop, loop.tools) + removed = await provider.reload() assert removed["ok"] is True assert removed["removed"] == ["browserbase"] - assert not loop.tools.has("mcp_browserbase_navigate") - assert "browserbase" not in loop._mcp_stacks + assert not registry.has("mcp_browserbase_navigate") + assert provider.connected_server_names == set() assert closed == ["browserbase"] @pytest.mark.asyncio -async def test_request_mcp_reload_reaches_runtime_control_without_restart( +async def test_reload_is_a_direct_provider_operation_without_an_agent_loop( tmp_path, monkeypatch: pytest.MonkeyPatch, ): - config_path = tmp_path / "config.json" - monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path) - config = load_config() - config.tools.mcp_servers["browserbase"] = MCPServerConfig( + browserbase = MCPServerConfig( type="stdio", command="browserbase-mcp", ) - save_config(config) + configured: dict[str, MCPServerConfig] = {"browserbase": browserbase} closed: list[str] = [] @@ -364,34 +393,24 @@ async def _fake_connect(servers, registry): return stacks monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) - loop = _make_loop(tmp_path, mcp_servers={}) - - async def _handle_one_runtime_control() -> None: - msg = await loop.bus.consume_inbound() - handled = await mcp_runtime.handle_runtime_control(loop, msg, loop.tools) - assert handled is True + registry = ToolRegistry() + provider = MCPProvider({}, registry, server_loader=lambda: configured) - consumer = asyncio.create_task(_handle_one_runtime_control()) - result = await mcp_runtime.request_mcp_reload(loop.bus, timeout=2.0) - await consumer + result = await provider.reload() assert result["ok"] is True assert result["added"] == ["browserbase"] assert result["requires_restart"] is False - assert loop.tools.has("mcp_browserbase_navigate") + assert registry.has("mcp_browserbase_navigate") - config = load_config() - del config.tools.mcp_servers["browserbase"] - save_config(config) + configured = {} - consumer = asyncio.create_task(_handle_one_runtime_control()) - result = await mcp_runtime.request_mcp_reload(loop.bus, timeout=2.0) - await consumer + result = await provider.reload() assert result["ok"] is True assert result["removed"] == ["browserbase"] assert result["requires_restart"] is False - assert not loop.tools.has("mcp_browserbase_navigate") + assert not registry.has("mcp_browserbase_navigate") assert closed == ["browserbase"] @@ -419,16 +438,18 @@ async def _fake_connect(servers, registry): return stacks monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) - loop = _make_loop(tmp_path, mcp_servers={"browserbase": config.tools.mcp_servers["browserbase"]}) + provider, registry = _make_provider( + mcp_servers={"browserbase": config.tools.mcp_servers["browserbase"]} + ) - result = await mcp_runtime.reload_servers(loop, loop.tools) + result = await provider.reload() assert result["ok"] is True assert result["added"] == [] assert result["changed"] == [] assert result["retried"] == ["browserbase"] - assert loop.tools.has("mcp_browserbase_navigate") - await loop.close_mcp() + assert registry.has("mcp_browserbase_navigate") + await provider.aclose() @pytest.mark.asyncio @@ -465,16 +486,16 @@ async def _fake_connect(servers, _registry): "nanobot.agent.tools.mcp_oauth.mcp_oauth_has_credentials", lambda name, _url: name == "linear", ) - loop = _make_loop(tmp_path, mcp_servers={"notion": notion}) + provider, _registry = _make_provider(mcp_servers={"notion": notion}) - result = await mcp_runtime.reload_servers(loop, loop.tools) + result = await provider.reload() assert attempted == ["linear"] assert result["ok"] is True assert result["failed"] == [] assert result["retried"] == [] assert result["connected"] == ["linear"] - await loop.close_mcp() + await provider.aclose() @pytest.mark.asyncio @@ -482,7 +503,9 @@ async def test_mcp_tool_reconnects_after_session_terminated( tmp_path, monkeypatch: pytest.MonkeyPatch, ): - loop = _make_loop(tmp_path, mcp_servers={"remote": object()}) + provider, registry = _make_provider( + mcp_servers={"remote": _stdio_server("remote")} + ) closed: list[str] = [] sessions: list[Any] = [] connect_count = 0 @@ -525,8 +548,8 @@ async def _fake_connect(servers, registry): monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) - await loop._connect_mcp() - old_tool = loop.tools.get("mcp_remote_quote") + await provider.connect() + old_tool = registry.get("mcp_remote_quote") assert isinstance(old_tool, MCPToolWrapper) output = await old_tool.execute(symbol="AAPL") @@ -536,8 +559,8 @@ async def _fake_connect(servers, registry): assert closed == ["remote"] assert sessions[0].call_count == 1 assert sessions[1].call_count == 1 - assert "remote" in loop._mcp_stacks - assert loop.tools.get("mcp_remote_quote") is not old_tool + assert provider.connected_server_names == {"remote"} + assert registry.get("mcp_remote_quote") is not old_tool @pytest.mark.asyncio @@ -545,7 +568,9 @@ async def test_mcp_reconnect_handler_uses_sanitized_server_prefix( tmp_path, monkeypatch: pytest.MonkeyPatch, ): - loop = _make_loop(tmp_path, mcp_servers={"remote_": object()}) + provider, registry = _make_provider( + mcp_servers={"remote_": _stdio_server("remote")} + ) connect_count = 0 class _FakeSession: @@ -578,15 +603,15 @@ async def _fake_connect(servers, registry): monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) - await loop._connect_mcp() - old_tool = loop.tools.get("mcp_remote_quote") + await provider.connect() + old_tool = registry.get("mcp_remote_quote") assert isinstance(old_tool, MCPToolWrapper) output = await old_tool.execute() assert output == "recovered" assert connect_count == 2 - assert loop.tools.get("mcp_remote_quote") is not old_tool + assert registry.get("mcp_remote_quote") is not old_tool @pytest.mark.asyncio @@ -594,7 +619,9 @@ async def test_concurrent_mcp_reconnect_reuses_fresh_session( tmp_path, monkeypatch: pytest.MonkeyPatch, ): - loop = _make_loop(tmp_path, mcp_servers={"remote": object()}) + provider, registry = _make_provider( + mcp_servers={"remote": _stdio_server("remote")} + ) closed: list[str] = [] connect_count = 0 @@ -638,9 +665,9 @@ async def _fake_connect(servers, registry): monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) - await loop._connect_mcp() - old_alpha = loop.tools.get("mcp_remote_resource_alpha") - old_beta = loop.tools.get("mcp_remote_resource_beta") + await provider.connect() + old_alpha = registry.get("mcp_remote_resource_alpha") + old_beta = registry.get("mcp_remote_resource_beta") assert isinstance(old_alpha, MCPResourceWrapper) assert isinstance(old_beta, MCPResourceWrapper) diff --git a/tests/agent/test_mcp_reconnect_crash.py b/tests/agent/test_mcp_reconnect_crash.py index 5ca387cfbd9..86df1261148 100644 --- a/tests/agent/test_mcp_reconnect_crash.py +++ b/tests/agent/test_mcp_reconnect_crash.py @@ -15,15 +15,13 @@ import multiprocessing import socket import time -from unittest.mock import MagicMock import httpx import pytest -from nanobot.agent.loop import AgentLoop from nanobot.agent.tools import mcp as mcp_module -from nanobot.agent.tools.mcp import MCPToolWrapper -from nanobot.bus.queue import MessageBus +from nanobot.agent.tools.mcp import MCPProvider, MCPToolWrapper +from nanobot.agent.tools.registry import ToolRegistry from nanobot.config.schema import MCPServerConfig from nanobot.security import network as security_network @@ -113,18 +111,9 @@ def mcp_server_url(): process.join(timeout=2.0) -def _make_loop(tmp_path, *, mcp_servers: dict) -> AgentLoop: - bus = MessageBus() - provider = MagicMock() - provider.get_default_model.return_value = "test-model" - provider.generation.max_tokens = 4096 - return AgentLoop( - bus=bus, - provider=provider, - workspace=tmp_path, - model="test-model", - mcp_servers=mcp_servers, - ) +def _make_provider(*, mcp_servers: dict) -> tuple[MCPProvider, ToolRegistry]: + registry = ToolRegistry() + return MCPProvider(mcp_servers, registry), registry @pytest.fixture(autouse=True) @@ -170,12 +159,12 @@ async def test_mcp_reconnect_after_session_timeout(tmp_path, mcp_server_url): tool_timeout=_TOOL_TIMEOUT_SECONDS, enabled_tools=["*"], ) - loop = _make_loop(tmp_path, mcp_servers={"repro": cfg}) + provider, registry = _make_provider(mcp_servers={"repro": cfg}) - await asyncio.create_task(loop._connect_mcp()) - assert "repro" in loop._mcp_stacks + await asyncio.create_task(provider.connect()) + assert provider.connected_server_names == {"repro"} - tool = loop.tools.get("mcp_repro_greet") + tool = registry.get("mcp_repro_greet") assert isinstance(tool, MCPToolWrapper) output = await asyncio.create_task(tool.execute(name="first")) @@ -187,7 +176,7 @@ async def test_mcp_reconnect_after_session_timeout(tmp_path, mcp_server_url): output = await asyncio.create_task(tool.execute(name="second")) assert "Hello, second" in output - await asyncio.create_task(loop.close_mcp()) + await asyncio.create_task(provider.aclose()) @pytest.mark.asyncio @@ -203,10 +192,10 @@ async def test_mcp_reconnect_during_shutdown_does_not_crash( tool_timeout=_TOOL_TIMEOUT_SECONDS, enabled_tools=["*"], ) - loop = _make_loop(tmp_path, mcp_servers={"repro": cfg}) + provider, registry = _make_provider(mcp_servers={"repro": cfg}) - await asyncio.create_task(loop._connect_mcp()) - tool = loop.tools.get("mcp_repro_greet") + await asyncio.create_task(provider.connect()) + tool = registry.get("mcp_repro_greet") assert isinstance(tool, MCPToolWrapper) await asyncio.create_task(tool.execute(name="first")) @@ -224,7 +213,7 @@ async def gated_connect(*args, **kwargs): monkeypatch.setattr(mcp_module, "connect_mcp_servers", gated_connect) call_task = asyncio.create_task(tool.execute(name="second")) await asyncio.wait_for(reconnect_started.wait(), timeout=5) - close_task = asyncio.create_task(loop.close_mcp()) + close_task = asyncio.create_task(provider.aclose()) await asyncio.sleep(0) finish_reconnect.set() @@ -245,4 +234,4 @@ def capture_unhandled(_loop, context): unhandled.append(exc) assert not unhandled, f"Unhandled exception leaked during reconnect/shutdown: {unhandled[0]}" - assert loop._mcp_stacks == {} + assert provider.connected_server_names == set() diff --git a/tests/agent/test_task_cancel.py b/tests/agent/test_task_cancel.py index bb68183d936..683f8032d9c 100644 --- a/tests/agent/test_task_cancel.py +++ b/tests/agent/test_task_cancel.py @@ -56,7 +56,7 @@ async def test_stop_no_active_task(self): assert "No active task" in out.content @pytest.mark.asyncio - async def test_close_mcp_cancels_active_turn_before_resources(self): + async def test_aclose_cancels_active_turn_before_resources(self): loop, _bus = _make_loop() events: list[str] = [] @@ -76,14 +76,13 @@ async def close_subagents(): 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() + await loop.aclose() assert events == ["turn_cancelled", "resources_closed"] assert task.cancelled() @pytest.mark.asyncio - async def test_close_mcp_serializes_duplicate_cleanup(self): + async def test_aclose_serializes_duplicate_cleanup(self): loop, _bus = _make_loop() entered = asyncio.Event() release = asyncio.Event() @@ -100,14 +99,13 @@ async def close_subagents(): 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) + first = asyncio.create_task(loop.aclose()) + await entered.wait() + second = asyncio.create_task(loop.aclose()) + await asyncio.sleep(0) + assert not second.done() + release.set() + await asyncio.gather(first, second) assert max_concurrent == 1 @@ -172,8 +170,7 @@ class TestDispatch: @pytest.mark.asyncio async def test_run_logs_and_continues_after_leaked_cancelled_error(self, monkeypatch): loop, bus = _make_loop() - loop._connect_mcp = AsyncMock() - loop.close_mcp = AsyncMock() + loop.aclose = AsyncMock() loop.auto_compact.check_expired = MagicMock() warnings: list[str] = [] calls = 0 diff --git a/tests/agent/tools/test_self_tool.py b/tests/agent/tools/test_self_tool.py index 6f5cffa245b..42727f9e78e 100644 --- a/tests/agent/tools/test_self_tool.py +++ b/tests/agent/tools/test_self_tool.py @@ -493,20 +493,6 @@ async def test_modify_workspace_preserves_display_compatibility(self): assert "Set workspace" in result assert tool._runtime_control.snapshot().workspace == "/new/path" - @pytest.mark.asyncio - async def test_modify_mcp_servers_blocked(self): - """_mcp_servers contains API credentials — must be blocked.""" - tool = _make_tool() - result = await tool.execute(action="set", key="_mcp_servers", value={"evil": "leaked"}) - assert "protected" in result - - @pytest.mark.asyncio - async def test_modify_mcp_stacks_blocked(self): - """_mcp_stacks holds connection handles — must be blocked.""" - tool = _make_tool() - result = await tool.execute(action="set", key="_mcp_stacks", value={}) - assert "protected" in result - @pytest.mark.asyncio async def test_modify_pending_queues_blocked(self): """_pending_queues controls message routing — must be blocked.""" @@ -535,13 +521,6 @@ async def test_modify_background_tasks_blocked(self): result = await tool.execute(action="set", key="_background_tasks", value=[]) assert "protected" in result - @pytest.mark.asyncio - async def test_inspect_mcp_servers_blocked(self): - """_mcp_servers contains credentials — check must be blocked too.""" - tool = _make_tool() - result = await tool.execute(action="check", key="_mcp_servers") - assert "not accessible" in result - @pytest.mark.asyncio async def test_modify_wrapped_denied(self): """__wrapped__ allows decorator bypass — must be denied.""" diff --git a/tests/cli/test_agent_interactive.py b/tests/cli/test_agent_interactive.py index 98ba559a5b8..22df60977f9 100644 --- a/tests/cli/test_agent_interactive.py +++ b/tests/cli/test_agent_interactive.py @@ -64,7 +64,7 @@ def from_config(cls, _config, bus, **_kwargs): 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 @@ async def run(self) -> None: 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 @@ async def close_mcp(self) -> None: 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) diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index 55d77edb5ce..b55bb2fd6f9 100644 --- a/tests/cli/test_commands.py +++ b/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 __init__(self, *args, **kwargs) -> 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 __init__(self, *args, **kwargs) -> 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.cron.service.CronService", _FakeCron) @@ -1712,7 +1712,7 @@ def __init__(self, *args, **kwargs) -> 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.cron.service.CronService", _FakeCron) @@ -1768,7 +1768,7 @@ def __init__(self, *args, **kwargs) -> 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.cron.service.CronService", _FakeCron) @@ -2062,7 +2062,7 @@ def __init__(self, *args, **kwargs) -> None: 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 from_config(cls, config, bus=None, **extra): 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 _fake_create_app( 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 @@ async def process_direct(self, *_args, **_kwargs): 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 @@ async def submit_cron_turn(self, msg: InboundMessage): 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 @@ async def run(self) -> None: 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 llm_runtime(self) -> None: 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 @@ async def read(self, _size: int) -> bytes: 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 llm_runtime(self) -> None: 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 @@ async def _fake_start_server(_handler, _host: str, _port: int): 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 @@ async def _fake_start_server(_handler, _host: str, _port: int): 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 @@ async def run(self) -> None: 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 _restore() -> None: 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 diff --git a/tests/cli/test_gateway_runtime.py b/tests/cli/test_gateway_runtime.py index daf02235025..d5fb0c51242 100644 --- a/tests/cli/test_gateway_runtime.py +++ b/tests/cli/test_gateway_runtime.py @@ -22,7 +22,7 @@ def __init__(self, events: list[str] | None = None) -> None: 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 @@ async def close_mcp(self) -> None: 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 @@ async def stop_all(self) -> None: 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 background_work() -> 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 background_work() -> 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 diff --git a/tests/test_api_attachment.py b/tests/test_api_attachment.py index a12c77ba0b5..cb23f669f85 100644 --- a/tests/test_api_attachment.py +++ b/tests/test_api_attachment.py @@ -32,8 +32,7 @@ def _make_mock_agent(response_text: str = "mock response") -> MagicMock: agent = MagicMock() agent.process_direct = AsyncMock(return_value=response_text) - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} return agent diff --git a/tests/test_api_stream.py b/tests/test_api_stream.py index f3487e4f516..35847f3cee5 100644 --- a/tests/test_api_stream.py +++ b/tests/test_api_stream.py @@ -64,8 +64,7 @@ def test_sse_done_format() -> None: def _make_streaming_agent(tokens: list[str]) -> MagicMock: """Create a mock agent that streams tokens via on_stream callback.""" agent = MagicMock() - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() async def fake_process_direct(*, content="", media=None, session_key="", channel="", chat_id="", on_stream=None, @@ -136,8 +135,7 @@ async def test_stream_false_returns_json(aiohttp_client) -> None: """stream=false should still return regular JSON response.""" agent = MagicMock() agent.process_direct = AsyncMock(return_value="normal reply") - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -160,8 +158,7 @@ async def test_stream_default_is_false(aiohttp_client) -> None: """Omitting stream should behave like stream=false.""" agent = MagicMock() agent.process_direct = AsyncMock(return_value="default reply") - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -217,8 +214,7 @@ async def fake_process_direct(**kwargs): agent = MagicMock() agent.process_direct = fake_process_direct - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -251,8 +247,7 @@ async def fake_process_direct(*, on_stream=None, on_stream_end=None, **kwargs): return "planning final" agent.process_direct = fake_process_direct - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -291,8 +286,7 @@ async def fake_process_direct(*, on_stream=None, on_stream_end=None, **kwargs): return "plain final" agent.process_direct = fake_process_direct - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -334,8 +328,7 @@ async def fake_process_direct(*, session_key="", on_stream=None, on_stream_end=N agent = MagicMock() agent.process_direct = fake_process_direct - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -364,8 +357,7 @@ async def boom(**kwargs): raise RuntimeError("backend blew up") agent.process_direct = boom - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) diff --git a/tests/test_nanobot_facade.py b/tests/test_nanobot_facade.py index 9c9f90267aa..6a989f883f9 100644 --- a/tests/test_nanobot_facade.py +++ b/tests/test_nanobot_facade.py @@ -1637,37 +1637,40 @@ async def test_runtime_helpers_expose_model_workspace_and_compact(tmp_path): @pytest.mark.asyncio -async def test_aclose_delegates_to_loop_close_mcp(tmp_path): +async def test_aclose_releases_loop_and_mcp_provider(tmp_path): config_path = _write_config(tmp_path) bot = Nanobot.from_config(config_path, workspace=tmp_path) - bot._loop.close_mcp = AsyncMock() + bot._loop.aclose = AsyncMock() + assert bot._mcp_provider is not None + bot._mcp_provider.aclose = AsyncMock() await bot.aclose() - bot._loop.close_mcp.assert_awaited_once() + bot._loop.aclose.assert_awaited_once() + bot._mcp_provider.aclose.assert_awaited_once() @pytest.mark.asyncio async def test_context_manager_calls_aclose_on_exit(tmp_path): config_path = _write_config(tmp_path) bot = Nanobot.from_config(config_path, workspace=tmp_path) - bot._loop.close_mcp = AsyncMock() + bot._loop.aclose = AsyncMock() async with bot as b: assert b is bot - bot._loop.close_mcp.assert_awaited_once() + bot._loop.aclose.assert_awaited_once() @pytest.mark.asyncio async def test_context_manager_does_not_swallow_exceptions(tmp_path): config_path = _write_config(tmp_path) bot = Nanobot.from_config(config_path, workspace=tmp_path) - bot._loop.close_mcp = AsyncMock() + bot._loop.aclose = AsyncMock() with pytest.raises(ValueError): async with bot as b: assert b is bot raise ValueError("boom") - bot._loop.close_mcp.assert_awaited_once() + bot._loop.aclose.assert_awaited_once() diff --git a/tests/test_openai_api.py b/tests/test_openai_api.py index a07729511a9..013e84907ed 100644 --- a/tests/test_openai_api.py +++ b/tests/test_openai_api.py @@ -34,8 +34,7 @@ def _make_mock_agent(response_text: str = "mock response") -> MagicMock: agent = MagicMock() agent.process_direct = AsyncMock(return_value=response_text) - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {"prompt_tokens": 100, "completion_tokens": 50} return agent @@ -149,6 +148,59 @@ async def test_api_routes_allow_requests_without_configured_api_key(aiohttp_clie mock_agent.process_direct.assert_called_once() +@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed") +@pytest.mark.asyncio +async def test_api_prepares_application_resources_before_each_turn(aiohttp_client) -> None: + events: list[str] = [] + agent = _make_mock_agent() + + async def prepare_agent() -> None: + events.append("prepare") + + async def process_direct(**_kwargs): + events.append("process") + return "ready" + + agent.process_direct = process_direct + app = create_app(agent, prepare_agent=prepare_agent) + client = await aiohttp_client(app) + + response = await client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status == 200 + assert events == ["prepare", "process"] + + +@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed") +@pytest.mark.asyncio +async def test_api_preparation_is_bounded_by_request_timeout(aiohttp_client) -> None: + agent = _make_mock_agent() + started = asyncio.Event() + + async def prepare_agent() -> None: + started.set() + await asyncio.Event().wait() + + app = create_app( + agent, + request_timeout=0.01, + prepare_agent=prepare_agent, + ) + client = await aiohttp_client(app) + + response = await client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hello"}]}, + ) + + assert started.is_set() + assert response.status == 504 + agent.process_direct.assert_not_awaited() + + @pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed") @pytest.mark.asyncio async def test_no_user_message_returns_400(aiohttp_client, app) -> None: @@ -275,8 +327,7 @@ async def fake_process(content, session_key="", channel="", chat_id="", **kwargs agent = MagicMock() agent.process_direct = fake_process - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -315,8 +366,7 @@ async def slow_process(content, session_key="", channel="", chat_id="", **kwargs agent = MagicMock() agent.process_direct = slow_process - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -433,8 +483,7 @@ async def always_empty(content, session_key="", channel="", chat_id="", **kwargs agent = MagicMock() agent.process_direct = always_empty - agent._connect_mcp = AsyncMock() - agent.close_mcp = AsyncMock() + agent.aclose = AsyncMock() agent._last_usage = {} app = create_app(agent, model_name="m", api_key=API_KEY) @@ -457,7 +506,6 @@ async def test_process_direct_accepts_media() -> None: from nanobot.bus.runtime_events import RuntimeEventPublisher loop = AgentLoop.__new__(AgentLoop) - loop._connect_mcp = AsyncMock() loop._session_locks = {} loop.runtime_event_publisher = RuntimeEventPublisher() diff --git a/tests/tools/test_exec_session_tools.py b/tests/tools/test_exec_session_tools.py index a042a098a84..a68bab8678a 100644 --- a/tests/tools/test_exec_session_tools.py +++ b/tests/tools/test_exec_session_tools.py @@ -12,7 +12,6 @@ import pytest -from nanobot.agent import context as agent_context from nanobot.agent.loop import AgentLoop from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context from nanobot.agent.tools.exec_session import ( @@ -712,7 +711,7 @@ async def run() -> None: asyncio.run(run()) -def test_agent_loop_shutdown_closes_exec_sessions(tmp_path, monkeypatch): +def test_agent_loop_shutdown_closes_exec_sessions(tmp_path): async def run() -> None: manager = ExecSessionManager() tool = ExecTool(working_dir=str(tmp_path), timeout=30, session_manager=manager) @@ -723,14 +722,13 @@ async def run() -> None: sid = _session_id(initial) process = manager._sessions[sid].process - monkeypatch.setattr(agent_context, "close_mcp", lambda _state: asyncio.sleep(0)) loop = object.__new__(AgentLoop) loop._background_tasks = set() loop._exec_session_manager = manager loop.subagents = SimpleNamespace(close=AsyncMock()) - await loop.close_mcp() - await loop.close_mcp() + await loop.aclose() + await loop.aclose() assert process.returncode is not None assert manager._sessions == {} @@ -739,7 +737,7 @@ async def run() -> None: asyncio.run(run()) -def test_agent_loop_shutdown_attempts_all_cleanup_after_errors(monkeypatch): +def test_agent_loop_shutdown_attempts_all_cleanup_after_errors(): async def run() -> None: loop = object.__new__(AgentLoop) loop._background_tasks = set() @@ -749,16 +747,12 @@ async def run() -> None: loop._exec_session_manager = SimpleNamespace( close_all=AsyncMock(side_effect=OSError("exec cleanup failed")), ) - close_mcp = AsyncMock() - monkeypatch.setattr(agent_context, "close_mcp", close_mcp) - with pytest.raises(BaseExceptionGroup) as exc_info: - await loop.close_mcp() + await loop.aclose() assert len(exc_info.value.exceptions) == 2 loop.subagents.close.assert_awaited_once() loop._exec_session_manager.close_all.assert_awaited_once() - close_mcp.assert_awaited_once_with(loop) asyncio.run(run()) @@ -892,7 +886,7 @@ async def run() -> None: asyncio.run(run()) -def test_agent_loop_shutdown_preserves_single_cleanup_error(monkeypatch): +def test_agent_loop_shutdown_preserves_single_cleanup_error(): async def run() -> None: loop = object.__new__(AgentLoop) loop._background_tasks = set() @@ -900,13 +894,9 @@ async def run() -> None: close=AsyncMock(side_effect=RuntimeError("subagent cleanup failed")), ) loop._exec_session_manager = SimpleNamespace(close_all=AsyncMock()) - close_mcp = AsyncMock() - monkeypatch.setattr(agent_context, "close_mcp", close_mcp) - with pytest.raises(RuntimeError, match="subagent cleanup failed"): - await loop.close_mcp() + await loop.aclose() loop._exec_session_manager.close_all.assert_awaited_once() - close_mcp.assert_awaited_once_with(loop) asyncio.run(run()) diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 008c887f7da..2f0f8bbc82d 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -13,6 +13,7 @@ import nanobot.agent.tools.mcp as mcp_mod from nanobot.agent.tools.mcp import ( MCPPromptWrapper, + MCPProvider, MCPResourceWrapper, MCPToolWrapper, _normalize_windows_stdio_command, @@ -153,7 +154,7 @@ def _make_wrapper(session: object, *, timeout: float = 0.1) -> MCPToolWrapper: @pytest.mark.asyncio -async def test_connect_missing_servers_propagates_external_cancellation(monkeypatch) -> None: +async def test_mcp_provider_connect_propagates_external_cancellation(monkeypatch) -> None: started = asyncio.Event() async def connect_mcp_servers(_servers: dict, _registry: ToolRegistry) -> dict: @@ -161,24 +162,20 @@ async def connect_mcp_servers(_servers: dict, _registry: ToolRegistry) -> dict: await asyncio.sleep(60) return {} - class State: - pass - - state = State() - state._mcp_closing = False - state._mcp_servers = {"test": MCPServerConfig(command="fake")} - state._mcp_stacks = {} - state._mcp_connecting = False + provider = MCPProvider( + {"test": MCPServerConfig(command="fake")}, + ToolRegistry(), + ) monkeypatch.setattr(mcp_mod, "connect_mcp_servers", connect_mcp_servers) - task = asyncio.create_task(mcp_mod.connect_missing_servers(state, ToolRegistry())) + task = asyncio.create_task(provider.connect()) await asyncio.wait_for(started.wait(), timeout=1.0) task.cancel() with pytest.raises(asyncio.CancelledError): await task - assert state._mcp_connecting is False + assert provider.connected_server_names == set() @pytest.mark.asyncio @@ -223,25 +220,20 @@ async def reachable(_url: str) -> bool: rejected_streamable_http, ) - class State: - pass - - state = State() - state._mcp_closing = False - state._mcp_servers = { - "xmind": MCPServerConfig( + provider = MCPProvider( + { + "xmind": MCPServerConfig( type="streamableHttp", auth="oauth", url="https://app.xmind.com/api/mcp", - ) - } - state._mcp_stacks = {} - state._mcp_runtime_statuses = {} - state._mcp_connecting = False + ) + }, + ToolRegistry(), + ) - await mcp_mod.connect_missing_servers(state, ToolRegistry()) + await provider.connect() - snapshot = mcp_mod.runtime_status(state) + snapshot = provider.runtime_status() assert snapshot == {"xmind": "failed"} assert "saved-oauth-secret" not in str(snapshot) assert "app.xmind.com" not in str(snapshot) @@ -1900,7 +1892,11 @@ def test_long_server_name_tools_are_matched_by_server_name() -> None: assert len(wrapper.name) == 64 assert not wrapper.name.startswith(mcp_mod._tool_prefix(server_name)) - mcp_mod._attach_reconnect_handlers(SimpleNamespace(), registry, {server_name}) + provider = MCPProvider( + {server_name: MCPServerConfig(command="fake")}, + registry, + ) + provider._attach_reconnect_handlers({server_name}) assert wrapper._reconnect is not None assert other_wrapper._reconnect is None diff --git a/tests/webui/test_settings_routes.py b/tests/webui/test_settings_routes.py index c048d8ec114..b77910883a7 100644 --- a/tests/webui/test_settings_routes.py +++ b/tests/webui/test_settings_routes.py @@ -1,7 +1,8 @@ from __future__ import annotations +import asyncio import json -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from pathlib import Path from types import SimpleNamespace from unittest.mock import ANY, AsyncMock, MagicMock @@ -22,6 +23,7 @@ def _router( authorized: bool = True, config_path: Path | None = None, mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None, + mcp_reload: Callable[[], Awaitable[dict[str, object]]] | None = None, ) -> WebUISettingsRouter: return WebUISettingsRouter( settings=WebUISettingsServices.create(config_path or get_config_path()), @@ -37,6 +39,7 @@ def _router( runtime_surface="browser", runtime_capabilities={}, mcp_runtime_status=mcp_runtime_status, + mcp_reload=mcp_reload, mcp_oauth_redirect_uri=lambda _request: "https://gateway.example/auth/mcp/callback", ) @@ -89,6 +92,33 @@ def runtime_snapshot() -> Mapping[str, str]: assert snapshot_calls == 1 +@pytest.mark.asyncio +async def test_mcp_reload_callback_is_bounded( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started = asyncio.Event() + + async def reload_mcp() -> dict[str, object]: + started.set() + await asyncio.Event().wait() + return {"ok": True} + + monkeypatch.setattr( + "nanobot.webui.settings_routes._MCP_RELOAD_TIMEOUT_SECONDS", + 0.01, + ) + router = _router(mcp_reload=reload_mcp) + + result = await router._reload_mcp_runtime() + + assert started.is_set() + assert result == { + "ok": False, + "message": "MCP hot reload timed out. Restart nanobot to pick up changes.", + "requires_restart": True, + } + + @pytest.mark.asyncio async def test_mcp_oauth_start_uses_gateway_callback_and_requires_api_auth(monkeypatch) -> None: config = SimpleNamespace( From 96dfef235b9b984c9dd2cb3f69a11d79cbe9da7e Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Wed, 12 Aug 2026 01:29:22 +0800 Subject: [PATCH 2/6] fix: roll back cancelled MCP connection batches --- nanobot/agent/tools/mcp.py | 82 +++++++++++++++++++++--------------- tests/tools/test_mcp_tool.py | 53 +++++++++++++++++++++++ 2 files changed, 101 insertions(+), 34 deletions(-) diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index 5e736c46800..aef9a324d10 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -63,6 +63,28 @@ class MCPConnection(Protocol): async def aclose(self) -> None: ... +async def _close_mcp_connection(name: str, connection: MCPConnection) -> None: + try: + await connection.aclose() + except asyncio.CancelledError: + if task_is_cancelling(): + raise + logger.debug("MCP server '{}' cleanup error (can be ignored)", name) + except (RuntimeError, BaseExceptionGroup): + logger.debug("MCP server '{}' cleanup error (can be ignored)", name) + + +async def _close_mcp_connections(connections: Mapping[str, MCPConnection]) -> None: + cancellation: asyncio.CancelledError | None = None + for name, connection in connections.items(): + try: + await _close_mcp_connection(name, connection) + except asyncio.CancelledError as exc: + cancellation = cancellation or exc + if cancellation is not None: + raise cancellation + + class _OwnedMCPConnection: """Close an MCP transport from the task that originally opened it.""" @@ -1273,15 +1295,29 @@ async def own_connection() -> None: return name, connection server_stacks: dict[str, MCPConnection] = {} + attempted_names: list[str] = [] - for name, cfg in mcp_servers.items(): + try: + for name, cfg in mcp_servers.items(): + attempted_names.append(name) + try: + result = await connect_single_server(name, cfg) + except Exception as e: + _log_mcp_connection_failure(name, e) + continue + if result[1] is not None: + server_stacks[result[0]] = result[1] + except BaseException: + # Callers can bound readiness/reload with a timeout. If cancellation + # interrupts a later server, ownership of earlier connections has not + # transferred yet, so roll the whole batch back before propagating it. + for name in attempted_names: + _unregister_server_tools(registry, name) try: - result = await connect_single_server(name, cfg) - except Exception as e: - _log_mcp_connection_failure(name, e) - continue - if result[1] is not None: - server_stacks[result[0]] = result[1] + await _close_mcp_connections(server_stacks) + except BaseException as cleanup_exc: + logger.debug("MCP batch rollback cleanup error (can be ignored): {}", cleanup_exc) + raise return server_stacks @@ -1373,28 +1409,6 @@ def _record_connection_result( self._set_runtime_status(connected_names, "connected") self._set_runtime_status(attempted_names - connected_names, "failed") - @staticmethod - async def _close_connection(name: str, connection: MCPConnection) -> None: - try: - await connection.aclose() - except asyncio.CancelledError: - if task_is_cancelling(): - raise - logger.debug("MCP server '{}' cleanup error (can be ignored)", name) - except (RuntimeError, BaseExceptionGroup): - logger.debug("MCP server '{}' cleanup error (can be ignored)", name) - - @classmethod - async def _close_connections(cls, connections: Mapping[str, MCPConnection]) -> None: - cancellation: asyncio.CancelledError | None = None - for name, connection in connections.items(): - try: - await cls._close_connection(name, connection) - except asyncio.CancelledError as exc: - cancellation = cancellation or exc - if cancellation is not None: - raise cancellation - async def connect(self) -> None: """Connect configured servers that are not currently live.""" async with self._lock: @@ -1432,7 +1446,7 @@ async def connect(self) -> None: try: connected = await connect_mcp_servers(missing_servers, self._registry) if self._closing: - await self._close_connections(connected) + await _close_mcp_connections(connected) return self._connections.update(connected) self._record_connection_result(missing_servers, connected) @@ -1520,7 +1534,7 @@ async def reload(self) -> dict[str, Any]: self._set_runtime_status(to_connect, "connecting") connected = await connect_mcp_servers(to_connect, self._registry) if self._closing: - await self._close_connections(connected) + await _close_mcp_connections(connected) return self._closing_result() self._connections.update(connected) self._record_connection_result(to_connect, connected) @@ -1632,7 +1646,7 @@ async def _refresh_terminated_server( self._registry, ) if self._closing: - await self._close_connections(connected) + await _close_mcp_connections(connected) return None self._connections.update(connected) self._record_connection_result({server_name}, connected) @@ -1649,7 +1663,7 @@ async def _close_server(self, server_name: str) -> None: connection = self._connections.pop(server_name, None) if connection is None: return - await self._close_connection(server_name, connection) + await _close_mcp_connection(server_name, connection) async def aclose(self) -> None: """Close every connection while excluding reconnect and hot reload.""" @@ -1660,7 +1674,7 @@ async def aclose(self) -> None: self._runtime_statuses.clear() for name in self._servers: _unregister_server_tools(self._registry, name) - await self._close_connections(connections) + await _close_mcp_connections(connections) def _server_signature(cfg: Any) -> Any: diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 2f0f8bbc82d..e3913918bad 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -1255,6 +1255,59 @@ async def _blocking_stdio_client(_params: object): await asyncio.wait_for(closed.wait(), timeout=1.0) +@pytest.mark.asyncio +async def test_connect_mcp_servers_rolls_back_completed_batch_on_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + slow_started = asyncio.Event() + closed: list[str] = [] + sessions = {"fast": _make_fake_session(["demo"])} + + class _SelectiveClientSession: + def __init__(self, read: object, _write: object) -> None: + self._session = sessions[str(read)] + + async def __aenter__(self) -> object: + return self._session + + async def __aexit__(self, exc_type, exc, tb) -> bool: + return False + + @asynccontextmanager + async def _selective_stdio_client(params: object): + command = str(params.command) + try: + if command == "slow": + slow_started.set() + await asyncio.Event().wait() + yield command, object() + finally: + closed.append(command) + + monkeypatch.setattr(sys.modules["mcp"], "ClientSession", _SelectiveClientSession) + monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _selective_stdio_client) + + registry = ToolRegistry() + task = asyncio.create_task( + connect_mcp_servers( + { + "fast": MCPServerConfig(command="fast"), + "slow": MCPServerConfig(command="slow"), + }, + registry, + ) + ) + await asyncio.wait_for(slow_started.wait(), timeout=1.0) + assert registry.tool_names == ["mcp_fast_demo"] + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert registry.tool_names == [] + assert sorted(closed) == ["fast", "slow"] + + @pytest.mark.asyncio async def test_connect_mcp_servers_streamable_http_uses_finite_timeout( fake_mcp_runtime: dict[str, object | None], From cc8ba65ce2271aed640a05382a681d3601c62630 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 12 Aug 2026 11:24:22 +0800 Subject: [PATCH 3/6] fix(mcp): mark cancelled connection attempts failed --- nanobot/agent/tools/mcp.py | 8 ++++-- tests/agent/test_mcp_connection.py | 41 ++++++++++++++++++++++++++++++ tests/tools/test_mcp_tool.py | 1 + 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index aef9a324d10..f0b8bbe5cb5 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -1459,9 +1459,9 @@ async def connect(self) -> None: "(will retry on the next readiness check)" ) except asyncio.CancelledError: + self._set_runtime_status(missing_servers, "failed") if task_is_cancelling(): raise - self._set_runtime_status(missing_servers, "failed") logger.warning( "MCP connection cancelled (will retry on the next readiness check)" ) @@ -1532,7 +1532,11 @@ async def reload(self) -> dict[str, Any]: connected: dict[str, MCPConnection] = {} if to_connect: self._set_runtime_status(to_connect, "connecting") - connected = await connect_mcp_servers(to_connect, self._registry) + try: + connected = await connect_mcp_servers(to_connect, self._registry) + except BaseException: + self._set_runtime_status(to_connect, "failed") + raise if self._closing: await _close_mcp_connections(connected) return self._closing_result() diff --git a/tests/agent/test_mcp_connection.py b/tests/agent/test_mcp_connection.py index dde56650e3d..724e1af0c9a 100644 --- a/tests/agent/test_mcp_connection.py +++ b/tests/agent/test_mcp_connection.py @@ -414,6 +414,47 @@ async def _fake_connect(servers, registry): assert closed == ["browserbase"] +@pytest.mark.asyncio +async def test_reload_timeout_marks_attempted_server_failed_and_allows_retry( + monkeypatch: pytest.MonkeyPatch, +): + server = _stdio_server("slow-mcp") + started = asyncio.Event() + attempts = 0 + + async def _fake_connect(servers, _registry): + nonlocal attempts + attempts += 1 + if attempts == 1: + started.set() + await asyncio.Event().wait() + stack = AsyncExitStack() + await stack.__aenter__() + return {name: stack for name in servers} + + monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect) + provider = MCPProvider( + {"test": server}, + ToolRegistry(), + server_loader=lambda: {"test": server}, + ) + + reload_task = asyncio.create_task(provider.reload()) + await asyncio.wait_for(started.wait(), timeout=1.0) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(reload_task, timeout=0.01) + + assert provider.connected_server_names == set() + assert provider.runtime_status() == {"test": "failed"} + + result = await provider.reload() + + assert result["ok"] is True + assert provider.connected_server_names == {"test"} + assert provider.runtime_status() == {"test": "connected"} + await provider.aclose() + + @pytest.mark.asyncio async def test_reload_mcp_servers_retries_configured_server_without_live_stack( tmp_path, diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index e3913918bad..8ed4944f436 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -176,6 +176,7 @@ async def connect_mcp_servers(_servers: dict, _registry: ToolRegistry) -> dict: await task assert provider.connected_server_names == set() + assert provider.runtime_status() == {"test": "failed"} @pytest.mark.asyncio From e3a369f38105529a740825086530cbcfa2f6e3bd Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 12 Aug 2026 16:14:47 +0800 Subject: [PATCH 4/6] fix(mcp): make external lifecycle ownership explicit --- docs/architecture.md | 27 ++++++++++ nanobot/agent/loop.py | 28 ++++++++++ nanobot/cli/agent.py | 1 + nanobot/cli/commands.py | 1 + nanobot/cli/gateway_runtime.py | 1 + nanobot/nanobot.py | 1 + tests/agent/test_loop_tool_context.py | 76 +++++++++++++++++++++++++++ tests/test_nanobot_facade.py | 13 +++++ 8 files changed, 148 insertions(+) diff --git a/docs/architecture.md b/docs/architecture.md index 6599423f24d..1d5ed76c488 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -51,6 +51,33 @@ Main files: - feeds tool results back into the model; - stops when a final answer is produced or runtime limits are hit. +MCP connections are application-owned infrastructure. The CLI, gateway, API, +and `Nanobot.from_config()` create an `MCPProvider` and share its `ToolRegistry` +with `AgentLoop`; the loop does not connect, reload, or close MCP servers. Custom +low-level compositions using `AgentLoop.from_config()` must likewise own the +provider lifecycle: share its registry, pass +`mcp_runtime_owned_externally=True`, await `MCPProvider.connect()` before using +the loop, and await `MCPProvider.aclose()` during shutdown. The complete shape is: + +```python +registry = ToolRegistry() +mcp_provider = MCPProvider.from_config(config, registry) +loop = AgentLoop.from_config( + config, + tool_registry=registry, + mcp_runtime_owned_externally=True, +) + +try: + await mcp_provider.connect() + await loop.process_direct(message) +finally: + try: + await loop.aclose() + finally: + await mcp_provider.aclose() +``` + Keep this split in mind when debugging. If a problem is about channel routing, session keys, workspace selection, or outbound delivery, start in `agent/loop.py`. If it is about provider calls, tool calls, streaming, or iteration limits, start in `agent/runner.py`. ## Providers diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index aae9d1a376e..dd228a98f54 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -458,14 +458,42 @@ def from_config( cls, config: Config, bus: MessageBus | None = None, + *, + mcp_runtime_owned_externally: bool = False, **extra: Any, ) -> AgentLoop: """Create an AgentLoop from config with the common parameter set. + MCP lifecycle belongs to the application. Callers with configured MCP + servers must provide a shared ``tool_registry`` and explicitly confirm + that an application-owned ``MCPProvider`` manages that registry. + Extra keyword arguments are forwarded to ``AgentLoop.__init__``, allowing callers to override or extend the standard config-derived parameters (e.g. ``cron_service``, ``session_manager``). """ + if mcp_runtime_owned_externally: + if extra.get("tool_registry") is None: + raise ValueError( + "mcp_runtime_owned_externally=True requires a shared tool_registry" + ) + else: + from nanobot.agent.plugins import agent_plugin_mcp_servers + + configured_mcp = agent_plugin_mcp_servers( + config.workspace_path, + config.tools.mcp_servers, + ) + if configured_mcp: + raise ValueError( + "Configured MCP servers require an application-owned MCPProvider; " + "AgentLoop no longer manages MCP lifecycle. Use Nanobot.from_config(), " + "or create mcp_provider = MCPProvider.from_config(config, registry), " + "pass the registry to AgentLoop.from_config(..., " + "mcp_runtime_owned_externally=True), await mcp_provider.connect() " + "before using the loop, and await mcp_provider.aclose() during shutdown." + ) + from nanobot.providers.factory import make_provider if bus is None: diff --git a/nanobot/cli/agent.py b/nanobot/cli/agent.py index b8cd7647ba9..2e4a6254d7c 100644 --- a/nanobot/cli/agent.py +++ b/nanobot/cli/agent.py @@ -95,6 +95,7 @@ def agent( agent_loop = AgentLoop.from_config( runtime_config, bus, + mcp_runtime_owned_externally=True, provider=provider, cron_service=cron, image_generation_provider_configs=image_gen_provider_configs(runtime_config), diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index 318e55b5739..50aa0ca592c 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -358,6 +358,7 @@ def serve( try: agent_loop = AgentLoop.from_config( runtime_config, bus, + mcp_runtime_owned_externally=True, session_manager=session_manager, image_generation_provider_configs=image_gen_provider_configs(runtime_config), hook_factories=[create_file_edit_activity_hook], diff --git a/nanobot/cli/gateway_runtime.py b/nanobot/cli/gateway_runtime.py index f31cb0e90ee..d7835793807 100644 --- a/nanobot/cli/gateway_runtime.py +++ b/nanobot/cli/gateway_runtime.py @@ -422,6 +422,7 @@ def _load_gateway_provider_snapshot( # Create agent with cron service agent = AgentLoop.from_config( config, bus, + mcp_runtime_owned_externally=True, provider=provider_snapshot.provider, model=provider_snapshot.model, context_window_tokens=provider_snapshot.context_window_tokens, diff --git a/nanobot/nanobot.py b/nanobot/nanobot.py index 2720438386c..5cf35812058 100644 --- a/nanobot/nanobot.py +++ b/nanobot/nanobot.py @@ -133,6 +133,7 @@ def from_config( mcp_provider = MCPProvider.from_config(config, tools) loop = AgentLoop.from_config( config, + mcp_runtime_owned_externally=True, image_generation_provider_configs=image_gen_provider_configs(config), hook_factories=[create_file_edit_activity_hook], tool_registry=tools, diff --git a/tests/agent/test_loop_tool_context.py b/tests/agent/test_loop_tool_context.py index ddb6a764e97..9964a88ee83 100644 --- a/tests/agent/test_loop_tool_context.py +++ b/tests/agent/test_loop_tool_context.py @@ -11,9 +11,11 @@ current_request_context, reset_request_context, ) +from nanobot.agent.tools.mcp import MCPProvider from nanobot.agent.tools.registry import ToolRegistry from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus +from nanobot.config.schema import Config, MCPServerConfig from nanobot.providers.base import LLMResponse, ToolCallRequest from nanobot.session.turn_continuation import INTERNAL_CONTINUATION_META @@ -73,6 +75,80 @@ def test_loop_registers_default_tools_in_injected_registry(tmp_path: Path) -> No assert registry.has("read_file") +def _config_for_loop(tmp_path: Path, *, with_mcp: bool = False) -> Config: + payload: dict = {"agents": {"defaults": {"workspace": str(tmp_path)}}} + if with_mcp: + payload["tools"] = {"mcpServers": {"demo": {"command": "fake-mcp"}}} + return Config.model_validate(payload) + + +def _provider_for_loop() -> MagicMock: + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + return provider + + +def test_loop_from_config_rejects_silently_unmanaged_mcp(tmp_path: Path) -> None: + config = _config_for_loop(tmp_path, with_mcp=True) + + with pytest.raises(ValueError, match="application-owned MCPProvider") as exc_info: + AgentLoop.from_config(config, provider=_provider_for_loop()) + + message = str(exc_info.value) + assert "await mcp_provider.connect()" in message + assert "await mcp_provider.aclose()" in message + + +def test_loop_from_config_detects_plugin_only_mcp( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = _config_for_loop(tmp_path) + monkeypatch.setattr( + "nanobot.agent.plugins.agent_plugin_mcp_servers", + lambda _workspace, _configured: {"plugin-demo": MCPServerConfig(command="fake-mcp")}, + ) + + with pytest.raises(ValueError, match="application-owned MCPProvider"): + AgentLoop.from_config(config, provider=_provider_for_loop()) + + +def test_loop_from_config_requires_registry_for_external_mcp_owner(tmp_path: Path) -> None: + config = _config_for_loop(tmp_path, with_mcp=True) + + with pytest.raises(ValueError, match="requires a shared tool_registry"): + AgentLoop.from_config( + config, + provider=_provider_for_loop(), + mcp_runtime_owned_externally=True, + ) + + +def test_loop_from_config_accepts_explicit_external_mcp_owner(tmp_path: Path) -> None: + config = _config_for_loop(tmp_path, with_mcp=True) + registry = ToolRegistry() + mcp_provider = MCPProvider.from_config(config, registry) + + loop = AgentLoop.from_config( + config, + provider=_provider_for_loop(), + tool_registry=registry, + mcp_runtime_owned_externally=True, + ) + + assert loop.tools is registry + assert mcp_provider.configured_server_names == {"demo"} + + +def test_loop_from_config_without_mcp_keeps_low_level_usage(tmp_path: Path) -> None: + loop = AgentLoop.from_config( + _config_for_loop(tmp_path), + provider=_provider_for_loop(), + ) + + assert loop.tools.has("read_file") + + @pytest.mark.asyncio async def test_loop_binds_request_context_for_tool_execution(tmp_path: Path) -> None: provider = MagicMock() diff --git a/tests/test_nanobot_facade.py b/tests/test_nanobot_facade.py index 6a989f883f9..1d097ffdc49 100644 --- a/tests/test_nanobot_facade.py +++ b/tests/test_nanobot_facade.py @@ -101,6 +101,19 @@ def test_from_config_creates_instance(tmp_path): assert bot._loop.workspace == tmp_path +def test_from_config_composes_configured_mcp_outside_agent_loop(tmp_path): + config_path = _write_config( + tmp_path, + {"tools": {"mcpServers": {"demo": {"command": "fake-mcp"}}}}, + ) + + bot = Nanobot.from_config(config_path, workspace=tmp_path) + + assert bot._mcp_provider is not None + assert bot._mcp_provider.configured_server_names == {"demo"} + assert bot._mcp_provider._registry is bot._loop.tools + + def test_from_config_accepts_default_model_override(tmp_path): config_path = _write_config(tmp_path) From 0ebecb6363bcc0f80104016af85f8919cf037981 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 12 Aug 2026 16:53:36 +0800 Subject: [PATCH 5/6] docs: condense MCP lifecycle guidance --- docs/architecture.md | 31 +++++-------------------------- 1 file changed, 5 insertions(+), 26 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 1d5ed76c488..48a0827ca05 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -51,32 +51,11 @@ Main files: - feeds tool results back into the model; - stops when a final answer is produced or runtime limits are hit. -MCP connections are application-owned infrastructure. The CLI, gateway, API, -and `Nanobot.from_config()` create an `MCPProvider` and share its `ToolRegistry` -with `AgentLoop`; the loop does not connect, reload, or close MCP servers. Custom -low-level compositions using `AgentLoop.from_config()` must likewise own the -provider lifecycle: share its registry, pass -`mcp_runtime_owned_externally=True`, await `MCPProvider.connect()` before using -the loop, and await `MCPProvider.aclose()` during shutdown. The complete shape is: - -```python -registry = ToolRegistry() -mcp_provider = MCPProvider.from_config(config, registry) -loop = AgentLoop.from_config( - config, - tool_registry=registry, - mcp_runtime_owned_externally=True, -) - -try: - await mcp_provider.connect() - await loop.process_direct(message) -finally: - try: - await loop.aclose() - finally: - await mcp_provider.aclose() -``` +MCP connections are application-owned infrastructure. Composition roots create +an `MCPProvider`, share its `ToolRegistry` with `AgentLoop`, await `connect()` +before use, and guarantee `aclose()` during shutdown; the loop does not manage +that lifecycle. Custom `AgentLoop.from_config()` callers with configured MCP +must follow the same contract and pass `mcp_runtime_owned_externally=True`. Keep this split in mind when debugging. If a problem is about channel routing, session keys, workspace selection, or outbound delivery, start in `agent/loop.py`. If it is about provider calls, tool calls, streaming, or iteration limits, start in `agent/runner.py`. From fc2e507d414aa85c6d2a59fa33537d2e98bce94e Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 12 Aug 2026 17:20:28 +0800 Subject: [PATCH 6/6] refactor(agent): require caller-owned tool registry --- docs/architecture.md | 5 +- nanobot/agent/loop.py | 30 ++---------- nanobot/cli/agent.py | 1 - nanobot/cli/commands.py | 1 - nanobot/cli/gateway_runtime.py | 1 - nanobot/nanobot.py | 1 - tests/agent/test_auto_compact.py | 7 ++- tests/agent/test_loop_tool_context.py | 67 +++++---------------------- tests/agent/test_runtime_refresh.py | 7 ++- tests/agent/test_self_model_preset.py | 5 +- 10 files changed, 33 insertions(+), 92 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 48a0827ca05..afe23fc8d38 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -54,8 +54,9 @@ Main files: MCP connections are application-owned infrastructure. Composition roots create an `MCPProvider`, share its `ToolRegistry` with `AgentLoop`, await `connect()` before use, and guarantee `aclose()` during shutdown; the loop does not manage -that lifecycle. Custom `AgentLoop.from_config()` callers with configured MCP -must follow the same contract and pass `mcp_runtime_owned_externally=True`. +that lifecycle. `AgentLoop.from_config()` therefore requires a caller-owned +`ToolRegistry`; callers using MCP share it with their application-owned +`MCPProvider`. Keep this split in mind when debugging. If a problem is about channel routing, session keys, workspace selection, or outbound delivery, start in `agent/loop.py`. If it is about provider calls, tool calls, streaming, or iteration limits, start in `agent/runner.py`. diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index dd228a98f54..85cc23b4f60 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -459,41 +459,18 @@ def from_config( config: Config, bus: MessageBus | None = None, *, - mcp_runtime_owned_externally: bool = False, + tool_registry: ToolRegistry, **extra: Any, ) -> AgentLoop: """Create an AgentLoop from config with the common parameter set. - MCP lifecycle belongs to the application. Callers with configured MCP - servers must provide a shared ``tool_registry`` and explicitly confirm - that an application-owned ``MCPProvider`` manages that registry. + The tool registry is caller-owned so application composition can share + it with infrastructure such as an ``MCPProvider``. Extra keyword arguments are forwarded to ``AgentLoop.__init__``, allowing callers to override or extend the standard config-derived parameters (e.g. ``cron_service``, ``session_manager``). """ - if mcp_runtime_owned_externally: - if extra.get("tool_registry") is None: - raise ValueError( - "mcp_runtime_owned_externally=True requires a shared tool_registry" - ) - else: - from nanobot.agent.plugins import agent_plugin_mcp_servers - - configured_mcp = agent_plugin_mcp_servers( - config.workspace_path, - config.tools.mcp_servers, - ) - if configured_mcp: - raise ValueError( - "Configured MCP servers require an application-owned MCPProvider; " - "AgentLoop no longer manages MCP lifecycle. Use Nanobot.from_config(), " - "or create mcp_provider = MCPProvider.from_config(config, registry), " - "pass the registry to AgentLoop.from_config(..., " - "mcp_runtime_owned_externally=True), await mcp_provider.connect() " - "before using the loop, and await mcp_provider.aclose() during shutdown." - ) - from nanobot.providers.factory import make_provider if bus is None: @@ -536,6 +513,7 @@ def from_config( restart_mode=config.gateway.restart_mode, provider_snapshot_loader=provider_snapshot_loader, preset_snapshot_loader=preset_snapshot_loader, + tool_registry=tool_registry, **extra, ) diff --git a/nanobot/cli/agent.py b/nanobot/cli/agent.py index 2e4a6254d7c..b8cd7647ba9 100644 --- a/nanobot/cli/agent.py +++ b/nanobot/cli/agent.py @@ -95,7 +95,6 @@ def agent( agent_loop = AgentLoop.from_config( runtime_config, bus, - mcp_runtime_owned_externally=True, provider=provider, cron_service=cron, image_generation_provider_configs=image_gen_provider_configs(runtime_config), diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index 50aa0ca592c..318e55b5739 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -358,7 +358,6 @@ def serve( try: agent_loop = AgentLoop.from_config( runtime_config, bus, - mcp_runtime_owned_externally=True, session_manager=session_manager, image_generation_provider_configs=image_gen_provider_configs(runtime_config), hook_factories=[create_file_edit_activity_hook], diff --git a/nanobot/cli/gateway_runtime.py b/nanobot/cli/gateway_runtime.py index d7835793807..f31cb0e90ee 100644 --- a/nanobot/cli/gateway_runtime.py +++ b/nanobot/cli/gateway_runtime.py @@ -422,7 +422,6 @@ def _load_gateway_provider_snapshot( # Create agent with cron service agent = AgentLoop.from_config( config, bus, - mcp_runtime_owned_externally=True, provider=provider_snapshot.provider, model=provider_snapshot.model, context_window_tokens=provider_snapshot.context_window_tokens, diff --git a/nanobot/nanobot.py b/nanobot/nanobot.py index 5cf35812058..2720438386c 100644 --- a/nanobot/nanobot.py +++ b/nanobot/nanobot.py @@ -133,7 +133,6 @@ def from_config( mcp_provider = MCPProvider.from_config(config, tools) loop = AgentLoop.from_config( config, - mcp_runtime_owned_externally=True, image_generation_provider_configs=image_gen_provider_configs(config), hook_factories=[create_file_edit_activity_hook], tool_registry=tools, diff --git a/tests/agent/test_auto_compact.py b/tests/agent/test_auto_compact.py index 7f579366f9b..b576481f9d5 100644 --- a/tests/agent/test_auto_compact.py +++ b/tests/agent/test_auto_compact.py @@ -8,6 +8,7 @@ import pytest from nanobot.agent.loop import AgentLoop +from nanobot.agent.tools.registry import ToolRegistry from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus from nanobot.command import CommandContext @@ -193,7 +194,11 @@ def test_configured_idle_scan_interval_throttles_checks(self, tmp_path, monkeypa }) provider = MagicMock() provider.get_default_model.return_value = "test-model" - loop = AgentLoop.from_config(config, provider=provider) + loop = AgentLoop.from_config( + config, + tool_registry=ToolRegistry(), + provider=provider, + ) loop.auto_compact.check_expired = MagicMock() loop._check_expired_sessions_if_due() diff --git a/tests/agent/test_loop_tool_context.py b/tests/agent/test_loop_tool_context.py index 9964a88ee83..0d92d33b1d8 100644 --- a/tests/agent/test_loop_tool_context.py +++ b/tests/agent/test_loop_tool_context.py @@ -1,4 +1,5 @@ import asyncio +import inspect from pathlib import Path from unittest.mock import AsyncMock, MagicMock @@ -11,11 +12,10 @@ current_request_context, reset_request_context, ) -from nanobot.agent.tools.mcp import MCPProvider from nanobot.agent.tools.registry import ToolRegistry from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus -from nanobot.config.schema import Config, MCPServerConfig +from nanobot.config.schema import Config from nanobot.providers.base import LLMResponse, ToolCallRequest from nanobot.session.turn_continuation import INTERNAL_CONTINUATION_META @@ -75,11 +75,8 @@ def test_loop_registers_default_tools_in_injected_registry(tmp_path: Path) -> No assert registry.has("read_file") -def _config_for_loop(tmp_path: Path, *, with_mcp: bool = False) -> Config: - payload: dict = {"agents": {"defaults": {"workspace": str(tmp_path)}}} - if with_mcp: - payload["tools"] = {"mcpServers": {"demo": {"command": "fake-mcp"}}} - return Config.model_validate(payload) +def _config_for_loop(tmp_path: Path) -> Config: + return Config.model_validate({"agents": {"defaults": {"workspace": str(tmp_path)}}}) def _provider_for_loop() -> MagicMock: @@ -88,64 +85,22 @@ def _provider_for_loop() -> MagicMock: return provider -def test_loop_from_config_rejects_silently_unmanaged_mcp(tmp_path: Path) -> None: - config = _config_for_loop(tmp_path, with_mcp=True) +def test_loop_from_config_requires_caller_owned_registry(tmp_path: Path) -> None: + signature = inspect.signature(AgentLoop.from_config) - with pytest.raises(ValueError, match="application-owned MCPProvider") as exc_info: - AgentLoop.from_config(config, provider=_provider_for_loop()) + with pytest.raises(TypeError, match="tool_registry"): + signature.bind(_config_for_loop(tmp_path)) - message = str(exc_info.value) - assert "await mcp_provider.connect()" in message - assert "await mcp_provider.aclose()" in message - -def test_loop_from_config_detects_plugin_only_mcp( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config = _config_for_loop(tmp_path) - monkeypatch.setattr( - "nanobot.agent.plugins.agent_plugin_mcp_servers", - lambda _workspace, _configured: {"plugin-demo": MCPServerConfig(command="fake-mcp")}, - ) - - with pytest.raises(ValueError, match="application-owned MCPProvider"): - AgentLoop.from_config(config, provider=_provider_for_loop()) - - -def test_loop_from_config_requires_registry_for_external_mcp_owner(tmp_path: Path) -> None: - config = _config_for_loop(tmp_path, with_mcp=True) - - with pytest.raises(ValueError, match="requires a shared tool_registry"): - AgentLoop.from_config( - config, - provider=_provider_for_loop(), - mcp_runtime_owned_externally=True, - ) - - -def test_loop_from_config_accepts_explicit_external_mcp_owner(tmp_path: Path) -> None: - config = _config_for_loop(tmp_path, with_mcp=True) +def test_loop_from_config_uses_caller_owned_registry(tmp_path: Path) -> None: registry = ToolRegistry() - mcp_provider = MCPProvider.from_config(config, registry) - - loop = AgentLoop.from_config( - config, - provider=_provider_for_loop(), - tool_registry=registry, - mcp_runtime_owned_externally=True, - ) - - assert loop.tools is registry - assert mcp_provider.configured_server_names == {"demo"} - - -def test_loop_from_config_without_mcp_keeps_low_level_usage(tmp_path: Path) -> None: loop = AgentLoop.from_config( _config_for_loop(tmp_path), + tool_registry=registry, provider=_provider_for_loop(), ) + assert loop.tools is registry assert loop.tools.has("read_file") diff --git a/tests/agent/test_runtime_refresh.py b/tests/agent/test_runtime_refresh.py index d8f24e9f309..f09b1286d0b 100644 --- a/tests/agent/test_runtime_refresh.py +++ b/tests/agent/test_runtime_refresh.py @@ -7,6 +7,7 @@ import pytest from nanobot.agent.loop import AgentLoop +from nanobot.agent.tools.registry import ToolRegistry from nanobot.bus.queue import MessageBus from nanobot.bus.runtime_events import RuntimeModelChanged from nanobot.config.errors import ConfigLoadError @@ -312,7 +313,11 @@ def test_settings_context_window_refreshes_runtime_state( def loader(*, preset_name: str | None = None) -> ProviderSnapshot: return load_provider_snapshot(config_path, preset_name=preset_name) - loop = AgentLoop.from_config(config, provider_snapshot_loader=loader) + loop = AgentLoop.from_config( + config, + tool_registry=ToolRegistry(), + provider_snapshot_loader=loader, + ) payload = update_agent_settings({"context_window_tokens": ["262144"]}) loop.runtime_resolver.invalidate() diff --git a/tests/agent/test_self_model_preset.py b/tests/agent/test_self_model_preset.py index 520755635ae..ed24fcf6dda 100644 --- a/tests/agent/test_self_model_preset.py +++ b/tests/agent/test_self_model_preset.py @@ -5,6 +5,7 @@ from nanobot.agent.loop import AgentLoop from nanobot.agent.tools.context import RequestContext, request_context +from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.runtime_control import AgentRuntimeControl from nanobot.agent.tools.self import MyTool from nanobot.bus.queue import MessageBus @@ -390,7 +391,7 @@ def test_from_config_injects_default_preset(tmp_path) -> None: }) fake_provider = _provider("openai/gpt-4.1") with patch("nanobot.providers.factory.make_provider", return_value=fake_provider): - loop = AgentLoop.from_config(config) + loop = AgentLoop.from_config(config, tool_registry=ToolRegistry()) assert loop.model == "openai/gpt-4.1" assert loop.model_preset is None assert "default" in loop.model_presets @@ -407,7 +408,7 @@ def test_from_config_static_preset_loader_does_not_enable_hot_reload(tmp_path) - }) fake_provider = _provider("openai/gpt-4.1") with patch("nanobot.providers.factory.make_provider", return_value=fake_provider): - loop = AgentLoop.from_config(config) + loop = AgentLoop.from_config(config, tool_registry=ToolRegistry()) default_runtime = loop.runtime_resolver.runtime resolved = loop.runtime_resolver.resolve_preset("fast") assert resolved.model == "openai/gpt-4.1-mini"