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 0ce437bd7b3..cccc489c569 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -30,6 +30,7 @@ _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 +_MAX_YAML_INT_LENGTH: Final = 1024 def sanitize_openapi_tool_name(raw_name: str) -> str: @@ -181,6 +182,7 @@ def _parse_openapi_spec(text: str) -> Mapping[str, Any]: 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: @@ -194,6 +196,16 @@ def _parse_openapi_spec(text: str) -> Mapping[str, Any]: 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): 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 38b2e08238a..71b13d0fa6c 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 @@ -347,6 +347,14 @@ components: ) +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.