From dff1861c02807b7680d0c9f2db15464773a2ebc0 Mon Sep 17 00:00:00 2001 From: Marlowe <321669285+cairn-intern@users.noreply.github.com> Date: Thu, 24 Sep 2026 00:37:57 -0400 Subject: [PATCH 1/2] fix(mcp): forward image tool results on the existing image channel Recreated from Gitlawb/zero#1012 (approved but unmerged). Original: euxaristia/zero:feat/823-forward-mcp-images --- internal/acp/agent.go | 57 +- internal/acp/agent_test.go | 241 ++++++++ internal/agent/loop.go | 73 ++- internal/agent/tool_result_images_test.go | 199 ++++++- internal/agent/types.go | 5 + internal/cli/exec.go | 3 + internal/cli/exec_spec.go | 15 +- internal/mcp/client.go | 176 +++++- internal/mcp/client_test.go | 2 +- internal/mcp/export_test.go | 19 + internal/mcp/network_client.go | 9 +- internal/mcp/network_client_test.go | 78 +++ internal/mcp/non_text_content_test.go | 526 +++++++++++++++++- internal/mcp/registry.go | 74 ++- internal/modelregistry/vision.go | 11 +- internal/modelregistry/vision_name_test.go | 5 + internal/providermodeldiscovery/discovery.go | 38 ++ .../providermodeldiscovery/discovery_test.go | 43 ++ internal/tui/image_attach.go | 54 +- internal/tui/image_attach_test.go | 357 ++++++++++++ internal/tui/model.go | 36 ++ internal/tui/picker.go | 141 ++++- 22 files changed, 2043 insertions(+), 119 deletions(-) create mode 100644 internal/mcp/export_test.go diff --git a/internal/acp/agent.go b/internal/acp/agent.go index 7c301044c..442ae4b40 100644 --- a/internal/acp/agent.go +++ b/internal/acp/agent.go @@ -15,6 +15,7 @@ import ( "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" @@ -508,6 +509,30 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string, } queue(messageEvent("user", userText)) + 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 + } + effectivePrompt := userText + 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") + effectivePrompt = fmt.Sprintf("[Note: %s]\n\n%s", msg, userText) + images = nil + } + opts := agent.Options{ Cwd: sess.cwd, SessionID: sess.id, @@ -519,6 +544,7 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string, DeferThreshold: workspace.DeferThreshold, MaxTurns: resolved.MaxTurns, Images: images, + SupportsVision: supportsVision, OnText: note.text, OnReasoning: note.thought, OnToolCall: func(call agent.ToolCall) { @@ -537,7 +563,7 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string, }, } - agentPrompt := buildPrompt(sess.snapshotHistory(), userText) + agentPrompt := buildPrompt(sess.snapshotHistory(), effectivePrompt) result, runErr := a.deps.RunAgent(ctx, agentPrompt, provider, opts) if result.FinalAnswer != "" { queue(messageEvent("assistant", result.FinalAnswer)) @@ -1305,3 +1331,32 @@ func sameWorkspace(left, right string) bool { } return os.SameFile(leftInfo, rightInfo) } + +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) +} diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index 41457549e..564f5a073 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -1,7 +1,9 @@ package acp import ( + "bytes" "context" + "encoding/base64" "encoding/json" "errors" "fmt" @@ -847,6 +849,48 @@ func TestACPRunTurnClosesWorkspaceAfterSuccessAndFailure(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) { @@ -2418,3 +2462,200 @@ func payloadString(payload any, key string) string { value, _ := decoded[key].(string) return value } + +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) + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 0d06d82f4..4e75afcdb 100644 --- a/internal/agent/loop.go +++ b/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 @@ -3464,8 +3476,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 { @@ -3486,9 +3500,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. Switch to a vision-capable model before retrying image delivery, including any images skipped by the tool's image budget.", + }, 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) +} diff --git a/internal/agent/tool_result_images_test.go b/internal/agent/tool_result_images_test.go index f6a5d0335..6a357b1b2 100644 --- a/internal/agent/tool_result_images_test.go +++ b/internal/agent/tool_result_images_test.go @@ -10,6 +10,24 @@ import ( "github.com/Gitlawb/zero/internal/zeroruntime" ) +func TestNonVisionImageNoticeQualifiesBudgetRecovery(t *testing.T) { + result := ToolResult{ + Name: "mcp_screenshot", + Output: "One image exceeded this result's remaining image budget. Retrying with fewer images can recover this payload.", + Images: []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte("image")}}, + } + notice, ok := toolResultImageMessage(result, Options{ + SupportsVision: func(string) bool { return false }, + }) + if !ok || len(notice.Images) != 0 { + t.Fatal("non-vision delivery must produce a text-only notice") + } + if !strings.Contains(notice.Content, "Switch to a vision-capable model before retrying") || + !strings.Contains(notice.Content, "budget") { + t.Fatalf("notice does not qualify budget recovery: %s", notice.Content) + } +} + // imageTool returns a tool result carrying an image, the way a screenshot tool // should be able to. type imageTool struct { @@ -62,6 +80,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) @@ -86,11 +105,16 @@ func TestRunDeliversToolResultImagesToTheModel(t *testing.T) { // The tool-result pairing must be untouched: one tool result per tool call. toolResults := 0 + var toolText string for _, message := range result.Messages { if message.Role == zeroruntime.MessageRoleTool { toolResults++ + toolText = message.Content } } + if toolText != "Captured a screenshot." { + t.Fatalf("tool text = %q, want preserved", toolText) + } if toolResults != 1 { t.Errorf("got %d tool-result messages for 1 tool call", toolResults) } @@ -136,7 +160,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 +229,176 @@ func messageShape(messages []zeroruntime.Message) string { } return "[" + strings.Join(parts, " ") + "]" } + +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 + var toolOutputSeen 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 returned by tool]") { + toolOutputSeen = true + } + } + if !noticed { + t.Fatal("expected drop notice when switching to non-vision model") + } + if !toolOutputSeen { + t.Fatal("expected tool output preserved in message history") + } + }) +} diff --git a/internal/agent/types.go b/internal/agent/types.go index 19ff955bc..a373d4f46 100644 --- a/internal/agent/types.go +++ b/internal/agent/types.go @@ -341,6 +341,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 diff --git a/internal/cli/exec.go b/internal/cli/exec.go index d47964c11..2820d03b9 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -660,6 +660,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, diff --git a/internal/cli/exec_spec.go b/internal/cli/exec_spec.go index fc22eed35..45ba2f6ca 100644 --- a/internal/cli/exec_spec.go +++ b/internal/cli/exec_spec.go @@ -118,12 +118,15 @@ func runExecSpecDraft(run execSpecDraftRun) int { hookDispatcher, hookSkip := newHookDispatcher(run.workspaceRoot, run.trustRoot, execution.NewRunner(run.sandboxEngine)) emitTrustNotice(run.stderr, hookSkip, run.mcpSkip) result, err := agent.Run(runCtx, run.prompt, run.provider, agent.Options{ - MaxTurns: run.resolved.MaxTurns, - ContextWindow: resolveAgentContextWindow(runCtx, run.modelRegistry, run.resolved.Provider), - SessionID: draftSession.SessionID, - SessionTitle: run.sessionTitle, - ProviderName: run.resolved.Provider.Name, - Model: run.resolved.Provider.Model, + MaxTurns: run.resolved.MaxTurns, + ContextWindow: resolveAgentContextWindow(runCtx, run.modelRegistry, run.resolved.Provider), + SessionID: draftSession.SessionID, + SessionTitle: run.sessionTitle, + ProviderName: run.resolved.Provider.Name, + Model: run.resolved.Provider.Model, + SupportsVision: func(modelID string) bool { + return modelregistry.SupportsVision(run.modelRegistry, modelID) + }, ReasoningEffort: run.reasoningEffort, Profile: run.profilePolicy, Cwd: run.workspaceRoot, diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 932b4c371..5b4851c47 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -3,11 +3,13 @@ package mcp import ( "bytes" "context" + "encoding/base64" "encoding/json" "errors" "fmt" "io" "math" + "net/http" "os" "os/exec" "strconv" @@ -16,6 +18,8 @@ import ( "time" "github.com/Gitlawb/zero/internal/execution" + "github.com/Gitlawb/zero/internal/imageinput" + "github.com/Gitlawb/zero/internal/zeroruntime" ) type RemoteTool struct { @@ -27,10 +31,16 @@ type RemoteTool struct { type Content struct { Type string `json:"type"` Text string `json:"text,omitempty"` - // MimeType names what a non-text block holds. Decoded but not yet forwarded: - // it is what lets a dropped block be described to the model instead of - // vanishing (#823). Servers that omit it still decode fine. + // MimeType names what a non-text block holds. Additive and omitempty, so + // servers that never send it decode exactly as before. Image blocks with + // valid data are forwarded on Result.Images; MimeType is what lets a + // remaining dropped block (audio, resource, failed decode) be described + // instead of vanishing (#823). MimeType string `json:"mimeType,omitempty"` + // Data is the MCP image block's base64 payload. Additive and omitempty so a + // result that never sent it still unmarshals. Decoded only for type + // "image"; other types leave it unread. + Data string `json:"data,omitempty"` } type CallToolResult struct { @@ -381,10 +391,6 @@ func (client *Client) request(ctx context.Context, method string, params any, ta } } -func (client *Client) ensureWriter() { - _ = client.startWriter() -} - func (client *Client) startWriter() error { client.writeMu.Lock() defer client.writeMu.Unlock() @@ -746,24 +752,95 @@ func TextContent(content []Content) string { return strings.TrimSpace(strings.Join(parts, "\n")) } -// DroppedContentSummary describes the blocks TextContent discards, e.g. -// "1 image/png block" or "2 resource blocks, 1 audio/wav block". It returns "" -// when a result is entirely text, so a caller adds nothing to the ordinary case. -// -// This exists because dropping silently is the worst available behaviour. A -// screenshot server returns a valid image, TextContent keeps nothing, and the -// call is reported as "(empty MCP tool result)" — so the model concludes the -// tool produced nothing and usually retries, burning another call on the same -// empty answer. Naming what came back costs nothing and ends that loop even -// though the payload still cannot be forwarded. -// -// Counts are grouped by mime type and ordered by first appearance, so the same -// result always produces the same sentence. -func DroppedContentSummary(content []Content) string { +// itemDisp is the per-item forwarding/drop disposition produced by the +// single-pass conversion. DroppedContentSummary is built from this so a +// valid image is never base64-decoded a second time just to name what was +// kept versus dropped. +type itemDisp uint8 + +const ( + dispText itemDisp = iota + dispForwarded + dispDropped + dispBudgetExceeded + dispUninspected +) + +// decodeImageBase64 is the MCP image payload decoder. Tests replace it to +// count decode attempts; production uses standard base64. +var decodeImageBase64 = base64.StdEncoding.DecodeString + +// Bound provider content-block overhead independently of decoded image size. +const maxForwardedImages = 16 + +// Bound decode work even when invalid or oversized-for-the-residue candidates +// do not consume the forwarding budgets. +const maxInspectedImages = 32 + +type imageLimitReason int + +const ( + imageLimitNone imageLimitReason = iota + imageLimitBudget + imageLimitForwardCount + imageLimitInspectionCount +) + +func forwardImages(content []Content) ([]zeroruntime.ImageBlock, []itemDisp, imageLimitReason) { + disp := make([]itemDisp, len(content)) + var images []zeroruntime.ImageBlock + remaining := imageinput.MaxImageBytes + inspected := 0 + var limitReason imageLimitReason + for i, item := range content { + if item.Type == "text" { + disp[i] = dispText + continue + } + if item.Type == "image" { + if remaining == 0 { + disp[i] = dispUninspected + if limitReason == imageLimitNone { + limitReason = imageLimitBudget + } + continue + } + if len(images) >= maxForwardedImages { + disp[i] = dispUninspected + if limitReason == imageLimitNone { + limitReason = imageLimitForwardCount + } + continue + } + if inspected >= maxInspectedImages { + disp[i] = dispUninspected + if limitReason == imageLimitNone { + limitReason = imageLimitInspectionCount + } + continue + } + inspected++ + if image, ok := imageBlockFromContent(item); ok { + if len(image.Data) <= remaining { + images = append(images, image) + remaining -= len(image.Data) + disp[i] = dispForwarded + continue + } + disp[i] = dispBudgetExceeded + continue + } + } + disp[i] = dispDropped + } + return images, disp, limitReason +} + +func droppedContentNote(content []Content, disp []itemDisp, kinds ...itemDisp) string { labels := make([]string, 0, len(content)) counts := make(map[string]int, len(content)) - for _, item := range content { - if item.Type == "text" { + for i, item := range content { + if i >= len(disp) || !dispKind(disp[i], kinds) { continue } // Prefer the mime type: "image/png" tells the reader more than "image". @@ -794,3 +871,56 @@ func DroppedContentSummary(content []Content) string { } return strings.Join(parts, ", ") } + +func dispKind(got itemDisp, kinds []itemDisp) bool { + for _, kind := range kinds { + if got == kind { + return true + } + } + return false +} + +func dispCount(disp []itemDisp, kinds ...itemDisp) int { + count := 0 + for _, d := range disp { + if dispKind(d, kinds) { + count++ + } + } + return count +} + +func imageBlockFromContent(item Content) (zeroruntime.ImageBlock, bool) { + if item.Type != "image" { + return zeroruntime.ImageBlock{}, false + } + raw := strings.TrimSpace(item.Data) + if raw == "" { + return zeroruntime.ImageBlock{}, false + } + // EncodedLen(MaxImageBytes) is the encoded size of an image that decodes + // to exactly the inclusive cap, including "==" padding. DecodedLen is an + // upper bound and reports cap+2 for that input, so using it here would + // reject a valid at-limit PNG. The post-decode len(data) check is the + // exact backstop. + if len(raw) > base64.StdEncoding.EncodedLen(imageinput.MaxImageBytes) { + return zeroruntime.ImageBlock{}, false + } + data, err := decodeImageBase64(raw) + if err != nil { + return zeroruntime.ImageBlock{}, false + } + if len(data) == 0 || len(data) > imageinput.MaxImageBytes { + return zeroruntime.ImageBlock{}, false + } + sniffLen := len(data) + if sniffLen > 512 { + sniffLen = 512 + } + mediaType := zeroruntime.NormalizeImageMediaType(http.DetectContentType(data[:sniffLen])) + if mediaType == "" { + return zeroruntime.ImageBlock{}, false + } + return zeroruntime.ImageBlock{MediaType: mediaType, Data: data}, true +} diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 938344f7f..b2ea41ab1 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -1420,7 +1420,7 @@ func TestWriterLoopExitsOnRepeatedClose(t *testing.T) { reader: newMessageReader(inReader), writer: newMessageWriter(outWriter), } - client.ensureWriter() + _ = client.startWriter() if err := client.writeMessage(context.Background(), rpcMessage{Method: "notifications/ping"}); err != nil { t.Fatalf("writeMessage: %v", err) } diff --git a/internal/mcp/export_test.go b/internal/mcp/export_test.go new file mode 100644 index 000000000..f08ff5d5c --- /dev/null +++ b/internal/mcp/export_test.go @@ -0,0 +1,19 @@ +package mcp + +import "github.com/Gitlawb/zero/internal/zeroruntime" + +// DroppedContentSummary describes the blocks that were not forwarded, e.g. +// "1 audio/wav block" or "2 resource blocks, 1 image/png block". It returns "" +// when every block is text or an image that ImageBlocks successfully +// forwarded, so a caller adds nothing to the ordinary case. +func DroppedContentSummary(content []Content) string { + _, disp, _ := forwardImages(content) + return droppedContentNote(content, disp, dispDropped, dispBudgetExceeded, dispUninspected) +} + +// ImageBlocks converts MCP image content into the same ImageBlock channel +// capture tools already use. +func ImageBlocks(content []Content) []zeroruntime.ImageBlock { + images, _, _ := forwardImages(content) + return images +} diff --git a/internal/mcp/network_client.go b/internal/mcp/network_client.go index ce0b41162..d993f747a 100644 --- a/internal/mcp/network_client.go +++ b/internal/mcp/network_client.go @@ -657,11 +657,10 @@ func decodeSSERPCMessage(reader io.Reader) (rpcMessage, error) { return rpcMessage{}, fmt.Errorf("missing MCP SSE response data") } -// maxSSEEventBytes bounds a single SSE line/event. The previous 1 MiB cap made a -// large but legitimate MCP message (e.g. a big tool result) hit bufio.ErrTooLong, -// which failed the request permanently with no recovery. Raise it to a generous -// bound that still protects against an unbounded remote server. -const maxSSEEventBytes = 8 * 1024 * 1024 +// maxSSEEventBytes bounds a single SSE line/event. It accommodates multi-image +// responses (such as 8 MiB + 4 MiB) in base64 (~16.8 MiB) plus JSON-RPC envelope, +// metadata, and SSE framing overhead. +const maxSSEEventBytes = 32 * 1024 * 1024 func scanSSEEvents(reader io.Reader, handle func(sseEvent) bool) error { scanner := bufio.NewScanner(reader) diff --git a/internal/mcp/network_client_test.go b/internal/mcp/network_client_test.go index 8cb0d2f13..5e1e410ce 100644 --- a/internal/mcp/network_client_test.go +++ b/internal/mcp/network_client_test.go @@ -2,6 +2,7 @@ package mcp import ( "context" + "encoding/json" "net/http" "net/http/httptest" "net/url" @@ -440,3 +441,80 @@ func TestDeliverEventMessageSkipsMethodPresence(t *testing.T) { t.Fatal("true response must complete pending") } } + +func TestScanSSEEventsLargeImagePayload(t *testing.T) { + // Construct a 10 MiB image payload (~13.98 MiB in base64). + tenMiB := 10 * 1024 * 1024 + imgB64 := paddedPNGBase64(tenMiB) + resultJSON, err := json.Marshal(map[string]any{ + "content": []map[string]any{ + {"type": "text", "text": "analysis screenshot"}, + {"type": "image", "mimeType": "image/png", "data": imgB64}, + }, + }) + if err != nil { + t.Fatal(err) + } + rpcJSON, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "result": json.RawMessage(resultJSON), + }) + if err != nil { + t.Fatal(err) + } + + t.Run("Single data line with 10 MiB image", func(t *testing.T) { + stream := "event: message\ndata: " + string(rpcJSON) + "\n\n" + msg, err := decodeSSERPCMessage(strings.NewReader(stream)) + if err != nil { + t.Fatalf("decodeSSERPCMessage failed on single-line 10 MiB image: %v", err) + } + if !rpcIDMatches(msg.ID, 1) { + t.Fatalf("expected id 1, got %#v", msg.ID) + } + var decoded CallToolResult + if err := json.Unmarshal(msg.Result, &decoded); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if len(decoded.Content) != 2 || len(decoded.Content[1].Data) != len(imgB64) { + t.Fatalf("image payload did not survive the SSE round trip: got %d bytes, want %d", + len(decoded.Content[1].Data), len(imgB64)) + } + }) + + t.Run("Multi data lines with 10 MiB image", func(t *testing.T) { + stream := "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\ndata: \"result\":" + string(resultJSON) + "}\n\n" + msg, err := decodeSSERPCMessage(strings.NewReader(stream)) + if err != nil { + t.Fatalf("decodeSSERPCMessage failed on multi-line 10 MiB image: %v", err) + } + if !rpcIDMatches(msg.ID, 1) { + t.Fatalf("expected id 1, got %#v", msg.ID) + } + var decoded CallToolResult + if err := json.Unmarshal(msg.Result, &decoded); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if len(decoded.Content) != 2 || len(decoded.Content[1].Data) != len(imgB64) { + t.Fatalf("image payload did not survive the SSE round trip: got %d bytes, want %d", + len(decoded.Content[1].Data), len(imgB64)) + } + }) + + t.Run("Oversized event exceeding 32 MiB rejected cleanly", func(t *testing.T) { + oversizedRPC, marshalErr := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "result": map[string]any{"padding": strings.Repeat("A", 33*1024*1024)}, + }) + if marshalErr != nil { + t.Fatal(marshalErr) + } + stream := "event: message\ndata: " + string(oversizedRPC) + "\n\n" + _, err := decodeSSERPCMessage(strings.NewReader(stream)) + if err == nil { + t.Fatal("expected error for 33 MiB event, got nil") + } + }) +} diff --git a/internal/mcp/non_text_content_test.go b/internal/mcp/non_text_content_test.go index a1082962f..7608d630e 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -2,20 +2,35 @@ package mcp import ( "context" + "encoding/base64" + "encoding/json" "strings" "testing" + "github.com/Gitlawb/zero/internal/imageinput" "github.com/Gitlawb/zero/internal/tools" ) -// A server that returns only an image currently reports "(empty MCP tool -// result)": TextContent keeps text blocks and drops the rest, so a successful -// call looks like it produced nothing. The model then usually retries, which is -// the worst outcome, and the user is never told an image existed (#823). +// tinyPNGBase64 is a 1x1 PNG. http.DetectContentType sniffs it as image/png. +const tinyPNGBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" + +// paddedPNGBase64 is a decodable image/png of decoded length n. The 8-byte +// PNG magic is enough for http.DetectContentType; the rest is padding so +// tests can hit size caps without committing a multi-mebibyte fixture. +func paddedPNGBase64(n int) string { + raw := make([]byte, n) + copy(raw, "\x89PNG\r\n\x1a\n") + return base64.StdEncoding.EncodeToString(raw) +} + +// A server that returns an image WITHOUT payload still cannot be forwarded, so +// the delivered result must name the block rather than report "(empty MCP tool +// result)". The model then usually retries, which is the worst outcome (#823). // -// Carrying the payload is a separate change. Naming what was dropped is what -// stops the retry loop, and it has to be true of the DELIVERED result, so this -// drives registryTool.Run rather than the helper alone. +// Image blocks that DO carry data are forwarded on Result.Images (see the +// payload tests below). This case is the remaining drop path, and it has to be +// true of the DELIVERED result, so this drives registryTool.Run rather than +// the helper alone. func TestAnImageOnlyResultSaysWhatItReturned(t *testing.T) { tool := registryTool{ client: &nonTextClient{content: []Content{ @@ -134,6 +149,24 @@ func TestDroppedContentSummaryNamesTheBlocks(t *testing.T) { }, want: "2 resource blocks, 1 audio/wav block", }, + { + name: "successfully forwarded image is not named as dropped", + content: []Content{{Type: "image", MimeType: "image/png", Data: tinyPNGBase64}}, + want: "", + }, + { + name: "malformed image data is still named", + content: []Content{{Type: "image", MimeType: "image/png", Data: "%%%not-base64%%%"}}, + want: "1 image/png block", + }, + { + name: "forwarded image plus audio names only the audio", + content: []Content{ + {Type: "image", MimeType: "image/png", Data: tinyPNGBase64}, + {Type: "audio", MimeType: "audio/wav"}, + }, + want: "1 audio/wav block", + }, } for _, test := range tests { @@ -156,3 +189,482 @@ func (client *nonTextClient) CallTool(context.Context, string, map[string]any) ( } func (client *nonTextClient) Close() error { return nil } + +func TestAnImageWithPayloadIsForwarded(t *testing.T) { + tool := registryTool{ + client: &nonTextClient{content: []Content{ + {Type: "image", MimeType: "image/png", Data: tinyPNGBase64}, + }}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + } + + result := tool.Run(context.Background(), map[string]any{}) + + if result.Output != "[image returned by tool]" { + t.Fatalf("image-only Output = %q, want [image returned by tool] so the tool_result is not an empty body", result.Output) + } + if strings.Contains(result.Output, "cannot forward") { + t.Fatalf("a forwarded image is still described as unforwardable:\n%s", result.Output) + } + if len(result.Images) != 1 { + t.Fatalf("Images len = %d, want 1", len(result.Images)) + } + if result.Images[0].MediaType != "image/png" { + t.Errorf("MediaType = %q, want image/png", result.Images[0].MediaType) + } + if len(result.Images[0].Data) == 0 { + t.Fatal("forwarded image has empty Data") + } + if DroppedContentSummary([]Content{{Type: "image", MimeType: "image/png", Data: tinyPNGBase64}}) != "" { + t.Fatal("drop summary named a successfully forwarded image") + } + if result.Status != tools.StatusOK { + t.Errorf("status = %v, want OK", result.Status) + } +} + +func TestAudioIsStillDroppedAndNamed(t *testing.T) { + tool := registryTool{ + client: &nonTextClient{content: []Content{ + {Type: "audio", MimeType: "audio/wav"}, + }}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "clip"}, + } + + result := tool.Run(context.Background(), map[string]any{}) + + if len(result.Images) != 0 { + t.Fatalf("audio was forwarded as an image: %#v", result.Images) + } + if !strings.Contains(result.Output, "audio/wav") { + t.Errorf("the output does not name the dropped audio:\n%s", result.Output) + } + if !strings.Contains(result.Output, "cannot forward yet") { + t.Errorf("the output does not say the audio cannot be forwarded:\n%s", result.Output) + } +} + +func TestTextAndImageKeepsTextAndForwardsImage(t *testing.T) { + tool := registryTool{ + client: &nonTextClient{content: []Content{ + {Type: "text", Text: "captured the page"}, + {Type: "image", MimeType: "image/png", Data: tinyPNGBase64}, + }}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + } + + result := tool.Run(context.Background(), map[string]any{}) + + if !strings.Contains(result.Output, "captured the page") { + t.Errorf("the text block was lost:\n%s", result.Output) + } + if strings.Contains(result.Output, "cannot forward") { + t.Errorf("a forwarded image is still described as unforwardable:\n%s", result.Output) + } + if strings.Contains(result.Output, "[image forwarded]") { + t.Errorf("text+image result substituted a placeholder over the text:\n%s", result.Output) + } + if len(result.Images) != 1 { + t.Fatalf("Images len = %d, want 1", len(result.Images)) + } + if result.Images[0].MediaType != "image/png" { + t.Errorf("MediaType = %q, want image/png", result.Images[0].MediaType) + } +} + +func TestMalformedImageDataDoesNotPanic(t *testing.T) { + tool := registryTool{ + client: &nonTextClient{content: []Content{ + {Type: "image", MimeType: "image/png", Data: "%%%not-base64%%%"}, + }}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + } + + result := tool.Run(context.Background(), map[string]any{}) + + if len(result.Images) != 0 { + t.Fatalf("malformed image was forwarded: %#v", result.Images) + } + if !strings.Contains(result.Output, "image/png") { + t.Errorf("malformed image was not named in the drop summary:\n%s", result.Output) + } + if strings.Contains(result.Output, "(empty MCP tool result)") { + t.Errorf("malformed image still reported as empty:\n%s", result.Output) + } +} + +func TestImageContentJSONDecodesDataAndStaysCompatibleWithoutIt(t *testing.T) { + var withData CallToolResult + raw := []byte(`{"content":[{"type":"image","mimeType":"image/png","data":"` + tinyPNGBase64 + `"}]}`) + if err := json.Unmarshal(raw, &withData); err != nil { + t.Fatalf("unmarshal image content: %v", err) + } + if len(withData.Content) != 1 { + t.Fatalf("content len = %d, want 1", len(withData.Content)) + } + if withData.Content[0].Type != "image" || withData.Content[0].MimeType != "image/png" { + t.Fatalf("decoded fields = %+v", withData.Content[0]) + } + if withData.Content[0].Data != tinyPNGBase64 { + t.Fatalf("data = %q, want tiny PNG base64", withData.Content[0].Data) + } + + var withoutData CallToolResult + if err := json.Unmarshal([]byte(`{"content":[{"type":"image","mimeType":"image/png"}]}`), &withoutData); err != nil { + t.Fatalf("unmarshal image content without data: %v", err) + } + if withoutData.Content[0].Data != "" { + t.Fatalf("absent data decoded as %q, want empty", withoutData.Content[0].Data) + } +} + +func TestAnExactlyMaxImageBytesPaddedPNGIsForwarded(t *testing.T) { + payload := paddedPNGBase64(imageinput.MaxImageBytes) + if got := base64.StdEncoding.DecodedLen(len(payload)); got <= imageinput.MaxImageBytes { + t.Fatalf("fixture DecodedLen = %d, want > %d so the old bound would reject it", got, imageinput.MaxImageBytes) + } + tool := registryTool{ + client: &nonTextClient{content: []Content{ + {Type: "image", MimeType: "image/png", Data: payload}, + }}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + } + + result := tool.Run(context.Background(), map[string]any{}) + if len(result.Images) != 1 { + t.Fatalf("at-limit padded PNG was dropped: images=%d output=%q", len(result.Images), result.Output) + } + if got := len(result.Images[0].Data); got != imageinput.MaxImageBytes { + t.Fatalf("forwarded size = %d, want %d", got, imageinput.MaxImageBytes) + } +} + +func TestAnOversizedImageIsDroppedAndNamed(t *testing.T) { + tool := registryTool{ + client: &nonTextClient{content: []Content{ + {Type: "image", MimeType: "image/png", Data: paddedPNGBase64(imageinput.MaxImageBytes + 1)}, + }}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + } + + result := tool.Run(context.Background(), map[string]any{}) + + if len(result.Images) != 0 { + t.Fatalf("oversized image was forwarded: %#v", result.Images) + } + if !strings.Contains(result.Output, "image/png") { + t.Errorf("oversized image was not named in the drop summary:\n%s", result.Output) + } + if !strings.Contains(result.Output, "cannot forward yet") { + t.Errorf("individually oversized image should stay unforwardable, not a budget skip:\n%s", result.Output) + } + if strings.Contains(result.Output, "image budget") { + t.Errorf("individually oversized image was described as a budget skip:\n%s", result.Output) + } +} + +func TestAggregateImageBudgetForwardsTheFirstAndNamesTheRest(t *testing.T) { + // Each payload is under the per-image cap; together they exceed the + // aggregate MaxImageBytes budget for one result. Identical bytes are + // deliberate: DroppedContentSummary must name the second even though + // imageBlockFromContent would accept it in isolation. + payload := paddedPNGBase64(imageinput.MaxImageBytes/2 + 1) + tool := registryTool{ + client: &nonTextClient{content: []Content{ + {Type: "image", MimeType: "image/png", Data: payload}, + {Type: "image", MimeType: "image/png", Data: payload}, + }}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + } + + result := tool.Run(context.Background(), map[string]any{}) + + if len(result.Images) != 1 { + t.Fatalf("Images len = %d, want 1 (first fits the aggregate budget)", len(result.Images)) + } + if got := len(result.Images[0].Data); got != imageinput.MaxImageBytes/2+1 { + t.Errorf("forwarded image size = %d, want %d", got, imageinput.MaxImageBytes/2+1) + } + if !strings.Contains(result.Output, "[image returned by tool]") { + t.Errorf("the forwarded first image has no placeholder:\n%s", result.Output) + } + if !strings.Contains(result.Output, "image/png") { + t.Errorf("the dropped second image was not named:\n%s", result.Output) + } + if !strings.Contains(result.Output, "which exceeded this result's remaining image budget") { + t.Errorf("the dropped second image is not described as a budget exceeded:\n%s", result.Output) + } + if !strings.Contains(result.Output, "Retrying with fewer images can recover this payload.") { + t.Errorf("the output does not tell the model a retry can recover the payload:\n%s", result.Output) + } + if strings.Contains(result.Output, "cannot forward yet") { + t.Errorf("a budget-exceeded image is described as unforwardable:\n%s", result.Output) + } + if strings.Contains(result.Output, "Retrying cannot recover this payload.") { + t.Errorf("a budget-exceeded image is described as unrecoverable:\n%s", result.Output) + } +} + +func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) { + orig := decodeImageBase64 + t.Cleanup(func() { decodeImageBase64 = orig }) + var n int + decodeImageBase64 = func(s string) ([]byte, error) { + n++ + return orig(s) + } + + // Two images fill the 10 MiB aggregate exactly, so remaining hits 0 and + // the later candidates must not be decoded at all. + half := imageinput.MaxImageBytes / 2 + payload := paddedPNGBase64(half) + content := []Content{ + {Type: "image", MimeType: "image/png", Data: payload}, + {Type: "image", MimeType: "image/png", Data: payload}, + {Type: "image", MimeType: "image/png", Data: payload}, + {Type: "image", MimeType: "image/png", Data: "not-even-valid-base64!"}, + {Type: "image", MimeType: "image/png", Data: ""}, + {Type: "audio", MimeType: "audio/wav"}, + } + + n = 0 + images := ImageBlocks(content) + if n != 2 { + t.Fatalf("ImageBlocks decoded %d payloads, want 2 (budget fills after two %d-byte images)", n, half) + } + if len(images) != 2 { + t.Fatalf("ImageBlocks len = %d, want 2", len(images)) + } + n = 0 + if got := DroppedContentSummary(content); got != "3 image/png blocks, 1 audio/wav block" { + t.Fatalf("DroppedContentSummary() = %q, want the three skipped images and the audio", got) + } + if n != 2 { + t.Fatalf("DroppedContentSummary decoded %d payloads, want 2 (same single pass as ImageBlocks)", n) + } + + n = 0 + result := registryTool{ + client: &nonTextClient{content: content}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + }.Run(context.Background(), map[string]any{}) + if n != 2 { + t.Fatalf("Run decoded %d payloads, want 2 (one pass; drop note must not decode again)", n) + } + if len(result.Images) != 2 { + t.Fatalf("Images len = %d, want 2", len(result.Images)) + } + if !strings.Contains(result.Output, "[image returned by tool]") { + t.Errorf("forwarded images have no placeholder:\n%s", result.Output) + } + if !strings.Contains(result.Output, "image/png") { + t.Errorf("skipped images were not named:\n%s", result.Output) + } + if !strings.Contains(result.Output, "which were not inspected because the aggregate image budget was reached") { + t.Errorf("uninspected images are not described correctly:\n%s", result.Output) + } + if strings.Contains(result.Output, "Retrying with fewer images can recover this payload.") { + t.Errorf("uninspected images must not make unsupported recovery claims:\n%s", result.Output) + } + if !strings.Contains(result.Output, "audio/wav") { + t.Errorf("audio was not named:\n%s", result.Output) + } + if !strings.Contains(result.Output, "cannot forward yet") { + t.Errorf("audio is not described as unforwardable:\n%s", result.Output) + } + if strings.Contains(result.Output, "(empty MCP tool result)") { + t.Errorf("forwarded images still reported as empty:\n%s", result.Output) + } + + n = 0 + one := []Content{{Type: "image", MimeType: "image/png", Data: payload}} + oneResult := registryTool{ + client: &nonTextClient{content: one}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + }.Run(context.Background(), map[string]any{}) + if n != 1 { + t.Fatalf("one image decoded %d times, want 1", n) + } + if len(oneResult.Images) != 1 { + t.Fatalf("one-image Images len = %d, want 1", len(oneResult.Images)) + } + if oneResult.Output != "[image returned by tool]" { + t.Fatalf("one-image Output = %q, want [image returned by tool]", oneResult.Output) + } + if strings.Contains(oneResult.Output, "cannot forward") { + t.Fatalf("a forwarded image is still described as unforwardable:\n%s", oneResult.Output) + } +} + +func TestImageBudgetNonZeroResidueAllowsSmallerLaterImage(t *testing.T) { + // First image: 8 MiB (fits, 2 MiB left) + // Second image: 3 MiB (exceeds remaining 2 MiB, budgetExceeded) + // Third image: 1 MiB (fits in remaining 2 MiB, forwarded, 1 MiB left) + img8 := paddedPNGBase64(8 * 1024 * 1024) + img3 := paddedPNGBase64(3 * 1024 * 1024) + img1 := paddedPNGBase64(1 * 1024 * 1024) + + content := []Content{ + {Type: "image", MimeType: "image/png", Data: img8}, + {Type: "image", MimeType: "image/png", Data: img3}, + {Type: "image", MimeType: "image/png", Data: img1}, + } + + images, disp, _ := forwardImages(content) + if len(images) != 2 { + t.Fatalf("forwardImages len = %d, want 2 (8 MiB + 1 MiB)", len(images)) + } + if disp[0] != dispForwarded || disp[1] != dispBudgetExceeded || disp[2] != dispForwarded { + t.Fatalf("dispositions = %v, want [forwarded, budgetExceeded, forwarded]", disp) + } + + result := registryTool{ + client: &nonTextClient{content: content}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + }.Run(context.Background(), map[string]any{}) + + if len(result.Images) != 2 { + t.Fatalf("result Images len = %d, want 2", len(result.Images)) + } + if !strings.Contains(result.Output, "exceeded this result's remaining image budget") { + t.Fatalf("expected remaining budget notice in output:\n%s", result.Output) + } + if !strings.Contains(result.Output, "Retrying with fewer images can recover this payload.") { + t.Fatalf("expected retry recovery guidance for validated exceeded image:\n%s", result.Output) + } +} + +func TestImageCountBudgetPreservesTextAndSkipsFurtherDecoding(t *testing.T) { + content := []Content{{Type: "image", Data: "invalid"}} + for range 17 { + content = append(content, Content{Type: "image", MimeType: "image/png", Data: tinyPNGBase64}) + } + content = append(content, Content{Type: "text", Text: "trailing text"}) + previous := decodeImageBase64 + decodes := 0 + decodeImageBase64 = func(s string) ([]byte, error) { + decodes++ + return previous(s) + } + t.Cleanup(func() { decodeImageBase64 = previous }) + result := registryTool{ + client: &nonTextClient{content: content}, + server: Server{Name: "shots"}, remote: RemoteTool{Name: "screenshot"}, + }.Run(context.Background(), map[string]any{}) + if len(result.Images) != 16 || decodes != 17 { + t.Fatalf("forwarded %d images with %d decodes, want 16 images and 17 attempts including invalid input", len(result.Images), decodes) + } + if !strings.Contains(result.Output, "trailing text") || !strings.Contains(result.Output, "1 image/png block, which was not inspected") { + t.Fatalf("missing text or count-budget notice: %s", result.Output) + } +} + +func TestImageInspectionBudgetBoundsRejectedCandidates(t *testing.T) { + // Use the decoder seam to exercise accounting without allocating a large + // payload for every rejected candidate. + previous := decodeImageBase64 + decodes := 0 + large, _ := previous(paddedPNGBase64(8 * 1024 * 1024)) + rejected, _ := previous(paddedPNGBase64(3 * 1024 * 1024)) + decodeImageBase64 = func(string) ([]byte, error) { + decodes++ + if decodes == 1 { + return large, nil + } + return rejected, nil + } + t.Cleanup(func() { decodeImageBase64 = previous }) + content := make([]Content, 40) + for i := range content { + content[i] = Content{Type: "image", MimeType: "image/png", Data: tinyPNGBase64} + } + content = append(content, Content{Type: "text", Text: "trailing text"}) + images, disp, _ := forwardImages(content) + if decodes != 32 || len(images) != 1 { + t.Fatalf("got %d decodes and %d forwarded images, want 32 and 1", decodes, len(images)) + } + for i := 1; i < 40; i++ { + want := dispBudgetExceeded + if i >= 32 { + want = dispUninspected + } + if disp[i] != want { + t.Errorf("disposition[%d] = %v, want %v", i, disp[i], want) + } + } + if disp[40] != dispText { + t.Fatal("inspection limit discarded trailing text") + } +} + +func BenchmarkForwardImagesFourHalfBudget(b *testing.B) { + payload := paddedPNGBase64(imageinput.MaxImageBytes / 2) + content := []Content{ + {Type: "image", MimeType: "image/png", Data: payload}, + {Type: "image", MimeType: "image/png", Data: payload}, + {Type: "image", MimeType: "image/png", Data: payload}, + {Type: "image", MimeType: "image/png", Data: payload}, + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _, _ = forwardImages(content) + } +} + +func TestMCPResultUninspectedImageLimitingReasonNotes(t *testing.T) { + t.Run("forwarded count limit reached", func(t *testing.T) { + content := make([]Content, 20) + for i := range content { + content[i] = Content{Type: "image", MimeType: "image/png", Data: tinyPNGBase64} + } + result := registryTool{ + client: &nonTextClient{content: content}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + }.Run(context.Background(), map[string]any{}) + if len(result.Images) != 16 { + t.Fatalf("expected 16 forwarded images, got %d", len(result.Images)) + } + if !strings.Contains(result.Output, "which were not inspected because the maximum forwarded image count was reached.") { + t.Fatalf("expected maximum forwarded image count message, got:\n%s", result.Output) + } + }) + + t.Run("inspection limit reached", func(t *testing.T) { + previous := decodeImageBase64 + large := make([]byte, 1024) + rejected := make([]byte, imageinput.MaxImageBytes+1) + decodes := 0 + decodeImageBase64 = func(s string) ([]byte, error) { + decodes++ + if decodes == 1 { + return large, nil + } + return rejected, nil + } + t.Cleanup(func() { decodeImageBase64 = previous }) + content := make([]Content, 35) + for i := range content { + content[i] = Content{Type: "image", MimeType: "image/png", Data: tinyPNGBase64} + } + result := registryTool{ + client: &nonTextClient{content: content}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + }.Run(context.Background(), map[string]any{}) + if !strings.Contains(result.Output, "which were not inspected because the maximum image inspection limit was reached.") { + t.Fatalf("expected maximum image inspection limit message, got:\n%s", result.Output) + } + }) +} diff --git a/internal/mcp/registry.go b/internal/mcp/registry.go index d65d6f764..b6970b054 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -331,23 +331,56 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res status = tools.StatusError } output := TextContent(result.Content) - // Say what was thrown away. Without this an image-only result reads as - // "(empty MCP tool result)", the model concludes the call produced nothing - // and retries, and the user never learns an image came back (#823). The note - // is appended only when something was actually dropped, so a text-only - // result is byte-for-byte what it was before. + images, disp, limitReason := forwardImages(result.Content) + // Image blocks with valid data ride Result.Images, the same channel capture + // tools already use. Everything else non-text is still named rather than + // silently dropped, because Zero still has nowhere to put audio, embedded + // resources, or a block whose payload could not be decoded (#823). // - // It says retrying cannot RECOVER the payload rather than that a retry - // returns the same thing. Each retry is a fresh call, so the server may well - // answer differently; what cannot change is that Zero still has nowhere to - // put a non-text block. Claiming the response would be identical would be a - // promise this code is in no position to make. - if dropped := DroppedContentSummary(result.Content); dropped != "" { - note := "[zero] this server also returned " + dropped + ", which Zero cannot forward yet. Retrying cannot recover this payload." - if output == "" { - note = "[zero] this server returned " + dropped + ", which Zero cannot forward yet. Retrying cannot recover this payload." + // Conversion happens once: the drop note is built from the same pass's + // per-item disposition, so a valid image is not decoded again just to + // decide whether it was forwarded. + // + // Image-only success has no text block, so a one-line placeholder keeps + // the tool_result self-describing. ModelOutput and finalizeToolOutcome + // copy Output through verbatim; an empty string would hand the model an + // empty body next to the image. + // + // Notes are appended only when something was actually dropped, so a + // text-only result is byte-for-byte what it was before, and a successfully + // forwarded image is not described as unforwardable. + // + // Unforwardable blocks (audio, resource, failed decode) say retrying + // cannot RECOVER the payload: Zero still has nowhere to put them. + // Budget-skipped images are different — Zero can forward them, and a + // retry with fewer images can recover the payload — so they get their + // own sentence. + if output == "" && len(images) > 0 { + output = "[image returned by tool]" + } + if exceeded := droppedContentNote(result.Content, disp, dispBudgetExceeded); exceeded != "" { + output = appendServerNote(output, exceeded, + ", which exceeded this result's remaining image budget. Retrying with fewer images can recover this payload.") + } + if uninspected := droppedContentNote(result.Content, disp, dispUninspected); uninspected != "" { + verb := ", which were not inspected" + if dispCount(disp, dispUninspected) == 1 { + verb = ", which was not inspected" + } + var cause string + switch limitReason { + case imageLimitForwardCount: + cause = " because the maximum forwarded image count was reached." + case imageLimitInspectionCount: + cause = " because the maximum image inspection limit was reached." + default: + cause = " because the aggregate image budget was reached." } - output = strings.TrimSpace(output + "\n\n" + note) + output = appendServerNote(output, uninspected, verb+cause) + } + if dropped := droppedContentNote(result.Content, disp, dispDropped); dropped != "" { + output = appendServerNote(output, dropped, + ", which Zero cannot forward yet. Retrying cannot recover this payload.") } if output == "" { output = "(empty MCP tool result)" @@ -355,6 +388,7 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res return tools.Result{ Status: status, Output: output, + Images: images, Meta: tool.meta(), } } @@ -399,3 +433,13 @@ func isPersistentlyApproved(store *PermissionStore, server Server, toolName stri }) return err == nil && approved } + +// appendServerNote joins a [zero] note to output, choosing "returned" or +// "also returned" based on whether anything precedes it. +func appendServerNote(output, subject, tail string) string { + verb := "this server also returned " + if output == "" { + verb = "this server returned " + } + return strings.TrimSpace(output + "\n\n[zero] " + verb + subject + tail) +} diff --git a/internal/modelregistry/vision.go b/internal/modelregistry/vision.go index de75d7bd8..722e91574 100644 --- a/internal/modelregistry/vision.go +++ b/internal/modelregistry/vision.go @@ -1,6 +1,13 @@ package modelregistry -import "strings" +import ( + "regexp" + "strings" +) + +// Match explicit multimodal local families, including common Ollama spellings. +// Gemma 3 requires a vision-capable size: 270M and 1B are text-only. +var localVisionFamily = regexp.MustCompile(`^(gemma-?3[-:](4b|12b|27b)|llama-?4|mistral-small-?3\.1|phi-?4-multimodal)($|[-:_.])`) // SupportsVision reports whether the model identified by modelID accepts image // input. A model in the curated catalog is authoritative (its declared @@ -36,6 +43,8 @@ func VisionCapableByName(modelID string) bool { id = id[slash+1:] // drop a "provider/" prefix } switch { + case localVisionFamily.MatchString(id): + return true case strings.Contains(id, "gemini"): return true // every Gemini model is multimodal case strings.Contains(id, "claude-3"), strings.Contains(id, "claude-4"), diff --git a/internal/modelregistry/vision_name_test.go b/internal/modelregistry/vision_name_test.go index 7447506d3..242d1c8c0 100644 --- a/internal/modelregistry/vision_name_test.go +++ b/internal/modelregistry/vision_name_test.go @@ -10,10 +10,15 @@ func TestVisionCapableByName(t *testing.T) { "claude-sonnet-4.5", "claude-3-haiku", "MiniMax-M3", "llava:13b", "qwen2.5-vl-7b", "llama3.2-vision", "pixtral-12b", "moondream", + "google/gemma-3-27b-it", "gemma3:27b", "gemma3:4b", "gemma-3-12b-it", + "llama-4-scout", "llama4:17b", "mistral-small-3.1", "mistral-small3.1:24b", + "phi-4-multimodal", "phi4-multimodal-instruct", } textOnly := []string{ "gpt-oss:120b", "kimi-for-coding", "deepseek-coder", "qwen2.5-coder", "codestral", "llama3.1-8b", "grok-text-only", "mistral-large", + "gemma3:1b", "gemma-3-1b-it", "gemma3:270m", "gemma-3-270m-it", + "gemma2:27b", "mistral-small3:24b", "phi4:14b", // Negated "vision" names must NOT match the bare-"vision" fallback. "my-custom-vision-less-model", "no-vision-model", "grok-vision-less", } diff --git a/internal/providermodeldiscovery/discovery.go b/internal/providermodeldiscovery/discovery.go index d4b7d8dd9..4b421b92b 100644 --- a/internal/providermodeldiscovery/discovery.go +++ b/internal/providermodeldiscovery/discovery.go @@ -375,6 +375,14 @@ type modelsResponseItem struct { AdditionalSpeedTiers []string `json:"additional_speed_tiers"` DefaultServiceTier string `json:"default_service_tier"` SupportedParameters []string `json:"supported_parameters"` + Modalities struct { + Input []string `json:"input"` + Output []string `json:"output"` + } `json:"modalities"` + Architecture struct { + InputModalities []string `json:"input_modalities"` + OutputModalities []string `json:"output_modalities"` + } `json:"architecture"` } type modelsResponse struct { @@ -428,6 +436,14 @@ func parseModelsResponse(body []byte) ([]Model, error) { discoveryContainsFold(item.SupportedParameters, "include_reasoning") efforts := reasoningEfforts(item) serviceTiers := modelServiceTiers(item) + inputModalities := cleanDiscoveryStrings(item.Modalities.Input) + if len(inputModalities) == 0 { + inputModalities = cleanDiscoveryStrings(item.Architecture.InputModalities) + } + outputModalities := cleanDiscoveryStrings(item.Modalities.Output) + if len(outputModalities) == 0 { + outputModalities = cleanDiscoveryStrings(item.Architecture.OutputModalities) + } models = append(models, Model{ ID: id, Description: description, @@ -438,6 +454,8 @@ func parseModelsResponse(body []byte) ([]Model, error) { DefaultReasoningEffort: strings.TrimSpace(item.DefaultReasoning), ServiceTiers: serviceTiers, DefaultServiceTier: normalizeServiceTier(item.DefaultServiceTier), + InputModalities: inputModalities, + OutputModalities: outputModalities, Source: "live", }) } @@ -574,6 +592,12 @@ func mergeLiveModels(provider providercatalog.Descriptor, liveModels []Model, ca if live.DefaultServiceTier != "" { catalog.DefaultServiceTier = live.DefaultServiceTier } + if len(catalog.InputModalities) == 0 && len(live.InputModalities) > 0 { + catalog.InputModalities = append([]string{}, live.InputModalities...) + } + if len(catalog.OutputModalities) == 0 && len(live.OutputModalities) > 0 { + catalog.OutputModalities = append([]string{}, live.OutputModalities...) + } catalog.Source = firstDiscoverySource(catalog.Source, "live") result = append(result, catalog) continue @@ -600,6 +624,20 @@ func mergeLiveModels(provider providercatalog.Descriptor, liveModels []Model, ca return result } +func cleanDiscoveryStrings(values []string) []string { + result := make([]string, 0, len(values)) + seen := map[string]bool{} + for _, value := range values { + value = strings.TrimSpace(value) + if value == "" || seen[value] { + continue + } + seen[value] = true + result = append(result, value) + } + return result +} + func discoveryContainsFold(values []string, want string) bool { for _, value := range values { if strings.EqualFold(strings.TrimSpace(value), want) { diff --git a/internal/providermodeldiscovery/discovery_test.go b/internal/providermodeldiscovery/discovery_test.go index f4b11bafb..587ec91da 100644 --- a/internal/providermodeldiscovery/discovery_test.go +++ b/internal/providermodeldiscovery/discovery_test.go @@ -109,6 +109,49 @@ func TestParseModelsResponseSupportsChatGPTCatalog(t *testing.T) { } } +func TestParseModelsResponseCapturesModalities(t *testing.T) { + models, err := parseModelsResponse([]byte(`{ + "data": [ + { + "id": "openrouter/custom-multimodal", + "architecture": { + "input_modalities": ["text", "image"], + "output_modalities": ["text"] + } + }, + { + "id": "opengateway/flat-multimodal", + "modalities": { + "input": ["text", "image"], + "output": ["text", "image"] + } + }, + { + "id": "text-only", + "architecture": { + "input_modalities": ["text"] + } + } + ] + }`)) + if err != nil { + t.Fatalf("parseModelsResponse: %v", err) + } + byID := map[string]Model{} + for _, m := range models { + byID[m.ID] = m + } + if got := strings.Join(byID["openrouter/custom-multimodal"].InputModalities, ","); got != "text,image" { + t.Fatalf("openrouter input modalities = %q, want text,image", got) + } + if got := strings.Join(byID["opengateway/flat-multimodal"].InputModalities, ","); got != "text,image" { + t.Fatalf("opengateway input modalities = %q, want text,image", got) + } + if got := strings.Join(byID["text-only"].InputModalities, ","); got != "text" { + t.Fatalf("text-only input modalities = %q, want text", got) + } +} + func TestParseModelsResponseNormalizesLegacyFastTier(t *testing.T) { models, err := parseModelsResponse([]byte(`{"data":[{"id":"gpt-test","additional_speed_tiers":["fast","priority"]}]}`)) if err != nil { diff --git a/internal/tui/image_attach.go b/internal/tui/image_attach.go index afc02c05b..73efd7b2f 100644 --- a/internal/tui/image_attach.go +++ b/internal/tui/image_attach.go @@ -14,6 +14,7 @@ import ( "github.com/Gitlawb/zero/internal/imageinput" "github.com/Gitlawb/zero/internal/modelregistry" + "github.com/Gitlawb/zero/internal/providermodeldiscovery" "github.com/Gitlawb/zero/internal/terminalpet" "github.com/Gitlawb/zero/internal/zeroruntime" _ "golang.org/x/image/webp" @@ -102,41 +103,44 @@ func stripMatchingQuotes(s string) (string, bool) { // fetched it) — this carries InputModalities from models.dev, which // includes "image" for vision-capable models // 3. Falls back to the name heuristic for unknown models -func (m model) modelSupportsVisionTUI() bool { - trimmed := strings.TrimSpace(m.modelName) +func (m model) modelSupportsVisionFor(modelID string) bool { + trimmed := strings.TrimSpace(modelID) if trimmed == "" { return false } - // The curated catalog is authoritative only when it knows the model. + // Check the discovered model list for the ACTIVE provider route first. + if models, ok := m.discoveredModelsForActiveRoute(); ok { + if supported, ok := discoveredVisionSupport(models, trimmed); ok { + return supported + } + } + // The curated catalog is authoritative when active-provider discovery is absent or inconclusive. if entry, known := m.modelCatalog.Resolve(trimmed); known { return entry.Supports(modelregistry.ModelCapabilityVision) } - // Check the discovered model list (from models.dev) for InputModalities - // containing "image". This covers custom/ollama/cloud models not in the - // curated catalog — models.dev knows their capabilities. - for _, models := range m.modelPickerLiveByProvider { - for _, dm := range models { - if strings.EqualFold(strings.TrimSpace(dm.ID), trimmed) { - // A provider's authenticated listing may only establish which models - // are available, without repeating modalities. Treat an empty list as - // unknown rather than as an explicit image-input denial, so the - // curated registry/name capability fallback remains available while - // models.dev metadata is temporarily unavailable. - if len(dm.InputModalities) == 0 { - continue - } - for _, modality := range dm.InputModalities { - if strings.EqualFold(strings.TrimSpace(modality), "image") { - return true - } + // Fall back to curated catalog or the name heuristic. + return modelregistry.SupportsVision(m.modelCatalog, trimmed) +} + +func discoveredVisionSupport(models []providermodeldiscovery.Model, modelID string) (bool, bool) { + for _, dm := range models { + if strings.EqualFold(strings.TrimSpace(dm.ID), modelID) { + if len(dm.InputModalities) == 0 { + return false, false + } + for _, modality := range dm.InputModalities { + if strings.EqualFold(strings.TrimSpace(modality), "image") { + return true, true } - return false // found the model in discovered list, no image modality } + return false, true // found model with explicit modalities, no image modality } } - // Fall back to the name heuristic for models not in the catalog or - // discovered list. - return modelregistry.VisionCapableByName(trimmed) + return false, false +} + +func (m model) modelSupportsVisionTUI() bool { + return m.modelSupportsVisionFor(m.modelName) } // attachClipboardImage attaches an image read from the OS clipboard (a diff --git a/internal/tui/image_attach_test.go b/internal/tui/image_attach_test.go index 84ba3e69f..df85a5b1a 100644 --- a/internal/tui/image_attach_test.go +++ b/internal/tui/image_attach_test.go @@ -10,6 +10,10 @@ import ( "strings" "testing" + tea "charm.land/bubbletea/v2" + "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/providercatalog" + "github.com/Gitlawb/zero/internal/providermodeldiscovery" "github.com/Gitlawb/zero/internal/tools" "github.com/Gitlawb/zero/internal/zeroruntime" ) @@ -572,3 +576,356 @@ func TestRetryResendsAttachments(t *testing.T) { t.Fatalf("retried prompt should include the remembered user text, got:\n%s", last.Content) } } + +func TestModelSupportsVision_ActiveProviderPrecedence(t *testing.T) { + m := newModel(t.Context(), Options{ModelName: "gpt-4.1", ProviderName: "openai"}) + // gpt-4.1 is in curated catalog as supporting vision. + // Override active provider to explicitly report it as text-only. + m.modelPickerLiveByProvider = map[string][]providermodeldiscovery.Model{ + "openai": { + { + ID: "gpt-4.1", + InputModalities: []string{"text"}, + }, + }, + "other-provider": { + { + ID: "custom-text-model", + InputModalities: []string{"image", "text"}, + }, + }, + } + // Active provider explicit modalities must win over catalog. + if m.modelSupportsVisionFor("gpt-4.1") { + t.Fatal("expected active provider text-only explicit modality to override catalog vision support") + } + + // Foreign provider discovery record must NOT decide capabilities for active provider. + if m.modelSupportsVisionFor("custom-text-model") { + t.Fatal("expected foreign provider discovery record to be ignored for active provider") + } +} + +func TestRunAgentWithOptions_DiscoverySnapshotRace(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + m := newModel(ctx, Options{ModelName: "gpt-4.1", ProviderName: "openai"}) + m.modelPickerLiveByProvider = map[string][]providermodeldiscovery.Model{ + "openai": {{ID: "gpt-4.1", InputModalities: []string{"text"}}}, + } + + // Schedule the command on the main thread (synchronously snapshotting discovered models) + _ = m.runAgentWithOptions(1, ctx, "hello", nil, tuiAgentRunOptions{}) + + // Concurrently simulate discovery resolution via applyModelPickerModelsDiscovered + done := make(chan struct{}) + go func() { + defer close(done) + localM := m + for i := 0; i < 200; i++ { + localM = localM.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{ + providerID: fmt.Sprintf("provider-%d", i%10), + models: []providermodeldiscovery.Model{ + {ID: "discovered-model", InputModalities: []string{"image"}}, + }, + }) + } + }() + + for i := 0; i < 50; i++ { + cmd := m.runAgentWithOptions(i+2, ctx, "hello", nil, tuiAgentRunOptions{}) + if cmd == nil { + t.Fatal("expected non-nil cmd") + } + } + <-done +} + +func TestModelSupportsVision_CustomEndpointScoping(t *testing.T) { + desc, ok := providercatalog.Get("custom-openai-compatible") + if !ok { + t.Fatal("custom-openai-compatible descriptor missing from catalog") + } + + profileA := config.ProviderProfile{ + Name: "custom-endpoint-a", + CatalogID: desc.ID, + BaseURL: "https://a.example.com/v1", + } + profileB := config.ProviderProfile{ + Name: "custom-endpoint-b", + CatalogID: desc.ID, + BaseURL: "https://b.example.com/v1", + } + + // Discovery returns gpt-4.1 as text-only on endpoint A, but multimodal on endpoint B. + endpointAModels := []providermodeldiscovery.Model{ + {ID: "gpt-4.1", InputModalities: []string{"text"}}, + {ID: "endpoint-a-model", InputModalities: []string{"text"}}, + } + endpointBModels := []providermodeldiscovery.Model{ + {ID: "gpt-4.1", InputModalities: []string{"text", "image"}}, + {ID: "endpoint-b-model", InputModalities: []string{"text", "image"}}, + } + + provider := &fakeProvider{events: []zeroruntime.StreamEvent{{Type: zeroruntime.StreamEventDone}}} + + // Case 1: Endpoint B is active. + m := newModel(t.Context(), Options{ + ProviderName: profileB.Name, + ModelName: "gpt-4.1", + ProviderProfile: profileB, + Provider: provider, + SavedProviders: []config.ProviderProfile{profileA, profileB}, + }) + // Simulate discovery delivery for both endpoints. + m = m.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{ + providerID: desc.ID, + endpointKey: providerEndpointKey(profileA, desc), + models: endpointAModels, + }) + m = m.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{ + providerID: desc.ID, + endpointKey: providerEndpointKey(profileB, desc), + models: endpointBModels, + }) + + // Synchronous gate on active endpoint B must see B's multimodal discovery, + // ignoring endpoint A's text-only result. + if !m.modelSupportsVisionFor("gpt-4.1") { + t.Fatal("expected gpt-4.1 to support vision on endpoint B") + } + if !m.modelSupportsVisionFor("endpoint-b-model") { + t.Fatal("expected endpoint-b-model to support vision on endpoint B") + } + if m.modelSupportsVisionFor("endpoint-a-model") { + t.Fatal("endpoint A model must not be recognized on active endpoint B") + } + + // Assert options.SupportsVision wired in runAgentWithOptions also reports true for endpoint B. + var capturedVision func(string) bool + m.captureRunSupportsVision = func(fn func(string) bool) { + capturedVision = fn + } + _ = execCmd(m.runAgentWithOptions(1, t.Context(), "test", nil, tuiAgentRunOptions{})) + if capturedVision == nil { + t.Fatal("expected captureRunSupportsVision to be called") + } + if !capturedVision("gpt-4.1") { + t.Fatal("expected run agent options.SupportsVision to report true for gpt-4.1 on endpoint B") + } + if !capturedVision("endpoint-b-model") { + t.Fatal("expected run agent options.SupportsVision to report true for endpoint-b-model on endpoint B") + } + if capturedVision("endpoint-a-model") { + t.Fatal("run agent options.SupportsVision must not inherit endpoint A model capabilities") + } + + // Case 2: Switch to Endpoint A. + m.providerProfile = profileA + m.providerName = profileA.Name + + // Synchronous gate on active endpoint A must restrict gpt-4.1 to text-only. + if m.modelSupportsVisionFor("gpt-4.1") { + t.Fatal("expected gpt-4.1 to be text-only on endpoint A") + } + if m.modelSupportsVisionFor("endpoint-b-model") { + t.Fatal("endpoint B model must not be recognized on active endpoint A") + } + + capturedVision = nil + _ = execCmd(m.runAgentWithOptions(2, t.Context(), "test", nil, tuiAgentRunOptions{})) + if capturedVision == nil { + t.Fatal("expected captureRunSupportsVision to be called") + } + if capturedVision("gpt-4.1") { + t.Fatal("expected run agent options.SupportsVision to report false for gpt-4.1 on endpoint A") + } +} + +func TestModelSupportsVision_UnnamedEndpointScoping(t *testing.T) { + desc, ok := providercatalog.Get("custom-openai-compatible") + if !ok { + t.Fatal("custom-openai-compatible descriptor missing from catalog") + } + + profile1 := config.ProviderProfile{ + CatalogID: desc.ID, + BaseURL: "https://host-one.example.com/v1", + } + profile2 := config.ProviderProfile{ + CatalogID: desc.ID, + BaseURL: "https://host-two.example.com/v1", + } + + key1 := providerEndpointKey(profile1, desc) + key2 := providerEndpointKey(profile2, desc) + if key1 == key2 || key1 == "" || key2 == "" { + t.Fatalf("unnamed custom profiles must have distinct non-empty endpoint keys, got %q vs %q", key1, key2) + } + + m := newModel(t.Context(), Options{ + ModelName: "gpt-4.1", + ProviderProfile: profile1, + }) + m.modelPickerLiveByProvider = map[string][]providermodeldiscovery.Model{ + key1: {{ID: "unnamed-model-1", InputModalities: []string{"image"}}}, + key2: {{ID: "unnamed-model-2", InputModalities: []string{"image"}}}, + } + + if !m.modelSupportsVisionFor("unnamed-model-1") { + t.Fatal("active endpoint 1 should support unnamed-model-1") + } + if m.modelSupportsVisionFor("unnamed-model-2") { + t.Fatal("active endpoint 1 should NOT support unnamed-model-2 from endpoint 2") + } +} + +func TestModelPickerDiscoveryCmds_OrderIndependence(t *testing.T) { + desc, ok := providercatalog.Get("custom-openai-compatible") + if !ok { + t.Fatal("custom-openai-compatible descriptor missing from catalog") + } + profileA := config.ProviderProfile{ + Name: "custom-a", + CatalogID: desc.ID, + BaseURL: "https://a.example.com/v1", + } + profileB := config.ProviderProfile{ + Name: "custom-b", + CatalogID: desc.ID, + BaseURL: "https://b.example.com/v1", + } + + for _, order := range [][]config.ProviderProfile{ + {profileA, profileB}, + {profileB, profileA}, + } { + m := newModel(t.Context(), Options{ + SavedProviders: order, + }) + cmd := m.modelPickerDiscoveryCmds() + if cmd == nil { + t.Fatalf("expected discovery commands for order %v", order) + } + msg := cmd() + batch, ok := msg.(tea.BatchMsg) + if !ok { + t.Fatalf("expected tea.BatchMsg, got %T", msg) + } + if len(batch) != 2 { + t.Fatalf("expected 2 discovery commands (one per custom endpoint), got %d", len(batch)) + } + } +} + +func TestModelSupportsVision_DelayedDiscoveryAfterProfileSwitch(t *testing.T) { + desc, ok := providercatalog.Get("custom-openai-compatible") + if !ok { + t.Fatal("custom-openai-compatible descriptor missing from catalog") + } + profileA := config.ProviderProfile{ + Name: "endpoint-a", + CatalogID: desc.ID, + BaseURL: "https://a.example.com/v1", + } + profileB := config.ProviderProfile{ + Name: "endpoint-b", + CatalogID: desc.ID, + BaseURL: "https://b.example.com/v1", + } + + // Start with endpoint A active. + m := newModel(t.Context(), Options{ + ProviderName: profileA.Name, + ModelName: "gpt-4.1", + ProviderProfile: profileA, + SavedProviders: []config.ProviderProfile{profileA, profileB}, + }) + + // User switches active provider to endpoint B before discovery completes. + m.providerProfile = profileB + m.providerName = profileB.Name + + // Now delayed discovery for endpoint A arrives with text-only gpt-4.1. + m = m.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{ + providerID: desc.ID, + endpointKey: providerEndpointKey(profileA, desc), + models: []providermodeldiscovery.Model{ + {ID: "gpt-4.1", InputModalities: []string{"text"}}, + }, + }) + + // Active endpoint B has not discovered yet, so it should fall back to catalog authority + // (where gpt-4.1 supports vision), NOT get contaminated by endpoint A's text-only result. + if !m.modelSupportsVisionFor("gpt-4.1") { + t.Fatal("expected active endpoint B to fall back to catalog vision support, not contaminated by endpoint A") + } + + // Generic descriptor key must NOT have been populated with endpoint A's discovery. + if _, genericFound := m.modelPickerLiveByProvider[desc.ID]; genericFound { + t.Fatalf("generic provider ID %q must not be stored in discovery map when endpointKey is set", desc.ID) + } +} + +func TestRunAgentWithOptions_SnapshotIsolation(t *testing.T) { + desc, ok := providercatalog.Get("custom-openai-compatible") + if !ok { + t.Fatal("custom-openai-compatible descriptor missing from catalog") + } + profileB := config.ProviderProfile{ + Name: "endpoint-b", + CatalogID: desc.ID, + BaseURL: "https://b.example.com/v1", + } + + provider := &fakeProvider{events: []zeroruntime.StreamEvent{{Type: zeroruntime.StreamEventDone}}} + m := newModel(t.Context(), Options{ + ProviderName: profileB.Name, + ModelName: "gpt-4.1", + ProviderProfile: profileB, + Provider: provider, + }) + + // Launch run 1 when endpoint B has NO discovery yet. + var capturedVisionRun1 func(string) bool + m.captureRunSupportsVision = func(fn func(string) bool) { + capturedVisionRun1 = fn + } + _ = execCmd(m.runAgentWithOptions(1, t.Context(), "run1", nil, tuiAgentRunOptions{})) + if capturedVisionRun1 == nil { + t.Fatal("expected captureRunSupportsVision for run 1") + } + // Before discovery: catalog authority confirms gpt-4.1 vision support. + if !capturedVisionRun1("gpt-4.1") { + t.Fatal("run 1 should support vision via catalog fallback") + } + + // Discovery arrives for endpoint B reporting gpt-4.1 as text-only. + m = m.applyModelPickerModelsDiscovered(modelPickerModelsDiscoveredMsg{ + providerID: desc.ID, + endpointKey: providerEndpointKey(profileB, desc), + models: []providermodeldiscovery.Model{ + {ID: "gpt-4.1", InputModalities: []string{"text"}}, + }, + }) + + // Newly launched run 2 must see text-only discovery. + var capturedVisionRun2 func(string) bool + m.captureRunSupportsVision = func(fn func(string) bool) { + capturedVisionRun2 = fn + } + _ = execCmd(m.runAgentWithOptions(2, t.Context(), "run2", nil, tuiAgentRunOptions{})) + if capturedVisionRun2 == nil { + t.Fatal("expected captureRunSupportsVision for run 2") + } + if capturedVisionRun2("gpt-4.1") { + t.Fatal("run 2 must report false for gpt-4.1 due to live text-only discovery") + } + + // CRITICAL: Run 1's captured predicate MUST STILL report true (snapshot isolation!). + if !capturedVisionRun1("gpt-4.1") { + t.Fatal("run 1 must retain its snapshot and still report true for gpt-4.1") + } +} diff --git a/internal/tui/model.go b/internal/tui/model.go index 3c69814df..3c06d6c4d 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -589,6 +589,11 @@ type model struct { // with. Nil in production; used by tests to assert image threading without a // real provider round-trip. captureRunImages func([]zeroruntime.ImageBlock) + + // captureRunSupportsVision, when set, is invoked with the vision-support + // predicate wired into the run's agent options. Nil in production; used by + // tests to assert snapshot isolation and discovery scoping. + captureRunSupportsVision func(func(string) bool) } type agentTextMsg struct { @@ -5481,6 +5486,19 @@ func selfCorrectAutonomyForMode(mode agent.PermissionMode) string { } func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt string, images []zeroruntime.ImageBlock, runOptions tuiAgentRunOptions) tea.Cmd { + var activeDiscovered []providermodeldiscovery.Model + if models, ok := m.discoveredModelsForActiveRoute(); ok { + activeDiscovered = make([]providermodeldiscovery.Model, len(models)) + for i, dm := range models { + copied := dm + if len(dm.InputModalities) > 0 { + copied.InputModalities = append([]string{}, dm.InputModalities...) + } + activeDiscovered[i] = copied + } + } + catalog := m.modelCatalog + return func() tea.Msg { started := m.now() if m.turnTimer != nil { @@ -5596,6 +5614,24 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str func(modelID string) { usageModelID = modelID }, ) } + options.SupportsVision = func(modelID string) bool { + trimmed := strings.TrimSpace(modelID) + if trimmed == "" { + return false + } + if len(activeDiscovered) > 0 { + if supported, ok := discoveredVisionSupport(activeDiscovered, trimmed); ok { + return supported + } + } + if entry, known := catalog.Resolve(trimmed); known { + return entry.Supports(modelregistry.ModelCapabilityVision) + } + return modelregistry.SupportsVision(catalog, trimmed) + } + if m.captureRunSupportsVision != nil { + m.captureRunSupportsVision(options.SupportsVision) + } // Post-edit self-correction is on by default in the TUI but kept FAST: it // runs LSP diagnostics over the changed files only — cheap, change-scoped, diff --git a/internal/tui/picker.go b/internal/tui/picker.go index b18b405b6..7e19c5528 100644 --- a/internal/tui/picker.go +++ b/internal/tui/picker.go @@ -289,10 +289,19 @@ func (m model) savedProviderModelPickerItems(profile config.ProviderProfile, act descriptor, hasDescriptor := m.descriptorForProfile(profile) group := modelPickerProviderGroup(profile, descriptor, hasDescriptor) + endpointKey := providerEndpointKey(profile, descriptor) + var discovered []providermodeldiscovery.Model + if endpointKey != "" { + discovered = m.modelPickerLiveByProvider[endpointKey] + } + if len(discovered) == 0 && hasDescriptor && !genericProviderCatalogID(descriptor.ID) { + discovered = m.modelPickerLiveByProvider[descriptor.ID] + } + var raw []pickerItem switch { - case hasDescriptor && len(m.modelPickerLiveByProvider[descriptor.ID]) > 0: - for _, model := range m.modelPickerLiveByProvider[descriptor.ID] { + case len(discovered) > 0: + for _, model := range discovered { if strings.TrimSpace(model.ID) == "" { continue } @@ -375,17 +384,24 @@ func (m model) openModelPicker() (model, tea.Cmd) { } // modelPickerDiscoveryCmds dispatches a live model-discovery command for each -// usable provider (deduped by catalog descriptor), so /model shows the same real +// usable provider (deduped by endpoint key), so /model shows the same real // models the provider-setup wizard discovers. func (m model) modelPickerDiscoveryCmds() tea.Cmd { cmds := []tea.Cmd{} seen := map[string]bool{} for _, profile := range m.modelPickerProviders() { descriptor, ok := m.descriptorForProfile(profile) - if !ok || seen[descriptor.ID] { + if !ok { continue } - seen[descriptor.ID] = true + endpointKey := providerEndpointKey(profile, descriptor) + if endpointKey == "" { + endpointKey = descriptor.ID + } + if seen[endpointKey] { + continue + } + seen[endpointKey] = true if cmd := m.modelPickerProviderDiscoveryCmd(descriptor, profile); cmd != nil { cmds = append(cmds, cmd) } @@ -417,6 +433,7 @@ func (m model) modelPickerProviderDiscoveryCmd(descriptor providercatalog.Descri } } providerID := descriptor.ID + endpointKey := providerEndpointKey(profile, descriptor) return func() tea.Msg { ctx, cancel := context.WithTimeout(m.ctx, 8*time.Second) defer cancel() @@ -427,7 +444,12 @@ func (m model) modelPickerProviderDiscoveryCmd(descriptor providercatalog.Descri } } models, err := discover(ctx, providerWizardDiscoveryProfile(descriptor, k, authed.BaseURL)) - return modelPickerModelsDiscoveredMsg{providerID: providerID, models: models, err: err} + return modelPickerModelsDiscoveredMsg{ + providerID: providerID, + endpointKey: endpointKey, + models: models, + err: err, + } } } @@ -711,7 +733,15 @@ func modelPickerTitleWord(word string) string { } func (m model) activeProviderDescriptor() (providercatalog.Descriptor, bool) { - return m.descriptorForProfile(m.providerProfile) + if descriptor, ok := m.descriptorForProfile(m.providerProfile); ok { + return descriptor, true + } + if name := strings.TrimSpace(m.providerName); name != "" { + if descriptor, ok := providercatalog.Get(name); ok { + return descriptor, true + } + } + return providercatalog.Descriptor{}, false } func customProviderDescriptorForProfile(profile config.ProviderProfile) (providercatalog.Descriptor, bool) { @@ -782,10 +812,87 @@ func genericProviderCatalogID(id string) bool { return strings.HasPrefix(strings.TrimSpace(id), "custom-") } +// providerEndpointKey derives an endpoint route identity for profile discovery. +// Profiles with an explicit Name use that name so distinct endpoints sharing a +// generic descriptor (e.g. custom-openai-compatible) never collide. +// Unnamed profiles qualify by base URL or descriptor ID. +func providerEndpointKey(profile config.ProviderProfile, descriptor providercatalog.Descriptor) string { + if name := strings.TrimSpace(profile.Name); name != "" { + return name + } + if baseURL := normalizeProviderBaseURL(profile.BaseURL); baseURL != "" { + if descriptor.ID != "" { + return descriptor.ID + "@" + baseURL + } + return baseURL + } + if descriptor.ID != "" { + return descriptor.ID + } + if catID := strings.TrimSpace(profile.CatalogID); catID != "" { + return catID + } + return "" +} + +// activeDiscoveryKeys returns candidate lookup keys for live-discovered models +// belonging to the active provider endpoint route, in order of specificity. +// Generic descriptors (custom-*) are excluded as fallback keys so one custom +// endpoint's discovery cannot determine another endpoint's capabilities. +func (m model) activeDiscoveryKeys() []string { + var keys []string + seen := map[string]bool{} + add := func(k string) { + k = strings.TrimSpace(k) + if k != "" && !seen[k] { + seen[k] = true + keys = append(keys, k) + } + } + + descriptor, hasDescriptor := m.activeProviderDescriptor() + if key := providerEndpointKey(m.providerProfile, descriptor); key != "" { + add(key) + } + if m.providerProfile.Name != "" { + add(m.providerProfile.Name) + } + if m.providerName != "" { + add(m.providerName) + } + if m.providerProfile.CatalogID != "" && !genericProviderCatalogID(m.providerProfile.CatalogID) { + add(m.providerProfile.CatalogID) + } + if hasDescriptor && !genericProviderCatalogID(descriptor.ID) { + add(descriptor.ID) + } + if len(keys) == 0 && len(m.modelPickerLiveByProvider) == 1 { + // When no active provider or profile is configured (e.g. dummy models in unit tests), + // accept the single available provider entry. + for id := range m.modelPickerLiveByProvider { + add(id) + break + } + } + return keys +} + +// discoveredModelsForActiveRoute returns the live-discovered models for the active +// endpoint route, if available in modelPickerLiveByProvider. +func (m model) discoveredModelsForActiveRoute() ([]providermodeldiscovery.Model, bool) { + for _, key := range m.activeDiscoveryKeys() { + if models, ok := m.modelPickerLiveByProvider[key]; ok && len(models) > 0 { + return models, true + } + } + return nil, false +} + type modelPickerModelsDiscoveredMsg struct { - providerID string - models []providermodeldiscovery.Model - err error + providerID string + endpointKey string + models []providermodeldiscovery.Model + err error } // ollamaContextWindowDiscoveredMsg carries the result of an async /api/show @@ -876,10 +983,18 @@ func (m model) applyModelPickerModelsDiscovered(msg modelPickerModelsDiscoveredM if msg.err != nil || len(msg.models) == 0 { return m } - if m.modelPickerLiveByProvider == nil { - m.modelPickerLiveByProvider = map[string][]providermodeldiscovery.Model{} + newMap := make(map[string][]providermodeldiscovery.Model, len(m.modelPickerLiveByProvider)+2) + for k, v := range m.modelPickerLiveByProvider { + newMap[k] = v + } + modelsCopy := append([]providermodeldiscovery.Model{}, msg.models...) + if msg.endpointKey != "" { + newMap[msg.endpointKey] = modelsCopy + } + if msg.providerID != "" && (!genericProviderCatalogID(msg.providerID) || msg.endpointKey == "") { + newMap[msg.providerID] = modelsCopy } - m.modelPickerLiveByProvider[msg.providerID] = append([]providermodeldiscovery.Model{}, msg.models...) + m.modelPickerLiveByProvider = newMap // Rebuild the open picker so this provider's section shows its live models, // preserving the current query + selection. if m.picker != nil && m.picker.kind == pickerModel { From cca50749a32e22a21f7234b4757c6fbb9f5e27fd Mon Sep 17 00:00:00 2001 From: Marlowe <321669285+cairn-intern@users.noreply.github.com> Date: Thu, 24 Sep 2026 07:04:44 -0400 Subject: [PATCH 2/2] test(mcp): pin the pre-decode encoded cap and MIME sniffing guards Adds the two reviewer-requested tests for the untrusted-input boundary: - TestOverCapEncodedPayloadIsRejectedBeforeDecoding: a payload one byte over EncodedLen(MaxImageBytes) is rejected and the decoder is never called, so a hostile server cannot force arbitrary decode work. - TestDeclaredMimeTypeDisagreeingWithSniffedBytesIsDropped: an HTML payload declared as image/png is sniffed from its bytes and dropped rather than forwarded to the provider as an image. Both mutations (removing the cap, trusting the declared mimeType) were verified to turn their test red locally. --- internal/mcp/non_text_content_test.go | 62 +++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/internal/mcp/non_text_content_test.go b/internal/mcp/non_text_content_test.go index 7608d630e..3383c4ba5 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -668,3 +668,65 @@ func TestMCPResultUninspectedImageLimitingReasonNotes(t *testing.T) { } }) } + +// The pre-decode encoded-length cap is the decode-DoS guard: it keeps a +// hostile server from making Zero decode a payload of any size (up to the +// inspection cap times per result) before the post-decode backstop drops it. +// Removing the cap changes nothing that is forwarded, so no existing test +// notices; this pins the guard on the decode hook instead. A payload one +// byte over the encoded cap must be rejected without the decoder ever being +// called. +func TestOverCapEncodedPayloadIsRejectedBeforeDecoding(t *testing.T) { + previous := decodeImageBase64 + decodes := 0 + decodeImageBase64 = func(s string) ([]byte, error) { + decodes++ + return previous(s) + } + t.Cleanup(func() { decodeImageBase64 = previous }) + + // An at-limit PNG would decode to exactly MaxImageBytes and forward; a + // single extra base64 character pushes it one byte over the encoded cap, + // so only the pre-decode guard can be responsible for rejecting it. + payload := paddedPNGBase64(imageinput.MaxImageBytes) + "A" + content := []Content{{Type: "image", MimeType: "image/png", Data: payload}} + + if images := ImageBlocks(content); len(images) != 0 { + t.Fatalf("ImageBlocks forwarded %d images from an over-cap payload, want 0", len(images)) + } + if decodes != 0 { + t.Fatalf("decoder ran %d times, want 0: the pre-decode cap must reject before any decode work", decodes) + } + if got := DroppedContentSummary(content); got != "1 image/png block" { + t.Fatalf("DroppedContentSummary() = %q, want the over-cap payload named as dropped", got) + } +} + +// The media type is sniffed from the decoded bytes, never taken from the +// block's declared mimeType: a hostile server must not be able to relabel +// HTML or a script as image/png and have it forwarded to the provider as an +// image. Every other fixture pairs real PNG bytes with image/png, so +// sniffing and trusting the declared type agree there and no test can tell +// them apart; this fixture makes them disagree. +func TestDeclaredMimeTypeDisagreeingWithSniffedBytesIsDropped(t *testing.T) { + html := []byte("") + tool := registryTool{ + client: &nonTextClient{content: []Content{ + {Type: "image", MimeType: "image/png", Data: base64.StdEncoding.EncodeToString(html)}, + }}, + server: Server{Name: "shots"}, + remote: RemoteTool{Name: "screenshot"}, + } + + result := tool.Run(context.Background(), map[string]any{}) + + if len(result.Images) != 0 { + t.Fatalf("forwarded %d images from an HTML payload declared as image/png, want 0", len(result.Images)) + } + if strings.Contains(result.Output, "[image returned by tool]") { + t.Errorf("a non-image payload declared as image/png is still announced as a returned image:\n%s", result.Output) + } + if !strings.Contains(result.Output, "image/png") { + t.Errorf("the dropped block is not named in the drop summary:\n%s", result.Output) + } +}