From b4b53716409b227fc1b5c62d28ce434ec0e5a484 Mon Sep 17 00:00:00 2001 From: 251392463 <19240165+251392463@users.noreply.github.com> Date: Fri, 28 Aug 2026 15:35:20 +0800 Subject: [PATCH] fix(grok): preserve Responses vision tool output images --- .../openai_gateway_grok_model_input.go | 188 +++++++++++++++++- .../openai_gateway_grok_model_input_test.go | 55 +++++ .../service/openai_gateway_grok_test.go | 41 ++++ 3 files changed, 281 insertions(+), 3 deletions(-) diff --git a/backend/internal/service/openai_gateway_grok_model_input.go b/backend/internal/service/openai_gateway_grok_model_input.go index 6ad74dda42..dc1b522a9c 100644 --- a/backend/internal/service/openai_gateway_grok_model_input.go +++ b/backend/internal/service/openai_gateway_grok_model_input.go @@ -38,9 +38,12 @@ func sanitizeGrokResponsesModelInput(body []byte) ([]byte, error) { } callIDs, outputIDs := pairGrokReplayCallIDs(items) filtered := make([]any, 0, len(items)) + pendingOutputImages := make([]grokToolOutputImage, 0) for index, rawItem := range items { item, ok := rawItem.(map[string]any) if !ok { + filtered = appendGrokToolOutputImageMessage(filtered, pendingOutputImages) + pendingOutputImages = pendingOutputImages[:0] if text, ok := rawItem.(string); ok && strings.TrimSpace(text) != "" { filtered = append(filtered, map[string]any{"type": "message", "role": "user", "content": text}) } @@ -50,14 +53,20 @@ func sanitizeGrokResponsesModelInput(body []byte) ([]byte, error) { itemType := strings.ToLower(strings.TrimSpace(grokStringValue(item["type"]))) role := strings.ToLower(strings.TrimSpace(grokStringValue(item["role"]))) if role == "tool" || role == "function" || isGrokReplayOutputType(itemType) { - output := firstNonNilGrokJSONValue(item["output"], item["content"], item["results"]) + callID := outputIDs[index] + rawOutput := firstNonNilGrokJSONValue(item["output"], item["content"], item["results"]) + output, nestedImages := normalizeGrokToolOutput(rawOutput, callID) filtered = append(filtered, map[string]any{ "type": "function_call_output", - "call_id": outputIDs[index], - "output": grokModelInputString(output, "(empty)"), + "call_id": callID, + "output": output, }) + pendingOutputImages = append(pendingOutputImages, extractGrokToolOutputImages(item, callID)...) + pendingOutputImages = append(pendingOutputImages, nestedImages...) continue } + filtered = appendGrokToolOutputImageMessage(filtered, pendingOutputImages) + pendingOutputImages = pendingOutputImages[:0] switch itemType { case "text", "input_text", "output_text": @@ -139,6 +148,7 @@ func sanitizeGrokResponsesModelInput(body []byte) ([]byte, error) { } filtered = append(filtered, item) } + filtered = appendGrokToolOutputImageMessage(filtered, pendingOutputImages) encoded, err := json.Marshal(filtered) if err != nil { @@ -151,6 +161,178 @@ func sanitizeGrokResponsesModelInput(body []byte) ([]byte, error) { return updated, nil } +type grokToolOutputImage struct { + callID string + url string +} + +// Grok 1.0.5 returns read_file media as structured output parts rather than +// top-level images. xAI requires function_call_output.output to be a string, +// so remove those media parts before stringifying and lift them separately. +func normalizeGrokToolOutput(value any, callID string) (string, []grokToolOutputImage) { + stripped, images, keep := stripGrokToolOutputImages(value, callID) + if len(images) == 0 { + return grokModelInputString(value, "(empty)"), nil + } + if !keep { + return "(empty)", images + } + return grokStructuredToolOutputString(stripped, "(empty)"), images +} + +func stripGrokToolOutputImages(value any, callID string) (any, []grokToolOutputImage, bool) { + switch typed := value.(type) { + case []any: + filtered := make([]any, 0, len(typed)) + images := make([]grokToolOutputImage, 0) + for _, item := range typed { + stripped, nested, keep := stripGrokToolOutputImages(item, callID) + images = append(images, nested...) + if keep { + filtered = append(filtered, stripped) + } + } + return filtered, images, len(filtered) > 0 + case map[string]any: + partType := strings.ToLower(strings.TrimSpace(grokStringValue(typed["type"]))) + if partType == "image" || partType == "image_url" || partType == "input_image" { + url := grokToolOutputImageURL(typed) + if url == "" || isEmptyBase64DataURI(url) { + return nil, nil, false + } + return nil, []grokToolOutputImage{{callID: callID, url: url}}, false + } + + filtered := make(map[string]any, len(typed)) + for key, item := range typed { + filtered[key] = item + } + images := make([]grokToolOutputImage, 0) + if rawImages, exists := filtered["images"]; exists { + if values, ok := rawImages.([]any); ok { + for _, rawImage := range values { + url := grokToolOutputImageURL(rawImage) + if url != "" && !isEmptyBase64DataURI(url) { + images = append(images, grokToolOutputImage{callID: callID, url: url}) + } + } + delete(filtered, "images") + } + } + for _, field := range []string{"content", "output", "results"} { + item, exists := filtered[field] + if !exists { + continue + } + stripped, nested, keep := stripGrokToolOutputImages(item, callID) + images = append(images, nested...) + if keep { + filtered[field] = stripped + } else { + delete(filtered, field) + } + } + return filtered, images, len(filtered) > 0 + default: + return value, nil, value != nil + } +} + +func grokStructuredToolOutputString(value any, fallback string) string { + parts, ok := value.([]any) + if !ok || len(parts) == 0 { + return grokModelInputString(value, fallback) + } + + texts := make([]string, 0, len(parts)) + for _, rawPart := range parts { + part, ok := rawPart.(map[string]any) + if !ok { + return grokModelInputString(value, fallback) + } + partType := strings.ToLower(strings.TrimSpace(grokStringValue(part["type"]))) + if partType != "text" && partType != "input_text" && partType != "output_text" { + return grokModelInputString(value, fallback) + } + if text := strings.TrimSpace(grokStringValue(part["text"])); text != "" { + texts = append(texts, text) + } + } + if len(texts) == 0 { + return fallback + } + return strings.Join(texts, "\n") +} + +// Grok Shell attaches local image results to function_call_output.images. +// xAI's ModelInput accepts only a string output, so keep the tool replies +// consecutive and lift their media into one following user input_image turn. +func extractGrokToolOutputImages(item map[string]any, callID string) []grokToolOutputImage { + rawImages, ok := item["images"].([]any) + if !ok || len(rawImages) == 0 { + return nil + } + + images := make([]grokToolOutputImage, 0, len(rawImages)) + for _, rawImage := range rawImages { + url := grokToolOutputImageURL(rawImage) + if url == "" || isEmptyBase64DataURI(url) { + continue + } + images = append(images, grokToolOutputImage{callID: callID, url: url}) + } + return images +} + +func grokToolOutputImageURL(value any) string { + switch image := value.(type) { + case string: + return strings.TrimSpace(image) + case map[string]any: + for _, field := range []string{"url", "image_url", "file_url"} { + raw := image[field] + switch typed := raw.(type) { + case string: + if url := strings.TrimSpace(typed); url != "" { + return url + } + case map[string]any: + if url := strings.TrimSpace(grokStringValue(typed["url"])); url != "" { + return url + } + } + } + } + return "" +} + +func appendGrokToolOutputImageMessage(items []any, images []grokToolOutputImage) []any { + if len(images) == 0 { + return items + } + + content := make([]any, 0, len(images)*2) + lastCallID := "" + for _, image := range images { + if image.callID != lastCallID { + content = append(content, map[string]any{ + "type": "input_text", + "text": fmt.Sprintf("[Tool output media for call %s]", image.callID), + }) + lastCallID = image.callID + } + content = append(content, map[string]any{ + "type": "input_image", + "image_url": image.url, + }) + } + return append(items, map[string]any{ + "type": "message", + "role": "user", + "content": content, + }) +} + func pairGrokReplayCallIDs(items []any) (map[int]string, map[int]string) { callIDs := make(map[int]string) outputIDs := make(map[int]string) diff --git a/backend/internal/service/openai_gateway_grok_model_input_test.go b/backend/internal/service/openai_gateway_grok_model_input_test.go index 4696c6feed..4533bc646c 100644 --- a/backend/internal/service/openai_gateway_grok_model_input_test.go +++ b/backend/internal/service/openai_gateway_grok_model_input_test.go @@ -112,6 +112,61 @@ func TestSanitizeGrokResponsesModelInputMapsAllItemAliasesAndRejectsConflicts(t require.NotEqual(t, thirdCallID, ambiguousOutputID) } +func TestSanitizeGrokResponsesModelInputPreservesGrokShellOutputImages(t *testing.T) { + body := []byte(`{"input":[ + {"type":"function_call","call_id":"call_a","name":"read_file","arguments":"{}"}, + {"type":"function_call","call_id":"call_b","name":"read_file","arguments":"{}"}, + {"type":"function_call_output","call_id":"call_a","content":"Read image file: /tmp/a.png","images":[{"type":"image","url":"data:image/png;base64,QUE="}]}, + {"type":"function_call_output","call_id":"call_b","content":"Read image file: /tmp/b.png","images":[{"type":"image","url":"data:image/jpeg;base64,QkI="}]} + ]}`) + + patched, err := sanitizeGrokResponsesModelInput(body) + require.NoError(t, err) + require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.2.type").String()) + require.Equal(t, "Read image file: /tmp/a.png", gjson.GetBytes(patched, "input.2.output").String()) + require.False(t, gjson.GetBytes(patched, "input.2.images").Exists()) + require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.3.type").String()) + require.Equal(t, "message", gjson.GetBytes(patched, "input.4.type").String()) + require.Equal(t, "user", gjson.GetBytes(patched, "input.4.role").String()) + require.Equal(t, "[Tool output media for call call_a]", gjson.GetBytes(patched, "input.4.content.0.text").String()) + require.Equal(t, "input_image", gjson.GetBytes(patched, "input.4.content.1.type").String()) + require.Equal(t, "data:image/png;base64,QUE=", gjson.GetBytes(patched, "input.4.content.1.image_url").String()) + require.Equal(t, "[Tool output media for call call_b]", gjson.GetBytes(patched, "input.4.content.2.text").String()) + require.Equal(t, "data:image/jpeg;base64,QkI=", gjson.GetBytes(patched, "input.4.content.3.image_url").String()) +} + +func TestSanitizeGrokResponsesModelInputPreservesGrok105StructuredOutputImages(t *testing.T) { + body := []byte(`{"input":[ + {"type":"function_call","call_id":"call_read","name":"read_file","arguments":"{}"}, + {"type":"function_call_output","call_id":"call_read","output":[ + {"type":"input_text","text":"Read image file: /tmp/example.png"}, + {"type":"input_image","detail":"auto","image_url":"data:image/png;base64,QUE="} + ]} + ]}`) + + patched, err := sanitizeGrokResponsesModelInput(body) + require.NoError(t, err) + require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.1.type").String()) + require.Equal(t, "Read image file: /tmp/example.png", gjson.GetBytes(patched, "input.1.output").String()) + require.NotContains(t, gjson.GetBytes(patched, "input.1.output").String(), "base64") + require.Equal(t, "message", gjson.GetBytes(patched, "input.2.type").String()) + require.Equal(t, "input_image", gjson.GetBytes(patched, "input.2.content.1.type").String()) + require.Equal(t, "data:image/png;base64,QUE=", gjson.GetBytes(patched, "input.2.content.1.image_url").String()) +} + +func TestSanitizeGrokResponsesModelInputSkipsInvalidOutputImages(t *testing.T) { + body := []byte(`{"input":[ + {"type":"function_call_output","call_id":"call_empty","output":"done","images":[{"type":"image","url":"data:image/png;base64,"},{}]}, + {"role":"user","content":"continue"} + ]}`) + + patched, err := sanitizeGrokResponsesModelInput(body) + require.NoError(t, err) + require.Len(t, gjson.GetBytes(patched, "input").Array(), 2) + require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.0.type").String()) + require.Equal(t, "message", gjson.GetBytes(patched, "input.1.type").String()) +} + func TestPatchGrokResponsesBodyCombinesAdditionalToolsAndDropsOrphanControls(t *testing.T) { body := []byte(`{ "input":[{"type":"additional_tools","tools":[{"type":"function","name":"lookup"}]},{"role":"user","content":"hi"}], diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index d1c966ac07..e04e56de99 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -144,6 +144,47 @@ func TestPatchGrokResponsesBodyDropsViewImageOnlyToolMetadata(t *testing.T) { require.False(t, gjson.GetBytes(patched, "parallel_tool_calls").Exists()) } +func TestPatchGrokResponsesBodyPreservesGrokShellFunctionOutputImages(t *testing.T) { + t.Parallel() + + body := []byte(`{ + "model":"grok-4.6", + "input":[ + {"type":"function_call","call_id":"call_read","name":"read_file","arguments":"{}"}, + {"type":"function_call_output","call_id":"call_read","content":"Read image file: /tmp/example.png","images":[{"type":"image","url":"data:image/png;base64,QUE="}]} + ] + }`) + + patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.6") + require.NoError(t, err) + require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.1.type").String()) + require.Equal(t, "Read image file: /tmp/example.png", gjson.GetBytes(patched, "input.1.output").String()) + require.Equal(t, "message", gjson.GetBytes(patched, "input.2.type").String()) + require.Equal(t, "input_image", gjson.GetBytes(patched, "input.2.content.1.type").String()) + require.Equal(t, "data:image/png;base64,QUE=", gjson.GetBytes(patched, "input.2.content.1.image_url").String()) +} + +func TestPatchGrokResponsesBodyPreservesGrok105StructuredFunctionOutputImages(t *testing.T) { + t.Parallel() + + body := []byte(`{ + "model":"grok-4.6", + "input":[ + {"type":"function_call","call_id":"call_read","name":"read_file","arguments":"{}"}, + {"type":"function_call_output","call_id":"call_read","output":[ + {"type":"input_text","text":"Read image file: /tmp/example.png"}, + {"type":"input_image","detail":"auto","image_url":"data:image/png;base64,QUE="} + ]} + ] + }`) + + patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.6") + require.NoError(t, err) + require.Equal(t, "Read image file: /tmp/example.png", gjson.GetBytes(patched, "input.1.output").String()) + require.Equal(t, "input_image", gjson.GetBytes(patched, "input.2.content.1.type").String()) + require.Equal(t, "data:image/png;base64,QUE=", gjson.GetBytes(patched, "input.2.content.1.image_url").String()) +} + func TestPatchGrokResponsesBodySanitizesComposerReasoningParameters(t *testing.T) { t.Parallel()