From 36dc5e0e6808e2a72c74020a72236c9b6b59c233 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 28 Oct 2026 13:34:31 -0400 Subject: [PATCH] fix(mcp): bound OAuth discovery when the resource GET opens an event stream A Streamable HTTP MCP server may answer GET on the resource URL with the optional server-to-client SSE stream and hold it open. OAuth discovery ran that GET buffered, so it waited for a body that never ends, and MCP_METADATA_TIMEOUT (httpx's per-read timeout) does not bound it because any keepalive resets the clock. Only the status line and any WWW-Authenticate header matter at this step, so open the GET with httpx's stream() context manager and close the response unread on every path. Each metadata fetch also gets its own MCP_METADATA_TIMEOUT bound, so startup discovery cannot hang boot on a URL that keeps sending bytes while a slow but healthy chain still resolves. The discovery tests now run a real AsyncHTTPHandler over httpx.MockTransport, and a new test holds an open event stream to check discovery neither reads nor leaks it. Fixes #37499 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../proxy/mcp_server/mcp_server_manager.py | 8 +- .../mcp_server/test_mcp_server_manager.py | 183 +++++++++++++----- 2 files changed, 138 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index ceac2b2a8e1..1b2612de5bf 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -5153,8 +5153,8 @@ class MCPServerManager: # internal MCP servers. Do not run them through user URL # validation, but also do not follow redirects because the # redirect target would not inherit the same-authority guarantee. - return await client.get(url, follow_redirects=False) - return await async_safe_get(client, url) + return await asyncio.wait_for(client.get(url, follow_redirects=False), timeout=MCP_METADATA_TIMEOUT) + return await asyncio.wait_for(async_safe_get(client, url), timeout=MCP_METADATA_TIMEOUT) async def _descovery_metadata( self, @@ -5201,8 +5201,8 @@ class MCPServerManager: llm_provider=httpxSpecialProvider.MCP, params={"timeout": MCP_METADATA_TIMEOUT}, ) - response: Final = await client.get(server_url) - response.raise_for_status() + async with client.client.stream("GET", server_url) as response: + response.raise_for_status() ( authorization_servers, resource_scopes, diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py index 13e4704bd90..4bba16cd411 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -12,7 +12,7 @@ import time from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from datetime import datetime from pathlib import Path -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Any, Dict, Final, Literal, Optional, TypeVar, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -515,6 +515,28 @@ def test_custom_admission_header_controls_authorization_forwarding( assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected +def _handler_answering(respond: Callable[[httpx.Request], httpx.Response]) -> AsyncHTTPHandler: + return AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + + +def _bearer_scope_read_challenge(request: httpx.Request) -> httpx.Response: + return httpx.Response(401, headers={"WWW-Authenticate": 'Bearer scope="read"'}) + + +class _OpenEventStream(httpx.AsyncByteStream): + def __init__(self) -> None: + self.read_started: Final = asyncio.Event() + self.closed: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + self.read_started.set() + await asyncio.Event().wait() + yield b": ping\n\n" + + async def aclose(self) -> None: + self.closed.set() + + class TestMCPServerManager: """Test MCP Server Manager stdio functionality""" @@ -4807,12 +4829,6 @@ class TestMCPServerManager: ): manager = MCPServerManager() - mock_response = MagicMock() - mock_response.raise_for_status = MagicMock() - - mock_client = MagicMock() - mock_client.get = AsyncMock(return_value=mock_response) - mock_metadata = MCPOAuthMetadata( scopes=None, authorization_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize", @@ -4823,7 +4839,7 @@ class TestMCPServerManager: with ( patch( "litellm.proxy.mcp_server.mcp_server_manager.get_async_httpx_client", - return_value=mock_client, + return_value=_handler_answering(lambda request: httpx.Response(200)), ) as get_client, patch.object( manager, @@ -4864,11 +4880,6 @@ class TestMCPServerManager: full capability list.""" manager = MCPServerManager() - mock_response = MagicMock() - mock_response.raise_for_status = MagicMock() - mock_client = MagicMock() - mock_client.get = AsyncMock(return_value=mock_response) - authorization_server_metadata = MCPOAuthMetadata( scopes=["as.read", "as.write"], authorization_url="https://idp.example.com/authorize", @@ -4878,7 +4889,7 @@ class TestMCPServerManager: with ( patch( "litellm.proxy.mcp_server.mcp_server_manager.get_async_httpx_client", - return_value=mock_client, + return_value=_handler_answering(lambda request: httpx.Response(200)), ), patch.object( manager, @@ -4896,6 +4907,107 @@ class TestMCPServerManager: assert result is not None assert result.scopes == ["resource.only"] + @pytest.mark.asyncio + async def test_descovery_metadata_closes_an_open_event_stream_unread(self): + body = _OpenEventStream() + + def respond(request: httpx.Request) -> httpx.Response: + if request.url.path == "/mcp": + return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, stream=body) + return httpx.Response(404) + + with patch( + "litellm.proxy.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=_handler_answering(respond), + ): + discovery = asyncio.create_task( + MCPServerManager()._discover_metadata_recording_attempts( + "https://stream.example.com/mcp", allow_origin_fallback=False + ) + ) + body_read = asyncio.create_task(body.read_started.wait()) + await asyncio.wait((discovery, body_read), return_when=asyncio.FIRST_COMPLETED) + body_read.cancel() + discovery.cancel() + + assert not body.read_started.is_set() + assert body.closed.is_set() + metadata, attempts = await discovery + assert metadata is None + assert "HTTP 200 (no RFC 9728 challenge)" in attempts[0] + + @pytest.mark.asyncio + @pytest.mark.parametrize("source", ["config", "database"]) + async def test_startup_discovery_moves_past_metadata_that_never_finishes( + self, source: str, monkeypatch: pytest.MonkeyPatch + ): + body: Final = _OpenEventStream() + documents: Final = MappingProxyType( + { + "/.well-known/oauth-protected-resource": {"authorization_servers": ["https://stuck.example.com"]}, + "/.well-known/oauth-authorization-server": { + "issuer": "https://stuck.example.com", + "authorization_endpoint": "https://stuck.example.com/authorize", + "token_endpoint": "https://stuck.example.com/token", + }, + } + ) + + def respond(request: httpx.Request) -> httpx.Response: + if request.url.path == "/mcp": + return httpx.Response(200) + if request.url.path == "/.well-known/oauth-protected-resource/mcp": + return httpx.Response(200, headers={"Content-Type": "application/json"}, stream=body) + document: Final = documents.get(request.url.path) + return httpx.Response(404) if document is None else httpx.Response(200, json=document) + + monkeypatch.setenv("LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP", "1") + manager: Final = MCPServerManager() + with ( + patch( + "litellm.proxy.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=_handler_answering(respond), + ), + patch("litellm.proxy.mcp_server.mcp_server_manager.MCP_METADATA_TIMEOUT", 0.05), + ): + if source == "config": + await asyncio.wait_for( + manager.load_servers_from_config( + { + "stuck": { + "url": "https://stuck.example.com/mcp", + "transport": "http", + "auth_type": "oauth2", + "oauth2_flow": "authorization_code", + "client_id": "qa-client", + } + } + ), + timeout=5, + ) + server = next(iter(manager.config_mcp_servers.values())) + else: + server = await asyncio.wait_for( + manager.build_mcp_server_from_table( + LiteLLM_MCPServerTable( + server_id="stuck", + alias="stuck", + url="https://stuck.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + created_at=datetime.now(), + updated_at=datetime.now(), + ), + credentials_are_encrypted=False, + ), + timeout=5, + ) + + assert body.read_started.is_set() + assert server.authorization_url == "https://stuck.example.com/authorize" + assert server.token_url == "https://stuck.example.com/token" + @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path( self, @@ -5081,21 +5193,6 @@ class TestMCPServerManager: manager = MCPServerManager() server_url = "https://example.com/public/mcp" - request = httpx.Request("GET", server_url) - response_obj = httpx.Response( - status_code=401, - request=request, - headers={"WWW-Authenticate": 'Bearer scope="read"'}, - ) - - def raise_http_error(): - raise httpx.HTTPStatusError("unauthorized", request=request, response=response_obj) - - response_obj.raise_for_status = MagicMock(side_effect=raise_http_error) - - mock_client = MagicMock() - mock_client.get = AsyncMock(return_value=response_obj) - mock_metadata = MCPOAuthMetadata( scopes=None, authorization_url="https://example.com/auth", @@ -5106,7 +5203,7 @@ class TestMCPServerManager: with ( patch( "litellm.proxy.mcp_server.mcp_server_manager.get_async_httpx_client", - return_value=mock_client, + return_value=_handler_answering(_bearer_scope_read_challenge), ), patch.object( manager, @@ -12662,22 +12759,11 @@ class TestOBOEndpointDiscovery: # resource origin is never assumed to be the IdP, so no token endpoint is invented. manager = MCPServerManager() server_url = "https://example.com/public/mcp" - request = httpx.Request("GET", server_url) - response_obj = httpx.Response( - status_code=401, request=request, headers={"WWW-Authenticate": 'Bearer scope="read"'} - ) - response_obj.raise_for_status = MagicMock( - side_effect=lambda: (_ for _ in ()).throw( - httpx.HTTPStatusError("unauthorized", request=request, response=response_obj) - ) - ) - mock_client = MagicMock() - mock_client.get = AsyncMock(return_value=response_obj) with ( patch( "litellm.proxy.mcp_server.mcp_server_manager.get_async_httpx_client", - return_value=mock_client, + return_value=_handler_answering(_bearer_scope_read_challenge), ), patch.object(manager, "_fetch_oauth_metadata_from_resource", AsyncMock(return_value=([], None))), patch.object(manager, "_attempt_well_known_discovery", AsyncMock(return_value=([], None))), @@ -13424,10 +13510,11 @@ class TestDiscoveryFailureLogging: Discovery failures used to die at debug level and the config-load path emitted no warning at all, so the only operator-facing signal was the bare 400 at /authorize.""" - def _connect_error_client(self, url: str) -> MagicMock: - client = MagicMock() - client.get = AsyncMock(side_effect=httpx.ConnectError(f"[Errno 8] nodename nor servname provided for {url}")) - return client + def _connect_error_client(self, url: str) -> AsyncHTTPHandler: + def refuse(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError(f"[Errno 8] nodename nor servname provided for {url}", request=request) + + return _handler_answering(refuse) @pytest.mark.asyncio async def test_descovery_metadata_warns_with_redacted_attempts_on_connect_error(self, caplog): @@ -13468,12 +13555,10 @@ class TestDiscoveryFailureLogging: async def test_descovery_metadata_attempt_trail_names_each_failed_step(self, caplog): manager = MCPServerManager() url = "https://real-host.example.com/mcp-typo" - client = MagicMock() - client.get = AsyncMock(return_value=httpx.Response(404, request=httpx.Request("GET", url))) with ( patch( "litellm.proxy.mcp_server.mcp_server_manager.get_async_httpx_client", - return_value=client, + return_value=_handler_answering(lambda request: httpx.Response(404)), ), caplog.at_level(logging.WARNING, logger="LiteLLM"), ):