From 129ce743c73e279bcd21c2f9ccc2049d00a8ac1c Mon Sep 17 00:00:00 2001 From: David Gageot Date: Thu, 21 May 2026 13:14:09 +0200 Subject: [PATCH 1/4] fix: wait for supervisor watcher shutdown --- pkg/tools/lifecycle/supervisor.go | 38 ++++++++++++++++++++++---- pkg/tools/lifecycle/supervisor_test.go | 15 ++++++++++ 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/pkg/tools/lifecycle/supervisor.go b/pkg/tools/lifecycle/supervisor.go index ef3c2c4ae..911e8af41 100644 --- a/pkg/tools/lifecycle/supervisor.go +++ b/pkg/tools/lifecycle/supervisor.go @@ -140,6 +140,10 @@ type Supervisor struct { // fresh channel by Start when transitioning out of a terminal state. done chan struct{} + // watchDone is closed by the current watcher goroutine. Stop waits on it + // after closing the session so no transport goroutines are left behind. + watchDone chan struct{} + // randFloat is the jitter source; tests may override. randFloat func() float64 } @@ -214,6 +218,9 @@ func (s *Supervisor) Start(ctx context.Context) error { } s.session = sess spawnWatcher := !s.watcherAlive + if spawnWatcher { + s.watchDone = make(chan struct{}) + } s.watcherAlive = true // Recovering from a terminal state (Failed → Start, or a watcher // that previously exited): refresh `done` so RestartAndWait callers @@ -244,24 +251,40 @@ func (s *Supervisor) Start(ctx context.Context) error { func (s *Supervisor) Stop(ctx context.Context) error { s.mu.Lock() if s.stopping { + watchDone := s.watchDone s.mu.Unlock() - return nil + return waitForWatcher(ctx, watchDone) } s.stopping = true sess := s.session s.session = nil + watchDone := s.watchDone s.mu.Unlock() s.tracker.Set(StateStopped) s.signalDone() - if sess == nil { + var closeErr error + if sess != nil { + closeErr = sess.Close(context.WithoutCancel(ctx)) + } + waitErr := waitForWatcher(ctx, watchDone) + if closeErr != nil && ctx.Err() == nil { + return closeErr + } + return waitErr +} + +func waitForWatcher(ctx context.Context, done <-chan struct{}) error { + if done == nil { return nil } - if err := sess.Close(context.WithoutCancel(ctx)); err != nil && ctx.Err() == nil { - return err + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() } - return nil } // RestartAndWait closes the current session (if any) so the watcher @@ -326,7 +349,12 @@ func (s *Supervisor) watch(ctx context.Context) { defer func() { s.mu.Lock() s.watcherAlive = false + watchDone := s.watchDone + s.watchDone = nil s.mu.Unlock() + if watchDone != nil { + close(watchDone) + } }() log := s.policy.logger() diff --git a/pkg/tools/lifecycle/supervisor_test.go b/pkg/tools/lifecycle/supervisor_test.go index ac8dfd6ba..4fb587348 100644 --- a/pkg/tools/lifecycle/supervisor_test.go +++ b/pkg/tools/lifecycle/supervisor_test.go @@ -458,3 +458,18 @@ func TestBackoff_Jitter(t *testing.T) { d = lifecycle.ExportedBackoffDelay(b, 0, func() float64 { return 0 }) assert.Check(t, d == 50*time.Millisecond) } + +func TestSupervisor_StopWaitsForWatcher(t *testing.T) { + t.Parallel() + + sess := newFakeSession() + c := newScriptedConnector(scriptStep{session: sess}) + s := lifecycle.New("test", c, lifecycle.Policy{}) + + assert.NilError(t, s.Start(t.Context())) + + start := time.Now() + assert.NilError(t, s.Stop(t.Context())) + assert.Check(t, time.Since(start) < time.Second) + assert.Check(t, is.Equal(s.State().State, lifecycle.StateStopped)) +} From 53f55bf7e26d020695470e5fdedae655221c946d Mon Sep 17 00:00:00 2001 From: David Gageot Date: Thu, 21 May 2026 13:21:22 +0200 Subject: [PATCH 2/4] fix: spool large mcp media to disk --- pkg/runtime/toolexec/dispatcher.go | 22 ++++++++---- pkg/tools/mcp/mcp.go | 55 +++++++++++++++++++++++++++--- pkg/tools/mcp/mcp_test.go | 19 +++++++++++ pkg/tools/tools.go | 7 ++-- 4 files changed, 89 insertions(+), 14 deletions(-) diff --git a/pkg/runtime/toolexec/dispatcher.go b/pkg/runtime/toolexec/dispatcher.go index bf901243b..4820a878c 100644 --- a/pkg/runtime/toolexec/dispatcher.go +++ b/pkg/runtime/toolexec/dispatcher.go @@ -724,13 +724,21 @@ func buildMultiContent(text string, images []tools.MediaContent) []chat.MessageP parts := make([]chat.MessagePart, 0, 1+len(images)) parts = append(parts, chat.MessagePart{Type: chat.MessagePartTypeText, Text: text}) for _, img := range images { - parts = append(parts, chat.MessagePart{ - Type: chat.MessagePartTypeImageURL, - ImageURL: &chat.MessageImageURL{ - URL: "data:" + img.MimeType + ";base64," + img.Data, - Detail: chat.ImageURLDetailAuto, - }, - }) + switch { + case img.FilePath != "": + parts = append(parts, chat.MessagePart{ + Type: chat.MessagePartTypeText, + Text: fmt.Sprintf("[image saved to %s (%s)]", img.FilePath, img.MimeType), + }) + case img.Data != "": + parts = append(parts, chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{ + URL: "data:" + img.MimeType + ";base64," + img.Data, + Detail: chat.ImageURLDetailAuto, + }, + }) + } } return parts } diff --git a/pkg/tools/mcp/mcp.go b/pkg/tools/mcp/mcp.go index 42208e4fc..8a44efc4a 100644 --- a/pkg/tools/mcp/mcp.go +++ b/pkg/tools/mcp/mcp.go @@ -12,6 +12,7 @@ import ( "log/slog" "net/url" "os" + "path/filepath" "strings" "sync" "time" @@ -714,6 +715,8 @@ func isInitNotificationSendError(err error) bool { return false } +const maxInlineMediaBytes = 256 * 1024 + func processMCPContent(toolResult *mcp.CallToolResult) *tools.ToolCallResult { var text strings.Builder var images, audios []tools.MediaContent @@ -760,12 +763,54 @@ func processMCPContent(toolResult *mcp.CallToolResult) *tools.ToolCallResult { } } -// encodeMedia re-encodes raw bytes (as decoded by the MCP SDK) back to base64 -// for our internal MediaContent representation. +// encodeMedia keeps small payloads inline and spools larger ones to disk so the +// session and TUI do not retain duplicate base64 copies. func encodeMedia(data []byte, mimeType string) tools.MediaContent { - return tools.MediaContent{ - Data: base64.StdEncoding.EncodeToString(data), - MimeType: mimeType, + media := tools.MediaContent{MimeType: mimeType} + if len(data) <= maxInlineMediaBytes { + media.Data = base64.StdEncoding.EncodeToString(data) + return media + } + + path, err := writeMediaFile(data, mimeType) + if err != nil { + slog.Warn("failed to spool MCP media to disk", "mime_type", mimeType, "bytes", len(data), "error", err) + media.Data = base64.StdEncoding.EncodeToString(data) + return media + } + media.FilePath = path + return media +} + +func writeMediaFile(data []byte, mimeType string) (string, error) { + dir, err := os.MkdirTemp("", "docker-agent-mcp-media-*") + if err != nil { + return "", err + } + path := filepath.Join(dir, "media"+mediaExtension(mimeType)) + if err := os.WriteFile(path, data, 0o600); err != nil { + _ = os.RemoveAll(dir) + return "", err + } + return path, nil +} + +func mediaExtension(mimeType string) string { + switch mimeType { + case "image/png": + return ".png" + case "image/jpeg": + return ".jpg" + case "image/gif": + return ".gif" + case "image/webp": + return ".webp" + case "audio/wav", "audio/wave", "audio/x-wav": + return ".wav" + case "audio/mpeg", "audio/mp3": + return ".mp3" + default: + return ".bin" } } diff --git a/pkg/tools/mcp/mcp_test.go b/pkg/tools/mcp/mcp_test.go index 997ac7c8f..4860037b7 100644 --- a/pkg/tools/mcp/mcp_test.go +++ b/pkg/tools/mcp/mcp_test.go @@ -1,9 +1,12 @@ package mcp import ( + "bytes" "context" "fmt" "iter" + "os" + "path/filepath" "sync" "sync/atomic" "testing" @@ -536,3 +539,19 @@ func TestCallToolRecoversFromErrSessionMissing(t *testing.T) { assert.Equal(t, "recovered", result.Output) assert.Equal(t, int32(2), callCount.Load(), "expected exactly 2 CallTool invocations (1 failed + 1 retry)") } + +func TestProcessMCPContentSpoolsLargeMedia(t *testing.T) { + large := bytes.Repeat([]byte("x"), maxInlineMediaBytes+1) + result := processMCPContent(callToolResult(&mcp.ImageContent{Data: large, MIMEType: "image/png"})) + + require.Len(t, result.Images, 1) + img := result.Images[0] + assert.Empty(t, img.Data) + assert.Equal(t, "image/png", img.MimeType) + require.NotEmpty(t, img.FilePath) + defer os.RemoveAll(filepath.Dir(img.FilePath)) + + got, err := os.ReadFile(img.FilePath) + require.NoError(t, err) + assert.Equal(t, large, got) +} diff --git a/pkg/tools/tools.go b/pkg/tools/tools.go index 2d0185943..1cc6f8026 100644 --- a/pkg/tools/tools.go +++ b/pkg/tools/tools.go @@ -73,8 +73,11 @@ type FunctionCall struct { // MediaContent represents base64-encoded binary data (image, audio, etc.) // returned by a tool. type MediaContent struct { - // Data is the base64-encoded payload. - Data string `json:"data"` + // Data is the base64-encoded payload. It is kept only for small media; large + // MCP payloads are spooled to FilePath to avoid retaining duplicate base64. + Data string `json:"data,omitempty"` + // FilePath is an optional local file containing the decoded media payload. + FilePath string `json:"filePath,omitempty"` // MimeType identifies the content type (e.g. "image/png", "audio/wav"). MimeType string `json:"mimeType"` } From 4bacbbacdd3e435bc11b5b1d74dd43b6c0f48106 Mon Sep 17 00:00:00 2001 From: David Gageot Date: Thu, 21 May 2026 13:23:28 +0200 Subject: [PATCH 3/4] fix: slim retained tui tool results --- pkg/tools/tools.go | 10 +++++++++ pkg/tools/tools_test.go | 21 +++++++++++++++++++ pkg/tui/components/messages/messages.go | 4 ++-- .../reasoningblock/reasoningblock.go | 2 +- 4 files changed, 34 insertions(+), 3 deletions(-) diff --git a/pkg/tools/tools.go b/pkg/tools/tools.go index 1cc6f8026..ca0cf84e3 100644 --- a/pkg/tools/tools.go +++ b/pkg/tools/tools.go @@ -102,6 +102,16 @@ type ToolCallResult struct { StructuredContent any `json:"structuredContent,omitempty"` } +func (r *ToolCallResult) WithoutPayload() *ToolCallResult { + if r == nil { + return nil + } + return &ToolCallResult{ + IsError: r.IsError, + Meta: r.Meta, + } +} + func ResultError(output string) *ToolCallResult { return &ToolCallResult{ Output: output, diff --git a/pkg/tools/tools_test.go b/pkg/tools/tools_test.go index 1f5338dbe..0ada71e0d 100644 --- a/pkg/tools/tools_test.go +++ b/pkg/tools/tools_test.go @@ -78,3 +78,24 @@ func TestNewHandler_InvalidArguments(t *testing.T) { }) require.Error(t, err) } + +func TestToolCallResultWithoutPayload(t *testing.T) { + result := &ToolCallResult{ + Output: "large output", + IsError: true, + Meta: "metadata", + Images: []MediaContent{{Data: "image", MimeType: "image/png"}}, + Audios: []MediaContent{{Data: "audio", MimeType: "audio/wav"}}, + StructuredContent: map[string]any{"key": "value"}, + } + + slim := result.WithoutPayload() + + require.NotNil(t, slim) + assert.Empty(t, slim.Output) + assert.True(t, slim.IsError) + assert.Equal(t, "metadata", slim.Meta) + assert.Nil(t, slim.Images) + assert.Nil(t, slim.Audios) + assert.Nil(t, slim.StructuredContent) +} diff --git a/pkg/tui/components/messages/messages.go b/pkg/tui/components/messages/messages.go index c5c5748f8..178d7461f 100644 --- a/pkg/tui/components/messages/messages.go +++ b/pkg/tui/components/messages/messages.go @@ -1475,7 +1475,7 @@ func (m *model) AddToolResult(msg *runtime.ToolCallResponseEvent, status types.T if m.messages[i].Type == types.MessageTypeAssistantReasoningBlock { if block, ok := m.views[i].(*reasoningblock.Model); ok { if block.HasToolCall(msg.ToolCallID) { - cmd := block.UpdateToolResult(msg.ToolCallID, msg.Response, status, msg.Result) + cmd := block.UpdateToolResult(msg.ToolCallID, msg.Response, status, msg.Result.WithoutPayload()) m.invalidateItem(i) return cmd } @@ -1489,7 +1489,7 @@ func (m *model) AddToolResult(msg *runtime.ToolCallResponseEvent, status types.T if toolMessage.Type == types.MessageTypeToolCall && toolMessage.ToolCall.ID == msg.ToolCallID { toolMessage.Content = strings.ReplaceAll(msg.Response, "\t", " ") toolMessage.ToolStatus = status - toolMessage.ToolResult = msg.Result + toolMessage.ToolResult = msg.Result.WithoutPayload() m.invalidateItem(i) view := m.createToolCallView(toolMessage) diff --git a/pkg/tui/components/reasoningblock/reasoningblock.go b/pkg/tui/components/reasoningblock/reasoningblock.go index aff776a07..5385d3a41 100644 --- a/pkg/tui/components/reasoningblock/reasoningblock.go +++ b/pkg/tui/components/reasoningblock/reasoningblock.go @@ -240,7 +240,7 @@ func (m *Model) UpdateToolResult(toolCallID, content string, status types.ToolSt entry.msg.Content = strings.ReplaceAll(content, "\t", " ") entry.msg.ToolStatus = status - entry.msg.ToolResult = result + entry.msg.ToolResult = result.WithoutPayload() // Set grace period if transitioning from in-progress to completed // Total visible time = completedToolVisibleDuration + completedToolFadeDuration From bb2f838f05a76502cde0ce1f10312ec0f143ac0c Mon Sep 17 00:00:00 2001 From: David Gageot Date: Thu, 21 May 2026 14:08:12 +0200 Subject: [PATCH 4/4] fix: avoid retaining file contents in metadata --- pkg/tools/builtin/filesystem/filesystem.go | 2 -- pkg/tools/builtin/filesystem/filesystem_test.go | 2 ++ 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/tools/builtin/filesystem/filesystem.go b/pkg/tools/builtin/filesystem/filesystem.go index 95293737d..811c31b45 100644 --- a/pkg/tools/builtin/filesystem/filesystem.go +++ b/pkg/tools/builtin/filesystem/filesystem.go @@ -267,7 +267,6 @@ type ReadFileArgs struct { type ReadFileMeta struct { Path string `json:"path"` - Content string `json:"content"` LineCount int `json:"lineCount"` Error string `json:"error,omitempty"` } @@ -1086,7 +1085,6 @@ func (t *ToolSet) handleReadMultipleFiles(ctx context.Context, args ReadMultiple Path: path, Content: text, }) - entry.Content = text entry.LineCount = strings.Count(text, "\n") + 1 meta.Files = append(meta.Files, entry) } diff --git a/pkg/tools/builtin/filesystem/filesystem_test.go b/pkg/tools/builtin/filesystem/filesystem_test.go index 9c2139641..f1c7398ee 100644 --- a/pkg/tools/builtin/filesystem/filesystem_test.go +++ b/pkg/tools/builtin/filesystem/filesystem_test.go @@ -104,6 +104,7 @@ func TestFilesystemTool_ReadFile_TildePath(t *testing.T) { require.NoError(t, err) assert.False(t, result.IsError) assert.Equal(t, content, result.Output) + assert.Equal(t, ReadFileMeta{LineCount: 1}, result.Meta) } func TestFilesystemTool_WriteFile(t *testing.T) { @@ -166,6 +167,7 @@ func TestFilesystemTool_ReadFile(t *testing.T) { }) require.NoError(t, err) assert.Equal(t, content, result.Output) + assert.Equal(t, ReadFileMeta{LineCount: 1}, result.Meta) result, err = tool.handleReadFile(t.Context(), ReadFileArgs{ Path: "nonexistent.txt",