From 054a6cf3c23948d83dabe20056a9353e4377737d Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Mon, 27 Jul 2026 14:51:13 -0400
Subject: [PATCH] Keep $defs when the reference scan hits its depth limit

---
 fastmcp_slim/fastmcp/utilities/json_schema.py | 15 +++++++++++++
 tests/utilities/test_json_schema.py           | 21 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/fastmcp_slim/fastmcp/utilities/json_schema.py b/fastmcp_slim/fastmcp/utilities/json_schema.py
index 9f12dc246..59715e8ed 100644
--- a/fastmcp_slim/fastmcp/utilities/json_schema.py
+++ b/fastmcp_slim/fastmcp/utilities/json_schema.py
@@ -544,6 +544,11 @@ def _single_pass_optimize(
     )  # def A references def B
     defs = schema.get("$defs")
 
+    # Set when the traversal below gives up at its depth limit. Once that
+    # happens the reference scan is incomplete, so we can no longer tell which
+    # definitions are genuinely unused.
+    reference_scan_truncated = False
+
     def traverse_and_clean(
         node: object,
         current_def_name: str | None = None,
@@ -561,7 +566,10 @@ def _single_pass_optimize(
         about) but we skip all cleanups so we don't mutate user data that
         happens to look metadata-shaped.
         """
+        nonlocal reference_scan_truncated
+
         if depth > 50:  # Prevent infinite recursion
+            reference_scan_truncated = True
             return
 
         if isinstance(node, dict):
@@ -685,6 +693,13 @@ def _single_pass_optimize(
         for def_name, def_schema in defs.items():
             traverse_and_clean(def_schema, current_def_name=def_name, in_schema=True)
 
+        # An incomplete scan has not seen every $ref, so a definition that looks
+        # unused may simply be referenced below the cutoff. Keeping an unused
+        # definition is harmless; dropping a referenced one leaves a dangling
+        # $ref and an invalid schema.
+        if reference_scan_truncated:
+            return schema
+
         # Phase 4: Remove unused definitions
         def is_def_used(def_name: str, visiting: set[str] | None = None) -> bool:
             """Check if a definition is used, handling circular references."""
diff --git a/tests/utilities/test_json_schema.py b/tests/utilities/test_json_schema.py
index 1ae37216c..b97153b02 100644
--- a/tests/utilities/test_json_schema.py
+++ b/tests/utilities/test_json_schema.py
@@ -442,6 +442,27 @@ class TestCompressSchema:
         # ...and the copy really was pruned as deep as the traversal reaches.
         assert _count_titles(result) < _count_titles(schema)
 
+    def test_keeps_defs_referenced_below_the_traversal_cutoff(self):
+        """A $ref deeper than the traversal walks must still pin its definition.
+
+        The reference scan stops at its depth guard, so past that point it
+        cannot prove a definition is unused. Dropping one anyway would leave a
+        dangling $ref — an invalid schema is worse than an unpruned one.
+        """
+        schema: dict[str, Any] = {"$ref": "#/$defs/Leaf"}
+        for _ in range(60):
+            schema = {"type": "array", "items": schema}
+        schema["$defs"] = {"Leaf": {"type": "string"}}
+
+        result = compress_schema(schema)
+
+        assert result["$defs"] == {"Leaf": {"type": "string"}}
+
+        node: Any = result
+        while isinstance(node.get("items"), dict):
+            node = node["items"]
+        assert node == {"$ref": "#/$defs/Leaf"}
+
     def test_preserves_refs_by_default(self):
         """Test that compress_schema preserves $refs by default."""
         schema = {