mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add file/image attachment support to chat input (#22604)
This change adds support for image attachments to chat via add button and clipboard paste. Files are stored in a new `chat_files` table and referenced by ID in message content. File data is resolved from storage at LLM dispatch time, keeping the message content column small. Upload validates MIME types via content type or content sniffing against an allowlist (png, jpeg, gif, webp). The retrieval endpoint serves files with immutable caching headers. On the frontend, uploads start eagerly on attach with a background fetch to pre-warm the browser HTTP cache so the timeline renders instantly after send.
This commit is contained in:
+35
-11
@@ -165,6 +165,9 @@ type CreateOptions struct {
|
||||
ModelConfigID uuid.UUID
|
||||
SystemPrompt string
|
||||
InitialUserContent []fantasy.Content
|
||||
// ContentFileIDs maps content block indices to their chat_files IDs
|
||||
// so the file_id can be preserved in the stored message JSON.
|
||||
ContentFileIDs map[int]uuid.UUID
|
||||
}
|
||||
|
||||
// SendMessageBusyBehavior controls what happens when a chat is already active.
|
||||
@@ -180,10 +183,11 @@ const (
|
||||
|
||||
// SendMessageOptions controls user message insertion with busy-state behavior.
|
||||
type SendMessageOptions struct {
|
||||
ChatID uuid.UUID
|
||||
Content []fantasy.Content
|
||||
ModelConfigID *uuid.UUID
|
||||
BusyBehavior SendMessageBusyBehavior
|
||||
ChatID uuid.UUID
|
||||
Content []fantasy.Content
|
||||
ContentFileIDs map[int]uuid.UUID
|
||||
ModelConfigID *uuid.UUID
|
||||
BusyBehavior SendMessageBusyBehavior
|
||||
}
|
||||
|
||||
// SendMessageResult contains the outcome of user message processing.
|
||||
@@ -199,6 +203,7 @@ type EditMessageOptions struct {
|
||||
ChatID uuid.UUID
|
||||
EditedMessageID int64
|
||||
Content []fantasy.Content
|
||||
ContentFileIDs map[int]uuid.UUID
|
||||
}
|
||||
|
||||
// EditMessageResult contains the updated user message and chat status.
|
||||
@@ -278,7 +283,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C
|
||||
}
|
||||
}
|
||||
|
||||
userContent, err := chatprompt.MarshalContent(opts.InitialUserContent)
|
||||
userContent, err := chatprompt.MarshalContent(opts.InitialUserContent, opts.ContentFileIDs)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("marshal initial user content: %w", err)
|
||||
}
|
||||
@@ -345,7 +350,7 @@ func (p *Server) SendMessage(
|
||||
return SendMessageResult{}, xerrors.Errorf("invalid busy behavior %q", opts.BusyBehavior)
|
||||
}
|
||||
|
||||
content, err := chatprompt.MarshalContent(opts.Content)
|
||||
content, err := chatprompt.MarshalContent(opts.Content, opts.ContentFileIDs)
|
||||
if err != nil {
|
||||
return SendMessageResult{}, xerrors.Errorf("marshal message content: %w", err)
|
||||
}
|
||||
@@ -448,7 +453,7 @@ func (p *Server) EditMessage(
|
||||
return EditMessageResult{}, xerrors.New("content is required")
|
||||
}
|
||||
|
||||
content, err := chatprompt.MarshalContent(opts.Content)
|
||||
content, err := chatprompt.MarshalContent(opts.Content, opts.ContentFileIDs)
|
||||
if err != nil {
|
||||
return EditMessageResult{}, xerrors.Errorf("marshal message content: %w", err)
|
||||
}
|
||||
@@ -1607,6 +1612,25 @@ func (p *Server) subscribeChatControl(
|
||||
return controlCancel
|
||||
}
|
||||
|
||||
// chatFileResolver returns a FileResolver that fetches chat file
|
||||
// content from the database by ID.
|
||||
func (p *Server) chatFileResolver() chatprompt.FileResolver {
|
||||
return func(ctx context.Context, ids []uuid.UUID) (map[uuid.UUID]chatprompt.FileData, error) {
|
||||
files, err := p.db.GetChatFilesByIDs(ctx, ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make(map[uuid.UUID]chatprompt.FileData, len(files))
|
||||
for _, f := range files {
|
||||
result[f.ID] = chatprompt.FileData{
|
||||
Data: f.Data,
|
||||
MediaType: f.Mimetype,
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Server) processChat(ctx context.Context, chat database.Chat) {
|
||||
logger := p.logger.With(slog.F("chat_id", chat.ID))
|
||||
logger.Info(ctx, "processing chat request")
|
||||
@@ -1922,7 +1946,7 @@ func (p *Server) runChat(
|
||||
p.maybeGenerateChatTitle(context.WithoutCancel(ctx), chat, messages, model, providerKeys, logger)
|
||||
}()
|
||||
|
||||
prompt, err := chatprompt.ConvertMessages(messages)
|
||||
prompt, err := chatprompt.ConvertMessagesWithFiles(ctx, messages, p.chatFileResolver())
|
||||
if err != nil {
|
||||
return xerrors.Errorf("build chat prompt: %w", err)
|
||||
}
|
||||
@@ -2064,7 +2088,7 @@ func (p *Server) runChat(
|
||||
}
|
||||
|
||||
if len(assistantBlocks) > 0 {
|
||||
assistantContent, err := chatprompt.MarshalContent(assistantBlocks)
|
||||
assistantContent, err := chatprompt.MarshalContent(assistantBlocks, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -2270,7 +2294,7 @@ func (p *Server) runChat(
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("reload chat messages: %w", err)
|
||||
}
|
||||
reloadedPrompt, err := chatprompt.ConvertMessages(reloadedMsgs)
|
||||
reloadedPrompt, err := chatprompt.ConvertMessagesWithFiles(reloadCtx, reloadedMsgs, p.chatFileResolver())
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("convert reloaded messages: %w", err)
|
||||
}
|
||||
@@ -2363,7 +2387,7 @@ func (p *Server) persistChatContextSummary(
|
||||
ToolName: "chat_summarized",
|
||||
Input: string(args),
|
||||
},
|
||||
})
|
||||
}, nil)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("encode summary tool call: %w", err)
|
||||
}
|
||||
|
||||
@@ -1,12 +1,14 @@
|
||||
package chatprompt
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"charm.land/fantasy"
|
||||
fantasyopenai "charm.land/fantasy/providers/openai"
|
||||
"github.com/google/uuid"
|
||||
"github.com/sqlc-dev/pqtype"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
@@ -16,12 +18,156 @@ import (
|
||||
|
||||
var toolCallIDSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_-]`)
|
||||
|
||||
// FileData holds resolved file content for LLM prompt building.
|
||||
type FileData struct {
|
||||
Data []byte
|
||||
MediaType string
|
||||
}
|
||||
|
||||
// FileResolver fetches file content by ID for LLM prompt building.
|
||||
type FileResolver func(ctx context.Context, ids []uuid.UUID) (map[uuid.UUID]FileData, error)
|
||||
|
||||
// ExtractFileID parses the file_id from a serialized file content
|
||||
// block envelope. Returns uuid.Nil and an error when the block is
|
||||
// not a file-type block or has no file_id.
|
||||
func ExtractFileID(raw json.RawMessage) (uuid.UUID, error) {
|
||||
var envelope struct {
|
||||
Type string `json:"type"`
|
||||
Data struct {
|
||||
FileID string `json:"file_id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &envelope); err != nil {
|
||||
return uuid.Nil, xerrors.Errorf("unmarshal content block: %w", err)
|
||||
}
|
||||
if !strings.EqualFold(envelope.Type, string(fantasy.ContentTypeFile)) {
|
||||
return uuid.Nil, xerrors.Errorf("not a file content block: %s", envelope.Type)
|
||||
}
|
||||
if envelope.Data.FileID == "" {
|
||||
return uuid.Nil, xerrors.New("no file_id")
|
||||
}
|
||||
return uuid.Parse(envelope.Data.FileID)
|
||||
}
|
||||
|
||||
// extractFileIDs scans raw message content for file_id references.
|
||||
// Returns a map of block index to file ID. Returns nil for
|
||||
// non-array content or content with no file references.
|
||||
func extractFileIDs(raw pqtype.NullRawMessage) map[int]uuid.UUID {
|
||||
if !raw.Valid || len(raw.RawMessage) == 0 {
|
||||
return nil
|
||||
}
|
||||
var rawBlocks []json.RawMessage
|
||||
if err := json.Unmarshal(raw.RawMessage, &rawBlocks); err != nil {
|
||||
return nil
|
||||
}
|
||||
var result map[int]uuid.UUID
|
||||
for i, block := range rawBlocks {
|
||||
fid, err := ExtractFileID(block)
|
||||
if err == nil {
|
||||
if result == nil {
|
||||
result = make(map[int]uuid.UUID)
|
||||
}
|
||||
result[i] = fid
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// patchFileContent fills in empty Data on FileContent blocks from
|
||||
// resolved file data. Blocks that already have inline data (backward
|
||||
// compat) or have no resolved data are left unchanged.
|
||||
func patchFileContent(
|
||||
content []fantasy.Content,
|
||||
fileIDs map[int]uuid.UUID,
|
||||
resolved map[uuid.UUID]FileData,
|
||||
) {
|
||||
for blockIdx, fid := range fileIDs {
|
||||
if blockIdx >= len(content) {
|
||||
continue
|
||||
}
|
||||
switch fc := content[blockIdx].(type) {
|
||||
case fantasy.FileContent:
|
||||
if len(fc.Data) > 0 {
|
||||
continue
|
||||
}
|
||||
if data, found := resolved[fid]; found {
|
||||
fc.Data = data.Data
|
||||
content[blockIdx] = fc
|
||||
}
|
||||
case *fantasy.FileContent:
|
||||
if len(fc.Data) > 0 {
|
||||
continue
|
||||
}
|
||||
if data, found := resolved[fid]; found {
|
||||
fc.Data = data.Data
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ConvertMessages converts persisted chat messages into LLM prompt
|
||||
// messages without resolving file references from storage. Inline
|
||||
// file data is preserved when present (backward compat).
|
||||
func ConvertMessages(
|
||||
messages []database.ChatMessage,
|
||||
) ([]fantasy.Message, error) {
|
||||
return ConvertMessagesWithFiles(context.Background(), messages, nil)
|
||||
}
|
||||
|
||||
// ConvertMessagesWithFiles converts persisted chat messages into LLM
|
||||
// prompt messages, resolving file references via the provided
|
||||
// resolver. When resolver is nil, file blocks without inline data
|
||||
// are passed through as-is (same behavior as ConvertMessages).
|
||||
func ConvertMessagesWithFiles(
|
||||
ctx context.Context,
|
||||
messages []database.ChatMessage,
|
||||
resolver FileResolver,
|
||||
) ([]fantasy.Message, error) {
|
||||
// Phase 1: Pre-scan user messages for file_id references.
|
||||
var allFileIDs []uuid.UUID
|
||||
seenFileIDs := make(map[uuid.UUID]struct{})
|
||||
fileIDsByMsg := make(map[int]map[int]uuid.UUID)
|
||||
|
||||
if resolver != nil {
|
||||
for i, msg := range messages {
|
||||
visibility := msg.Visibility
|
||||
if visibility == "" {
|
||||
visibility = database.ChatMessageVisibilityBoth
|
||||
}
|
||||
if visibility != database.ChatMessageVisibilityModel &&
|
||||
visibility != database.ChatMessageVisibilityBoth {
|
||||
continue
|
||||
}
|
||||
if msg.Role != string(fantasy.MessageRoleUser) {
|
||||
continue
|
||||
}
|
||||
fids := extractFileIDs(msg.Content)
|
||||
if len(fids) > 0 {
|
||||
fileIDsByMsg[i] = fids
|
||||
for _, fid := range fids {
|
||||
if _, seen := seenFileIDs[fid]; !seen {
|
||||
seenFileIDs[fid] = struct{}{}
|
||||
allFileIDs = append(allFileIDs, fid)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2: Batch resolve file data.
|
||||
var resolved map[uuid.UUID]FileData
|
||||
if len(allFileIDs) > 0 {
|
||||
var err error
|
||||
resolved, err = resolver(ctx, allFileIDs)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("resolve chat files: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 3: Convert messages, patching file content as needed.
|
||||
prompt := make([]fantasy.Message, 0, len(messages))
|
||||
toolNameByCallID := make(map[string]string)
|
||||
for _, message := range messages {
|
||||
for i, message := range messages {
|
||||
visibility := message.Visibility
|
||||
if visibility == "" {
|
||||
visibility = database.ChatMessageVisibilityBoth
|
||||
@@ -51,6 +197,9 @@ func ConvertMessages(
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if fids, ok := fileIDsByMsg[i]; ok {
|
||||
patchFileContent(content, fids, resolved)
|
||||
}
|
||||
prompt = append(prompt, fantasy.Message{
|
||||
Role: fantasy.MessageRoleUser,
|
||||
Content: ToMessageParts(content),
|
||||
@@ -400,7 +549,10 @@ func ExtractToolCalls(parts []fantasy.MessagePart) []fantasy.ToolCallContent {
|
||||
}
|
||||
|
||||
// MarshalContent encodes message content blocks for persistence.
|
||||
func MarshalContent(blocks []fantasy.Content) (pqtype.NullRawMessage, error) {
|
||||
// fileIDs optionally maps block indices to chat_files IDs, which
|
||||
// are injected into the JSON envelope for file-type blocks so
|
||||
// the reference survives round-trips through storage.
|
||||
func MarshalContent(blocks []fantasy.Content, fileIDs map[int]uuid.UUID) (pqtype.NullRawMessage, error) {
|
||||
if len(blocks) == 0 {
|
||||
return pqtype.NullRawMessage{}, nil
|
||||
}
|
||||
@@ -415,6 +567,16 @@ func MarshalContent(blocks []fantasy.Content) (pqtype.NullRawMessage, error) {
|
||||
err,
|
||||
)
|
||||
}
|
||||
if fid, ok := fileIDs[i]; ok {
|
||||
encoded, err = injectFileID(encoded, fid)
|
||||
if err != nil {
|
||||
return pqtype.NullRawMessage{}, xerrors.Errorf(
|
||||
"inject file_id into content block %d: %w",
|
||||
i,
|
||||
err,
|
||||
)
|
||||
}
|
||||
}
|
||||
encodedBlocks = append(encodedBlocks, encoded)
|
||||
}
|
||||
|
||||
@@ -425,6 +587,27 @@ func MarshalContent(blocks []fantasy.Content) (pqtype.NullRawMessage, error) {
|
||||
return pqtype.NullRawMessage{RawMessage: data, Valid: true}, nil
|
||||
}
|
||||
|
||||
// injectFileID adds a file_id field into the data sub-object of a
|
||||
// serialized content block envelope. This follows the same pattern
|
||||
// as the reasoning title injection in marshalContentBlock.
|
||||
func injectFileID(encoded json.RawMessage, fileID uuid.UUID) (json.RawMessage, error) {
|
||||
var envelope struct {
|
||||
Type string `json:"type"`
|
||||
Data struct {
|
||||
MediaType string `json:"media_type"`
|
||||
Data json.RawMessage `json:"data,omitempty"`
|
||||
FileID string `json:"file_id,omitempty"`
|
||||
ProviderMetadata *json.RawMessage `json:"provider_metadata,omitempty"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(encoded, &envelope); err != nil {
|
||||
return encoded, err
|
||||
}
|
||||
envelope.Data.FileID = fileID.String()
|
||||
envelope.Data.Data = nil // Strip inline data; resolved at LLM dispatch time.
|
||||
return json.Marshal(envelope)
|
||||
}
|
||||
|
||||
// MarshalToolResult encodes a single tool result for persistence as
|
||||
// an opaque JSON blob. The stored shape is
|
||||
// [{"tool_call_id":…,"tool_name":…,"result":…,"is_error":…}].
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
package chatprompt_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/google/uuid"
|
||||
"github.com/sqlc-dev/pqtype"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/chatd/chatprompt"
|
||||
@@ -52,7 +55,7 @@ func TestConvertMessages_NormalizesAssistantToolCallInput(t *testing.T) {
|
||||
ToolName: "execute",
|
||||
Input: tc.input,
|
||||
},
|
||||
})
|
||||
}, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
toolContent, err := chatprompt.MarshalToolResult(
|
||||
@@ -89,3 +92,139 @@ func TestConvertMessages_NormalizesAssistantToolCallInput(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertMessagesWithFiles_ResolvesFileData(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fileID := uuid.New()
|
||||
fileData := []byte("fake-image-bytes")
|
||||
|
||||
// Build a user message with file_id but no inline data, as
|
||||
// would be stored after injectFileID strips the data.
|
||||
rawContent := mustJSON(t, []json.RawMessage{
|
||||
mustJSON(t, map[string]any{
|
||||
"type": "file",
|
||||
"data": map[string]any{
|
||||
"media_type": "image/png",
|
||||
"file_id": fileID.String(),
|
||||
},
|
||||
}),
|
||||
})
|
||||
|
||||
resolver := func(_ context.Context, ids []uuid.UUID) (map[uuid.UUID]chatprompt.FileData, error) {
|
||||
result := make(map[uuid.UUID]chatprompt.FileData)
|
||||
for _, id := range ids {
|
||||
if id == fileID {
|
||||
result[id] = chatprompt.FileData{
|
||||
Data: fileData,
|
||||
MediaType: "image/png",
|
||||
}
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
prompt, err := chatprompt.ConvertMessagesWithFiles(
|
||||
context.Background(),
|
||||
[]database.ChatMessage{
|
||||
{
|
||||
Role: string(fantasy.MessageRoleUser),
|
||||
Visibility: database.ChatMessageVisibilityBoth,
|
||||
Content: pqtype.NullRawMessage{RawMessage: rawContent, Valid: true},
|
||||
},
|
||||
},
|
||||
resolver,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, prompt, 1)
|
||||
require.Equal(t, fantasy.MessageRoleUser, prompt[0].Role)
|
||||
require.Len(t, prompt[0].Content, 1)
|
||||
|
||||
filePart, ok := fantasy.AsMessagePart[fantasy.FilePart](prompt[0].Content[0])
|
||||
require.True(t, ok, "expected FilePart")
|
||||
require.Equal(t, fileData, filePart.Data)
|
||||
require.Equal(t, "image/png", filePart.MediaType)
|
||||
}
|
||||
|
||||
func TestConvertMessagesWithFiles_BackwardCompat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// A message with inline data and a file_id should use the
|
||||
// inline data even when the resolver returns nothing.
|
||||
fileID := uuid.New()
|
||||
inlineData := []byte("inline-image-data")
|
||||
|
||||
rawContent := mustJSON(t, []json.RawMessage{
|
||||
mustJSON(t, map[string]any{
|
||||
"type": "file",
|
||||
"data": map[string]any{
|
||||
"media_type": "image/png",
|
||||
"data": inlineData,
|
||||
"file_id": fileID.String(),
|
||||
},
|
||||
}),
|
||||
})
|
||||
|
||||
prompt, err := chatprompt.ConvertMessagesWithFiles(
|
||||
context.Background(),
|
||||
[]database.ChatMessage{
|
||||
{
|
||||
Role: string(fantasy.MessageRoleUser),
|
||||
Visibility: database.ChatMessageVisibilityBoth,
|
||||
Content: pqtype.NullRawMessage{RawMessage: rawContent, Valid: true},
|
||||
},
|
||||
},
|
||||
nil, // No resolver.
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, prompt, 1)
|
||||
require.Len(t, prompt[0].Content, 1)
|
||||
|
||||
filePart, ok := fantasy.AsMessagePart[fantasy.FilePart](prompt[0].Content[0])
|
||||
require.True(t, ok, "expected FilePart")
|
||||
require.Equal(t, inlineData, filePart.Data)
|
||||
}
|
||||
|
||||
func TestInjectFileID_StripsInlineData(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fileID := uuid.New()
|
||||
imageData := []byte("raw-image-bytes")
|
||||
|
||||
// Marshal a file content block with inline data, then inject
|
||||
// a file_id. The result should have file_id but no data.
|
||||
content, err := chatprompt.MarshalContent([]fantasy.Content{
|
||||
fantasy.FileContent{
|
||||
MediaType: "image/png",
|
||||
Data: imageData,
|
||||
},
|
||||
}, map[int]uuid.UUID{0: fileID})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Parse the stored content to verify shape.
|
||||
var blocks []json.RawMessage
|
||||
require.NoError(t, json.Unmarshal(content.RawMessage, &blocks))
|
||||
require.Len(t, blocks, 1)
|
||||
|
||||
var envelope struct {
|
||||
Type string `json:"type"`
|
||||
Data struct {
|
||||
MediaType string `json:"media_type"`
|
||||
Data *json.RawMessage `json:"data,omitempty"`
|
||||
FileID string `json:"file_id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(blocks[0], &envelope))
|
||||
require.Equal(t, "file", envelope.Type)
|
||||
require.Equal(t, "image/png", envelope.Data.MediaType)
|
||||
require.Equal(t, fileID.String(), envelope.Data.FileID)
|
||||
// Data should be nil (omitted) since injectFileID strips it.
|
||||
require.Nil(t, envelope.Data.Data, "inline data should be stripped")
|
||||
}
|
||||
|
||||
func mustJSON(t *testing.T, v any) json.RawMessage {
|
||||
t.Helper()
|
||||
data, err := json.Marshal(v)
|
||||
require.NoError(t, err)
|
||||
return data
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user