fix(mcp): forward image tool results on the existing image channel

internal/agent/loop.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/Gitlawb/zero/internal/execution"
 	"github.com/Gitlawb/zero/internal/hooks"
+	"github.com/Gitlawb/zero/internal/modelregistry"
 	"github.com/Gitlawb/zero/internal/redaction"
 	"github.com/Gitlawb/zero/internal/sandbox"
 	"github.com/Gitlawb/zero/internal/streamjson"
@@ -646,9 +647,25 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
 		// recorded. They travel as user messages, and a user message between two
 		// tool_results breaks strict provider replay — Anthropic coalesces them
 		// into one user block list and requires the tool_result blocks first, so
-		// interleaving yields [tool_result, text, image, tool_result] and a 400.
-		// Same reason the self-correction feedback below is deferred.
-		var toolImageMessages []zeroruntime.Message
+		// Images ride a following USER message rather than the tool result
+		// above. Every provider drops images on a tool-role message —
+		// Anthropic's tool_result content is a string, Gemini's is a
+		// functionResponse, and OpenAI guards its image parts to the user role
+		// — so attaching them there would silently deliver nothing. A separate
+		// message also keeps the one-tool-result-per-tool-call pairing intact,
+		// which the providers validate. Vision-gate evaluation is deferred until
+		// after any model switch this turn resolves, so images can reach an
+		// escalated vision-capable model.
+		var toolResultsWithImages []ToolResult
+		buildToolImageMessages := func() []zeroruntime.Message {
+			var out []zeroruntime.Message
+			for _, tr := range toolResultsWithImages {
+				if imageMessage, ok := toolResultImageMessage(tr, options); ok {
+					out = append(out, imageMessage)
+				}
+			}
+			return out
+		}
 		// Parallel read-ahead state: results for calls[precomputedStart:precomputedEnd]
 		// executed concurrently, consumed strictly in order below.
 		var precomputed []precomputedToolResult
@@ -705,15 +722,8 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
 				IsError:            toolResult.Status == tools.StatusError,
 				ChangedFiles:       append([]string(nil), toolResult.ChangedFiles...),
 			})
-			// Images ride a following USER message rather than the tool result
-			// above. Every provider drops images on a tool-role message —
-			// Anthropic's tool_result content is a string, Gemini's is a
-			// functionResponse, and OpenAI guards its image parts to the user role
-			// — so attaching them there would silently deliver nothing. A separate
-			// message also keeps the one-tool-result-per-tool-call pairing intact,
-			// which the providers validate.
-			if imageMessage, ok := toolResultImageMessage(toolResult); ok {
-				toolImageMessages = append(toolImageMessages, imageMessage)
+			if len(toolResult.Images) > 0 {
+				toolResultsWithImages = append(toolResultsWithImages, toolResult)
 			}
 
 			// A tool may demand the run ABORT — a canceled/timed-out ask_user prompt
@@ -725,13 +735,13 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
 			}
 			if abortErr != nil {
 				messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
-				messages = append(messages, toolImageMessages...)
+				messages = append(messages, buildToolImageMessages()...)
 				result.Messages = copyMessages(messages)
 				return result, abortErr
 			}
 			if stopReason := stopReasonFromToolResult(toolResult); stopReason != "" {
 				messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
-				messages = append(messages, toolImageMessages...)
+				messages = append(messages, buildToolImageMessages()...)
 				result.FinalAnswer = toolResult.ModelOutput()
 				result.StopReason = stopReason
 				result.Messages = copyMessages(messages)
@@ -755,7 +765,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
 				// messages stay valid for a strict provider replay (Anthropic
 				// rejects a tool_use with no answering tool_result).
 				messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
-				messages = append(messages, toolImageMessages...)
+				messages = append(messages, buildToolImageMessages()...)
 				result.FinalAnswer = toolFailureStopAnswer(call.Name, outcome.Count)
 				result.Messages = copyMessages(messages)
 				return result, nil
@@ -774,10 +784,6 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
 				postEditDiagnostics.enqueue(ctx, toolResult.ChangedFiles)
 			}
 		}
-		// Every tool_result for this turn is now recorded, including aborted
-		// placeholders, so the images can follow without splitting them.
-		messages = append(messages, toolImageMessages...)
-		toolImageMessages = nil
 
 		// Run post-edit self-correction once over the union of files this turn
 		// changed, then append any feedback after every tool_result is recorded so
@@ -865,6 +871,12 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
 			}
 		}
 
+		// Every tool_result, self-correct notice, and model-switch update for this
+		// turn is now in place; evaluate and append image messages against the
+		// effective (potentially escalated) model.
+		messages = append(messages, buildToolImageMessages()...)
+		toolResultsWithImages = nil
+
 		// A turn can mix valid tool calls with a dropped (nameless) one. The valid
 		// calls executed above; surface the dropped call too so it is never
 		// silently ignored just because the turn also did real work. This is
@@ -3457,8 +3469,10 @@ func copyMessages(messages []Message) []Message {
 //
 // The text names the tool so the model can tell which call an image came from
 // when several ran in one turn — the images arrive detached from their tool
-// result, so nothing else associates them.
-func toolResultImageMessage(result ToolResult) (zeroruntime.Message, bool) {
+// result, so nothing else associates them. If the effective model cannot
+// accept images, the tool's text result is left unchanged and this follow-up
+// is a notice with no image parts, matching the CLI/TUI vision gate.
+func toolResultImageMessage(result ToolResult, options Options) (zeroruntime.Message, bool) {
 	images := make([]zeroruntime.ImageBlock, 0, len(result.Images))
 	for _, image := range result.Images {
 		if len(image.Data) == 0 {
@@ -3479,9 +3493,26 @@ func toolResultImageMessage(result ToolResult) (zeroruntime.Message, bool) {
 	if label == "" {
 		label = "tool"
 	}
+	if !modelAcceptsToolImages(options) {
+		return zeroruntime.Message{
+			Role:    zeroruntime.MessageRoleUser,
+			Content: "Image output from " + label + " was not sent because the current model does not support image input.",
+		}, true
+	}
 	return zeroruntime.Message{
 		Role:    zeroruntime.MessageRoleUser,
 		Content: "Image output from " + label + ":",
 		Images:  images,
 	}, true
 }
+
+func modelAcceptsToolImages(options Options) bool {
+	if options.SupportsVision != nil {
+		return options.SupportsVision(options.Model)
+	}
+	registry, err := modelregistry.DefaultRegistry()
+	if err != nil {
+		return modelregistry.VisionCapableByName(options.Model)
+	}
+	return modelregistry.SupportsVision(registry, options.Model)
+}

internal/agent/tool_result_images_test.go

@@ -62,6 +62,7 @@ func TestRunDeliversToolResultImagesToTheModel(t *testing.T) {
 	result, err := Run(context.Background(), "screenshot please", provider, Options{
 		Registry: registry,
 		MaxTurns: 2,
+		Model:    "gpt-4o",
 	})
 	if err != nil {
 		t.Fatalf("Run: %v", err)
@@ -136,7 +137,7 @@ func TestRunKeepsToolResultsContiguousWhenAToolReturnsAnImage(t *testing.T) {
 		{{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}},
 	}}
 
-	result, err := Run(context.Background(), "two calls", provider, Options{Registry: registry, MaxTurns: 2})
+	result, err := Run(context.Background(), "two calls", provider, Options{Registry: registry, MaxTurns: 2, Model: "gpt-4o"})
 	if err != nil {
 		t.Fatalf("Run: %v", err)
 	}
@@ -205,3 +206,211 @@ func messageShape(messages []zeroruntime.Message) string {
 	}
 	return "[" + strings.Join(parts, " ") + "]"
 }
+
+func TestRunDeliversToolResultImagesToAVisionModel(t *testing.T) {
+	registry := tools.NewRegistry()
+	registry.Register(imageTool{media: "image/png", data: []byte("\x89PNG\r\n\x1a\nfake")})
+	provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
+		{
+			{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "capture"},
+			{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`},
+			{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"},
+			{Type: zeroruntime.StreamEventDone},
+		},
+		{{Type: zeroruntime.StreamEventText, Content: "I can see it."}, {Type: zeroruntime.StreamEventDone}},
+	}}
+
+	result, err := Run(context.Background(), "screenshot please", provider, Options{
+		Registry: registry,
+		MaxTurns: 2,
+		Model:    "gpt-4o",
+	})
+	if err != nil {
+		t.Fatalf("Run: %v", err)
+	}
+	var carrier *zeroruntime.Message
+	var toolText string
+	for index := range result.Messages {
+		if result.Messages[index].Role == zeroruntime.MessageRoleTool {
+			toolText = result.Messages[index].Content
+		}
+		if len(result.Messages[index].Images) > 0 {
+			carrier = &result.Messages[index]
+		}
+	}
+	if toolText != "Captured a screenshot." {
+		t.Fatalf("tool text = %q, want preserved", toolText)
+	}
+	if carrier == nil {
+		t.Fatal("vision model must receive the tool image")
+	}
+}
+
+func TestRunDropsToolResultImagesForANonVisionModel(t *testing.T) {
+	registry := tools.NewRegistry()
+	registry.Register(imageTool{media: "image/png", data: []byte("\x89PNG\r\n\x1a\nfake")})
+	provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
+		{
+			{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "capture"},
+			{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`},
+			{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"},
+			{Type: zeroruntime.StreamEventDone},
+		},
+		{{Type: zeroruntime.StreamEventText, Content: "ok"}, {Type: zeroruntime.StreamEventDone}},
+	}}
+
+	result, err := Run(context.Background(), "screenshot please", provider, Options{
+		Registry: registry,
+		MaxTurns: 2,
+		Model:    "totally-made-up-model",
+	})
+	if err != nil {
+		t.Fatalf("Run: %v", err)
+	}
+	var toolText string
+	var noticed bool
+	for _, message := range result.Messages {
+		if message.Role == zeroruntime.MessageRoleTool {
+			toolText = message.Content
+		}
+		if len(message.Images) > 0 {
+			t.Fatalf("non-vision model received image bytes on %q: %s", message.Role, messageShape(result.Messages))
+		}
+		if strings.Contains(message.Content, "does not support image input") {
+			noticed = true
+		}
+	}
+	if toolText != "Captured a screenshot." {
+		t.Fatalf("tool text = %q, want preserved on the non-vision path", toolText)
+	}
+	if !noticed {
+		t.Fatalf("non-vision model was not told the image was dropped; recorded %s", messageShape(result.Messages))
+	}
+}
+
+type switchAndCaptureTool struct {
+	targetModel string
+}
+
+func (switchAndCaptureTool) Name() string        { return "switch_and_capture" }
+func (switchAndCaptureTool) Description() string { return "Switches model and captures image" }
+func (switchAndCaptureTool) Parameters() tools.Schema {
+	return tools.Schema{Type: "object", Properties: map[string]tools.PropertySchema{}}
+}
+func (switchAndCaptureTool) Safety() tools.Safety {
+	return tools.Safety{Permission: tools.PermissionAllow}
+}
+func (t switchAndCaptureTool) Run(context.Context, map[string]any) tools.Result {
+	return tools.Result{
+		Status: tools.StatusOK,
+		Output: "[image returned by tool]",
+		Images: []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte("\x89PNG\r\n\x1a\nfake")}},
+		Meta:   map[string]string{"escalate_to_model": t.targetModel},
+	}
+}
+
+func TestRunToolImagesRespectsModelSwitch(t *testing.T) {
+	t.Run("Non-vision to vision model escalation forwards images", func(t *testing.T) {
+		registry := tools.NewRegistry()
+		registry.Register(switchAndCaptureTool{targetModel: "gpt-4o"})
+
+		provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
+			{
+				{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "switch_and_capture"},
+				{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`},
+				{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"},
+				{Type: zeroruntime.StreamEventDone},
+			},
+			{
+				{Type: zeroruntime.StreamEventText, Content: "I see it on gpt-4o."},
+				{Type: zeroruntime.StreamEventDone},
+			},
+		}}
+
+		switched := false
+		switcher := func(_ context.Context, modelID string) (Provider, error) {
+			switched = true
+			return provider, nil
+		}
+
+		result, err := Run(context.Background(), "test", provider, Options{
+			Registry:      registry,
+			MaxTurns:      2,
+			Model:         "non-vision-initial",
+			ModelSwitcher: switcher,
+			SupportsVision: func(modelID string) bool {
+				return modelID == "gpt-4o"
+			},
+		})
+		if err != nil {
+			t.Fatalf("Run: %v", err)
+		}
+		if !switched {
+			t.Fatal("expected model switch to occur")
+		}
+
+		var carrier *zeroruntime.Message
+		for index := range result.Messages {
+			if len(result.Messages[index].Images) > 0 {
+				carrier = &result.Messages[index]
+			}
+			if strings.Contains(result.Messages[index].Content, "does not support image input") {
+				t.Fatalf("images were unexpectedly dropped after switch to vision model: %v", result.Messages[index])
+			}
+		}
+		if carrier == nil {
+			t.Fatal("expected image carrier message after escalation to vision model")
+		}
+	})
+
+	t.Run("Vision to non-vision model switch drops images with notice", func(t *testing.T) {
+		registry := tools.NewRegistry()
+		registry.Register(switchAndCaptureTool{targetModel: "text-only-dest"})
+
+		provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
+			{
+				{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "switch_and_capture"},
+				{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`},
+				{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"},
+				{Type: zeroruntime.StreamEventDone},
+			},
+			{
+				{Type: zeroruntime.StreamEventText, Content: "ok."},
+				{Type: zeroruntime.StreamEventDone},
+			},
+		}}
+
+		switcher := func(_ context.Context, modelID string) (Provider, error) {
+			return provider, nil
+		}
+
+		result, err := Run(context.Background(), "test", provider, Options{
+			Registry:      registry,
+			MaxTurns:      2,
+			Model:         "gpt-4o",
+			ModelSwitcher: switcher,
+			SupportsVision: func(modelID string) bool {
+				return modelID == "gpt-4o"
+			},
+		})
+		if err != nil {
+			t.Fatalf("Run: %v", err)
+		}
+
+		var noticed bool
+		for _, m := range result.Messages {
+			if len(m.Images) > 0 {
+				t.Fatalf("image delivered to non-vision destination model: %v", m)
+			}
+			if strings.Contains(m.Content, "does not support image input") {
+				noticed = true
+			}
+			if strings.Contains(m.Content, "[image forwarded]") {
+				t.Fatalf("contradictory forwarded text found in message: %v", m)
+			}
+		}
+		if !noticed {
+			t.Fatal("expected drop notice when switching to non-vision model")
+		}
+	})
+}

internal/agent/types.go

@@ -340,6 +340,11 @@ type Options struct {
 	// nil for text-only runs (the seeded message then carries no images, exactly
 	// as before).
 	Images []zeroruntime.ImageBlock
+	// SupportsVision, when set, reports whether the effective model accepts
+	// image input. Tool-produced images are dropped at the shared delivery
+	// boundary when this is false. nil uses the curated catalog plus the
+	// name heuristic via modelregistry.SupportsVision.
+	SupportsVision func(modelID string) bool
 	// ContextWindow is the model's maximum input token budget. When > 0 the agent
 	// loop compacts long conversations once the estimated size crosses a fraction
 	// of this window. 0 DISABLES compaction entirely (every existing caller/test

internal/cli/exec.go

@@ -686,6 +686,9 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
 		ContextWindowFor: func(modelID string) int {
 			return modelregistry.AgentContextWindow(modelContextWindow(modelRegistry, modelID))
 		},
+		SupportsVision: func(modelID string) bool {
+			return modelregistry.SupportsVision(modelRegistry, modelID)
+		},
 		ReasoningEffort: forwardEffort,
 		Trace:           traceRecorder,
 		Cwd:             workspaceRoot,