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 = {