feat(executor): add root union type normalization and extended tests

- Introduced `normalizeXAIObjectRootUnionBranchTypes` for ensuring object-only branches in root-level unions.
- Updated XAI executor logic to handle additional schema adjustments, ensuring compatibility with stricter validations.
- Extended test coverage to validate schema normalization, including both functional and custom tool scenarios.

Closes: #4397

internal/runtime/executor/xai_executor.go

@@ -1666,11 +1666,25 @@ func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bo
 	if toolType == xaiToolSearchType || toolType == xaiImageGenerationToolType {
 		return nil, true, true
 	}
+	if toolType == xaiCustomToolType && tool.Get("name").String() == "apply_patch" {
+		return nil, true, true
+	}
+
 	raw := []byte(tool.Raw)
-	if toolType == xaiCustomToolType {
-		if tool.Get("name").String() == "apply_patch" {
-			return nil, true, true
+	schemaTool := tool
+	if toolType == xaiFunctionToolType || toolType == xaiCustomToolType {
+		updatedTool, schemaChanged, ok := normalizeXAIObjectRootUnionBranchTypes(raw)
+		if !ok {
+			return nil, false, false
 		}
+		raw = updatedTool
+		if schemaChanged {
+			schemaTool = gjson.ParseBytes(raw)
+			changed = true
+			log.Debugf("xai: added object types to root union branches for tool %s.%s", namespaceName, tool.Get("name").String())
+		}
+	}
+	if toolType == xaiCustomToolType {
 		updatedTool, errSet := sjson.SetBytes(raw, "type", xaiFunctionToolType)
 		if errSet != nil {
 			return nil, false, false
@@ -1687,7 +1701,7 @@ func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bo
 		raw = updatedTool
 		changed = true
 	}
-	if toolType == xaiFunctionToolType && !tool.Get("parameters").Exists() {
+	if toolType == xaiFunctionToolType && !schemaTool.Get("parameters").Exists() {
 		updatedTool, errSet := sjson.SetRawBytes(raw, "parameters", []byte(`{"type":"object","properties":{}}`))
 		if errSet != nil {
 			return nil, false, false
@@ -1697,7 +1711,7 @@ func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bo
 	}
 	// Simplify the Codex Desktop automation schema and root unions that xAI
 	// rejects because function parameters must resolve exclusively to objects.
-	if toolType == xaiFunctionToolType && xaiFunctionParametersNeedSimplification(tool, namespaceName) {
+	if toolType == xaiFunctionToolType && xaiFunctionParametersNeedSimplification(schemaTool, namespaceName) {
 		updatedTool, errSet := sjson.SetRawBytes(raw, "parameters", []byte(xaiSafeFunctionParameters))
 		if errSet != nil {
 			return nil, false, false
@@ -2201,6 +2215,57 @@ func restoreXAINamespaceToolCallAtPath(data []byte, path string, refs map[string
 	return updated
 }
 
+// normalizeXAIObjectRootUnionBranchTypes makes untyped root union branches
+// explicitly object-only when the parameter root already permits only objects.
+// This preserves the original schema semantics while satisfying xAI validation.
+func normalizeXAIObjectRootUnionBranchTypes(tool []byte) ([]byte, bool, bool) {
+	parameters := gjson.GetBytes(tool, "parameters")
+	rootType := parameters.Get("type")
+	if rootType.Type != gjson.String || rootType.String() != "object" {
+		return tool, false, true
+	}
+
+	original := tool
+	changed := false
+	for _, unionName := range []string{"anyOf", "oneOf"} {
+		union := parameters.Get(unionName)
+		if !union.IsArray() {
+			continue
+		}
+		for index, branch := range union.Array() {
+			if !branch.IsObject() || branch.Get("type").Exists() {
+				continue
+			}
+			updated, errSet := sjson.SetBytes(tool, fmt.Sprintf("parameters.%s.%d.type", unionName, index), "object")
+			if errSet != nil {
+				return original, false, false
+			}
+			tool = updated
+			changed = true
+		}
+	}
+	return tool, changed, true
+}
+
+func xaiSchemaTypeIsObjectOnly(schemaType gjson.Result) bool {
+	if schemaType.Type == gjson.String {
+		return strings.EqualFold(strings.TrimSpace(schemaType.String()), "object")
+	}
+	if !schemaType.IsArray() {
+		return false
+	}
+	types := schemaType.Array()
+	if len(types) == 0 {
+		return false
+	}
+	for _, schemaTypeItem := range types {
+		if schemaTypeItem.Type != gjson.String || !strings.EqualFold(strings.TrimSpace(schemaTypeItem.String()), "object") {
+			return false
+		}
+	}
+	return true
+}
+
 // xaiFunctionParametersNeedSimplification reports whether a function tool, or
 // a custom tool normalized to a function, has a schema that xAI cannot accept.
 func xaiFunctionParametersNeedSimplification(tool gjson.Result, namespaceName string) bool {
@@ -2226,27 +2291,9 @@ func xaiFunctionParametersNeedSimplification(tool gjson.Result, namespaceName st
 			continue
 		}
 		for _, branch := range union.Array() {
-			branchType := branch.Get("type")
-			if branchType.Type == gjson.String {
-				if !strings.EqualFold(strings.TrimSpace(branchType.String()), "object") {
-					return true
-				}
-				continue
-			}
-			if !branchType.IsArray() {
-				// Without an explicit object type, the branch may accept non-object
-				// values even when it only declares object-specific keywords.
+			if !xaiSchemaTypeIsObjectOnly(branch.Get("type")) {
 				return true
 			}
-			allowedTypes := branchType.Array()
-			if len(allowedTypes) == 0 {
-				return true
-			}
-			for _, allowedType := range allowedTypes {
-				if allowedType.Type != gjson.String || !strings.EqualFold(strings.TrimSpace(allowedType.String()), "object") {
-					return true
-				}
-			}
 		}
 	}
 	return false

internal/runtime/executor/xai_executor_test.go

@@ -642,6 +642,111 @@ func TestXAIExecutorPrepareDropsOrphanedToolChoiceBeforeXSearchInject(t *testing
 	}
 }
 
+func TestXAIExecutorPrepareResponsesRequestAddsObjectTypeToRootUnionBranches(t *testing.T) {
+	t.Parallel()
+
+	cropParameters := `{
+		"type":"object",
+		"additionalProperties":false,
+		"required":["imagePath","point"],
+		"oneOf":[
+			{"required":["radius"],"not":{"required":["size"]}},
+			{"required":["size"],"not":{"required":["radius"]}}
+		],
+		"properties":{
+			"imagePath":{"type":"string"},
+			"point":{"type":"array"},
+			"radius":{"type":"number"},
+			"size":{"type":"object"}
+		}
+	}`
+	tests := []struct {
+		name         string
+		sourceFormat sdktranslator.Format
+		payload      []byte
+	}{
+		{
+			name:         "OpenAI Responses",
+			sourceFormat: sdktranslator.FormatOpenAIResponse,
+			payload: []byte(`{
+				"model":"grok-4.5",
+				"input":"crop a region",
+				"tools":[{
+					"type":"function",
+					"name":"crop_around_point",
+					"parameters":` + cropParameters + `
+				}]
+			}`),
+		},
+		{
+			name:         "OpenAI Chat Completions",
+			sourceFormat: sdktranslator.FormatOpenAI,
+			payload: []byte(`{
+				"model":"grok-4.5",
+				"messages":[{"role":"user","content":"crop a region"}],
+				"tools":[{
+					"type":"function",
+					"function":{
+						"name":"crop_around_point",
+						"parameters":` + cropParameters + `
+					}
+				}]
+			}`),
+		},
+	}
+
+	exec := NewXAIExecutor(&config.Config{})
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			prepared, err := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{
+				Model:   "grok-4.5",
+				Payload: tt.payload,
+			}, cliproxyexecutor.Options{
+				SourceFormat: tt.sourceFormat,
+				Stream:       true,
+			}, true)
+			if err != nil {
+				t.Fatalf("prepareResponsesRequest() error = %v", err)
+			}
+
+			var cropTool gjson.Result
+			for _, tool := range gjson.GetBytes(prepared.body, "tools").Array() {
+				if tool.Get("type").String() == xaiFunctionToolType && tool.Get("name").String() == "crop_around_point" {
+					cropTool = tool
+					break
+				}
+			}
+			if !cropTool.Exists() {
+				t.Fatalf("crop_around_point missing from upstream tools: %s", prepared.body)
+			}
+
+			parameters := cropTool.Get("parameters")
+			branches := parameters.Get("oneOf").Array()
+			if len(branches) != 2 {
+				t.Fatalf("oneOf branch count = %d, want 2; parameters=%s", len(branches), parameters.Raw)
+			}
+			for index, branch := range branches {
+				if got := branch.Get("type").String(); got != "object" {
+					t.Fatalf("oneOf.%d.type = %q, want object; parameters=%s", index, got, parameters.Raw)
+				}
+			}
+			for _, propertyName := range []string{"imagePath", "point", "radius", "size"} {
+				if !parameters.Get("properties." + propertyName).Exists() {
+					t.Fatalf("properties.%s missing: %s", propertyName, parameters.Raw)
+				}
+			}
+			if parameters.Get("additionalProperties").Type != gjson.False {
+				t.Fatalf("additionalProperties changed: %s", parameters.Raw)
+			}
+			if !branches[0].Get("not.required").Exists() || !branches[1].Get("not.required").Exists() {
+				t.Fatalf("oneOf constraints changed: %s", parameters.Raw)
+			}
+		})
+	}
+}
+
 func TestXAIExecutorPrepareAllowedToolsSyncsInjectedXSearch(t *testing.T) {
 	t.Parallel()
 
@@ -2608,7 +2713,7 @@ func TestXAIExecutorExecuteVideosUsesNativeEndpointFromRequestPath(t *testing.T)
 
 func TestNormalizeXAITools_SimplifiesCodexAppAutomationUpdateSchema(t *testing.T) {
 	// Large oneOf+$ref schema mimicking Codex Desktop codex_app.automation_update.
-	params := `{"oneOf":[{"type":"object","properties":{"mode":{"type":"string"}}}],"$defs":{"a":{"type":"string"}},"x":"` + strings.Repeat("y", 1600) + `"}`
+	params := `{"type":"object","oneOf":[{"properties":{"mode":{"type":"string"}}}],"$defs":{"a":{"type":"string"}},"x":"` + strings.Repeat("y", 1600) + `"}`
 	body := []byte(`{"model":"grok-4.5","tools":[{"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"automation_update","description":"sched","strict":true,"parameters":` + params + `}]},{"type":"function","name":"exec_command","parameters":{"type":"object","properties":{"cmd":{"type":"string"}}}}]}`)
 	out := normalizeXAITools(body)
 
@@ -2651,14 +2756,14 @@ func TestNormalizeXAITools_SimplifiesCodexAppAutomationUpdateSchema(t *testing.T
 }
 
 func TestNormalizeXAITools_SimplifiesFlattenedAndInvalidRootSchemas(t *testing.T) {
-	body := []byte(`{"tools":[{"type":"function","name":"codex_app__automation_update","strict":true,"parameters":{"oneOf":[{"type":"object","properties":{"action":{"type":"string"}},"required":["action"]},{"type":"null"}]}},{"type":"function","name":"nullable_lookup","strict":true,"parameters":{"anyOf":[{"type":"object","properties":{"query":{"type":"string"}}},{"type":["object","null"]}]}},{"type":"custom","name":"nullable_custom","strict":true,"parameters":{"oneOf":[{"type":"object"},{"type":"null"}]}},{"type":"function","name":"echo_tool","strict":true,"parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}]}`)
+	body := []byte(`{"tools":[{"type":"function","name":"codex_app__automation_update","strict":true,"parameters":{"oneOf":[{"type":"object","properties":{"action":{"type":"string"}},"required":["action"]},{"type":"null"}]}},{"type":"function","name":"nullable_lookup","strict":true,"parameters":{"anyOf":[{"type":"object","properties":{"query":{"type":"string"}}},{"type":["object","null"]}]}},{"type":"custom","name":"nullable_custom","strict":true,"parameters":{"oneOf":[{"type":"object"},{"type":"null"}]}},{"type":"function","name":"mixed_nullable","strict":true,"parameters":{"type":"object","oneOf":[{"required":["query"]},{"type":"null"}],"properties":{"query":{"type":"string"}}}},{"type":"function","name":"array_root_union","strict":true,"parameters":{"type":["object"],"anyOf":[{"required":["query"]},{"required":["id"]}],"properties":{"query":{"type":"string"},"id":{"type":"integer"}}}},{"type":"function","name":"echo_tool","strict":true,"parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}]}`)
 	out := normalizeXAITools(body)
 
 	tools := gjson.GetBytes(out, "tools").Array()
-	if len(tools) != 4 {
-		t.Fatalf("tools length = %d, want 4; body=%s", len(tools), string(out))
+	if len(tools) != 6 {
+		t.Fatalf("tools length = %d, want 6; body=%s", len(tools), string(out))
 	}
-	for index, wantName := range []string{"codex_app__automation_update", "nullable_lookup", "nullable_custom"} {
+	for index, wantName := range []string{"codex_app__automation_update", "nullable_lookup", "nullable_custom", "mixed_nullable", "array_root_union"} {
 		tool := tools[index]
 		if got := tool.Get("name").String(); got != wantName {
 			t.Fatalf("tools.%d.name = %q, want %q; body=%s", index, got, wantName, string(out))
@@ -2677,7 +2782,7 @@ func TestNormalizeXAITools_SimplifiesFlattenedAndInvalidRootSchemas(t *testing.T
 		}
 	}
 
-	echoTool := tools[3]
+	echoTool := tools[5]
 	if got := echoTool.Get("parameters.properties.message.type").String(); got != "string" {
 		t.Fatalf("echo_tool schema changed, message type = %q; body=%s", got, string(out))
 	}
@@ -2689,6 +2794,98 @@ func TestNormalizeXAITools_SimplifiesFlattenedAndInvalidRootSchemas(t *testing.T
 	}
 }
 
+func TestNormalizeXAITools_AddsObjectTypeToRootUnionBranches(t *testing.T) {
+	body := []byte(`{
+		"tools":[
+			{
+				"type":"function",
+				"name":"crop_around_point",
+				"strict":true,
+				"parameters":{
+					"type":"object",
+					"additionalProperties":false,
+					"required":["imagePath","point"],
+					"oneOf":[
+						{"required":["radius"],"not":{"required":["size"]}},
+						{"required":["size"],"not":{"required":["radius"]}}
+					],
+					"properties":{
+						"imagePath":{"type":"string"},
+						"point":{"type":"array"},
+						"radius":{"type":"number"},
+						"size":{"type":"object"},
+						"nested":{"oneOf":[{"required":["value"]},{}]}
+					}
+				}
+			},
+			{
+				"type":"function",
+				"name":"lookup",
+				"strict":true,
+				"parameters":{
+					"type":"object",
+					"anyOf":[{"required":["query"]},{"required":["id"]}],
+					"properties":{"query":{"type":"string"},"id":{"type":"integer"}}
+				}
+			},
+			{
+				"type":"custom",
+				"name":"custom_lookup",
+				"strict":true,
+				"parameters":{
+					"type":"object",
+					"oneOf":[{"required":["query"]},{"required":["id"]}],
+					"properties":{"query":{"type":"string"},"id":{"type":"integer"}}
+				}
+			}
+		]
+	}`)
+	out := normalizeXAITools(body)
+
+	for toolIndex, unionName := range []string{"oneOf", "anyOf"} {
+		tool := gjson.GetBytes(out, fmt.Sprintf("tools.%d", toolIndex))
+		branches := tool.Get("parameters." + unionName).Array()
+		if len(branches) != 2 {
+			t.Fatalf("tools.%d %s branch count = %d, want 2; body=%s", toolIndex, unionName, len(branches), string(out))
+		}
+		for branchIndex, branch := range branches {
+			if got := branch.Get("type").String(); got != "object" {
+				t.Fatalf("tools.%d parameters.%s.%d.type = %q, want object; body=%s", toolIndex, unionName, branchIndex, got, string(out))
+			}
+		}
+		if tool.Get("strict").Type != gjson.True {
+			t.Fatalf("tools.%d strict changed: %s", toolIndex, string(out))
+		}
+	}
+
+	cropParameters := gjson.GetBytes(out, "tools.0.parameters")
+	if cropParameters.Get("additionalProperties").Type != gjson.False {
+		t.Fatalf("crop additionalProperties changed: %s", cropParameters.Raw)
+	}
+	if got := cropParameters.Get("required.#").Int(); got != 2 {
+		t.Fatalf("crop required length = %d, want 2; parameters=%s", got, cropParameters.Raw)
+	}
+	if !cropParameters.Get("oneOf.0.not.required").Exists() || !cropParameters.Get("oneOf.1.not.required").Exists() {
+		t.Fatalf("crop oneOf constraints changed: %s", cropParameters.Raw)
+	}
+	if cropParameters.Get("properties.nested.oneOf.0.type").Exists() {
+		t.Fatalf("nested union branch must not be changed: %s", cropParameters.Raw)
+	}
+
+	customTool := gjson.GetBytes(out, "tools.2")
+	if got := customTool.Get("type").String(); got != xaiFunctionToolType {
+		t.Fatalf("custom tool type = %q, want function; body=%s", got, string(out))
+	}
+	for branchIndex, branch := range customTool.Get("parameters.oneOf").Array() {
+		if got := branch.Get("type").String(); got != "object" {
+			t.Fatalf("custom tool oneOf.%d.type = %q, want object; body=%s", branchIndex, got, string(out))
+		}
+	}
+	if customTool.Get("strict").Type != gjson.True {
+		t.Fatalf("custom tool strict changed: %s", string(out))
+	}
+}
+
 func TestNormalizeXAITools_QualifiesSameNamedNamespaceTools(t *testing.T) {
 	body := []byte(`{
 		"tools":[