Merge pull request #897 from owtaylor/openapi-recursive-components

src/fastmcp/utilities/json_schema.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import copy
+from collections import defaultdict
 
 
 def _prune_param(schema: dict, param: str) -> dict:
@@ -24,25 +25,77 @@ def _prune_param(schema: dict, param: str) -> dict:
     return schema
 
 
+def _prune_unused_defs(schema: dict) -> dict:
+    """Walk the schema and prune unused defs."""
+
+    root_defs: set[str] = set()
+    referenced_by: defaultdict[str, list] = defaultdict(list)
+
+    defs = schema.get("$defs")
+    if defs is None:
+        return schema
+
+    def walk(
+        node: object, current_def: str | None = None, skip_defs: bool = False
+    ) -> None:
+        if isinstance(node, dict):
+            # Process $ref for definition tracking
+            ref = node.get("$ref")
+            if isinstance(ref, str) and ref.startswith("#/$defs/"):
+                def_name = ref.split("/")[-1]
+                if current_def:
+                    referenced_by[def_name].append(current_def)
+                else:
+                    root_defs.add(def_name)
+
+            # Walk children
+            for k, v in node.items():
+                if skip_defs and k == "$defs":
+                    continue
+
+                walk(v, current_def=current_def)
+
+        elif isinstance(node, list):
+            for v in node:
+                walk(v)
+
+    # Traverse the schema once, skipping the $defs
+    walk(schema, skip_defs=True)
+
+    # Now figure out what defs reference other defs
+    for def_name, value in defs.items():
+        walk(value, current_def=def_name)
+
+    # Figure out what defs were referenced directly or recursively
+    def def_is_referenced(def_name):
+        if def_name in root_defs:
+            return True
+        references = referenced_by.get(def_name)
+        if references:
+            for reference in references:
+                if def_is_referenced(reference):
+                    return True
+        return False
+
+    # Remove orphaned definitions if requested
+    for def_name in list(defs):
+        if not def_is_referenced(def_name):
+            defs.pop(def_name)
+    if not defs:
+        schema.pop("$defs", None)
+
+    return schema
+
+
 def _walk_and_prune(
     schema: dict,
-    prune_defs: bool = False,
     prune_titles: bool = False,
     prune_additional_properties: bool = False,
 ) -> dict:
-    """Walk the schema and optionally prune titles, unused definitions, and additionalProperties: false."""
-
-    # Will only be used if prune_defs is True
-    used_defs: set[str] = set()
+    """Walk the schema and optionally prune titles and additionalProperties: false."""
 
     def walk(node: object) -> None:
         if isinstance(node, dict):
-            # Process $ref for definition tracking
-            if prune_defs:
-                ref = node.get("$ref")
-                if isinstance(ref, str) and ref.startswith("#/$defs/"):
-                    used_defs.add(ref.split("/")[-1])
-
             # Remove title if requested
             if prune_titles and "title" in node:
                 node.pop("title")
@@ -62,18 +115,8 @@ def _walk_and_prune(
             for v in node:
                 walk(v)
 
-    # Traverse the schema once
     walk(schema)
 
-    # Remove orphaned definitions if requested
-    if prune_defs:
-        defs = schema.get("$defs", {})
-        for def_name in list(defs):
-            if def_name not in used_defs:
-                defs.pop(def_name)
-        if not defs:
-            schema.pop("$defs", None)
-
     return schema
 
 
@@ -109,12 +152,13 @@ def compress_schema(
         schema = _prune_param(schema, param=param)
 
     # Do a single walk to handle pruning operations
-    if prune_defs or prune_titles or prune_additional_properties:
+    if prune_titles or prune_additional_properties:
         schema = _walk_and_prune(
             schema,
-            prune_defs=prune_defs,
             prune_titles=prune_titles,
             prune_additional_properties=prune_additional_properties,
         )
+    if prune_defs:
+        schema = _prune_unused_defs(schema)
 
     return schema

src/fastmcp/utilities/openapi.py

@@ -262,16 +262,18 @@ class OpenAPIParser(
 
             if isinstance(resolved_schema, (self.schema_cls)):
                 # Convert schema to dictionary
-                return resolved_schema.model_dump(
+                result = resolved_schema.model_dump(
                     mode="json", by_alias=True, exclude_none=True
                 )
             elif isinstance(resolved_schema, dict):
-                return resolved_schema
+                result = resolved_schema
             else:
                 logger.warning(
                     f"Expected Schema after resolving, got {type(resolved_schema)}. Returning empty dict."
                 )
-                return {}
+                result = {}
+
+            return _replace_ref_with_defs(result)
         except Exception as e:
             logger.error(f"Failed to extract schema as dict: {e}", exc_info=False)
             return {}

tests/utilities/openapi/test_openapi_advanced.py

@@ -294,6 +294,28 @@ def test_complex_schema_route_count(parsed_complex_routes):
     assert len(parsed_complex_routes) == 3
 
 
+def test_complex_schema_ref_rewriting(parsed_complex_routes):
+    """Test that all #/components references have been rewritten."""
+
+    def no_components(value):
+        if isinstance(value, dict):
+            for k, v in value.items():
+                if k == "$ref":
+                    assert not v.startswith("#/components/"), (
+                        f"reference '{v}' was not rewritten"
+                    )
+                else:
+                    no_components(v)
+        elif isinstance(value, list):
+            for v in value:
+                no_components(v)
+
+    for route in parsed_complex_routes:
+        no_components(route.schema_definitions)
+        for param in route.parameters:
+            no_components(param.schema_)
+
+
 def test_complex_schema_list_users_query_param_limit(complex_route_map):
     """Test that a reference to a limit query parameter is correctly resolved."""
     list_users = complex_route_map["listUsers"]

tests/utilities/test_json_schema.py

@@ -1,14 +1,11 @@
 from fastmcp.utilities.json_schema import (
     _prune_param,
+    _prune_unused_defs,
     _walk_and_prune,
     compress_schema,
 )
 
-
-# Create wrappers for backward compatibility with tests
-def _prune_unused_defs(schema):
-    """Wrapper for _walk_and_prune that only prunes definitions."""
-    return _walk_and_prune(schema, prune_defs=True)
+# Wrapper for backward compatibility with tests
 
 
 def _prune_additional_properties(schema):
@@ -95,6 +92,21 @@ class TestPruneUnusedDefs:
         assert "nested_def" in result["$defs"]
         assert "unused_def" not in result["$defs"]
 
+    def test_nested_references_removed(self):
+        """Test that definitions referenced via nesting in unused defs are removed."""
+        schema = {
+            "properties": {},
+            "$defs": {
+                "foo_def": {
+                    "type": "object",
+                    "properties": {"nested": {"$ref": "#/$defs/nested_def"}},
+                },
+                "nested_def": {"type": "string"},
+            },
+        }
+        result = _prune_unused_defs(schema)
+        assert "$defs" not in result
+
     def test_array_references_kept(self):
         """Test that definitions referenced in array items are kept."""
         schema = {