diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 12e515e19a8..f4883b57fbc 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/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, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 4177e1b16dc..d0b0c83a23c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/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 diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index ba72f7bea4d..d115eb8b3c1 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/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() 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 3d4ba264c1c..f8868cfaf83 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/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() 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 7eff19f3ed9..e3dd92485fc 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 @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index af0b188d22d..5fa202224e3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/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