Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
chengyongru
2026-08-12 17:51:04 +08:00
committed by GitHub
parent 686dd0603e
commit 19997d20bb
39 changed files with 1192 additions and 846 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/context.py
+1 -19
View File
@@ -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:
nanobot/agent/loop.py
+18 -33
View File
@@ -95,11 +95,9 @@ from nanobot.utils.runtime import (
)
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 @@ class AgentLoop:
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 @@ class AgentLoop:
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 @@ class AgentLoop:
)
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()
)
@@ -464,10 +458,15 @@ class AgentLoop:
cls,
config: Config,
bus: MessageBus | None = None,
*,
tool_registry: ToolRegistry,
**extra: Any,
) -> AgentLoop:
"""Create an AgentLoop from config with the common parameter set.
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``).
@@ -486,8 +485,6 @@ class AgentLoop:
config,
provider_snapshot_loader,
)
from nanobot.agent.plugins import agent_plugin_mcp_servers
return cls(
bus=bus,
provider=provider,
@@ -502,7 +499,6 @@ class AgentLoop:
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,
@@ -517,6 +513,7 @@ class AgentLoop:
restart_mode=config.gateway.restart_mode,
provider_snapshot_loader=provider_snapshot_loader,
preset_snapshot_loader=preset_snapshot_loader,
tool_registry=tool_registry,
**extra,
)
@@ -643,14 +640,6 @@ class AgentLoop:
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 +1151,6 @@ class AgentLoop:
"""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 +1241,7 @@ class AgentLoop:
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 +1359,24 @@ class AgentLoop:
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 +1399,6 @@ class AgentLoop:
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 +2287,6 @@ class AgentLoop:
"""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
nanobot/agent/tools/mcp.py
+377 -396
View File
@@ -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 os
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 @@ if TYPE_CHECKING:
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,18 +54,37 @@ _WINDOWS_SHELL_LAUNCHERS: frozenset[str] = frozenset(("npx", "npm", "pnpm", "yar
# 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):
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."""
@@ -492,11 +505,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 +599,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 +761,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 +865,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 +998,10 @@ class MCPPromptWrapper(_MCPWrapperBase):
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 +1015,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 +1257,7 @@ async def connect_mcp_servers(
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()
@@ -1282,15 +1295,29 @@ async def connect_mcp_servers(
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
@@ -1301,369 +1328,357 @@ 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 _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 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 _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()))
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
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
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")
@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)
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
@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
}
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:
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")
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
}
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 _close_mcp_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:
self._set_runtime_status(missing_servers, "failed")
if task_is_cancelling():
raise
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,
)
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": "Could not reload MCP config. Restart nanobot to pick up changes.",
"requires_restart": True,
"error": str(exc),
}
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
authorization_pending = {
name
for name, cfg in oauth_servers.items()
if not mcp_oauth_has_credentials(name, cfg.url)
for name, cfg in next_servers.items()
if cfg.auth == "oauth" and 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
}
if state._mcp_connecting or not missing_servers:
return
state._mcp_connecting = True
_set_runtime_status(state, missing_servers, "connecting")
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)")
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
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)
for name in [*removed, *authorization_pending]:
self._runtime_statuses.pop(name, None)
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")
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()
self._connections.update(connected)
self._record_connection_result(to_connect, connected)
self._attach_reconnect_handlers(connected)
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,
}
try:
from nanobot.agent.plugins import agent_plugin_mcp_servers
from nanobot.config.loader import load_config, resolve_config_env_vars
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."
config = resolve_config_env_vars(load_config())
next_servers = agent_plugin_mcp_servers(
config.workspace_path,
config.tools.mcp_servers,
logger.info(
"MCP hot reload: added={} changed={} removed={} retried={} "
"connected={} failed={} tools_removed={}",
added,
changed,
removed,
retry_missing,
sorted(connected),
failed,
tools_removed,
)
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),
"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,
}
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()
return {
"ok": False,
"message": "MCP connections are shutting down.",
"requires_restart": True,
}
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,
}
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:
@staticmethod
def _closing_result() -> dict[str, Any]:
return {
"ok": False,
"message": "MCP hot reload timed out. Restart nanobot to pick up changes.",
"message": "MCP connections are shutting down.",
"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,
}
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,
)
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 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)
ack = metadata.get(RUNTIME_CONTROL_ACK)
try:
result = await reload_servers(state, registry)
except Exception as exc:
logger.exception("MCP hot reload failed")
result = {
"ok": False,
"message": "MCP hot reload failed. Restart nanobot to pick up changes.",
"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
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
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(
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(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(
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
self._set_runtime_status({server_name}, "connecting")
connected = await connect_mcp_servers(
{server_name: cfg},
self._registry,
)
if self._closing:
await _close_mcp_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)
logger.warning("MCP server '{}' session terminated; refreshing connection", server_name)
_unregister_server_tools(registry, server_name)
await _close_server(state, server_name)
async def _close_server(self, server_name: str) -> None:
connection = self._connections.pop(server_name, None)
if connection is None:
return
await _close_mcp_connection(server_name, connection)
_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)
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 _close_mcp_connections(connections)
def _server_signature(cfg: Any) -> Any:
@@ -1690,37 +1705,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)
nanobot/agent/tools/self.py
+1 -1
View File
@@ -78,7 +78,7 @@ class MyTool(Tool):
"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",
nanobot/api/server.py
+23 -10
View File
@@ -48,6 +48,7 @@ _AGENT_LOOP_KEY = web.AppKey[Any]("agent_loop")
_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 handle_chat_completions(request: web.Request) -> web.Response | web.St
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 handle_chat_completions(request: web.Request) -> web.Response | web.St
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 handle_chat_completions(request: web.Request) -> web.Response | web.St
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(
nanobot/bus/events.py
-1
View File
@@ -16,7 +16,6 @@ OUTBOUND_META_AGENT_UI = "_agent_ui"
# 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"
nanobot/channels/manager.py
+4 -1
View File
@@ -5,7 +5,7 @@ from __future__ import annotations
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 @@ class ChannelManager:
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 @@ class ChannelManager:
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 @@ class ChannelManager:
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,
)
nanobot/channels/websocket/tests/test_websocket_http_routes.py
+10 -5
View File
@@ -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 @@ async def test_mcp_presets_routes_require_token_and_return_payload(
_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")
nanobot/cli/agent.py
+38 -23
View File
@@ -13,6 +13,8 @@ from rich.console import Console
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 @@ def agent(
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 agent(
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 @@ def agent(
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())
nanobot/cli/commands.py
+11 -2
View File
@@ -49,6 +49,8 @@ from nanobot import __logo__, __version__ # noqa: E402
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)
nanobot/cli/gateway_runtime.py
+39 -17
View File
@@ -14,6 +14,8 @@ from rich.console import Console
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 _run_gateway(
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 _run_gateway(
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 @@ def _run_gateway(
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 @@ def _run_gateway(
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 _run_gateway(
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 @@ def _run_gateway(
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 @@ def _run_gateway(
),
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 @@ def _run_gateway(
agent.stop()
# Cancel runtime tasks first, then deterministically close
# exec/MCP resources while the event loop is still alive.
await _close_gateway_runtime(agent, channels, tasks, runtime_tasks)
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.).
nanobot/nanobot.py
+24 -4
View File
@@ -10,6 +10,8 @@ from typing import Any
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 @@ class Nanobot:
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 @@ class Nanobot:
)
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 @@ class Nanobot:
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 @@ class Nanobot:
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
nanobot/webui/gateway_services.py
+3 -1
View File
@@ -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,
)
nanobot/webui/settings_routes.py
+33 -4
View File
@@ -5,14 +5,13 @@ from __future__ import annotations
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_PAYLOAD_ATTR = "_nanobot_webui_mutation_payload"
_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 @@ class WebUISettingsRouter:
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 @@ class WebUISettingsRouter:
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 @@ class WebUISettingsRouter:
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 @@ class WebUISettingsRouter:
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 @@ class WebUISettingsRouter:
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:
nanobot/webui/ws_http.py
+3 -1
View File
@@ -14,7 +14,7 @@ import json
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 @@ class GatewayHTTPHandler:
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 @@ class GatewayHTTPHandler:
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,
)