From 85616ba5e9ebb84cda17e9bbd515b45fc440b0c2 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Mon, 6 Jul 2026 14:40:22 +0530 Subject: [PATCH 1/3] refactor(tui): extract attachments from document parts This commit updates extractAttachmentsFromSession to properly read modern MessagePartTypeDocument attachments instead of relying solely on the legacy 'Contents of' string matching fallback, addressing the TODO in chat.go. --- pkg/tui/page/chat/chat.go | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/pkg/tui/page/chat/chat.go b/pkg/tui/page/chat/chat.go index 4551c81ffb..613a043729 100644 --- a/pkg/tui/page/chat/chat.go +++ b/pkg/tui/page/chat/chat.go @@ -2,6 +2,7 @@ package chat import ( "context" + "encoding/base64" "fmt" "log/slog" "strings" @@ -785,8 +786,8 @@ func (p *chatPage) handleInlineEditCancelled(msg messages.InlineEditCancelledMsg } // extractAttachmentsFromSession extracts attachments from a session message at the given position. -// Attachments are stored as text parts in MultiContent with format "Contents of : ". -// TODO(krisetto): meh we can store and retrieve attachments better in the session itself +// Legacy attachments are stored as text parts in MultiContent with format "Contents of : ". +// New attachments are stored as Document parts. func (p *chatPage) extractAttachmentsFromSession(position int) []msgtypes.Attachment { sess := p.app.Session() if sess == nil || position < 0 || position >= len(sess.Messages) { @@ -805,20 +806,35 @@ func (p *chatPage) extractAttachmentsFromSession(position int) []msgtypes.Attach } var attachments []msgtypes.Attachment - const prefix = "Contents of " + const legacyPrefix = "Contents of " // Skip the first part (main text content), look for attachment parts for i := 1; i < len(msg.MultiContent); i++ { part := msg.MultiContent[i] + + if part.Type == chat.MessagePartTypeDocument && part.Document != nil { + content := part.Document.Source.InlineText + if content == "" && len(part.Document.Source.InlineData) > 0 { + content = "data:" + part.Document.MimeType + ";base64," + base64.StdEncoding.EncodeToString(part.Document.Source.InlineData) + } + if content != "" { + attachments = append(attachments, msgtypes.Attachment{ + Name: part.Document.Name, + Content: content, + }) + } + continue + } + if part.Type != chat.MessagePartTypeText { continue } text := part.Text - if !strings.HasPrefix(text, prefix) { + if !strings.HasPrefix(text, legacyPrefix) { continue } // Parse "Contents of : " - rest := text[len(prefix):] + rest := text[len(legacyPrefix):] before, after, ok := strings.Cut(rest, ": ") if !ok { continue From 5b3f376918c290032fbf7e54acaf90083310df15 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Mon, 6 Jul 2026 18:52:48 +0530 Subject: [PATCH 2/3] fix(tui): natively process binary attachments on resend path This fixes a context overflow issue where binary attachments (like images) were being serialized as massive base64 text strings when a message was edited and resent. We now extend messages.Attachment to properly hold MimeType and Data, and rebuild the Document part natively in App.Run. --- pkg/app/app.go | 32 ++++++++++++++++++++++++++++++-- pkg/tui/messages/session.go | 4 ++++ pkg/tui/page/chat/chat.go | 15 ++++++++------- 3 files changed, 42 insertions(+), 9 deletions(-) diff --git a/pkg/app/app.go b/pkg/app/app.go index a14133bd88..b44ce0a3ff 100644 --- a/pkg/app/app.go +++ b/pkg/app/app.go @@ -488,8 +488,35 @@ func (a *App) Run(ctx context.Context, cancel context.CancelFunc, message string case att.Content != "": // Inline content attachment (e.g. pasted text). a.processInlineAttachment(att, &textBuilder) + case len(att.Data) > 0: + // Binary inline content. + doc, resizeMeta, procErr := chat.ProcessAttachmentWithMetadata(chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: att.Name, + MimeType: att.MimeType, + Source: chat.DocumentSource{ + InlineData: att.Data, + }, + }, + }) + if procErr != nil { + slog.WarnContext(ctx, "skipping inline attachment: processing failed", "name", att.Name, "error", procErr) + a.sendEvent(ctx, runtime.Warning(fmt.Sprintf("Skipped attachment %s: %s", att.Name, procErr), "")) + continue + } + if resizeMeta != nil { + if note := chat.FormatDimensionNote(resizeMeta); note != "" { + textBuilder.WriteString("\n") + textBuilder.WriteString(note) + } + } + binaryParts = append(binaryParts, chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: &doc, + }) default: - slog.DebugContext(ctx, "skipping attachment with no file path or content", "name", att.Name) + slog.DebugContext(ctx, "skipping attachment with no file path, content, or data", "name", att.Name) } } @@ -599,7 +626,8 @@ func (a *App) processFileAttachment(ctx context.Context, att messages.Attachment // For images, emit a dimension note so the model can map coordinates back to the original. if resizeMeta != nil { if note := chat.FormatDimensionNote(resizeMeta); note != "" { - textBuilder.WriteString("\n" + note) + textBuilder.WriteString("\n") + textBuilder.WriteString(note) } } *binaryParts = append(*binaryParts, chat.MessagePart{ diff --git a/pkg/tui/messages/session.go b/pkg/tui/messages/session.go index fd882c1cb9..6daf19d984 100644 --- a/pkg/tui/messages/session.go +++ b/pkg/tui/messages/session.go @@ -18,6 +18,10 @@ type Attachment struct { // backing temp file is cleaned up before the message reaches the app layer. // Empty for file-reference attachments that are read from disk. Content string + // MimeType is the MIME type of the binary data, if Data is present. + MimeType string + // Data holds raw binary content for non-text inline attachments. + Data []byte } // Session lifecycle messages control session state and persistence. diff --git a/pkg/tui/page/chat/chat.go b/pkg/tui/page/chat/chat.go index 613a043729..44817a7cd7 100644 --- a/pkg/tui/page/chat/chat.go +++ b/pkg/tui/page/chat/chat.go @@ -2,7 +2,6 @@ package chat import ( "context" - "encoding/base64" "fmt" "log/slog" "strings" @@ -814,13 +813,15 @@ func (p *chatPage) extractAttachmentsFromSession(position int) []msgtypes.Attach if part.Type == chat.MessagePartTypeDocument && part.Document != nil { content := part.Document.Source.InlineText - if content == "" && len(part.Document.Source.InlineData) > 0 { - content = "data:" + part.Document.MimeType + ";base64," + base64.StdEncoding.EncodeToString(part.Document.Source.InlineData) - } - if content != "" { + data := part.Document.Source.InlineData + mimeType := part.Document.MimeType + + if content != "" || len(data) > 0 { attachments = append(attachments, msgtypes.Attachment{ - Name: part.Document.Name, - Content: content, + Name: part.Document.Name, + Content: content, + MimeType: mimeType, + Data: data, }) } continue From 1e98435508fe6ae86d641e545d3b26e463116fb1 Mon Sep 17 00:00:00 2001 From: piyush0049 Date: Wed, 8 Jul 2026 20:24:53 +0530 Subject: [PATCH 3/3] fix: address PR review feedback for attachment extraction --- pkg/app/app.go | 3 +- pkg/chat/attach_test.go | 26 +++++++++++ pkg/tui/messages/session.go | 11 ++--- pkg/tui/page/chat/chat_test.go | 83 ++++++++++++++++++++++++++++++++++ 4 files changed, 115 insertions(+), 8 deletions(-) create mode 100644 pkg/tui/page/chat/chat_test.go diff --git a/pkg/app/app.go b/pkg/app/app.go index b44ce0a3ff..9f85453c9b 100644 --- a/pkg/app/app.go +++ b/pkg/app/app.go @@ -626,8 +626,7 @@ func (a *App) processFileAttachment(ctx context.Context, att messages.Attachment // For images, emit a dimension note so the model can map coordinates back to the original. if resizeMeta != nil { if note := chat.FormatDimensionNote(resizeMeta); note != "" { - textBuilder.WriteString("\n") - textBuilder.WriteString(note) + textBuilder.WriteString("\n" + note) } } *binaryParts = append(*binaryParts, chat.MessagePart{ diff --git a/pkg/chat/attach_test.go b/pkg/chat/attach_test.go index e36b8fabfe..07fa0bcac4 100644 --- a/pkg/chat/attach_test.go +++ b/pkg/chat/attach_test.go @@ -141,6 +141,32 @@ func TestProcessAttachment_ImageTooLarge_Resized(t *testing.T) { assert.LessOrEqual(t, b.Dy(), chat.MaxImageDimension) } +func TestProcessAttachmentWithMetadata_IdempotentResend(t *testing.T) { + t.Parallel() + // An image that is already within the dimension limits + smallData := encodeJPEGBytes(50, 50) + + doc, resizeMeta, err := chat.ProcessAttachmentWithMetadata(chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: "photo.jpg", + MimeType: "image/jpeg", + Source: chat.DocumentSource{InlineData: smallData}, + }, + }) + require.NoError(t, err) + assert.NotEmpty(t, doc.Source.InlineData) + + // Since it fits within the limit, it shouldn't be resized + if assert.NotNil(t, resizeMeta) { + assert.False(t, resizeMeta.Resized) + } + + // FormatDimensionNote should return empty string for un-resized images + note := chat.FormatDimensionNote(resizeMeta) + assert.Empty(t, note, "expected no dimension note for an image that wasn't resized") +} + func TestProcessAttachment_PDF_Passthrough(t *testing.T) { t.Parallel() pdfBytes := []byte("%PDF-1.4 fake pdf content for testing") diff --git a/pkg/tui/messages/session.go b/pkg/tui/messages/session.go index 6daf19d984..d7423e3bd9 100644 --- a/pkg/tui/messages/session.go +++ b/pkg/tui/messages/session.go @@ -2,12 +2,11 @@ package messages import "github.com/docker/docker-agent/pkg/session" -// Attachment represents content attached to a message. It is either a reference -// to a file on disk (FilePath is set) or inline content already in memory -// (Content is set, e.g. pasted text). When FilePath is set the consumer reads -// and classifies the file at send time; when only Content is set the consumer -// uses it directly as inline text. This design lets us add binary-file support -// (images, PDFs, …) in the future by extending the struct with a MimeType hint. +// Attachment represents content attached to a message. It can be a reference +// to a file on disk (FilePath is set), inline text content (Content is set), +// or inline binary content (Data and MimeType are set). When FilePath is set, +// the consumer reads and classifies the file at send time; inline content +// is used directly. type Attachment struct { // Name is the human-readable label (e.g. "paste-1", "main.go"). Name string diff --git a/pkg/tui/page/chat/chat_test.go b/pkg/tui/page/chat/chat_test.go new file mode 100644 index 0000000000..8afb3ee3da --- /dev/null +++ b/pkg/tui/page/chat/chat_test.go @@ -0,0 +1,83 @@ +package chat + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/app" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" +) + +func TestExtractAttachmentsFromSession(t *testing.T) { + t.Parallel() + + sess := session.New() + p := newTestChatPage(t) + p.app = app.New(t.Context(), queueTestRuntime{}, sess) + + msg := session.Message{ + Message: chat.Message{ + Role: chat.MessageRoleUser, + MultiContent: []chat.MessagePart{ + {Type: chat.MessagePartTypeText, Text: "main user prompt"}, + // 1. Document with InlineText + { + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: "text_doc.txt", + Source: chat.DocumentSource{ + InlineText: "text content", + }, + }, + }, + // 2. Document with InlineData + { + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: "binary_doc.png", + MimeType: "image/png", + Source: chat.DocumentSource{ + InlineData: []byte("binary data"), + }, + }, + }, + // 3. Legacy Text Contents + { + Type: chat.MessagePartTypeText, + Text: "Contents of legacy_doc.txt: legacy content", + }, + // 4. Empty Document + { + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: "empty_doc.txt", + }, + }, + }, + }, + } + sess.Messages = append(sess.Messages, session.Item{Message: &msg}) + + attachments := p.extractAttachmentsFromSession(0) + require.Len(t, attachments, 3, "expected exactly 3 extracted attachments") + + // 1. Document with InlineText + assert.Equal(t, "text_doc.txt", attachments[0].Name) + assert.Equal(t, "text content", attachments[0].Content) + assert.Empty(t, attachments[0].Data) + + // 2. Document with InlineData + assert.Equal(t, "binary_doc.png", attachments[1].Name) + assert.Equal(t, "image/png", attachments[1].MimeType) + assert.Equal(t, []byte("binary data"), attachments[1].Data) + assert.Empty(t, attachments[1].Content) + + // 3. Legacy Text Contents + assert.Equal(t, "legacy_doc.txt", attachments[2].Name) + assert.Equal(t, "legacy content", attachments[2].Content) + assert.Empty(t, attachments[2].Data) + assert.Empty(t, attachments[2].MimeType) +}