diff --git a/fastmcp_slim/fastmcp/utilities/json_schema.py b/fastmcp_slim/fastmcp/utilities/json_schema.py index 2cd604e78..59715e8ed 100644 --- a/fastmcp_slim/fastmcp/utilities/json_schema.py +++ b/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", {}) @@ -498,6 +531,11 @@ def _single_pass_optimize( if not (prune_defs or prune_titles or prune_additional_properties): return schema # Nothing to do + # 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_schema(schema) + # Phase 1: Collect references and apply simple cleanups # Track which $defs are referenced from the main schema and from other $defs root_refs: set[str] = set() # $defs referenced directly from main schema @@ -506,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, @@ -523,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): @@ -647,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 1620127e1..b97153b02 100644 --- a/tests/utilities/test_json_schema.py +++ b/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.""" @@ -358,6 +385,84 @@ class TestDereferenceRefs: class TestCompressSchema: """Tests for the compress_schema function.""" + def test_does_not_mutate_input(self): + """compress_schema must return a new dict and leave the caller's schema + untouched, even when it prunes titles, additionalProperties and unused + $defs (a live Tool.input_schema is passed straight in at some call sites).""" + schema = { + "type": "object", + "title": "MySchema", + "additionalProperties": False, + "properties": { + "a": {"type": "string", "title": "A"}, + "b": { + "type": "object", + "title": "B", + "properties": {"c": {"type": "integer", "title": "C"}}, + }, + }, + "$defs": {"Unused": {"type": "string", "title": "Unused"}}, + } + original = copy.deepcopy(schema) + + result = compress_schema( + schema, prune_titles=True, prune_additional_properties=True + ) + + # The input is untouched... + assert schema == original + assert result is not schema + # ...and the returned copy really was optimized (so it is not a no-op). + assert "title" not in result + assert "title" not in result["properties"]["b"]["properties"]["c"] + 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 = {