diff --git a/internal/acp/agent.go b/internal/acp/agent.go index 6050c7c8b..0d9128ca2 100644 --- a/internal/acp/agent.go +++ b/internal/acp/agent.go @@ -5,12 +5,14 @@ import ( "encoding/base64" "encoding/json" "errors" + "fmt" "log" "strings" "sync" "github.com/Gitlawb/zero/internal/agent" "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/modelregistry" "github.com/Gitlawb/zero/internal/providercatalog" "github.com/Gitlawb/zero/internal/providermodelcatalog" "github.com/Gitlawb/zero/internal/providermodeldiscovery" @@ -249,6 +251,29 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string, } note := ¬ifier{conn: a.conn, sessionID: sess.id} + visionCache := make(map[string]bool) + var visionCacheMu sync.Mutex + supportsVision := func(modelID string) bool { + modelID = strings.TrimSpace(modelID) + if modelID == "" { + return false + } + visionCacheMu.Lock() + defer visionCacheMu.Unlock() + if cached, ok := visionCache[modelID]; ok { + return cached + } + supported := a.modelSupportsVision(ctx, resolved.Provider, modelID) + visionCache[modelID] = supported + return supported + } + if len(images) > 0 && !supportsVision(resolved.Provider.Model) { + msg := fmt.Sprintf("Model %s does not support image input; ignoring %d prompt image(s).", resolved.Provider.Model, len(images)) + note.text("[zero] " + msg + "\n\n") + userText = fmt.Sprintf("[Note: %s]\n\n%s", msg, userText) + images = nil + } + opts := agent.Options{ Cwd: sess.cwd, SessionID: sess.id, @@ -259,6 +284,7 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string, PermissionMode: sess.currentMode(), MaxTurns: resolved.MaxTurns, Images: images, + SupportsVision: supportsVision, OnText: note.text, OnReasoning: note.thought, OnToolCall: note.toolCall, @@ -767,3 +793,32 @@ func (s *acpSession) snapshotHistory() []turnRecord { defer s.mu.Unlock() return append([]turnRecord(nil), s.history...) } + +func (a *Agent) modelSupportsVision(ctx context.Context, profile config.ProviderProfile, modelID string) bool { + trimmed := strings.TrimSpace(modelID) + if trimmed == "" { + return false + } + if a.deps.DiscoverModels != nil { + if discovered, err := a.deps.DiscoverModels(ctx, profile); err == nil { + for _, dm := range discovered { + if strings.EqualFold(strings.TrimSpace(dm.ID), trimmed) { + if len(dm.InputModalities) > 0 { + for _, mod := range dm.InputModalities { + if strings.EqualFold(strings.TrimSpace(mod), "image") { + return true + } + } + return false + } + break + } + } + } + } + reg, _ := modelregistry.DefaultRegistry() + if entry, known := reg.Resolve(trimmed); known { + return entry.Supports(modelregistry.ModelCapabilityVision) + } + return modelregistry.SupportsVision(reg, trimmed) +} diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index 4fa97a258..5a3bd824a 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" @@ -9,6 +11,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "time" @@ -497,6 +500,48 @@ func TestACPRunTurnWiresSandboxAndScopedRegistry(t *testing.T) { } } +func TestACPWiresSupportsVision(t *testing.T) { + deps := testDeps(t) + var captured agent.Options + deps.RunAgent = func(_ context.Context, _ string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) { + captured = opts + return agent.Result{FinalAnswer: "done"}, nil + } + deps.DiscoverModels = func(_ context.Context, _ config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + return []providermodeldiscovery.Model{ + {ID: "custom-vision", InputModalities: []string{"text", "image"}}, + {ID: "custom-text", InputModalities: []string{"text"}}, + }, nil + } + h := newHarness(t, deps) + defer h.stop() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var newRes NewSessionResult + if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil { + t.Fatalf("session/new: %v", err) + } + var promptRes PromptResult + if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{ + SessionID: newRes.SessionID, + Prompt: []ContentBlock{ + {Type: "text", Text: "hello"}, + }, + }, &promptRes); err != nil { + t.Fatalf("session/prompt: %v", err) + } + if captured.SupportsVision == nil { + t.Fatal("SupportsVision was not wired into agent.Options") + } + if !captured.SupportsVision("custom-vision") { + t.Fatal("SupportsVision(custom-vision) = false, want true") + } + if captured.SupportsVision("custom-text") { + t.Fatal("SupportsVision(custom-text) = true, want false") + } +} + // TestACPRejectsInvalidCwd confirms session/new fails when the workspace root // resolver rejects the client cwd (e.g. filesystem root). func TestACPRejectsInvalidCwd(t *testing.T) { @@ -599,3 +644,200 @@ func drainTextUntil(t *testing.T, ch <-chan string, done func(string) bool) stri } } } + +func TestACPPromptUnsupportedModelDropsImagesAndNotifies(t *testing.T) { + deps := testDeps(t) + deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + return []providermodeldiscovery.Model{ + {ID: "fake-model", InputModalities: []string{"text"}}, + }, nil + } + var capturedOpts agent.Options + var capturedPrompt string + deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) { + capturedPrompt = prompt + capturedOpts = opts + return agent.Result{FinalAnswer: "Answered without images"}, nil + } + + h := newHarness(t, deps) + defer h.stop() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var initRes InitializeResult + if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil { + t.Fatalf("initialize: %v", err) + } + + var newRes NewSessionResult + if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil { + t.Fatalf("session/new: %v", err) + } + + rawPng := "\x89PNG\r\n\x1a\nfake-image-bytes" + b64Png := base64.StdEncoding.EncodeToString([]byte(rawPng)) + var promptRes PromptResult + if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{ + SessionID: newRes.SessionID, + Prompt: []ContentBlock{ + TextBlock("look at this image"), + ImageBlock(b64Png, "image/png"), + }, + }, &promptRes); err != nil { + t.Fatalf("session/prompt: %v", err) + } + if promptRes.StopReason != StopEndTurn { + t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn) + } + + got := drainTextUntil(t, h.updates, func(text string) bool { + return strings.Contains(text, "does not support image input") + }) + if !strings.Contains(got, "Model fake-model does not support image input; ignoring 1 prompt image(s).") { + t.Fatalf("streamed text = %q, want drop notice", got) + } + if len(capturedOpts.Images) != 0 { + t.Fatalf("captured %d images, want 0 (images should be withheld)", len(capturedOpts.Images)) + } + if !strings.Contains(capturedPrompt, "[Note: Model fake-model does not support image input; ignoring 1 prompt image(s).]") { + t.Fatalf("captured prompt = %q, want prompt note", capturedPrompt) + } +} + +func TestACPSessionPrompt_SupportsVisionMemoizedPerRun(t *testing.T) { + deps := testDeps(t) + var discoverCount int + var discoverMu sync.Mutex + deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + discoverMu.Lock() + discoverCount++ + discoverMu.Unlock() + return []providermodeldiscovery.Model{ + {ID: "fake-model", InputModalities: []string{"text", "image"}}, + }, nil + } + deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) { + for i := 0; i < 5; i++ { + if !opts.SupportsVision(opts.Model) { + t.Error("expected SupportsVision to report true") + } + } + return agent.Result{FinalAnswer: "Done"}, nil + } + + h := newHarness(t, deps) + defer h.stop() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var initRes InitializeResult + if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil { + t.Fatalf("initialize: %v", err) + } + + var newRes NewSessionResult + if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil { + t.Fatalf("session/new: %v", err) + } + + discoverMu.Lock() + discoverCount = 0 + discoverMu.Unlock() + + rawPng := "\x89PNG\r\n\x1a\nfake-image-bytes" + b64Png := base64.StdEncoding.EncodeToString([]byte(rawPng)) + var promptRes PromptResult + if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{ + SessionID: newRes.SessionID, + Prompt: []ContentBlock{ + TextBlock("look at this"), + ImageBlock(b64Png, "image/png"), + }, + }, &promptRes); err != nil { + t.Fatalf("session/prompt: %v", err) + } + if promptRes.StopReason != StopEndTurn { + t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn) + } + + discoverMu.Lock() + count := discoverCount + discoverMu.Unlock() + if count != 1 { + t.Fatalf("DiscoverModels was called %d times, want exactly 1 call (memoized per run)", count) + } +} + +func TestACPEndToEndImagePromptSupportedModel(t *testing.T) { + deps := testDeps(t) + deps.DiscoverModels = func(ctx context.Context, p config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + return []providermodeldiscovery.Model{ + {ID: "discovered-vision-model", InputModalities: []string{"text", "image"}}, + }, nil + } + deps.ResolveConfig = func(_ string, o config.Overrides) (config.ResolvedConfig, error) { + model := "discovered-vision-model" + if o.Provider.Model != "" { + model = o.Provider.Model + } + return config.ResolvedConfig{ + Provider: config.ProviderProfile{Name: "fake", Model: model}, + MaxTurns: 4, + }, nil + } + var capturedOpts agent.Options + var capturedPrompt string + deps.RunAgent = func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) { + capturedPrompt = prompt + capturedOpts = opts + return agent.Result{FinalAnswer: "I can see the image!"}, nil + } + + h := newHarness(t, deps) + defer h.stop() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var initRes InitializeResult + if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil { + t.Fatalf("initialize: %v", err) + } + + var newRes NewSessionResult + if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil { + t.Fatalf("session/new: %v", err) + } + + rawPng := []byte("\x89PNG\r\n\x1a\nreal-png-bytes") + b64Png := base64.StdEncoding.EncodeToString(rawPng) + var promptRes PromptResult + if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{ + SessionID: newRes.SessionID, + Prompt: []ContentBlock{ + TextBlock("analyze this image"), + ImageBlock(b64Png, "image/png"), + }, + }, &promptRes); err != nil { + t.Fatalf("session/prompt: %v", err) + } + if promptRes.StopReason != StopEndTurn { + t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn) + } + + if len(capturedOpts.Images) != 1 { + t.Fatalf("captured %d images, want 1", len(capturedOpts.Images)) + } + if capturedOpts.Images[0].MediaType != "image/png" { + t.Fatalf("mediaType = %q, want image/png", capturedOpts.Images[0].MediaType) + } + if !bytes.Equal(capturedOpts.Images[0].Data, rawPng) { + t.Fatalf("image bytes mismatch: got %v, want %v", capturedOpts.Images[0].Data, rawPng) + } + if strings.Contains(capturedPrompt, "does not support image input") { + t.Fatalf("prompt unexpectedly contains refusal note: %q", capturedPrompt) + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index fe691ac4c..0a081105d 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 @@ -3457,8 +3469,10 @@ func copyMessages(messages []Message) []Message { // // The text names the tool so the model can tell which call an image came from // when several ran in one turn — the images arrive detached from their tool -// result, so nothing else associates them. -func toolResultImageMessage(result ToolResult) (zeroruntime.Message, bool) { +// result, so nothing else associates them. If the effective model cannot +// accept images, the tool's text result is left unchanged and this follow-up +// is a notice with no image parts, matching the CLI/TUI vision gate. +func toolResultImageMessage(result ToolResult, options Options) (zeroruntime.Message, bool) { images := make([]zeroruntime.ImageBlock, 0, len(result.Images)) for _, image := range result.Images { if len(image.Data) == 0 { @@ -3479,9 +3493,26 @@ func toolResultImageMessage(result ToolResult) (zeroruntime.Message, bool) { if label == "" { label = "tool" } + if !modelAcceptsToolImages(options) { + return zeroruntime.Message{ + Role: zeroruntime.MessageRoleUser, + Content: "Image output from " + label + " was not sent because the current model does not support image input.", + }, true + } return zeroruntime.Message{ Role: zeroruntime.MessageRoleUser, Content: "Image output from " + label + ":", Images: images, }, true } + +func modelAcceptsToolImages(options Options) bool { + if options.SupportsVision != nil { + return options.SupportsVision(options.Model) + } + registry, err := modelregistry.DefaultRegistry() + if err != nil { + return modelregistry.VisionCapableByName(options.Model) + } + return modelregistry.SupportsVision(registry, options.Model) +} diff --git a/internal/agent/tool_result_images_test.go b/internal/agent/tool_result_images_test.go index f6a5d0335..62ead0dc5 100644 --- a/internal/agent/tool_result_images_test.go +++ b/internal/agent/tool_result_images_test.go @@ -62,6 +62,7 @@ func TestRunDeliversToolResultImagesToTheModel(t *testing.T) { result, err := Run(context.Background(), "screenshot please", provider, Options{ Registry: registry, MaxTurns: 2, + Model: "gpt-4o", }) if err != nil { t.Fatalf("Run: %v", err) @@ -86,11 +87,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 +142,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 +211,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 511ea7140..f74138781 100644 --- a/internal/agent/types.go +++ b/internal/agent/types.go @@ -340,6 +340,11 @@ type Options struct { // nil for text-only runs (the seeded message then carries no images, exactly // as before). Images []zeroruntime.ImageBlock + // SupportsVision, when set, reports whether the effective model accepts + // image input. Tool-produced images are dropped at the shared delivery + // boundary when this is false. nil uses the curated catalog plus the + // name heuristic via modelregistry.SupportsVision. + SupportsVision func(modelID string) bool // ContextWindow is the model's maximum input token budget. When > 0 the agent // loop compacts long conversations once the estimated size crosses a fraction // of this window. 0 DISABLES compaction entirely (every existing caller/test diff --git a/internal/cli/exec.go b/internal/cli/exec.go index 63d22f8bf..1689f2435 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -686,6 +686,9 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in ContextWindowFor: func(modelID string) int { return modelregistry.AgentContextWindow(modelContextWindow(modelRegistry, modelID)) }, + SupportsVision: func(modelID string) bool { + return modelregistry.SupportsVision(modelRegistry, modelID) + }, ReasoningEffort: forwardEffort, Trace: traceRecorder, Cwd: workspaceRoot, 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 064e7f213..e81a13297 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -3,10 +3,12 @@ package mcp import ( "bytes" "context" + "encoding/base64" "encoding/json" "errors" "fmt" "io" + "net/http" "os" "os/exec" "strconv" @@ -15,6 +17,8 @@ import ( "time" "github.com/Gitlawb/zero/internal/execution" + "github.com/Gitlawb/zero/internal/imageinput" + "github.com/Gitlawb/zero/internal/zeroruntime" ) type RemoteTool struct { @@ -26,10 +30,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 { @@ -500,24 +510,103 @@ 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. +// 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. // -// 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. +// Image payloads ride Result.Images. Audio, embedded resources, structured +// content, and image blocks whose data cannot be decoded still have nowhere +// to go. Images skipped because they would exceed the aggregate byte budget +// are also named so the model hears that a valid screenshot was dropped. +// +// Only images actually kept by ImageBlocks are omitted from the note. An +// image that would decode in isolation but was skipped by the aggregate cap +// is still named, otherwise the model would not hear that a valid screenshot +// was dropped. // // 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 { + _, disp := forwardImages(content) + return droppedContentNote(content, disp, dispDropped, dispBudgetExceeded, dispUninspected) +} + +// ImageBlocks converts MCP image content into the same ImageBlock channel +// capture tools already use. Blocks that cannot be decoded, exceed +// imageinput.MaxImageBytes individually, sniff to a type outside the provider +// allow-list, or would push the result over an aggregate +// imageinput.MaxImageBytes budget or maxForwardedImages count, are left for +// DroppedContentSummary to name. +// +// The aggregate cap is the same 10 MiB as the per-image cap: a server that +// returns many individually valid images must not retain all of them in +// Result.Images. Once the next valid image would exceed the remaining +// budget it is skipped; a later smaller image may still fit. Later image +// payloads are not decoded once remaining is zero or the count cap is reached. +// A leftover residue still fully decodes the next candidate before the length +// check rejects it. +func ImageBlocks(content []Content) []zeroruntime.ImageBlock { + images, _ := forwardImages(content) + return images +} + +// 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 + +func forwardImages(content []Content) ([]zeroruntime.ImageBlock, []itemDisp) { + disp := make([]itemDisp, len(content)) + var images []zeroruntime.ImageBlock + remaining := imageinput.MaxImageBytes + for i, item := range content { + if item.Type == "text" { + disp[i] = dispText + continue + } + if item.Type == "image" { + if remaining == 0 || len(images) >= maxForwardedImages { + disp[i] = dispUninspected + continue + } + 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 +} + +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". @@ -548,3 +637,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/network_client.go b/internal/mcp/network_client.go index b422b3c28..ed920723d 100644 --- a/internal/mcp/network_client.go +++ b/internal/mcp/network_client.go @@ -651,11 +651,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 fba92f1ec..005222080 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" @@ -339,3 +340,80 @@ func TestDecodeSSERPCMessageSkipsNotifications(t *testing.T) { t.Fatalf("expected a result payload, got %#v", msg) } } + +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..a116cda02 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,397 @@ 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 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) + } +} diff --git a/internal/mcp/registry.go b/internal/mcp/registry.go index d1a2978dc..62b71e7e0 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -327,23 +327,48 @@ 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 := 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" } - output = strings.TrimSpace(output + "\n\n" + note) + output = appendServerNote(output, uninspected, + verb+" because the aggregate image budget was reached.") + } + 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)" @@ -351,6 +376,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(), } } @@ -395,3 +421,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..727b4a0b5 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,55 @@ 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 first. + activeID := "" + if descriptor, ok := m.activeProviderDescriptor(); ok && descriptor.ID != "" { + activeID = descriptor.ID + } else if len(m.modelPickerLiveByProvider) == 1 { + for id := range m.modelPickerLiveByProvider { + activeID = id + break + } + } + if activeID != "" { + if models, ok := m.modelPickerLiveByProvider[activeID]; 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..c4b2820c1 100644 --- a/internal/tui/image_attach_test.go +++ b/internal/tui/image_attach_test.go @@ -10,6 +10,7 @@ import ( "strings" "testing" + "github.com/Gitlawb/zero/internal/providermodeldiscovery" "github.com/Gitlawb/zero/internal/tools" "github.com/Gitlawb/zero/internal/zeroruntime" ) @@ -572,3 +573,68 @@ 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 +} diff --git a/internal/tui/model.go b/internal/tui/model.go index 9473de06a..cfcb260e0 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -5405,6 +5405,29 @@ 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 activeDescriptorID string + if descriptor, ok := m.activeProviderDescriptor(); ok { + activeDescriptorID = descriptor.ID + } else if len(m.modelPickerLiveByProvider) == 1 { + for id := range m.modelPickerLiveByProvider { + activeDescriptorID = id + break + } + } + discoveredSnapshot := make(map[string][]providermodeldiscovery.Model, len(m.modelPickerLiveByProvider)) + for pID, list := range m.modelPickerLiveByProvider { + copiedList := make([]providermodeldiscovery.Model, len(list)) + for i, dm := range list { + copied := dm + if len(dm.InputModalities) > 0 { + copied.InputModalities = append([]string{}, dm.InputModalities...) + } + copiedList[i] = copied + } + discoveredSnapshot[pID] = copiedList + } + catalog := m.modelCatalog + return func() tea.Msg { started := m.now() if m.turnTimer != nil { @@ -5496,6 +5519,23 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str options.ContextWindowFor = func(modelID string) int { return modelregistry.AgentContextWindow(m.modelContextWindow(modelID)) } + options.SupportsVision = func(modelID string) bool { + trimmed := strings.TrimSpace(modelID) + if trimmed == "" { + return false + } + if activeDescriptorID != "" { + if models, ok := discoveredSnapshot[activeDescriptorID]; ok { + if supported, ok := discoveredVisionSupport(models, trimmed); ok { + return supported + } + } + } + if entry, known := catalog.Resolve(trimmed); known { + return entry.Supports(modelregistry.ModelCapabilityVision) + } + return modelregistry.SupportsVision(catalog, trimmed) + } // 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 563dffe62..10f6200d6 100644 --- a/internal/tui/picker.go +++ b/internal/tui/picker.go @@ -708,7 +708,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) { @@ -873,10 +881,12 @@ 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)+1) + for k, v := range m.modelPickerLiveByProvider { + newMap[k] = v } - m.modelPickerLiveByProvider[msg.providerID] = append([]providermodeldiscovery.Model{}, msg.models...) + newMap[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 {