diff --git a/pkg/app/app.go b/pkg/app/app.go index a14133bd88..9f85453c9b 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) } } 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 fd882c1cb9..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 @@ -18,6 +17,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 4551c81ffb..44817a7cd7 100644 --- a/pkg/tui/page/chat/chat.go +++ b/pkg/tui/page/chat/chat.go @@ -785,8 +785,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 +805,37 @@ 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 + 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, + MimeType: mimeType, + Data: data, + }) + } + 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 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) +}