diff --git a/internal/acp/agent.go b/internal/acp/agent.go index 6050c7c8b..047d4b01e 100644 --- a/internal/acp/agent.go +++ b/internal/acp/agent.go @@ -11,6 +11,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" @@ -249,6 +250,13 @@ func (a *Agent) runTurn(ctx context.Context, sess *acpSession, userText string, } note := ¬ifier{conn: a.conn, sessionID: sess.id} + supportsVision := func(modelID string) bool { + return a.modelSupportsVision(ctx, resolved.Provider, modelID) + } + if len(images) > 0 && !supportsVision(resolved.Provider.Model) { + images = nil + } + opts := agent.Options{ Cwd: sess.cwd, SessionID: sess.id, @@ -259,6 +267,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 +776,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 + } + reg, _ := modelregistry.DefaultRegistry() + if entry, known := reg.Resolve(trimmed); known { + return entry.Supports(modelregistry.ModelCapabilityVision) + } + 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 + } + } + } + } + return modelregistry.SupportsVision(reg, trimmed) +} diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index 4fa97a258..abcfb0267 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -497,6 +497,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) { 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..8c5df1b53 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) @@ -136,7 +137,7 @@ func TestRunKeepsToolResultsContiguousWhenAToolReturnsAnImage(t *testing.T) { {{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}}, }} - result, err := Run(context.Background(), "two calls", provider, Options{Registry: registry, MaxTurns: 2}) + result, err := Run(context.Background(), "two calls", provider, Options{Registry: registry, MaxTurns: 2, Model: "gpt-4o"}) if err != nil { t.Fatalf("Run: %v", err) } @@ -205,3 +206,211 @@ func messageShape(messages []zeroruntime.Message) string { } return "[" + strings.Join(parts, " ") + "]" } + +func TestRunDeliversToolResultImagesToAVisionModel(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(imageTool{media: "image/png", data: []byte("\x89PNG\r\n\x1a\nfake")}) + provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{ + { + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "capture"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"}, + {Type: zeroruntime.StreamEventDone}, + }, + {{Type: zeroruntime.StreamEventText, Content: "I can see it."}, {Type: zeroruntime.StreamEventDone}}, + }} + + result, err := Run(context.Background(), "screenshot please", provider, Options{ + Registry: registry, + MaxTurns: 2, + Model: "gpt-4o", + }) + if err != nil { + t.Fatalf("Run: %v", err) + } + var carrier *zeroruntime.Message + var toolText string + for index := range result.Messages { + if result.Messages[index].Role == zeroruntime.MessageRoleTool { + toolText = result.Messages[index].Content + } + if len(result.Messages[index].Images) > 0 { + carrier = &result.Messages[index] + } + } + if toolText != "Captured a screenshot." { + t.Fatalf("tool text = %q, want preserved", toolText) + } + if carrier == nil { + t.Fatal("vision model must receive the tool image") + } +} + +func TestRunDropsToolResultImagesForANonVisionModel(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(imageTool{media: "image/png", data: []byte("\x89PNG\r\n\x1a\nfake")}) + provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{ + { + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "capture"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"}, + {Type: zeroruntime.StreamEventDone}, + }, + {{Type: zeroruntime.StreamEventText, Content: "ok"}, {Type: zeroruntime.StreamEventDone}}, + }} + + result, err := Run(context.Background(), "screenshot please", provider, Options{ + Registry: registry, + MaxTurns: 2, + Model: "totally-made-up-model", + }) + if err != nil { + t.Fatalf("Run: %v", err) + } + var toolText string + var noticed bool + for _, message := range result.Messages { + if message.Role == zeroruntime.MessageRoleTool { + toolText = message.Content + } + if len(message.Images) > 0 { + t.Fatalf("non-vision model received image bytes on %q: %s", message.Role, messageShape(result.Messages)) + } + if strings.Contains(message.Content, "does not support image input") { + noticed = true + } + } + if toolText != "Captured a screenshot." { + t.Fatalf("tool text = %q, want preserved on the non-vision path", toolText) + } + if !noticed { + t.Fatalf("non-vision model was not told the image was dropped; recorded %s", messageShape(result.Messages)) + } +} + +type switchAndCaptureTool struct { + targetModel string +} + +func (switchAndCaptureTool) Name() string { return "switch_and_capture" } +func (switchAndCaptureTool) Description() string { return "Switches model and captures image" } +func (switchAndCaptureTool) Parameters() tools.Schema { + return tools.Schema{Type: "object", Properties: map[string]tools.PropertySchema{}} +} +func (switchAndCaptureTool) Safety() tools.Safety { + return tools.Safety{Permission: tools.PermissionAllow} +} +func (t switchAndCaptureTool) Run(context.Context, map[string]any) tools.Result { + return tools.Result{ + Status: tools.StatusOK, + Output: "[image returned by tool]", + Images: []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte("\x89PNG\r\n\x1a\nfake")}}, + Meta: map[string]string{"escalate_to_model": t.targetModel}, + } +} + +func TestRunToolImagesRespectsModelSwitch(t *testing.T) { + t.Run("Non-vision to vision model escalation forwards images", func(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(switchAndCaptureTool{targetModel: "gpt-4o"}) + + provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{ + { + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "switch_and_capture"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"}, + {Type: zeroruntime.StreamEventDone}, + }, + { + {Type: zeroruntime.StreamEventText, Content: "I see it on gpt-4o."}, + {Type: zeroruntime.StreamEventDone}, + }, + }} + + switched := false + switcher := func(_ context.Context, modelID string) (Provider, error) { + switched = true + return provider, nil + } + + result, err := Run(context.Background(), "test", provider, Options{ + Registry: registry, + MaxTurns: 2, + Model: "non-vision-initial", + ModelSwitcher: switcher, + SupportsVision: func(modelID string) bool { + return modelID == "gpt-4o" + }, + }) + if err != nil { + t.Fatalf("Run: %v", err) + } + if !switched { + t.Fatal("expected model switch to occur") + } + + var carrier *zeroruntime.Message + for index := range result.Messages { + if len(result.Messages[index].Images) > 0 { + carrier = &result.Messages[index] + } + if strings.Contains(result.Messages[index].Content, "does not support image input") { + t.Fatalf("images were unexpectedly dropped after switch to vision model: %v", result.Messages[index]) + } + } + if carrier == nil { + t.Fatal("expected image carrier message after escalation to vision model") + } + }) + + t.Run("Vision to non-vision model switch drops images with notice", func(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(switchAndCaptureTool{targetModel: "text-only-dest"}) + + provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{ + { + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_1", ToolName: "switch_and_capture"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_1", ArgumentsFragment: `{}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_1"}, + {Type: zeroruntime.StreamEventDone}, + }, + { + {Type: zeroruntime.StreamEventText, Content: "ok."}, + {Type: zeroruntime.StreamEventDone}, + }, + }} + + switcher := func(_ context.Context, modelID string) (Provider, error) { + return provider, nil + } + + result, err := Run(context.Background(), "test", provider, Options{ + Registry: registry, + MaxTurns: 2, + Model: "gpt-4o", + ModelSwitcher: switcher, + SupportsVision: func(modelID string) bool { + return modelID == "gpt-4o" + }, + }) + if err != nil { + t.Fatalf("Run: %v", err) + } + + var noticed bool + for _, m := range result.Messages { + if len(m.Images) > 0 { + t.Fatalf("image delivered to non-vision destination model: %v", m) + } + if strings.Contains(m.Content, "does not support image input") { + noticed = true + } + if strings.Contains(m.Content, "[image forwarded]") { + t.Fatalf("contradictory forwarded text found in message: %v", m) + } + } + if !noticed { + t.Fatal("expected drop notice when switching to non-vision model") + } + }) +} 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..841fb0444 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,98 @@ 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, 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 only once remaining is zero; 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 + +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 { + 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 +632,46 @@ 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 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..7b4cb2864 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,64 @@ 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) + } + }) + + 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) + } + }) + + 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..c591b2a89 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,372 @@ 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 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..08fb03fcc 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -327,18 +327,52 @@ 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 != "" { + // 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 != "" { + note := "[zero] this server also returned " + exceeded + ", which exceeded this result's remaining image budget. Retrying with fewer images can recover this payload." + if output == "" { + note = "[zero] this server returned " + exceeded + ", which exceeded this result's remaining image budget. Retrying with fewer images can recover this payload." + } + output = strings.TrimSpace(output + "\n\n" + note) + } + if uninspected := droppedContentNote(result.Content, disp, dispUninspected); uninspected != "" { + verb := "which was not inspected" + if !strings.HasPrefix(uninspected, "1 ") { + verb = "which were not inspected" + } + note := "[zero] this server also returned " + uninspected + ", " + verb + " because the aggregate image budget was reached." + if output == "" { + note = "[zero] this server returned " + uninspected + ", " + verb + " because the aggregate image budget was reached." + } + output = strings.TrimSpace(output + "\n\n" + note) + } + if dropped := droppedContentNote(result.Content, disp, dispDropped); 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." @@ -351,6 +385,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(), } } diff --git a/internal/tui/image_attach.go b/internal/tui/image_attach.go index afc02c05b..ccabb37b7 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,8 +103,8 @@ 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 } @@ -111,32 +112,42 @@ func (m model) modelSupportsVisionTUI() bool { 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. + // Check the discovered model list, preferring the ACTIVE provider's models. + if descriptor, ok := m.activeProviderDescriptor(); ok && descriptor.ID != "" { + if models, ok := m.modelPickerLiveByProvider[descriptor.ID]; ok { + if supported, ok := discoveredVisionSupport(models, trimmed); ok { + return supported + } + } + } 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 - } + if supported, ok := discoveredVisionSupport(models, trimmed); ok { + return supported + } + } + // 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/model.go b/internal/tui/model.go index 9473de06a..6f023da50 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -5496,6 +5496,45 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str options.ContextWindowFor = func(modelID string) int { return modelregistry.AgentContextWindow(m.modelContextWindow(modelID)) } + var activeDescriptorID string + if descriptor, ok := m.activeProviderDescriptor(); ok { + activeDescriptorID = descriptor.ID + } + 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 + options.SupportsVision = func(modelID string) bool { + trimmed := strings.TrimSpace(modelID) + if trimmed == "" { + return false + } + if entry, known := catalog.Resolve(trimmed); known { + return entry.Supports(modelregistry.ModelCapabilityVision) + } + if activeDescriptorID != "" { + if models, ok := discoveredSnapshot[activeDescriptorID]; ok { + if supported, ok := discoveredVisionSupport(models, trimmed); ok { + return supported + } + } + } + for _, models := range discoveredSnapshot { + if supported, ok := discoveredVisionSupport(models, trimmed); ok { + return supported + } + } + 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,