mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-12 23:29:16 +03:00
refactor: move MCP lifecycle out of AgentLoop (#5343)
This commit is contained in:
@@ -51,6 +51,13 @@ 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. 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. `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`.
|
||||
|
||||
## Providers
|
||||
|
||||
@@ -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:
|
||||
|
||||
+18
-33
@@ -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
|
||||
|
||||
+377
-396
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
+23
-10
@@ -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(
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
|
||||
+38
-23
@@ -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())
|
||||
|
||||
+11
-2
@@ -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)
|
||||
|
||||
@@ -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.).
|
||||
|
||||
+24
-4
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock
|
||||
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 @@ class TestIdleScanThrottling:
|
||||
})
|
||||
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()
|
||||
@@ -310,7 +315,7 @@ class TestAutoCompact:
|
||||
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 +325,7 @@ class TestAutoCompact:
|
||||
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 +348,7 @@ class TestAutoCompact:
|
||||
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 +372,7 @@ class TestAutoCompact:
|
||||
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 +403,7 @@ class TestAutoCompact:
|
||||
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 +427,7 @@ class TestAutoCompact:
|
||||
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 +441,7 @@ class TestAutoCompact:
|
||||
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 +460,7 @@ class TestAutoCompact:
|
||||
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 +479,7 @@ class TestAutoCompactIdleDetection:
|
||||
|
||||
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 +508,7 @@ class TestAutoCompactIdleDetection:
|
||||
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 +522,7 @@ class TestAutoCompactIdleDetection:
|
||||
|
||||
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 +545,7 @@ class TestAutoCompactIdleDetection:
|
||||
# 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 +567,7 @@ class TestAutoCompactIdleDetection:
|
||||
|
||||
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 +586,7 @@ class TestAutoCompactIdleDetection:
|
||||
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 +602,7 @@ class TestAutoCompactIdleDetection:
|
||||
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 +633,7 @@ class TestAutoCompactSystemMessages:
|
||||
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 +661,7 @@ class TestAutoCompactEdgeCases:
|
||||
# "(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 +682,7 @@ class TestAutoCompactEdgeCases:
|
||||
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 +714,7 @@ class TestAutoCompactEdgeCases:
|
||||
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 +784,7 @@ class TestAutoCompactIntegration:
|
||||
# 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 +812,7 @@ class TestAutoCompactIntegration:
|
||||
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 +875,7 @@ class TestProactiveAutoCompact:
|
||||
|
||||
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 +902,7 @@ class TestProactiveAutoCompact:
|
||||
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 +923,7 @@ class TestProactiveAutoCompact:
|
||||
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 +937,7 @@ class TestProactiveAutoCompact:
|
||||
|
||||
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 +973,7 @@ class TestProactiveAutoCompact:
|
||||
# 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 +994,7 @@ class TestProactiveAutoCompact:
|
||||
|
||||
# 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 +1010,7 @@ class TestProactiveAutoCompact:
|
||||
|
||||
# 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 +1031,7 @@ class TestProactiveAutoCompact:
|
||||
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 +1052,7 @@ class TestProactiveAutoCompact:
|
||||
# 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 +1088,7 @@ class TestProactiveAutoCompact:
|
||||
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 +1109,7 @@ class TestProactiveAutoCompact:
|
||||
# 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 +1129,7 @@ class TestProactiveAutoCompact:
|
||||
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 +1160,7 @@ class TestProactiveAutoCompact:
|
||||
# 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 +1187,7 @@ class TestSummaryPersistence:
|
||||
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 +1223,7 @@ class TestSummaryPersistence:
|
||||
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 +1251,7 @@ class TestSummaryPersistence:
|
||||
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 +1277,7 @@ class TestSummaryPersistence:
|
||||
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 +1319,7 @@ class TestSummaryPersistence:
|
||||
_, 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 +1347,4 @@ class TestSummaryPersistence:
|
||||
# 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()
|
||||
|
||||
@@ -538,7 +538,7 @@ class TestNewCommandArchival:
|
||||
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 @@ class TestNewCommandArchival:
|
||||
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 @@ class TestNewCommandArchival:
|
||||
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 @@ class TestNewCommandArchival:
|
||||
|
||||
assert not archived.is_set()
|
||||
release_archive.set()
|
||||
await loop.close_mcp()
|
||||
await loop.aclose()
|
||||
assert archived.is_set()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -1519,13 +1519,11 @@ async def test_run_agent_loop_goal_continue_message_reads_latest_metadata(
|
||||
@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")
|
||||
|
||||
@@ -123,8 +123,7 @@ async def test_session_discard_control_cancels_active_turn(tmp_path, monkeypatch
|
||||
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,
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
@@ -11,8 +12,10 @@ from nanobot.agent.tools.context import (
|
||||
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.config.schema import Config
|
||||
from nanobot.providers.base import LLMResponse, ToolCallRequest
|
||||
from nanobot.session.turn_continuation import INTERNAL_CONTINUATION_META
|
||||
|
||||
@@ -56,6 +59,51 @@ class _Tools:
|
||||
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")
|
||||
|
||||
|
||||
def _config_for_loop(tmp_path: Path) -> Config:
|
||||
return Config.model_validate({"agents": {"defaults": {"workspace": str(tmp_path)}}})
|
||||
|
||||
|
||||
def _provider_for_loop() -> MagicMock:
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
return provider
|
||||
|
||||
|
||||
def test_loop_from_config_requires_caller_owned_registry(tmp_path: Path) -> None:
|
||||
signature = inspect.signature(AgentLoop.from_config)
|
||||
|
||||
with pytest.raises(TypeError, match="tool_registry"):
|
||||
signature.bind(_config_for_loop(tmp_path))
|
||||
|
||||
|
||||
def test_loop_from_config_uses_caller_owned_registry(tmp_path: Path) -> None:
|
||||
registry = ToolRegistry()
|
||||
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")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_loop_binds_request_context_for_tool_execution(tmp_path: Path) -> None:
|
||||
provider = MagicMock()
|
||||
|
||||
@@ -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 @@ import asyncio
|
||||
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.exceptions import McpError
|
||||
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 @@ class _FakeMcpTool(Tool):
|
||||
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 test_owned_mcp_connection_closes_from_its_owner_task():
|
||||
|
||||
@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 test_connect_mcp_retries_when_no_servers_connect(tmp_path, monkeypatch
|
||||
|
||||
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 test_agent_loop_run_closes_mcp_from_connection_owner_task(
|
||||
|
||||
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 @@ async def test_close_mcp_servers_continues_after_server_cancelled_error(tmp_path
|
||||
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 test_mcp_cleanup_re_raises_external_cancellation(tmp_path, close_all:
|
||||
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 test_reload_mcp_servers_adds_and_removes_tools_without_restart(
|
||||
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,37 +393,68 @@ async def test_request_mcp_reload_reaches_runtime_control_without_restart(
|
||||
return stacks
|
||||
|
||||
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
|
||||
loop = _make_loop(tmp_path, mcp_servers={})
|
||||
registry = ToolRegistry()
|
||||
provider = MCPProvider({}, registry, server_loader=lambda: configured)
|
||||
|
||||
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
|
||||
|
||||
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"]
|
||||
|
||||
|
||||
@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,
|
||||
@@ -419,16 +479,18 @@ async def test_reload_mcp_servers_retries_configured_server_without_live_stack(
|
||||
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 +527,16 @@ async def test_reload_mcp_servers_skips_oauth_server_waiting_for_authorization(
|
||||
"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 +544,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 +589,8 @@ async def test_mcp_tool_reconnects_after_session_terminated(
|
||||
|
||||
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 +600,8 @@ async def test_mcp_tool_reconnects_after_session_terminated(
|
||||
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 +609,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 +644,15 @@ async def test_mcp_reconnect_handler_uses_sanitized_server_prefix(
|
||||
|
||||
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 +660,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 +706,9 @@ async def test_concurrent_mcp_reconnect_reuses_fresh_session(
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@@ -15,15 +15,13 @@ import asyncio
|
||||
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 test_mcp_reconnect_during_shutdown_does_not_crash(
|
||||
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 @@ async def test_mcp_reconnect_during_shutdown_does_not_crash(
|
||||
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()
|
||||
|
||||
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock
|
||||
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()
|
||||
|
||||
@@ -5,6 +5,7 @@ import pytest
|
||||
|
||||
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"
|
||||
|
||||
@@ -56,7 +56,7 @@ class TestHandleStop:
|
||||
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 @@ class TestHandleStop:
|
||||
|
||||
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 @@ class TestHandleStop:
|
||||
|
||||
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
|
||||
|
||||
@@ -493,20 +493,6 @@ class TestModifyOpen:
|
||||
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 @@ class TestModifyOpen:
|
||||
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."""
|
||||
|
||||
@@ -64,7 +64,7 @@ def test_interactive_agent_routes_a_complete_user_turn(
|
||||
def __init__(self, bus) -> None:
|
||||
self.bus = bus
|
||||
self.stopped = asyncio.Event()
|
||||
self.close_mcp_calls = 0
|
||||
self.aclose_calls = 0
|
||||
|
||||
async def run(self) -> None:
|
||||
message = await self.bus.consume_inbound()
|
||||
@@ -97,8 +97,8 @@ def test_interactive_agent_routes_a_complete_user_turn(
|
||||
def stop(self) -> None:
|
||||
self.stopped.set()
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
self.close_mcp_calls += 1
|
||||
async def aclose(self) -> None:
|
||||
self.aclose_calls += 1
|
||||
|
||||
read_input = AsyncMock(side_effect=["hello nanobot", "exit"])
|
||||
print_response = MagicMock()
|
||||
@@ -136,7 +136,7 @@ def test_interactive_agent_routes_a_complete_user_turn(
|
||||
assert inbound.metadata == {"_wants_stream": True}
|
||||
loop = seen["loop"]
|
||||
assert isinstance(loop, _AgentLoop)
|
||||
assert loop.close_mcp_calls == 1
|
||||
assert loop.aclose_calls == 1
|
||||
assert len(renderers) == 1
|
||||
renderer = renderers[0]
|
||||
assert isinstance(renderer, _Renderer)
|
||||
|
||||
+50
-19
@@ -1543,7 +1543,7 @@ def mock_agent_runtime(tmp_path):
|
||||
agent_loop.process_direct = AsyncMock(
|
||||
return_value=OutboundMessage(channel="cli", chat_id="direct", content="mock-response"),
|
||||
)
|
||||
agent_loop.close_mcp = AsyncMock(return_value=None)
|
||||
agent_loop.aclose = AsyncMock(return_value=None)
|
||||
mock_from_config.return_value = agent_loop
|
||||
|
||||
yield {
|
||||
@@ -1621,7 +1621,7 @@ def test_agent_config_sets_active_path(monkeypatch, tmp_path: Path) -> None:
|
||||
async def process_direct(self, *_args, **_kwargs):
|
||||
return OutboundMessage(channel="cli", chat_id="direct", content="ok")
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("nanobot.cli.agent.AgentLoop", _FakeAgentLoop)
|
||||
@@ -1662,7 +1662,7 @@ def test_agent_uses_workspace_directory_for_cron_store(monkeypatch, tmp_path: Pa
|
||||
async def process_direct(self, *_args, **_kwargs):
|
||||
return OutboundMessage(channel="cli", chat_id="direct", content="ok")
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
|
||||
@@ -1712,7 +1712,7 @@ def test_agent_workspace_override_does_not_migrate_legacy_cron(
|
||||
async def process_direct(self, *_args, **_kwargs):
|
||||
return OutboundMessage(channel="cli", chat_id="direct", content="ok")
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
|
||||
@@ -1768,7 +1768,7 @@ def test_agent_custom_config_workspace_does_not_migrate_legacy_cron(
|
||||
async def process_direct(self, *_args, **_kwargs):
|
||||
return OutboundMessage(channel="cli", chat_id="direct", content="ok")
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
|
||||
@@ -2062,7 +2062,7 @@ def test_heartbeat_empty_response_still_retains_recent_messages(
|
||||
async def process_direct(self, *_args, **_kwargs):
|
||||
return SimpleNamespace(content="")
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
async def run(self) -> None:
|
||||
@@ -2738,10 +2738,7 @@ def _patch_serve_runtime(monkeypatch, config: Config, seen: dict[str, object]) -
|
||||
def __init__(self, **kwargs) -> None:
|
||||
seen["workspace"] = kwargs["workspace"]
|
||||
|
||||
async def _connect_mcp(self) -> None:
|
||||
return None
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
def _fake_create_app(
|
||||
@@ -2749,11 +2746,13 @@ def _patch_serve_runtime(monkeypatch, config: Config, seen: dict[str, object]) -
|
||||
model_name: str,
|
||||
request_timeout: float,
|
||||
api_key: str = "",
|
||||
prepare_agent=None,
|
||||
):
|
||||
seen["agent_loop"] = agent_loop
|
||||
seen["model_name"] = model_name
|
||||
seen["request_timeout"] = request_timeout
|
||||
seen["api_key"] = api_key
|
||||
seen["prepare_agent"] = prepare_agent
|
||||
return _FakeApiApp()
|
||||
|
||||
def _fake_run_app(api_app, host: str, port: int, print):
|
||||
@@ -2914,7 +2913,7 @@ def test_gateway_unbound_agent_cron_is_skipped(
|
||||
async def submit_cron_turn(self, _msg: InboundMessage):
|
||||
raise AssertionError("unbound cron job must not run as a bound cron turn")
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
async def run(self) -> None:
|
||||
@@ -3033,7 +3032,7 @@ def test_gateway_bound_cron_runs_as_session_turn(
|
||||
content="Checked the repo.",
|
||||
)
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
async def run(self) -> None:
|
||||
@@ -3253,7 +3252,7 @@ def test_gateway_local_trigger_queue_submits_agent_turns(
|
||||
self.runtime_resolver.invalidate.assert_called_once_with()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
def stop(self) -> None:
|
||||
@@ -3499,7 +3498,7 @@ def test_gateway_health_endpoint_binds_and_serves_expected_responses(
|
||||
async def run(self) -> None:
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
def stop(self) -> None:
|
||||
@@ -3668,7 +3667,7 @@ def test_gateway_health_endpoint_binds_and_serves_expected_responses(
|
||||
assert timed_out_writer.output == b""
|
||||
|
||||
|
||||
def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
|
||||
def test_gateway_agent_task_owns_initial_mcp_provider_close(
|
||||
monkeypatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
@@ -3696,17 +3695,41 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
|
||||
return None
|
||||
|
||||
async def run(self) -> None:
|
||||
seen["agent_task"] = asyncio.current_task()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
seen["agent_task_cleaned_up"] = True
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
raise AssertionError("gateway must not close MCP from the outer task")
|
||||
async def aclose(self) -> None:
|
||||
seen["agent_closed"] = True
|
||||
|
||||
def stop(self) -> None:
|
||||
seen["agent_stopped"] = True
|
||||
|
||||
class _FakeMCPProvider:
|
||||
def __init__(self) -> None:
|
||||
self.connect_task: asyncio.Task | None = None
|
||||
self.close_tasks: list[asyncio.Task | None] = []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, _config, _registry):
|
||||
provider = cls()
|
||||
seen["mcp_provider"] = provider
|
||||
return provider
|
||||
|
||||
async def connect(self) -> None:
|
||||
self.connect_task = asyncio.current_task()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close_tasks.append(asyncio.current_task())
|
||||
|
||||
def runtime_status(self) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
async def reload(self) -> dict[str, object]:
|
||||
return {"ok": True}
|
||||
|
||||
class _FakeChannelManager:
|
||||
def __init__(self, _config, _bus, **_kwargs) -> None:
|
||||
self.enabled_channels = ["telegram"]
|
||||
@@ -3753,6 +3776,7 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
|
||||
session_manager=lambda _workspace: object(),
|
||||
)
|
||||
monkeypatch.setattr("nanobot.cli.gateway_runtime.AgentLoop", _FakeAgentLoop)
|
||||
monkeypatch.setattr("nanobot.cli.gateway_runtime.MCPProvider", _FakeMCPProvider)
|
||||
monkeypatch.setattr("nanobot.channels.manager.ChannelManager", _FakeChannelManager)
|
||||
monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCronService)
|
||||
monkeypatch.setattr("asyncio.start_server", _fake_start_server)
|
||||
@@ -3761,9 +3785,15 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert seen["agent_stopped"] is True
|
||||
assert seen["agent_closed"] is True
|
||||
assert seen["agent_task_cleaned_up"] is True
|
||||
assert seen["channels_stopped"] is True
|
||||
assert seen["cron_stopped"] is True
|
||||
mcp_provider = seen["mcp_provider"]
|
||||
assert isinstance(mcp_provider, _FakeMCPProvider)
|
||||
assert mcp_provider.connect_task is seen["agent_task"]
|
||||
assert mcp_provider.close_tasks[0] is mcp_provider.connect_task
|
||||
assert len(mcp_provider.close_tasks) == 2
|
||||
|
||||
|
||||
def test_gateway_shutdown_event_exits_forever_runtime_tasks(
|
||||
@@ -3800,8 +3830,8 @@ def test_gateway_shutdown_event_exits_forever_runtime_tasks(
|
||||
finally:
|
||||
seen["agent_task_cleaned_up"] = True
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
raise AssertionError("gateway must not close MCP from the outer task")
|
||||
async def aclose(self) -> None:
|
||||
seen["agent_closed"] = True
|
||||
|
||||
def stop(self) -> None:
|
||||
seen["agent_stopped"] = True
|
||||
@@ -3881,6 +3911,7 @@ def test_gateway_shutdown_event_exits_forever_runtime_tasks(
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert seen["agent_stopped"] is True
|
||||
assert seen["agent_closed"] is True
|
||||
assert seen["agent_task_cleaned_up"] is True
|
||||
assert seen["channel_task_cleaned_up"] is True
|
||||
assert seen["channels_stopped"] is True
|
||||
|
||||
@@ -22,7 +22,7 @@ class _FakeAgent:
|
||||
self.raise_on_close = False
|
||||
self.background: asyncio.Task[None] | None = None
|
||||
|
||||
async def close_mcp(self) -> None:
|
||||
async def aclose(self) -> None:
|
||||
self.close_calls += 1
|
||||
if self.hang_on_close:
|
||||
await asyncio.sleep(3600)
|
||||
@@ -30,7 +30,7 @@ class _FakeAgent:
|
||||
raise RuntimeError("cleanup exploded")
|
||||
if self.background is not None:
|
||||
await self.background
|
||||
self.events.append("close_mcp")
|
||||
self.events.append("aclose")
|
||||
|
||||
|
||||
class _FakeChannels:
|
||||
@@ -43,6 +43,16 @@ class _FakeChannels:
|
||||
self.events.append("channels_stopped")
|
||||
|
||||
|
||||
class _FakeMCPProvider:
|
||||
def __init__(self, events: list[str] | None = None) -> None:
|
||||
self.close_calls = 0
|
||||
self.events = events if events is not None else []
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close_calls += 1
|
||||
self.events.append("mcp_closed")
|
||||
|
||||
|
||||
async def _cancellable_task(events: list[str]) -> None:
|
||||
try:
|
||||
await asyncio.sleep(3600)
|
||||
@@ -64,13 +74,14 @@ async def _stubborn_task(events: list[str]) -> None:
|
||||
async def test_runtime_tasks_cancelled_before_resources_closed() -> None:
|
||||
events: list[str] = []
|
||||
agent = _FakeAgent(events)
|
||||
provider = _FakeMCPProvider(events)
|
||||
channels = _FakeChannels()
|
||||
task = asyncio.create_task(_cancellable_task(events))
|
||||
await asyncio.sleep(0) # let the task start (cancellation pre-start skips its body)
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [task], None)
|
||||
await _close_gateway_runtime(agent, provider, channels, [task], None)
|
||||
|
||||
assert events == ["cancelled", "close_mcp"] # cancel happens before close
|
||||
assert events == ["cancelled", "aclose", "mcp_closed"]
|
||||
assert channels.stopped == 1
|
||||
assert agent.close_calls == 1
|
||||
assert task.cancelled()
|
||||
@@ -78,6 +89,7 @@ async def test_runtime_tasks_cancelled_before_resources_closed() -> None:
|
||||
|
||||
async def test_pending_background_work_is_drained_before_close_returns() -> None:
|
||||
agent = _FakeAgent()
|
||||
provider = _FakeMCPProvider()
|
||||
channels = _FakeChannels()
|
||||
done: dict[str, bool] = {"done": False}
|
||||
|
||||
@@ -87,7 +99,7 @@ async def test_pending_background_work_is_drained_before_close_returns() -> None
|
||||
|
||||
agent.background = asyncio.create_task(background_work())
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [], None)
|
||||
await _close_gateway_runtime(agent, provider, channels, [], None)
|
||||
|
||||
assert done["done"] is True
|
||||
assert agent.close_calls == 1
|
||||
@@ -95,6 +107,7 @@ async def test_pending_background_work_is_drained_before_close_returns() -> None
|
||||
|
||||
async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
|
||||
agent = _FakeAgent()
|
||||
provider = _FakeMCPProvider()
|
||||
channels = _FakeChannels()
|
||||
events: list[str] = []
|
||||
task = asyncio.create_task(_stubborn_task(events))
|
||||
@@ -104,6 +117,7 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
|
||||
start = time.monotonic()
|
||||
await _close_gateway_runtime(
|
||||
agent,
|
||||
provider,
|
||||
channels,
|
||||
[task],
|
||||
runtime_tasks,
|
||||
@@ -117,70 +131,88 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
|
||||
assert task.done() # the timed-out task received a second cancellation
|
||||
assert runtime_tasks.done()
|
||||
assert agent.close_calls == 1 # resources still closed underneath it
|
||||
assert provider.close_calls == 1
|
||||
assert elapsed < 1.0 # bounded, not held open by the stubborn task
|
||||
|
||||
|
||||
async def test_hanging_close_is_bounded_and_does_not_raise() -> None:
|
||||
agent = _FakeAgent()
|
||||
provider = _FakeMCPProvider()
|
||||
agent.hang_on_close = True
|
||||
channels = _FakeChannels()
|
||||
|
||||
start = time.monotonic()
|
||||
await _close_gateway_runtime(agent, channels, [], None, close_timeout=0.05)
|
||||
await _close_gateway_runtime(
|
||||
agent,
|
||||
provider,
|
||||
channels,
|
||||
[],
|
||||
None,
|
||||
close_timeout=0.05,
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
assert agent.close_calls == 1
|
||||
assert provider.close_calls == 1
|
||||
assert channels.stopped == 1
|
||||
assert elapsed < 1.0
|
||||
|
||||
|
||||
async def test_failing_close_is_logged_but_shutdown_proceeds() -> None:
|
||||
agent = _FakeAgent()
|
||||
provider = _FakeMCPProvider()
|
||||
agent.raise_on_close = True
|
||||
channels = _FakeChannels()
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [], None)
|
||||
await _close_gateway_runtime(agent, provider, channels, [], None)
|
||||
|
||||
assert agent.close_calls == 1
|
||||
assert provider.close_calls == 1
|
||||
assert channels.stopped == 1 # teardown continued past the failure
|
||||
|
||||
|
||||
async def test_duplicate_cleanup_is_idempotent() -> None:
|
||||
agent = _FakeAgent()
|
||||
provider = _FakeMCPProvider()
|
||||
channels = _FakeChannels()
|
||||
task = asyncio.create_task(_cancellable_task([]))
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [task], None)
|
||||
await _close_gateway_runtime(agent, channels, [task], None)
|
||||
await _close_gateway_runtime(agent, provider, channels, [task], None)
|
||||
await _close_gateway_runtime(agent, provider, channels, [task], None)
|
||||
|
||||
assert agent.close_calls == 2 # second pass is a clean no-op
|
||||
assert provider.close_calls == 2
|
||||
assert channels.stopped == 2
|
||||
assert task.cancelled()
|
||||
|
||||
|
||||
async def test_finished_runtime_tasks_gather_is_retrieved() -> None:
|
||||
agent = _FakeAgent()
|
||||
provider = _FakeMCPProvider()
|
||||
channels = _FakeChannels()
|
||||
finished = asyncio.get_running_loop().create_future()
|
||||
finished.set_result(None)
|
||||
runtime_tasks = asyncio.gather(finished)
|
||||
await asyncio.sleep(0) # let the gather observe the finished child
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [], runtime_tasks)
|
||||
await _close_gateway_runtime(agent, provider, channels, [], runtime_tasks)
|
||||
|
||||
assert runtime_tasks.done()
|
||||
assert agent.close_calls == 1
|
||||
assert provider.close_calls == 1
|
||||
|
||||
|
||||
async def test_cancelled_runtime_tasks_gather_does_not_raise() -> None:
|
||||
agent = _FakeAgent()
|
||||
provider = _FakeMCPProvider()
|
||||
channels = _FakeChannels()
|
||||
runtime_tasks = asyncio.gather(asyncio.sleep(3600))
|
||||
runtime_tasks.cancel()
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [], runtime_tasks)
|
||||
await _close_gateway_runtime(agent, provider, channels, [], runtime_tasks)
|
||||
with suppress(asyncio.CancelledError):
|
||||
await runtime_tasks # settle the cancelled gather without raising
|
||||
|
||||
assert runtime_tasks.done() # the cancelled gather was awaited without raising
|
||||
assert agent.close_calls == 1
|
||||
assert provider.close_calls == 1
|
||||
|
||||
@@ -32,8 +32,7 @@ AUTH_HEADERS = {"Authorization": f"Bearer {API_KEY}"}
|
||||
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
|
||||
|
||||
|
||||
@@ -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 test_stream_passes_on_stream_callbacks(aiohttp_client) -> None:
|
||||
|
||||
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 test_stream_segment_end_does_not_close_sse(aiohttp_client) -> None:
|
||||
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 test_stream_uses_final_response_when_no_deltas(aiohttp_client) -> None
|
||||
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 test_stream_with_session_id(aiohttp_client) -> None:
|
||||
|
||||
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 test_streaming_backend_failure_does_not_emit_success_terminator(aiohtt
|
||||
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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -1637,37 +1650,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()
|
||||
|
||||
@@ -34,8 +34,7 @@ AUTH_HEADERS = {"Authorization": f"Bearer {API_KEY}"}
|
||||
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 test_followup_requests_share_same_session_key(aiohttp_client) -> None:
|
||||
|
||||
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 test_fixed_session_requests_are_serialized(aiohttp_client) -> None:
|
||||
|
||||
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 test_empty_response_falls_back_without_retry(aiohttp_client) -> None:
|
||||
|
||||
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()
|
||||
|
||||
|
||||
@@ -12,7 +12,6 @@ from unittest.mock import AsyncMock
|
||||
|
||||
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 @@ def test_exec_session_manager_preserves_single_cleanup_error():
|
||||
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 @@ def test_agent_loop_shutdown_closes_exec_sessions(tmp_path, monkeypatch):
|
||||
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 @@ def test_agent_loop_shutdown_closes_exec_sessions(tmp_path, monkeypatch):
|
||||
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 @@ def test_agent_loop_shutdown_attempts_all_cleanup_after_errors(monkeypatch):
|
||||
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 @@ def test_terminate_by_owner_skips_sessions_without_owner_key(tmp_path):
|
||||
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 @@ def test_agent_loop_shutdown_preserves_single_cleanup_error(monkeypatch):
|
||||
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())
|
||||
|
||||
@@ -13,6 +13,7 @@ import pytest
|
||||
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,21 @@ async def test_connect_missing_servers_propagates_external_cancellation(monkeypa
|
||||
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()
|
||||
assert provider.runtime_status() == {"test": "failed"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -223,25 +221,20 @@ async def test_saved_oauth_http_403_projects_failed_runtime_without_details(
|
||||
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)
|
||||
@@ -1263,6 +1256,59 @@ async def test_connect_mcp_servers_propagates_external_cancellation(
|
||||
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],
|
||||
@@ -1900,7 +1946,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
|
||||
|
||||
|
||||
@@ -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 @@ async def test_mcp_list_serializes_local_runtime_failure_snapshot(tmp_path) -> N
|
||||
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(
|
||||
|
||||
Reference in New Issue
Block a user