fix(jsonschema): preserve multi-branch anyOf/oneOf unions

jsonschema/jsonschema.py

@@ -1,5 +1,5 @@
 # /// zerodep
-# version = "0.4.0"
+# version = "0.4.1"
 # deps = []
 # tier = "subsystem"
 # category = "validation"
@@ -407,8 +407,7 @@ def _simplify_node(schema: dict[str, Any]) -> dict[str, Any]:
         if len(non_null) == 1:
             base = _deep_merge_two(base, non_null[0])
         elif len(non_null) > 1:
-            # Lossy but safe for LLM tool schemas: keep first non-null variant.
-            base = _deep_merge_two(base, non_null[0])
+            base[keyword] = non_null
         # else: all null — base stays as-is
 
         if has_null:
@@ -451,8 +450,8 @@ def simplify_unions(schema: dict[str, Any]) -> dict[str, Any]:
 
     - Nullable pattern ``[{type: T}, {type: null}]`` → ``{type: T, nullable: true}``
     - Single-variant: unwrap.
-    - Multi-variant: keep first non-null variant (lossy but safe for LLM tool
-      schemas).
+    - Multi-variant with null: strip the ``{type: null}`` branch, keep remaining
+      branches under the original keyword (``anyOf`` or ``oneOf``).
 
     Args:
         schema: A JSON Schema dict.

jsonschema/test_jsonschema_correctness.py

@@ -574,16 +574,16 @@ def test_single_variant_no_null(self):
         assert "nullable" not in result
         assert "anyOf" not in result
 
-    def test_multi_variant_keeps_first(self):
+    def test_multi_variant_preserved(self):
         schema = {
             "anyOf": [
                 {"type": "string"},
                 {"type": "integer"},
             ]
         }
         result = simplify_unions(schema)
-        assert result["type"] == "string"
-        assert "anyOf" not in result
+        assert result["anyOf"] == [{"type": "string"}, {"type": "integer"}]
+        assert "nullable" not in result
 
     def test_multi_variant_with_null(self):
         schema = {
@@ -594,9 +594,40 @@ def test_multi_variant_with_null(self):
             ]
         }
         result = simplify_unions(schema)
-        assert result["type"] == "string"
+        assert result["anyOf"] == [{"type": "string"}, {"type": "integer"}]
+        assert result["nullable"] is True
+
+    def test_multi_variant_number_and_array(self):
+        schema = {
+            "anyOf": [
+                {"type": "number"},
+                {"type": "array", "items": {"type": "number"}},
+                {"type": "null"},
+            ]
+        }
+        result = simplify_unions(schema)
+        assert result["anyOf"] == [
+            {"type": "number"},
+            {"type": "array", "items": {"type": "number"}},
+        ]
         assert result["nullable"] is True
 
+    def test_multi_variant_no_null(self):
+        schema = {
+            "oneOf": [
+                {"type": "string"},
+                {"type": "number"},
+                {"type": "boolean"},
+            ]
+        }
+        result = simplify_unions(schema)
+        assert result["oneOf"] == [
+            {"type": "string"},
+            {"type": "number"},
+            {"type": "boolean"},
+        ]
+        assert "nullable" not in result
+
     def test_all_null_variants(self):
         schema = {"anyOf": [{"type": "null"}, {"type": "null"}]}
         result = simplify_unions(schema)