refactor(mcp): reuse the shared HTTP handler for bounded probes

litellm/llms/custom_httpx/http_handler.py

@@ -9,6 +9,7 @@ import threading
 import time
 from collections.abc import AsyncIterable, Callable, Iterable, Mapping
 from http.cookiejar import CookieJar, DefaultCookiePolicy
+from io import BytesIO
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict, TypeVar
 
@@ -505,6 +506,10 @@ async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> N
     raise MaskedHTTPStatusError(e, message=_text, text=_text) from None
 
 
+class HTTPResponseLimitError(ValueError):
+    pass
+
+
 class MaskedHTTPStatusError(httpx.HTTPStatusError):
     def __init__(self, original_error, message: str | None = None, text: str | None = None):
         # Create a new error with the masked URL
@@ -654,6 +659,7 @@ class AsyncHTTPHandler:
         headers: dict | None = None,
         follow_redirects: bool | None = None,
         timeout: float | httpx.Timeout | None = None,
+        max_response_bytes: int | None = None,
     ):
         # Set follow_redirects to UseClientDefault if None
         _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT
@@ -661,6 +667,16 @@ class AsyncHTTPHandler:
         params = params or {}
         params.update(HTTPHandler.extract_query_params(url))
 
+        if max_response_bytes is not None:
+            return await self._get_with_response_limit(
+                url,
+                params=httpx.QueryParams(params),
+                headers=httpx.Headers(headers),
+                max_bytes=max_response_bytes,
+                follow_redirects=self.client.follow_redirects if follow_redirects is None else follow_redirects,
+                timeout=self.client.timeout if timeout is None else httpx.Timeout(timeout),
+            )
+
         response: Final = await self.client.get(
             url,
             params=params,
@@ -670,6 +686,57 @@ class AsyncHTTPHandler:
         )
         return response
 
+    async def _get_with_response_limit(
+        self,
+        url: str,
+        *,
+        params: httpx.QueryParams,
+        headers: httpx.Headers,
+        timeout: httpx.Timeout,
+        max_bytes: int,
+        follow_redirects: bool,
+    ) -> httpx.Response:
+        request: Final = self.client.build_request(
+            "GET",
+            url,
+            headers=MappingProxyType({**headers, "accept-encoding": "identity"}),
+            params=params,
+            timeout=timeout,
+        )
+        response: Final = await self.client.send(request, stream=True, follow_redirects=False)
+        return await self._read_with_response_limit(response, max_bytes=max_bytes, follow_redirects=follow_redirects)
+
+    async def _read_with_response_limit(
+        self, response: httpx.Response, *, max_bytes: int, follow_redirects: bool, redirects_remaining: int = 10
+    ) -> httpx.Response:
+        try:
+            if response.next_request is not None and follow_redirects:
+                if redirects_remaining == 0:
+                    raise ValueError("Too many redirects")
+                await response.aclose()
+                following: Final = await self.client.send(
+                    response.next_request, auth=None, stream=True, follow_redirects=False
+                )
+                return await self._read_with_response_limit(
+                    following, max_bytes=max_bytes, follow_redirects=True, redirects_remaining=redirects_remaining - 1
+                )
+            if response.is_redirect or response.is_error:
+                return httpx.Response(response.status_code, headers=response.headers, request=response.request)
+            if response.headers.get("content-encoding", "identity").lower() != "identity":
+                raise HTTPResponseLimitError("Response size limits require an uncompressed response")
+            if int(response.headers.get("content-length", "0")) > max_bytes:
+                raise HTTPResponseLimitError("Response exceeds the configured size limit")
+            with BytesIO() as body:
+                async for chunk in response.aiter_bytes(chunk_size=65536):
+                    if body.tell() + len(chunk) > max_bytes:
+                        raise HTTPResponseLimitError("Response exceeds the configured size limit")
+                    body.write(chunk)
+                return httpx.Response(
+                    response.status_code, headers=response.headers, content=body.getvalue(), request=response.request
+                )
+        finally:
+            await response.aclose()
+
     @track_llm_api_timing()
     async def post(
         self,

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -888,10 +888,8 @@ async def _openapi_spec_health(
     spec_path: str, *, timeout: float
 ) -> tuple[Literal["healthy", "unhealthy", "unknown"], str | None]:
     """Check specification availability, not upstream operations or user credentials."""
-    from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
-        OpenAPISpecProbeLimitError,
-        load_openapi_spec_async,
-    )
+    from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError
+    from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async
 
     if not spec_path.startswith(("http://", "https://")):
         return "unknown", "OpenAPI servers have no protocol-level health probe"
@@ -903,8 +901,8 @@ async def _openapi_spec_health(
         return "unknown", "OpenAPI specification check was cancelled"
     except HTTPStatusError as exc:
         return "unhealthy", f"OpenAPI specification request failed (HTTP {exc.response.status_code})"
-    except OpenAPISpecProbeLimitError as exc:
-        return "unknown", str(exc)
+    except HTTPResponseLimitError as exc:
+        return "unknown", f"OpenAPI specification probe refused: {exc}"
     except (httpx.RequestError, ValueError, OSError) as exc:
         return "unhealthy", f"OpenAPI specification could not be loaded ({type(exc).__name__})"
     return "healthy", None

litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py

@@ -8,10 +8,7 @@ import json
 import os
 import re
 from collections.abc import Mapping, Sequence
-from dataclasses import dataclass
-from io import BytesIO
 from pathlib import PurePosixPath
-from types import MappingProxyType
 from typing import Any, Final, TypedDict
 from urllib.parse import quote
 
@@ -166,60 +163,13 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]:
     return asyncio.run(load_openapi_spec_async(filepath))
 
 
-class OpenAPISpecProbeLimitError(ValueError):
-    pass
-
-
-@dataclass(frozen=True)
-class _BoundedOpenAPIFetcher:
-    client: httpx.AsyncClient
-    max_bytes: int
-
-    async def get(
-        self,
-        url: str,
-        *,
-        headers: dict[str, str] | None = None,
-        follow_redirects: bool = False,
-        redirects_remaining: int = 10,
-    ) -> httpx.Response:
-        async with self.client.stream(
-            "GET",
-            url,
-            headers=MappingProxyType({**(headers or MappingProxyType({})), "Accept-Encoding": "identity"}),
-            follow_redirects=False,
-        ) as response:
-            if response.is_redirect and follow_redirects:
-                if redirects_remaining == 0:
-                    raise ValueError("Too many specification redirects")
-                await response.aclose()
-                return await self.get(
-                    str(response.url.join(response.headers["location"])),
-                    headers=headers,
-                    follow_redirects=True,
-                    redirects_remaining=redirects_remaining - 1,
-                )
-            if response.is_redirect or response.is_error:
-                return httpx.Response(response.status_code, headers=response.headers, request=response.request)
-            if response.headers.get("content-encoding", "identity").lower() != "identity":
-                raise OpenAPISpecProbeLimitError("OpenAPI specification probe requires an uncompressed response")
-            if int(response.headers.get("content-length", "0")) > self.max_bytes:
-                raise OpenAPISpecProbeLimitError("OpenAPI specification exceeds the health-check size limit")
-            with BytesIO() as body:
-                async for chunk in response.aiter_bytes(chunk_size=65536):
-                    if body.tell() + len(chunk) > self.max_bytes:
-                        raise OpenAPISpecProbeLimitError("OpenAPI specification exceeds the health-check size limit")
-                    body.write(chunk)
-                return httpx.Response(
-                    response.status_code, headers=response.headers, content=body.getvalue(), request=response.request
-                )
-
-
 async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> dict[str, Any]:
     if filepath.startswith("http://") or filepath.startswith("https://"):
         client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
-        r: Final[httpx.Response] = await async_safe_get(
-            client if max_bytes is None else _BoundedOpenAPIFetcher(client.client, max_bytes), filepath
+        r: Final[httpx.Response] = (
+            await async_safe_get(client, filepath)
+            if max_bytes is None
+            else await async_safe_get(client, filepath, max_response_bytes=max_bytes)
         )
         r.raise_for_status()
         return r.json()

tests/test_litellm/llms/custom_httpx/test_http_handler.py

@@ -1612,3 +1612,66 @@ async def test_a_retried_put_stays_a_put_and_still_refuses_redirects():
         assert attempts == [("PUT", "/first"), ("PUT", "/first")]
     finally:
         await handler.client.aclose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("target", ["https://example.com/final.json?next=1", "https://other.example/final.json?next=1"])
+async def test_bounded_get_preserves_sdk_redirect_auth_and_query_handling(respx_mock, monkeypatch, target):
+    monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
+    respx_mock.get("https://example.com/spec.json?original=1").respond(302, headers={"location": target})
+    destination = respx_mock.get(target).respond(200, json={"paths": {}})
+    handler = AsyncHTTPHandler()
+    try:
+        response = await handler.get(
+            "https://example.com/spec.json?original=1", max_response_bytes=100, follow_redirects=True,
+            headers={"Authorization": "Bearer sentinel", "Accept-Encoding": "gzip"}, timeout=2.0,
+        )
+    finally:
+        await handler.close()
+    assert response.json() == {"paths": {}}
+    request = destination.calls[0].request
+    assert request.headers.get("authorization") == (None if "other.example" in target else "Bearer sentinel")
+    assert request.headers["accept-encoding"] == "identity"
+    assert str(request.url) == target
+    assert request.extensions["timeout"]["read"] == 2.0
+
+
+@pytest.mark.asyncio
+async def test_bounded_get_stops_redirect_loops(respx_mock, monkeypatch):
+    monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
+    route = respx_mock.get("https://example.com/spec.json").respond(302, headers={"location": "/spec.json"})
+    handler = AsyncHTTPHandler()
+    try:
+        with pytest.raises(ValueError, match="Too many redirects"):
+            await handler.get("https://example.com/spec.json", max_response_bytes=100, follow_redirects=True)
+    finally:
+        await handler.close()
+    assert route.call_count == 11
+
+
+@pytest.mark.asyncio
+async def test_bounded_get_closes_stream_on_cancellation(respx_mock, monkeypatch):
+    monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
+    started = asyncio.Event()
+    closed = asyncio.Event()
+
+    class SlowStream(httpx.AsyncByteStream):
+        async def __aiter__(self):
+            yield b"x"
+            started.set()
+            await asyncio.Event().wait()
+
+        async def aclose(self):
+            closed.set()
+
+    respx_mock.get("https://example.com/slow.json").respond(200, stream=SlowStream())
+    handler = AsyncHTTPHandler()
+    try:
+        task = asyncio.create_task(handler.get("https://example.com/slow.json", max_response_bytes=100))
+        await asyncio.wait_for(started.wait(), timeout=1)
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    finally:
+        await handler.close()
+    assert closed.is_set()

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -12858,7 +12858,7 @@ async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(r
     result = await manager.health_check_server(server.server_id)
     cached = await manager.health_check_server(server.server_id)
     assert result.status == "unknown"
-    assert result.health_check_error == "OpenAPI specification exceeds the health-check size limit"
+    assert result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit"
     assert cached.health_check_error == result.health_check_error
     assert cached.last_health_check == result.last_health_check
     assert route.call_count == 1

tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py

@@ -1394,8 +1394,8 @@ class TestBoundedOpenAPISpecLoading:
     @pytest.mark.parametrize("headers", [{"content-length": "1000000"}, {"content-encoding": "gzip"}])
     async def test_unsafe_response_headers_reject_before_reading(self, respx_mock, monkeypatch, headers):
         import httpx
+        from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError
         from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
-            OpenAPISpecProbeLimitError,
             load_openapi_spec_async,
         )
 
@@ -1411,15 +1411,15 @@ class TestBoundedOpenAPISpecLoading:
                 closed.append(True)
 
         respx_mock.get("https://93.184.216.34/spec.json").respond(200, headers=headers, stream=UnreadableStream())
-        with pytest.raises(OpenAPISpecProbeLimitError):
+        with pytest.raises(HTTPResponseLimitError):
             await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12)
         assert closed == [True]
 
     @pytest.mark.asyncio
     async def test_chunked_response_is_bounded_and_closed(self, respx_mock, monkeypatch):
         import httpx
+        from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError
         from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
-            OpenAPISpecProbeLimitError,
             load_openapi_spec_async,
         )
 
@@ -1437,7 +1437,7 @@ class TestBoundedOpenAPISpecLoading:
                 closed.append(True)
 
         respx_mock.get("https://93.184.216.34/spec.json").respond(200, stream=ChunkedStream())
-        with pytest.raises(OpenAPISpecProbeLimitError, match="size limit"):
+        with pytest.raises(HTTPResponseLimitError, match="size limit"):
             await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=65536)
         assert consumed == [0, 1]
         assert closed == [True]
@@ -1458,25 +1458,3 @@ class TestBoundedOpenAPISpecLoading:
         else:
             assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}}
             assert destination.call_count == 1
-
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("loop", [False, True])
-    async def test_redirects_are_bounded_when_validation_is_disabled(self, respx_mock, loop):
-        import httpx
-        from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import _BoundedOpenAPIFetcher
-
-        source = respx_mock.get("https://example.com/spec.json").respond(
-            302, headers={"location": "/spec.json" if loop else "/final.json"}
-        )
-        destination = respx_mock.get("https://example.com/final.json").respond(200, json={"paths": {}})
-        async with httpx.AsyncClient() as client:
-            fetcher = _BoundedOpenAPIFetcher(client, 100)
-            if loop:
-                with pytest.raises(ValueError, match="Too many specification redirects"):
-                    await fetcher.get("https://example.com/spec.json", follow_redirects=True)
-                assert source.call_count == 11
-                assert not destination.called
-            else:
-                response = await fetcher.get("https://example.com/spec.json", follow_redirects=True)
-                assert response.json() == {"paths": {}}
-                assert destination.call_count == 1