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 5b23695d06d..6f945c2448a 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -28,6 +28,7 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( # which include '/'. Sanitize here so the same regex passes everywhere downstream. _OPENAPI_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") _OPENAPI_TOOL_NAME_MAX_LEN: Final = 128 +_MAX_YAML_INT_LENGTH: Final = 1024 def sanitize_openapi_tool_name(raw_name: str) -> str: @@ -147,7 +148,7 @@ def _sanitize_path_parameter_value(param_value: object, param_name: str) -> str: return quote(value_str, safe="") -def load_openapi_spec(filepath: str) -> dict[str, Any]: +def load_openapi_spec(filepath: str) -> Mapping[str, Any]: """ Sync wrapper. For URL specs, use the shared/custom MCP httpx client. """ @@ -165,6 +166,43 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]: return asyncio.run(load_openapi_spec_async(filepath)) +def _parse_openapi_spec(text: str) -> Mapping[str, Any]: + """Parse JSON or YAML OpenAPI documents into a mapping.""" + try: + parsed: Any = json.loads(text) + except json.JSONDecodeError: + import yaml + from yaml.nodes import ScalarNode + + class _NoMergeSafeLoader(yaml.SafeLoader): + def compose_node(self, parent: object, index: object) -> object: + if self.check_event(yaml.events.AliasEvent): + raise yaml.YAMLError("YAML aliases are not supported") + return super().compose_node(parent, index) + + def flatten_mapping(self, node: object) -> None: + node_values = getattr(node, "value", ()) + if any(key_node.tag == "tag:yaml.org,2002:merge" for key_node, _ in node_values): + raise yaml.YAMLError("YAML merge keys are not supported") + super().flatten_mapping(node) + + def construct_yaml_int(self, node: ScalarNode) -> int: + if len(node.value) > _MAX_YAML_INT_LENGTH: + raise yaml.YAMLError("YAML integer is too long") + return super().construct_yaml_int(node) + + _NoMergeSafeLoader.add_constructor( + "tag:yaml.org,2002:int", + _NoMergeSafeLoader.construct_yaml_int, + ) + + parsed = yaml.load(text, Loader=_NoMergeSafeLoader) + + if not isinstance(parsed, dict): + raise TypeError("OpenAPI spec must be a JSON or YAML object") + return parsed + + 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) @@ -174,14 +212,14 @@ 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 _parse_openapi_spec(r.text) # fallback: local file # Local filesystem path 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: Mapping[str, Any], spec_path: str | None = None) -> str: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 9aaba0e9356..1962789cfec 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -4936,7 +4936,7 @@ class TestMCPServerManager: ( httpx.Response(200, text="secret invalid JSON body"), "unhealthy", - "OpenAPI specification could not be loaded (JSONDecodeError)", + "OpenAPI specification could not be loaded (TypeError)", ), ], ) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/unit/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 6b0211c3866..5f5b8e438af 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -17,8 +17,11 @@ import pytest from fastapi import HTTPException from respx import MockRouter -from litellm.types.mcp import MCPAuth, MCPAuthType - +from litellm.proxy._experimental.mcp_server import openapi_to_mcp_generator as gen +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPOpenApiUpstreamError, + MCPUpstreamAuthError, +) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, @@ -33,18 +36,15 @@ from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( resolve_operation_params, ) from litellm.proxy._experimental.mcp_server.tool_outcome import JsonResult, TextResult - -from litellm.proxy._experimental.mcp_server.exceptions import ( - MCPOpenApiUpstreamError, - MCPUpstreamAuthError, -) +from litellm.types.mcp import MCPAuth, MCPAuthType GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client" @pytest.mark.asyncio async def test_unsupported_http_method_returns_text_without_sending_request( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") tool: Final = create_tool_function("/echo", "HEAD", {}, "https://upstream.example") @@ -58,19 +58,24 @@ async def test_unsupported_http_method_returns_text_without_sending_request( @pytest.mark.asyncio -@pytest.mark.parametrize("body,expected", [ - (' { "ok": true }\n', JsonResult({"ok": True}, ' { "ok": true }\n')), - (' [1, 2]\n', JsonResult([1, 2], ' [1, 2]\n')), - ('false', JsonResult(False, 'false')), - ('0', JsonResult(0, '0')), - ('""', JsonResult("", '""')), - ('null', TextResult('null')), - ('{"unfinished":', TextResult('{"unfinished":')), - ('', TextResult('')), -]) +@pytest.mark.parametrize( + "body,expected", + [ + (' { "ok": true }\n', JsonResult({"ok": True}, ' { "ok": true }\n')), + (" [1, 2]\n", JsonResult([1, 2], " [1, 2]\n")), + ("false", JsonResult(False, "false")), + ("0", JsonResult(0, "0")), + ('""', JsonResult("", '""')), + ("null", TextResult("null")), + ('{"unfinished":', TextResult('{"unfinished":')), + ("", TextResult("")), + ], +) async def test_http_response_preserves_body_and_classifies_json( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, - body: str, expected: TextResult | JsonResult, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + body: str, + expected: TextResult | JsonResult, ) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example") @@ -80,23 +85,42 @@ async def test_http_response_preserves_body_and_classifies_json( @pytest.mark.asyncio -@pytest.mark.parametrize("auth_type,value,accepted", [ - (MCPAuth.api_key, "Bearer Bearer", False), (MCPAuth.api_key, "ApiKey ApiKey", False), - (MCPAuth.api_key, "token token", False), (MCPAuth.api_key, "bEaReR BEARER", False), - (MCPAuth.api_key, "aPiKeY\tAPIKEY", False), (MCPAuth.api_key, "Bearer fixture-key", True), - (MCPAuth.api_key, "ApiKey fixture-key", True), (MCPAuth.api_key, "token fixture-key", True), - (MCPAuth.authorization, "Bearer", False), (MCPAuth.authorization, "basic", False), - (MCPAuth.authorization, "token", False), (MCPAuth.authorization, "ApiKey", False), - (MCPAuth.authorization, " bEaReR ", False), (MCPAuth.authorization, "\tTOKEN\t", False), - (MCPAuth.authorization, "opaque-secret-value", True), (MCPAuth.authorization, "Bearer abc", True), - (MCPAuth.authorization, "Custom abc", True), -]) +@pytest.mark.parametrize( + "auth_type,value,accepted", + [ + (MCPAuth.api_key, "Bearer Bearer", False), + (MCPAuth.api_key, "ApiKey ApiKey", False), + (MCPAuth.api_key, "token token", False), + (MCPAuth.api_key, "bEaReR BEARER", False), + (MCPAuth.api_key, "aPiKeY\tAPIKEY", False), + (MCPAuth.api_key, "Bearer fixture-key", True), + (MCPAuth.api_key, "ApiKey fixture-key", True), + (MCPAuth.api_key, "token fixture-key", True), + (MCPAuth.authorization, "Bearer", False), + (MCPAuth.authorization, "basic", False), + (MCPAuth.authorization, "token", False), + (MCPAuth.authorization, "ApiKey", False), + (MCPAuth.authorization, " bEaReR ", False), + (MCPAuth.authorization, "\tTOKEN\t", False), + (MCPAuth.authorization, "opaque-secret-value", True), + (MCPAuth.authorization, "Bearer abc", True), + (MCPAuth.authorization, "Custom abc", True), + ], +) async def test_authorization_validates_credentials_before_http( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuthType, value: str, accepted: bool, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, + value: str, + accepted: bool, ) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") tool: Final = create_tool_function( - "/echo", "get", {}, "https://upstream.example", auth_type=auth_type, + "/echo", + "get", + {}, + "https://upstream.example", + auth_type=auth_type, ) destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") caller_token: Final = _request_auth_header.set(value) @@ -115,20 +139,44 @@ async def test_authorization_validates_credentials_before_http( @pytest.mark.asyncio -@pytest.mark.parametrize("static,forwarded,caller,resolved,expected", [ - ({"Authorization": "Bearer configured"}, {"authorization": "Bearer forwarded"}, None, None, "Bearer configured"), - ({"Authorization": "Bearer configured"}, None, "Bearer caller", None, "Bearer caller"), - ({"Authorization": "Bearer configured"}, None, "Bearer", None, None), - ({"Authorization": "Bearer configured"}, None, "Bearer caller", {"authorization": " "}, None), - ({"Authorization": "Bearer configured"}, None, "Bearer", {"authorization": "Bearer resolved"}, "Bearer resolved"), -]) +@pytest.mark.parametrize( + "static,forwarded,caller,resolved,expected", + [ + ( + {"Authorization": "Bearer configured"}, + {"authorization": "Bearer forwarded"}, + None, + None, + "Bearer configured", + ), + ({"Authorization": "Bearer configured"}, None, "Bearer caller", None, "Bearer caller"), + ({"Authorization": "Bearer configured"}, None, "Bearer", None, None), + ({"Authorization": "Bearer configured"}, None, "Bearer caller", {"authorization": " "}, None), + ( + {"Authorization": "Bearer configured"}, + None, + "Bearer", + {"authorization": "Bearer resolved"}, + "Bearer resolved", + ), + ], +) async def test_static_auth_validates_headers_after_existing_precedence( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, - static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None, - resolved: dict[str, str] | None, expected: str | None, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], + forwarded: dict[str, str] | None, + caller: str | None, + resolved: dict[str, str] | None, + expected: str | None, ) -> None: tool: Final = create_tool_function( - "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.bearer_token, + "/echo", + "get", + {}, + "https://upstream.example", + headers=static, + auth_type=MCPAuth.bearer_token, ) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") @@ -154,12 +202,19 @@ async def test_static_auth_validates_headers_after_existing_precedence( @pytest.mark.asyncio @pytest.mark.parametrize("credential", ["custom-key", ""]) async def test_static_auth_uses_configured_custom_header( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + credential: str, ) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") tool: Final = create_tool_function( - "/echo", "get", {}, "https://upstream.example", headers={"x-custom": credential}, - auth_type=MCPAuth.api_key, upstream_token_header="X-Custom", + "/echo", + "get", + {}, + "https://upstream.example", + headers={"x-custom": credential}, + auth_type=MCPAuth.api_key, + upstream_token_header="X-Custom", ) destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") if credential: @@ -175,11 +230,18 @@ async def test_static_auth_uses_configured_custom_header( @pytest.mark.asyncio @pytest.mark.parametrize("credential", ["static-key", ""]) async def test_static_auth_accepts_api_key_carried_by_static_header( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + credential: str, ) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") tool: Final = create_tool_function( - "/echo", "get", {}, "https://upstream.example", headers={"apikey": credential}, auth_type=MCPAuth.api_key, + "/echo", + "get", + {}, + "https://upstream.example", + headers={"apikey": credential}, + auth_type=MCPAuth.api_key, ) destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") if credential: @@ -193,13 +255,18 @@ async def test_static_auth_accepts_api_key_carried_by_static_header( @pytest.mark.asyncio -@pytest.mark.parametrize("auth_type,resolved", [ - (MCPAuth.none, None), - (MCPAuth.oauth2, {"Authorization": "Bearer user-oauth"}), -]) +@pytest.mark.parametrize( + "auth_type,resolved", + [ + (MCPAuth.none, None), + (MCPAuth.oauth2, {"Authorization": "Bearer user-oauth"}), + ], +) async def test_static_validation_preserves_no_auth_and_resolved_oauth( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, - auth_type: MCPAuthType, resolved: dict[str, str] | None, + respx_mock: MockRouter, + monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, + resolved: dict[str, str] | None, ) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example", auth_type=auth_type) @@ -213,6 +280,81 @@ async def test_static_validation_preserves_no_auth_and_resolved_oauth( _request_resolved_auth_headers.reset(token) +def test_load_openapi_spec_supports_yaml_http_url(monkeypatch: pytest.MonkeyPatch) -> None: + url = "http://example.local/openapi.yaml" + response = SimpleNamespace( + text=""" +openapi: 3.0.0 +info: + title: YAML API + version: 1.0.0 +paths: {} +""", + status_code=200, + headers={}, + raise_for_status=lambda: None, + ) + client = AsyncMock() + client.get = AsyncMock(return_value=response) + + monkeypatch.setattr(gen, "get_async_httpx_client", lambda **kwargs: client) + monkeypatch.setattr(gen, "async_safe_get", lambda client, url, **kwargs: client.get(url)) + + assert gen.load_openapi_spec(url) == { + "openapi": "3.0.0", + "info": {"title": "YAML API", "version": "1.0.0"}, + "paths": {}, + } + + +def test_parse_openapi_spec_rejects_non_object_document() -> None: + with pytest.raises(TypeError, match="OpenAPI spec must be a JSON or YAML object"): + gen._parse_openapi_spec("- not an OpenAPI object\n") + + +def test_parse_openapi_spec_rejects_yaml_merge_aliases() -> None: + import yaml + + with pytest.raises(yaml.YAMLError, match="YAML aliases are not supported"): + gen._parse_openapi_spec( + """ +base: &base + openapi: 3.0.0 + info: + title: Base + version: '1.0.0' + paths: {} +merged: + <<: *base +""" + ) + + +def test_parse_openapi_spec_rejects_yaml_aliases() -> None: + import yaml + + with pytest.raises(yaml.YAMLError, match="YAML aliases are not supported"): + gen._parse_openapi_spec( + """ +shared: &shared + type: string +components: + schemas: + Request: + properties: + first: *shared +""" + ) + + +def test_parse_openapi_spec_rejects_oversized_yaml_integer() -> None: + import yaml + + oversized_integer: Final = "1" + ":11" * gen._MAX_YAML_INT_LENGTH + with pytest.raises(yaml.YAMLError, match="YAML integer is too long"): + gen._parse_openapi_spec(f"openapi: 3.0.0\nx-expensive: {oversized_integer}\n") + + def _create_mock_client(method: str, response_text: str, status_code: int = 200) -> AsyncMock: """Utility to create a mocked async httpx client for the given method. @@ -263,9 +405,7 @@ class TestCreateToolFunction: # Verify URL was constructed correctly call_args = async_client.get.call_args - assert "repository-id" in str(call_args[0][0]) or "test-repo" in str( - call_args[0][0] - ) + assert "repository-id" in str(call_args[0][0]) or "test-repo" in str(call_args[0][0]) @pytest.mark.asyncio async def test_leading_digit_parameter(self): @@ -610,9 +750,7 @@ class TestExtractParameters: {"name": "filter", "in": "query"}, {"name": "data", "in": "body"}, ], - "requestBody": { - "content": {"application/json": {"schema": {"type": "object"}}} - }, + "requestBody": {"content": {"application/json": {"schema": {"type": "object"}}}}, } path_params, query_params, body_params = extract_parameters(operation) @@ -814,7 +952,6 @@ class TestGetBaseUrl: base_url = get_base_url(spec, spec_path) assert base_url == "https://production.example.com" - def test_fallback_with_nested_path(self): """Test fallback with deeply nested spec path.""" spec = {"openapi": "3.0.0", "paths": {}} @@ -1042,9 +1179,7 @@ class TestResolveOperationParams: ], } path_item = {"parameters": path_level_params, "get": operation} - result = resolve_operation_params( - operation, path_item, {"parameters": component_params} - ) + result = resolve_operation_params(operation, path_item, {"parameters": component_params}) names = [p["name"] for p in result["parameters"]] assert "owner" in names assert "repo" in names @@ -1147,22 +1282,16 @@ class TestRegisterToolsFromOpenAPI: } } - openapi_to_mcp_generator.register_tools_from_openapi( - spec, base_url="https://api.example.com" - ) + openapi_to_mcp_generator.register_tools_from_openapi(spec, base_url="https://api.example.com") assert registered, "expected at least one registered tool" anthropic_re = re.compile(r"^[a-zA-Z0-9_-]{1,128}$") for name in registered: - assert anthropic_re.match( - name - ), f"tool name {name!r} violates ^[a-zA-Z0-9_-]+$" + assert anthropic_re.match(name), f"tool name {name!r} violates ^[a-zA-Z0-9_-]+$" assert "actions_download-job-logs-for-workflow-run" in registered assert "pulls_list-files" in registered - def test_missing_operation_id_uses_sanitized_method_path_fallback( - self, monkeypatch - ): + def test_missing_operation_id_uses_sanitized_method_path_fallback(self, monkeypatch): import re from litellm.proxy._experimental.mcp_server import openapi_to_mcp_generator @@ -1185,15 +1314,11 @@ class TestRegisterToolsFromOpenAPI: } } } - openapi_to_mcp_generator.register_tools_from_openapi( - spec, base_url="https://api.example.com" - ) + openapi_to_mcp_generator.register_tools_from_openapi(spec, base_url="https://api.example.com") assert registered for name in registered: - assert re.match( - r"^[a-zA-Z0-9_-]+$", name - ), f"fallback tool name {name!r} not sanitized" + assert re.match(r"^[a-zA-Z0-9_-]+$", name), f"fallback tool name {name!r} not sanitized" class TestRequestExtraHeaders: @@ -1349,9 +1474,7 @@ class TestRequestExtraHeaders: async_client = _create_mock_client("get", "secure-data") mock_client.return_value = async_client - extra_token = _request_extra_headers.set( - {"Authorization": "Bearer extra", "X-TOKEN": "token-value"} - ) + extra_token = _request_extra_headers.set({"Authorization": "Bearer extra", "X-TOKEN": "token-value"}) auth_token = _request_auth_header.set("Bearer byok-credential") try: result = await func() @@ -1479,7 +1602,9 @@ class TestUpstreamStatusIsClassified: @pytest.mark.asyncio async def test_401_raises_the_reauth_signal_carrying_the_challenge(self): - tool, client = self._tool(401, text='{"error":"invalid_token"}', headers={"www-authenticate": 'Bearer realm="x"'}) + tool, client = self._tool( + 401, text='{"error":"invalid_token"}', headers={"www-authenticate": 'Bearer realm="x"'} + ) with patch(GET_ASYNC_CLIENT_TARGET, return_value=client): with pytest.raises(MCPUpstreamAuthError) as exc: await tool() @@ -1557,6 +1682,7 @@ class TestUpstreamStatusIsClassified: 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]) @@ -1572,6 +1698,7 @@ 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 ( load_openapi_spec_async, @@ -1596,6 +1723,7 @@ class TestBoundedOpenAPISpecLoading: @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,