feat: #479 support for streamable http MCP servers for MCPToolset

Author: omarbenhamid

pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
   "google-cloud-storage>=2.18.0, <3.0.0",    # For GCS Artifact service
   "google-genai>=1.12.1",                    # Google GenAI SDK
   "graphviz>=0.20.2",                        # Graphviz for graph rendering
-  "mcp>=1.5.0;python_version>='3.10'",       # For MCP Toolset
+  "mcp>=1.8.0;python_version>='3.10'",       # For MCP Toolset
   "opentelemetry-api>=1.31.0",               # OpenTelemetry
   "opentelemetry-exporter-gcp-trace>=1.9.0",
   "opentelemetry-sdk>=1.31.0",

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

@@ -15,6 +15,7 @@
 from contextlib import AsyncExitStack
 import functools
 import sys
+from datetime import timedelta
 from typing import Any, TextIO
 import anyio
 from pydantic import BaseModel
@@ -23,6 +24,7 @@
   from mcp import ClientSession, 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
 
@@ -48,6 +50,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.
 
@@ -117,7 +132,7 @@ class MCPSessionManager:
 
   def __init__(
       self,
-      connection_params: StdioServerParameters | SseServerParams,
+      connection_params: StdioServerParameters | SseServerParams | StreamableHTTPServerParams,
       exit_stack: AsyncExitStack,
       errlog: TextIO = sys.stderr,
   ) -> ClientSession:
@@ -153,7 +168,7 @@ async def create_session(self) -> ClientSession:
   async def initialize_session(
       cls,
       *,
-      connection_params: StdioServerParameters | SseServerParams,
+      connection_params: StdioServerParameters | SseServerParams | StreamableHTTPServerParams,
       exit_stack: AsyncExitStack,
       errlog: TextIO = sys.stderr,
   ) -> ClientSession:
@@ -177,6 +192,14 @@ async def initialize_session(
           timeout=connection_params.timeout,
           sse_read_timeout=connection_params.sse_read_timeout,
       )
+    elif isinstance(connection_params, StreamableHTTPServerParams):
+      client = streamablehttp_client(
+          url=connection_params.url,
+          headers=connection_params.headers,
+          timeout=timedelta(seconds=connection_params.timeout),
+          sse_read_timeout=timedelta(seconds=connection_params.sse_read_timeout),
+          terminate_on_close=connection_params.terminate_on_close,
+      )
     else:
       raise ValueError(
           'Unable to initialize connection. Connection should be'
@@ -185,6 +208,6 @@ async def initialize_session(
       )
 
     transports = await exit_stack.enter_async_context(client)
-    session = await exit_stack.enter_async_context(ClientSession(*transports))
+    session = await exit_stack.enter_async_context(ClientSession(*transports[:2]))
     await session.initialize()
     return session