From 64e71e16fe558481e28df8e7a68bd354aafddd94 Mon Sep 17 00:00:00 2001
From: roboomp <omp@can.ac>
Date: Sun, 7 Jun 2026 06:50:06 +0000
Subject: [PATCH] fix(ai): scoped union branch tagging to combinator path

Only mark Zod and JSON Schema validation issues as union-branch when their own path matches the combinator's path so nested array fields inside a tag-selected branch keep their singleton-wrap repair.

Fixes #2026
---
 .../src/utils/schema/json-schema-validator.ts | 10 +++-
 packages/ai/src/utils/validation.ts           |  6 +-
 .../ai/test/tool-argument-coercion.test.ts    | 55 +++++++++++++++++++
 3 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/packages/ai/src/utils/schema/json-schema-validator.ts b/packages/ai/src/utils/schema/json-schema-validator.ts
index fa8c82157..79a93052e 100644
--- a/packages/ai/src/utils/schema/json-schema-validator.ts
+++ b/packages/ai/src/utils/schema/json-schema-validator.ts
@@ -250,8 +250,16 @@ function validateSchemaNode(
 		const branchValid = keyword === "anyOf" ? matches > 0 : matches === 1;
 		if (!branchValid) {
 			if (matches === 0 && firstIssues && firstIssues.length > 0) {
+				// Only tag issues that sit at the combinator's own path as
+				// union-branch; deeper issues describe a specific field within
+				// the failed branch and should remain individually repairable.
+				const unionDepth = path.length;
 				for (const branchIssue of firstIssues) {
-					issues.push({ ...branchIssue, fromUnionBranch: true });
+					if (branchIssue.path.length === unionDepth) {
+						issues.push({ ...branchIssue, fromUnionBranch: true });
+					} else {
+						issues.push(branchIssue);
+					}
 				}
 			} else {
 				pushIssue(
diff --git a/packages/ai/src/utils/validation.ts b/packages/ai/src/utils/validation.ts
index e39e7d188..506c72978 100644
--- a/packages/ai/src/utils/validation.ts
+++ b/packages/ai/src/utils/validation.ts
@@ -784,9 +784,13 @@ function flattenIssues(issues: ReadonlyArray<ZodIssue>): FlatIssue[] {
 		if (issue.code === "invalid_union") {
 			const inner = (issue as unknown as { errors?: ReadonlyArray<ReadonlyArray<ZodIssue>> }).errors;
 			if (inner) {
+				// A union-branch issue only competes with a sibling branch when it
+				// sits at the union node's own path. Issues whose own path is
+				// non-empty live on a deeper field that an already-identified
+				// branch owns, so the singleton-array repair should still apply.
 				for (const branch of inner) {
 					for (const child of branch) {
-						walk(child, fullPath, true);
+						walk(child, fullPath, child.path.length === 0);
 					}
 				}
 			}
diff --git a/packages/ai/test/tool-argument-coercion.test.ts b/packages/ai/test/tool-argument-coercion.test.ts
index 586ae2942..e0ca64378 100644
--- a/packages/ai/test/tool-argument-coercion.test.ts
+++ b/packages/ai/test/tool-argument-coercion.test.ts
@@ -195,6 +195,61 @@ describe("Tool argument coercion", () => {
 		).toThrow("Validation failed");
 	});
 
+	it("still wraps nested array fields inside a tag-selected Zod union branch", () => {
+		const tool: Tool = {
+			name: "tagged_union",
+			description: "",
+			parameters: z.union([
+				z.object({ type: z.literal("indices"), indices: z.array(z.number()) }),
+				z.object({ type: z.literal("all") }),
+			]),
+		};
+
+		const result = validateToolArguments(tool, {
+			type: "toolCall",
+			id: "call-tagged-union",
+			name: "tagged_union",
+			arguments: { type: "indices", indices: 1 },
+		});
+
+		expect(result).toEqual({ type: "indices", indices: [1] });
+	});
+
+	it("still wraps nested array fields inside a tag-selected JSON Schema anyOf branch", () => {
+		const tool: Tool = {
+			name: "tagged_json_union",
+			description: "",
+			parameters: {
+				anyOf: [
+					{
+						type: "object",
+						properties: {
+							type: { const: "indices" },
+							indices: { type: "array", items: { type: "number" } },
+						},
+						required: ["type", "indices"],
+						additionalProperties: false,
+					},
+					{
+						type: "object",
+						properties: { type: { const: "all" } },
+						required: ["type"],
+						additionalProperties: false,
+					},
+				],
+			},
+		};
+
+		const result = validateToolArguments(tool, {
+			type: "toolCall",
+			id: "call-tagged-json-union",
+			name: "tagged_json_union",
+			arguments: { type: "indices", indices: 1 },
+		});
+
+		expect(result).toEqual({ type: "indices", indices: [1] });
+	});
+
 	it("parses JSON objects in string values when schema expects object", () => {
 		const tool: Tool = {
 			name: "t4",