From 45ef6684352e3c8082958bece8610df60048f4a3 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Fri, 30 May 2025 14:38:40 -0700 Subject: [PATCH] fix: timeout issues for mcpstdio server when mcp tools are incorrect. Fixes https://github.com/google/adk-python/issues/643 PiperOrigin-RevId: 765342572 --- .../adk/tools/mcp_tool/mcp_session_manager.py | 23 +++++++++++++++---- src/google/adk/tools/mcp_tool/mcp_toolset.py | 8 +++---- 2 files changed, 23 insertions(+), 8 deletions(-) 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 074fdcb..9130148 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -149,7 +149,8 @@ class MCPSessionManager: """Initializes the MCP session manager. Args: - connection_params: Parameters for the MCP connection (Stdio, SSE or Streamable HTTP). + connection_params: Parameters for the MCP connection (Stdio, SSE or + Streamable HTTP). errlog: (Optional) TextIO stream for error logging. Use only for initializing a local stdio MCP session. """ @@ -203,9 +204,23 @@ class MCPSessionManager: transports = await self._exit_stack.enter_async_context(client) # The streamable http client returns a GetSessionCallback in addition to the read/write MemoryObjectStreams # needed to build the ClientSession, we limit then to the two first values to be compatible with all clients. - session = await self._exit_stack.enter_async_context( - ClientSession(*transports[:2]) - ) + # The StdioServerParameters does not provide a timeout parameter for the + # session, so we need to set a default timeout for it. Other clients + # (SseServerParams and StreamableHTTPServerParams) already provide a + # timeout parameter in their configuration. + if isinstance(self._connection_params, StdioServerParameters): + # Default timeout for MCP session is 5 seconds, same as SseServerParams + # and StreamableHTTPServerParams. + session = await self._exit_stack.enter_async_context( + ClientSession( + *transports[:2], + read_timeout_seconds=timedelta(seconds=5), + ) + ) + else: + session = await self._exit_stack.enter_async_context( + ClientSession(*transports[:2]) + ) await session.initialize() self._session = session diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index f4837ec..d05bb8f 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -98,10 +98,10 @@ class MCPToolset(BaseToolset): `StdioServerParameters` for using local mcp server (e.g. using `npx` or `python3`); or `SseServerParams` for a local/remote SSE server; or `StreamableHTTPServerParams` for local/remote Streamable http server. - tool_filter: Optional filter to select specific tools. Can be either: - - A list of tool names to include - - A ToolPredicate function for custom filtering logic - errlog: TextIO stream for error logging. + tool_filter: Optional filter to select specific tools. Can be either: - A + list of tool names to include - A ToolPredicate function for custom + filtering logic + errlog: TextIO stream for error logging. """ super().__init__(tool_filter=tool_filter)