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 6ab7d5956b9..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: @@ -171,6 +172,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: @@ -184,6 +186,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 5253d0653c0..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 @@ -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.