Merge pull request #5174 from HuiCheng/fix/xai-keep-image-generation-grok-46

fix(xai): keep image_generation on grok-4.6+ conversation requests

internal/runtime/executor/xai_executor_execute.go

@@ -148,6 +148,10 @@ func (e *XAIExecutor) executeCompactRequest(ctx context.Context, auth *cliproxya
 	}
 	prepared.body, _ = sjson.DeleteBytes(prepared.body, "stream")
 	prepared.body, _ = sjson.DeleteBytes(prepared.body, "tools")
+	// Compact deletes tools after prepareResponsesRequestTo, which can now keep
+	// image_generation and rewrite its forced choice to allowed_tools on grok-4.6+.
+	// Drop the leftover selection so compact does not send tool_choice without tools.
+	prepared.body = normalizeXAIToolChoiceForTools(prepared.body)
 	for _, field := range []string{"max_output_tokens", "temperature", "top_p", "top_k", "stop"} {
 		prepared.body, _ = sjson.DeleteBytes(prepared.body, field)
 	}

internal/runtime/executor/xai_executor_request.go

@@ -98,6 +98,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox
 	// configured x_search injection, so no surviving choice references a deleted tool.
 	body = normalizeXAINamespaceToolChoice(body)
 	body = normalizeXAIForcedWebSearchToolChoice(body)
+	body = normalizeXAIForcedImageGenerationToolChoice(body)
 	body = pruneXAIOrphanedToolChoice(body)
 	body = normalizeXAIToolChoiceForTools(body)
 	if e.cfg != nil && e.cfg.XAI.InjectXSearch {
@@ -530,6 +531,91 @@ func preserveXAIResponsesOutputControls(body, source []byte, from sdktranslator.
 	return body
 }
 
+// xaiGrokImageGenerationMinVersion is the first Grok line that accepts xAI's
+// native Responses image_generation tool. Older conversation models still
+// reject that hosted type, so the executor keeps stripping it there.
+var xaiGrokImageGenerationMinVersion = xaiGrokVersion{major: 4, minor: 6}
+
+type xaiGrokVersion struct {
+	major int
+	minor int
+}
+
+// xaiSupportsNativeImageGeneration reports whether the Grok model accepts
+// xAI's native Responses image_generation tool. grok-4.20-* is an older
+// product line whose dotted minor is not comparable to grok-4.6.
+func xaiSupportsNativeImageGeneration(model string) bool {
+	name := strings.ToLower(strings.TrimSpace(thinking.ParseSuffix(model).ModelName))
+	if idx := strings.LastIndex(name, "/"); idx >= 0 {
+		name = name[idx+1:]
+	}
+	if name == "" || !strings.HasPrefix(name, "grok-") {
+		return false
+	}
+	rest := strings.TrimPrefix(name, "grok-")
+	if rest == "4.20" || strings.HasPrefix(rest, "4.20-") {
+		return false
+	}
+	ver, ok := xaiParseGrokVersionPrefix(rest)
+	if !ok {
+		return false
+	}
+	return xaiCompareGrokVersion(ver, xaiGrokImageGenerationMinVersion) >= 0
+}
+
+func xaiParseGrokVersionPrefix(rest string) (xaiGrokVersion, bool) {
+	i := 0
+	for i < len(rest) && rest[i] >= '0' && rest[i] <= '9' {
+		i++
+	}
+	if i == 0 {
+		return xaiGrokVersion{}, false
+	}
+	major, err := strconv.Atoi(rest[:i])
+	if err != nil {
+		return xaiGrokVersion{}, false
+	}
+	if i == len(rest) || rest[i] != '.' {
+		return xaiGrokVersion{major: major, minor: -1}, true
+	}
+	j := i + 1
+	for j < len(rest) && rest[j] >= '0' && rest[j] <= '9' {
+		j++
+	}
+	if j == i+1 {
+		return xaiGrokVersion{major: major, minor: -1}, true
+	}
+	minor, err := strconv.Atoi(rest[i+1 : j])
+	if err != nil {
+		return xaiGrokVersion{}, false
+	}
+	return xaiGrokVersion{major: major, minor: minor}, true
+}
+
+func xaiCompareGrokVersion(a, b xaiGrokVersion) int {
+	if a.major != b.major {
+		if a.major < b.major {
+			return -1
+		}
+		return 1
+	}
+	aMinor := a.minor
+	if aMinor < 0 {
+		aMinor = 0
+	}
+	bMinor := b.minor
+	if bMinor < 0 {
+		bMinor = 0
+	}
+	if aMinor < bMinor {
+		return -1
+	}
+	if aMinor > bMinor {
+		return 1
+	}
+	return 0
+}
+
 func sanitizeXAIResponsesBody(body []byte, model string) []byte {
 	// stop is supported by Chat Completions but not by xAI's Responses API.
 	body, _ = sjson.DeleteBytes(body, "stop")
@@ -590,8 +676,18 @@ func ensureXAINativeXSearchAllowedTools(body []byte) []byte {
 // normalizeXAIForcedWebSearchToolChoice rewrites Codex's hosted-tool choice
 // into the allowed_tools form accepted by xAI's ModelToolChoice schema.
 func normalizeXAIForcedWebSearchToolChoice(body []byte) []byte {
+	return normalizeXAIForcedHostedToolChoice(body, xaiWebSearchToolType)
+}
+
+// normalizeXAIForcedImageGenerationToolChoice rewrites a forced image_generation
+// choice into the same allowed_tools form used for web_search.
+func normalizeXAIForcedImageGenerationToolChoice(body []byte) []byte {
+	return normalizeXAIForcedHostedToolChoice(body, xaiImageGenerationToolType)
+}
+
+func normalizeXAIForcedHostedToolChoice(body []byte, toolType string) []byte {
 	choice := gjson.GetBytes(body, "tool_choice")
-	if !choice.IsObject() || strings.TrimSpace(choice.Get("type").String()) != xaiWebSearchToolType {
+	if !choice.IsObject() || strings.TrimSpace(choice.Get("type").String()) != toolType {
 		return body
 	}
 
@@ -731,13 +827,14 @@ func normalizeXAITools(body []byte) []byte {
 	if !gjson.ValidBytes(body) {
 		return body
 	}
+	keepImageGeneration := xaiSupportsNativeImageGeneration(gjson.GetBytes(body, "model").String())
 	original := body
 	normalizeAtPath := func(path string) bool {
 		tools := gjson.GetBytes(body, path)
 		if !tools.Exists() || !tools.IsArray() {
 			return true
 		}
-		filtered, changed, ok := normalizeXAIToolArray(tools)
+		filtered, changed, ok := normalizeXAIToolArray(tools, keepImageGeneration)
 		if !ok {
 			return false
 		}
@@ -827,7 +924,7 @@ func promoteXAIAdditionalTools(body []byte) []byte {
 	return updated
 }
 
-func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) {
+func normalizeXAIToolArray(tools gjson.Result, keepImageGeneration bool) ([]byte, bool, bool) {
 	toolItems := tools.Array()
 	filtered := make([][]byte, 0, len(toolItems))
 	changed := false
@@ -838,7 +935,7 @@ func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) {
 			namespaceName := tool.Get("name").String()
 			if namespaceTools := tool.Get("tools"); namespaceTools.IsArray() {
 				for _, nestedTool := range namespaceTools.Array() {
-					nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName)
+					nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName, keepImageGeneration)
 					if !ok {
 						return nil, false, false
 					}
@@ -850,7 +947,7 @@ func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) {
 			}
 			continue
 		}
-		raw, toolChanged, ok := normalizeXAITool(tool, "")
+		raw, toolChanged, ok := normalizeXAITool(tool, "", keepImageGeneration)
 		if !ok {
 			return nil, false, false
 		}
@@ -944,10 +1041,13 @@ func normalizeXAINamespaceToolChoice(body []byte) []byte {
 	return body
 }
 
-func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bool) {
+func normalizeXAITool(tool gjson.Result, namespaceName string, keepImageGeneration bool) ([]byte, bool, bool) {
 	toolType := tool.Get("type").String()
 	changed := false
-	if toolType == xaiToolSearchType || toolType == xaiImageGenerationToolType {
+	if toolType == xaiToolSearchType {
+		return nil, true, true
+	}
+	if toolType == xaiImageGenerationToolType && !keepImageGeneration {
 		return nil, true, true
 	}
 	if toolType == xaiCustomToolType && tool.Get("name").String() == "apply_patch" {

internal/runtime/executor/xai_executor_test.go

@@ -980,6 +980,149 @@ func TestPruneXAIOrphanedToolChoice(t *testing.T) {
 	}
 }
 
+func TestXAISupportsNativeImageGeneration(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		model string
+		want  bool
+	}{
+		{model: "", want: false},
+		{model: "grok-4.5", want: false},
+		{model: "grok-4.3", want: false},
+		{model: "grok-4", want: false},
+		{model: "grok-4.20-0309-reasoning", want: false},
+		{model: "grok-4.20-multi-agent-0309", want: false},
+		{model: "grok-build-0.1", want: false},
+		{model: "grok-composer-2.5-fast", want: false},
+		{model: "grok-3-mini", want: false},
+		{model: "gpt-5.6", want: false},
+		{model: "grok-4.6", want: true},
+		{model: "grok-4.6(high)", want: true},
+		{model: "xai/grok-4.6", want: true},
+		{model: "grok-4.7", want: true},
+		{model: "grok-5", want: true},
+		{model: "grok-5.0", want: true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.model, func(t *testing.T) {
+			t.Parallel()
+			if got := xaiSupportsNativeImageGeneration(tt.model); got != tt.want {
+				t.Fatalf("xaiSupportsNativeImageGeneration(%q) = %t, want %t", tt.model, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestNormalizeXAITools_ImageGenerationByModel(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		body       []byte
+		wantKeep   bool
+		wantAction string
+	}{
+		{
+			name:     "missing model still strips",
+			body:     []byte(`{"tools":[{"type":"image_generation"},{"type":"web_search"}]}`),
+			wantKeep: false,
+		},
+		{
+			name:     "grok-4.5 strips",
+			body:     []byte(`{"model":"grok-4.5","tools":[{"type":"image_generation"},{"type":"web_search"}]}`),
+			wantKeep: false,
+		},
+		{
+			name:     "grok-4.20 strips despite larger minor",
+			body:     []byte(`{"model":"grok-4.20-0309-reasoning","tools":[{"type":"image_generation"},{"type":"web_search"}]}`),
+			wantKeep: false,
+		},
+		{
+			name:       "grok-4.6 keeps action",
+			body:       []byte(`{"model":"grok-4.6","tools":[{"type":"image_generation","action":"generate"},{"type":"web_search"}]}`),
+			wantKeep:   true,
+			wantAction: "generate",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			out := normalizeXAITools(tt.body)
+			tools := gjson.GetBytes(out, "tools").Array()
+			foundImage := false
+			foundWebSearch := false
+			var imageTool gjson.Result
+			for _, tool := range tools {
+				switch tool.Get("type").String() {
+				case "image_generation":
+					foundImage = true
+					imageTool = tool
+				case "web_search":
+					foundWebSearch = true
+				}
+			}
+			if !foundWebSearch {
+				t.Fatalf("web_search missing; body=%s", out)
+			}
+			if foundImage != tt.wantKeep {
+				t.Fatalf("image_generation kept=%t, want %t; body=%s", foundImage, tt.wantKeep, out)
+			}
+			if tt.wantKeep && tt.wantAction != "" {
+				if got := imageTool.Get("action").String(); got != tt.wantAction {
+					t.Fatalf("image_generation.action = %q, want %q; body=%s", got, tt.wantAction, out)
+				}
+			}
+		})
+	}
+}
+
+func TestXAIExecutorPrepareKeepsNativeImageGenerationForGrok46(t *testing.T) {
+	t.Parallel()
+
+	exec := NewXAIExecutor(&config.Config{})
+	prepared, err := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{
+		Model: "grok-4.6",
+		Payload: []byte(`{
+			"model":"grok-4.6",
+			"input":"draw a red circle",
+			"tools":[{"type":"image_generation","action":"generate"}],
+			"tool_choice":{"type":"image_generation"}
+		}`),
+	}, cliproxyexecutor.Options{
+		SourceFormat: sdktranslator.FormatOpenAIResponse,
+		Stream:       false,
+	}, false)
+	if err != nil {
+		t.Fatalf("prepareResponsesRequest() error = %v", err)
+	}
+
+	tools := gjson.GetBytes(prepared.body, "tools").Array()
+	if len(tools) != 1 {
+		t.Fatalf("tools length = %d, want 1; body=%s", len(tools), prepared.body)
+	}
+	if got := tools[0].Get("type").String(); got != "image_generation" {
+		t.Fatalf("tools.0.type = %q, want image_generation; body=%s", got, prepared.body)
+	}
+	if got := tools[0].Get("action").String(); got != "generate" {
+		t.Fatalf("tools.0.action = %q, want generate; body=%s", got, prepared.body)
+	}
+	choice := gjson.GetBytes(prepared.body, "tool_choice")
+	if got := choice.Get("type").String(); got != "allowed_tools" {
+		t.Fatalf("tool_choice.type = %q, want allowed_tools; body=%s", got, prepared.body)
+	}
+	if got := choice.Get("mode").String(); got != "required" {
+		t.Fatalf("tool_choice.mode = %q, want required; body=%s", got, prepared.body)
+	}
+	allowed := choice.Get("tools").Array()
+	if len(allowed) != 1 {
+		t.Fatalf("tool_choice.tools length = %d, want 1; body=%s", len(allowed), prepared.body)
+	}
+	if got := allowed[0].Get("type").String(); got != "image_generation" {
+		t.Fatalf("tool_choice.tools.0.type = %q, want image_generation; body=%s", got, prepared.body)
+	}
+}
+
 func TestXAIExecutorPrepareDropsOrphanedToolChoiceBeforeXSearchInject(t *testing.T) {
 	t.Parallel()
 
@@ -2057,6 +2200,55 @@ func TestXAIExecutorCompactUsesCompactEndpoint(t *testing.T) {
 	}
 }
 
+func TestXAIExecutorCompactDropsOrphanedImageGenerationToolChoice(t *testing.T) {
+	var gotBody []byte
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var errRead error
+		gotBody, errRead = io.ReadAll(r.Body)
+		if errRead != nil {
+			t.Fatalf("read body: %v", errRead)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"id":"resp_1","object":"response.compaction","output":[{"type":"compaction","encrypted_content":"opaque-out"}]}`))
+	}))
+	defer server.Close()
+
+	exec := NewXAIExecutor(&config.Config{})
+	auth := &cliproxyauth.Auth{
+		Provider: "xai",
+		Attributes: map[string]string{
+			"base_url": server.URL,
+			"api_key":  "xai-token",
+		},
+	}
+
+	_, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
+		Model: "grok-4.6",
+		Payload: []byte(`{
+			"model":"grok-4.6",
+			"input":"compact this",
+			"tools":[{"type":"image_generation","action":"generate"}],
+			"tool_choice":{"type":"image_generation"}
+		}`),
+	}, cliproxyexecutor.Options{
+		SourceFormat: sdktranslator.FormatOpenAIResponse,
+		Alt:          "responses/compact",
+	})
+	if err != nil {
+		t.Fatalf("Execute compact error: %v", err)
+	}
+	if gjson.GetBytes(gotBody, "tools").Exists() {
+		t.Fatalf("tools exists in compact body: %s", gotBody)
+	}
+	if gjson.GetBytes(gotBody, "tool_choice").Exists() {
+		t.Fatalf("orphaned tool_choice leaked into compact body: %s", gotBody)
+	}
+	if gjson.GetBytes(gotBody, "parallel_tool_calls").Exists() {
+		t.Fatalf("parallel_tool_calls exists in compact body: %s", gotBody)
+	}
+}
+
 func TestXAIExecutorCompactOAuthUsesOfficialAPIHeadersNotCLIProxy(t *testing.T) {
 	var gotPath string
 	var gotHost string