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 af25b0e919a..17d3098bf4c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -25,6 +25,7 @@ from collections.abc import ( ) from contextlib import asynccontextmanager from dataclasses import dataclass, replace +from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast from urllib.parse import ParseResult, urlparse @@ -883,6 +884,53 @@ def _sanitized_error_text(exc: Exception) -> str: return re.sub(r"https?://\S+", "", str(exc))[:200] +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.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" + try: + await asyncio.wait_for(load_openapi_spec_async(spec_path, max_bytes=10 * 1024 * 1024), timeout=timeout) + except asyncio.TimeoutError: + return "unhealthy", f"OpenAPI specification check timed out after {timeout} seconds" + except HTTPStatusError as exc: + return "unhealthy", f"OpenAPI specification request failed (HTTP {exc.response.status_code})" + 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 + + +class _OpenAPIHealthProbe: + def __init__(self, spec_path: str, clock: Callable[[], float] = time.monotonic) -> None: + self.spec_path = spec_path + self.clock = clock + self.lock = asyncio.Lock() + self.checked_at = float("-inf") + self.result: tuple[Literal["healthy", "unhealthy", "unknown"], str | None, datetime.datetime] | None = None + + async def check(self) -> tuple[Literal["healthy", "unhealthy", "unknown"], str | None, datetime.datetime]: + async with self.lock: + if self.result is not None and self.clock() - self.checked_at < 30.0: + return self.result + try: + status, error = await _openapi_spec_health(self.spec_path, timeout=MCP_HEALTH_CHECK_TIMEOUT) + except asyncio.CancelledError: + return ( + "unknown", + "OpenAPI specification check was cancelled", + datetime.datetime.now(datetime.timezone.utc), + ) + self.result = (status, error, datetime.datetime.now(datetime.timezone.utc)) + self.checked_at = self.clock() + return self.result + + def _discovery_failure_leaves_needs_unresolved( *, needs_authorization_url: bool, @@ -1749,6 +1797,7 @@ class MCPServerManager: token_exchanger=build_token_exchanger(), ) self.registry: dict[str, MCPServer] = {} + self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe) self.config_mcp_servers: dict[str, MCPServer] = {} """ eg. @@ -6679,6 +6728,18 @@ class MCPServerManager: last_health_check=datetime.now(), ) + if server.spec_path: + spec_status, spec_error, spec_checked_at = await self._openapi_health_probes(server.spec_path).check() + return self._build_mcp_server_table(server).model_copy( + update=MappingProxyType( + { + "status": spec_status, + "health_check_error": spec_error, + "last_health_check": spec_checked_at, + } + ) + ) + status: Literal["healthy", "unhealthy", "unknown"] = "unknown" health_check_error = 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 16f58ef5b76..d115eb8b3c1 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -163,10 +163,14 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]: return asyncio.run(load_openapi_spec_async(filepath)) -async def load_openapi_spec_async(filepath: str) -> dict[str, Any]: +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, 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 adf985e9a21..a6e43686e2c 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 @@ -4473,6 +4473,116 @@ class TestMCPServerManager: assert len(result) == 1 assert result[0].name == "github_tool_1" + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.api_key, MCPAuth.oauth2]) + @pytest.mark.parametrize("is_byok", [False, True]) + @pytest.mark.parametrize("scheme", ["http", "https"]) + async def test_openapi_health_loads_spec_without_mcp_handshake(self, respx_mock, monkeypatch, auth_type, is_byok, scheme): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="openapi-health", + name="openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path=f"{scheme}://93.184.216.34/openapi.json", + auth_type=auth_type, + is_byok=is_byok, + authentication_token=None if is_byok else "shared-secret", + static_headers={"Authorization": "Bearer static-secret"}, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).respond(200, json={"openapi": "3.0.0", "paths": {}}) + result = await manager.health_check_server(server.server_id, mcp_auth_header="caller-secret") + assert result.status == "healthy" + assert result.health_check_error is None + assert result.last_health_check is not None + assert result.spec_path == server.spec_path + assert route.call_count == 1 + assert "authorization" not in route.calls[0].request.headers + assert "x-api-key" not in route.calls[0].request.headers + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token]) + @pytest.mark.parametrize("spec_path", ["/config/openapi.json", "relative/openapi.json"]) + async def test_openapi_local_spec_health_is_unknown(self, respx_mock, auth_type, spec_path): + manager = MCPServerManager() + server = MCPServer( + server_id="local-openapi-health", + name="local-openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path=spec_path, + auth_type=auth_type, + is_byok=True, + ) + manager.registry = {server.server_id: server} + result = await manager.health_check_server(server.server_id) + assert result.status == "unknown" + assert result.health_check_error == "OpenAPI servers have no protocol-level health probe" + assert result.last_health_check is not None + assert not respx_mock.calls + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("failure", "expected_status", "expected_error"), + [ + (httpx.Response(401, text="secret response content"), "unhealthy", "OpenAPI specification request failed (HTTP 401)"), + (httpx.Response(404), "unhealthy", "OpenAPI specification request failed (HTTP 404)"), + (httpx.Response(500), "unhealthy", "OpenAPI specification request failed (HTTP 500)"), + (httpx.ConnectError("secret network details"), "unhealthy", "OpenAPI specification could not be loaded (ConnectError)"), + (httpx.Response(200, text="secret invalid JSON body"), "unhealthy", "OpenAPI specification could not be loaded (JSONDecodeError)"), + ], + ) + async def test_openapi_health_reports_safe_failures(self, respx_mock, monkeypatch, failure, expected_status, expected_error): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="failed-openapi-health", + name="failed-openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path="https://93.184.216.34/key-secret?token=query-secret", + auth_type=MCPAuth.bearer_token, + is_byok=True, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).mock(side_effect=[failure]) + result = await manager.health_check_server(server.server_id) + assert result.status == expected_status + assert result.health_check_error == expected_error + assert result.last_health_check is not None + assert route.call_count == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("cancel", [False, True]) + async def test_openapi_health_timeout_and_cancellation_cleanup(self, respx_mock, monkeypatch, cancel): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_spec_health + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + started = asyncio.Event() + cancelled = asyncio.Event() + + async def slow_load(request): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + respx_mock.get("https://93.184.216.34/slow.json").mock(side_effect=slow_load) + task = asyncio.create_task(_openapi_spec_health("https://93.184.216.34/slow.json", timeout=0.1)) + await asyncio.wait_for(started.wait(), timeout=1) + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + status, error = await task + assert status == "unhealthy" + assert error == "OpenAPI specification check timed out after 0.1 seconds" + assert cancelled.is_set() + @pytest.mark.asyncio async def test_health_check_server_healthy(self): """Test health check for a healthy server""" @@ -12676,3 +12786,118 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Li assert "Credential=AKIDEXAMPLE/" in request.headers["Authorization"] finally: request_ctx.reset(token) + +@pytest.mark.asyncio +async def test_openapi_health_coalesces_concurrent_checks_and_reuses_results(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="coalesced", + name="coalesced", + transport=MCPTransport.http, + spec_path="https://93.184.216.34/coalesced.json", + auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + started = asyncio.Event() + release = asyncio.Event() + + async def serve(request): + started.set() + await release.wait() + return httpx.Response(200, json={"paths": {}}) + + route = respx_mock.get(server.spec_path).mock(side_effect=serve) + tasks = [asyncio.create_task(manager.health_check_server(server.server_id)) for _ in range(4)] + await asyncio.wait_for(started.wait(), timeout=1) + release.set() + results = await asyncio.gather(*tasks) + cached = await manager.health_check_server(server.server_id) + assert [result.status for result in results] == ["healthy"] * 4 + assert cached.status == "healthy" + assert {result.last_health_check for result in [*results, cached]} == {results[0].last_health_check} + assert route.call_count == 1 + + +@pytest.mark.asyncio +async def test_openapi_health_cache_expires_at_thirty_seconds(respx_mock, monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _OpenAPIHealthProbe + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + clock = iter([0.0, 29.0, 30.0, 30.0]) + probe = _OpenAPIHealthProbe("https://93.184.216.34/expiry.json", clock=clock.__next__) + route = respx_mock.get(probe.spec_path).mock( + side_effect=[ + httpx.Response(200, json={"paths": {}}), + httpx.Response(503), + ] + ) + first = await probe.check() + assert first[0] == "healthy" + assert await probe.check() == first + refreshed = await probe.check() + assert refreshed[0] == "unhealthy" + assert refreshed[1] == "OpenAPI specification request failed (HTTP 503)" + assert refreshed[2] >= first[2] + assert route.call_count == 2 + + +@pytest.mark.asyncio +async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="oversized", + name="oversized", + transport=MCPTransport.http, + spec_path="https://93.184.216.34/large.json", + auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).respond(200, headers={"content-length": str(12 * 1024 * 1024)}) + 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 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 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("already_waiting", [False, True]) +async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, monkeypatch, already_waiting): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="cancelled-cache", name="cancelled-cache", transport=MCPTransport.http, + spec_path="https://93.184.216.34/cancelled-cache.json", auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + started = asyncio.Event() + attempts = [] + + async def serve(request): + attempts.append(request.url) + if not started.is_set(): + started.set() + await asyncio.Event().wait() + return httpx.Response(200, json={"paths": {}}) + + route = respx_mock.get(server.spec_path).mock(side_effect=serve) + leader = asyncio.create_task(manager.health_check_server(server.server_id)) + await asyncio.wait_for(started.wait(), timeout=1) + follower = asyncio.create_task(manager.health_check_server(server.server_id)) if already_waiting else None + await asyncio.sleep(0) + leader.cancel() + cancelled = await leader + assert cancelled.status == "unknown" + assert cancelled.health_check_error == "OpenAPI specification check was cancelled" + recovered = await follower if follower is not None else await manager.health_check_server(server.server_id) + assert recovered.status == "healthy" + assert recovered.health_check_error is None + cached = await manager.health_check_server(server.server_id) + assert cached.last_health_check == recovered.last_health_check + assert cached.status == "healthy" + assert len(attempts) == 2 + 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 e59616e53c1..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 @@ -1378,3 +1378,83 @@ class TestUpstreamStatusIsClassified: assert exc.value.status_code == status_code assert secret_body not in str(exc.value) assert str(exc.value) == f"upstream returned HTTP {status_code}" + +class TestBoundedOpenAPISpecLoading: + @pytest.mark.asyncio + @pytest.mark.parametrize("max_bytes", [12, 13]) + async def test_exact_size_and_smaller_specs_load(self, respx_mock, monkeypatch, max_bytes): + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + route = respx_mock.get("https://93.184.216.34/spec.json").respond(200, content=b'{"paths":{}}') + assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=max_bytes) == {"paths": {}} + assert route.calls[0].request.headers["accept-encoding"] == "identity" + + @pytest.mark.asyncio + @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 ( + load_openapi_spec_async, + ) + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + closed = [] + + class UnreadableStream(httpx.AsyncByteStream): + async def __aiter__(self): + pytest.fail("Oversized or compressed response must not be consumed") + yield b"" + + async def aclose(self): + closed.append(True) + + respx_mock.get("https://93.184.216.34/spec.json").respond(200, headers=headers, stream=UnreadableStream()) + 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 ( + load_openapi_spec_async, + ) + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + consumed = [] + closed = [] + + class ChunkedStream(httpx.AsyncByteStream): + async def __aiter__(self): + for index in range(10): + consumed.append(index) + yield b"x" * 65536 + + async def aclose(self): + closed.append(True) + + respx_mock.get("https://93.184.216.34/spec.json").respond(200, stream=ChunkedStream()) + 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] + + @pytest.mark.asyncio + @pytest.mark.parametrize("target", ["https://93.184.216.35/final.json", "http://127.0.0.1/private.json"]) + async def test_bounded_spec_redirects_preserve_ssrf_protection(self, respx_mock, monkeypatch, target): + from litellm.litellm_core_utils.url_utils import SSRFError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + respx_mock.get("https://93.184.216.34/spec.json").respond(302, headers={"location": target}) + destination = respx_mock.get(target).respond(200, json={"paths": {}}) + if "127.0.0.1" in target: + with pytest.raises(SSRFError): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) + assert not destination.called + else: + assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}} + assert destination.call_count == 1