from unittest.mock import patch from jsonref import replace_refs from fastmcp.utilities.json_schema import ( _prune_param, _strip_remote_refs, compress_schema, dereference_refs, resolve_root_ref, ) class TestPruneParam: """Tests for the _prune_param function.""" def test_nonexistent(self): """Test pruning a parameter that doesn't exist.""" schema = {"properties": {"foo": {"type": "string"}}} result = _prune_param(schema, "bar") assert result == schema # Schema should be unchanged def test_exists(self): """Test pruning a parameter that exists.""" schema = {"properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}}} result = _prune_param(schema, "bar") assert result["properties"] == {"foo": {"type": "string"}} def test_last_property(self): """Test pruning the only/last parameter, should leave empty properties object.""" schema = {"properties": {"foo": {"type": "string"}}} result = _prune_param(schema, "foo") assert "properties" in result assert result["properties"] == {} def test_from_required(self): """Test pruning a parameter that's in the required list.""" schema = { "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}}, "required": ["foo", "bar"], } result = _prune_param(schema, "bar") assert result["required"] == ["foo"] def test_last_required(self): """Test pruning the last required parameter, should remove required field.""" schema = { "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}}, "required": ["foo"], } result = _prune_param(schema, "foo") assert "required" not in result class TestDereferenceRefs: """Tests for the dereference_refs function.""" def test_dereferences_simple_ref(self): """Test that simple $ref is dereferenced.""" schema = { "properties": { "foo": {"$ref": "#/$defs/foo_def"}, }, "$defs": { "foo_def": {"type": "string"}, }, } result = dereference_refs(schema) # $ref should be inlined assert result["properties"]["foo"] == {"type": "string"} # $defs should be removed assert "$defs" not in result def test_dereferences_nested_refs(self): """Test that nested $refs are dereferenced.""" schema = { "properties": { "foo": {"$ref": "#/$defs/foo_def"}, }, "$defs": { "foo_def": { "type": "object", "properties": {"nested": {"$ref": "#/$defs/nested_def"}}, }, "nested_def": {"type": "string"}, }, } result = dereference_refs(schema) # All refs should be inlined assert result["properties"]["foo"]["properties"]["nested"] == {"type": "string"} # $defs should be removed assert "$defs" not in result def test_falls_back_for_circular_refs(self): """Test that circular references fall back to resolve_root_ref.""" schema = { "$defs": { "Node": { "type": "object", "properties": { "children": { "type": "array", "items": {"$ref": "#/$defs/Node"}, } }, } }, "$ref": "#/$defs/Node", } result = dereference_refs(schema) # Should fall back to resolve_root_ref behavior # Root should be resolved but nested refs preserved assert result.get("type") == "object" assert "$defs" in result # $defs preserved for circular refs def test_preserves_sibling_keywords(self): """Test that sibling keywords (default, description) are preserved. Pydantic places description, default, examples as siblings to $ref. These should not be lost during dereferencing. """ schema = { "$defs": { "Status": {"type": "string", "enum": ["active", "inactive"]}, }, "properties": { "status": { "$ref": "#/$defs/Status", "default": "active", "description": "The user status", }, }, "type": "object", } result = dereference_refs(schema) # $ref should be inlined with siblings preserved status = result["properties"]["status"] assert status["type"] == "string" assert status["enum"] == ["active", "inactive"] assert status["default"] == "active" assert status["description"] == "The user status" # $defs should be removed assert "$defs" not in result def test_preserves_siblings_in_lists(self): """Test that siblings are preserved for $refs inside lists (allOf, anyOf, etc).""" schema = { "$defs": { "StringType": {"type": "string"}, "IntType": {"type": "integer"}, }, "properties": { "field": { "anyOf": [ {"$ref": "#/$defs/StringType", "description": "As string"}, {"$ref": "#/$defs/IntType", "description": "As integer"}, ] }, }, } result = dereference_refs(schema) # Both items in anyOf should have their siblings preserved any_of = result["properties"]["field"]["anyOf"] assert any_of[0]["type"] == "string" assert any_of[0]["description"] == "As string" assert any_of[1]["type"] == "integer" assert any_of[1]["description"] == "As integer" assert "$defs" not in result def test_preserves_nested_siblings(self): """Test that siblings on nested $refs are preserved.""" schema = { "$defs": { "Address": { "type": "object", "properties": { "country": {"$ref": "#/$defs/Country", "default": "US"}, }, }, "Country": {"type": "string", "enum": ["US", "UK", "CA"]}, }, "properties": { "home_address": {"$ref": "#/$defs/Address"}, }, } result = dereference_refs(schema) # The nested $ref's sibling (default) should be preserved country = result["properties"]["home_address"]["properties"]["country"] assert country["type"] == "string" assert country["enum"] == ["US", "UK", "CA"] assert country["default"] == "US" assert "$defs" not in result def test_strips_discriminator_mapping_after_inlining(self): """Discriminator.mapping refs dangle after $defs are inlined (#3679).""" schema = { "$defs": { "IdentifyPerson": { "type": "object", "properties": { "action": {"const": "identify", "type": "string"}, "name": {"type": "string"}, }, "required": ["action", "name"], }, "PersonDelete": { "type": "object", "properties": { "action": {"const": "delete", "type": "string"}, }, "required": ["action"], }, }, "anyOf": [ {"$ref": "#/$defs/IdentifyPerson"}, {"$ref": "#/$defs/PersonDelete"}, ], "discriminator": { "mapping": { "identify": "#/$defs/IdentifyPerson", "delete": "#/$defs/PersonDelete", }, "propertyName": "action", }, } result = dereference_refs(schema) assert "$defs" not in result assert "discriminator" not in result # The anyOf variants should be inlined with their const values intact assert len(result["anyOf"]) == 2 actions = {v["properties"]["action"]["const"] for v in result["anyOf"]} assert actions == {"identify", "delete"} def test_preserves_property_named_discriminator(self): """A field *named* 'discriminator' inside properties must survive.""" schema = { "$defs": { "Inner": { "type": "object", "properties": { "discriminator": {"type": "string"}, }, }, }, "properties": { "item": {"$ref": "#/$defs/Inner"}, }, } result = dereference_refs(schema) assert "$defs" not in result assert "discriminator" in result["properties"]["item"]["properties"] class TestCompressSchema: """Tests for the compress_schema function.""" def test_preserves_refs_by_default(self): """Test that compress_schema preserves $refs by default.""" schema = { "properties": { "foo": {"$ref": "#/$defs/foo_def"}, }, "$defs": { "foo_def": {"type": "string"}, }, } result = compress_schema(schema) # $ref should be preserved (dereferencing is handled by middleware) assert result["properties"]["foo"] == {"$ref": "#/$defs/foo_def"} assert "$defs" in result def test_prune_params(self): """Test pruning parameters with compress_schema.""" schema = { "properties": { "foo": {"type": "string"}, "bar": {"type": "integer"}, "baz": {"type": "boolean"}, }, "required": ["foo", "bar"], } result = compress_schema(schema, prune_params=["foo", "baz"]) assert result["properties"] == {"bar": {"type": "integer"}} assert result["required"] == ["bar"] def test_pruning_additional_properties(self): """Test pruning additionalProperties when explicitly enabled.""" schema = { "type": "object", "properties": {"foo": {"type": "string"}}, "additionalProperties": False, } # Must explicitly enable pruning now (default changed for MCP compatibility) result = compress_schema(schema, prune_additional_properties=True) assert "additionalProperties" not in result def test_disable_pruning_additional_properties(self): """Test disabling pruning of additionalProperties.""" schema = { "type": "object", "properties": {"foo": {"type": "string"}}, "additionalProperties": False, } result = compress_schema(schema, prune_additional_properties=False) assert "additionalProperties" in result assert result["additionalProperties"] is False def test_combined_operations(self): """Test all pruning operations together.""" schema = { "type": "object", "properties": { "keep": {"type": "string"}, "remove": {"$ref": "#/$defs/remove_def"}, }, "required": ["keep", "remove"], "additionalProperties": False, "$defs": { "remove_def": {"type": "string"}, "unused_def": {"type": "number"}, }, } result = compress_schema( schema, prune_params=["remove"], prune_additional_properties=True ) # Check that parameter was removed assert "remove" not in result["properties"] # Check that required list was updated assert result["required"] == ["keep"] # All $defs entries are now unreferenced after pruning "remove", so they're cleaned up assert "$defs" not in result # Check that additionalProperties was removed assert "additionalProperties" not in result def test_prune_titles(self): """Test pruning title fields.""" schema = { "title": "Root Schema", "type": "object", "properties": { "foo": {"title": "Foo Property", "type": "string"}, "bar": { "title": "Bar Property", "type": "object", "properties": { "nested": {"title": "Nested Property", "type": "string"} }, }, }, } result = compress_schema(schema, prune_titles=True) assert "title" not in result assert "title" not in result["properties"]["foo"] assert "title" not in result["properties"]["bar"] assert "title" not in result["properties"]["bar"]["properties"]["nested"] def test_prune_nested_additional_properties(self): """Test pruning additionalProperties: false at all levels when explicitly enabled.""" schema = { "type": "object", "additionalProperties": False, "properties": { "foo": { "type": "object", "additionalProperties": False, "properties": { "nested": { "type": "object", "additionalProperties": False, } }, }, }, } result = compress_schema(schema, prune_additional_properties=True) assert "additionalProperties" not in result assert "additionalProperties" not in result["properties"]["foo"] assert ( "additionalProperties" not in result["properties"]["foo"]["properties"]["nested"] ) def test_title_pruning_preserves_parameter_named_title(self): """Test that a parameter named 'title' is not removed during title pruning. This is a critical edge case - we want to remove title metadata but preserve actual parameters that happen to be named 'title'. """ from typing import Annotated from pydantic import Field, TypeAdapter def greet( name: Annotated[str, Field(description="The name to greet")], title: Annotated[str, Field(description="Optional title", default="")], ) -> str: """A greeting function.""" return f"Hello {title} {name}" adapter = TypeAdapter(greet) schema = adapter.json_schema() # Compress with title pruning compressed = compress_schema(schema, prune_titles=True) # The 'title' parameter should be preserved assert "title" in compressed["properties"] assert compressed["properties"]["title"]["description"] == "Optional title" assert compressed["properties"]["title"]["default"] == "" # But title metadata should be removed assert "title" not in compressed["properties"]["name"] assert "title" not in compressed["properties"]["title"] def test_title_pruning_preserves_title_property_when_type_property_exists(self): """Regression test for #3576: properties dict containing both 'title' and 'type' as parameter names caused the heuristic to treat 'title' as schema metadata and strip the entire property definition.""" schema = { "type": "object", "properties": { "dashboard_id": {"type": "string", "title": "Dashboard Id"}, "title": {"type": "string", "title": "Title"}, "type": {"type": "string", "title": "Type", "default": "vis"}, }, "required": ["dashboard_id", "title"], } compressed = compress_schema(schema, prune_titles=True) # All three properties must survive assert "dashboard_id" in compressed["properties"] assert "title" in compressed["properties"] assert "type" in compressed["properties"] # 'title' is still required assert "title" in compressed["required"] # But metadata title strings inside each property schema are removed assert "title" not in compressed["properties"]["dashboard_id"] assert "title" not in compressed["properties"]["title"] assert "title" not in compressed["properties"]["type"] def test_title_pruning_with_nested_properties(self): """Test that nested property structures are handled correctly.""" schema = { "type": "object", "title": "OuterObject", "properties": { "title": { # This is a property named "title", not metadata "type": "object", "title": "TitleObject", # This is metadata "properties": { "subtitle": { "type": "string", "title": "SubTitle", # This is metadata } }, }, "normal_field": { "type": "string", "title": "NormalField", # This is metadata }, }, } compressed = compress_schema(schema, prune_titles=True) # Root title should be removed assert "title" not in compressed # The property named "title" should be preserved assert "title" in compressed["properties"] # But its metadata title should be removed assert "title" not in compressed["properties"]["title"] # Nested metadata titles should be removed assert ( "title" not in compressed["properties"]["title"]["properties"]["subtitle"] ) assert "title" not in compressed["properties"]["normal_field"] def test_mcp_client_compatibility_requires_additional_properties(self): """Test that compress_schema preserves additionalProperties: false for MCP clients. MCP clients like Claude require strict JSON schemas with additionalProperties: false. When tools use Pydantic models with extra="forbid", this constraint must be preserved. Without this, MCP clients return: "Invalid schema for function 'X': In context=('properties', 'Y'), 'additionalProperties' is required to be supplied and to be false" See: https://github.com/PrefectHQ/fastmcp/issues/3008 """ # Schema representing a Pydantic model with extra="forbid" schema = { "type": "object", "properties": { "graph_table": { "type": "object", "properties": { "name": {"type": "string"}, "columns": {"type": "array", "items": {"type": "string"}}, }, "required": ["name"], "additionalProperties": False, } }, "required": ["graph_table"], "additionalProperties": False, } # By default, compress_schema should NOT strip additionalProperties: false # This is the new expected behavior for MCP compatibility result = compress_schema(schema) # Root level should preserve additionalProperties: false assert result.get("additionalProperties") is False, ( "Root additionalProperties: false was removed, breaking MCP compatibility" ) # Nested object should also preserve additionalProperties: false graph_table = result["properties"]["graph_table"] assert graph_table.get("additionalProperties") is False, ( "Nested additionalProperties: false was removed, breaking MCP compatibility" ) class TestCompressSchemaDereference: """Tests for the dereference parameter of compress_schema.""" SCHEMA_WITH_REFS = { "properties": { "foo": {"$ref": "#/$defs/foo_def"}, }, "$defs": { "foo_def": {"type": "string"}, }, } def test_dereference_true_inlines_refs(self): result = compress_schema(self.SCHEMA_WITH_REFS, dereference=True) assert result["properties"]["foo"] == {"type": "string"} assert "$defs" not in result def test_dereference_false_preserves_refs(self): result = compress_schema(self.SCHEMA_WITH_REFS, dereference=False) assert result["properties"]["foo"] == {"$ref": "#/$defs/foo_def"} assert "$defs" in result def test_other_optimizations_still_apply_without_dereference(self): schema = { "properties": { "foo": {"$ref": "#/$defs/foo_def"}, "bar": {"type": "integer", "title": "Bar"}, }, "$defs": { "foo_def": {"type": "string"}, }, } result = compress_schema( schema, dereference=False, prune_params=["bar"], prune_titles=True ) assert "bar" not in result["properties"] assert "$ref" in result["properties"]["foo"] assert "$defs" in result class TestResolveRootRef: """Tests for the resolve_root_ref function. This function resolves $ref at root level to meet MCP spec requirements. MCP specification requires outputSchema to have "type": "object" at root. """ def test_resolves_simple_root_ref(self): """Test that simple $ref at root is resolved.""" schema = { "$defs": { "Node": { "type": "object", "properties": { "id": {"type": "string"}, "name": {"type": "string"}, }, "required": ["id"], } }, "$ref": "#/$defs/Node", } result = resolve_root_ref(schema) # Should have type: object at root now assert result.get("type") == "object" assert "properties" in result assert "id" in result["properties"] assert "name" in result["properties"] # Should still have $defs for nested references assert "$defs" in result # Should NOT have $ref at root assert "$ref" not in result def test_resolves_self_referential_model(self): """Test resolving schema for self-referential models like Issue.""" # This is the exact schema Pydantic generates for self-referential models schema = { "$defs": { "Issue": { "type": "object", "properties": { "id": {"type": "string"}, "title": {"type": "string"}, "dependencies": { "type": "array", "items": {"$ref": "#/$defs/Issue"}, }, "dependents": { "type": "array", "items": {"$ref": "#/$defs/Issue"}, }, }, "required": ["id", "title"], } }, "$ref": "#/$defs/Issue", } result = resolve_root_ref(schema) # Should have type: object at root assert result.get("type") == "object" assert "properties" in result assert "id" in result["properties"] assert "dependencies" in result["properties"] # Nested $refs should still point to $defs assert result["properties"]["dependencies"]["items"]["$ref"] == "#/$defs/Issue" # Should have $defs preserved for nested references assert "$defs" in result assert "Issue" in result["$defs"] def test_does_not_modify_schema_with_type_at_root(self): """Test that schemas already having type at root are not modified.""" schema = { "type": "object", "properties": {"id": {"type": "string"}}, "$defs": {"SomeType": {"type": "string"}}, "$ref": "#/$defs/SomeType", # This would be unusual but possible } result = resolve_root_ref(schema) # Schema should be unchanged (returned as-is) assert result is schema def test_does_not_modify_schema_without_ref(self): """Test that schemas without $ref are not modified.""" schema = { "type": "object", "properties": {"id": {"type": "string"}}, } result = resolve_root_ref(schema) assert result is schema def test_does_not_modify_schema_without_defs(self): """Test that schemas with $ref but without $defs are not modified.""" schema = { "$ref": "#/$defs/Missing", } result = resolve_root_ref(schema) assert result is schema def test_does_not_modify_external_ref(self): """Test that external $refs (not pointing to $defs) are not resolved.""" schema = { "$defs": {"Node": {"type": "object"}}, "$ref": "https://example.com/schema.json#/definitions/Node", } result = resolve_root_ref(schema) assert result is schema def test_preserves_all_defs_for_nested_references(self): """Test that $defs are preserved even if multiple definitions exist.""" schema = { "$defs": { "Node": { "type": "object", "properties": { "child": {"$ref": "#/$defs/ChildNode"}, }, }, "ChildNode": { "type": "object", "properties": {"value": {"type": "string"}}, }, }, "$ref": "#/$defs/Node", } result = resolve_root_ref(schema) # Both defs should be preserved assert "$defs" in result assert "Node" in result["$defs"] assert "ChildNode" in result["$defs"] def test_handles_missing_def_gracefully(self): """Test that missing definition in $defs doesn't cause error.""" schema = { "$defs": {"OtherType": {"type": "string"}}, "$ref": "#/$defs/Missing", } result = resolve_root_ref(schema) # Should return original schema unchanged assert result is schema class TestStripRemoteRefs: """Tests for _strip_remote_refs which prevents SSRF/LFI via $ref.""" def test_preserves_local_ref(self): schema = {"$ref": "#/$defs/Foo"} assert _strip_remote_refs(schema) == {"$ref": "#/$defs/Foo"} def test_strips_http_ref(self): schema = {"$ref": "http://evil.com/schema.json"} assert _strip_remote_refs(schema) == {} def test_strips_https_ref(self): schema = {"$ref": "https://evil.com/schema.json"} assert _strip_remote_refs(schema) == {} def test_strips_file_ref(self): schema = {"$ref": "file:///etc/passwd"} assert _strip_remote_refs(schema) == {} def test_preserves_siblings_when_stripping(self): schema = { "$ref": "http://evil.com/schema.json", "description": "keep me", "default": 42, } result = _strip_remote_refs(schema) assert result == {"description": "keep me", "default": 42} def test_strips_nested_remote_refs(self): schema = { "properties": { "safe": {"$ref": "#/$defs/Safe"}, "evil": {"$ref": "http://169.254.169.254/latest/meta-data/"}, } } result = _strip_remote_refs(schema) assert result["properties"]["safe"] == {"$ref": "#/$defs/Safe"} assert "$ref" not in result["properties"]["evil"] def test_strips_remote_refs_in_lists(self): schema = { "anyOf": [ {"$ref": "#/$defs/Good"}, {"$ref": "file:///etc/credentials.json"}, ] } result = _strip_remote_refs(schema) assert result["anyOf"][0] == {"$ref": "#/$defs/Good"} assert "$ref" not in result["anyOf"][1] def test_deep_nesting(self): schema = { "properties": { "a": { "type": "object", "properties": {"b": {"$ref": "https://internal-service/secret"}}, } } } result = _strip_remote_refs(schema) assert "$ref" not in result["properties"]["a"]["properties"]["b"] class TestDereferenceRefsRemoteRefSafety: """Verify dereference_refs never fetches remote URIs.""" def test_http_ref_not_fetched(self): schema = { "type": "object", "properties": { "name": {"$ref": "http://evil.com/schema.json"}, }, } with patch( "fastmcp.utilities.json_schema.replace_refs", wraps=replace_refs ) as mock: result = dereference_refs(schema) # The remote $ref should have been stripped before replace_refs if mock.called: call_schema = mock.call_args[0][0] assert "$ref" not in call_schema.get("properties", {}).get("name", {}) # Result should not contain the remote $ref assert "$ref" not in result.get("properties", {}).get("name", {}) def test_file_ref_not_fetched(self): schema = { "type": "object", "properties": { "secret": {"$ref": "file:///etc/passwd"}, }, } result = dereference_refs(schema) assert "$ref" not in result.get("properties", {}).get("secret", {}) def test_cloud_metadata_ref_not_fetched(self): schema = { "type": "object", "properties": { "creds": { "$ref": "http://169.254.169.254/latest/meta-data/iam/security-credentials/" }, }, } result = dereference_refs(schema) assert "$ref" not in result.get("properties", {}).get("creds", {}) def test_local_refs_still_resolved(self): schema = { "$defs": {"Status": {"type": "string", "enum": ["a", "b"]}}, "type": "object", "properties": { "status": {"$ref": "#/$defs/Status"}, "evil": {"$ref": "https://evil.com/inject"}, }, } result = dereference_refs(schema) # Local ref should be resolved assert result["properties"]["status"] == {"type": "string", "enum": ["a", "b"]} # Remote ref should be stripped assert "$ref" not in result["properties"]["evil"] assert "$defs" not in result