fix(adk-middleware): map oneOf → anyOf so genai keeps discriminated-union tool params

integrations/adk-middleware/python/src/ag_ui_adk/client_proxy_tool.py

@@ -61,11 +61,12 @@
 def _clean_schema_for_genai(schema: Any) -> Any:
     """Recursively clean a JSON Schema dict for google.genai.types.Schema.
 
-    Three transformations:
+    Transformations:
     1. Strip ``$``-prefixed keys (``$schema``, ``$id``, ``$ref``, ``$defs``,
        ``$comment``) — these are JSON Schema infrastructure, never in genai.
-    2. Map ``examples`` → ``example`` (first element only) and
-       ``const`` → ``enum`` (single-value list), preserving useful context.
+    2. Map ``examples`` → ``example`` (first element only), ``const`` →
+       ``enum`` (single-value list), and ``oneOf`` → ``anyOf`` (genai accepts
+       ``anyOf`` but not ``oneOf``), preserving useful structure.
     3. Filter remaining keys to only those accepted by ``types.Schema``,
        using an allowlist derived from ``types.Schema.model_fields``.
     """
@@ -87,6 +88,16 @@ def _clean_schema_for_genai(schema: Any) -> Any:
             if k == "const":
                 result["enum"] = [v if isinstance(v, str) else json.dumps(v)]
                 continue
+            # Map oneOf -> anyOf. genai.types.Schema accepts ``anyOf`` but not
+            # ``oneOf``, so an unmapped ``oneOf`` would be dropped by the
+            # allowlist below — silently erasing the structure of discriminated
+            # unions (zod ``discriminatedUnion`` serializes to ``oneOf`` via
+            # zod-to-json-schema, used by CopilotKit / AG-UI frontend tools).
+            # For tool-argument schemas the distinction is immaterial: the model
+            # emits exactly one branch either way.
+            if k == "oneOf":
+                result["anyOf"] = _clean_schema_for_genai(v)
+                continue
             # Only keep keys that genai.types.Schema accepts
             if k not in _ALLOWED_SCHEMA_KEYS:
                 continue

integrations/adk-middleware/python/tests/test_client_proxy_tool.py

@@ -600,6 +600,61 @@ def test_maps_const_structured_to_enum_json(self):
         result = _clean_schema_for_genai(schema)
         assert result["enum"] == ['{"foo": 1}']
 
+    # --- Mapping tests: oneOf -> anyOf ---
+
+    def test_maps_oneof_to_anyof(self):
+        """oneOf is mapped to anyOf (genai accepts anyOf but not oneOf)."""
+        schema = {
+            "oneOf": [
+                {"type": "string"},
+                {"type": "number"},
+            ]
+        }
+        result = _clean_schema_for_genai(schema)
+        assert "oneOf" not in result
+        assert len(result["anyOf"]) == 2
+        assert result["anyOf"][0]["type"] == "string"
+        assert result["anyOf"][1]["type"] == "number"
+
+    def test_maps_oneof_recursively_and_cleans_branches(self):
+        """A nested oneOf (e.g. from a zod discriminatedUnion) is mapped to
+        anyOf and each branch is cleaned, so the union structure survives
+        instead of being dropped by the allowlist."""
+        schema = {
+            "type": "object",
+            "properties": {
+                "config": {
+                    "description": "discriminated union",
+                    "oneOf": [
+                        {
+                            "$comment": "branch A",
+                            "type": "object",
+                            "additionalProperties": False,
+                            "properties": {
+                                "kind": {"const": "a"},
+                            },
+                        },
+                        {
+                            "type": "object",
+                            "properties": {"kind": {"const": "b"}},
+                        },
+                    ],
+                }
+            },
+        }
+        result = _clean_schema_for_genai(schema)
+        config = result["properties"]["config"]
+        assert "oneOf" not in config
+        any_of = config["anyOf"]
+        assert len(any_of) == 2
+        # branches are recursively cleaned: $-keys / rejected keys stripped,
+        # const mapped to enum
+        assert "$comment" not in any_of[0]
+        assert "additionalProperties" not in any_of[0]
+        assert any_of[0]["properties"]["kind"]["enum"] == ["a"]
+        # the sibling description is preserved alongside the mapped anyOf
+        assert config["description"] == "discriminated union"
+
     # --- Edge cases ---
 
     def test_handles_non_dict_input(self):