Refactor TransformedTool

- Remove _find_referenced_defs method
- Update merging of defs in _merge_schema_with_precedence.
- Adjust tests in test_tool_transform.py to reflect these changes.

src/fastmcp/tools/tool_transform.py

@@ -787,45 +787,6 @@ class TransformedTool(Tool):
 
         return new_name, new_schema, is_required
 
-    @staticmethod
-    def _find_referenced_defs(
-        schema: dict[str, Any], available_defs: dict[str, Any]
-    ) -> dict[str, Any]:
-        """Find all $defs that are actually referenced in the schema.
-
-        Args:
-            schema: The schema to search for references
-            available_defs: All available definitions to check against
-
-        Returns:
-            Dictionary containing only the referenced definitions
-        """
-        referenced = set()
-        visited = set()  # Track visited definitions to prevent infinite recursion
-
-        def find_refs(obj):
-            if isinstance(obj, dict):
-                if "$ref" in obj:
-                    ref = obj["$ref"]
-                    if ref.startswith("#/$defs/"):
-                        def_name = ref[8:]  # Remove "#/$defs/" prefix
-                        if def_name in available_defs and def_name not in visited:
-                            referenced.add(def_name)
-                            visited.add(def_name)  # Mark as visited before recursing
-                            # Recursively check the referenced definition
-                            find_refs(available_defs[def_name])
-                for value in obj.values():
-                    find_refs(value)
-            elif isinstance(obj, list):
-                for item in obj:
-                    find_refs(item)
-
-        find_refs(schema)
-
-        return {
-            name: available_defs[name] for name in referenced if name in available_defs
-        }
-
     @staticmethod
     def _merge_schema_with_precedence(
         base_schema: dict[str, Any], override_schema: dict[str, Any]
@@ -883,7 +844,14 @@ class TransformedTool(Tool):
         # Merge $defs from both schemas, with override taking precedence
         merged_defs = base_schema.get("$defs", {}).copy()
         override_defs = override_schema.get("$defs", {})
-        merged_defs.update(override_defs)
+
+        for def_name, def_schema in override_defs.items():
+            if def_name in merged_defs:
+                base_def = merged_defs[def_name].copy()
+                base_def.update(def_schema)
+                merged_defs[def_name] = base_def
+            else:
+                merged_defs[def_name] = def_schema.copy()
 
         result = {
             "type": "object",
@@ -891,11 +859,9 @@ class TransformedTool(Tool):
             "required": list(final_required),
         }
 
-        # Only include $defs that are actually referenced in the schema
         if merged_defs:
-            referenced_defs = TransformedTool._find_referenced_defs(result, merged_defs)
-            if referenced_defs:
-                result["$defs"] = referenced_defs
+            result["$defs"] = merged_defs
+            result = compress_schema(result, prune_defs=True)
 
         return result
 

tests/tools/test_tool_transform.py

@@ -1508,84 +1508,6 @@ def test_tool_transform_config_removes_meta(sample_tool):
 class TestDefsAndReferences:
     """Test schema definition handling and reference finding."""
 
-    def test_find_referenced_defs_simple_reference(self):
-        """Test _find_referenced_defs with a simple reference."""
-        schema = {"type": "object", "properties": {"field1": {"$ref": "#/$defs/TypeA"}}}
-        available_defs = {
-            "TypeA": {"type": "string"},
-            "TypeB": {"type": "integer"},  # Not referenced
-        }
-
-        result = TransformedTool._find_referenced_defs(schema, available_defs)
-        assert result == {"TypeA": {"type": "string"}}
-        assert "TypeB" not in result
-
-    def test_find_referenced_defs_nested_references(self):
-        """Test _find_referenced_defs with nested references."""
-        schema = {"type": "object", "properties": {"field1": {"$ref": "#/$defs/TypeA"}}}
-        available_defs = {
-            "TypeA": {
-                "type": "object",
-                "properties": {"nested": {"$ref": "#/$defs/TypeB"}},
-            },
-            "TypeB": {"type": "string"},
-            "TypeC": {"type": "integer"},  # Not referenced
-        }
-
-        result = TransformedTool._find_referenced_defs(schema, available_defs)
-        assert result == {
-            "TypeA": {
-                "type": "object",
-                "properties": {"nested": {"$ref": "#/$defs/TypeB"}},
-            },
-            "TypeB": {"type": "string"},
-        }
-        assert "TypeC" not in result
-
-    def test_find_referenced_defs_circular_references(self):
-        """Test _find_referenced_defs handles circular references."""
-        schema = {"type": "object", "properties": {"field1": {"$ref": "#/$defs/TypeA"}}}
-        available_defs = {
-            "TypeA": {
-                "type": "object",
-                "properties": {"circular": {"$ref": "#/$defs/TypeB"}},
-            },
-            "TypeB": {
-                "type": "object",
-                "properties": {"back_ref": {"$ref": "#/$defs/TypeA"}},
-            },
-            "TypeC": {"type": "string"},  # Not referenced
-        }
-
-        result = TransformedTool._find_referenced_defs(schema, available_defs)
-        assert "TypeA" in result
-        assert "TypeB" in result
-        assert "TypeC" not in result
-
-    def test_find_referenced_defs_array_references(self):
-        """Test _find_referenced_defs with references in arrays."""
-        schema = {
-            "type": "object",
-            "properties": {
-                "field1": {"type": "array", "items": {"$ref": "#/$defs/TypeA"}}
-            },
-        }
-        available_defs = {
-            "TypeA": {"type": "string"},
-            "TypeB": {"type": "integer"},  # Not referenced
-        }
-
-        result = TransformedTool._find_referenced_defs(schema, available_defs)
-        assert result == {"TypeA": {"type": "string"}}
-
-    def test_find_referenced_defs_no_references(self):
-        """Test _find_referenced_defs with no references."""
-        schema = {"type": "object", "properties": {"field1": {"type": "string"}}}
-        available_defs = {"TypeA": {"type": "string"}, "TypeB": {"type": "integer"}}
-
-        result = TransformedTool._find_referenced_defs(schema, available_defs)
-        assert result == {}
-
     def test_merge_schema_with_defs_precedence(self):
         """Test _merge_schema_with_precedence merges $defs correctly."""
         base_schema = {