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>
This commit is contained in:
Teknium
2026-03-28 18:20:49 -07:00
committed by GitHub
parent 1a032ccf79
commit 3e1157080a

View File

@@ -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."""