From 0536358bf0453ccdcaad667e6506d8203c0cba63 Mon Sep 17 00:00:00 2001 From: bunnysayzz Date: Mon, 14 Sep 2026 23:14:22 +0530 Subject: [PATCH] fix(mcp): yaml loader returns mapping-only, health stays unhealthy on garbage bodies --- .../mcp_server/mcp_server_manager.py | 2 +- .../mcp_server/openapi_to_mcp_generator.py | 59 +++++++++++++------ tests/mcp_tests/test_openapi_spec_path_url.py | 24 ++++++++ 3 files changed, 65 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index fb0c623473a..539c53b1998 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -905,7 +905,7 @@ async def _openapi_spec_health( 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: + except (httpx.RequestError, TypeError, 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 6931c79d144..0cf35711e60 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -9,6 +9,7 @@ import os import re from collections.abc import Mapping, Sequence from pathlib import PurePosixPath +from types import ModuleType from typing import Any, Final, TypedDict from urllib.parse import quote @@ -57,7 +58,7 @@ from litellm.proxy._experimental.mcp_server.tool_registry import ( from litellm.types.mcp import credential_redirect_hook, custom_credential_slot -def _import_yaml(): +def _import_yaml() -> ModuleType: """Import and return the yaml module, raising a clear error if missing.""" try: import yaml as _yaml @@ -69,6 +70,18 @@ def _import_yaml(): ) from None +def _load_yaml_mapping(text: str) -> dict[str, Any]: + """Parse YAML text, requiring a mapping at the document root.""" + yaml_mod = _import_yaml() + try: + parsed = yaml_mod.safe_load(text) + except yaml_mod.YAMLError as exc: + raise ValueError(f"Invalid YAML OpenAPI spec: {exc}") from exc + if not isinstance(parsed, dict): + raise TypeError("Invalid OpenAPI spec: expected a JSON/YAML mapping at the document root") + return parsed + + class _OpenAPIJSONSchema(TypedDict, total=False): properties: Mapping[str, object] type: ReadOnly[str] @@ -185,6 +198,23 @@ def _is_yaml_content(filepath: str, content_type: str | None = None) -> bool: return bool(content_type and "yaml" in content_type) +def _load_local_openapi_spec(filepath: str) -> dict[str, Any]: + """Read a local OpenAPI spec file, parsing YAML or JSON.""" + if not os.path.exists(filepath): + raise FileNotFoundError(f"OpenAPI spec not found at {filepath}") + with open(filepath, "r", encoding="utf-8") as f: + content = f.read() + if _is_yaml_content(filepath): + return _load_yaml_mapping(content) + try: + return json.loads(content) + except ValueError as json_exc: + try: + return _load_yaml_mapping(content) + except (TypeError, ValueError): + raise json_exc from None + + 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) @@ -197,29 +227,20 @@ async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None content_type = r.headers.get("content-type", "") if _is_yaml_content(filepath, content_type): - return _import_yaml().safe_load(r.text) + return _load_yaml_mapping(r.text) # Try JSON first; fall back to YAML for specs served without # proper Content-Type headers (common with raw GitHub URLs). try: return r.json() - except ValueError: - return _import_yaml().safe_load(r.text) + except ValueError as json_exc: + try: + return _load_yaml_mapping(r.text) + except (TypeError, ValueError): + raise json_exc from None - # fallback: local file - # Local filesystem path - if not os.path.exists(filepath): - raise FileNotFoundError(f"OpenAPI spec not found at {filepath}") - - if _is_yaml_content(filepath): - with open(filepath, "r", encoding="utf-8") as f: - return _import_yaml().safe_load(f) - - with open(filepath, "r", encoding="utf-8") as f: - try: - return json.load(f) - except ValueError: - f.seek(0) - return _import_yaml().safe_load(f) + # Local files go through a worker thread: the async path must not + # perform blocking disk I/O directly (ruff ASYNC230). + return await asyncio.to_thread(_load_local_openapi_spec, filepath) def get_base_url(spec: Mapping[str, Any], spec_path: str | None = None) -> str: diff --git a/tests/mcp_tests/test_openapi_spec_path_url.py b/tests/mcp_tests/test_openapi_spec_path_url.py index 3e91bde0333..5b231aac09d 100644 --- a/tests/mcp_tests/test_openapi_spec_path_url.py +++ b/tests/mcp_tests/test_openapi_spec_path_url.py @@ -173,3 +173,27 @@ def test_load_openapi_spec_url_yaml_fallback( spec = gen.load_openapi_spec(url) assert spec["info"]["title"] == "Firefly" + + +def test_load_openapi_spec_url_plain_text_body_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A 200 body that is neither JSON nor a YAML mapping must raise, not return a string.""" + url = "https://example.local/openapi.json" + + req = httpx.Request("GET", url) + resp = httpx.Response( + status_code=200, + content=b"secret invalid JSON body", + headers={"content-type": "text/plain"}, + request=req, + ) + + def fake_get_async_httpx_client(*args, **kwargs): + return _FakeAsyncHTTPHandler(resp, expected_url=url) + + monkeypatch.setattr(gen, "get_async_httpx_client", fake_get_async_httpx_client) + monkeypatch.setattr(gen, "async_safe_get", lambda client, url, **kw: client.get(url)) + + with pytest.raises(ValueError, match="Expecting value"): + gen.load_openapi_spec(url)