diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5cec763bb5d..1e96747882d 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -581,20 +581,21 @@ class AsyncHTTPHandler: headers: Optional[dict] = None, follow_redirects: Optional[bool] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + stream: bool = False, ): - # Set follow_redirects to UseClientDefault if None _follow_redirects = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT params = params or {} params.update(HTTPHandler.extract_query_params(url)) - response = await self.client.get( + req = self.client.build_request( + "GET", url, params=params, - headers=headers, # type: ignore - follow_redirects=_follow_redirects, # type: ignore + headers=headers, timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT, ) + response = await self.client.send(req, stream=stream, follow_redirects=_follow_redirects) # type: ignore[arg-type] # UseClientDefault sentinel not in httpx stubs return response @track_llm_api_timing() diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1d681b43b9e..c17e83afcb9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2922,15 +2922,17 @@ class MCPServerManager: """ try: - client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - response = await client.get(server_url) + handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) + response = await handler.get(server_url, stream=True) + status_code = response.status_code response.raise_for_status() + await response.aclose() ( authorization_servers, resource_scopes, ) = await self._attempt_well_known_discovery(server_url) metadata = await self._fetch_authorization_server_metadata(authorization_servers, server_url) - if metadata is None and not resource_scopes and authorization_servers and response.status_code == 200: + if metadata is None and not resource_scopes and authorization_servers and status_code == 200: verbose_logger.warning( "MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.", server_url, diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 7bd1d7a6031..b9cc66b9a76 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -904,3 +904,45 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout: litellm.in_memory_llm_clients_cache = LLMClientCache() client = get_async_httpx_client(llm_provider=LlmProviders.BEDROCK) assert client.timeout.read == 300.0 + + +@pytest.mark.asyncio +async def test_async_http_handler_get_stream_does_not_buffer_body(): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + handler = AsyncHTTPHandler() + + streaming_response = MagicMock(spec=httpx.Response) + streaming_response.status_code = 200 + + mock_send = AsyncMock(return_value=streaming_response) + + with patch.object(handler.client, "send", mock_send): + result = await handler.get("https://example.com/mcp", stream=True) + + _, send_kwargs = mock_send.call_args + assert send_kwargs["stream"] is True + assert result is streaming_response + + +@pytest.mark.asyncio +async def test_async_http_handler_get_non_stream_buffers_body(): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + handler = AsyncHTTPHandler() + + buffered_response = MagicMock(spec=httpx.Response) + buffered_response.status_code = 200 + + mock_send = AsyncMock(return_value=buffered_response) + + with patch.object(handler.client, "send", mock_send): + result = await handler.get("https://example.com/mcp", stream=False) + + _, send_kwargs = mock_send.call_args + assert send_kwargs["stream"] is False + assert result is buffered_response diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7c55bd4560f..2e030d54245 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -7097,6 +7097,33 @@ class TestOBOEndpointDiscovery: mock_fetch_auth.assert_not_awaited() assert result is None or result.token_url is None + @pytest.mark.asyncio + async def test_descovery_metadata_returns_none_for_no_auth_sse_server(self): + manager = MCPServerManager() + server_url = "https://example.com/mcp" + + streaming_response = MagicMock() + streaming_response.status_code = 200 + streaming_response.raise_for_status = MagicMock() + streaming_response.aclose = AsyncMock() + + mock_handler = MagicMock() + mock_handler.get = AsyncMock(return_value=streaming_response) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_handler, + ), + patch.object(manager, "_attempt_well_known_discovery", AsyncMock(return_value=([], None))), + patch.object(manager, "_fetch_authorization_server_metadata", AsyncMock(return_value=None)), + ): + result = await manager._descovery_metadata(server_url) + + mock_handler.get.assert_awaited_once_with(server_url, stream=True) + streaming_response.aclose.assert_awaited_once() + assert result is None + if __name__ == "__main__": pytest.main([__file__])