mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-07 01:48:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore(mcp): migrate client integration to SDK v2

Browse Source
This commit is contained in:
chengyongru 2026-07-30 17:07:24 +08:00
parent 606ac56e8f
commit c2349a0bfc
8 changed files with 166 additions and 85 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
nanobot/agent/tools/mcp.py
View File

@ -12,7 +12,7 @@ from contextlib import AsyncExitStack, suppress
from typing import TYPE_CHECKING, Any, Mapping, Protocol, cast from typing import TYPE_CHECKING, Any, Mapping, Protocol, cast
from weakref import WeakKeyDictionary from weakref import WeakKeyDictionary
import httpx import httpx2 as httpx
from loguru import logger from loguru import logger
from nanobot.agent.tools.base import Tool, ToolResult from nanobot.agent.tools.base import Tool, ToolResult
@ -25,9 +25,9 @@ from nanobot.bus.events import (
) )
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.security.network import ( from nanobot.security.network import (
PinnedDNSAsyncTransport, Httpx2PinnedDNSAsyncTransport,
env_proxy_applies_to_url, env_proxy_applies_to_url,
httpx_env_proxy_mounts, httpx2_env_proxy_mounts,
resolve_url_target, resolve_url_target,
validate_url_target, validate_url_target,
) )
@ -194,7 +194,7 @@ def _is_session_terminated(exc: BaseException) -> bool:
messages.append(str(getattr(error, "message", ""))) messages.append(str(getattr(error, "message", "")))
return any( return any(
marker in message.lower() marker in message.lower()
for marker in ("session terminated", "connection closed") for marker in ("session terminated", "session not found", "connection closed")
for message in messages for message in messages
) )
@ -252,8 +252,8 @@ def _redact_url(url: str) -> str:
def _pinned_transport_kwargs() -> dict[str, Any]: def _pinned_transport_kwargs() -> dict[str, Any]:
kwargs: dict[str, Any] = {"transport": PinnedDNSAsyncTransport()} kwargs: dict[str, Any] = {"transport": Httpx2PinnedDNSAsyncTransport()}
mounts = httpx_env_proxy_mounts() mounts = httpx2_env_proxy_mounts()
if mounts: if mounts:
kwargs["mounts"] = mounts kwargs["mounts"] = mounts
return kwargs return kwargs
@ -518,7 +518,7 @@ def _image_block_data_url(block: Any, types: Any) -> str | None:
""" """
image_cls = getattr(types, "ImageContent", None) image_cls = getattr(types, "ImageContent", None)
if image_cls is not None and isinstance(block, image_cls): if image_cls is not None and isinstance(block, image_cls):
mime = getattr(block, "mimeType", None) or "image/png" mime = getattr(block, "mime_type", None) or "image/png"
return f"data:{mime};base64,{block.data}" return f"data:{mime};base64,{block.data}"
embedded_cls = getattr(types, "EmbeddedResource", None) embedded_cls = getattr(types, "EmbeddedResource", None)
@ -527,7 +527,7 @@ def _image_block_data_url(block: Any, types: Any) -> str | None:
resource = getattr(block, "resource", None) resource = getattr(block, "resource", None)
if blob_cls is not None and isinstance(resource, blob_cls): if blob_cls is not None and isinstance(resource, blob_cls):
blob_resource = cast(Any, resource) blob_resource = cast(Any, resource)
mime = getattr(blob_resource, "mimeType", None) or "" mime = getattr(blob_resource, "mime_type", None) or ""
if isinstance(mime, str) and mime.startswith("image/"): if isinstance(mime, str) and mime.startswith("image/"):
return f"data:{mime};base64,{blob_resource.blob}" return f"data:{mime};base64,{blob_resource.blob}"
return None return None
@ -571,7 +571,7 @@ class MCPToolWrapper(_MCPWrapperBase):
self._original_name = tool_def.name self._original_name = tool_def.name
self._name = _sanitize_mcp_tool_name(f"mcp_{server_name}_{tool_def.name}") self._name = _sanitize_mcp_tool_name(f"mcp_{server_name}_{tool_def.name}")
self._description = tool_def.description or tool_def.name self._description = tool_def.description or tool_def.name
raw_schema = tool_def.inputSchema or {"type": "object", "properties": {}} raw_schema = tool_def.input_schema or {"type": "object", "properties": {}}
self._parameters = _normalize_schema_for_openai(raw_schema) self._parameters = _normalize_schema_for_openai(raw_schema)
self._tool_timeout = tool_timeout self._tool_timeout = tool_timeout
@ -650,7 +650,7 @@ class MCPToolWrapper(_MCPWrapperBase):
# Success — extract text and persist any image content as artifacts. # Success — extract text and persist any image content as artifacts.
try: try:
rendered = self._render_call_result(result.content, kwargs) rendered = self._render_call_result(result.content, kwargs)
if getattr(result, "isError", False): if getattr(result, "is_error", False):
return ToolResult.error(rendered) return ToolResult.error(rendered)
return rendered return rendered
except Exception as exc: except Exception as exc:
@ -876,8 +876,7 @@ class MCPPromptWrapper(_MCPWrapperBase):
return True return True
async def execute(self, **kwargs: Any) -> str: async def execute(self, **kwargs: Any) -> str:
from mcp import types from mcp import MCPError, types
from mcp.shared.exceptions import McpError
retried_transient = False retried_transient = False
refreshed_session = False refreshed_session = False
@ -897,7 +896,7 @@ class MCPPromptWrapper(_MCPWrapperBase):
raise raise
logger.warning("MCP prompt '{}' was cancelled by server/SDK", self._name) logger.warning("MCP prompt '{}' was cancelled by server/SDK", self._name)
return "(MCP prompt call was cancelled)" return "(MCP prompt call was cancelled)"
except McpError as exc: except MCPError as exc:
if await self._refresh_session_after_termination( if await self._refresh_session_after_termination(
exc, exc,
refreshed_session, refreshed_session,
@ -1062,7 +1061,7 @@ async def connect_mcp_servers(
**_pinned_transport_kwargs(), **_pinned_transport_kwargs(),
) )
) )
read, write, _ = await server_stack.enter_async_context( read, write = await server_stack.enter_async_context(
streamable_http_client(cfg.url, http_client=http_client) streamable_http_client(cfg.url, http_client=http_client)
) )
else: else:

59
nanobot/security/network.py
View File

@ -12,6 +12,7 @@ from urllib.parse import urlparse
from urllib.request import getproxies, proxy_bypass from urllib.request import getproxies, proxy_bypass
import httpx import httpx
import httpx2
_BLOCKED_NETWORKS = [ _BLOCKED_NETWORKS = [
ipaddress.ip_network("0.0.0.0/8"), ipaddress.ip_network("0.0.0.0/8"),
@ -29,6 +30,7 @@ _BLOCKED_NETWORKS = [
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE) _URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = [] _allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
_DNS_PIN_RESOLVER_LOCK = asyncio.Lock()
def is_loopback_host(host: str) -> bool: def is_loopback_host(host: str) -> bool:
@ -195,6 +197,30 @@ def httpx_env_proxy_mounts() -> dict[str, httpx.AsyncBaseTransport | None]:
return mounts return mounts
def httpx2_env_proxy_mounts() -> dict[str, httpx2.AsyncBaseTransport | None]:
"""Build HTTPX2 proxy mounts while leaving direct routes to the base transport."""
proxies = getproxies()
mounts: dict[str, httpx2.AsyncBaseTransport | None] = {}
for scheme in ("http", "https", "all"):
proxy_url = proxies.get(scheme)
if proxy_url:
if "://" not in proxy_url:
proxy_url = f"http://{proxy_url}"
mounts[f"{scheme}://"] = httpx2.AsyncHTTPTransport(proxy=httpx2.Proxy(proxy_url))
if not mounts:
return {}
no_proxy = proxies.get("no", "")
if no_proxy == "*":
return {}
for entry in no_proxy.split(","):
pattern = _no_proxy_mount_pattern(entry.strip())
if pattern:
mounts[pattern] = None
return mounts
def _no_proxy_mount_pattern(hostname: str) -> str | None: def _no_proxy_mount_pattern(hostname: str) -> str | None:
if not hostname: if not hostname:
return None return None
@ -264,7 +290,7 @@ class UnsafeURLRequestError(httpx.RequestError):
class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport): class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
"""HTTPX transport that pins each request to the IPs validated for its URL.""" """HTTPX transport that pins each request to the IPs validated for its URL."""
_resolver_lock = asyncio.Lock() _resolver_lock = _DNS_PIN_RESOLVER_LOCK
def __init__( def __init__(
self, self,
@ -288,6 +314,37 @@ class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
await self._inner.aclose() await self._inner.aclose()
class Httpx2UnsafeURLRequestError(httpx2.RequestError):
"""Raised when an HTTPX2 request is rejected by URL safety validation."""
class Httpx2PinnedDNSAsyncTransport(httpx2.AsyncBaseTransport):
"""HTTPX2 transport that pins each request to the IPs validated for its URL."""
_resolver_lock = _DNS_PIN_RESOLVER_LOCK
def __init__(
self,
*,
allow_loopback: bool = False,
inner: httpx2.AsyncBaseTransport | None = None,
) -> None:
self._allow_loopback = allow_loopback
self._inner = inner or httpx2.AsyncHTTPTransport()
async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
url = str(request.url)
ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback)
if not ok:
raise Httpx2UnsafeURLRequestError(error, request=request)
async with self._resolver_lock:
with pin_resolved_url_dns(url, resolved_ips):
return await self._inner.handle_async_request(request)
async def aclose(self) -> None:
await self._inner.aclose()
def validate_resolved_url(url: str) -> tuple[bool, str]: def validate_resolved_url(url: str) -> tuple[bool, str]:
"""Validate an already-fetched URL (e.g. after redirect). Only checks the IP, skips DNS.""" """Validate an already-fetched URL (e.g. after redirect). Only checks the IP, skips DNS."""
try: try:

4
pyproject.toml
View File

@ -31,6 +31,8 @@ dependencies = [
"websockets>=15.0,<17.0", "websockets>=15.0,<17.0",
"websocket-client>=1.9.0,<2.0.0", "websocket-client>=1.9.0,<2.0.0",
"httpx>=0.28.0,<1.0.0", "httpx>=0.28.0,<1.0.0",
# MCP v2 uses the independently versioned httpx2 package for HTTP transports.
"httpx2>=2.5.0,<3.0.0",
"ddgs>=9.5.5,<10.0.0", "ddgs>=9.5.5,<10.0.0",
"oauth-cli-kit>=0.1.6,<1.0.0", "oauth-cli-kit>=0.1.6,<1.0.0",
"loguru>=0.7.3,<1.0.0", "loguru>=0.7.3,<1.0.0",
@ -40,7 +42,7 @@ dependencies = [
"croniter>=6.0.0,<7.0.0", "croniter>=6.0.0,<7.0.0",
"prompt-toolkit>=3.0.50,<4.0.0", "prompt-toolkit>=3.0.50,<4.0.0",
"questionary>=2.0.0,<3.0.0", "questionary>=2.0.0,<3.0.0",
"mcp>=1.26.0,<2.0.0", "mcp>=2.0.0,<3.0.0",
"json-repair>=0.57.0,<1.0.0", "json-repair>=0.57.0,<1.0.0",
"chardet>=3.0.2,<6.0.0", "chardet>=3.0.2,<6.0.0",
"openai>=2.8.0", "openai>=2.8.0",

17
tests/agent/test_mcp_connection.py
View File

@ -10,10 +10,9 @@ from unittest.mock import MagicMock
import anyio import anyio
import pytest import pytest
from mcp import MCPError
from mcp import types as mcp_types from mcp import types as mcp_types
from mcp.shared.exceptions import McpError
from mcp.shared.message import SessionMessage from mcp.shared.message import SessionMessage
from mcp.types import ErrorData
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools import mcp as mcp_runtime from nanobot.agent.tools import mcp as mcp_runtime
@ -26,14 +25,12 @@ from nanobot.config.schema import MCPServerConfig
def _mcp_notification(method: str, params: dict[str, Any] | None = None) -> SessionMessage: def _mcp_notification(method: str, params: dict[str, Any] | None = None) -> SessionMessage:
return SessionMessage( return SessionMessage(
message=mcp_types.JSONRPCMessage( message=mcp_types.JSONRPCNotification(
mcp_types.JSONRPCNotification(
jsonrpc="2.0", jsonrpc="2.0",
method=method, method=method,
params=params, params=params,
) )
) )
)
def test_mcp_progress_detection_accepts_flattened_sdk_message_shape(): def test_mcp_progress_detection_accepts_flattened_sdk_message_shape():
@ -428,7 +425,7 @@ async def test_mcp_tool_reconnects_after_session_terminated(
self.call_count += 1 self.call_count += 1
assert arguments == {"symbol": "AAPL"} assert arguments == {"symbol": "AAPL"}
if self.index == 1: if self.index == 1:
raise McpError(ErrorData(code=-32000, message="Session terminated")) raise MCPError(-32000, "Session terminated")
return SimpleNamespace( return SimpleNamespace(
content=[mcp_types.TextContent(type="text", text="recovered")] content=[mcp_types.TextContent(type="text", text="recovered")]
) )
@ -443,7 +440,7 @@ async def test_mcp_tool_reconnects_after_session_terminated(
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="quote", name="quote",
description="quote tool", description="quote tool",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
registry.register(MCPToolWrapper(session, name, tool_def, tool_timeout=5)) registry.register(MCPToolWrapper(session, name, tool_def, tool_timeout=5))
stack = AsyncExitStack() stack = AsyncExitStack()
@ -484,7 +481,7 @@ async def test_mcp_reconnect_handler_uses_sanitized_server_prefix(
async def call_tool(self, _name: str, arguments: dict[str, Any]) -> Any: async def call_tool(self, _name: str, arguments: dict[str, Any]) -> Any:
assert arguments == {} assert arguments == {}
if self.index == 1: if self.index == 1:
raise McpError(ErrorData(code=-32000, message="Session terminated")) raise MCPError(-32000, "Session terminated")
return SimpleNamespace( return SimpleNamespace(
content=[mcp_types.TextContent(type="text", text="recovered")] content=[mcp_types.TextContent(type="text", text="recovered")]
) )
@ -497,7 +494,7 @@ async def test_mcp_reconnect_handler_uses_sanitized_server_prefix(
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="quote", name="quote",
description="quote tool", description="quote tool",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
registry.register(MCPToolWrapper(_FakeSession(connect_count), name, tool_def)) registry.register(MCPToolWrapper(_FakeSession(connect_count), name, tool_def))
stack = AsyncExitStack() stack = AsyncExitStack()
@ -532,7 +529,7 @@ async def test_concurrent_mcp_reconnect_reuses_fresh_session(
class _DeadSession: class _DeadSession:
async def read_resource(self, _uri: str) -> Any: async def read_resource(self, _uri: str) -> Any:
raise McpError(ErrorData(code=-32000, message="Session terminated")) raise MCPError(-32000, "Session terminated")
class _LiveSession: class _LiveSession:
async def read_resource(self, uri: str) -> Any: async def read_resource(self, uri: str) -> Any:

32
tests/agent/test_mcp_reconnect_crash.py
View File

@ -17,7 +17,7 @@ import socket
import time import time
from unittest.mock import MagicMock from unittest.mock import MagicMock
import httpx import httpx2 as httpx
import pytest import pytest
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
@ -39,31 +39,29 @@ def _free_port() -> int:
def _run_mcp_server(port: int, ready_event: multiprocessing.Event) -> None: def _run_mcp_server(port: int, ready_event: multiprocessing.Event) -> None:
"""FastMCP server target for ``multiprocessing.Process``. """MCPServer target for ``multiprocessing.Process``.
The server exposes a single ``greet`` tool and terminates idle sessions The server exposes a single ``greet`` tool and terminates idle sessions
after ``_IDLE_TIMEOUT_SECONDS``. after ``_IDLE_TIMEOUT_SECONDS``.
""" """
from mcp.server.fastmcp import FastMCP import uvicorn
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from mcp.server import MCPServer
mcp = FastMCP("IdleTimeoutDemo", json_response=True, port=port) mcp = MCPServer("IdleTimeoutDemo")
@mcp.tool() @mcp.tool()
def greet(name: str = "World") -> str: # noqa: N802 def greet(name: str = "World") -> str: # noqa: N802
"""Greet someone.""" """Greet someone."""
return f"Hello, {name}!" return f"Hello, {name}!"
mcp._session_manager = StreamableHTTPSessionManager( app = mcp.streamable_http_app(
app=mcp._mcp_server, json_response=True,
json_response=mcp.settings.json_response, host="127.0.0.1",
stateless=mcp.settings.stateless_http,
security_settings=mcp.settings.transport_security,
session_idle_timeout=_IDLE_TIMEOUT_SECONDS,
) )
mcp.session_manager.session_idle_timeout = _IDLE_TIMEOUT_SECONDS
ready_event.set() ready_event.set()
mcp.run(transport="streamable-http") uvicorn.run(app, host="127.0.0.1", port=port, log_level="warning")
async def _wait_for_server(url: str, timeout: float = 10.0) -> bool: async def _wait_for_server(url: str, timeout: float = 10.0) -> bool:
@ -128,10 +126,14 @@ def _make_loop(tmp_path, *, mcp_servers: dict) -> AgentLoop:
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True)
def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch): def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch):
"""The repro server runs on 127.0.0.1; allow nanobot to talk to it.""" """The repro server runs on 127.0.0.1; allow nanobot to talk to it."""
class TestPinnedDNSAsyncTransport(security_network.PinnedDNSAsyncTransport): class TestPinnedDNSAsyncTransport(security_network.Httpx2PinnedDNSAsyncTransport):
_resolver_lock = asyncio.Lock() _resolver_lock = asyncio.Lock()
monkeypatch.setattr(mcp_module, "PinnedDNSAsyncTransport", TestPinnedDNSAsyncTransport) monkeypatch.setattr(
mcp_module,
"Httpx2PinnedDNSAsyncTransport",
TestPinnedDNSAsyncTransport,
)
monkeypatch.setattr( monkeypatch.setattr(
mcp_module, mcp_module,
"validate_url_target", "validate_url_target",
@ -154,7 +156,7 @@ def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch):
) )
monkeypatch.setattr( monkeypatch.setattr(
mcp_module, mcp_module,
"httpx_env_proxy_mounts", "httpx2_env_proxy_mounts",
lambda: {}, lambda: {},
) )

27
tests/agent/test_mcp_transient_retry.py
View File

@ -5,9 +5,8 @@ from types import SimpleNamespace
from unittest.mock import AsyncMock, patch from unittest.mock import AsyncMock, patch
import pytest import pytest
from mcp import MCPError
from mcp import types as mcp_types from mcp import types as mcp_types
from mcp.shared.exceptions import McpError
from mcp.types import ErrorData
from nanobot.agent.tools.mcp import ( from nanobot.agent.tools.mcp import (
MCPPromptWrapper, MCPPromptWrapper,
@ -37,12 +36,16 @@ class _FakeEndOfStreamError(Exception):
_FakeEndOfStreamError.__name__ = "EndOfStream" _FakeEndOfStreamError.__name__ = "EndOfStream"
def _session_terminated_error() -> McpError: def _session_terminated_error() -> MCPError:
return McpError(ErrorData(code=-32000, message="Session terminated")) return MCPError(-32000, "Session terminated")
def _connection_closed_error() -> McpError: def _connection_closed_error() -> MCPError:
return McpError(ErrorData(code=-32000, message="Connection closed")) return MCPError(-32000, "Connection closed")
def _session_not_found_error() -> MCPError:
return MCPError(-32600, "Session not found")
def test_is_transient_recognizes_closed_resource(): def test_is_transient_recognizes_closed_resource():
@ -85,6 +88,10 @@ def test_is_session_terminated_recognizes_connection_closed_mcp_error():
assert _is_session_terminated(_connection_closed_error()) assert _is_session_terminated(_connection_closed_error())
def test_is_session_terminated_recognizes_v2_session_not_found_error():
assert _is_session_terminated(_session_not_found_error())
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# MCPToolWrapper retry behaviour # MCPToolWrapper retry behaviour
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@ -94,7 +101,7 @@ def _make_tool_def(name="test_tool"):
return SimpleNamespace( return SimpleNamespace(
name=name, name=name,
description="A test tool", description="A test tool",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
@ -415,10 +422,10 @@ async def test_prompt_fails_after_retry_exhausted():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_prompt_no_retry_on_mcp_error(): async def test_prompt_no_retry_on_mcp_error():
"""McpError (application-level) should NOT trigger retry.""" """MCPError (application-level) should NOT trigger retry."""
session = AsyncMock() session = AsyncMock()
session.get_prompt = AsyncMock( session.get_prompt = AsyncMock(
side_effect=McpError(ErrorData(code=-1, message="not found")) side_effect=MCPError(-1, "not found")
) )
wrapper = MCPPromptWrapper(session, "test_server", _make_prompt_def()) wrapper = MCPPromptWrapper(session, "test_server", _make_prompt_def())
@ -443,7 +450,7 @@ async def test_prompt_no_retry_on_non_transient():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_prompt_reconnects_on_session_terminated(): async def test_prompt_reconnects_on_session_terminated():
"""Prompt should reconnect once before falling back to McpError handling.""" """Prompt should reconnect once before falling back to MCPError handling."""
old_session = AsyncMock() old_session = AsyncMock()
old_session.get_prompt = AsyncMock(side_effect=_session_terminated_error()) old_session.get_prompt = AsyncMock(side_effect=_session_terminated_error())
new_session = AsyncMock() new_session = AsyncMock()

14
tests/security/test_security_network.py
View File

@ -9,9 +9,12 @@ from unittest.mock import patch
import pytest import pytest
from nanobot.security.network import ( from nanobot.security.network import (
Httpx2PinnedDNSAsyncTransport,
PinnedDNSAsyncTransport,
configure_ssrf_whitelist, configure_ssrf_whitelist,
contains_internal_url, contains_internal_url,
env_proxy_applies_to_url, env_proxy_applies_to_url,
httpx2_env_proxy_mounts,
httpx_env_proxy_mounts, httpx_env_proxy_mounts,
is_loopback_host, is_loopback_host,
pin_resolved_url_dns, pin_resolved_url_dns,
@ -264,6 +267,17 @@ def test_env_proxy_helpers_respect_no_proxy(monkeypatch):
assert any(transport is None for transport in mounts.values()) assert any(transport is None for transport in mounts.values())
assert any(transport is not None for transport in mounts.values()) assert any(transport is not None for transport in mounts.values())
httpx2_mounts = httpx2_env_proxy_mounts()
assert any(transport is None for transport in httpx2_mounts.values())
assert any(transport is not None for transport in httpx2_mounts.values())
def test_httpx_transports_share_global_dns_pin_lock():
assert (
Httpx2PinnedDNSAsyncTransport._resolver_lock
is PinnedDNSAsyncTransport._resolver_lock
)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# contains_internal_url — shell command scanning # contains_internal_url — shell command scanning

65
tests/tools/test_mcp_tool.py
View File

@ -7,7 +7,7 @@ from contextlib import asynccontextmanager
from pathlib import Path from pathlib import Path
from types import ModuleType, SimpleNamespace from types import ModuleType, SimpleNamespace
import httpx import httpx2 as httpx
import pytest import pytest
import nanobot.agent.tools.mcp as mcp_mod import nanobot.agent.tools.mcp as mcp_mod
@ -52,7 +52,7 @@ class _FakeBlobResourceContents:
class _FakeImageContent: class _FakeImageContent:
def __init__(self, data: str, mime_type: str = "image/png") -> None: def __init__(self, data: str, mime_type: str = "image/png") -> None:
self.data = data self.data = data
self.mimeType = mime_type self.mime_type = mime_type
@pytest.fixture @pytest.fixture
@ -111,7 +111,7 @@ def _fake_mcp_module(
@asynccontextmanager @asynccontextmanager
async def _fake_streamable_http_client(_url: str, http_client=None): async def _fake_streamable_http_client(_url: str, http_client=None):
yield object(), object(), object() yield object(), object()
mod.ClientSession = _FakeClientSession mod.ClientSession = _FakeClientSession
mod.StdioServerParameters = _FakeStdioServerParameters mod.StdioServerParameters = _FakeStdioServerParameters
@ -133,12 +133,13 @@ def _fake_mcp_module(
shared_mod = ModuleType("mcp.shared") shared_mod = ModuleType("mcp.shared")
exc_mod = ModuleType("mcp.shared.exceptions") exc_mod = ModuleType("mcp.shared.exceptions")
class _FakeMcpError(Exception): class _FakeMCPError(Exception):
def __init__(self, code: int = -1, message: str = "error"): def __init__(self, code: int = -1, message: str = "error"):
self.error = SimpleNamespace(code=code, message=message) self.error = SimpleNamespace(code=code, message=message)
super().__init__(message) super().__init__(message)
exc_mod.McpError = _FakeMcpError mod.MCPError = _FakeMCPError
exc_mod.MCPError = _FakeMCPError
monkeypatch.setitem(sys.modules, "mcp.shared", shared_mod) monkeypatch.setitem(sys.modules, "mcp.shared", shared_mod)
monkeypatch.setitem(sys.modules, "mcp.shared.exceptions", exc_mod) monkeypatch.setitem(sys.modules, "mcp.shared.exceptions", exc_mod)
@ -147,7 +148,7 @@ def _make_wrapper(session: object, *, timeout: float = 0.1) -> MCPToolWrapper:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="demo", name="demo",
description="demo tool", description="demo tool",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
return MCPToolWrapper(session, "test", tool_def, tool_timeout=timeout) return MCPToolWrapper(session, "test", tool_def, tool_timeout=timeout)
@ -185,7 +186,7 @@ def test_wrapper_preserves_non_nullable_unions() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="demo", name="demo",
description="demo tool", description="demo tool",
inputSchema={ input_schema={
"type": "object", "type": "object",
"properties": { "properties": {
"value": { "value": {
@ -207,7 +208,7 @@ def test_wrapper_normalizes_nullable_property_type_union() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="demo", name="demo",
description="demo tool", description="demo tool",
inputSchema={ input_schema={
"type": "object", "type": "object",
"properties": { "properties": {
"name": {"type": ["string", "null"]}, "name": {"type": ["string", "null"]},
@ -224,7 +225,7 @@ def test_wrapper_normalizes_nullable_property_anyof() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="demo", name="demo",
description="demo tool", description="demo tool",
inputSchema={ input_schema={
"type": "object", "type": "object",
"properties": { "properties": {
"name": { "name": {
@ -249,7 +250,7 @@ def test_wrapper_hoists_recursive_local_refs_into_defs() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="search_dataset", name="search_dataset",
description="search tool", description="search tool",
inputSchema={ input_schema={
"type": "object", "type": "object",
"properties": { "properties": {
"filter": { "filter": {
@ -282,7 +283,7 @@ def test_wrapper_hoists_root_self_ref_into_defs() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="tree", name="tree",
description="tree tool", description="tree tool",
inputSchema={ input_schema={
"type": "object", "type": "object",
"properties": { "properties": {
"children": {"type": "array", "items": {"$ref": "#"}}, "children": {"type": "array", "items": {"$ref": "#"}},
@ -304,7 +305,7 @@ def test_wrapper_preserves_existing_defs_refs() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="demo", name="demo",
description="demo tool", description="demo tool",
inputSchema={ input_schema={
"type": "object", "type": "object",
"$defs": {"value": {"type": "string"}}, "$defs": {"value": {"type": "string"}},
"properties": {"value": {"$ref": "#/$defs/value"}}, "properties": {"value": {"$ref": "#/$defs/value"}},
@ -321,7 +322,7 @@ def test_wrapper_resolves_uri_encoded_json_pointer() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="demo", name="demo",
description="demo tool", description="demo tool",
inputSchema={ input_schema={
"type": "object", "type": "object",
"properties": { "properties": {
"space name/value": {"type": "string"}, "space name/value": {"type": "string"},
@ -449,7 +450,7 @@ async def test_execute_wraps_mcp_is_error_result() -> None:
async def call_tool(_name: str, arguments: dict) -> object: async def call_tool(_name: str, arguments: dict) -> object:
return SimpleNamespace( return SimpleNamespace(
content=[_FakeTextContent("Error: server-side MCP failure")], content=[_FakeTextContent("Error: server-side MCP failure")],
isError=True, is_error=True,
) )
wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool)) wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))
@ -494,7 +495,7 @@ async def test_execute_preserves_success_text_that_starts_with_error() -> None:
async def call_tool(_name: str, arguments: dict) -> object: async def call_tool(_name: str, arguments: dict) -> object:
return SimpleNamespace( return SimpleNamespace(
content=[_FakeTextContent("Error: generated report successfully")], content=[_FakeTextContent("Error: generated report successfully")],
isError=False, is_error=False,
) )
wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool)) wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))
@ -622,7 +623,7 @@ def _make_tool_def(name: str) -> SimpleNamespace:
return SimpleNamespace( return SimpleNamespace(
name=name, name=name,
description=f"{name} tool", description=f"{name} tool",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
@ -936,7 +937,7 @@ async def test_connect_mcp_servers_env_proxy_adds_proxy_mounts_and_keeps_pinned_
@asynccontextmanager @asynccontextmanager
async def _capturing_streamable_http_client(_url: str, http_client=None): async def _capturing_streamable_http_client(_url: str, http_client=None):
assert http_client is not None assert http_client is not None
yield object(), object(), object() yield object(), object()
monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080") monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080")
monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1,::1") monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1,::1")
@ -944,11 +945,11 @@ async def test_connect_mcp_servers_env_proxy_adds_proxy_mounts_and_keeps_pinned_
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable) monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
monkeypatch.setattr( monkeypatch.setattr(
mcp_mod, mcp_mod,
"PinnedDNSAsyncTransport", "Httpx2PinnedDNSAsyncTransport",
lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)), lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
) )
monkeypatch.setattr( monkeypatch.setattr(
"nanobot.security.network.httpx.AsyncHTTPTransport", "nanobot.security.network.httpx2.AsyncHTTPTransport",
lambda **_kwargs: httpx.MockTransport( lambda **_kwargs: httpx.MockTransport(
lambda request: httpx.Response(200, request=request) lambda request: httpx.Response(200, request=request)
), ),
@ -976,11 +977,11 @@ def test_mcp_http_clients_no_proxy_env_keeps_pinned_direct_route(monkeypatch):
monkeypatch.setenv("NO_PROXY", "mcp.example.com") monkeypatch.setenv("NO_PROXY", "mcp.example.com")
monkeypatch.setattr( monkeypatch.setattr(
mcp_mod, mcp_mod,
"PinnedDNSAsyncTransport", "Httpx2PinnedDNSAsyncTransport",
lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)), lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
) )
monkeypatch.setattr( monkeypatch.setattr(
"nanobot.security.network.httpx.AsyncHTTPTransport", "nanobot.security.network.httpx2.AsyncHTTPTransport",
lambda **_kwargs: httpx.MockTransport( lambda **_kwargs: httpx.MockTransport(
lambda request: httpx.Response(200, request=request) lambda request: httpx.Response(200, request=request)
), ),
@ -1050,13 +1051,15 @@ async def test_connect_mcp_servers_http_clients_reject_unsafe_redirect_targets(
assert http_client is not None assert http_client is not None
used_transports.append("streamableHttp") used_transports.append("streamableHttp")
await http_client.get("https://example.com/start") await http_client.get("https://example.com/start")
yield object(), object(), object() yield object(), object()
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate) monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable) monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
# Keep the redirect exercise isolated from host-level proxy settings.
monkeypatch.setattr(mcp_mod, "httpx2_env_proxy_mounts", lambda: {})
monkeypatch.setattr( monkeypatch.setattr(
mcp_mod, mcp_mod,
"PinnedDNSAsyncTransport", "Httpx2PinnedDNSAsyncTransport",
lambda **_kwargs: httpx.MockTransport(_handler), lambda **_kwargs: httpx.MockTransport(_handler),
) )
monkeypatch.setattr(mcp_mod.httpx, "AsyncClient", _async_client_with_mock_transport) monkeypatch.setattr(mcp_mod.httpx, "AsyncClient", _async_client_with_mock_transport)
@ -1138,13 +1141,13 @@ async def test_connect_mcp_servers_streamable_http_uses_finite_timeout(
@asynccontextmanager @asynccontextmanager
async def _capturing_streamable_http_client(_url: str, http_client=None): async def _capturing_streamable_http_client(_url: str, http_client=None):
captured["timeout"] = http_client.timeout captured["timeout"] = http_client.timeout
yield object(), object(), object() yield object(), object()
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate) monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable) monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
monkeypatch.setattr( monkeypatch.setattr(
mcp_mod, mcp_mod,
"PinnedDNSAsyncTransport", "Httpx2PinnedDNSAsyncTransport",
lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)), lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
) )
monkeypatch.setattr( monkeypatch.setattr(
@ -1385,10 +1388,10 @@ async def test_prompt_wrapper_execute_handles_timeout() -> None:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_prompt_wrapper_execute_handles_mcp_error() -> None: async def test_prompt_wrapper_execute_handles_mcp_error() -> None:
from mcp.shared.exceptions import McpError from mcp import MCPError
async def get_prompt(name: str, arguments: dict | None = None) -> object: async def get_prompt(name: str, arguments: dict | None = None) -> object:
raise McpError(code=42, message="invalid argument") raise MCPError(code=42, message="invalid argument")
wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt)) wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt))
result = await wrapper.execute() result = await wrapper.execute()
@ -1510,7 +1513,7 @@ def test_tool_wrapper_sanitizes_name() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="My Tool", name="My Tool",
description="tool with spaces", description="tool with spaces",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "srv", tool_def) wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "srv", tool_def)
assert wrapper.name == "mcp_srv_My_Tool" assert wrapper.name == "mcp_srv_My_Tool"
@ -1541,7 +1544,7 @@ def test_tool_wrapper_preserves_original_name_for_mcp_call() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="My Tool", name="My Tool",
description="tool with spaces", description="tool with spaces",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "srv", tool_def) wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "srv", tool_def)
# The sanitized API-facing name differs from the original MCP name # The sanitized API-facing name differs from the original MCP name
@ -1619,12 +1622,12 @@ def test_long_server_name_tools_are_matched_by_server_name() -> None:
tool_def = SimpleNamespace( tool_def = SimpleNamespace(
name="search", name="search",
description="search tool", description="search tool",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
other_tool_def = SimpleNamespace( other_tool_def = SimpleNamespace(
name="search", name="search",
description="other search tool", description="other search tool",
inputSchema={"type": "object", "properties": {}}, input_schema={"type": "object", "properties": {}},
) )
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), server_name, tool_def) wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), server_name, tool_def)
other_wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "other", other_tool_def) other_wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "other", other_tool_def)