From b7e1eecabd585a2ca17e5cdc901421362ee9fc50 Mon Sep 17 00:00:00 2001 From: jp Date: Mon, 15 Jun 2026 15:24:36 -0400 Subject: [PATCH] fix(adk-middleware): map oneOf -> anyOf in _clean_schema_for_genai MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit google.genai's Schema accepts `anyOf` but not `oneOf`, so the schema cleaner's allowlist silently dropped `oneOf` — erasing the structure of discriminated unions. CopilotKit / AG-UI frontend tools serialize a zod `discriminatedUnion` to `oneOf` via zod-to-json-schema, so any frontend tool with such a parameter reached Gemini as a bare `{description}` object with no fields, and the model would invent a shape. Map `oneOf` -> `anyOf` (recursively, cleaning each branch) alongside the existing `examples`/`const` mappings. For tool-argument schemas the distinction is immaterial — the model emits exactly one branch either way. Mirrors the prior Gemini-compat fixes (#1919 additionalProperties, #1664 nested required). Co-Authored-By: Claude Opus 4.8 --- .../python/src/ag_ui_adk/client_proxy_tool.py | 17 +++++- .../python/tests/test_client_proxy_tool.py | 55 +++++++++++++++++++ 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/integrations/adk-middleware/python/src/ag_ui_adk/client_proxy_tool.py b/integrations/adk-middleware/python/src/ag_ui_adk/client_proxy_tool.py index ec23a760d9..73867d3135 100644 --- a/integrations/adk-middleware/python/src/ag_ui_adk/client_proxy_tool.py +++ b/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 diff --git a/integrations/adk-middleware/python/tests/test_client_proxy_tool.py b/integrations/adk-middleware/python/tests/test_client_proxy_tool.py index 227027b419..139c92042f 100644 --- a/integrations/adk-middleware/python/tests/test_client_proxy_tool.py +++ b/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):