chore: switch experimental client to streamable_http_client API

litellm/experimental_mcp_client/client.py

@@ -4,14 +4,13 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
 
 import asyncio
 import base64
-from datetime import timedelta
 from typing import Awaitable, Callable, Dict, List, Optional, TypeVar, Union
 
 import httpx
 from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
 from mcp.client.sse import sse_client
 from mcp.client.stdio import stdio_client
-from mcp.client.streamable_http import streamablehttp_client
+from mcp.client.streamable_http import streamable_http_client
 from mcp.types import (
     CallToolRequestParams as MCPCallToolRequestParams,
     GetPromptRequestParams,
@@ -80,6 +79,7 @@ class MCPClient:
     ) -> TSessionResult:
         """Open a session, run the provided coroutine, and clean up."""
         transport_ctx = None
+        http_client: Optional[httpx.AsyncClient] = None
 
         try:
             if self.transport_type == MCPTransport.stdio:
@@ -105,13 +105,15 @@ class MCPClient:
                 headers = self._get_auth_headers()
                 httpx_client_factory = self._create_httpx_client_factory()
                 verbose_logger.debug(
-                    "litellm headers for streamablehttp_client: %s", headers
+                    "litellm headers for streamable_http_client: %s", headers
                 )
-                transport_ctx = streamablehttp_client(
-                    url=self.server_url,
-                    timeout=timedelta(seconds=self.timeout),
+                http_client = httpx_client_factory(
                     headers=headers,
-                    httpx_client_factory=httpx_client_factory,
+                    timeout=httpx.Timeout(self.timeout),
+                )
+                transport_ctx = streamable_http_client(
+                    url=self.server_url,
+                    http_client=http_client,
                 )
 
             if transport_ctx is None:
@@ -128,6 +130,9 @@ class MCPClient:
                 "MCP client run_with_session failed for %s", self.server_url or "stdio"
             )
             raise
+        finally:
+            if http_client is not None:
+                await http_client.aclose()
 
     def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]):
         """

tests/mcp_tests/test_mcp_client_unit.py

@@ -82,7 +82,7 @@ class TestMCPClientUnitTests:
         assert headers == {}
 
     @pytest.mark.asyncio
-    @patch("litellm.experimental_mcp_client.client.streamablehttp_client")
+    @patch("litellm.experimental_mcp_client.client.streamable_http_client")
     @patch("litellm.experimental_mcp_client.client.ClientSession")
     async def test_run_with_session(self, mock_session_class, mock_transport):
         """Test run_with_session establishes session with auth headers."""
@@ -110,15 +110,14 @@ class TestMCPClientUnitTests:
 
         # Verify transport was created with auth headers
         call_args = mock_transport.call_args
-        assert call_args[1]["headers"] == {
-            "Authorization": "Bearer test_token",
-        }
+        http_client = call_args[1]["http_client"]
+        assert http_client.headers.get("Authorization") == "Bearer test_token"
 
         # Verify session was initialized
         mock_session_instance.initialize.assert_called_once()
 
     @pytest.mark.asyncio
-    @patch("litellm.experimental_mcp_client.client.streamablehttp_client")
+    @patch("litellm.experimental_mcp_client.client.streamable_http_client")
     @patch("litellm.experimental_mcp_client.client.ClientSession")
     async def test_list_tools(self, mock_session_class, mock_transport):
         """Test listing tools from the server."""
@@ -156,7 +155,7 @@ class TestMCPClientUnitTests:
         mock_session_instance.list_tools.assert_called_once()
 
     @pytest.mark.asyncio
-    @patch("litellm.experimental_mcp_client.client.streamablehttp_client")
+    @patch("litellm.experimental_mcp_client.client.streamable_http_client")
     @patch("litellm.experimental_mcp_client.client.ClientSession")
     async def test_call_tool(self, mock_session_class, mock_transport):
         """Test calling a tool."""

tests/test_litellm/experimental_mcp_client/test_mcp_client.py

@@ -81,7 +81,7 @@ class TestMCPClient:
         assert call_args.env == {"DEBUG": "1"}
 
     @pytest.mark.asyncio
-    @patch("litellm.experimental_mcp_client.client.streamablehttp_client")
+    @patch("litellm.experimental_mcp_client.client.streamable_http_client")
     @patch.dict(
         os.environ,
         {
@@ -90,12 +90,12 @@ class TestMCPClient:
         },
     )
     async def test_mcp_client_ssl_configuration_from_env(
-        self, mock_streamablehttp_client
+        self, mock_streamable_http_client
     ):
         """Test that MCP client uses SSL configuration from environment variables"""
         # Setup mocks
         mock_transport = (MagicMock(), MagicMock())
-        mock_streamablehttp_client.return_value.__aenter__ = AsyncMock(
+        mock_streamable_http_client.return_value.__aenter__ = AsyncMock(
             return_value=mock_transport
         )
 
@@ -121,23 +121,19 @@ class TestMCPClient:
             await client.run_with_session(_operation)
 
             # Verify streamablehttp_client was called
-            mock_streamablehttp_client.assert_called_once()
-            call_kwargs = mock_streamablehttp_client.call_args[1]
+            mock_streamable_http_client.assert_called_once()
+            call_kwargs = mock_streamable_http_client.call_args[1]
+            assert "http_client" in call_kwargs
+            http_client = call_kwargs["http_client"]
+            assert isinstance(http_client, httpx.AsyncClient)
 
-            # Verify httpx_client_factory was passed
-            assert "httpx_client_factory" in call_kwargs
-            httpx_factory = call_kwargs["httpx_client_factory"]
-
-            # Test the factory creates a client with proper SSL config
-            # When SSL_CERT_FILE is set, the factory should use get_ssl_configuration
+            # Test the factory still creates a client with proper SSL config
+            httpx_factory = client._create_httpx_client_factory()
             test_client = httpx_factory(headers={"test": "header"})
 
-            # Verify the client was created successfully with SSL configuration
             assert test_client is not None
             assert isinstance(test_client, httpx.AsyncClient)
-            # Verify it has the expected properties
             assert test_client.headers is not None
-            # Clean up
             await test_client.aclose()
 
     @pytest.mark.asyncio
@@ -192,12 +188,12 @@ class TestMCPClient:
             await test_client.aclose()
 
     @pytest.mark.asyncio
-    @patch("litellm.experimental_mcp_client.client.streamablehttp_client")
-    async def test_mcp_client_ssl_verify_custom_path(self, mock_streamablehttp_client):
+    @patch("litellm.experimental_mcp_client.client.streamable_http_client")
+    async def test_mcp_client_ssl_verify_custom_path(self, mock_streamable_http_client):
         """Test that MCP client uses custom CA bundle path from ssl_verify parameter"""
         # Setup mocks
         mock_transport = (MagicMock(), MagicMock())
-        mock_streamablehttp_client.return_value.__aenter__ = AsyncMock(
+        mock_streamable_http_client.return_value.__aenter__ = AsyncMock(
             return_value=mock_transport
         )
 
@@ -226,23 +222,18 @@ class TestMCPClient:
             await client.run_with_session(_operation)
 
             # Verify streamablehttp_client was called
-            mock_streamablehttp_client.assert_called_once()
-            call_kwargs = mock_streamablehttp_client.call_args[1]
+            mock_streamable_http_client.assert_called_once()
+            call_kwargs = mock_streamable_http_client.call_args[1]
+            assert "http_client" in call_kwargs
+            http_client = call_kwargs["http_client"]
+            assert isinstance(http_client, httpx.AsyncClient)
 
-            # Verify httpx_client_factory was passed
-            assert "httpx_client_factory" in call_kwargs
-            httpx_factory = call_kwargs["httpx_client_factory"]
-
-            # Test the factory creates a client with custom CA bundle path
-            # When ssl_verify is a path, the factory should use that path for SSL verification
+            httpx_factory = client._create_httpx_client_factory()
             test_client = httpx_factory(headers={"test": "header"})
 
-            # Verify the client was created successfully
             assert test_client is not None
             assert isinstance(test_client, httpx.AsyncClient)
-            # Verify it has the expected properties
             assert test_client.headers is not None
-            # Clean up
             await test_client.aclose()