Preserve nullable types in MCP tool schemas

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":

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"]}

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",