feat: google/adk-python#479 support for streamable http MCP servers for MCPToolset

Copybara import of the project:

--
c5b9d49d7b6d858ff0a93bd690e6d653b7c32221 by Omar BENHAMID <omar.benhamid@smart-gts.com>:

feat: google/adk-python#479 support for streamable http MCP servers for MCPToolset

--
9431bc19e6538c1b814aba0b24ff564acf046075 by Omar BENHAMID <omar.benhamid@smart-gts.com>:

feat: google/adk-python#479 streamable http added to right package

--
8b4aabed45a6f0dc828beb61f12985dc7b14f3d0 by Omar BENHAMID <omar.benhamid@smart-gts.com>:

feat: google/adk-python#479 streamable http : review feedbacks + sample agent
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/650 from omarbenhamid:feature/mcp-streamable-http 625f028784c216401d45cb1b5d4d998535ebcb00
PiperOrigin-RevId: 764419586
This commit is contained in:
Omar BENHAMID
2025-05-28 14:32:42 -07:00
committed by Copybara-Service
parent c7ce987676
commit d232e6216d
7 changed files with 198 additions and 10 deletions

View File

@@ -17,10 +17,10 @@ from contextlib import AsyncExitStack
import functools
import logging
import sys
from datetime import timedelta
from typing import Any
from typing import Optional
from typing import TextIO
import anyio
from pydantic import BaseModel
@@ -29,6 +29,7 @@ try:
from mcp import StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamablehttp_client
except ImportError as e:
import sys
@@ -56,6 +57,19 @@ class SseServerParams(BaseModel):
sse_read_timeout: float = 60 * 5
class StreamableHTTPServerParams(BaseModel):
"""Parameters for the MCP SSE connection.
See MCP SSE Client documentation for more details.
https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/streamable_http.py
"""
url: str
headers: dict[str, Any] | None = None
timeout: float = 5
sse_read_timeout: float = 60 * 5
terminate_on_close: bool = True
def retry_on_closed_resource(async_reinit_func_name: str):
"""Decorator to automatically reinitialize session and retry action.
@@ -123,13 +137,13 @@ class MCPSessionManager:
def __init__(
self,
connection_params: StdioServerParameters | SseServerParams,
connection_params: StdioServerParameters | SseServerParams | StreamableHTTPServerParams,
errlog: TextIO = sys.stderr,
):
"""Initializes the MCP session manager.
Args:
connection_params: Parameters for the MCP connection (Stdio or SSE).
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.
"""
@@ -163,6 +177,14 @@ class MCPSessionManager:
timeout=self._connection_params.timeout,
sse_read_timeout=self._connection_params.sse_read_timeout,
)
elif isinstance(self._connection_params, StreamableHTTPServerParams):
client = streamablehttp_client(
url=self._connection_params.url,
headers=self._connection_params.headers,
timeout=timedelta(seconds=self._connection_params.timeout),
sse_read_timeout=timedelta(seconds=self._connection_params.sse_read_timeout),
terminate_on_close=self._connection_params.terminate_on_close,
)
else:
raise ValueError(
'Unable to initialize connection. Connection should be'
@@ -171,8 +193,10 @@ 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)
ClientSession(*transports[:2])
)
await session.initialize()

View File

@@ -26,6 +26,7 @@ from ..base_toolset import ToolPredicate
from .mcp_session_manager import MCPSessionManager
from .mcp_session_manager import retry_on_closed_resource
from .mcp_session_manager import SseServerParams
from .mcp_session_manager import StreamableHTTPServerParams
# Attempt to import MCP Tool from the MCP library, and hints user to upgrade
# their Python version to 3.10 if it fails.
@@ -82,17 +83,18 @@ class MCPToolset(BaseToolset):
def __init__(
self,
*,
connection_params: StdioServerParameters | SseServerParams,
connection_params: StdioServerParameters | SseServerParams | StreamableHTTPServerParams,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
errlog: TextIO = sys.stderr,
):
"""Initializes the MCPToolset.
Args:
connection_params: The connection parameters to the MCP server. Can be:
`StdioServerParameters` for using local mcp server (e.g. using `npx` or
`python3`); or `SseServerParams` for a local/remote SSE server.
tool_filter: Optional filter to select specific tools. Can be either:
connection_params: The connection parameters to the MCP server. Can be:
`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.
@@ -110,6 +112,7 @@ class MCPToolset(BaseToolset):
connection_params=self._connection_params,
errlog=self._errlog,
)
self._session = None
@retry_on_closed_resource("_reinitialize_session")