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

internal/acp/agent.go

@@ -5,12 +5,14 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"log"
 	"strings"
 	"sync"
 
 	"github.com/Gitlawb/zero/internal/agent"
 	"github.com/Gitlawb/zero/internal/config"
+	"github.com/Gitlawb/zero/internal/modelregistry"
 	"github.com/Gitlawb/zero/internal/providercatalog"
 	"github.com/Gitlawb/zero/internal/providermodelcatalog"
 	"github.com/Gitlawb/zero/internal/providermodeldiscovery"
@@ -249,6 +251,29 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
 	}
 	note := &notifier{conn: a.conn, sessionID: sess.id}
 
+	visionCache := make(map[string]bool)
+	var visionCacheMu sync.Mutex
+	supportsVision := func(modelID string) bool {
+		modelID = strings.TrimSpace(modelID)
+		if modelID == "" {
+			return false
+		}
+		visionCacheMu.Lock()
+		defer visionCacheMu.Unlock()
+		if cached, ok := visionCache[modelID]; ok {
+			return cached
+		}
+		supported := a.modelSupportsVision(ctx, resolved.Provider, modelID)
+		visionCache[modelID] = supported
+		return supported
+	}
+	if len(images) > 0 && !supportsVision(resolved.Provider.Model) {
+		msg := fmt.Sprintf("Model %s does not support image input; ignoring %d prompt image(s).", resolved.Provider.Model, len(images))
+		note.text("[zero] " + msg + "\n\n")
+		userText = fmt.Sprintf("[Note: %s]\n\n%s", msg, userText)
+		images = nil
+	}
+
 	opts := agent.Options{
 		Cwd:            sess.cwd,
 		SessionID:      sess.id,
@@ -259,6 +284,7 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string,
 		PermissionMode: sess.currentMode(),
 		MaxTurns:       resolved.MaxTurns,
 		Images:         images,
+		SupportsVision: supportsVision,
 		OnText:         note.text,
 		OnReasoning:    note.thought,
 		OnToolCall:     note.toolCall,
@@ -767,3 +793,32 @@ func (s *acpSession) snapshotHistory() []turnRecord {
 	defer s.mu.Unlock()
 	return append([]turnRecord(nil), s.history...)
 }
+
+func (a *Agent) modelSupportsVision(ctx context.Context, profile config.ProviderProfile, modelID string) bool {
+	trimmed := strings.TrimSpace(modelID)
+	if trimmed == "" {
+		return false
+	}
+	if a.deps.DiscoverModels != nil {
+		if discovered, err := a.deps.DiscoverModels(ctx, profile); err == nil {
+			for _, dm := range discovered {
+				if strings.EqualFold(strings.TrimSpace(dm.ID), trimmed) {
+					if len(dm.InputModalities) > 0 {
+						for _, mod := range dm.InputModalities {
+							if strings.EqualFold(strings.TrimSpace(mod), "image") {
+								return true
+							}
+						}
+						return false
+					}
+					break
+				}
+			}
+		}
+	}
+	reg, _ := modelregistry.DefaultRegistry()
+	if entry, known := reg.Resolve(trimmed); known {
+		return entry.Supports(modelregistry.ModelCapabilityVision)
+	}
+	return modelregistry.SupportsVision(reg, trimmed)
+}

internal/acp/agent_test.go

@@ -1,14 +1,17 @@
 package acp
 
 import (
+	"bytes"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -497,6 +500,48 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) {
 	}
 }
 
+func TestACPWiresSupportsVision(t *testing.T) {
+	deps := testDeps(t)
+	var captured agent.Options
+	deps.RunAgent = func(_ context.Context, _ string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
+		captured = opts
+		return agent.Result{FinalAnswer: "done"}, nil
+	}
+	deps.DiscoverModels = func(_ context.Context, _ config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
+		return []providermodeldiscovery.Model{
+			{ID: "custom-vision", InputModalities: []string{"text", "image"}},
+			{ID: "custom-text", InputModalities: []string{"text"}},
+		}, nil
+	}
+	h := newHarness(t, deps)
+	defer h.stop()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	var newRes NewSessionResult
+	if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
+		t.Fatalf("session/new: %v", err)
+	}
+	var promptRes PromptResult
+	if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
+		SessionID: newRes.SessionID,
+		Prompt: []ContentBlock{
+			{Type: "text", Text: "hello"},
+		},
+	}, &promptRes); err != nil {
+		t.Fatalf("session/prompt: %v", err)
+	}
+	if captured.SupportsVision == nil {
+		t.Fatal("SupportsVision was not wired into agent.Options")
+	}
+	if !captured.SupportsVision("custom-vision") {
+		t.Fatal("SupportsVision(custom-vision) = false, want true")
+	}
+	if captured.SupportsVision("custom-text") {
+		t.Fatal("SupportsVision(custom-text) = true, want false")
+	}
+}
+
 // TestACPRejectsInvalidCwd confirms session/new fails when the workspace root
 // resolver rejects the client cwd (e.g. filesystem root).
 func TestACPRejectsInvalidCwd(t *testing.T) {
@@ -599,3 +644,200 @@ func drainTextUntil(t *testing.T, ch <-chan string, done func(string) bool) stri
 		}
 	}
 }
+
+func TestACPPromptUnsupportedModelDropsImagesAndNotifies(t *testing.T) {
+	deps := testDeps(t)
+	deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
+		return []providermodeldiscovery.Model{
+			{ID: "fake-model", InputModalities: []string{"text"}},
+		}, nil
+	}
+	var capturedOpts agent.Options
+	var capturedPrompt string
+	deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
+		capturedPrompt = prompt
+		capturedOpts = opts
+		return agent.Result{FinalAnswer: "Answered without images"}, nil
+	}
+
+	h := newHarness(t, deps)
+	defer h.stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	var initRes InitializeResult
+	if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil {
+		t.Fatalf("initialize: %v", err)
+	}
+
+	var newRes NewSessionResult
+	if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
+		t.Fatalf("session/new: %v", err)
+	}
+
+	rawPng := "\x89PNG\r\n\x1a\nfake-image-bytes"
+	b64Png := base64.StdEncoding.EncodeToString([]byte(rawPng))
+	var promptRes PromptResult
+	if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
+		SessionID: newRes.SessionID,
+		Prompt: []ContentBlock{
+			TextBlock("look at this image"),
+			ImageBlock(b64Png, "image/png"),
+		},
+	}, &promptRes); err != nil {
+		t.Fatalf("session/prompt: %v", err)
+	}
+	if promptRes.StopReason != StopEndTurn {
+		t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn)
+	}
+
+	got := drainTextUntil(t, h.updates, func(text string) bool {
+		return strings.Contains(text, "does not support image input")
+	})
+	if !strings.Contains(got, "Model fake-model does not support image input; ignoring 1 prompt image(s).") {
+		t.Fatalf("streamed text = %q, want drop notice", got)
+	}
+	if len(capturedOpts.Images) != 0 {
+		t.Fatalf("captured %d images, want 0 (images should be withheld)", len(capturedOpts.Images))
+	}
+	if !strings.Contains(capturedPrompt, "[Note: Model fake-model does not support image input; ignoring 1 prompt image(s).]") {
+		t.Fatalf("captured prompt = %q, want prompt note", capturedPrompt)
+	}
+}
+
+func TestACPSessionPrompt_SupportsVisionMemoizedPerRun(t *testing.T) {
+	deps := testDeps(t)
+	var discoverCount int
+	var discoverMu sync.Mutex
+	deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
+		discoverMu.Lock()
+		discoverCount++
+		discoverMu.Unlock()
+		return []providermodeldiscovery.Model{
+			{ID: "fake-model", InputModalities: []string{"text", "image"}},
+		}, nil
+	}
+	deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
+		for i := 0; i < 5; i++ {
+			if !opts.SupportsVision(opts.Model) {
+				t.Error("expected SupportsVision to report true")
+			}
+		}
+		return agent.Result{FinalAnswer: "Done"}, nil
+	}
+
+	h := newHarness(t, deps)
+	defer h.stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	var initRes InitializeResult
+	if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil {
+		t.Fatalf("initialize: %v", err)
+	}
+
+	var newRes NewSessionResult
+	if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
+		t.Fatalf("session/new: %v", err)
+	}
+
+	discoverMu.Lock()
+	discoverCount = 0
+	discoverMu.Unlock()
+
+	rawPng := "\x89PNG\r\n\x1a\nfake-image-bytes"
+	b64Png := base64.StdEncoding.EncodeToString([]byte(rawPng))
+	var promptRes PromptResult
+	if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
+		SessionID: newRes.SessionID,
+		Prompt: []ContentBlock{
+			TextBlock("look at this"),
+			ImageBlock(b64Png, "image/png"),
+		},
+	}, &promptRes); err != nil {
+		t.Fatalf("session/prompt: %v", err)
+	}
+	if promptRes.StopReason != StopEndTurn {
+		t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn)
+	}
+
+	discoverMu.Lock()
+	count := discoverCount
+	discoverMu.Unlock()
+	if count != 1 {
+		t.Fatalf("DiscoverModels was called %d times, want exactly 1 call (memoized per run)", count)
+	}
+}
+
+func TestACPEndToEndImagePromptSupportedModel(t *testing.T) {
+	deps := testDeps(t)
+	deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
+		return []providermodeldiscovery.Model{
+			{ID: "discovered-vision-model", InputModalities: []string{"text", "image"}},
+		}, nil
+	}
+	deps.ResolveConfig = func(_ string, o config.Overrides) (config.ResolvedConfig, error) {
+		model := "discovered-vision-model"
+		if o.Provider.Model != "" {
+			model = o.Provider.Model
+		}
+		return config.ResolvedConfig{
+			Provider: config.ProviderProfile{Name: "fake", Model: model},
+			MaxTurns: 4,
+		}, nil
+	}
+	var capturedOpts agent.Options
+	var capturedPrompt string
+	deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
+		capturedPrompt = prompt
+		capturedOpts = opts
+		return agent.Result{FinalAnswer: "I can see the image!"}, nil
+	}
+
+	h := newHarness(t, deps)
+	defer h.stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	var initRes InitializeResult
+	if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil {
+		t.Fatalf("initialize: %v", err)
+	}
+
+	var newRes NewSessionResult
+	if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
+		t.Fatalf("session/new: %v", err)
+	}
+
+	rawPng := []byte("\x89PNG\r\n\x1a\nreal-png-bytes")
+	b64Png := base64.StdEncoding.EncodeToString(rawPng)
+	var promptRes PromptResult
+	if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
+		SessionID: newRes.SessionID,
+		Prompt: []ContentBlock{
+			TextBlock("analyze this image"),
+			ImageBlock(b64Png, "image/png"),
+		},
+	}, &promptRes); err != nil {
+		t.Fatalf("session/prompt: %v", err)
+	}
+	if promptRes.StopReason != StopEndTurn {
+		t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn)
+	}
+
+	if len(capturedOpts.Images) != 1 {
+		t.Fatalf("captured %d images, want 1", len(capturedOpts.Images))
+	}
+	if capturedOpts.Images[0].MediaType != "image/png" {
+		t.Fatalf("mediaType = %q, want image/png", capturedOpts.Images[0].MediaType)
+	}
+	if !bytes.Equal(capturedOpts.Images[0].Data, rawPng) {
+		t.Fatalf("image bytes mismatch: got %v, want %v", capturedOpts.Images[0].Data, rawPng)
+	}
+	if strings.Contains(capturedPrompt, "does not support image input") {
+		t.Fatalf("prompt unexpectedly contains refusal note: %q", capturedPrompt)
+	}
+}