diff --git a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py index f2527ec1..cc52c0f2 100644 --- a/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py +++ b/atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py @@ -18,6 +18,7 @@ "boolean": bool, "array": list, "object": dict, + "null": type(None), } @@ -111,6 +112,13 @@ def json_to_pydantic_field( python_type = union_types[0] else: python_type = Union[tuple(union_types)] + # A JSON Schema type array is a union of its listed types. + elif isinstance(prop_schema.get("type"), list): + union_types = [ + SchemaTransformer.json_to_pydantic_field({**prop_schema, "type": json_type}, True, root_schema, model_cache)[0] + for json_type in prop_schema["type"] + ] + python_type = Union[tuple(union_types)] # Handle regular types else: json_type = prop_schema.get("type") @@ -119,15 +127,7 @@ def json_to_pydantic_field( if json_type == "array": items_schema = prop_schema.get("items", {}) - if "$ref" in items_schema: - item_type = SchemaTransformer._resolve_ref(items_schema["$ref"], root_schema, model_cache) - elif "oneOf" in items_schema or "anyOf" in items_schema: - # Handle arrays of unions - item_type, _ = SchemaTransformer.json_to_pydantic_field(items_schema, True, root_schema, model_cache) - elif items_schema.get("type") in JSON_TYPE_MAP: - item_type = JSON_TYPE_MAP[items_schema["type"]] - else: - item_type = Any + item_type, _ = SchemaTransformer.json_to_pydantic_field(items_schema, True, root_schema, model_cache) python_type = List[item_type] elif json_type == "object": diff --git a/atomic-agents/tests/connectors/mcp/test_mcp_factory.py b/atomic-agents/tests/connectors/mcp/test_mcp_factory.py index 1b158187..f13873cc 100644 --- a/atomic-agents/tests/connectors/mcp/test_mcp_factory.py +++ b/atomic-agents/tests/connectors/mcp/test_mcp_factory.py @@ -93,6 +93,33 @@ def test_fetch_mcp_tools_with_definitions_http(monkeypatch): assert tool_cls._has_typed_output_schema is False +@pytest.mark.parametrize("nullable_schema_location", ["input", "output"]) +def test_fetch_mcp_tools_preserves_nullable_schema(monkeypatch, nullable_schema_location): + """A tool must not disappear when its input or output contains a type array.""" + nullable_schema = { + "type": "object", + "properties": {"query": {"type": ["string", "null"]}}, + "required": ["query"], + } + definition = MCPToolDefinition( + name="SearchTool", + description="Search tool", + input_schema=nullable_schema if nullable_schema_location == "input" else {"type": "object"}, + output_schema=nullable_schema if nullable_schema_location == "output" else None, + ) + monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: [definition]) + + tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) + + assert len(tools) == 1 + tool_cls = tools[0] + if nullable_schema_location == "input": + assert tool_cls.input_schema(tool_name="SearchTool", query=None).query is None + else: + assert tool_cls._has_typed_output_schema is True + assert tool_cls.output_schema(query=None).query is None + + def test_fetch_mcp_tools_with_typed_output_schema(monkeypatch): """Test that tools with outputSchema get typed output models""" input_schema = {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]} diff --git a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py index cbfd4be5..0e2a4254 100644 --- a/atomic-agents/tests/connectors/mcp/test_schema_transformer.py +++ b/atomic-agents/tests/connectors/mcp/test_schema_transformer.py @@ -71,6 +71,77 @@ def test_no_type(self): class TestCreateModelFromSchema: + @pytest.mark.parametrize( + "prop_schema, valid_value", + [ + ({"type": "null"}, None), + ({"type": ["string", "null"]}, "query"), + ({"type": ["null", "string"]}, "query"), + ({"anyOf": [{"type": "string"}, {"type": "null"}]}, "query"), + ({"oneOf": [{"type": "string"}, {"type": "null"}]}, "query"), + ], + ) + def test_required_nullable_field(self, prop_schema, valid_value): + """A nullable MCP parameter still rejects other types and must be present.""" + schema = {"type": "object", "properties": {"query": prop_schema}, "required": ["query"]} + model = SchemaTransformer.create_model_from_schema(schema, "SearchInput", "search") + + assert model(tool_name="search", query=valid_value).query == valid_value + assert model(tool_name="search", query=None).query is None + with pytest.raises(ValueError): + model(tool_name="search", query={"unexpected": "object"}) + with pytest.raises(ValueError): + model(tool_name="search") + + @pytest.mark.parametrize("types", [["string"], ["string", "integer"]]) + def test_type_array_without_null(self, types): + """Type arrays preserve the allowed types without implicitly allowing null.""" + schema = {"type": "object", "properties": {"value": {"type": types}}, "required": ["value"]} + model = SchemaTransformer.create_model_from_schema(schema, "ValueOutput", "value", is_output_schema=True) + + assert model(value="text").value == "text" + if "integer" in types: + assert model(value=42).value == 42 + with pytest.raises(ValueError): + model(value=None) + with pytest.raises(ValueError): + model(value={}) + + @pytest.mark.parametrize("default", [None, "fallback"]) + def test_optional_nullable_field_default(self, default): + """Converting a type array retains the property's default and description.""" + schema = { + "type": "object", + "properties": {"query": {"type": ["string", "null"], "default": default, "description": "Search query"}}, + } + model = SchemaTransformer.create_model_from_schema(schema, "SearchInput", "search") + + assert model(tool_name="search").query == default + assert model.model_fields["query"].description == "Search query" + assert model(tool_name="search", query=None).query is None + + @pytest.mark.parametrize( + "items, valid, invalid", + [ + ({"type": ["string", "null"]}, ["text", None], [{}]), + ({"type": "null"}, [None], ["text"]), + ({"type": "array", "items": {"type": ["string", "null"]}}, [["text", None]], [[{}]]), + ], + ) + def test_nullable_array_items(self, items, valid, invalid): + """Nullable array types retain their item types at each array depth.""" + schema = { + "type": "object", + "properties": {"values": {"type": ["array", "null"], "items": items}}, + "required": ["values"], + } + model = SchemaTransformer.create_model_from_schema(schema, "ValuesOutput", "values", is_output_schema=True) + + assert model(values=valid).values == valid + assert model(values=None).values is None + with pytest.raises(ValueError): + model(values=invalid) + def test_basic_model_creation(self): schema = { "type": "object",