From 3e1157080a01ad2375a2be30bdd6954b5d1eaa2f Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Sat, 28 Mar 2026 18:20:49 -0700
Subject: [PATCH] fix(tools): use non-deprecated streamable_http_client for MCP
 HTTP transport (#3646)

Switch MCP HTTP transport from the deprecated streamablehttp_client()
(mcp < 1.24.0) to the new streamable_http_client() API that accepts a
pre-built httpx.AsyncClient.

Changes vs the original PR #3391:
- Separate try/except imports so mcp < 1.24.0 doesn't break (graceful
  fallback to deprecated API instead of losing HTTP MCP entirely)
- Wrap httpx.AsyncClient in async-with for proper lifecycle management
  (the new SDK API explicitly skips closing caller-provided clients)
- Match SDK's own create_mcp_http_client defaults: follow_redirects=True,
  Timeout(connect_timeout, read=300.0)
- Keep deprecated code path as fallback for older SDK versions

Co-authored-by: HenkDz <HenkDz@users.noreply.github.com>
---
 tools/mcp_tool.py | 66 ++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 51 insertions(+), 15 deletions(-)

diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py
index 2b68ff4bf..5ce1ee192 100644
--- a/tools/mcp_tool.py
+++ b/tools/mcp_tool.py
@@ -98,6 +98,13 @@ try:
         _MCP_HTTP_AVAILABLE = True
     except ImportError:
         _MCP_HTTP_AVAILABLE = False
+    # Prefer the non-deprecated API (mcp >= 1.24.0); fall back to the
+    # deprecated wrapper for older SDK versions.
+    try:
+        from mcp.client.streamable_http import streamable_http_client
+        _MCP_NEW_HTTP = True
+    except ImportError:
+        _MCP_NEW_HTTP = False
     # Sampling types -- separated so older SDK versions don't break MCP support
     try:
         from mcp.types import (
@@ -762,21 +769,50 @@ class MCPServerTask:
                 logger.warning("MCP OAuth setup failed for '%s': %s", self.name, exc)
 
         sampling_kwargs = self._sampling.session_kwargs() if self._sampling else {}
-        _http_kwargs: dict = {
-            "headers": headers,
-            "timeout": float(connect_timeout),
-        }
-        if _oauth_auth is not None:
-            _http_kwargs["auth"] = _oauth_auth
-        async with streamablehttp_client(url, **_http_kwargs) as (
-            read_stream, write_stream, _get_session_id,
-        ):
-            async with ClientSession(read_stream, write_stream, **sampling_kwargs) as session:
-                await session.initialize()
-                self.session = session
-                await self._discover_tools()
-                self._ready.set()
-                await self._shutdown_event.wait()
+
+        if _MCP_NEW_HTTP:
+            # New API (mcp >= 1.24.0): build an explicit httpx.AsyncClient
+            # matching the SDK's own create_mcp_http_client defaults.
+            import httpx
+
+            client_kwargs: dict = {
+                "follow_redirects": True,
+                "timeout": httpx.Timeout(float(connect_timeout), read=300.0),
+            }
+            if headers:
+                client_kwargs["headers"] = headers
+            if _oauth_auth is not None:
+                client_kwargs["auth"] = _oauth_auth
+
+            # Caller owns the client lifecycle — the SDK skips cleanup when
+            # http_client is provided, so we wrap in async-with.
+            async with httpx.AsyncClient(**client_kwargs) as http_client:
+                async with streamable_http_client(url, http_client=http_client) as (
+                    read_stream, write_stream, _get_session_id,
+                ):
+                    async with ClientSession(read_stream, write_stream, **sampling_kwargs) as session:
+                        await session.initialize()
+                        self.session = session
+                        await self._discover_tools()
+                        self._ready.set()
+                        await self._shutdown_event.wait()
+        else:
+            # Deprecated API (mcp < 1.24.0): manages httpx client internally.
+            _http_kwargs: dict = {
+                "headers": headers,
+                "timeout": float(connect_timeout),
+            }
+            if _oauth_auth is not None:
+                _http_kwargs["auth"] = _oauth_auth
+            async with streamablehttp_client(url, **_http_kwargs) as (
+                read_stream, write_stream, _get_session_id,
+            ):
+                async with ClientSession(read_stream, write_stream, **sampling_kwargs) as session:
+                    await session.initialize()
+                    self.session = session
+                    await self._discover_tools()
+                    self._ready.set()
+                    await self._shutdown_event.wait()
 
     async def _discover_tools(self):
         """Discover tools from the connected session."""