diff --git a/docs/src/content/docs/agents/mcp-tool-provider.mdx b/docs/src/content/docs/agents/mcp-tool-provider.mdx index 816ddbfd..5fe97f1c 100644 --- a/docs/src/content/docs/agents/mcp-tool-provider.mdx +++ b/docs/src/content/docs/agents/mcp-tool-provider.mdx @@ -30,7 +30,12 @@ npm install @modelcontextprotocol/sdk ```bash pip install "agent-squad[mcp]" ``` -The `mcp` extra is optional — it is not pulled in when you install the base package. It installs the MCP SDK 1.x (`mcp>=1.0.0,<2`); the SDK's 2.x line is a breaking release and is not supported yet (tracked in [#634](https://github.com/2FastLabs/agent-squad/issues/634)). +The `mcp` extra is optional — it is not pulled in when you install the base package. Both SDK majors are supported: + +| Installed `mcp` version | MCP protocol versions | Notes | +|---|---|---| +| `mcp` >= 2.0 | 2026-07-28 **and** all legacy versions | auto-negotiated per server (`server/discover` probe with `initialize` fallback) | +| `mcp` 1.x | 2025-11-25 and older | the 1.x line will not receive 2026-07-28 support | @@ -153,7 +158,7 @@ provider = await MCPToolProvider.create([ ) ]) ``` -Requires `mcp` >= 1.9 (already satisfied by a fresh `pip install "agent-squad[mcp]"`). +Requires `mcp` >= 1.9 (already satisfied by a fresh `pip install "agent-squad[mcp]"`). On `mcp` 2.x, headers are delivered through the SDK's httpx client factory — same config, no code change. diff --git a/python/setup.cfg b/python/setup.cfg index 219a3c71..9bacd2b8 100644 --- a/python/setup.cfg +++ b/python/setup.cfg @@ -34,7 +34,7 @@ strands-agents = dakera = dakera>=0.12.8 mcp = - mcp>=1.0.0,<2 + mcp>=1.0.0 all = anthropic>=0.40.0 @@ -42,7 +42,7 @@ all = boto3>=1.36.18 libsql-client>=0.3.1 dakera>=0.12.8 - mcp>=1.0.0,<2 + mcp>=1.0.0 [options.packages.find] where = src diff --git a/python/src/agent_squad/tools/mcp_tool_provider.py b/python/src/agent_squad/tools/mcp_tool_provider.py index 57be6f9b..aaed1aeb 100644 --- a/python/src/agent_squad/tools/mcp_tool_provider.py +++ b/python/src/agent_squad/tools/mcp_tool_provider.py @@ -50,6 +50,26 @@ except ImportError: streamablehttp_client = None +# mcp 2.x detection: the first-class Client only exists there. When present, all +# connections go through it (mode="auto" probes server/discover for protocol +# 2026-07-28 and falls back to the legacy initialize handshake per server). +try: + from mcp import Client as _V2Client +except ImportError: + _V2Client = None + +# v2 name of the streamable HTTP transport (also present in late 1.x). +try: + from mcp.client.streamable_http import streamable_http_client +except ImportError: + streamable_http_client = None + +# Blessed httpx2/httpx client factory (headers, MCP-recommended timeouts). +try: + from mcp.shared._httpx_utils import create_mcp_http_client +except ImportError: + create_mcp_http_client = None + from pydantic import AnyUrl # mcp depends on pydantic, so it's available whenever the import above succeeds from dataclasses import dataclass, field @@ -86,6 +106,23 @@ class MCPServerConfig: headers: Optional[dict[str, str]] = None +_UNSET = object() + + +def _field(obj: Any, *names: str, default: Any = None) -> Any: + """Read the first present attribute among ``names``. + + mcp 1.x exposes camelCase attributes (``inputSchema``, ``isError``) while + 2.x is strictly snake_case (``input_schema``, ``is_error``) with no alias + attribute access — both spellings must be tried. + """ + for name in names: + value = getattr(obj, name, _UNSET) + if value is not _UNSET: + return value + return default + + def _meta_dict(mcp_tool: Any) -> Optional[dict[str, Any]]: """The tool's ``_meta`` (MCP Apps UI metadata); the SDK exposes it as ``.meta``.""" return getattr(mcp_tool, "meta", None) or getattr(mcp_tool, "_meta", None) @@ -204,45 +241,13 @@ async def _ensure_connected(self) -> None: return for server_cfg in self._servers: - if server_cfg.type == "stdio": - if not server_cfg.command: - raise ValueError("MCPServerConfig with type='stdio' requires a 'command'") - params = StdioServerParameters( - command=server_cfg.command, - args=server_cfg.args or [], - env=server_cfg.env, - ) - cm = stdio_client(params) - elif server_cfg.type == "sse": - if not server_cfg.url: - raise ValueError("MCPServerConfig with type='sse' requires a 'url'") - cm = sse_client(server_cfg.url, headers=server_cfg.headers or {}) - elif server_cfg.type == "streamable-http": - if streamablehttp_client is None: - raise ImportError( - "The streamable-http transport requires mcp>=1.9. " - "Upgrade it with: pip install -U 'mcp>=1.9,<2'" - ) - if not server_cfg.url: - raise ValueError("MCPServerConfig with type='streamable-http' requires a 'url'") - cm = streamablehttp_client(server_cfg.url, headers=server_cfg.headers or {}) + if _V2Client is not None: + session = await self._connect_v2(server_cfg) else: - raise ValueError( - f"Unsupported MCPServerConfig type: '{server_cfg.type}'. " - "Use 'stdio', 'streamable-http' or 'sse'." - ) - - # stdio/sse yield (read, write); streamable-http yields (read, write, get_session_id) - read, write, *_ = await cm.__aenter__() - self._cm_stack.append(cm) + session = await self._connect_v1(server_cfg) - session = ClientSession(read, write) - await session.__aenter__() - self._sessions.append(session) - await session.initialize() - - tools_result = await session.list_tools() - for mcp_tool in tools_result.tools: + tools = await self._list_all_tools(session) + for mcp_tool in tools: meta = _meta_dict(mcp_tool) self._tool_map[mcp_tool.name] = _MCPToolEntry( session=session, @@ -253,6 +258,85 @@ async def _ensure_connected(self) -> None: self._connected = True + def _transport_cm(self, server_cfg: MCPServerConfig, v2: bool) -> Any: + """The transport async context manager for a server config (shared by both majors).""" + if server_cfg.type == "stdio": + if not server_cfg.command: + raise ValueError("MCPServerConfig with type='stdio' requires a 'command'") + params = StdioServerParameters( + command=server_cfg.command, + args=server_cfg.args or [], + env=server_cfg.env, + ) + return stdio_client(params) + if server_cfg.type == "sse": + if not server_cfg.url: + raise ValueError("MCPServerConfig with type='sse' requires a 'url'") + return sse_client(server_cfg.url, headers=server_cfg.headers or {}) + if server_cfg.type == "streamable-http": + if not server_cfg.url: + raise ValueError("MCPServerConfig with type='streamable-http' requires a 'url'") + if v2: + # v2 renamed the function and moved headers onto an httpx client. + if server_cfg.headers: + if create_mcp_http_client is not None: + http_client = create_mcp_http_client(headers=server_cfg.headers) + else: + # The factory lives in a private mcp module; fall back to a + # plain client (httpx2 is a hard dependency of mcp 2.x). + import httpx2 + + http_client = httpx2.AsyncClient( + headers=server_cfg.headers, follow_redirects=True + ) + # v2 does not manage a caller-provided client; we own its lifecycle. + self._cm_stack.append(http_client) + return streamable_http_client(server_cfg.url, http_client=http_client) + return streamable_http_client(server_cfg.url) + if streamablehttp_client is None: + raise ImportError( + "The streamable-http transport requires mcp>=1.9. " + "Upgrade it with: pip install -U mcp" + ) + return streamablehttp_client(server_cfg.url, headers=server_cfg.headers or {}) + raise ValueError( + f"Unsupported MCPServerConfig type: '{server_cfg.type}'. " + "Use 'stdio', 'streamable-http' or 'sse'." + ) + + async def _connect_v1(self, server_cfg: MCPServerConfig) -> Any: + """mcp 1.x: transport streams + ClientSession + legacy initialize handshake.""" + cm = self._transport_cm(server_cfg, v2=False) + # stdio/sse yield (read, write); streamable-http yields (read, write, get_session_id) + read, write, *_ = await cm.__aenter__() + self._cm_stack.append(cm) + + session = ClientSession(read, write) + await session.__aenter__() + self._sessions.append(session) + await session.initialize() + return session + + async def _connect_v2(self, server_cfg: MCPServerConfig) -> Any: + """mcp 2.x: one Client per server; mode="auto" negotiates the protocol era.""" + client = _V2Client(self._transport_cm(server_cfg, v2=True), mode="auto") + await client.__aenter__() + self._sessions.append(client) + return client + + async def _list_all_tools(self, session: Any) -> list[Any]: + """All tools from a session, following pagination cursors when present.""" + result = await session.list_tools() + tools = list(result.tools) + cursor = _field(result, "nextCursor", "next_cursor") + # v1 deliberately keeps its pre-existing single-call behavior (no behavior + # change for existing users); only the v2 path follows pagination cursors. + while cursor and _V2Client is not None and isinstance(session, _V2Client): + result = await session.list_tools(cursor=cursor) + tools.extend(result.tools) + cursor = _field(result, "nextCursor", "next_cursor") + return tools + async def disconnect(self) -> None: """Disconnect from all MCP servers and release resources. @@ -373,11 +457,11 @@ async def _call_mcp_tool(self, tool_name: str, input_data: dict) -> ToolResult: ] text = "\n".join(parts) - if getattr(call_result, "isError", False): + if _field(call_result, "isError", "is_error", default=False): # Surface the error text back to the model so it can react. return ToolResult(content=f"Tool error: {text}" if text else "Tool returned an error") - structured = getattr(call_result, "structuredContent", None) or {} + structured = _field(call_result, "structuredContent", "structured_content") or {} ui: Optional[UIPayload] = None if entry.ui: @@ -402,14 +486,18 @@ async def _template_for(self, session: Any, resource_uri: str) -> Optional[tuple if cache_key in self._template_cache: return self._template_cache[cache_key] try: - read_result = await session.read_resource(AnyUrl(resource_uri)) + # v2's read_resource takes a plain str; v1's ClientSession wants AnyUrl. + if _V2Client is not None and isinstance(session, _V2Client): + read_result = await session.read_resource(resource_uri) + else: + read_result = await session.read_resource(AnyUrl(resource_uri)) except Exception: # noqa: BLE001 return None contents = getattr(read_result, "contents", None) or [] if not contents: return None first = contents[0] - mime_type = getattr(first, "mimeType", None) or "text/html;profile=mcp-app" + mime_type = _field(first, "mimeType", "mime_type") or "text/html;profile=mcp-app" body = getattr(first, "text", None) if body is None: blob = getattr(first, "blob", None) @@ -441,9 +529,10 @@ def to_bedrock_format(self) -> list[dict[str, Any]]: if not entry.model_visible: continue # app-only tool: callable by the UI, never advertised to the model mcp_tool = entry.tool + raw_schema = _field(mcp_tool, "inputSchema", "input_schema") input_schema = ( - mcp_tool.inputSchema - if isinstance(mcp_tool.inputSchema, dict) + raw_schema + if isinstance(raw_schema, dict) else {"type": "object", "properties": {}} ) result.append( @@ -464,9 +553,10 @@ def to_claude_format(self) -> list[dict[str, Any]]: if not entry.model_visible: continue # app-only tool: callable by the UI, never advertised to the model mcp_tool = entry.tool + raw_schema = _field(mcp_tool, "inputSchema", "input_schema") input_schema = ( - mcp_tool.inputSchema - if isinstance(mcp_tool.inputSchema, dict) + raw_schema + if isinstance(raw_schema, dict) else {"type": "object", "properties": {}} ) result.append( @@ -489,9 +579,10 @@ def to_openai_format(self) -> list[dict[str, Any]]: if not entry.model_visible: continue # app-only tool: callable by the UI, never advertised to the model mcp_tool = entry.tool + raw_schema = _field(mcp_tool, "inputSchema", "input_schema") input_schema = ( - mcp_tool.inputSchema - if isinstance(mcp_tool.inputSchema, dict) + raw_schema + if isinstance(raw_schema, dict) else {"type": "object", "properties": {}} ) # Ensure required field is present for strict mode compatibility diff --git a/python/src/tests/tools/test_mcp_tool_provider.py b/python/src/tests/tools/test_mcp_tool_provider.py index e3759dda..1e248cb5 100644 --- a/python/src/tests/tools/test_mcp_tool_provider.py +++ b/python/src/tests/tools/test_mcp_tool_provider.py @@ -70,6 +70,8 @@ def mock_mcp_modules(): patch("agent_squad.tools.mcp_tool_provider.sse_client", mock_sse_client), patch("agent_squad.tools.mcp_tool_provider.streamablehttp_client", mock_streamablehttp_client), patch("agent_squad.tools.mcp_tool_provider.StdioServerParameters", mock_stdio_params_cls), + # Pin these tests to the mcp 1.x code path regardless of the installed major. + patch("agent_squad.tools.mcp_tool_provider._V2Client", None), ): yield { "ClientSession": mock_client_session_cls, @@ -716,3 +718,253 @@ async def test_app_only_tool_still_callable(mock_mcp_modules): result = await provider._call_mcp_tool("refresh_order", {}) assert result.content == "refreshed" assert result.structured_content == {"s": 2} + + +# --------------------------------------------------------------------------- +# mcp 2.x code path (v2 Client, snake_case types) +# --------------------------------------------------------------------------- + +def _make_mcp_tool_v2(name: str, description: str = "", meta: dict | None = None): + """A mock shaped like an mcp-types 2.x Tool: strictly snake_case attributes.""" + return SimpleNamespace( + name=name, + description=description, + input_schema={"type": "object", "properties": {"q": {"type": "string"}}, "required": ["q"]}, + meta=meta, + ) + + +def _make_call_result_v2(text: str, is_error: bool = False, structured: dict | None = None): + """A mock shaped like an mcp-types 2.x CallToolResult: strictly snake_case.""" + return SimpleNamespace( + content=[SimpleNamespace(text=text)], + is_error=is_error, + structured_content=structured, + meta=None, + ) + + +class _FakeV2Client: + """Stands in for mcp.Client (2.x): async CM, constructor records transport/mode.""" + + instances: list["_FakeV2Client"] = [] + + def __init__(self, transport, *, mode=None): + self.transport = transport + self.mode = mode + self.entered = False + self.exited = False + self.list_tools = AsyncMock() + self.call_tool = AsyncMock() + self.read_resource = AsyncMock() + _FakeV2Client.instances.append(self) + + async def __aenter__(self): + self.entered = True + return self + + async def __aexit__(self, *args): + self.exited = True + return False + + +@pytest.fixture() +def mock_mcp_v2_modules(): + """Patch the provider onto the mcp 2.x code path.""" + _FakeV2Client.instances = [] + mock_stdio_client = MagicMock(return_value="stdio-transport-cm") + mock_sse_client = MagicMock(return_value="sse-transport-cm") + mock_streamable = MagicMock(return_value="streamable-transport-cm") + mock_httpx_factory = MagicMock(return_value=MagicMock(name="httpx2-client")) + mock_stdio_params_cls = MagicMock() + + with ( + patch("agent_squad.tools.mcp_tool_provider._V2Client", _FakeV2Client), + patch("agent_squad.tools.mcp_tool_provider.stdio_client", mock_stdio_client), + patch("agent_squad.tools.mcp_tool_provider.sse_client", mock_sse_client), + patch("agent_squad.tools.mcp_tool_provider.streamable_http_client", mock_streamable), + patch("agent_squad.tools.mcp_tool_provider.create_mcp_http_client", mock_httpx_factory), + patch("agent_squad.tools.mcp_tool_provider.StdioServerParameters", mock_stdio_params_cls), + patch("agent_squad.tools.mcp_tool_provider.ClientSession", MagicMock()), + ): + yield { + "stdio_client": mock_stdio_client, + "sse_client": mock_sse_client, + "streamable_http_client": mock_streamable, + "create_mcp_http_client": mock_httpx_factory, + "StdioServerParameters": mock_stdio_params_cls, + } + + +def _seed_v2_list_tools(pages): + """Context manager: every _FakeV2Client built inside returns these list_tools pages.""" + orig_init = _FakeV2Client.__init__ + + def seeded_init(self, transport, *, mode=None): + orig_init(self, transport, mode=mode) + if len(pages) == 1: + self.list_tools.return_value = pages[0] + else: + self.list_tools.side_effect = list(pages) + + return patch.object(_FakeV2Client, "__init__", seeded_init) + + +@pytest.mark.asyncio +async def test_v2_stdio_connection(mock_mcp_v2_modules): + from agent_squad.tools.mcp_tool_provider import MCPToolProvider, MCPServerConfig + + tool = _make_mcp_tool_v2("search", "Search") + provider = MCPToolProvider([MCPServerConfig(type="stdio", command="uvx", args=["s"])]) + + with _seed_v2_list_tools([SimpleNamespace(tools=[tool], next_cursor=None)]): + await provider._ensure_connected() + + assert provider._connected + assert "search" in provider._tool_map + client = _FakeV2Client.instances[0] + assert client.entered + assert client.mode == "auto" + assert client.transport == "stdio-transport-cm" + # v1 machinery untouched on the v2 path + assert provider._cm_stack == [] + + +@pytest.mark.asyncio +async def test_v2_streamable_http_headers_via_httpx_client(mock_mcp_v2_modules): + from agent_squad.tools.mcp_tool_provider import MCPToolProvider, MCPServerConfig + + tool = _make_mcp_tool_v2("search", "Search") + provider = MCPToolProvider([ + MCPServerConfig(type="streamable-http", url="http://h/mcp", headers={"x-k": "v"}) + ]) + + with _seed_v2_list_tools([SimpleNamespace(tools=[tool], next_cursor=None)]): + await provider._ensure_connected() + + mock_mcp_v2_modules["create_mcp_http_client"].assert_called_once_with(headers={"x-k": "v"}) + http_client = mock_mcp_v2_modules["create_mcp_http_client"].return_value + mock_mcp_v2_modules["streamable_http_client"].assert_called_once_with( + "http://h/mcp", http_client=http_client + ) + # We own the httpx client's lifecycle: it must sit on the cm stack for disconnect. + assert http_client in provider._cm_stack + + +@pytest.mark.asyncio +async def test_v2_list_tools_pagination(mock_mcp_v2_modules): + from agent_squad.tools.mcp_tool_provider import MCPToolProvider, MCPServerConfig + + t1 = _make_mcp_tool_v2("alpha") + t2 = _make_mcp_tool_v2("beta") + provider = MCPToolProvider([MCPServerConfig(type="stdio", command="uvx")]) + + with _seed_v2_list_tools([ + SimpleNamespace(tools=[t1], next_cursor="page2"), + SimpleNamespace(tools=[t2], next_cursor=None), + ]): + await provider._ensure_connected() + + assert set(provider._tool_map) == {"alpha", "beta"} + client = _FakeV2Client.instances[0] + assert client.list_tools.call_args_list[1].kwargs == {"cursor": "page2"} + + +@pytest.mark.asyncio +async def test_v2_snake_case_call_result(mock_mcp_v2_modules): + """is_error / structured_content must be honored on 2.x result shapes.""" + from agent_squad.tools.mcp_tool_provider import ( + MCPToolProvider, MCPServerConfig, _MCPToolEntry, _meta_dict, _ui_resource_uri, _model_visible, + ) + + tool = _make_mcp_tool_v2("weather") + provider = MCPToolProvider([MCPServerConfig(type="stdio", command="uvx")]) + session = _FakeV2Client("t", mode="auto") + meta = _meta_dict(tool) + provider._tool_map["weather"] = _MCPToolEntry( + session=session, tool=tool, ui=_ui_resource_uri(meta), model_visible=_model_visible(meta), + ) + provider._connected = True + + session.call_tool.return_value = _make_call_result_v2("Sunny", structured={"temp": 25}) + ok = await provider._call_mcp_tool("weather", {"q": "Paris"}) + assert ok.content == "Sunny" + assert ok.structured_content == {"temp": 25} + + session.call_tool.return_value = _make_call_result_v2("boom", is_error=True) + err = await provider._call_mcp_tool("weather", {"q": "Paris"}) + assert "Tool error" in err.content + + +def test_v2_snake_case_input_schema(mock_mcp_v2_modules): + """to_*_format must read input_schema on 2.x tool shapes.""" + from agent_squad.tools.mcp_tool_provider import ( + MCPToolProvider, MCPServerConfig, _MCPToolEntry, + ) + + tool = _make_mcp_tool_v2("calc", "Math") + provider = MCPToolProvider([MCPServerConfig(type="stdio", command="uvx")]) + provider._tool_map["calc"] = _MCPToolEntry(session=MagicMock(), tool=tool) + provider._connected = True + + bedrock = provider.to_bedrock_format() + assert bedrock[0]["toolSpec"]["inputSchema"]["json"]["properties"] == {"q": {"type": "string"}} + claude = provider.to_claude_format() + assert claude[0]["input_schema"]["required"] == ["q"] + + +@pytest.mark.asyncio +async def test_v2_read_resource_receives_plain_str(mock_mcp_v2_modules): + """2.x read_resource rejects AnyUrl — the provider must pass the URI as str.""" + from agent_squad.tools.mcp_tool_provider import MCPToolProvider, MCPServerConfig + + provider = MCPToolProvider([MCPServerConfig(type="stdio", command="uvx")]) + session = _FakeV2Client("t", mode="auto") + session.read_resource.return_value = SimpleNamespace( + contents=[SimpleNamespace(uri="ui://x", mime_type="text/html;profile=mcp-app", text="")] + ) + + template = await provider._template_for(session, "ui://x") + + assert template == ("text/html;profile=mcp-app", "") + session.read_resource.assert_called_once_with("ui://x") + assert isinstance(session.read_resource.call_args.args[0], str) + + +@pytest.mark.asyncio +async def test_v2_disconnect_closes_client_and_owned_httpx_client(mock_mcp_v2_modules): + """disconnect must exit the v2 Client first, then close the httpx client we own.""" + from agent_squad.tools.mcp_tool_provider import MCPToolProvider, MCPServerConfig + + tool = _make_mcp_tool_v2("search") + # An AsyncMock __aexit__ is essential: a bare MagicMock would raise on await + # and be swallowed by disconnect's except, making this test trivially green. + httpx_client = MagicMock() + httpx_client.__aexit__ = AsyncMock(return_value=False) + mock_mcp_v2_modules["create_mcp_http_client"].return_value = httpx_client + + provider = MCPToolProvider([ + MCPServerConfig(type="streamable-http", url="http://h/mcp", headers={"x-k": "v"}) + ]) + with _seed_v2_list_tools([SimpleNamespace(tools=[tool], next_cursor=None)]): + await provider._ensure_connected() + client = _FakeV2Client.instances[0] + + await provider.disconnect() + + assert client.exited is True + httpx_client.__aexit__.assert_awaited_once() + assert provider._sessions == [] + assert provider._cm_stack == [] + assert provider._tool_map == {} + assert provider._connected is False + + +def test_field_helper_reads_both_spellings(): + from agent_squad.tools.mcp_tool_provider import _field + + v1 = SimpleNamespace(isError=True) + v2 = SimpleNamespace(is_error=True) + assert _field(v1, "isError", "is_error", default=False) is True + assert _field(v2, "isError", "is_error", default=False) is True + assert _field(SimpleNamespace(), "isError", "is_error", default=False) is False