Copy schemas iteratively so deep nesting still compresses (#4671)

* Copy schemas iteratively so deep nesting still compresses

* Keep $defs when the reference scan hits its depth limit

fastmcp_slim/fastmcp/utilities/json_schema.py

@@ -1,12 +1,45 @@
 from __future__ import annotations
 
-import copy
 from collections import defaultdict
 from typing import Any
 
 from jsonref import JsonRefError, replace_refs
 
 
+def _copy_schema(schema: dict[str, Any]) -> dict[str, Any]:
+    """Return a deep copy of a JSON schema without recursing.
+
+    `copy.deepcopy` consumes stack frames in proportion to nesting depth, so a
+    deeply nested schema raises `RecursionError` before the traversals in this
+    module can apply their own depth guards — turning a schema that used to
+    compress into one that fails outright. Schemas are plain JSON, so an
+    explicit stack copies the containers at any depth and shares the immutable
+    scalars at the leaves.
+    """
+    root: dict[str, Any] = {}
+    stack: list[tuple[Any, Any]] = [(schema, root)]
+
+    while stack:
+        source, target = stack.pop()
+        if isinstance(source, dict):
+            pairs: list[tuple[Any, Any]] = list(source.items())
+        else:
+            pairs = list(enumerate(source))
+
+        for key, value in pairs:
+            if isinstance(value, dict):
+                child: Any = {}
+                stack.append((value, child))
+            elif isinstance(value, list):
+                child = [None] * len(value)
+                stack.append((value, child))
+            else:
+                child = value
+            target[key] = child
+
+    return root
+
+
 def _defs_have_cycles(defs: dict[str, Any]) -> bool:
     """Check whether any definitions in ``$defs`` form a reference cycle.
 
@@ -348,7 +381,7 @@ def _prune_param(schema: dict[str, Any], param: str) -> dict[str, Any]:
     """Return a new schema with *param* removed from `properties`, `required`,
     and (if no longer referenced) `$defs`.
     """
-    schema = copy.deepcopy(schema)
+    schema = _copy_schema(schema)
 
     # ── 1. drop from properties/required ──────────────────────────────
     props = schema.get("properties", {})
@@ -501,7 +534,7 @@ def _single_pass_optimize(
     # Work on a copy so the caller's schema is never mutated (see docstring). The
     # pruning phases below pop keys/$defs in place, which would otherwise corrupt a
     # shared dict such as a live Tool.input_schema passed straight to compress_schema.
-    schema = copy.deepcopy(schema)
+    schema = _copy_schema(schema)
 
     # Phase 1: Collect references and apply simple cleanups
     # Track which $defs are referenced from the main schema and from other $defs
@@ -511,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,
@@ -528,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):
@@ -652,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."""

tests/utilities/test_json_schema.py

@@ -1,4 +1,6 @@
 import copy
+import sys
+from typing import Any
 from unittest.mock import patch
 
 from jsonref import replace_refs
@@ -12,6 +14,31 @@ from fastmcp.utilities.json_schema import (
 )
 
 
+def _measure_depth(schema: dict[str, Any]) -> int:
+    """Return how many `items` levels deep an array-nested schema goes.
+
+    Walks iteratively so the assertion helpers cannot themselves hit the
+    recursion limit the tests are probing.
+    """
+    depth = 0
+    node: Any = schema
+    while isinstance(node.get("items"), dict):
+        node = node["items"]
+        depth += 1
+    return depth
+
+
+def _count_titles(schema: dict[str, Any]) -> int:
+    """Count the `title` keys down an array-nested schema, iteratively."""
+    count = 0
+    node: Any = schema
+    while isinstance(node, dict):
+        if "title" in node:
+            count += 1
+        node = node.get("items")
+    return count
+
+
 class TestPruneParam:
     """Tests for the _prune_param function."""
 
@@ -391,6 +418,51 @@ class TestCompressSchema:
         assert "additionalProperties" not in result
         assert "$defs" not in result
 
+    def test_compresses_schema_nested_far_beyond_the_recursion_limit(self):
+        """Deeply nested schemas must compress rather than raise RecursionError.
+
+        Copying the schema is what sets the depth ceiling, so it must not
+        recurse: schemas this deep arrive from proxied or remote MCP servers,
+        and failing to compress them is worse than compressing them partially.
+        """
+        depth = sys.getrecursionlimit() * 2
+
+        schema: dict[str, Any] = {"type": "string", "title": "Leaf"}
+        for _ in range(depth):
+            schema = {"type": "array", "title": "Level", "items": schema}
+
+        original_depth = _measure_depth(schema)
+
+        result = compress_schema(schema, prune_titles=True)
+
+        assert result is not schema
+        assert _measure_depth(result) == original_depth
+        # The caller's schema is still intact at every level...
+        assert _count_titles(schema) == original_depth + 1
+        # ...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 = {