mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
Ethan
68808c015e
Agents can now attach any file type as a downloadable chat artifact, where previously the stored-file allowlist rejected types like `.zip`. The reason arbitrary types were blocked is that a single media-type list (`codersdk.AllChatAttachmentMediaTypes`) was doing three different jobs at once: gating what users may upload as prompt input, deciding what is safe to render inline in the browser, and admitting what the agent's `attach_file` could store. Because the agent storage path reused that same list as an admission gate, any artifact outside it was rejected even though agent artifacts are only ever downloaded by the user and are never forwarded to the model, so the prompt-input and inline-render constraints did not actually apply to them. This splits those concerns. `PrepareStoredFile` now only normalizes the name and classifies the bytes, and the prompt-input allowlist is enforced inline at `postChatFile` instead, which is the correct layer for user-provided input. User uploads are unchanged and still limited to the allowed prompt-input media types, and unsafe or unknown types remain download-only because `IsInlineRenderableStoredMediaType` still refuses to render them inline. Model replay is also unchanged: assistant and tool attachments are never forwarded to the LLM. Closes CODAGT-654
290 lines
10 KiB
Go
290 lines
10 KiB
Go
package chatd
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/mock/gomock"
|
|
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbmock"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
|
|
"github.com/coder/coder/v2/codersdk"
|
|
)
|
|
|
|
func TestStoreChatAttachment_Success(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
tx := dbmock.NewMockStore(ctrl)
|
|
server := &Server{db: db}
|
|
|
|
chatID := uuid.New()
|
|
ownerID := uuid.New()
|
|
workspaceID := uuid.New()
|
|
orgID := uuid.New()
|
|
fileID := uuid.New()
|
|
chatSnapshot := database.Chat{
|
|
ID: chatID,
|
|
OwnerID: ownerID,
|
|
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
|
}
|
|
|
|
expectStoreChatAttachmentTx(t, db, tx)
|
|
tx.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{ID: workspaceID, OrganizationID: orgID}, nil)
|
|
tx.EXPECT().InsertChatFile(gomock.Any(), gomock.AssignableToTypeOf(database.InsertChatFileParams{})).DoAndReturn(
|
|
func(_ context.Context, arg database.InsertChatFileParams) (database.InsertChatFileRow, error) {
|
|
require.Equal(t, ownerID, arg.OwnerID)
|
|
require.Equal(t, orgID, arg.OrganizationID)
|
|
require.Equal(t, "build.log", arg.Name)
|
|
require.Equal(t, "text/plain", arg.Mimetype)
|
|
require.Equal(t, []byte("build output"), arg.Data)
|
|
return database.InsertChatFileRow{ID: fileID}, nil
|
|
},
|
|
)
|
|
tx.EXPECT().LinkChatFiles(gomock.Any(), database.LinkChatFilesParams{
|
|
ChatID: chatID,
|
|
MaxFileLinks: int32(codersdk.MaxChatFileIDs),
|
|
FileIds: []uuid.UUID{fileID},
|
|
}).Return(int32(0), nil)
|
|
|
|
attachment, err := server.storeChatAttachment(context.Background(), chatSnapshot, "build.log", "build.log", []byte("build output"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, chattool.AttachmentMetadata{
|
|
FileID: fileID,
|
|
MediaType: "text/plain",
|
|
Name: "build.log",
|
|
}, attachment)
|
|
}
|
|
|
|
func TestStoreChatAttachment_UsesDetectNameForClassification(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
tx := dbmock.NewMockStore(ctrl)
|
|
server := &Server{db: db}
|
|
|
|
chatID := uuid.New()
|
|
ownerID := uuid.New()
|
|
workspaceID := uuid.New()
|
|
orgID := uuid.New()
|
|
fileID := uuid.New()
|
|
chatSnapshot := database.Chat{
|
|
ID: chatID,
|
|
OwnerID: ownerID,
|
|
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
|
}
|
|
|
|
expectStoreChatAttachmentTx(t, db, tx)
|
|
tx.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{ID: workspaceID, OrganizationID: orgID}, nil)
|
|
tx.EXPECT().InsertChatFile(gomock.Any(), gomock.AssignableToTypeOf(database.InsertChatFileParams{})).DoAndReturn(
|
|
func(_ context.Context, arg database.InsertChatFileParams) (database.InsertChatFileRow, error) {
|
|
require.Equal(t, "payload.txt", arg.Name)
|
|
require.Equal(t, "application/json", arg.Mimetype)
|
|
return database.InsertChatFileRow{ID: fileID}, nil
|
|
},
|
|
)
|
|
tx.EXPECT().LinkChatFiles(gomock.Any(), database.LinkChatFilesParams{
|
|
ChatID: chatID,
|
|
MaxFileLinks: int32(codersdk.MaxChatFileIDs),
|
|
FileIds: []uuid.UUID{fileID},
|
|
}).Return(int32(0), nil)
|
|
|
|
attachment, err := server.storeChatAttachment(context.Background(), chatSnapshot, "payload.txt", "report.json", []byte(`{"ok":true}`))
|
|
require.NoError(t, err)
|
|
require.Equal(t, "payload.txt", attachment.Name)
|
|
require.Equal(t, "application/json", attachment.MediaType)
|
|
}
|
|
|
|
func TestStoreChatAttachment_AllowsUnsupportedPromptInputType(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
tx := dbmock.NewMockStore(ctrl)
|
|
server := &Server{db: db}
|
|
|
|
chatID := uuid.New()
|
|
ownerID := uuid.New()
|
|
workspaceID := uuid.New()
|
|
orgID := uuid.New()
|
|
fileID := uuid.New()
|
|
chatSnapshot := database.Chat{
|
|
ID: chatID,
|
|
OwnerID: ownerID,
|
|
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
|
}
|
|
data := []byte(`<svg xmlns="http://www.w3.org/2000/svg"><rect/></svg>`)
|
|
|
|
expectStoreChatAttachmentTx(t, db, tx)
|
|
tx.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{ID: workspaceID, OrganizationID: orgID}, nil)
|
|
tx.EXPECT().InsertChatFile(gomock.Any(), gomock.AssignableToTypeOf(database.InsertChatFileParams{})).DoAndReturn(
|
|
func(_ context.Context, arg database.InsertChatFileParams) (database.InsertChatFileRow, error) {
|
|
require.Equal(t, ownerID, arg.OwnerID)
|
|
require.Equal(t, orgID, arg.OrganizationID)
|
|
require.Equal(t, "evil.svg", arg.Name)
|
|
require.Equal(t, "image/svg+xml", arg.Mimetype)
|
|
require.Equal(t, data, arg.Data)
|
|
return database.InsertChatFileRow{ID: fileID}, nil
|
|
},
|
|
)
|
|
tx.EXPECT().LinkChatFiles(gomock.Any(), database.LinkChatFilesParams{
|
|
ChatID: chatID,
|
|
MaxFileLinks: int32(codersdk.MaxChatFileIDs),
|
|
FileIds: []uuid.UUID{fileID},
|
|
}).Return(int32(0), nil)
|
|
|
|
attachment, err := server.storeChatAttachment(context.Background(), chatSnapshot, "evil.svg", "evil.svg", data)
|
|
require.NoError(t, err)
|
|
require.Equal(t, chattool.AttachmentMetadata{
|
|
FileID: fileID,
|
|
MediaType: "image/svg+xml",
|
|
Name: "evil.svg",
|
|
}, attachment)
|
|
}
|
|
|
|
func TestStoreChatAttachment_NoWorkspace(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
server := &Server{db: db}
|
|
|
|
attachment, err := server.storeChatAttachment(context.Background(), database.Chat{}, "build.log", "build.log", []byte("build output"))
|
|
require.ErrorContains(t, err, "no workspace is associated")
|
|
require.Equal(t, chattool.AttachmentMetadata{}, attachment)
|
|
}
|
|
|
|
func TestStoreChatAttachment_WorkspaceLookupError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
tx := dbmock.NewMockStore(ctrl)
|
|
server := &Server{db: db}
|
|
|
|
workspaceID := uuid.New()
|
|
chatSnapshot := database.Chat{
|
|
ID: uuid.New(),
|
|
OwnerID: uuid.New(),
|
|
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
|
}
|
|
|
|
expectStoreChatAttachmentTx(t, db, tx)
|
|
tx.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{}, context.DeadlineExceeded)
|
|
|
|
attachment, err := server.storeChatAttachment(context.Background(), chatSnapshot, "build.log", "build.log", []byte("build output"))
|
|
require.ErrorContains(t, err, "resolve workspace")
|
|
require.ErrorIs(t, err, context.DeadlineExceeded)
|
|
require.Equal(t, chattool.AttachmentMetadata{}, attachment)
|
|
}
|
|
|
|
func TestStoreChatAttachment_InsertError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
tx := dbmock.NewMockStore(ctrl)
|
|
server := &Server{db: db}
|
|
|
|
workspaceID := uuid.New()
|
|
chatSnapshot := database.Chat{
|
|
ID: uuid.New(),
|
|
OwnerID: uuid.New(),
|
|
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
|
}
|
|
|
|
expectStoreChatAttachmentTx(t, db, tx)
|
|
tx.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{ID: workspaceID, OrganizationID: uuid.New()}, nil)
|
|
tx.EXPECT().InsertChatFile(gomock.Any(), gomock.Any()).Return(database.InsertChatFileRow{}, context.DeadlineExceeded)
|
|
|
|
attachment, err := server.storeChatAttachment(context.Background(), chatSnapshot, "build.log", "build.log", []byte("build output"))
|
|
require.ErrorContains(t, err, "insert chat file")
|
|
require.ErrorIs(t, err, context.DeadlineExceeded)
|
|
require.Equal(t, chattool.AttachmentMetadata{}, attachment)
|
|
}
|
|
|
|
func TestStoreChatAttachment_StrictCapError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
tx := dbmock.NewMockStore(ctrl)
|
|
server := &Server{db: db}
|
|
|
|
chatID := uuid.New()
|
|
ownerID := uuid.New()
|
|
workspaceID := uuid.New()
|
|
orgID := uuid.New()
|
|
fileID := uuid.New()
|
|
chatSnapshot := database.Chat{
|
|
ID: chatID,
|
|
OwnerID: ownerID,
|
|
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
|
}
|
|
|
|
expectStoreChatAttachmentTx(t, db, tx)
|
|
tx.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{ID: workspaceID, OrganizationID: orgID}, nil)
|
|
tx.EXPECT().InsertChatFile(gomock.Any(), gomock.AssignableToTypeOf(database.InsertChatFileParams{})).Return(database.InsertChatFileRow{ID: fileID}, nil)
|
|
tx.EXPECT().LinkChatFiles(gomock.Any(), database.LinkChatFilesParams{
|
|
ChatID: chatID,
|
|
MaxFileLinks: int32(codersdk.MaxChatFileIDs),
|
|
FileIds: []uuid.UUID{fileID},
|
|
}).Return(int32(1), nil)
|
|
|
|
attachment, err := server.storeChatAttachment(context.Background(), chatSnapshot, "build.log", "build.log", []byte("build output"))
|
|
require.ErrorContains(t, err, fmt.Sprintf("chat already has the maximum of %d linked files", codersdk.MaxChatFileIDs))
|
|
require.Equal(t, chattool.AttachmentMetadata{}, attachment)
|
|
}
|
|
|
|
func TestStoreChatAttachment_LinkError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
tx := dbmock.NewMockStore(ctrl)
|
|
server := &Server{db: db}
|
|
|
|
chatID := uuid.New()
|
|
ownerID := uuid.New()
|
|
workspaceID := uuid.New()
|
|
orgID := uuid.New()
|
|
fileID := uuid.New()
|
|
chatSnapshot := database.Chat{
|
|
ID: chatID,
|
|
OwnerID: ownerID,
|
|
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
|
}
|
|
|
|
expectStoreChatAttachmentTx(t, db, tx)
|
|
tx.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{ID: workspaceID, OrganizationID: orgID}, nil)
|
|
tx.EXPECT().InsertChatFile(gomock.Any(), gomock.Any()).Return(database.InsertChatFileRow{ID: fileID}, nil)
|
|
tx.EXPECT().LinkChatFiles(gomock.Any(), database.LinkChatFilesParams{
|
|
ChatID: chatID,
|
|
MaxFileLinks: int32(codersdk.MaxChatFileIDs),
|
|
FileIds: []uuid.UUID{fileID},
|
|
}).Return(int32(0), context.DeadlineExceeded)
|
|
|
|
attachment, err := server.storeChatAttachment(context.Background(), chatSnapshot, "build.log", "build.log", []byte("build output"))
|
|
require.ErrorContains(t, err, "link chat file")
|
|
require.ErrorIs(t, err, context.DeadlineExceeded)
|
|
require.Equal(t, chattool.AttachmentMetadata{}, attachment)
|
|
}
|
|
|
|
func expectStoreChatAttachmentTx(t *testing.T, db, tx *dbmock.MockStore) {
|
|
t.Helper()
|
|
|
|
db.EXPECT().InTx(gomock.Any(), gomock.AssignableToTypeOf(&database.TxOptions{})).DoAndReturn(
|
|
func(fn func(database.Store) error, opts *database.TxOptions) error {
|
|
require.NotNil(t, opts)
|
|
require.Equal(t, "store_chat_attachment", opts.TxIdentifier)
|
|
return fn(tx)
|
|
},
|
|
)
|
|
}
|