| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "strings" |
| 7 | "testing" |
| 8 | |
| 9 | "reasonix/internal/event" |
| 10 | "reasonix/internal/provider" |
| 11 | "reasonix/internal/tool" |
| 12 | ) |
| 13 | |
| 14 | // fakeImageTool implements tool.ImageTool: text and images travel on separate |
| 15 | // channels, like an MCP remote tool returning a screenshot. |
| 16 | type fakeImageTool struct { |
| 17 | text string |
| 18 | images []string |
| 19 | } |
| 20 | |
| 21 | func (f *fakeImageTool) Name() string { return "shot" } |
| 22 | func (f *fakeImageTool) Description() string { return "returns a screenshot" } |
| 23 | func (f *fakeImageTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) } |
| 24 | func (f *fakeImageTool) ReadOnly() bool { return true } |
| 25 | func (f *fakeImageTool) Execute(ctx context.Context, args json.RawMessage) (string, error) { |
| 26 | text, _, err := f.ExecuteWithImages(ctx, args) |
| 27 | return text, err |
| 28 | } |
| 29 | func (f *fakeImageTool) ExecuteWithImages(ctx context.Context, args json.RawMessage) (string, []string, error) { |
| 30 | return f.text, f.images, nil |
| 31 | } |
| 32 | |
| 33 | // Tool-result images must reach the session message intact even when the text |
| 34 | // output blows the truncation budget: the head+tail splice that trims tool text |
| 35 | // would corrupt a base64 payload, so images ride outside the truncated text. |
| 36 | func TestToolResultImagesBypassTruncation(t *testing.T) { |
| 37 | dataURL := "data:image/png;base64," + strings.Repeat("QUFB", 20000) // ~80KB payload, alone over the text budget |
| 38 | longText := strings.Repeat("x", maxToolOutputBytes+1024) + "[image: image/png]" |
| 39 | reg := tool.NewRegistry() |
| 40 | reg.Add(&fakeImageTool{text: longText, images: []string{dataURL}}) |
| 41 | prov := &scriptedProvider{name: "p", turns: [][]provider.Chunk{ |
| 42 | {toolCallChunk("c1", "shot", `{}`), {Type: provider.ChunkDone}}, |
| 43 | {{Type: provider.ChunkText, Text: "done"}, {Type: provider.ChunkDone}}, |
| 44 | }} |
| 45 | a := New(prov, reg, NewSession(""), Options{}, event.Discard) |
| 46 | if err := a.Run(context.Background(), "take a screenshot"); err != nil { |
| 47 | t.Fatalf("Run: %v", err) |
| 48 | } |
| 49 | var msg *provider.Message |
| 50 | for i := range a.session.Messages { |
| 51 | if a.session.Messages[i].Role == provider.RoleTool && a.session.Messages[i].Name == "shot" { |
| 52 | msg = &a.session.Messages[i] |
| 53 | break |
| 54 | } |
| 55 | } |
| 56 | if msg == nil { |
| 57 | t.Fatal("no tool message recorded for shot") |
| 58 | } |
| 59 | if len(msg.Images) != 1 || msg.Images[0] != dataURL { |
| 60 | t.Fatalf("tool message images corrupted or missing: got %d images", len(msg.Images)) |
| 61 | } |
| 62 | if len(msg.Content) > maxToolOutputBytes+1024 || !strings.Contains(msg.Content, "truncated") { |
| 63 | t.Fatalf("tool text should be head+tail truncated, len=%d", len(msg.Content)) |
| 64 | } |
| 65 | if strings.Contains(msg.Content, dataURL) { |
| 66 | t.Fatal("image payload must not be embedded in the tool text") |
| 67 | } |
| 68 | } |
| 69 |