From 51cf782292f8b4a23b05abdd536ab6507ddc0c16 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Tue, 20 Jan 2026 07:37:50 +0900 Subject: [PATCH] chore: switch experimental client to streamable_http_client API --- litellm/experimental_mcp_client/client.py | 19 +++++--- tests/mcp_tests/test_mcp_client_unit.py | 11 ++--- .../test_mcp_client.py | 47 ++++++++----------- 3 files changed, 36 insertions(+), 41 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 943cc6b2d5..582d044023 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/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]]): """ diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index 451c377437..301234ec98 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/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.""" diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 14b747a570..a03090d9c5 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/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()