Optimize OpenAPI parser performance with single-pass schema processing (#1214)

Author: jlowin

src/fastmcp/experimental/utilities/openapi/schemas.py

@@ -3,8 +3,6 @@
 import logging
 from typing import Any, cast
 
-from fastmcp.utilities.json_schema import compress_schema
-
 from .models import HTTPRoute, JsonSchema, ResponseInfo
 
 logger = logging.getLogger(__name__)
@@ -314,10 +312,42 @@ def _combine_schemas_and_map_params(
     }
     # Add schema definitions if available
     if route.schema_definitions:
-        result["$defs"] = route.schema_definitions
-
-    # Use compress_schema to remove unused definitions
-    result = compress_schema(result)
+        result["$defs"] = route.schema_definitions.copy()
+
+    # Use lightweight compression - prune additionalProperties and unused definitions
+    if result.get("additionalProperties") is False:
+        result.pop("additionalProperties")
+
+    # Remove unused definitions (lightweight approach - just check direct $ref usage)
+    if "$defs" in result:
+        used_refs = set()
+
+        def find_refs_in_value(value):
+            if isinstance(value, dict):
+                if "$ref" in value and isinstance(value["$ref"], str):
+                    ref = value["$ref"]
+                    if ref.startswith("#/$defs/"):
+                        used_refs.add(ref.split("/")[-1])
+                for v in value.values():
+                    find_refs_in_value(v)
+            elif isinstance(value, list):
+                for item in value:
+                    find_refs_in_value(item)
+
+        # Find refs in the main schema (excluding $defs section)
+        for key, value in result.items():
+            if key != "$defs":
+                find_refs_in_value(value)
+
+        # Remove unused definitions
+        if used_refs:
+            result["$defs"] = {
+                name: def_schema
+                for name, def_schema in result["$defs"].items()
+                if name in used_refs
+            }
+        else:
+            result.pop("$defs")
 
     return result, parameter_map
 
@@ -339,17 +369,63 @@ def _combine_schemas(route: HTTPRoute) -> dict[str, Any]:
     return schema
 
 
+def _has_one_of(obj: dict[str, Any] | list[Any]) -> bool:
+    """Quickly check if schema contains any 'oneOf' keys without deep traversal."""
+    if isinstance(obj, dict):
+        if "oneOf" in obj:
+            return True
+        # Only check likely schema containers, skip examples/defaults
+        for k, v in obj.items():
+            if k in [
+                "properties",
+                "items",
+                "allOf",
+                "anyOf",
+                "additionalProperties",
+            ] and isinstance(v, dict | list):
+                if _has_one_of(v):
+                    return True
+    elif isinstance(obj, list):
+        for item in obj:
+            if isinstance(item, dict | list) and _has_one_of(item):
+                return True
+    return False
+
+
 def _adjust_union_types(
-    schema: dict[str, Any] | list[Any],
+    schema: dict[str, Any] | list[Any], _depth: int = 0
 ) -> dict[str, Any] | list[Any]:
     """Recursively replace 'oneOf' with 'anyOf' in schema to handle overlapping unions."""
+    # MAJOR OPTIMIZATION: Skip entirely if schema has no oneOf keys
+    if _depth == 0 and not _has_one_of(schema):
+        return schema
+
+    # OPTIMIZATION: Early termination for very deep structures to prevent exponential slowdown
+    if _depth > 30:  # Reduced from 50 for better performance
+        return schema
+
     if isinstance(schema, dict):
-        if "oneOf" in schema:
-            schema["anyOf"] = schema.pop("oneOf")
-        for k, v in schema.items():
-            schema[k] = _adjust_union_types(v)
+        # Work on a copy to avoid mutating the input
+        result = schema.copy()
+        if "oneOf" in result:
+            result["anyOf"] = result.pop("oneOf")
+        # OPTIMIZATION: Only recurse into values that could contain more schemas
+        for k, v in result.items():
+            if isinstance(v, dict | list) and k not in [
+                "examples",
+                "example",
+                "default",
+            ]:
+                result[k] = _adjust_union_types(v, _depth + 1)
+        return result
     elif isinstance(schema, list):
-        return [_adjust_union_types(item) for item in schema]
+        # Process list items without mutating the input list
+        return [
+            _adjust_union_types(item, _depth + 1)
+            if isinstance(item, dict | list)
+            else item
+            for item in schema
+        ]
     return schema
 
 
@@ -436,10 +512,42 @@ def extract_output_schema_from_responses(
 
     # Add schema definitions if available
     if schema_definitions:
-        output_schema["$defs"] = schema_definitions
-
-    # Use compress_schema to remove unused definitions
-    output_schema = compress_schema(output_schema)
+        output_schema["$defs"] = schema_definitions.copy()
+
+    # Use lightweight compression - prune additionalProperties and unused definitions
+    if output_schema.get("additionalProperties") is False:
+        output_schema.pop("additionalProperties")
+
+    # Remove unused definitions (lightweight approach - just check direct $ref usage)
+    if "$defs" in output_schema:
+        used_refs = set()
+
+        def find_refs_in_value(value):
+            if isinstance(value, dict):
+                if "$ref" in value and isinstance(value["$ref"], str):
+                    ref = value["$ref"]
+                    if ref.startswith("#/$defs/"):
+                        used_refs.add(ref.split("/")[-1])
+                for v in value.values():
+                    find_refs_in_value(v)
+            elif isinstance(value, list):
+                for item in value:
+                    find_refs_in_value(item)
+
+        # Find refs in the main schema (excluding $defs section)
+        for key, value in output_schema.items():
+            if key != "$defs":
+                find_refs_in_value(value)
+
+        # Remove unused definitions
+        if used_refs:
+            output_schema["$defs"] = {
+                name: def_schema
+                for name, def_schema in output_schema["$defs"].items()
+                if name in used_refs
+            }
+        else:
+            output_schema.pop("$defs")
 
     # Adjust union types to handle overlapping unions
     output_schema = cast(dict[str, Any], _adjust_union_types(output_schema))