From 68f264ec0186cbc3d7f2a31d1ca13a971c210244 Mon Sep 17 00:00:00 2001 From: pldesch-chift Date: Fri, 22 Aug 2025 10:53:59 +0200 Subject: [PATCH] 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}"