diff --git a/pyproject.toml b/pyproject.toml index 21cd6a1f11..10e3e16984 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,7 +83,7 @@ optional-dependencies.all = [ "google-cloud-spanner>=3.56,<4", "google-cloud-speech>=2.30,<3", "google-cloud-storage>=2.18,<4", - "mcp>=1.24,<2", + "mcp>=2.0.0,<3", "opentelemetry-exporter-gcp-logging>=1.9.0a0,<=1.12.0a0", "opentelemetry-exporter-gcp-monitoring>=1.9.0a0,<2", "opentelemetry-exporter-gcp-trace>=1.9,<2", @@ -195,7 +195,7 @@ optional-dependencies.gcp = [ ] optional-dependencies.mcp = [ "anyio>=4.9,<5", - "mcp>=1.24,<2", + "mcp>=2.0.0,<3", ] optional-dependencies.oci = [ "oci>=2.126", # OCI Generative AI native SDK (OCIGenAILlm) @@ -243,7 +243,7 @@ optional-dependencies.test = [ "litellm>=1.84", "llama-index-readers-file>=0.4", "lxml>=5.3", - "mcp>=1.24,<2", + "mcp>=2.0.0,<3", "openai>=2.20,<3", "opentelemetry-exporter-gcp-logging>=1.9.0a0,<=1.12.0a0", "opentelemetry-exporter-gcp-monitoring>=1.9.0a0,<2", diff --git a/src/google/adk/tools/mcp_tool/_agent_to_mcp.py b/src/google/adk/tools/mcp_tool/_agent_to_mcp.py index 5c521927d5..ccf9bf7a58 100644 --- a/src/google/adk/tools/mcp_tool/_agent_to_mcp.py +++ b/src/google/adk/tools/mcp_tool/_agent_to_mcp.py @@ -23,8 +23,8 @@ from google.genai import types from mcp import types as mcp_types -from mcp.server.fastmcp import Context -from mcp.server.fastmcp import FastMCP +from mcp.server.mcpserver import Context +from mcp.server.mcpserver import MCPServer from ...agents.base_agent import BaseAgent from ...artifacts.in_memory_artifact_service import InMemoryArtifactService @@ -68,18 +68,37 @@ def _part_to_content(part: types.Part) -> Optional[mcp_types.ContentBlock]: data = base64.b64encode(blob.data).decode("ascii") mime = blob.mime_type or "application/octet-stream" if mime.startswith("image/"): - return mcp_types.ImageContent(type="image", data=data, mimeType=mime) + return mcp_types.ImageContent(type="image", data=data, mime_type=mime) if mime.startswith("audio/"): - return mcp_types.AudioContent(type="audio", data=data, mimeType=mime) + return mcp_types.AudioContent(type="audio", data=data, mime_type=mime) return mcp_types.EmbeddedResource( type="resource", resource=mcp_types.BlobResourceContents( - uri=_INLINE_RESOURCE_URI, blob=data, mimeType=mime + uri=_INLINE_RESOURCE_URI, blob=data, mime_type=mime ), ) return None +def _connection_key(ctx: Context) -> object: + """Returns a stable per-connection key for an MCP tool call context. + + In mcp 2.0, ``ctx.session`` is a new ``ServerSession`` object on every + request even over a single connection, so it can no longer key the + per-connection ADK session map. The underlying ``Connection`` object is + shared by every request on one connection, so we use it when available and + fall back to ``ctx.session`` otherwise. + + Args: + ctx: The MCP tool call context. + + Returns: + A hashable object that is stable across all requests on one connection. + """ + connection = getattr(ctx.session, "_connection", None) + return connection if connection is not None else ctx.session + + async def _run_agent( runner: Runner, request: str, @@ -106,14 +125,14 @@ async def _run_agent( """ session_id: Optional[str] = None if ctx is not None and sessions is not None: - session_id = sessions.get(ctx.session) + session_id = sessions.get(_connection_key(ctx)) if session_id is None: session = await runner.session_service.create_session( app_name=runner.app_name, user_id=_MCP_USER_ID ) session_id = session.id if ctx is not None and sessions is not None: - sessions[ctx.session] = session_id + sessions[_connection_key(ctx)] = session_id new_message = types.Content(role="user", parts=[types.Part(text=request)]) final_content: list[mcp_types.ContentBlock] = [] async for event in runner.run_async( @@ -142,7 +161,7 @@ def to_mcp_server( name: Optional[str] = None, instructions: Optional[str] = None, runner: Optional[Runner] = None, -) -> FastMCP: +) -> MCPServer: """Exposes an ADK agent as an MCP server. The returned server registers a single MCP tool that runs the agent: an MCP @@ -166,7 +185,7 @@ def to_mcp_server( services. Returns: - A ``FastMCP`` server exposing the agent as a single tool. + A ``MCPServer`` server exposing the agent as a single tool. Example:: @@ -175,7 +194,7 @@ def to_mcp_server( server.run(transport="stdio") """ tool_name = name or agent.name or "adk_agent" - server = FastMCP(name=tool_name, instructions=instructions) + server = MCPServer(name=tool_name, instructions=instructions) agent_runner = runner if runner is not None else _build_runner(agent) # Maps each MCP connection to its ADK session; WeakKeyDictionary drops the # entry when the connection is garbage-collected. pylint wrongly flags the diff --git a/src/google/adk/tools/mcp_tool/conversion_utils.py b/src/google/adk/tools/mcp_tool/conversion_utils.py index ddf5d3aa34..0e87023ede 100644 --- a/src/google/adk/tools/mcp_tool/conversion_utils.py +++ b/src/google/adk/tools/mcp_tool/conversion_utils.py @@ -56,7 +56,7 @@ def adk_to_mcp_tool_type(tool: BaseTool) -> mcp_types.Tool: return mcp_types.Tool( name=tool.name, description=tool.description, - inputSchema=input_schema, + input_schema=input_schema, ) diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 4d130c59bd..9d11d0f68c 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -64,7 +64,6 @@ class AsyncAuthorizedSession: # pylint: disable=g-bad-classes from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client from mcp.client.streamable_http import create_mcp_http_client as _create_mcp_http_client -from mcp.client.streamable_http import McpHttpClientFactory from mcp.client.streamable_http import streamable_http_client from pydantic import BaseModel from pydantic import ConfigDict @@ -215,8 +214,21 @@ class SseConnectionParams(BaseModel): @runtime_checkable -class CheckableMcpHttpClientFactory(McpHttpClientFactory, Protocol): - pass +class CheckableMcpHttpClientFactory(Protocol): + """Factory protocol for creating custom HTTPX async clients. + + In mcp 2.0 the upstream ``McpHttpClientFactory`` protocol lives in the + private ``mcp.shared._httpx_utils`` module, so we declare the equivalent + shape locally rather than depend on a private import. + """ + + def __call__( + self, + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + ... class _DebugHttpxClientFactory: diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 3be223af84..213ec23533 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -28,8 +28,8 @@ from fastapi.openapi.models import APIKeyIn from google.genai.types import FunctionDeclaration -from mcp.shared.exceptions import McpError -from mcp.shared.session import ProgressFnT +from mcp.shared.dispatcher import ProgressFnT +from mcp.shared.exceptions import MCPError from mcp.types import Tool as McpBaseTool from opentelemetry import propagate from typing_extensions import override @@ -201,8 +201,8 @@ def _get_declaration(self) -> FunctionDeclaration: Returns: FunctionDeclaration: The Gemini function declaration for the tool. """ - input_schema = self._mcp_tool.inputSchema - output_schema = self._mcp_tool.outputSchema + input_schema = self._mcp_tool.input_schema + output_schema = self._mcp_tool.output_schema if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL): function_decl = FunctionDeclaration( name=self.name, @@ -375,8 +375,8 @@ async def run_async( # any AGW policy) returns a 403 mid-tool-call. try: return await super().run_async(args=args, tool_context=tool_context) - except McpError as e: - logger.warning("MCP tool execution failed with McpError: %s", e) + except MCPError as e: + logger.warning("MCP tool execution failed with MCPError: %s", e) return {"error": f"MCP tool execution failed: {e}"} except Exception as e: # pylint: disable=broad-exception-caught logger.warning( @@ -489,7 +489,7 @@ async def _run_async_impl( def _detect_error_in_response(self, response: Any) -> str | None: """Telemetry hook: returns an error type if the response indicates an error.""" - if isinstance(response, dict) and response.get("isError"): + if isinstance(response, dict) and response.get("is_error"): return "MCP_TOOL_ERROR" return None diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index e8531fcaa6..8180925854 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -34,7 +34,7 @@ from mcp import StdioServerParameters from mcp.client.session import ElicitationFnT from mcp.client.session import SamplingFnT -from mcp.shared.session import ProgressFnT +from mcp.shared.dispatcher import ProgressFnT from mcp.types import ListResourcesResult from mcp.types import ListToolsResult from pydantic import model_validator diff --git a/src/google/adk/tools/mcp_tool/session_context.py b/src/google/adk/tools/mcp_tool/session_context.py index f08b03da3c..af6fa539ba 100644 --- a/src/google/adk/tools/mcp_tool/session_context.py +++ b/src/google/adk/tools/mcp_tool/session_context.py @@ -17,7 +17,6 @@ import asyncio from contextlib import AbstractAsyncContextManager from contextlib import AsyncExitStack -from datetime import timedelta import logging from types import TracebackType from typing import Any @@ -321,7 +320,7 @@ async def _run(self) -> None: session = await exit_stack.enter_async_context( ClientSession( *transports[:2], - read_timeout_seconds=timedelta(seconds=self._timeout) + read_timeout_seconds=float(self._timeout) if self._timeout is not None else None, sampling_callback=self._sampling_callback, @@ -335,7 +334,7 @@ async def _run(self) -> None: session = await exit_stack.enter_async_context( ClientSession( *transports[:2], - read_timeout_seconds=timedelta(seconds=self._sse_read_timeout) + read_timeout_seconds=float(self._sse_read_timeout) if self._sse_read_timeout is not None else None, sampling_callback=self._sampling_callback, diff --git a/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py b/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py index ceefcf1342..ed8c5cfb01 100644 --- a/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py +++ b/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py @@ -15,19 +15,57 @@ from __future__ import annotations import base64 +from contextlib import asynccontextmanager from types import SimpleNamespace from typing import AsyncGenerator +import anyio from google.adk.agents.base_agent import BaseAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events.event import Event from google.adk.tools.mcp_tool._agent_to_mcp import _run_agent from google.adk.tools.mcp_tool._agent_to_mcp import to_mcp_server from google.genai import types -from mcp.shared.memory import create_connected_server_and_client_session +from mcp import ClientSession +from mcp.server.mcpserver import MCPServer +from mcp.shared.memory import create_client_server_memory_streams import pytest +@asynccontextmanager +async def _connected_client_session(server: MCPServer): + """Connects an in-memory ClientSession to an MCPServer for testing. + + This replaces the ``create_connected_server_and_client_session`` helper that + was removed in mcp 2.0. It runs the server's low-level transport on one end + of an in-memory stream pair and yields a connected, initialized + ClientSession on the other. + + Args: + server: The MCPServer to connect to. + + Yields: + An initialized ClientSession connected to the server. + """ + async with create_client_server_memory_streams() as ( + client_streams, + server_streams, + ): + client_read, client_write = client_streams + server_read, server_write = server_streams + lowlevel_server = server._lowlevel_server # pylint: disable=protected-access + async with anyio.create_task_group() as task_group: + task_group.start_soon( + lowlevel_server.run, + server_read, + server_write, + lowlevel_server.create_initialization_options(), + ) + async with ClientSession(client_read, client_write) as session: + await session.initialize() + yield session + + class _EchoAgent(BaseAgent): """Minimal agent that emits a single final text event.""" @@ -111,7 +149,7 @@ async def test_to_mcp_server_registers_agent_as_single_tool(): assert len(tools) == 1 assert tools[0].name == "my_agent" assert tools[0].description == "does useful things" - assert "request" in tools[0].inputSchema["properties"] + assert "request" in tools[0].input_schema["properties"] @pytest.mark.asyncio @@ -129,10 +167,10 @@ async def test_call_tool_runs_agent_end_to_end(): agent = _EchoAgent(name="assistant") server = to_mcp_server(agent) - async with create_connected_server_and_client_session(server) as client: + async with _connected_client_session(server) as client: result = await client.call_tool("assistant", {"request": "hi"}) - assert not result.isError + assert not result.is_error assert "hello from the agent" in result.content[0].text @@ -174,7 +212,7 @@ async def test_run_agent_maps_image_output_to_image_content(): assert len(result) == 1 assert result[0].type == "image" - assert result[0].mimeType == "image/png" + assert result[0].mime_type == "image/png" assert base64.b64decode(result[0].data) == png @@ -209,7 +247,7 @@ async def test_call_tool_reuses_session_across_calls_on_one_connection(): runner = _FakeRunner([_text_event("ok")]) server = to_mcp_server(agent, runner=runner) - async with create_connected_server_and_client_session(server) as client: + async with _connected_client_session(server) as client: await client.call_tool("assistant", {"request": "first"}) await client.call_tool("assistant", {"request": "second"}) diff --git a/tests/unittests/tools/mcp_tool/test_conversion_utils.py b/tests/unittests/tools/mcp_tool/test_conversion_utils.py index d37c7546a7..34a5c2dc34 100644 --- a/tests/unittests/tools/mcp_tool/test_conversion_utils.py +++ b/tests/unittests/tools/mcp_tool/test_conversion_utils.py @@ -39,7 +39,7 @@ def test_tool_with_no_declaration(self): assert isinstance(result, mcp_types.Tool) assert result.name == "test_tool" assert result.description == "Test tool" - assert result.inputSchema == {} + assert result.input_schema == {} def test_tool_with_parameters_schema(self): """Test conversion when tool has parameters Schema object.""" @@ -72,14 +72,14 @@ def test_tool_with_parameters_schema(self): assert isinstance(result, mcp_types.Tool) assert result.name == "get_weather" assert result.description == "Gets weather information" - assert "type" in result.inputSchema - assert result.inputSchema["type"] == "object" - assert "properties" in result.inputSchema - assert "location" in result.inputSchema["properties"] - assert "units" in result.inputSchema["properties"] - assert result.inputSchema["properties"]["location"]["type"] == "string" - assert "required" in result.inputSchema - assert "location" in result.inputSchema["required"] + assert "type" in result.input_schema + assert result.input_schema["type"] == "object" + assert "properties" in result.input_schema + assert "location" in result.input_schema["properties"] + assert "units" in result.input_schema["properties"] + assert result.input_schema["properties"]["location"]["type"] == "string" + assert "required" in result.input_schema + assert "location" in result.input_schema["required"] def test_tool_with_parameters_json_schema(self): """Test conversion when tool has parameters_json_schema.""" @@ -115,7 +115,7 @@ def test_tool_with_parameters_json_schema(self): assert result.name == "search_database" assert result.description == "Searches a database" # Should use the JSON schema directly - assert result.inputSchema == json_schema + assert result.input_schema == json_schema def test_tool_with_no_parameters(self): """Test conversion when tool has declaration but no parameters.""" @@ -134,7 +134,7 @@ def test_tool_with_no_parameters(self): assert isinstance(result, mcp_types.Tool) assert result.name == "get_current_time" assert result.description == "Gets the current time" - assert not result.inputSchema + assert not result.input_schema def test_tool_prefers_json_schema_over_parameters(self): """Test that parameters_json_schema is preferred over parameters.""" @@ -166,9 +166,9 @@ def test_tool_prefers_json_schema_over_parameters(self): result = adk_to_mcp_tool_type(mock_tool) # Should use parameters_json_schema, not parameters - assert result.inputSchema == json_schema - assert "json_param" in result.inputSchema["properties"] - assert "schema_param" not in result.inputSchema["properties"] + assert result.input_schema == json_schema + assert "json_param" in result.input_schema["properties"] + assert "schema_param" not in result.input_schema["properties"] def test_tool_with_complex_nested_schema(self): """Test conversion with complex nested parameters_json_schema.""" @@ -206,4 +206,4 @@ def test_tool_with_complex_nested_schema(self): result = adk_to_mcp_tool_type(mock_tool) assert isinstance(result, mcp_types.Tool) - assert result.inputSchema == json_schema + assert result.input_schema == json_schema diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index 916f7b52ef..2ad8fb85a6 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -34,6 +34,7 @@ from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams import httpx +import httpx2 from mcp import StdioServerParameters import pytest @@ -244,7 +245,8 @@ def test_init_with_streamable_http_default_httpx_factory( kwargs = mock_streamable_http_client.call_args.kwargs assert kwargs["url"] == "https://example.com/mcp" assert kwargs["terminate_on_close"] is True - assert isinstance(kwargs["http_client"], httpx.AsyncClient) + # mcp 2.0's default create_mcp_http_client builds a vendored httpx2 client. + assert isinstance(kwargs["http_client"], httpx2.AsyncClient) @patch( "google.adk.tools.mcp_tool.mcp_session_manager.HTTPXClientInstrumentor", diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index fefd04f190..e791f16378 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -49,13 +49,13 @@ def __init__( self, name="test_tool", description="Test tool description", - outputSchema=None, + output_schema=None, meta=None, ): self.name = name self.description = description self.meta = meta - self.inputSchema = { + self.input_schema = { "type": "object", "properties": { "param1": {"type": "string", "description": "First parameter"}, @@ -63,7 +63,7 @@ def __init__( }, "required": ["param1"], } - self.outputSchema = outputSchema + self.output_schema = output_schema class TestMCPToolLegacy: @@ -152,7 +152,7 @@ def test_get_declaration_with_output_schema_and_json_schema_for_func_decl_enable } tool = MCPTool( - mcp_tool=MockMCPTool(outputSchema=output_schema), + mcp_tool=MockMCPTool(output_schema=output_schema), mcp_session_manager=self.mock_session_manager, ) @@ -170,7 +170,7 @@ def test_get_declaration_with_empty_output_schema_and_json_schema_for_func_decl_ ): """Test function declaration with an empty output schema and json schema for func decl enabled.""" tool = MCPTool( - mcp_tool=MockMCPTool(outputSchema={}), + mcp_tool=MockMCPTool(output_schema={}), mcp_session_manager=self.mock_session_manager, ) @@ -1384,10 +1384,9 @@ async def mock_call_tool(*args, **kwargs): async def test_run_async_captures_http_debug_info_on_graceful_error( self, mock_is_enabled ): - """Test that run_async captures HTTP debug info when tool call fails gracefully with McpError.""" + """Test that run_async captures HTTP debug info when tool call fails gracefully with MCPError.""" from google.adk.tools.mcp_tool.mcp_session_manager import _http_debug_var - from mcp.shared.exceptions import McpError - from mcp.types import ErrorData + from mcp.shared.exceptions import MCPError tool = MCPTool( mcp_tool=self.mock_mcp_tool, @@ -1400,7 +1399,7 @@ async def mock_call_tool(*args, **kwargs): debug_list.append( {"url": "https://example.com/api", "status_code": 403} ) - raise McpError(ErrorData(code=-32000, message="Forbidden")) + raise MCPError(code=-32000, message="Forbidden") self.mock_session.call_tool = mock_call_tool @@ -1447,17 +1446,19 @@ def setup_method(self): @pytest.mark.asyncio async def test_run_async_returns_dict_on_mcp_error_when_flag_on(self): - """When the flag is on, McpError surfaces as `{"error": "..."}`.""" - from mcp.shared.exceptions import McpError - from mcp.types import ErrorData + """When the flag is on, MCPError surfaces as `{"error": "..."}`.""" + from mcp.shared.exceptions import MCPError tool = MCPTool( mcp_tool=self.mock_mcp_tool, mcp_session_manager=self.mock_session_manager, ) - error_data = ErrorData(code=-32000, message="Client error '403 Forbidden'") - tool._run_async_impl = AsyncMock(side_effect=McpError(error_data)) + tool._run_async_impl = AsyncMock( + side_effect=MCPError( + code=-32000, message="Client error '403 Forbidden'" + ) + ) tool_context = Mock(spec=ToolContext) args = {"param1": "test_value"} @@ -1507,16 +1508,18 @@ async def test_run_async_propagates_mcp_error_when_flag_off(self): This protects downstream consumers that haven't migrated yet from a silent behavior change. """ - from mcp.shared.exceptions import McpError - from mcp.types import ErrorData + from mcp.shared.exceptions import MCPError tool = MCPTool( mcp_tool=self.mock_mcp_tool, mcp_session_manager=self.mock_session_manager, ) - error_data = ErrorData(code=-32000, message="Client error '403 Forbidden'") - tool._run_async_impl = AsyncMock(side_effect=McpError(error_data)) + tool._run_async_impl = AsyncMock( + side_effect=MCPError( + code=-32000, message="Client error '403 Forbidden'" + ) + ) tool_context = Mock(spec=ToolContext) args = {"param1": "test_value"} @@ -1524,7 +1527,7 @@ async def test_run_async_propagates_mcp_error_when_flag_off(self): with temporary_feature_override( FeatureName._MCP_GRACEFUL_ERROR_HANDLING, False ): - with pytest.raises(McpError): + with pytest.raises(MCPError): await tool.run_async(args=args, tool_context=tool_context) @pytest.mark.asyncio diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index ceff08918a..d262f98b31 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -58,7 +58,7 @@ class MockMCPTool: def __init__(self, name, description="Test tool description"): self.name = name self.description = description - self.inputSchema = { + self.input_schema = { "type": "object", "properties": {"param": {"type": "string"}}, } @@ -737,11 +737,11 @@ async def test_read_resource(self, name, mime_type, content, encoding): # Mock read_resource if encoding == "base64": contents = [ - BlobResourceContents(uri=uri, mimeType=mime_type, blob=content) + BlobResourceContents(uri=uri, mime_type=mime_type, blob=content) ] else: contents = [ - TextResourceContents(uri=uri, mimeType=mime_type, text=content) + TextResourceContents(uri=uri, mime_type=mime_type, text=content) ] read_resource_result = ReadResourceResult(contents=contents) diff --git a/tests/unittests/tools/mcp_tool/test_session_context.py b/tests/unittests/tools/mcp_tool/test_session_context.py index bc3391f65e..fbe5517c9b 100644 --- a/tests/unittests/tools/mcp_tool/test_session_context.py +++ b/tests/unittests/tools/mcp_tool/test_session_context.py @@ -16,7 +16,6 @@ import asyncio from contextlib import AsyncExitStack -from datetime import timedelta from unittest.mock import AsyncMock from unittest.mock import Mock from unittest.mock import patch @@ -420,7 +419,7 @@ async def test_stdio_client_with_read_timeout(self): # Verify ClientSession was called with read_timeout_seconds for stdio call_args = mock_session_class.call_args assert 'read_timeout_seconds' in call_args.kwargs - assert call_args.kwargs['read_timeout_seconds'] == timedelta(seconds=5.0) + assert call_args.kwargs['read_timeout_seconds'] == 5.0 await session_context.close() @@ -469,9 +468,7 @@ async def test_sse_read_timeout_passed_to_client_session(self): # Verify ClientSession was called with sse_read_timeout call_args = mock_session_class.call_args assert 'read_timeout_seconds' in call_args.kwargs - assert call_args.kwargs['read_timeout_seconds'] == timedelta( - seconds=300.0 - ) + assert call_args.kwargs['read_timeout_seconds'] == 300.0 await session_context.close()