From a7617c588b73bab14acc96b807e866ab3c5a55e6 Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Fri, 28 Aug 2026 01:53:37 +0000 Subject: [PATCH 1/9] fix(mcp): forward image tool results on the existing image channel Part 1 of #823 named dropped non-text blocks. Screenshot servers still could not hand the model the picture. Decode MCP image blocks onto tools.Result.Images so the agent loop can emit them, and only name the block types still not forwarded. Fixes Gitlawb/zero#823 --- internal/mcp/client.go | 79 ++++++++++-- internal/mcp/non_text_content_test.go | 165 ++++++++++++++++++++++++-- internal/mcp/registry.go | 21 ++-- 3 files changed, 238 insertions(+), 27 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 064e7f213..60ce16bef 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,16 +510,15 @@ 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: naming them costs nothing and stops the model from treating a successful +// call as empty and retrying. // // Counts are grouped by mime type and ordered by first appearance, so the same // result always produces the same sentence. @@ -520,6 +529,9 @@ func DroppedContentSummary(content []Content) string { if item.Type == "text" { continue } + if _, ok := imageBlockFromContent(item); ok { + continue + } // Prefer the mime type: "image/png" tells the reader more than "image". // A server may omit it, so fall back to the block type rather than // printing an empty label. @@ -548,3 +560,46 @@ func DroppedContentSummary(content []Content) string { } return strings.Join(parts, ", ") } + +// ImageBlocks converts MCP image content into the same ImageBlock channel +// capture tools already use. Blocks that cannot be decoded, exceed +// imageinput.MaxImageBytes, or sniff to a type outside the provider +// allow-list are left for DroppedContentSummary to name. +func ImageBlocks(content []Content) []zeroruntime.ImageBlock { + var images []zeroruntime.ImageBlock + for _, item := range content { + if image, ok := imageBlockFromContent(item); ok { + images = append(images, image) + } + } + return images +} + +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 + } + if base64.StdEncoding.DecodedLen(len(raw)) > imageinput.MaxImageBytes { + return zeroruntime.ImageBlock{}, false + } + data, err := base64.StdEncoding.DecodeString(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/non_text_content_test.go b/internal/mcp/non_text_content_test.go index a1082962f..ac5f1482b 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -2,20 +2,24 @@ package mcp import ( "context" + "encoding/json" "strings" "testing" "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==" + +// 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 +138,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 +178,132 @@ 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 strings.Contains(result.Output, "(empty MCP tool result)") { + t.Fatalf("image payload still reported as empty:\n%s", 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 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) + } +} diff --git a/internal/mcp/registry.go b/internal/mcp/registry.go index d1a2978dc..c1abb1be8 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -327,17 +327,21 @@ 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 := ImageBlocks(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). + // + // The note is 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. // // 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. + // put the remaining 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 == "" { @@ -345,12 +349,13 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res } output = strings.TrimSpace(output + "\n\n" + note) } - if output == "" { + if output == "" && len(images) == 0 { output = "(empty MCP tool result)" } return tools.Result{ Status: status, Output: output, + Images: images, Meta: tool.meta(), } } From 571e4832cd47e2b4fd964ebf49d6b1145d4af0ea Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Fri, 28 Aug 2026 02:18:24 +0000 Subject: [PATCH 2/9] fix(mcp): bound total forwarded image bytes in one MCP result ImageBlocks applied MaxImageBytes per block only, so many 10 MiB images could exhaust memory. Cap the sum at MaxImageBytes and skip the next valid image once it would exceed the remaining budget. DroppedContentSummary now omits only images ImageBlocks actually kept, so aggregate-skipped payloads are named rather than silently dropped. --- internal/mcp/client.go | 40 +++++++++++++++--- internal/mcp/non_text_content_test.go | 61 +++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 7 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 60ce16bef..ec036fc46 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -516,21 +516,34 @@ func TextContent(content []Content) string { // forwarded, so a caller adds nothing to the ordinary case. // // Image payloads ride Result.Images. Audio, embedded resources, structured -// content, and image blocks whose data cannot be decoded still have nowhere to +// content, image blocks whose data cannot be decoded, and images skipped +// because they would exceed the aggregate byte budget still have nowhere to // go: naming them costs nothing and stops the model from treating a successful // call as empty and retrying. // +// 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 { + forwarded := ImageBlocks(content) + next := 0 labels := make([]string, 0, len(content)) counts := make(map[string]int, len(content)) for _, item := range content { if item.Type == "text" { continue } - if _, ok := imageBlockFromContent(item); ok { - continue + if image, ok := imageBlockFromContent(item); ok { + if next < len(forwarded) && + image.MediaType == forwarded[next].MediaType && + bytes.Equal(image.Data, forwarded[next].Data) { + next++ + continue + } } // Prefer the mime type: "image/png" tells the reader more than "image". // A server may omit it, so fall back to the block type rather than @@ -563,14 +576,27 @@ func DroppedContentSummary(content []Content) string { // ImageBlocks converts MCP image content into the same ImageBlock channel // capture tools already use. Blocks that cannot be decoded, exceed -// imageinput.MaxImageBytes, or sniff to a type outside the provider -// allow-list are left for DroppedContentSummary to name. +// 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. func ImageBlocks(content []Content) []zeroruntime.ImageBlock { var images []zeroruntime.ImageBlock + remaining := imageinput.MaxImageBytes for _, item := range content { - if image, ok := imageBlockFromContent(item); ok { - images = append(images, image) + image, ok := imageBlockFromContent(item) + if !ok { + continue + } + if len(image.Data) > remaining { + continue } + images = append(images, image) + remaining -= len(image.Data) } return images } diff --git a/internal/mcp/non_text_content_test.go b/internal/mcp/non_text_content_test.go index ac5f1482b..b52628907 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -2,16 +2,27 @@ package mcp import ( "context" + "encoding/base64" "encoding/json" "strings" "testing" + "github.com/Gitlawb/zero/internal/imageinput" "github.com/Gitlawb/zero/internal/tools" ) // 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). @@ -307,3 +318,53 @@ func TestImageContentJSONDecodesDataAndStaysCompatibleWithoutIt(t *testing.T) { t.Fatalf("absent data decoded as %q, want empty", withoutData.Content[0].Data) } } + +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) + } +} + +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/png") { + t.Errorf("the dropped second image was not named:\n%s", result.Output) + } + if !strings.Contains(result.Output, "cannot forward") { + t.Errorf("the dropped second image is not described as unforwardable:\n%s", result.Output) + } +} From 06387982fb8fa49e0d4543af8cac884d65f80aa8 Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Fri, 28 Aug 2026 06:40:04 +0000 Subject: [PATCH 3/9] fix(mcp): decode each MCP image payload at most once registryTool.Run called ImageBlocks then DroppedContentSummary, which decoded every accepted image two more times and kept decoding after the aggregate budget was spent. Classify content in one pass, build the drop note from that disposition, and skip later image payloads once no budget remains. --- internal/mcp/client.go | 99 +++++++++++++++++---------- internal/mcp/non_text_content_test.go | 92 +++++++++++++++++++++++++ internal/mcp/registry.go | 8 ++- 3 files changed, 159 insertions(+), 40 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index ec036fc46..0b59365a0 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -529,22 +529,72 @@ func TextContent(content []Content) string { // 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 { - forwarded := ImageBlocks(content) - next := 0 - labels := make([]string, 0, len(content)) - counts := make(map[string]int, len(content)) - for _, item := range content { + _, disp := forwardImages(content) + return droppedContentNote(content, disp) +} + +// 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. Once no +// budget remains, later image payloads are not decoded — the cap is a +// work limit, not only a retention limit. +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 +) + +// 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 image, ok := imageBlockFromContent(item); ok { - if next < len(forwarded) && - image.MediaType == forwarded[next].MediaType && - bytes.Equal(image.Data, forwarded[next].Data) { - next++ + if item.Type == "image" && remaining > 0 { + if image, ok := imageBlockFromContent(item); ok && len(image.Data) <= remaining { + images = append(images, image) + remaining -= len(image.Data) + disp[i] = dispForwarded continue } } + disp[i] = dispDropped + } + return images, disp +} + +func droppedContentNote(content []Content, disp []itemDisp) string { + labels := make([]string, 0, len(content)) + counts := make(map[string]int, len(content)) + for i, item := range content { + if i >= len(disp) || disp[i] != dispDropped { + continue + } // Prefer the mime type: "image/png" tells the reader more than "image". // A server may omit it, so fall back to the block type rather than // printing an empty label. @@ -574,33 +624,6 @@ func DroppedContentSummary(content []Content) string { return strings.Join(parts, ", ") } -// 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. -func ImageBlocks(content []Content) []zeroruntime.ImageBlock { - var images []zeroruntime.ImageBlock - remaining := imageinput.MaxImageBytes - for _, item := range content { - image, ok := imageBlockFromContent(item) - if !ok { - continue - } - if len(image.Data) > remaining { - continue - } - images = append(images, image) - remaining -= len(image.Data) - } - return images -} - func imageBlockFromContent(item Content) (zeroruntime.ImageBlock, bool) { if item.Type != "image" { return zeroruntime.ImageBlock{}, false @@ -612,7 +635,7 @@ func imageBlockFromContent(item Content) (zeroruntime.ImageBlock, bool) { if base64.StdEncoding.DecodedLen(len(raw)) > imageinput.MaxImageBytes { return zeroruntime.ImageBlock{}, false } - data, err := base64.StdEncoding.DecodeString(raw) + data, err := decodeImageBase64(raw) if err != nil { return zeroruntime.ImageBlock{}, false } diff --git a/internal/mcp/non_text_content_test.go b/internal/mcp/non_text_content_test.go index b52628907..057eaa01d 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -368,3 +368,95 @@ func TestAggregateImageBudgetForwardsTheFirstAndNamesTheRest(t *testing.T) { t.Errorf("the dropped second image is not described as unforwardable:\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: payload}, + {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 != "2 image/png blocks, 1 audio/wav block" { + t.Fatalf("DroppedContentSummary() = %q, want the two 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/png") { + t.Errorf("skipped images were not named:\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, "(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 strings.Contains(oneResult.Output, "cannot forward") { + t.Fatalf("a forwarded image is still described as unforwardable:\n%s", oneResult.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 c1abb1be8..2d1cbdc74 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -327,12 +327,16 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res status = tools.StatusError } output := TextContent(result.Content) - images := ImageBlocks(result.Content) + 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). // + // 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. + // // The note is 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. @@ -342,7 +346,7 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res // answer differently; what cannot change is that Zero still has nowhere to // put the remaining 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 != "" { + if dropped := droppedContentNote(result.Content, disp); 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." From 3b496d19753c094e5a3de573f8152f768a75d09b Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Fri, 28 Aug 2026 18:21:04 +0000 Subject: [PATCH 4/9] fix(mcp): placeholder for image-only results and distinct budget-drop note An image-only tool result left Output empty, so the model got a blank tool_result next to the image. Budget-skipped images reused the unrecoverable drop sentence even though a retry with fewer images would recover them. --- internal/mcp/client.go | 47 ++++++++++++++++++--------- internal/mcp/non_text_content_test.go | 44 ++++++++++++++++++++++--- internal/mcp/registry.go | 31 +++++++++++++----- 3 files changed, 95 insertions(+), 27 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 0b59365a0..03ca64a90 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -516,10 +516,9 @@ func TextContent(content []Content) string { // forwarded, so a caller adds nothing to the ordinary case. // // Image payloads ride Result.Images. Audio, embedded resources, structured -// content, image blocks whose data cannot be decoded, and images skipped -// because they would exceed the aggregate byte budget still have nowhere to -// go: naming them costs nothing and stops the model from treating a successful -// call as empty and retrying. +// 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 @@ -530,7 +529,7 @@ func TextContent(content []Content) string { // result always produces the same sentence. func DroppedContentSummary(content []Content) string { _, disp := forwardImages(content) - return droppedContentNote(content, disp) + return droppedContentNote(content, disp, dispDropped, dispBudgetSkipped) } // ImageBlocks converts MCP image content into the same ImageBlock channel @@ -542,9 +541,9 @@ func DroppedContentSummary(content []Content) string { // 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. Once no -// budget remains, later image payloads are not decoded — the cap is a -// work limit, not only a retention limit. +// 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 @@ -560,6 +559,7 @@ const ( dispText itemDisp = iota dispForwarded dispDropped + dispBudgetSkipped ) // decodeImageBase64 is the MCP image payload decoder. Tests replace it to @@ -575,11 +575,19 @@ func forwardImages(content []Content) ([]zeroruntime.ImageBlock, []itemDisp) { disp[i] = dispText continue } - if item.Type == "image" && remaining > 0 { - if image, ok := imageBlockFromContent(item); ok && len(image.Data) <= remaining { - images = append(images, image) - remaining -= len(image.Data) - disp[i] = dispForwarded + if item.Type == "image" { + if remaining == 0 { + disp[i] = dispBudgetSkipped + 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] = dispBudgetSkipped continue } } @@ -588,11 +596,11 @@ func forwardImages(content []Content) ([]zeroruntime.ImageBlock, []itemDisp) { return images, disp } -func droppedContentNote(content []Content, disp []itemDisp) string { +func droppedContentNote(content []Content, disp []itemDisp, kinds ...itemDisp) string { labels := make([]string, 0, len(content)) counts := make(map[string]int, len(content)) for i, item := range content { - if i >= len(disp) || disp[i] != dispDropped { + if i >= len(disp) || !dispKind(disp[i], kinds) { continue } // Prefer the mime type: "image/png" tells the reader more than "image". @@ -624,6 +632,15 @@ func droppedContentNote(content []Content, disp []itemDisp) 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 diff --git a/internal/mcp/non_text_content_test.go b/internal/mcp/non_text_content_test.go index 057eaa01d..a3a1fe65c 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -201,8 +201,8 @@ func TestAnImageWithPayloadIsForwarded(t *testing.T) { result := tool.Run(context.Background(), map[string]any{}) - if strings.Contains(result.Output, "(empty MCP tool result)") { - t.Fatalf("image payload still reported as empty:\n%s", result.Output) + if result.Output != "[image forwarded]" { + t.Fatalf("image-only Output = %q, want [image forwarded] 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) @@ -264,6 +264,9 @@ func TestTextAndImageKeepsTextAndForwardsImage(t *testing.T) { 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)) } @@ -336,6 +339,12 @@ func TestAnOversizedImageIsDroppedAndNamed(t *testing.T) { 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) { @@ -361,11 +370,23 @@ func TestAggregateImageBudgetForwardsTheFirstAndNamesTheRest(t *testing.T) { 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 forwarded]") { + 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, "cannot forward") { - t.Errorf("the dropped second image is not described as unforwardable:\n%s", result.Output) + if !strings.Contains(result.Output, "which exceeded this result's image budget") { + t.Errorf("the dropped second image is not described as a budget skip:\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-skipped image is described as unforwardable:\n%s", result.Output) + } + if strings.Contains(result.Output, "Retrying cannot recover this payload.") { + t.Errorf("a budget-skipped image is described as unrecoverable:\n%s", result.Output) } } @@ -418,12 +439,24 @@ func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) { if len(result.Images) != 2 { t.Fatalf("Images len = %d, want 2", len(result.Images)) } + if !strings.Contains(result.Output, "[image forwarded]") { + 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 exceeded this result's image budget") { + t.Errorf("budget-skipped images are not described as a budget drop:\n%s", result.Output) + } + if !strings.Contains(result.Output, "Retrying with fewer images can recover this payload.") { + t.Errorf("budget-skipped images do not say a retry can recover them:\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) } @@ -441,6 +474,9 @@ func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) { if len(oneResult.Images) != 1 { t.Fatalf("one-image Images len = %d, want 1", len(oneResult.Images)) } + if oneResult.Output != "[image forwarded]" { + t.Fatalf("one-image Output = %q, want [image forwarded]", oneResult.Output) + } if strings.Contains(oneResult.Output, "cannot forward") { t.Fatalf("a forwarded image is still described as unforwardable:\n%s", oneResult.Output) } diff --git a/internal/mcp/registry.go b/internal/mcp/registry.go index 2d1cbdc74..a015e8e2a 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -337,23 +337,38 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res // per-item disposition, so a valid image is not decoded again just to // decide whether it was forwarded. // - // The note is appended only when something was actually dropped, so a + // 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. // - // 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 the remaining non-text block. Claiming the response would be identical - // would be a promise this code is in no position to make. - if dropped := droppedContentNote(result.Content, disp); dropped != "" { + // 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 forwarded]" + } + if skipped := droppedContentNote(result.Content, disp, dispBudgetSkipped); skipped != "" { + note := "[zero] this server also returned " + skipped + ", which exceeded this result's image budget. Retrying with fewer images can recover this payload." + if output == "" { + note = "[zero] this server returned " + skipped + ", which exceeded this result's image budget. Retrying with fewer images can recover this payload." + } + 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." } output = strings.TrimSpace(output + "\n\n" + note) } - if output == "" && len(images) == 0 { + if output == "" { output = "(empty MCP tool result)" } return tools.Result{ From 896b7ac63c8ae0e9eb9e8b4b7308735a7422a23f Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Tue, 1 Sep 2026 03:40:45 -0400 Subject: [PATCH 5/9] fix(mcp): gate tool images on vision support and accept at-limit padding Tool-produced images were always attached to the next user message, so a text-only model could have its following completion rejected. Drop those attachments at the shared delivery boundary when the effective model cannot accept images, and keep a notice without changing the tool text. The MCP pre-decode size check used DecodedLen, which reports cap+2 for a standard-base64 PNG of exactly MaxImageBytes. Bound on EncodedLen instead so an at-limit padded image still forwards. --- internal/agent/loop.go | 26 ++++++- internal/agent/tool_result_images_test.go | 84 ++++++++++++++++++++++- internal/agent/types.go | 5 ++ internal/mcp/client.go | 7 +- internal/mcp/non_text_content_test.go | 22 ++++++ 5 files changed, 139 insertions(+), 5 deletions(-) diff --git a/internal/agent/loop.go b/internal/agent/loop.go index fe691ac4c..49e9d90f5 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" @@ -712,7 +713,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // — 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 { + if imageMessage, ok := toolResultImageMessage(toolResult, options); ok { toolImageMessages = append(toolImageMessages, imageMessage) } @@ -3457,8 +3458,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 +3482,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..008a5585d 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,84 @@ 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)) + } +} 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/mcp/client.go b/internal/mcp/client.go index 03ca64a90..053e00a21 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -649,7 +649,12 @@ func imageBlockFromContent(item Content) (zeroruntime.ImageBlock, bool) { if raw == "" { return zeroruntime.ImageBlock{}, false } - if base64.StdEncoding.DecodedLen(len(raw)) > imageinput.MaxImageBytes { + // 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) diff --git a/internal/mcp/non_text_content_test.go b/internal/mcp/non_text_content_test.go index a3a1fe65c..feecb1337 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -322,6 +322,28 @@ func TestImageContentJSONDecodesDataAndStaysCompatibleWithoutIt(t *testing.T) { } } +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{ From 49163360eb1c609b772f8d48ab8a758e0428ab0b Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Tue, 1 Sep 2026 16:13:39 -0400 Subject: [PATCH 6/9] fix(agent): evaluate tool image vision gate after turn model switch --- internal/agent/loop.go | 49 ++++++++++++++++++++++++++---------------- 1 file changed, 30 insertions(+), 19 deletions(-) diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 49e9d90f5..0a081105d 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -647,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 @@ -706,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, options); 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 @@ -726,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) @@ -756,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 @@ -775,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 @@ -866,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 From dd8843c0eeafa60ea7a00a6df98885cd69e3d738 Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Tue, 1 Sep 2026 18:55:26 -0400 Subject: [PATCH 7/9] fix(mcp): align SSE payload limits, separate uninspected budget blocks, and route vision capability --- internal/agent/tool_result_images_test.go | 127 ++++++++++++++++++++++ internal/cli/exec.go | 3 + internal/cli/exec_spec.go | 15 ++- internal/mcp/client.go | 9 +- internal/mcp/network_client.go | 9 +- internal/mcp/network_client_test.go | 55 ++++++++++ internal/mcp/non_text_content_test.go | 74 ++++++++++--- internal/mcp/registry.go | 15 ++- internal/tui/image_attach.go | 13 ++- internal/tui/model.go | 3 + 10 files changed, 282 insertions(+), 41 deletions(-) diff --git a/internal/agent/tool_result_images_test.go b/internal/agent/tool_result_images_test.go index 008a5585d..8c5df1b53 100644 --- a/internal/agent/tool_result_images_test.go +++ b/internal/agent/tool_result_images_test.go @@ -287,3 +287,130 @@ func TestRunDropsToolResultImagesForANonVisionModel(t *testing.T) { 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/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 053e00a21..841fb0444 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -529,7 +529,7 @@ func TextContent(content []Content) string { // result always produces the same sentence. func DroppedContentSummary(content []Content) string { _, disp := forwardImages(content) - return droppedContentNote(content, disp, dispDropped, dispBudgetSkipped) + return droppedContentNote(content, disp, dispDropped, dispBudgetExceeded, dispUninspected) } // ImageBlocks converts MCP image content into the same ImageBlock channel @@ -559,7 +559,8 @@ const ( dispText itemDisp = iota dispForwarded dispDropped - dispBudgetSkipped + dispBudgetExceeded + dispUninspected ) // decodeImageBase64 is the MCP image payload decoder. Tests replace it to @@ -577,7 +578,7 @@ func forwardImages(content []Content) ([]zeroruntime.ImageBlock, []itemDisp) { } if item.Type == "image" { if remaining == 0 { - disp[i] = dispBudgetSkipped + disp[i] = dispUninspected continue } if image, ok := imageBlockFromContent(item); ok { @@ -587,7 +588,7 @@ func forwardImages(content []Content) ([]zeroruntime.ImageBlock, []itemDisp) { disp[i] = dispForwarded continue } - disp[i] = dispBudgetSkipped + disp[i] = dispBudgetExceeded continue } } diff --git a/internal/mcp/network_client.go b/internal/mcp/network_client.go index b422b3c28..12c218a99 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 an accepted +// 10 MiB image payload in base64 (~13.98 MiB) plus JSON-RPC envelope, metadata, +// and SSE framing overhead. +const maxSSEEventBytes = 16 * 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..2f1cb55e2 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,57 @@ 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 16 MiB rejected cleanly", func(t *testing.T) { + huge := strings.Repeat("A", 17*1024*1024) + stream := "event: message\ndata: " + huge + "\n\n" + _, err := decodeSSERPCMessage(strings.NewReader(stream)) + if err == nil { + t.Fatal("expected error for 17 MiB event, got nil") + } + }) +} diff --git a/internal/mcp/non_text_content_test.go b/internal/mcp/non_text_content_test.go index feecb1337..d7bc68fc8 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -201,8 +201,8 @@ func TestAnImageWithPayloadIsForwarded(t *testing.T) { result := tool.Run(context.Background(), map[string]any{}) - if result.Output != "[image forwarded]" { - t.Fatalf("image-only Output = %q, want [image forwarded] so the tool_result is not an empty body", result.Output) + 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) @@ -392,23 +392,23 @@ func TestAggregateImageBudgetForwardsTheFirstAndNamesTheRest(t *testing.T) { 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 forwarded]") { + 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 image budget") { - t.Errorf("the dropped second image is not described as a budget skip:\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-skipped image is described as unforwardable:\n%s", result.Output) + 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-skipped image is described as unrecoverable:\n%s", result.Output) + t.Errorf("a budget-exceeded image is described as unrecoverable:\n%s", result.Output) } } @@ -429,7 +429,8 @@ func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) { {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}, + {Type: "image", MimeType: "image/png", Data: "not-even-valid-base64!"}, + {Type: "image", MimeType: "image/png", Data: ""}, {Type: "audio", MimeType: "audio/wav"}, } @@ -442,8 +443,8 @@ func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) { t.Fatalf("ImageBlocks len = %d, want 2", len(images)) } n = 0 - if got := DroppedContentSummary(content); got != "2 image/png blocks, 1 audio/wav block" { - t.Fatalf("DroppedContentSummary() = %q, want the two skipped images and the audio", got) + 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) @@ -461,17 +462,17 @@ func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) { if len(result.Images) != 2 { t.Fatalf("Images len = %d, want 2", len(result.Images)) } - if !strings.Contains(result.Output, "[image forwarded]") { + 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 exceeded this result's image budget") { - t.Errorf("budget-skipped images are not described as a budget drop:\n%s", result.Output) + if !strings.Contains(result.Output, "which was 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("budget-skipped images do not say a retry can recover them:\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) @@ -496,14 +497,53 @@ func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) { if len(oneResult.Images) != 1 { t.Fatalf("one-image Images len = %d, want 1", len(oneResult.Images)) } - if oneResult.Output != "[image forwarded]" { - t.Fatalf("one-image Output = %q, want [image forwarded]", oneResult.Output) + 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{ diff --git a/internal/mcp/registry.go b/internal/mcp/registry.go index a015e8e2a..1c295dd08 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -352,12 +352,19 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res // retry with fewer images can recover the payload — so they get their // own sentence. if output == "" && len(images) > 0 { - output = "[image forwarded]" + output = "[image returned by tool]" } - if skipped := droppedContentNote(result.Content, disp, dispBudgetSkipped); skipped != "" { - note := "[zero] this server also returned " + skipped + ", which exceeded this result's image budget. Retrying with fewer images can recover this payload." + 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 " + skipped + ", which exceeded this result's image budget. Retrying with fewer images can recover this payload." + 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 != "" { + note := "[zero] this server also returned " + uninspected + ", which was not inspected because the aggregate image budget was reached." + if output == "" { + note = "[zero] this server returned " + uninspected + ", which was not inspected because the aggregate image budget was reached." } output = strings.TrimSpace(output + "\n\n" + note) } diff --git a/internal/tui/image_attach.go b/internal/tui/image_attach.go index afc02c05b..2492ca7ff 100644 --- a/internal/tui/image_attach.go +++ b/internal/tui/image_attach.go @@ -102,8 +102,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 } @@ -134,9 +134,12 @@ func (m model) modelSupportsVisionTUI() bool { } } } - // Fall back to the name heuristic for models not in the catalog or - // discovered list. - return modelregistry.VisionCapableByName(trimmed) + // Fall back to curated catalog or the name heuristic. + return modelregistry.SupportsVision(m.modelCatalog, trimmed) +} + +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..500e0f1bf 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -5496,6 +5496,9 @@ 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 { + return m.modelSupportsVisionFor(modelID) + } // 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, From 8bbe8627eca2dbcade004b280a3d36144da41ada Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Tue, 1 Sep 2026 19:18:57 -0400 Subject: [PATCH 8/9] test(mcp): use valid JSON-RPC structure for oversized SSE event test --- internal/mcp/network_client_test.go | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/internal/mcp/network_client_test.go b/internal/mcp/network_client_test.go index 2f1cb55e2..ece6234c2 100644 --- a/internal/mcp/network_client_test.go +++ b/internal/mcp/network_client_test.go @@ -386,8 +386,15 @@ func TestScanSSEEventsLargeImagePayload(t *testing.T) { }) t.Run("Oversized event exceeding 16 MiB rejected cleanly", func(t *testing.T) { - huge := strings.Repeat("A", 17*1024*1024) - stream := "event: message\ndata: " + huge + "\n\n" + oversizedRPC, marshalErr := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "result": map[string]any{"padding": strings.Repeat("A", 17*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 17 MiB event, got nil") From 1aa8806b2632b78766e1d4e23d96a595048b0931 Mon Sep 17 00:00:00 2001 From: cairn-intern Date: Wed, 2 Sep 2026 04:56:37 -0400 Subject: [PATCH 9/9] fix(mcp): snapshot vision capabilities, expand SSE payload bound, and unify ACP image authority --- internal/acp/agent.go | 38 +++++++++++++++++++++ internal/acp/agent_test.go | 42 +++++++++++++++++++++++ internal/mcp/network_client.go | 8 ++--- internal/mcp/network_client_test.go | 6 ++-- internal/mcp/non_text_content_test.go | 2 +- internal/mcp/registry.go | 8 +++-- internal/tui/image_attach.go | 48 ++++++++++++++++----------- internal/tui/model.go | 38 ++++++++++++++++++++- 8 files changed, 159 insertions(+), 31 deletions(-) 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/mcp/network_client.go b/internal/mcp/network_client.go index 12c218a99..ed920723d 100644 --- a/internal/mcp/network_client.go +++ b/internal/mcp/network_client.go @@ -651,10 +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. It accommodates an accepted -// 10 MiB image payload in base64 (~13.98 MiB) plus JSON-RPC envelope, metadata, -// and SSE framing overhead. -const maxSSEEventBytes = 16 * 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 ece6234c2..7b4cb2864 100644 --- a/internal/mcp/network_client_test.go +++ b/internal/mcp/network_client_test.go @@ -385,11 +385,11 @@ func TestScanSSEEventsLargeImagePayload(t *testing.T) { } }) - t.Run("Oversized event exceeding 16 MiB rejected cleanly", func(t *testing.T) { + 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", 17*1024*1024)}, + "result": map[string]any{"padding": strings.Repeat("A", 33*1024*1024)}, }) if marshalErr != nil { t.Fatal(marshalErr) @@ -397,7 +397,7 @@ func TestScanSSEEventsLargeImagePayload(t *testing.T) { stream := "event: message\ndata: " + string(oversizedRPC) + "\n\n" _, err := decodeSSERPCMessage(strings.NewReader(stream)) if err == nil { - t.Fatal("expected error for 17 MiB event, got 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 d7bc68fc8..c591b2a89 100644 --- a/internal/mcp/non_text_content_test.go +++ b/internal/mcp/non_text_content_test.go @@ -468,7 +468,7 @@ func TestImagePayloadsAreDecodedOnceAndNotPastTheBudget(t *testing.T) { if !strings.Contains(result.Output, "image/png") { t.Errorf("skipped images were not named:\n%s", result.Output) } - if !strings.Contains(result.Output, "which was not inspected because the aggregate image budget was reached") { + 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.") { diff --git a/internal/mcp/registry.go b/internal/mcp/registry.go index 1c295dd08..08fb03fcc 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -362,9 +362,13 @@ func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Res output = strings.TrimSpace(output + "\n\n" + note) } if uninspected := droppedContentNote(result.Content, disp, dispUninspected); uninspected != "" { - note := "[zero] this server also returned " + uninspected + ", which was not inspected because the aggregate image budget was reached." + 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 + ", which was not inspected because the aggregate image budget was reached." + note = "[zero] this server returned " + uninspected + ", " + verb + " because the aggregate image budget was reached." } output = strings.TrimSpace(output + "\n\n" + note) } diff --git a/internal/tui/image_attach.go b/internal/tui/image_attach.go index 2492ca7ff..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" @@ -111,33 +112,40 @@ func (m model) modelSupportsVisionFor(modelID string) 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. - 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 - } - } - return false // found the model in discovered list, no image modality + // 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 { + 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, true // found model with explicit modalities, no image modality + } + } + return false, false +} + func (m model) modelSupportsVisionTUI() bool { return m.modelSupportsVisionFor(m.modelName) } diff --git a/internal/tui/model.go b/internal/tui/model.go index 500e0f1bf..6f023da50 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -5496,8 +5496,44 @@ 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 { - return m.modelSupportsVisionFor(modelID) + 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