Skip to content

Commit 5e2057d

Browse files
feat: migrate MCP tooling to mcp 2.x
Bump the mcp dependency to >=2.0.0,<3 and adapt the MCP tool integration to the 2.0 SDK surface: - McpError -> MCPError; ProgressFnT now lives in mcp.shared.dispatcher (mcp.shared.session was removed). - Read snake_case model attributes (input_schema, output_schema, is_error, mime_type) instead of the removed camelCase aliases. - ClientSession read_timeout_seconds now takes float seconds, not a timedelta. - streamable_http_client / streamable-http layer: the public McpHttpClientFactory moved to a private module, so declare the equivalent factory Protocol locally. - to_mcp_server: mcp.server.fastmcp was removed; build the server with mcp.server.mcpserver.MCPServer instead. In mcp 2.0 the server-side Context.session is a fresh ServerSession per request rather than per connection, so key the per-connection ADK session map on the shared underlying Connection object (with a fallback to the session) to keep one conversation per MCP connection. Refs #6532
1 parent 73ecb5b commit 5e2057d

13 files changed

Lines changed: 147 additions & 77 deletions

pyproject.toml

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ optional-dependencies.all = [
8383
"google-cloud-spanner>=3.56,<4",
8484
"google-cloud-speech>=2.30,<3",
8585
"google-cloud-storage>=2.18,<4",
86-
"mcp>=1.24,<2",
86+
"mcp>=2.0.0,<3",
8787
"opentelemetry-exporter-gcp-logging>=1.9.0a0,<=1.12.0a0",
8888
"opentelemetry-exporter-gcp-monitoring>=1.9.0a0,<2",
8989
"opentelemetry-exporter-gcp-trace>=1.9,<2",
@@ -195,7 +195,7 @@ optional-dependencies.gcp = [
195195
]
196196
optional-dependencies.mcp = [
197197
"anyio>=4.9,<5",
198-
"mcp>=1.24,<2",
198+
"mcp>=2.0.0,<3",
199199
]
200200
optional-dependencies.oci = [
201201
"oci>=2.126", # OCI Generative AI native SDK (OCIGenAILlm)
@@ -243,7 +243,7 @@ optional-dependencies.test = [
243243
"litellm>=1.84",
244244
"llama-index-readers-file>=0.4",
245245
"lxml>=5.3",
246-
"mcp>=1.24,<2",
246+
"mcp>=2.0.0,<3",
247247
"openai>=2.20,<3",
248248
"opentelemetry-exporter-gcp-logging>=1.9.0a0,<=1.12.0a0",
249249
"opentelemetry-exporter-gcp-monitoring>=1.9.0a0,<2",

src/google/adk/tools/mcp_tool/_agent_to_mcp.py

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,8 @@
2323

2424
from google.genai import types
2525
from mcp import types as mcp_types
26-
from mcp.server.fastmcp import Context
27-
from mcp.server.fastmcp import FastMCP
26+
from mcp.server.mcpserver import Context
27+
from mcp.server.mcpserver import MCPServer
2828

2929
from ...agents.base_agent import BaseAgent
3030
from ...artifacts.in_memory_artifact_service import InMemoryArtifactService
@@ -68,18 +68,37 @@ def _part_to_content(part: types.Part) -> Optional[mcp_types.ContentBlock]:
6868
data = base64.b64encode(blob.data).decode("ascii")
6969
mime = blob.mime_type or "application/octet-stream"
7070
if mime.startswith("image/"):
71-
return mcp_types.ImageContent(type="image", data=data, mimeType=mime)
71+
return mcp_types.ImageContent(type="image", data=data, mime_type=mime)
7272
if mime.startswith("audio/"):
73-
return mcp_types.AudioContent(type="audio", data=data, mimeType=mime)
73+
return mcp_types.AudioContent(type="audio", data=data, mime_type=mime)
7474
return mcp_types.EmbeddedResource(
7575
type="resource",
7676
resource=mcp_types.BlobResourceContents(
77-
uri=_INLINE_RESOURCE_URI, blob=data, mimeType=mime
77+
uri=_INLINE_RESOURCE_URI, blob=data, mime_type=mime
7878
),
7979
)
8080
return None
8181

8282

83+
def _connection_key(ctx: Context) -> object:
84+
"""Returns a stable per-connection key for an MCP tool call context.
85+
86+
In mcp 2.0, ``ctx.session`` is a new ``ServerSession`` object on every
87+
request even over a single connection, so it can no longer key the
88+
per-connection ADK session map. The underlying ``Connection`` object is
89+
shared by every request on one connection, so we use it when available and
90+
fall back to ``ctx.session`` otherwise.
91+
92+
Args:
93+
ctx: The MCP tool call context.
94+
95+
Returns:
96+
A hashable object that is stable across all requests on one connection.
97+
"""
98+
connection = getattr(ctx.session, "_connection", None)
99+
return connection if connection is not None else ctx.session
100+
101+
83102
async def _run_agent(
84103
runner: Runner,
85104
request: str,
@@ -106,14 +125,14 @@ async def _run_agent(
106125
"""
107126
session_id: Optional[str] = None
108127
if ctx is not None and sessions is not None:
109-
session_id = sessions.get(ctx.session)
128+
session_id = sessions.get(_connection_key(ctx))
110129
if session_id is None:
111130
session = await runner.session_service.create_session(
112131
app_name=runner.app_name, user_id=_MCP_USER_ID
113132
)
114133
session_id = session.id
115134
if ctx is not None and sessions is not None:
116-
sessions[ctx.session] = session_id
135+
sessions[_connection_key(ctx)] = session_id
117136
new_message = types.Content(role="user", parts=[types.Part(text=request)])
118137
final_content: list[mcp_types.ContentBlock] = []
119138
async for event in runner.run_async(
@@ -142,7 +161,7 @@ def to_mcp_server(
142161
name: Optional[str] = None,
143162
instructions: Optional[str] = None,
144163
runner: Optional[Runner] = None,
145-
) -> FastMCP:
164+
) -> MCPServer:
146165
"""Exposes an ADK agent as an MCP server.
147166
148167
The returned server registers a single MCP tool that runs the agent: an MCP
@@ -166,7 +185,7 @@ def to_mcp_server(
166185
services.
167186
168187
Returns:
169-
A ``FastMCP`` server exposing the agent as a single tool.
188+
A ``MCPServer`` server exposing the agent as a single tool.
170189
171190
Example::
172191
@@ -175,7 +194,7 @@ def to_mcp_server(
175194
server.run(transport="stdio")
176195
"""
177196
tool_name = name or agent.name or "adk_agent"
178-
server = FastMCP(name=tool_name, instructions=instructions)
197+
server = MCPServer(name=tool_name, instructions=instructions)
179198
agent_runner = runner if runner is not None else _build_runner(agent)
180199
# Maps each MCP connection to its ADK session; WeakKeyDictionary drops the
181200
# entry when the connection is garbage-collected. pylint wrongly flags the

src/google/adk/tools/mcp_tool/conversion_utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ def adk_to_mcp_tool_type(tool: BaseTool) -> mcp_types.Tool:
5656
return mcp_types.Tool(
5757
name=tool.name,
5858
description=tool.description,
59-
inputSchema=input_schema,
59+
input_schema=input_schema,
6060
)
6161

6262

src/google/adk/tools/mcp_tool/mcp_session_manager.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,6 @@ class AsyncAuthorizedSession: # pylint: disable=g-bad-classes
6464
from mcp.client.sse import sse_client
6565
from mcp.client.stdio import stdio_client
6666
from mcp.client.streamable_http import create_mcp_http_client as _create_mcp_http_client
67-
from mcp.client.streamable_http import McpHttpClientFactory
6867
from mcp.client.streamable_http import streamable_http_client
6968
from pydantic import BaseModel
7069
from pydantic import ConfigDict
@@ -215,8 +214,21 @@ class SseConnectionParams(BaseModel):
215214

216215

217216
@runtime_checkable
218-
class CheckableMcpHttpClientFactory(McpHttpClientFactory, Protocol):
219-
pass
217+
class CheckableMcpHttpClientFactory(Protocol):
218+
"""Factory protocol for creating custom HTTPX async clients.
219+
220+
In mcp 2.0 the upstream ``McpHttpClientFactory`` protocol lives in the
221+
private ``mcp.shared._httpx_utils`` module, so we declare the equivalent
222+
shape locally rather than depend on a private import.
223+
"""
224+
225+
def __call__(
226+
self,
227+
headers: dict[str, str] | None = None,
228+
timeout: httpx.Timeout | None = None,
229+
auth: httpx.Auth | None = None,
230+
) -> httpx.AsyncClient:
231+
...
220232

221233

222234
class _DebugHttpxClientFactory:

src/google/adk/tools/mcp_tool/mcp_tool.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@
2828

2929
from fastapi.openapi.models import APIKeyIn
3030
from google.genai.types import FunctionDeclaration
31-
from mcp.shared.exceptions import McpError
32-
from mcp.shared.session import ProgressFnT
31+
from mcp.shared.dispatcher import ProgressFnT
32+
from mcp.shared.exceptions import MCPError
3333
from mcp.types import Tool as McpBaseTool
3434
from opentelemetry import propagate
3535
from typing_extensions import override
@@ -201,8 +201,8 @@ def _get_declaration(self) -> FunctionDeclaration:
201201
Returns:
202202
FunctionDeclaration: The Gemini function declaration for the tool.
203203
"""
204-
input_schema = self._mcp_tool.inputSchema
205-
output_schema = self._mcp_tool.outputSchema
204+
input_schema = self._mcp_tool.input_schema
205+
output_schema = self._mcp_tool.output_schema
206206
if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
207207
function_decl = FunctionDeclaration(
208208
name=self.name,
@@ -375,8 +375,8 @@ async def run_async(
375375
# any AGW policy) returns a 403 mid-tool-call.
376376
try:
377377
return await super().run_async(args=args, tool_context=tool_context)
378-
except McpError as e:
379-
logger.warning("MCP tool execution failed with McpError: %s", e)
378+
except MCPError as e:
379+
logger.warning("MCP tool execution failed with MCPError: %s", e)
380380
return {"error": f"MCP tool execution failed: {e}"}
381381
except Exception as e: # pylint: disable=broad-exception-caught
382382
logger.warning(
@@ -489,7 +489,7 @@ async def _run_async_impl(
489489

490490
def _detect_error_in_response(self, response: Any) -> str | None:
491491
"""Telemetry hook: returns an error type if the response indicates an error."""
492-
if isinstance(response, dict) and response.get("isError"):
492+
if isinstance(response, dict) and response.get("is_error"):
493493
return "MCP_TOOL_ERROR"
494494
return None
495495

src/google/adk/tools/mcp_tool/mcp_toolset.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
from mcp import StdioServerParameters
3535
from mcp.client.session import ElicitationFnT
3636
from mcp.client.session import SamplingFnT
37-
from mcp.shared.session import ProgressFnT
37+
from mcp.shared.dispatcher import ProgressFnT
3838
from mcp.types import ListResourcesResult
3939
from mcp.types import ListToolsResult
4040
from pydantic import model_validator

src/google/adk/tools/mcp_tool/session_context.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import asyncio
1818
from contextlib import AbstractAsyncContextManager
1919
from contextlib import AsyncExitStack
20-
from datetime import timedelta
2120
import logging
2221
from types import TracebackType
2322
from typing import Any
@@ -321,7 +320,7 @@ async def _run(self) -> None:
321320
session = await exit_stack.enter_async_context(
322321
ClientSession(
323322
*transports[:2],
324-
read_timeout_seconds=timedelta(seconds=self._timeout)
323+
read_timeout_seconds=float(self._timeout)
325324
if self._timeout is not None
326325
else None,
327326
sampling_callback=self._sampling_callback,
@@ -335,7 +334,7 @@ async def _run(self) -> None:
335334
session = await exit_stack.enter_async_context(
336335
ClientSession(
337336
*transports[:2],
338-
read_timeout_seconds=timedelta(seconds=self._sse_read_timeout)
337+
read_timeout_seconds=float(self._sse_read_timeout)
339338
if self._sse_read_timeout is not None
340339
else None,
341340
sampling_callback=self._sampling_callback,

tests/unittests/tools/mcp_tool/test_agent_to_mcp.py

Lines changed: 44 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,57 @@
1515
from __future__ import annotations
1616

1717
import base64
18+
from contextlib import asynccontextmanager
1819
from types import SimpleNamespace
1920
from typing import AsyncGenerator
2021

22+
import anyio
2123
from google.adk.agents.base_agent import BaseAgent
2224
from google.adk.agents.invocation_context import InvocationContext
2325
from google.adk.events.event import Event
2426
from google.adk.tools.mcp_tool._agent_to_mcp import _run_agent
2527
from google.adk.tools.mcp_tool._agent_to_mcp import to_mcp_server
2628
from google.genai import types
27-
from mcp.shared.memory import create_connected_server_and_client_session
29+
from mcp import ClientSession
30+
from mcp.server.mcpserver import MCPServer
31+
from mcp.shared.memory import create_client_server_memory_streams
2832
import pytest
2933

3034

35+
@asynccontextmanager
36+
async def _connected_client_session(server: MCPServer):
37+
"""Connects an in-memory ClientSession to an MCPServer for testing.
38+
39+
This replaces the ``create_connected_server_and_client_session`` helper that
40+
was removed in mcp 2.0. It runs the server's low-level transport on one end
41+
of an in-memory stream pair and yields a connected, initialized
42+
ClientSession on the other.
43+
44+
Args:
45+
server: The MCPServer to connect to.
46+
47+
Yields:
48+
An initialized ClientSession connected to the server.
49+
"""
50+
async with create_client_server_memory_streams() as (
51+
client_streams,
52+
server_streams,
53+
):
54+
client_read, client_write = client_streams
55+
server_read, server_write = server_streams
56+
lowlevel_server = server._lowlevel_server # pylint: disable=protected-access
57+
async with anyio.create_task_group() as task_group:
58+
task_group.start_soon(
59+
lowlevel_server.run,
60+
server_read,
61+
server_write,
62+
lowlevel_server.create_initialization_options(),
63+
)
64+
async with ClientSession(client_read, client_write) as session:
65+
await session.initialize()
66+
yield session
67+
68+
3169
class _EchoAgent(BaseAgent):
3270
"""Minimal agent that emits a single final text event."""
3371

@@ -111,7 +149,7 @@ async def test_to_mcp_server_registers_agent_as_single_tool():
111149
assert len(tools) == 1
112150
assert tools[0].name == "my_agent"
113151
assert tools[0].description == "does useful things"
114-
assert "request" in tools[0].inputSchema["properties"]
152+
assert "request" in tools[0].input_schema["properties"]
115153

116154

117155
@pytest.mark.asyncio
@@ -129,10 +167,10 @@ async def test_call_tool_runs_agent_end_to_end():
129167
agent = _EchoAgent(name="assistant")
130168
server = to_mcp_server(agent)
131169

132-
async with create_connected_server_and_client_session(server) as client:
170+
async with _connected_client_session(server) as client:
133171
result = await client.call_tool("assistant", {"request": "hi"})
134172

135-
assert not result.isError
173+
assert not result.is_error
136174
assert "hello from the agent" in result.content[0].text
137175

138176

@@ -174,7 +212,7 @@ async def test_run_agent_maps_image_output_to_image_content():
174212

175213
assert len(result) == 1
176214
assert result[0].type == "image"
177-
assert result[0].mimeType == "image/png"
215+
assert result[0].mime_type == "image/png"
178216
assert base64.b64decode(result[0].data) == png
179217

180218

@@ -209,7 +247,7 @@ async def test_call_tool_reuses_session_across_calls_on_one_connection():
209247
runner = _FakeRunner([_text_event("ok")])
210248
server = to_mcp_server(agent, runner=runner)
211249

212-
async with create_connected_server_and_client_session(server) as client:
250+
async with _connected_client_session(server) as client:
213251
await client.call_tool("assistant", {"request": "first"})
214252
await client.call_tool("assistant", {"request": "second"})
215253

0 commit comments

Comments
 (0)