mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-04 08:28:36 +00:00
chore(mcp): migrate client integration to SDK v2
This commit is contained in:
chengyongru
parent
606ac56e8f
commit
c2349a0bfc
@ -12,7 +12,7 @@ from contextlib import AsyncExitStack, suppress
|
||||
from typing import TYPE_CHECKING, Any, Mapping, Protocol, cast
|
||||
from weakref import WeakKeyDictionary
|
||||
|
||||
import httpx
|
||||
import httpx2 as httpx
|
||||
from loguru import logger
|
||||
|
||||
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.security.network import (
|
||||
PinnedDNSAsyncTransport,
|
||||
Httpx2PinnedDNSAsyncTransport,
|
||||
env_proxy_applies_to_url,
|
||||
httpx_env_proxy_mounts,
|
||||
httpx2_env_proxy_mounts,
|
||||
resolve_url_target,
|
||||
validate_url_target,
|
||||
)
|
||||
@ -194,7 +194,7 @@ def _is_session_terminated(exc: BaseException) -> bool:
|
||||
messages.append(str(getattr(error, "message", "")))
|
||||
return any(
|
||||
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
|
||||
)
|
||||
|
||||
@ -252,8 +252,8 @@ def _redact_url(url: str) -> str:
|
||||
|
||||
|
||||
def _pinned_transport_kwargs() -> dict[str, Any]:
|
||||
kwargs: dict[str, Any] = {"transport": PinnedDNSAsyncTransport()}
|
||||
mounts = httpx_env_proxy_mounts()
|
||||
kwargs: dict[str, Any] = {"transport": Httpx2PinnedDNSAsyncTransport()}
|
||||
mounts = httpx2_env_proxy_mounts()
|
||||
if mounts:
|
||||
kwargs["mounts"] = mounts
|
||||
return kwargs
|
||||
@ -518,7 +518,7 @@ def _image_block_data_url(block: Any, types: Any) -> str | None:
|
||||
"""
|
||||
image_cls = getattr(types, "ImageContent", None)
|
||||
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}"
|
||||
|
||||
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)
|
||||
if blob_cls is not None and isinstance(resource, blob_cls):
|
||||
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/"):
|
||||
return f"data:{mime};base64,{blob_resource.blob}"
|
||||
return None
|
||||
@ -571,7 +571,7 @@ class MCPToolWrapper(_MCPWrapperBase):
|
||||
self._original_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
|
||||
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._tool_timeout = tool_timeout
|
||||
|
||||
@ -650,7 +650,7 @@ class MCPToolWrapper(_MCPWrapperBase):
|
||||
# Success — extract text and persist any image content as artifacts.
|
||||
try:
|
||||
rendered = self._render_call_result(result.content, kwargs)
|
||||
if getattr(result, "isError", False):
|
||||
if getattr(result, "is_error", False):
|
||||
return ToolResult.error(rendered)
|
||||
return rendered
|
||||
except Exception as exc:
|
||||
@ -876,8 +876,7 @@ class MCPPromptWrapper(_MCPWrapperBase):
|
||||
return True
|
||||
|
||||
async def execute(self, **kwargs: Any) -> str:
|
||||
from mcp import types
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp import MCPError, types
|
||||
|
||||
retried_transient = False
|
||||
refreshed_session = False
|
||||
@ -897,7 +896,7 @@ class MCPPromptWrapper(_MCPWrapperBase):
|
||||
raise
|
||||
logger.warning("MCP prompt '{}' was cancelled by server/SDK", self._name)
|
||||
return "(MCP prompt call was cancelled)"
|
||||
except McpError as exc:
|
||||
except MCPError as exc:
|
||||
if await self._refresh_session_after_termination(
|
||||
exc,
|
||||
refreshed_session,
|
||||
@ -1062,7 +1061,7 @@ async def connect_mcp_servers(
|
||||
**_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)
|
||||
)
|
||||
else:
|
||||
|
||||
@ -12,6 +12,7 @@ from urllib.parse import urlparse
|
||||
from urllib.request import getproxies, proxy_bypass
|
||||
|
||||
import httpx
|
||||
import httpx2
|
||||
|
||||
_BLOCKED_NETWORKS = [
|
||||
ipaddress.ip_network("0.0.0.0/8"),
|
||||
@ -29,6 +30,7 @@ _BLOCKED_NETWORKS = [
|
||||
|
||||
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
|
||||
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
|
||||
_DNS_PIN_RESOLVER_LOCK = asyncio.Lock()
|
||||
|
||||
|
||||
def is_loopback_host(host: str) -> bool:
|
||||
@ -195,6 +197,30 @@ def httpx_env_proxy_mounts() -> dict[str, httpx.AsyncBaseTransport | None]:
|
||||
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:
|
||||
if not hostname:
|
||||
return None
|
||||
@ -264,7 +290,7 @@ class UnsafeURLRequestError(httpx.RequestError):
|
||||
class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
|
||||
"""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__(
|
||||
self,
|
||||
@ -288,6 +314,37 @@ class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
|
||||
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]:
|
||||
"""Validate an already-fetched URL (e.g. after redirect). Only checks the IP, skips DNS."""
|
||||
try:
|
||||
|
||||
@ -31,6 +31,8 @@ dependencies = [
|
||||
"websockets>=15.0,<17.0",
|
||||
"websocket-client>=1.9.0,<2.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",
|
||||
"oauth-cli-kit>=0.1.6,<1.0.0",
|
||||
"loguru>=0.7.3,<1.0.0",
|
||||
@ -40,7 +42,7 @@ dependencies = [
|
||||
"croniter>=6.0.0,<7.0.0",
|
||||
"prompt-toolkit>=3.0.50,<4.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",
|
||||
"chardet>=3.0.2,<6.0.0",
|
||||
"openai>=2.8.0",
|
||||
|
||||
@ -10,10 +10,9 @@ from unittest.mock import MagicMock
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from mcp import MCPError
|
||||
from mcp import types as mcp_types
|
||||
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
|
||||
@ -26,12 +25,10 @@ from nanobot.config.schema import MCPServerConfig
|
||||
|
||||
def _mcp_notification(method: str, params: dict[str, Any] | None = None) -> SessionMessage:
|
||||
return SessionMessage(
|
||||
message=mcp_types.JSONRPCMessage(
|
||||
mcp_types.JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method=method,
|
||||
params=params,
|
||||
)
|
||||
message=mcp_types.JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method=method,
|
||||
params=params,
|
||||
)
|
||||
)
|
||||
|
||||
@ -428,7 +425,7 @@ async def test_mcp_tool_reconnects_after_session_terminated(
|
||||
self.call_count += 1
|
||||
assert arguments == {"symbol": "AAPL"}
|
||||
if self.index == 1:
|
||||
raise McpError(ErrorData(code=-32000, message="Session terminated"))
|
||||
raise MCPError(-32000, "Session terminated")
|
||||
return SimpleNamespace(
|
||||
content=[mcp_types.TextContent(type="text", text="recovered")]
|
||||
)
|
||||
@ -443,7 +440,7 @@ async def test_mcp_tool_reconnects_after_session_terminated(
|
||||
tool_def = SimpleNamespace(
|
||||
name="quote",
|
||||
description="quote tool",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
)
|
||||
registry.register(MCPToolWrapper(session, name, tool_def, tool_timeout=5))
|
||||
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:
|
||||
assert arguments == {}
|
||||
if self.index == 1:
|
||||
raise McpError(ErrorData(code=-32000, message="Session terminated"))
|
||||
raise MCPError(-32000, "Session terminated")
|
||||
return SimpleNamespace(
|
||||
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(
|
||||
name="quote",
|
||||
description="quote tool",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
)
|
||||
registry.register(MCPToolWrapper(_FakeSession(connect_count), name, tool_def))
|
||||
stack = AsyncExitStack()
|
||||
@ -532,7 +529,7 @@ async def test_concurrent_mcp_reconnect_reuses_fresh_session(
|
||||
|
||||
class _DeadSession:
|
||||
async def read_resource(self, _uri: str) -> Any:
|
||||
raise McpError(ErrorData(code=-32000, message="Session terminated"))
|
||||
raise MCPError(-32000, "Session terminated")
|
||||
|
||||
class _LiveSession:
|
||||
async def read_resource(self, uri: str) -> Any:
|
||||
|
||||
@ -17,7 +17,7 @@ import socket
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import httpx2 as httpx
|
||||
import pytest
|
||||
|
||||
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:
|
||||
"""FastMCP server target for ``multiprocessing.Process``.
|
||||
"""MCPServer target for ``multiprocessing.Process``.
|
||||
|
||||
The server exposes a single ``greet`` tool and terminates idle sessions
|
||||
after ``_IDLE_TIMEOUT_SECONDS``.
|
||||
"""
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
||||
import uvicorn
|
||||
from mcp.server import MCPServer
|
||||
|
||||
mcp = FastMCP("IdleTimeoutDemo", json_response=True, port=port)
|
||||
mcp = MCPServer("IdleTimeoutDemo")
|
||||
|
||||
@mcp.tool()
|
||||
def greet(name: str = "World") -> str: # noqa: N802
|
||||
"""Greet someone."""
|
||||
return f"Hello, {name}!"
|
||||
|
||||
mcp._session_manager = StreamableHTTPSessionManager(
|
||||
app=mcp._mcp_server,
|
||||
json_response=mcp.settings.json_response,
|
||||
stateless=mcp.settings.stateless_http,
|
||||
security_settings=mcp.settings.transport_security,
|
||||
session_idle_timeout=_IDLE_TIMEOUT_SECONDS,
|
||||
app = mcp.streamable_http_app(
|
||||
json_response=True,
|
||||
host="127.0.0.1",
|
||||
)
|
||||
mcp.session_manager.session_idle_timeout = _IDLE_TIMEOUT_SECONDS
|
||||
|
||||
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:
|
||||
@ -128,10 +126,14 @@ def _make_loop(tmp_path, *, mcp_servers: dict) -> AgentLoop:
|
||||
@pytest.fixture(autouse=True)
|
||||
def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch):
|
||||
"""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()
|
||||
|
||||
monkeypatch.setattr(mcp_module, "PinnedDNSAsyncTransport", TestPinnedDNSAsyncTransport)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"Httpx2PinnedDNSAsyncTransport",
|
||||
TestPinnedDNSAsyncTransport,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"validate_url_target",
|
||||
@ -154,7 +156,7 @@ def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch):
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"httpx_env_proxy_mounts",
|
||||
"httpx2_env_proxy_mounts",
|
||||
lambda: {},
|
||||
)
|
||||
|
||||
|
||||
@ -5,9 +5,8 @@ from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from mcp import MCPError
|
||||
from mcp import types as mcp_types
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.types import ErrorData
|
||||
|
||||
from nanobot.agent.tools.mcp import (
|
||||
MCPPromptWrapper,
|
||||
@ -37,12 +36,16 @@ class _FakeEndOfStreamError(Exception):
|
||||
_FakeEndOfStreamError.__name__ = "EndOfStream"
|
||||
|
||||
|
||||
def _session_terminated_error() -> McpError:
|
||||
return McpError(ErrorData(code=-32000, message="Session terminated"))
|
||||
def _session_terminated_error() -> MCPError:
|
||||
return MCPError(-32000, "Session terminated")
|
||||
|
||||
|
||||
def _connection_closed_error() -> McpError:
|
||||
return McpError(ErrorData(code=-32000, message="Connection closed"))
|
||||
def _connection_closed_error() -> MCPError:
|
||||
return MCPError(-32000, "Connection closed")
|
||||
|
||||
|
||||
def _session_not_found_error() -> MCPError:
|
||||
return MCPError(-32600, "Session not found")
|
||||
|
||||
|
||||
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())
|
||||
|
||||
|
||||
def test_is_session_terminated_recognizes_v2_session_not_found_error():
|
||||
assert _is_session_terminated(_session_not_found_error())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCPToolWrapper retry behaviour
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -94,7 +101,7 @@ def _make_tool_def(name="test_tool"):
|
||||
return SimpleNamespace(
|
||||
name=name,
|
||||
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
|
||||
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.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())
|
||||
@ -443,7 +450,7 @@ async def test_prompt_no_retry_on_non_transient():
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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.get_prompt = AsyncMock(side_effect=_session_terminated_error())
|
||||
new_session = AsyncMock()
|
||||
|
||||
@ -9,9 +9,12 @@ from unittest.mock import patch
|
||||
import pytest
|
||||
|
||||
from nanobot.security.network import (
|
||||
Httpx2PinnedDNSAsyncTransport,
|
||||
PinnedDNSAsyncTransport,
|
||||
configure_ssrf_whitelist,
|
||||
contains_internal_url,
|
||||
env_proxy_applies_to_url,
|
||||
httpx2_env_proxy_mounts,
|
||||
httpx_env_proxy_mounts,
|
||||
is_loopback_host,
|
||||
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 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
|
||||
|
||||
@ -7,7 +7,7 @@ from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from types import ModuleType, SimpleNamespace
|
||||
|
||||
import httpx
|
||||
import httpx2 as httpx
|
||||
import pytest
|
||||
|
||||
import nanobot.agent.tools.mcp as mcp_mod
|
||||
@ -52,7 +52,7 @@ class _FakeBlobResourceContents:
|
||||
class _FakeImageContent:
|
||||
def __init__(self, data: str, mime_type: str = "image/png") -> None:
|
||||
self.data = data
|
||||
self.mimeType = mime_type
|
||||
self.mime_type = mime_type
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@ -111,7 +111,7 @@ def _fake_mcp_module(
|
||||
|
||||
@asynccontextmanager
|
||||
async def _fake_streamable_http_client(_url: str, http_client=None):
|
||||
yield object(), object(), object()
|
||||
yield object(), object()
|
||||
|
||||
mod.ClientSession = _FakeClientSession
|
||||
mod.StdioServerParameters = _FakeStdioServerParameters
|
||||
@ -133,12 +133,13 @@ def _fake_mcp_module(
|
||||
shared_mod = ModuleType("mcp.shared")
|
||||
exc_mod = ModuleType("mcp.shared.exceptions")
|
||||
|
||||
class _FakeMcpError(Exception):
|
||||
class _FakeMCPError(Exception):
|
||||
def __init__(self, code: int = -1, message: str = "error"):
|
||||
self.error = SimpleNamespace(code=code, message=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.exceptions", exc_mod)
|
||||
|
||||
@ -147,7 +148,7 @@ def _make_wrapper(session: object, *, timeout: float = 0.1) -> MCPToolWrapper:
|
||||
tool_def = SimpleNamespace(
|
||||
name="demo",
|
||||
description="demo tool",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
)
|
||||
return MCPToolWrapper(session, "test", tool_def, tool_timeout=timeout)
|
||||
|
||||
@ -185,7 +186,7 @@ def test_wrapper_preserves_non_nullable_unions() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="demo",
|
||||
description="demo tool",
|
||||
inputSchema={
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"value": {
|
||||
@ -207,7 +208,7 @@ def test_wrapper_normalizes_nullable_property_type_union() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="demo",
|
||||
description="demo tool",
|
||||
inputSchema={
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": ["string", "null"]},
|
||||
@ -224,7 +225,7 @@ def test_wrapper_normalizes_nullable_property_anyof() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="demo",
|
||||
description="demo tool",
|
||||
inputSchema={
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {
|
||||
@ -249,7 +250,7 @@ def test_wrapper_hoists_recursive_local_refs_into_defs() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="search_dataset",
|
||||
description="search tool",
|
||||
inputSchema={
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filter": {
|
||||
@ -282,7 +283,7 @@ def test_wrapper_hoists_root_self_ref_into_defs() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="tree",
|
||||
description="tree tool",
|
||||
inputSchema={
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"children": {"type": "array", "items": {"$ref": "#"}},
|
||||
@ -304,7 +305,7 @@ def test_wrapper_preserves_existing_defs_refs() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="demo",
|
||||
description="demo tool",
|
||||
inputSchema={
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"$defs": {"value": {"type": "string"}},
|
||||
"properties": {"value": {"$ref": "#/$defs/value"}},
|
||||
@ -321,7 +322,7 @@ def test_wrapper_resolves_uri_encoded_json_pointer() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="demo",
|
||||
description="demo tool",
|
||||
inputSchema={
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"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:
|
||||
return SimpleNamespace(
|
||||
content=[_FakeTextContent("Error: server-side MCP failure")],
|
||||
isError=True,
|
||||
is_error=True,
|
||||
)
|
||||
|
||||
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:
|
||||
return SimpleNamespace(
|
||||
content=[_FakeTextContent("Error: generated report successfully")],
|
||||
isError=False,
|
||||
is_error=False,
|
||||
)
|
||||
|
||||
wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))
|
||||
@ -622,7 +623,7 @@ def _make_tool_def(name: str) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
name=name,
|
||||
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
|
||||
async def _capturing_streamable_http_client(_url: str, http_client=None):
|
||||
assert http_client is not None
|
||||
yield object(), object(), object()
|
||||
yield object(), object()
|
||||
|
||||
monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080")
|
||||
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,
|
||||
"PinnedDNSAsyncTransport",
|
||||
"Httpx2PinnedDNSAsyncTransport",
|
||||
lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"nanobot.security.network.httpx.AsyncHTTPTransport",
|
||||
"nanobot.security.network.httpx2.AsyncHTTPTransport",
|
||||
lambda **_kwargs: httpx.MockTransport(
|
||||
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.setattr(
|
||||
mcp_mod,
|
||||
"PinnedDNSAsyncTransport",
|
||||
"Httpx2PinnedDNSAsyncTransport",
|
||||
lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"nanobot.security.network.httpx.AsyncHTTPTransport",
|
||||
"nanobot.security.network.httpx2.AsyncHTTPTransport",
|
||||
lambda **_kwargs: httpx.MockTransport(
|
||||
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
|
||||
used_transports.append("streamableHttp")
|
||||
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, "_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(
|
||||
mcp_mod,
|
||||
"PinnedDNSAsyncTransport",
|
||||
"Httpx2PinnedDNSAsyncTransport",
|
||||
lambda **_kwargs: httpx.MockTransport(_handler),
|
||||
)
|
||||
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
|
||||
async def _capturing_streamable_http_client(_url: str, http_client=None):
|
||||
captured["timeout"] = http_client.timeout
|
||||
yield object(), object(), object()
|
||||
yield object(), object()
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
|
||||
monkeypatch.setattr(
|
||||
mcp_mod,
|
||||
"PinnedDNSAsyncTransport",
|
||||
"Httpx2PinnedDNSAsyncTransport",
|
||||
lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
@ -1385,10 +1388,10 @@ async def test_prompt_wrapper_execute_handles_timeout() -> None:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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:
|
||||
raise McpError(code=42, message="invalid argument")
|
||||
raise MCPError(code=42, message="invalid argument")
|
||||
|
||||
wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt))
|
||||
result = await wrapper.execute()
|
||||
@ -1510,7 +1513,7 @@ def test_tool_wrapper_sanitizes_name() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="My Tool",
|
||||
description="tool with spaces",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
)
|
||||
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "srv", tool_def)
|
||||
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(
|
||||
name="My Tool",
|
||||
description="tool with spaces",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
)
|
||||
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "srv", tool_def)
|
||||
# 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(
|
||||
name="search",
|
||||
description="search tool",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
)
|
||||
other_tool_def = SimpleNamespace(
|
||||
name="search",
|
||||
description="other search tool",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
)
|
||||
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), server_name, tool_def)
|
||||
other_wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "other", other_tool_def)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user