From bbf8165636bcc4d8815bcc53df10ba967c5b40fb Mon Sep 17 00:00:00 2001
From: "Owen W. Taylor" <otaylor@fishsoup.net>
Date: Fri, 20 Jun 2025 14:29:06 -0400
Subject: [PATCH] openapi: Improve pruning of unused defs

Defs that are used only by other unused defs were counted as
used. Fix this by tracing what defs use other defs recursively.
---
 src/fastmcp/utilities/json_schema.py | 90 +++++++++++++++++++++-------
 tests/utilities/test_json_schema.py  | 22 +++++--
 2 files changed, 84 insertions(+), 28 deletions(-)

diff --git a/src/fastmcp/utilities/json_schema.py b/src/fastmcp/utilities/json_schema.py
index 87ea4a789..ca5fe37c9 100644
--- a/src/fastmcp/utilities/json_schema.py
+++ b/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
diff --git a/tests/utilities/test_json_schema.py b/tests/utilities/test_json_schema.py
index a1b5f1584..55c970224 100644
--- a/tests/utilities/test_json_schema.py
+++ b/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 = {