From 02df4b8dd57ad8a2b40a830cc137454b2c32bb5d Mon Sep 17 00:00:00 2001 From: ollie Date: Mon, 25 May 2026 18:02:21 +0200 Subject: [PATCH] agent: plumb image ContentBlocks from tool results through to backend messages toolExecutor now returns (string, []ContentBlock, error). extractToolResult parses image blocks from structured tool output. Session.update populates ContentBlocks on tool messages. file_read detects PNG/JPEG/GIF/WEBP and returns base64-encoded image blocks. --- pkg/agent/core.go | 52 +++++---- pkg/agent/core_test.go | 146 ++++++++++++++++-------- pkg/agent/loop.go | 8 +- pkg/agent/loop_cache_test.go | 28 ++--- pkg/agent/loop_errlimit_test.go | 35 +++--- pkg/agent/loop_maxsteps_helpers_test.go | 2 + pkg/agent/loop_maxsteps_test.go | 16 +-- pkg/agent/loop_parallel_test.go | 24 ++-- pkg/agent/loop_texttool_test.go | 26 +++-- pkg/agent/loop_truncate_test.go | 20 ++-- pkg/agent/loop_turnerror_test.go | 1 + pkg/agent/session.go | 7 +- pkg/agent/state.go | 11 +- 13 files changed, 229 insertions(+), 147 deletions(-) diff --git a/pkg/agent/core.go b/pkg/agent/core.go index 0509ee7..305ccfe 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -98,21 +98,21 @@ func BuildRuntime(cfg *config.Config, d tools.Dispatcher, cwd string, env []stri } } - exec := func(ctx context.Context, name string, args json.RawMessage) (string, error) { - server, ok := serverOf[name] - if !ok { - return "", fmt.Errorf("unknown tool: %s", name) + exec := func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) { + server, ok := serverOf[name] + if !ok { + return "", nil, fmt.Errorf("unknown tool: %s", name) + } + raw, err := d.Dispatch(ctx, server, name, args) + if err != nil { + return "", nil, err + } + text, blocks, isErr := extractToolResult(raw) + if isErr { + return "", nil, fmt.Errorf("%s", text) + } + return text, blocks, nil } - raw, err := d.Dispatch(ctx, server, name, args) - if err != nil { - return "", err - } - text, isErr := extractToolResult(raw) - if isErr { - return "", fmt.Errorf("%s", text) - } - return text, nil - } var classify toolClassifier if srv, ok := d.GetServer("execute"); ok { @@ -1209,22 +1209,34 @@ func toolInfosToBackend(infos []tools.ToolInfo) []backend.Tool { return out } -func extractToolResult(raw json.RawMessage) (text string, isError bool) { +func extractToolResult(raw json.RawMessage) (text string, contentBlocks []backend.ContentBlock, isError bool) { var result struct { IsError bool `json:"isError"` Content []struct { - Type string `json:"type"` - Text string `json:"text"` + Type string `json:"type"` + Text string `json:"text"` + MediaType string `json:"media_type"` + Data string `json:"data"` } `json:"content"` } if err := json.Unmarshal(raw, &result); err != nil { - return string(raw), false + return string(raw), nil, false } var parts []string for _, c := range result.Content { - if c.Type == "text" { + switch c.Type { + case "text": parts = append(parts, c.Text) + case "image": + contentBlocks = append(contentBlocks, backend.ContentBlock{ + Type: "image", + ImageSource: &backend.ImageSource{ + Type: "base64", + MediaType: c.MediaType, + Data: c.Data, + }, + }) } } - return strings.Join(parts, "\n"), result.IsError + return strings.Join(parts, "\n"), contentBlocks, result.IsError } diff --git a/pkg/agent/core_test.go b/pkg/agent/core_test.go index bcbde31..86c27c6 100644 --- a/pkg/agent/core_test.go +++ b/pkg/agent/core_test.go @@ -184,7 +184,8 @@ func TestSubmit_StateTransitions(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -221,17 +222,19 @@ func TestSubmit_ToolCallStateTransitions(t *testing.T) { StopReason: "tool_calls", } close(ch) - return ch, nil +return ch, nil + } - return textStream("done"), nil +return textStream("done"), nil + } c := newCore(t, be, nil) var stateAtExec string - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { stateAtExec = c.State() - return `{}`, nil +return `{}`, nil, nil } collectEvents(context.Background(), c, "run tool") @@ -252,7 +255,8 @@ func TestSubmit_PreTurnHookBlocks(t *testing.T) { be := defaultBE() be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { callCount++ - return textStream("should not be called"), nil +return textStream("should not be called"), nil + } c := newCore(t, be, Hooks{HookPreTurn: []string{"exit 2"}}) @@ -281,7 +285,7 @@ func TestSubmit_PreTurnHookContext(t *testing.T) { lastUserMsg = m.Content } } - return textStream("ok"), nil + return textStream("ok"), nil } c := newCore(t, be, Hooks{HookPreTurn: []string{`echo "extra context"`}}) @@ -304,7 +308,8 @@ func TestSubmit_PostTurnHookContinue(t *testing.T) { be := defaultBE() be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { callCount++ - return textStream(fmt.Sprintf("response %d", callCount)), nil +return textStream(fmt.Sprintf("response %d", callCount)), nil + } c := newCore(t, be, Hooks{HookPostTurn: []string{hookScript}}) @@ -322,7 +327,7 @@ func TestSubmit_FIFODrain(t *testing.T) { be := defaultBE() be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { callCount.Add(1) - return textStream("ok"), nil + return textStream("ok"), nil } c := newCore(t, be, nil) @@ -350,7 +355,7 @@ func TestSubmit_PendingInjectAsNextTurn(t *testing.T) { be := defaultBE() be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { callCount++ - return textStream("ok"), nil + return textStream("ok"), nil } c := newCore(t, be, nil) @@ -370,7 +375,8 @@ func TestInterrupt_Running(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -412,7 +418,8 @@ func TestSubmit_WhileRunning(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -449,7 +456,8 @@ func TestManualCompact(t *testing.T) { be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { callCount++ stateAtCompact = c.State() - return textStream("summary or answer"), nil +return textStream("summary or answer"), nil + } // Seed a session with more than hotTailSize+warmIndexSize messages so compact() doesn't short-circuit. @@ -500,9 +508,11 @@ func TestAutoCompact(t *testing.T) { callCount++ if callCount == 1 { stateAtCompact = c.State() - return textStream("summary text for compaction"), nil +return textStream("summary text for compaction"), nil + } - return textStream("answer"), nil +return textStream("answer"), nil + } c.session = newSession("goal") @@ -664,7 +674,7 @@ func TestCommand_I_SetsInject(t *testing.T) { done := make(chan struct{}) be.respond = func(_ context.Context, msgs []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { <-done - return textStream("ok"), nil + return textStream("ok"), nil } c := newCore(t, be, nil) go c.Submit(context.Background(), "hello") @@ -684,7 +694,7 @@ func TestCommand_I_SubmitsWhenIdle(t *testing.T) { be := defaultBE() be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { callCount.Add(1) - return textStream("ok"), nil + return textStream("ok"), nil } c := newCore(t, be, nil) c.Submit(context.Background(), "/i my prompt") @@ -728,7 +738,8 @@ func TestCommand_Compact_WhileRunning(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -796,7 +807,8 @@ func TestCommand_Clear_WhileRunning(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -877,7 +889,8 @@ func TestRun_StreamInterrupted(t *testing.T) { be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { ch := make(chan backend.StreamEvent) close(ch) // close without sending Done=true - return ch, nil +return ch, nil + } c := newCore(t, be, nil) evs := collectEvents(context.Background(), c, "hello") @@ -900,7 +913,8 @@ func TestRun_UnknownStopReason(t *testing.T) { ch := make(chan backend.StreamEvent, 1) ch <- backend.StreamEvent{Done: true, StopReason: "max_completion_tokens", Content: "partial"} close(ch) - return ch, nil +return ch, nil + } c := newCore(t, be, nil) evs := collectEvents(context.Background(), c, "hello") @@ -930,9 +944,11 @@ func TestRun_ToolEmptyName(t *testing.T) { StopReason: "tool_calls", } close(ch) - return ch, nil +return ch, nil + } - return textStream("done"), nil +return textStream("done"), nil + } c := newCore(t, be, nil) collectEvents(context.Background(), c, "run empty tool") @@ -959,9 +975,11 @@ func TestRun_NoExec(t *testing.T) { StopReason: "tool_calls", } close(ch) - return ch, nil +return ch, nil + } - return textStream("done"), nil +return textStream("done"), nil + } c := newCore(t, be, nil) c.runtime.Exec = nil // no executor @@ -978,7 +996,7 @@ func TestHook_NonZeroExitCode(t *testing.T) { be := defaultBE() be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { callCount++ - return textStream("ok"), nil + return textStream("ok"), nil } // exit 1 is a non-blocking warning: turn should still proceed. c := newCore(t, be, Hooks{HookPreTurn: []string{"exit 1"}}) @@ -1018,7 +1036,8 @@ func TestCommand_Backend_WhileRunning(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -1074,7 +1093,8 @@ func TestCommand_Model_WhileRunning(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -1156,7 +1176,8 @@ func TestCommand_Agent_WhileRunning(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -1448,7 +1469,7 @@ func TestRun_RateLimitRetry(t *testing.T) { if callCount == 1 { return nil, &backend.RateLimitError{RetryAfter: time.Millisecond} } - return textStream("ok"), nil + return textStream("ok"), nil } c := newCore(t, be, nil) var retryEvents []Event @@ -1484,7 +1505,8 @@ func TestRun_ToolCancelledBeforeExec(t *testing.T) { StopReason: "tool_calls", } close(ch) - return ch, nil +return ch, nil + } c := newCore(t, be, nil) c.Submit(ctx, "hello") @@ -1529,14 +1551,15 @@ func TestRun_ExecCancelledWithInject(t *testing.T) { StopReason: "tool_calls", } close(ch) - return ch, nil +return ch, nil + } c := newCore(t, be, nil) inject := "user interrupt" c.pendingInject.Store(&inject) - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { cancel() // ctx cancelled during exec - return "", fmt.Errorf("exec cancelled") + return "", nil, fmt.Errorf("exec cancelled") } c.Submit(ctx, "hello") if callCount != 1 { @@ -1559,9 +1582,11 @@ func TestAutoCompact_WithHookContext(t *testing.T) { be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { callCount++ if callCount == 1 { - return textStream("summary text for compaction"), nil +return textStream("summary text for compaction"), nil + } - return textStream("answer"), nil +return textStream("answer"), nil + } c.session = newSession("goal") for range 15 { @@ -1628,7 +1653,8 @@ func TestCore_SetGenerationParams_WhileRunning(t *testing.T) { unblock := make(chan struct{}) be := defaultBE() be.respond = func(ctx context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return blockedStream(ctx, unblock), nil +return blockedStream(ctx, unblock), nil + } c := newCore(t, be, nil) @@ -1855,7 +1881,8 @@ func TestManualCompact_EmptySummary(t *testing.T) { ch := make(chan backend.StreamEvent, 1) ch <- backend.StreamEvent{Done: true, StopReason: "stop", Content: " "} close(ch) - return ch, nil +return ch, nil + } c.session = newSession("goal") for i := range 15 { @@ -1882,7 +1909,8 @@ func TestManualCompact_WithToolMessages(t *testing.T) { be := defaultBE() c := newCore(t, be, nil) be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return textStream("summary"), nil +return textStream("summary"), nil + } c.session = newSession("goal") // Seed enough messages to pass the compaction threshold. @@ -1992,7 +2020,7 @@ func TestCommand_Backend_Error(t *testing.T) { func TestExtractToolResult_Success(t *testing.T) { raw := json.RawMessage(`{"isError":false,"content":[{"type":"text","text":"hello"}]}`) - text, isErr := extractToolResult(raw) + text, _, isErr := extractToolResult(raw) if text != "hello" { t.Errorf("text = %q; want %q", text, "hello") } @@ -2003,7 +2031,7 @@ func TestExtractToolResult_Success(t *testing.T) { func TestExtractToolResult_IsError(t *testing.T) { raw := json.RawMessage(`{"isError":true,"content":[{"type":"text","text":"something failed"}]}`) - text, isErr := extractToolResult(raw) + text, _, isErr := extractToolResult(raw) if text != "something failed" { t.Errorf("text = %q; want %q", text, "something failed") } @@ -2014,7 +2042,7 @@ func TestExtractToolResult_IsError(t *testing.T) { func TestExtractToolResult_MultipleContentItems(t *testing.T) { raw := json.RawMessage(`{"isError":false,"content":[{"type":"text","text":"a"},{"type":"text","text":"b"}]}`) - text, _ := extractToolResult(raw) + text, _, _ := extractToolResult(raw) if text != "a\nb" { t.Errorf("text = %q; want %q", text, "a\nb") } @@ -2022,15 +2050,33 @@ func TestExtractToolResult_MultipleContentItems(t *testing.T) { func TestExtractToolResult_NonTextItemsSkipped(t *testing.T) { raw := json.RawMessage(`{"isError":false,"content":[{"type":"image","text":"ignored"},{"type":"text","text":"kept"}]}`) - text, _ := extractToolResult(raw) + text, _, _ := extractToolResult(raw) if text != "kept" { t.Errorf("text = %q; want %q", text, "kept") } } +func TestExtractToolResult_ImageBlock(t *testing.T) { + raw := json.RawMessage(`{"isError":false,"content":[{"type":"image","media_type":"image/png","data":"iVBOR"},{"type":"text","text":"screenshot above"}]}`) + text, blocks, isErr := extractToolResult(raw) + if text != "screenshot above" { + t.Errorf("text = %q; want %q", text, "screenshot above") + } + if isErr { + t.Error("isError should be false") + } + if len(blocks) != 1 { + t.Fatalf("blocks len = %d; want 1", len(blocks)) + } + b := blocks[0] + if b.Type != "image" || b.ImageSource == nil || b.ImageSource.MediaType != "image/png" || b.ImageSource.Data != "iVBOR" { + t.Errorf("block = %+v", b) + } +} + func TestExtractToolResult_InvalidJSON(t *testing.T) { raw := json.RawMessage(`not json`) - text, isErr := extractToolResult(raw) + text, _, isErr := extractToolResult(raw) if text != "not json" { t.Errorf("text = %q; want raw input on parse failure", text) } @@ -2297,7 +2343,7 @@ func TestBuildRuntime_ExecUnknownTool(t *testing.T) { d := tools.NewDispatcher() env := BuildRuntime(nil, d, t.TempDir(), nil) - _, err := env.Exec(context.Background(), "no_such_tool", json.RawMessage(`{}`)) + _, _, err := env.Exec(context.Background(), "no_such_tool", json.RawMessage(`{}`)) if err == nil || !strings.Contains(err.Error(), "unknown tool") { t.Errorf("expected unknown tool error; got %v", err) } @@ -2394,11 +2440,11 @@ func TestBuildRuntime_ExecDispatchSuccess(t *testing.T) { d := &mockDispatcher{ tools: []tools.ToolInfo{{Server: "s1", Name: "mytool"}}, dispatch: func(ctx context.Context, server, tool string, args json.RawMessage) (json.RawMessage, error) { - return json.RawMessage(`{"content":[{"type":"text","text":"ok"}]}`), nil +return json.RawMessage(`{"content":[{"type":"text","text":"ok"}]}`), nil }, } env := BuildRuntime(nil, d, t.TempDir(), nil) - result, err := env.Exec(context.Background(), "mytool", json.RawMessage(`{}`)) + result, _, err := env.Exec(context.Background(), "mytool", json.RawMessage(`{}`)) if err != nil { t.Fatal(err) } @@ -2416,7 +2462,7 @@ func TestBuildRuntime_ExecDispatchError(t *testing.T) { }, } env := BuildRuntime(nil, d, t.TempDir(), nil) - _, err := env.Exec(context.Background(), "mytool", json.RawMessage(`{}`)) + _, _, err := env.Exec(context.Background(), "mytool", json.RawMessage(`{}`)) if err == nil || !strings.Contains(err.Error(), "dispatch failed") { t.Errorf("err = %v", err) } @@ -2427,11 +2473,11 @@ func TestBuildRuntime_ExecToolResultIsError(t *testing.T) { d := &mockDispatcher{ tools: []tools.ToolInfo{{Server: "s1", Name: "mytool"}}, dispatch: func(ctx context.Context, server, tool string, args json.RawMessage) (json.RawMessage, error) { - return json.RawMessage(`{"isError":true,"content":[{"type":"text","text":"bad thing"}]}`), nil +return json.RawMessage(`{"isError":true,"content":[{"type":"text","text":"bad thing"}]}`), nil }, } env := BuildRuntime(nil, d, t.TempDir(), nil) - _, err := env.Exec(context.Background(), "mytool", json.RawMessage(`{}`)) + _, _, err := env.Exec(context.Background(), "mytool", json.RawMessage(`{}`)) if err == nil || !strings.Contains(err.Error(), "bad thing") { t.Errorf("err = %v", err) } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 7552c12..2a0cc43 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -40,7 +40,7 @@ const stallThreshold = 5 // task state re-injection into the conversation. const planReinjectInterval = 10 -type toolExecutor func(ctx context.Context, name string, args json.RawMessage) (string, error) +type toolExecutor func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) type agentConfig struct { Backend backend.Backend @@ -287,9 +287,10 @@ func run(ctx context.Context, cfg agentConfig, state state) error { } } var result string + var resultBlocks []backend.ContentBlock var isErr bool if cfg.Exec != nil { - out, err := cfg.Exec(ctx, tc.Name, tc.Arguments) + out, blocks, err := cfg.Exec(ctx, tc.Name, tc.Arguments) if err != nil { isErr = true if ctx.Err() != nil { @@ -305,6 +306,7 @@ func run(ctx context.Context, cfg agentConfig, state state) error { result = fmt.Sprintf("error: %v", err) } else { result = out + resultBlocks = blocks } } else { result = "error: no tool executor configured" @@ -337,7 +339,7 @@ func run(ctx context.Context, cfg agentConfig, state state) error { if !isErr && cfg.ClassifyTier != nil { tier = cfg.ClassifyTier(tc.Name) } - return toolResult{ToolCallID: tc.ID, Name: tc.Name, Content: result, IsError: isErr, Tier: tier}, false + return toolResult{ToolCallID: tc.ID, Name: tc.Name, Content: result, ContentBlocks: resultBlocks, IsError: isErr, Tier: tier}, false } isParallelSafe := func(name string) bool { diff --git a/pkg/agent/loop_cache_test.go b/pkg/agent/loop_cache_test.go index ff6018b..b57b33c 100644 --- a/pkg/agent/loop_cache_test.go +++ b/pkg/agent/loop_cache_test.go @@ -21,9 +21,9 @@ func TestResultCache_HitSkipsExec(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = func(string) bool { return true } - c.runtime.Exec = func(_ context.Context, name string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, name string, _ json.RawMessage) (string, []backend.ContentBlock, error) { atomic.AddInt32(&execCount, 1) - return "contents of a.txt", nil +return "contents of a.txt", nil, nil } evs := collectEvents(context.Background(), c, "dup read") @@ -54,9 +54,9 @@ func TestResultCache_DifferentArgsMiss(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = func(string) bool { return true } - c.runtime.Exec = func(_ context.Context, name string, args json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) { atomic.AddInt32(&execCount, 1) - return "result for " + string(args), nil +return "result for " + string(args), nil, nil } collectEvents(context.Background(), c, "different args") @@ -78,9 +78,9 @@ func TestResultCache_SerialToolNotCached(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = func(name string) bool { return false } // all serial - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { atomic.AddInt32(&execCount, 1) - return "ok", nil +return "ok", nil, nil } collectEvents(context.Background(), c, "serial tool") @@ -101,9 +101,11 @@ func multiTurnToolsStream(turns [][]backend.ToolCall) func(context.Context, []ba ch := make(chan backend.StreamEvent, 1) ch <- backend.StreamEvent{ToolCalls: turns[i], Done: true, StopReason: "tool_calls"} close(ch) - return ch, nil +return ch, nil + } - return textStream("done"), nil +return textStream("done"), nil + } } @@ -121,12 +123,12 @@ func TestResultCache_ErrorNotCached(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = func(string) bool { return true } - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { n := atomic.AddInt32(&execCount, 1) if n == 1 { - return "", &mockErr{"transient failure"} + return "", nil, &mockErr{"transient failure"} } - return "ok now", nil +return "ok now", nil, nil } evs := collectEvents(context.Background(), c, "error then ok") @@ -156,9 +158,9 @@ func TestResultCache_IdenticalParallelReads(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = func(string) bool { return true } // both parallel-safe → same batch - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { atomic.AddInt32(&execCount, 1) - return "file contents", nil +return "file contents", nil, nil } evs := collectEvents(context.Background(), c, "identical parallel") diff --git a/pkg/agent/loop_errlimit_test.go b/pkg/agent/loop_errlimit_test.go index fbc4841..2122fe9 100644 --- a/pkg/agent/loop_errlimit_test.go +++ b/pkg/agent/loop_errlimit_test.go @@ -23,7 +23,8 @@ func alwaysFailStream() func(context.Context, []backend.Message, []backend.Tool, Usage: backend.Usage{InputTokens: 10, OutputTokens: 5}, } close(ch) - return ch, nil +return ch, nil + } } @@ -32,8 +33,8 @@ func TestConsecutiveErrors_HardLimit(t *testing.T) { be.respond = alwaysFailStream() c := newCore(t, be, nil) - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return "", fmt.Errorf("always fails") + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { + return "", nil, fmt.Errorf("always fails") } evs := collectEvents(context.Background(), c, "do something") @@ -62,7 +63,8 @@ func TestConsecutiveErrors_SoftLimitNudge(t *testing.T) { if n > consecutiveErrorSoftLimit { for _, m := range msgs { if m.Role == "user" && strings.Contains(m.Content, "your last several tool calls all failed") { - return textStream("giving up"), nil +return textStream("giving up"), nil + } } // Nudge not found — keep going (will hit hard limit if broken). @@ -74,7 +76,8 @@ func TestConsecutiveErrors_SoftLimitNudge(t *testing.T) { Usage: backend.Usage{InputTokens: 10, OutputTokens: 5}, } close(ch) - return ch, nil +return ch, nil + } ch := make(chan backend.StreamEvent, 1) ch <- backend.StreamEvent{ @@ -84,12 +87,13 @@ func TestConsecutiveErrors_SoftLimitNudge(t *testing.T) { Usage: backend.Usage{InputTokens: 10, OutputTokens: 5}, } close(ch) - return ch, nil +return ch, nil + } c := newCore(t, be, nil) - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return "", fmt.Errorf("always fails") + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { + return "", nil, fmt.Errorf("always fails") } evs := collectEvents(context.Background(), c, "do something") @@ -124,10 +128,12 @@ func TestConsecutiveErrors_ResetOnSuccess(t *testing.T) { Usage: backend.Usage{InputTokens: 10, OutputTokens: 5}, } close(ch) - return ch, nil +return ch, nil + } if n >= 11 { - return textStream("done"), nil +return textStream("done"), nil + } ch := make(chan backend.StreamEvent, 1) ch <- backend.StreamEvent{ @@ -137,15 +143,16 @@ func TestConsecutiveErrors_ResetOnSuccess(t *testing.T) { Usage: backend.Usage{InputTokens: 10, OutputTokens: 5}, } close(ch) - return ch, nil +return ch, nil + } c := newCore(t, be, nil) - c.runtime.Exec = func(_ context.Context, name string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, name string, _ json.RawMessage) (string, []backend.ContentBlock, error) { if name == "good_tool" { - return "ok", nil +return "ok", nil, nil } - return "", fmt.Errorf("fails") + return "", nil, fmt.Errorf("fails") } evs := collectEvents(context.Background(), c, "do something") diff --git a/pkg/agent/loop_maxsteps_helpers_test.go b/pkg/agent/loop_maxsteps_helpers_test.go index 1c74dc5..e9cda99 100644 --- a/pkg/agent/loop_maxsteps_helpers_test.go +++ b/pkg/agent/loop_maxsteps_helpers_test.go @@ -21,6 +21,7 @@ func sequentialStream(responses []mockResponse) func(context.Context, []backend. i := int(atomic.AddInt32(&n, 1)) - 1 if i >= len(responses) { return textStream("done"), nil + } r := responses[i] ch := make(chan backend.StreamEvent, 1) @@ -32,6 +33,7 @@ func sequentialStream(responses []mockResponse) func(context.Context, []backend. } close(ch) return ch, nil + } } diff --git a/pkg/agent/loop_maxsteps_test.go b/pkg/agent/loop_maxsteps_test.go index 79e1415..16f4a8b 100644 --- a/pkg/agent/loop_maxsteps_test.go +++ b/pkg/agent/loop_maxsteps_test.go @@ -25,9 +25,9 @@ func TestMaxStepsZeroUnlimited(t *testing.T) { Backend: mb, Tools: []backend.Tool{{Name: "tool"}}, MaxSteps: 0, - Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { steps++ - return "ok", nil + return "ok", nil, nil }, } if err := run(context.Background(), cfg, newState()); err != nil { @@ -60,8 +60,8 @@ func TestMaxStepsSoftNudge(t *testing.T) { Backend: mb, Tools: []backend.Tool{{Name: "tool"}}, MaxSteps: 1, - Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return "ok", nil + Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { + return "ok", nil, nil }, Output: func(ev Event) { if ev.Role == "maxsteps" { @@ -117,9 +117,9 @@ func TestMaxStepsExactBoundary(t *testing.T) { Backend: mb, Tools: []backend.Tool{{Name: "tool"}}, MaxSteps: limit, - Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { toolRounds++ - return "ok", nil + return "ok", nil, nil }, } if err := run(context.Background(), cfg, newState()); err != nil { @@ -149,8 +149,8 @@ func TestTaskStateReinjection(t *testing.T) { cfg := agentConfig{ Backend: mb, Tools: []backend.Tool{{Name: "tool"}}, - Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return "ok", nil + Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { + return "ok", nil, nil }, } diff --git a/pkg/agent/loop_parallel_test.go b/pkg/agent/loop_parallel_test.go index 92e5447..11082f6 100644 --- a/pkg/agent/loop_parallel_test.go +++ b/pkg/agent/loop_parallel_test.go @@ -21,9 +21,11 @@ func toolsStream(calls []backend.ToolCall) func(context.Context, []backend.Messa ch := make(chan backend.StreamEvent, 1) ch <- backend.StreamEvent{ToolCalls: calls, Done: true, StopReason: "tool_calls"} close(ch) - return ch, nil +return ch, nil + } - return textStream("done"), nil +return textStream("done"), nil + } } @@ -51,13 +53,13 @@ func TestParallel_ConcurrentExecution(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = func(string) bool { return true } - c.runtime.Exec = func(ctx context.Context, name string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(ctx context.Context, name string, _ json.RawMessage) (string, []backend.ContentBlock, error) { started <- struct{}{} select { case <-gate: - return name + "-result", nil +return name + "-result", nil, nil case <-ctx.Done(): - return "", ctx.Err() + return "", nil, ctx.Err() } } @@ -93,11 +95,11 @@ func TestParallel_SerialToolBreaksBatch(t *testing.T) { c.runtime.ClassifyTool = func(name string) bool { return name == "read_a" || name == "read_c" } - c.runtime.Exec = func(_ context.Context, name string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, name string, _ json.RawMessage) (string, []backend.ContentBlock, error) { mu.Lock() order = append(order, name) mu.Unlock() - return name + "-result", nil +return name + "-result", nil, nil } collectEvents(context.Background(), c, "mixed tools") @@ -134,10 +136,10 @@ func TestParallel_CancellationFillsRemaining(t *testing.T) { c.runtime.ClassifyTool = func(name string) bool { return name == "tool_a" || name == "tool_b" } - c.runtime.Exec = func(execCtx context.Context, name string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(execCtx context.Context, name string, _ json.RawMessage) (string, []backend.ContentBlock, error) { cancel() // cancel on first execution; propagates to all <-execCtx.Done() - return "", execCtx.Err() + return "", nil, execCtx.Err() } evs := collectEvents(ctx, c, "cancel mid-batch") @@ -175,11 +177,11 @@ func TestParallel_NilClassifyToolIsSerial(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = nil // no classifier → all serial - c.runtime.Exec = func(_ context.Context, name string, _ json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, name string, _ json.RawMessage) (string, []backend.ContentBlock, error) { mu.Lock() order = append(order, name) mu.Unlock() - return name + "-result", nil +return name + "-result", nil, nil } evs := collectEvents(context.Background(), c, "serial fallback") diff --git a/pkg/agent/loop_texttool_test.go b/pkg/agent/loop_texttool_test.go index c1063fe..67d846c 100644 --- a/pkg/agent/loop_texttool_test.go +++ b/pkg/agent/loop_texttool_test.go @@ -18,10 +18,12 @@ func TestTextToolCall_ParsedAndExecuted(t *testing.T) { for _, m := range msgs { if m.Role == "tool" { // Second call after tool result: respond normally - return textStream("done"), nil +return textStream("done"), nil + } } - return textStream(`call_tool: {"calls":[{"tool":"file_read","args":["/tmp/test.go"]}]}`), nil +return textStream(`call_tool: {"calls":[{"tool":"file_read","args":["/tmp/test.go"]}]}`), nil + } var execCalled bool @@ -30,12 +32,12 @@ func TestTextToolCall_ParsedAndExecuted(t *testing.T) { {Name: "call_tool", Description: "run tools"}, {Name: "execute_code", Description: "run code"}, } - c.runtime.Exec = func(_ context.Context, name string, args json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) { execCalled = true if name != "call_tool" { t.Errorf("expected tool name 'call_tool', got %q", name) } - return "file contents here", nil +return "file contents here", nil, nil } evs := collectEvents(context.Background(), c, "read a file") @@ -55,10 +57,12 @@ func TestTextToolCall_RelaxedJSON(t *testing.T) { be.respond = func(_ context.Context, msgs []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { for _, m := range msgs { if m.Role == "tool" { - return textStream("done"), nil +return textStream("done"), nil + } } - return textStream(`execute_code: steps=[{code: "date"}]`), nil +return textStream(`execute_code: steps=[{code: "date"}]`), nil + } var execCalled bool @@ -67,12 +71,12 @@ func TestTextToolCall_RelaxedJSON(t *testing.T) { {Name: "execute_code", Description: "run code"}, {Name: "call_tool", Description: "run tools"}, } - c.runtime.Exec = func(_ context.Context, name string, args json.RawMessage) (string, error) { + c.runtime.Exec = func(_ context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) { execCalled = true if name != "execute_code" { t.Errorf("expected tool name 'execute_code', got %q", name) } - return "Sat May 23 11:00:00 UTC 2026", nil +return "Sat May 23 11:00:00 UTC 2026", nil, nil } collectEvents(context.Background(), c, "what time is it") @@ -87,7 +91,8 @@ func TestTextToolCall_RelaxedJSON(t *testing.T) { func TestTextToolCall_NoFalsePositive(t *testing.T) { be := defaultBE() be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return textStream("Here is how execute_code works in general."), nil +return textStream("Here is how execute_code works in general."), nil + } c := newCore(t, be, nil) @@ -110,7 +115,8 @@ func TestTextToolCall_NoFalsePositive(t *testing.T) { func TestTextToolCall_NoToolsConfigured(t *testing.T) { be := defaultBE() be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { - return textStream(`execute_code: steps=[{code: "print('hi')"}]`), nil +return textStream(`execute_code: steps=[{code: "print('hi')"}]`), nil + } c := newCore(t, be, nil) diff --git a/pkg/agent/loop_truncate_test.go b/pkg/agent/loop_truncate_test.go index fc22f15..196889f 100644 --- a/pkg/agent/loop_truncate_test.go +++ b/pkg/agent/loop_truncate_test.go @@ -23,8 +23,8 @@ func TestTruncation_LargeResultTruncated(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = nil c.runtime.ToolResultMaxBytes = maxBytes - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return large, nil + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { +return large, nil, nil } evs := collectEvents(context.Background(), c, "big result") @@ -55,8 +55,8 @@ func TestTruncation_SmallResultNotTruncated(t *testing.T) { c := newCore(t, be, nil) c.runtime.ToolResultMaxBytes = 8192 - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return "short result", nil + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { +return "short result", nil, nil } evs := collectEvents(context.Background(), c, "small result") @@ -82,8 +82,8 @@ func TestTruncation_ZeroDisabled(t *testing.T) { c := newCore(t, be, nil) c.runtime.ToolResultMaxBytes = 0 // disabled - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return large, nil + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { +return large, nil, nil } evs := collectEvents(context.Background(), c, "no limit") @@ -110,8 +110,8 @@ func TestTruncation_ErrorNotTruncated(t *testing.T) { c := newCore(t, be, nil) c.runtime.ToolResultMaxBytes = maxBytes - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return "", &mockErr{longErr} + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { + return "", nil, &mockErr{longErr} } evs := collectEvents(context.Background(), c, "error result") @@ -140,8 +140,8 @@ func TestTruncation_CachedResultAlreadyTruncated(t *testing.T) { c := newCore(t, be, nil) c.runtime.ClassifyTool = func(string) bool { return true } // read-safe → cacheable c.runtime.ToolResultMaxBytes = maxBytes - c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { - return large, nil + c.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) { +return large, nil, nil } evs := collectEvents(context.Background(), c, "cached truncated") diff --git a/pkg/agent/loop_turnerror_test.go b/pkg/agent/loop_turnerror_test.go index a1b2db4..2c2871f 100644 --- a/pkg/agent/loop_turnerror_test.go +++ b/pkg/agent/loop_turnerror_test.go @@ -24,6 +24,7 @@ func errThenOKStream(err error) func(context.Context, []backend.Message, []backe return nil, err } return textStream("recovered"), nil + } } diff --git a/pkg/agent/session.go b/pkg/agent/session.go index 7f7e054..48a0151 100644 --- a/pkg/agent/session.go +++ b/pkg/agent/session.go @@ -218,9 +218,10 @@ func (s *Session) update(assistant backend.Message, results []toolResult) { content = summarizeToolResult(r.Name, content) } s.messages = append(s.messages, backend.Message{ - Role: "tool", - Content: content, - ToolCallID: r.ToolCallID, + Role: "tool", + Content: content, + ContentBlocks: r.ContentBlocks, + ToolCallID: r.ToolCallID, }) } } diff --git a/pkg/agent/state.go b/pkg/agent/state.go index bc3645f..f6d8ced 100644 --- a/pkg/agent/state.go +++ b/pkg/agent/state.go @@ -13,11 +13,12 @@ const ( // toolResult holds the output of a single tool call. type toolResult struct { - ToolCallID string // may be empty (Ollama does not always set this) - Name string - Content string - IsError bool - Tier ResultTier // default TierHot; tools may self-classify + ToolCallID string // may be empty (Ollama does not always set this) + Name string + Content string + ContentBlocks []backend.ContentBlock // image blocks extracted from tool output + IsError bool + Tier ResultTier // default TierHot; tools may self-classify } // state is the interface both ephemeral and bead-backed state must satisfy.