From 68f264ec0186cbc3d7f2a31d1ca13a971c210244 Mon Sep 17 00:00:00 2001
From: pldesch-chift <pl.deschaetzen@chift.eu>
Date: Fri, 22 Aug 2025 10:53:59 +0200
Subject: [PATCH 1/3] Implement _find_referenced_defs method for schema
 reference resolution and enhance _merge_schema_with_precedence to prune
 unused $defs. Add comprehensive tests for new functionality in
 test_tool_transform.py.

---
 src/fastmcp/tools/tool_transform.py |  54 ++++++-
 tests/tools/test_tool_transform.py  | 221 ++++++++++++++++++++++++++++
 2 files changed, 274 insertions(+), 1 deletion(-)

diff --git a/src/fastmcp/tools/tool_transform.py b/src/fastmcp/tools/tool_transform.py
index 063c32361..e8e5f88f7 100644
--- a/src/fastmcp/tools/tool_transform.py
+++ b/src/fastmcp/tools/tool_transform.py
@@ -787,6 +787,45 @@ 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]
@@ -841,12 +880,25 @@ class TransformedTool(Tool):
             if "default" in param_schema:
                 final_required.discard(param_name)
 
-        return {
+        # 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)
+
+        result = {
             "type": "object",
             "properties": merged_props,
             "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
+
+        return result
+
     @staticmethod
     def _function_has_kwargs(fn: Callable[..., Any]) -> bool:
         """Check if function accepts **kwargs.
diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py
index f0b21bded..ecb6c2997 100644
--- a/tests/tools/test_tool_transform.py
+++ b/tests/tools/test_tool_transform.py
@@ -1502,3 +1502,224 @@ def test_tool_transform_config_removes_meta(sample_tool):
     config = ToolTransformConfig(name="config_tool", meta=None)
     transformed = config.apply(sample_tool)
     assert transformed.meta is None
+
+
+# Tests for $defs and _find_referenced_defs functionality
+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 = {
+            "type": "object",
+            "properties": {"field1": {"$ref": "#/$defs/BaseType"}},
+            "$defs": {
+                "BaseType": {"type": "string", "description": "base"},
+                "SharedType": {"type": "integer", "minimum": 0},
+            },
+        }
+
+        override_schema = {
+            "type": "object",
+            "properties": {"field2": {"$ref": "#/$defs/OverrideType"}},
+            "$defs": {
+                "OverrideType": {"type": "boolean"},
+                "SharedType": {"type": "integer", "minimum": 10},  # Override
+            },
+        }
+
+        result = TransformedTool._merge_schema_with_precedence(
+            base_schema, override_schema
+        )
+
+        # Should have both field1 and field2
+        assert "field1" in result["properties"]
+        assert "field2" in result["properties"]
+
+        # Should only include referenced defs
+        defs = result.get("$defs", {})
+        assert "BaseType" in defs  # Referenced by field1
+        assert "OverrideType" in defs  # Referenced by field2
+
+        # SharedType should use override version, but only if referenced
+        # Since it's not referenced by either field, it shouldn't be included
+        assert "SharedType" not in defs
+
+    def test_transform_tool_with_complex_defs_pruning(self):
+        """Test that tool transformation properly prunes unused $defs."""
+
+        class UsedType(BaseModel):
+            value: str
+
+        class UnusedType(BaseModel):
+            other: int
+
+        @Tool.from_function
+        def complex_tool(
+            used_param: UsedType, unused_param: UnusedType | None = None
+        ) -> str:
+            return used_param.value
+
+        # Transform to hide unused_param
+        transformed = Tool.from_tool(
+            complex_tool, transform_args={"unused_param": ArgTransform(hide=True)}
+        )
+
+        # Only UsedType should be in $defs, not UnusedType
+        defs = transformed.parameters.get("$defs", {})
+        type_names = set(defs.keys())
+
+        # Should contain UsedType but not UnusedType
+        used_type_found = any("UsedType" in name for name in type_names)
+        unused_type_found = any("UnusedType" in name for name in type_names)
+
+        assert used_type_found, f"UsedType not found in defs: {type_names}"
+        assert not unused_type_found, f"UnusedType should not be in defs: {type_names}"
+
+    def test_transform_with_custom_function_preserves_needed_defs(self):
+        """Test that custom transform functions preserve necessary $defs."""
+
+        class InputType(BaseModel):
+            data: str
+
+        class OutputType(BaseModel):
+            result: str
+
+        @Tool.from_function
+        def base_tool(input_data: InputType) -> OutputType:
+            return OutputType(result=input_data.data.upper())
+
+        async def transform_function(renamed_input: InputType):
+            return await forward(renamed_input=renamed_input)
+
+        # Transform with custom function and argument rename
+        transformed = Tool.from_tool(
+            base_tool,
+            transform_fn=transform_function,
+            transform_args={"input_data": ArgTransform(name="renamed_input")},
+        )
+
+        # Both InputType and OutputType should be preserved in defs
+        defs = transformed.parameters.get("$defs", {})
+        type_names = set(defs.keys())
+
+        input_type_found = any("InputType" in name for name in type_names)
+        assert input_type_found, f"InputType not found in defs: {type_names}"
+
+    def test_chained_transforms_preserve_correct_defs(self):
+        """Test that chained transformations preserve correct $defs."""
+
+        class TypeA(BaseModel):
+            a: str
+
+        class TypeB(BaseModel):
+            b: int
+
+        class TypeC(BaseModel):
+            c: bool
+
+        @Tool.from_function
+        def base_tool(param_a: TypeA, param_b: TypeB, param_c: TypeC) -> str:
+            return f"{param_a.a}-{param_b.b}-{param_c.c}"
+
+        # First transform: hide param_c
+        transform1 = Tool.from_tool(
+            base_tool,
+            transform_args={"param_c": ArgTransform(hide=True, default=TypeC(c=True))},
+        )
+
+        # Second transform: hide param_b
+        transform2 = Tool.from_tool(
+            transform1,
+            transform_args={"param_b": ArgTransform(hide=True, default=TypeB(b=42))},
+        )
+
+        # Final schema should only have TypeA in $defs
+        defs = transform2.parameters.get("$defs", {})
+        type_names = set(defs.keys())
+
+        type_a_found = any("TypeA" in name for name in type_names)
+        type_b_found = any("TypeB" in name for name in type_names)
+        type_c_found = any("TypeC" in name for name in type_names)
+
+        assert type_a_found, f"TypeA should be in defs: {type_names}"
+        assert not type_b_found, f"TypeB should not be in defs: {type_names}"
+        assert not type_c_found, f"TypeC should not be in defs: {type_names}"

From 6575768d4d7f3d2902dae8b7f98535effe4c1f33 Mon Sep 17 00:00:00 2001
From: pldesch-chift <pl.deschaetzen@chift.eu>
Date: Mon, 25 Aug 2025 09:33:16 +0200
Subject: [PATCH 2/3] 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 | 54 ++++----------------
 tests/tools/test_tool_transform.py  | 78 -----------------------------
 2 files changed, 10 insertions(+), 122 deletions(-)

diff --git a/src/fastmcp/tools/tool_transform.py b/src/fastmcp/tools/tool_transform.py
index e8e5f88f7..56e0bfd82 100644
--- a/src/fastmcp/tools/tool_transform.py
+++ b/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
 
diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py
index ecb6c2997..12058c904 100644
--- a/tests/tools/test_tool_transform.py
+++ b/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 = {

From 0114f0ad3c08bf18bd1784ee2faca75d43020a0b Mon Sep 17 00:00:00 2001
From: William Easton <bill.easton@elastic.co>
Date: Mon, 25 Aug 2025 10:18:39 -0500
Subject: [PATCH 3/3] PR Clean-up

---
 tests/tools/test_tool_transform.py | 206 +++++++++++++++++++----------
 1 file changed, 134 insertions(+), 72 deletions(-)

diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py
index 12058c904..d01cf666d 100644
--- a/tests/tools/test_tool_transform.py
+++ b/tests/tools/test_tool_transform.py
@@ -4,6 +4,7 @@ from typing import Annotated, Any
 
 import pytest
 from dirty_equals import IsList
+from inline_snapshot import snapshot
 from mcp.types import TextContent
 from pydantic import BaseModel, Field, TypeAdapter
 from typing_extensions import TypedDict
@@ -1055,38 +1056,6 @@ class TestEnableDisable:
                 await client.call_tool("new_add", {"x": 1, "y": 2})
 
 
-def test_arg_transform_examples_in_schema(add_tool):
-    # Simple example
-    new_tool = Tool.from_tool(
-        add_tool,
-        transform_args={
-            "old_x": ArgTransform(examples=[1, 2, 3]),
-        },
-    )
-    prop = get_property(new_tool, "old_x")
-    assert prop["examples"] == [1, 2, 3]
-
-    # Nested example (e.g., for array type)
-    new_tool2 = Tool.from_tool(
-        add_tool,
-        transform_args={
-            "old_x": ArgTransform(examples=[["a", "b"], ["c", "d"]]),
-        },
-    )
-    prop2 = get_property(new_tool2, "old_x")
-    assert prop2["examples"] == [["a", "b"], ["c", "d"]]
-
-    # If not set, should not be present
-    new_tool3 = Tool.from_tool(
-        add_tool,
-        transform_args={
-            "old_x": ArgTransform(),
-        },
-    )
-    prop3 = get_property(new_tool3, "old_x")
-    assert "examples" not in prop3
-
-
 class TestTransformToolOutputSchema:
     """Test output schema handling in transformed tools."""
 
@@ -1504,10 +1473,40 @@ def test_tool_transform_config_removes_meta(sample_tool):
     assert transformed.meta is None
 
 
-# Tests for $defs and _find_referenced_defs functionality
-class TestDefsAndReferences:
+class TestInputSchema:
     """Test schema definition handling and reference finding."""
 
+    def test_arg_transform_examples_in_schema(self, add_tool: Tool):
+        # Simple example
+        new_tool = Tool.from_tool(
+            add_tool,
+            transform_args={
+                "old_x": ArgTransform(examples=[1, 2, 3]),
+            },
+        )
+        prop = get_property(new_tool, "old_x")
+        assert prop["examples"] == [1, 2, 3]
+
+        # Nested example (e.g., for array type)
+        new_tool2 = Tool.from_tool(
+            add_tool,
+            transform_args={
+                "old_x": ArgTransform(examples=[["a", "b"], ["c", "d"]]),
+            },
+        )
+        prop2 = get_property(new_tool2, "old_x")
+        assert prop2["examples"] == [["a", "b"], ["c", "d"]]
+
+        # If not set, should not be present
+        new_tool3 = Tool.from_tool(
+            add_tool,
+            transform_args={
+                "old_x": ArgTransform(),
+            },
+        )
+        prop3 = get_property(new_tool3, "old_x")
+        assert "examples" not in prop3
+
     def test_merge_schema_with_defs_precedence(self):
         """Test _merge_schema_with_precedence merges $defs correctly."""
         base_schema = {
@@ -1528,22 +1527,27 @@ class TestDefsAndReferences:
             },
         }
 
-        result = TransformedTool._merge_schema_with_precedence(
+        transformed_tool_schema = TransformedTool._merge_schema_with_precedence(
             base_schema, override_schema
         )
 
-        # Should have both field1 and field2
-        assert "field1" in result["properties"]
-        assert "field2" in result["properties"]
+        # SharedType should no longer be present on the schema
+        assert "SharedType" not in transformed_tool_schema["$defs"]
 
-        # Should only include referenced defs
-        defs = result.get("$defs", {})
-        assert "BaseType" in defs  # Referenced by field1
-        assert "OverrideType" in defs  # Referenced by field2
-
-        # SharedType should use override version, but only if referenced
-        # Since it's not referenced by either field, it shouldn't be included
-        assert "SharedType" not in defs
+        assert transformed_tool_schema == snapshot(
+            {
+                "type": "object",
+                "properties": {
+                    "field1": {"$ref": "#/$defs/BaseType"},
+                    "field2": {"$ref": "#/$defs/OverrideType"},
+                },
+                "required": [],
+                "$defs": {
+                    "BaseType": {"type": "string", "description": "base"},
+                    "OverrideType": {"type": "boolean"},
+                },
+            }
+        )
 
     def test_transform_tool_with_complex_defs_pruning(self):
         """Test that tool transformation properly prunes unused $defs."""
@@ -1561,20 +1565,29 @@ class TestDefsAndReferences:
             return used_param.value
 
         # Transform to hide unused_param
-        transformed = Tool.from_tool(
+        transformed_tool: TransformedTool = Tool.from_tool(
             complex_tool, transform_args={"unused_param": ArgTransform(hide=True)}
         )
 
-        # Only UsedType should be in $defs, not UnusedType
-        defs = transformed.parameters.get("$defs", {})
-        type_names = set(defs.keys())
+        assert "UnusedType" not in transformed_tool.parameters["$defs"]
 
-        # Should contain UsedType but not UnusedType
-        used_type_found = any("UsedType" in name for name in type_names)
-        unused_type_found = any("UnusedType" in name for name in type_names)
-
-        assert used_type_found, f"UsedType not found in defs: {type_names}"
-        assert not unused_type_found, f"UnusedType should not be in defs: {type_names}"
+        assert transformed_tool.parameters == snapshot(
+            {
+                "type": "object",
+                "properties": {
+                    "used_param": {"$ref": "#/$defs/UsedType", "title": "Used Param"}
+                },
+                "required": ["used_param"],
+                "$defs": {
+                    "UsedType": {
+                        "properties": {"value": {"title": "Value", "type": "string"}},
+                        "required": ["value"],
+                        "title": "UsedType",
+                        "type": "object",
+                    }
+                },
+            }
+        )
 
     def test_transform_with_custom_function_preserves_needed_defs(self):
         """Test that custom transform functions preserve necessary $defs."""
@@ -1599,12 +1612,26 @@ class TestDefsAndReferences:
             transform_args={"input_data": ArgTransform(name="renamed_input")},
         )
 
-        # Both InputType and OutputType should be preserved in defs
-        defs = transformed.parameters.get("$defs", {})
-        type_names = set(defs.keys())
-
-        input_type_found = any("InputType" in name for name in type_names)
-        assert input_type_found, f"InputType not found in defs: {type_names}"
+        assert transformed.parameters == snapshot(
+            {
+                "type": "object",
+                "properties": {
+                    "renamed_input": {
+                        "$ref": "#/$defs/InputType",
+                        "title": "Input Data",
+                    }
+                },
+                "required": ["renamed_input"],
+                "$defs": {
+                    "InputType": {
+                        "properties": {"data": {"title": "Data", "type": "string"}},
+                        "required": ["data"],
+                        "title": "InputType",
+                        "type": "object",
+                    }
+                },
+            }
+        )
 
     def test_chained_transforms_preserve_correct_defs(self):
         """Test that chained transformations preserve correct $defs."""
@@ -1628,20 +1655,55 @@ class TestDefsAndReferences:
             transform_args={"param_c": ArgTransform(hide=True, default=TypeC(c=True))},
         )
 
+        assert transform1.parameters == snapshot(
+            {
+                "type": "object",
+                "properties": {
+                    "param_a": {"$ref": "#/$defs/TypeA", "title": "Param A"},
+                    "param_b": {"$ref": "#/$defs/TypeB", "title": "Param B"},
+                },
+                "required": IsList("param_b", "param_a", check_order=False),
+                "$defs": {
+                    "TypeA": {
+                        "properties": {"a": {"title": "A", "type": "string"}},
+                        "required": ["a"],
+                        "title": "TypeA",
+                        "type": "object",
+                    },
+                    "TypeB": {
+                        "properties": {"b": {"title": "B", "type": "integer"}},
+                        "required": ["b"],
+                        "title": "TypeB",
+                        "type": "object",
+                    },
+                },
+            }
+        )
+
+        assert "TypeA" in transform1.parameters["$defs"]
+
         # Second transform: hide param_b
         transform2 = Tool.from_tool(
             transform1,
             transform_args={"param_b": ArgTransform(hide=True, default=TypeB(b=42))},
         )
 
-        # Final schema should only have TypeA in $defs
-        defs = transform2.parameters.get("$defs", {})
-        type_names = set(defs.keys())
+        assert "TypeB" not in transform2.parameters["$defs"]
 
-        type_a_found = any("TypeA" in name for name in type_names)
-        type_b_found = any("TypeB" in name for name in type_names)
-        type_c_found = any("TypeC" in name for name in type_names)
-
-        assert type_a_found, f"TypeA should be in defs: {type_names}"
-        assert not type_b_found, f"TypeB should not be in defs: {type_names}"
-        assert not type_c_found, f"TypeC should not be in defs: {type_names}"
+        assert transform2.parameters == snapshot(
+            {
+                "type": "object",
+                "properties": {
+                    "param_a": {"$ref": "#/$defs/TypeA", "title": "Param A"}
+                },
+                "required": ["param_a"],
+                "$defs": {
+                    "TypeA": {
+                        "properties": {"a": {"title": "A", "type": "string"}},
+                        "required": ["a"],
+                        "title": "TypeA",
+                        "type": "object",
+                    }
+                },
+            }
+        )