diff --git a/src/fastmcp/utilities/json_schema.py b/src/fastmcp/utilities/json_schema.py index ca5fe37c9..29bb36aa3 100644 --- a/src/fastmcp/utilities/json_schema.py +++ b/src/fastmcp/utilities/json_schema.py @@ -67,13 +67,24 @@ def _prune_unused_defs(schema: dict) -> dict: walk(value, current_def=def_name) # Figure out what defs were referenced directly or recursively - def def_is_referenced(def_name): + def def_is_referenced(def_name, parent_def_names: set[str] | None = None): if def_name in root_defs: return True references = referenced_by.get(def_name) if references: - for reference in references: - if def_is_referenced(reference): + if parent_def_names is None: + parent_def_names = set() + + # Handle recursion by excluding references already present in parent references + parent_def_names = parent_def_names | {def_name} + valid_references = [ + reference + for reference in references + if reference not in parent_def_names + ] + + for reference in valid_references: + if def_is_referenced(reference, parent_def_names): return True return False diff --git a/tests/utilities/test_json_schema.py b/tests/utilities/test_json_schema.py index 55c970224..420567dc5 100644 --- a/tests/utilities/test_json_schema.py +++ b/tests/utilities/test_json_schema.py @@ -107,6 +107,79 @@ class TestPruneUnusedDefs: result = _prune_unused_defs(schema) assert "$defs" not in result + def test_nested_references_with_recursion_kept(self): + """Test that definitions with recursion referenced via nesting are kept.""" + schema = { + "properties": { + "foo": {"$ref": "#/$defs/foo_def"}, + }, + "$defs": { + "foo_def": { + "type": "object", + "properties": {"nested": {"$ref": "#/$defs/foo_def"}}, + }, + "unused_def": {"type": "integer"}, + }, + } + result = _prune_unused_defs(schema) + assert "foo_def" in result["$defs"] + assert "unused_def" not in result["$defs"] + + def test_nested_references_with_recursion_removed(self): + """Test that definitions with recursion referenced via nesting in unused defs are removed.""" + schema = { + "properties": {}, + "$defs": { + "foo_def": { + "type": "object", + "properties": {"nested": {"$ref": "#/$defs/foo_def"}}, + }, + }, + } + result = _prune_unused_defs(schema) + assert "$defs" not in result + + def test_multiple_nested_references_with_recursion_kept(self): + """Test that definitions with multiple levels of recursion referenced via nesting are kept.""" + schema = { + "properties": { + "foo": {"$ref": "#/$defs/foo_def"}, + }, + "$defs": { + "foo_def": { + "type": "object", + "properties": {"nested": {"$ref": "#/$defs/nested_def"}}, + }, + "nested_def": { + "type": "object", + "properties": {"nested": {"$ref": "#/$defs/foo_def"}}, + }, + "unused_def": {"type": "integer"}, + }, + } + result = _prune_unused_defs(schema) + assert "foo_def" in result["$defs"] + assert "nested_def" in result["$defs"] + assert "unused_def" not in result["$defs"] + + def test_multiple_nested_references_with_recursion_removed(self): + """Test that definitions with multiple levels of recursion referenced via nesting in unused defs are removed.""" + schema = { + "properties": {}, + "$defs": { + "foo_def": { + "type": "object", + "properties": {"nested": {"$ref": "#/$defs/nested_def"}}, + }, + "nested_def": { + "type": "object", + "properties": {"nested": {"$ref": "#/$defs/foo_def"}}, + }, + }, + } + result = _prune_unused_defs(schema) + assert "$defs" not in result + def test_array_references_kept(self): """Test that definitions referenced in array items are kept.""" schema = {