diff --git a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py index 96d178dbeb1..0fa904cb582 100644 --- a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py @@ -5,15 +5,29 @@ This module is used to generate MCP tools from OpenAPI specs. import asyncio import contextvars import json +import math import os import re from collections.abc import Mapping, Sequence from pathlib import PurePosixPath -from typing import Final, Literal, Protocol, TypeAlias, TypedDict, overload +from typing import ( + ClassVar, + Final, + Literal, + Protocol, + TypeAlias, + TypedDict, + cast, # noqa: TID251 # existing OpenAPI Protocol fields have no runtime schema + overload, +) from urllib.parse import quote import httpx +import yaml +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, Required +from yaml.events import AliasEvent, ScalarEvent +from yaml.nodes import MappingNode, Node, ScalarNode from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError from litellm.proxy.mcp_server.exceptions import ( @@ -30,6 +44,18 @@ _OPENAPI_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") OPENAPI_TOOL_NAME_MAX_LEN: Final = 128 _OPENAPI_TOOL_NAME_MAX_LEN: Final = OPENAPI_TOOL_NAME_MAX_LEN +_OPENAPI_DOCUMENT: Final = TypeAdapter(dict[str, object]) +_MAX_YAML_INT_LENGTH: Final = 1024 +# YAML 1.2 core scalars; isolate these from PyYAML's global YAML 1.1 resolver. +_YAML_SCALARS: Final = { + "tag:yaml.org,2002:null": re.compile(r"^(?:null|Null|NULL|~|)$"), + "tag:yaml.org,2002:bool": re.compile(r"^(?:true|True|TRUE|false|False|FALSE)$"), + "tag:yaml.org,2002:int": re.compile(r"^(?:[-+]?[0-9]+|0o[0-7]+|0x[0-9a-fA-F]+)$"), + "tag:yaml.org,2002:float": re.compile( + r"^(?:[-+]?(?:\.[0-9]+|[0-9]+(?:\.[0-9]*)?)(?:[eE][-+]?[0-9]+)?" + r"|[-+]?\.(?:inf|Inf|INF)|\.(?:nan|NaN|NAN))$" + ), +} def sanitize_openapi_tool_name(raw_name: str) -> str: @@ -199,6 +225,70 @@ def load_openapi_spec(filepath: str) -> _OpenAPISpec: return asyncio.run(load_openapi_spec_async(filepath)) +class _OpenAPISafeLoader(yaml.SafeLoader): + yaml_implicit_resolvers: ClassVar[Mapping[str | None, Sequence[tuple[str, re.Pattern[str]]]]] = { # pyright: ignore[reportIncompatibleVariableOverride] # PyYAML resolver tables are class state + None: list(_YAML_SCALARS.items()) + } + + def compose_node(self, parent: Node | None, index: int) -> Node | None: + event: Final[object] = self.peek_event() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # incomplete PyYAML event stubs + if isinstance(event, AliasEvent): + raise ValueError("YAML aliases are not supported in OpenAPI specs") # noqa: TRY004 # callers handle malformed documents as ValueError + node: Final = super().compose_node(parent, index) + # PyYAML otherwise implicitly resolves even an explicit nonspecific ! tag. + if isinstance(event, ScalarEvent) and event.tag == "!" and isinstance(node, ScalarNode): + node.tag = "tag:yaml.org,2002:str" + return node + + def construct_object(self, node: Node, deep: bool = False) -> object: + pattern: Final = _YAML_SCALARS.get(node.tag) + if pattern is not None: + if not isinstance(node, ScalarNode): + raise ValueError("Invalid YAML OpenAPI spec") + scalar: Final[str] = node.value + if node.tag == "tag:yaml.org,2002:int" and len(scalar) > _MAX_YAML_INT_LENGTH: + raise ValueError("YAML integer is too long in OpenAPI spec") + if pattern.fullmatch(scalar) is None: + raise ValueError("Invalid YAML OpenAPI spec") + if node.tag == "tag:yaml.org,2002:int": + # YAML 1.2 decimal values with leading zeroes are not octal. + return int(scalar, 0 if scalar.startswith(("0o", "0x")) else 10) + if node.tag == "tag:yaml.org,2002:float": + value: Final = self.construct_yaml_float(node) + if not math.isfinite(value): + raise ValueError("Invalid YAML OpenAPI spec") + return value + elif node.tag not in {"tag:yaml.org,2002:str", "tag:yaml.org,2002:seq", "tag:yaml.org,2002:map"} or ( + node.tag == "tag:yaml.org,2002:str" and not isinstance(node, ScalarNode) + ): + raise ValueError("Invalid YAML OpenAPI spec") + if isinstance(node, MappingNode): + pairs: Final[Sequence[tuple[Node, Node]]] = node.value + if any(not isinstance(key, ScalarNode) or key.tag != "tag:yaml.org,2002:str" for key, _ in pairs): + raise ValueError("Invalid YAML OpenAPI spec") + return super().construct_object( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # incomplete PyYAML stubs + node, deep + ) + + +def _validate_openapi_spec(value: object) -> _OpenAPISpec: + try: + document: Final = _OPENAPI_DOCUMENT.validate_python(value) + except ValidationError as exc: + raise ValueError("OpenAPI spec must be a JSON or YAML object") from exc + return cast(_OpenAPISpec, document) # cast-ok: preserve existing lazy OpenAPI Protocol validation + + +def _parse_openapi_spec(content: str | bytes) -> _OpenAPISpec: + try: + return _validate_openapi_spec(json.loads(content)) + except json.JSONDecodeError: + try: + return _validate_openapi_spec(yaml.load(content, Loader=_OpenAPISafeLoader)) + except (yaml.YAMLError, RecursionError) as exc: + raise ValueError("Invalid YAML OpenAPI spec") from exc + + async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> _OpenAPISpec: if filepath.startswith("http://") or filepath.startswith("https://"): client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) @@ -208,14 +298,17 @@ async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None else await async_safe_get(client, filepath, max_response_bytes=max_bytes) ) r.raise_for_status() - return r.json() + return await asyncio.to_thread(_parse_openapi_spec, r.content) # fallback: local file - # Local filesystem path + return await asyncio.to_thread(_load_openapi_file, filepath) + + +def _load_openapi_file(filepath: str) -> _OpenAPISpec: if not os.path.exists(filepath): raise FileNotFoundError(f"OpenAPI spec not found at {filepath}") with open(filepath, "r", encoding="utf-8") as f: - return json.load(f) + return _parse_openapi_spec(f.read()) def get_base_url(spec: _OpenAPISpec, spec_path: str | None = None) -> str: 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 c42579dc7f7..ac79c10869b 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -5405,7 +5405,42 @@ class TestMCPServerManager: ( httpx.Response(200, text="secret invalid JSON body"), "unhealthy", - "OpenAPI specification could not be loaded (JSONDecodeError)", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="paths: [secret invalid YAML body"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="paths: {}\nx: \x01"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: &a []\ny: *a"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!bool invalid"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!int ''"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!float ''"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!timestamp invalid"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", ), ], ) diff --git a/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py b/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py index aac85160eba..0f29a3bf20d 100644 --- a/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py +++ b/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py @@ -1,9 +1,15 @@ from __future__ import annotations -from typing import Any, Dict +import asyncio +import json +import threading +from pathlib import Path +from typing import Any, Dict, Final import httpx import pytest +import yaml +from respx import MockRouter from litellm.proxy.mcp_server import openapi_to_mcp_generator as gen @@ -98,3 +104,147 @@ def test_load_openapi_spec_supports_local_file_path( spec = gen.load_openapi_spec(str(p)) assert spec == expected + + +@pytest.mark.parametrize("source", ["url", "file"]) +@pytest.mark.parametrize("suffix", [".yaml", ".json", ""]) +@pytest.mark.parametrize("document_format", ["json", "yaml"]) +@pytest.mark.parametrize("use_async", [False, True]) +def test_openapi_formats_preserve_document_across_sources( + tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + source: str, suffix: str, document_format: str, use_async: bool, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + expected: Final = { + "openapi": "3.0.0", + "paths": {"/ping": {"get": { + "operationId": "on", + "parameters": [ + {"name": "on", "in": "query", "schema": {"type": "string"}}, + {"name": "off", "in": "query", "schema": {"type": "string"}}, + ], + }}}, + "x-values": [100.0, "on", "off", "yes", "no", "2026-01-01", "1:20", True, False, None, 12, 10, 10], + } + content: Final = ( + json.dumps(expected).replace("100.0", "1e2") + if document_format == "json" + else "openapi: 3.0.0\npaths:\n /ping:\n get:\n operationId: on\n" + " parameters:\n" + " - {name: on, in: query, schema: {type: string}}\n" + " - {name: off, in: query, schema: {type: string}}\n" + "x-values: [1e2, on, off, yes, no, 2026-01-01, 1:20, TRUE, False, null, 012, 0o12, 0xA]\n" + ) + path: Final = tmp_path / ("spec" + suffix) + path.write_text(content, encoding="utf-8") + target: Final = str(path) if source == "file" else "https://93.184.216.34/spec" + suffix + if source == "url": + respx_mock.get(target).respond(200, text=content, headers={"content-type": "text/plain"}) + + spec: Final = asyncio.run(gen.load_openapi_spec_async(target)) if use_async else gen.load_openapi_spec(target) + + assert spec == expected + operation: Final = spec["paths"]["/ping"]["get"] + assert gen.sanitize_openapi_tool_name(operation["operationId"]) == "on" + assert gen.build_input_schema(operation) == { + "type": "object", + "properties": {"on": {"type": "string", "description": ""}, "off": {"type": "string", "description": ""}}, + "required": [], + } + + +_LONG_YAML_INTEGER: Final = "1" * 1025 + + +@pytest.mark.parametrize("content,message", [ + ("", "OpenAPI spec must be a JSON or YAML object"), + ("null", "OpenAPI spec must be a JSON or YAML object"), + ("[]", "OpenAPI spec must be a JSON or YAML object"), + ("- item", "OpenAPI spec must be a JSON or YAML object"), + ("plain text", "OpenAPI spec must be a JSON or YAML object"), + ("paths: [", "Invalid YAML OpenAPI spec"), + ("paths: {}\nx: \x01", "Invalid YAML OpenAPI spec"), + ("paths: " + "[" * 600 + "0" + "]" * 600, "Invalid YAML OpenAPI spec"), + ("x: !!python/object/apply:builtins.str [unsafe]", "Invalid YAML OpenAPI spec"), + ("x: &a []\ny: *a", "YAML aliases are not supported in OpenAPI specs"), + ("x: " + _LONG_YAML_INTEGER, "YAML integer is too long in OpenAPI spec"), + ("x: !!int " + _LONG_YAML_INTEGER, "YAML integer is too long in OpenAPI spec"), + ("x: !!int {=: '" + _LONG_YAML_INTEGER + "'}", "Invalid YAML OpenAPI spec"), + ("x: !!int [1]", "Invalid YAML OpenAPI spec"), + ("x: !!bool invalid", "Invalid YAML OpenAPI spec"), + ("x: !!bool []", "Invalid YAML OpenAPI spec"), + ("x: !!str []", "Invalid YAML OpenAPI spec"), + ("x: !!int ''", "Invalid YAML OpenAPI spec"), + ("x: !!float ''", "Invalid YAML OpenAPI spec"), + ("x: !!float {=: 1.2}", "Invalid YAML OpenAPI spec"), + ("x: !!timestamp invalid", "Invalid YAML OpenAPI spec"), + ("x: !!timestamp 2026-01-01", "Invalid YAML OpenAPI spec"), + ("x: !!set {a: null}", "Invalid YAML OpenAPI spec"), + ("x: !!binary YQ==", "Invalid YAML OpenAPI spec"), + ("x: {1: value}", "Invalid YAML OpenAPI spec"), + ("x: {[key]: value}", "Invalid YAML OpenAPI spec"), + ("x: .inf", "Invalid YAML OpenAPI spec"), + ("x: .nan", "Invalid YAML OpenAPI spec"), + ("x: 1e999", "Invalid YAML OpenAPI spec"), +], ids=[ + "empty", "null", "json-array", "yaml-sequence", "scalar", "syntax", "reader", "nesting", + "python-tag", "alias", "implicit-int", "explicit-int", "mapping-int", "sequence-int", + "invalid-bool", "sequence-bool", "sequence-str", "empty-int", "empty-float", "mapping-float", "invalid-timestamp", + "timestamp", "set", "binary", "non-string-key", "complex-key", "infinity", "nan", "float-overflow", +]) +def test_openapi_loader_rejects_invalid_or_unsafe_documents( + tmp_path: Path, content: str, message: str, +) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text(content, encoding="utf-8") + + with pytest.raises(ValueError, match="^" + message + "$"): + gen.load_openapi_spec(str(path)) + + +def test_openapi_yaml_integer_boundary_preserves_quoted_strings(tmp_path: Path) -> None: + integer: Final = "1" * 1024 + path: Final = tmp_path / "spec.yaml" + path.write_text("x-integer: " + integer + "\nx-string: '" + integer + "1'", encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x-integer": int(integer), "x-string": integer + "1"} + + +def test_openapi_json_integer_preserved_even_with_yaml_extension(tmp_path: Path) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text('{"x-integer":' + _LONG_YAML_INTEGER + "}", encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x-integer": int(_LONG_YAML_INTEGER)} + + +@pytest.mark.parametrize("value", ["! " + _LONG_YAML_INTEGER, "! '" + _LONG_YAML_INTEGER + "'"], ids=["plain", "quoted"]) +def test_openapi_nonspecific_tag_disables_scalar_resolution(tmp_path: Path, value: str) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text("x: " + value, encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x": _LONG_YAML_INTEGER} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["url", "file"]) +async def test_openapi_yaml_parses_off_the_event_loop( + tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, source: str, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + event_loop_thread: Final = threading.get_ident() + parser_threads: Final[list[int]] = [] + original_load: Final = yaml.load + + def observed_load(stream: str | bytes, Loader: type[yaml.SafeLoader]) -> object: + parser_threads.append(threading.get_ident()) + return original_load(stream, Loader=Loader) + + monkeypatch.setattr(yaml, "load", observed_load) + path: Final = tmp_path / "spec.yaml" + path.write_text("paths: {}", encoding="utf-8") + target: Final = str(path) if source == "file" else "https://93.184.216.34/spec.yaml" + if source == "url": + respx_mock.get(target).respond(200, text="paths: {}") + + assert await gen.load_openapi_spec_async(target) == {"paths": {}} + assert len(parser_threads) == 1 and parser_threads[0] != event_loop_thread