diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index 26a0b003a0c43..a7f437984ca39 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/coderd/database/dbpurge/dbpurge_test.go @@ -2775,27 +2775,6 @@ func TestDeleteOldChatFiles(t *testing.T) { now.Add(-10*24*time.Hour), recentArchivedChat.ID) require.NoError(t, err) - // File F: 31 days old, in BOTH an active chat AND an old archived chat -> should be retained. - fileF := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now.Add(-31*24*time.Hour)) - anotherOldArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) - _, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ - ChatID: anotherOldArchivedChat.ID, - MaxFileLinks: 100, - FileIds: []uuid.UUID{fileF}, - }) - require.NoError(t, err) - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET updated_at = $1 WHERE id = $2", - now.Add(-31*24*time.Hour), anotherOldArchivedChat.ID) - require.NoError(t, err) - - activeChatForF := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now) - _, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ - ChatID: activeChatForF.ID, - MaxFileLinks: 100, - FileIds: []uuid.UUID{fileF}, - }) - require.NoError(t, err) - done := awaitDoTick(ctx, t, clk) closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk)) defer closer.Close() @@ -2806,9 +2785,6 @@ func TestDeleteOldChatFiles(t *testing.T) { _, err = db.GetChatFileByID(ctx, fileE) require.NoError(t, err, "file E in recently archived chat should be retained") - - _, err = db.GetChatFileByID(ctx, fileF) - require.NoError(t, err, "file F in active + old archived chat should be retained") }, }, { diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 6e93c540739ac..abbfda38e5c9b 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -4426,6 +4426,9 @@ ALTER TABLE ONLY chat_diff_statuses ALTER TABLE ONLY chat_file_links ADD CONSTRAINT chat_file_links_chat_id_file_id_key UNIQUE (chat_id, file_id); +ALTER TABLE ONLY chat_file_links + ADD CONSTRAINT chat_file_links_file_id_key UNIQUE (file_id); + ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_pkey PRIMARY KEY (id); @@ -4904,8 +4907,6 @@ CREATE INDEX idx_chat_diff_statuses_url_lower ON chat_diff_statuses USING btree CREATE INDEX idx_chat_file_links_chat_id ON chat_file_links USING btree (chat_id); -CREATE INDEX idx_chat_file_links_file_id ON chat_file_links USING btree (file_id); - CREATE INDEX idx_chat_files_created_at ON chat_files USING btree (created_at); CREATE INDEX idx_chat_files_org ON chat_files USING btree (organization_id); diff --git a/coderd/database/migrations/000591_chat_file_links_unique_file_id.down.sql b/coderd/database/migrations/000591_chat_file_links_unique_file_id.down.sql new file mode 100644 index 0000000000000..a23d0713dca39 --- /dev/null +++ b/coderd/database/migrations/000591_chat_file_links_unique_file_id.down.sql @@ -0,0 +1,4 @@ +CREATE INDEX idx_chat_file_links_file_id ON chat_file_links (file_id); + +ALTER TABLE chat_file_links + DROP CONSTRAINT chat_file_links_file_id_key; diff --git a/coderd/database/migrations/000591_chat_file_links_unique_file_id.up.sql b/coderd/database/migrations/000591_chat_file_links_unique_file_id.up.sql new file mode 100644 index 0000000000000..d52d2a4d3ab0d --- /dev/null +++ b/coderd/database/migrations/000591_chat_file_links_unique_file_id.up.sql @@ -0,0 +1,14 @@ +-- A file belongs to one chat. Drop any duplicate links, keeping the link on +-- the oldest chat, before enforcing it. +DELETE FROM chat_file_links l +USING chat_file_links o, chats lc, chats oc +WHERE o.file_id = l.file_id + AND lc.id = l.chat_id + AND oc.id = o.chat_id + AND (oc.created_at, oc.id) < (lc.created_at, lc.id); + +ALTER TABLE chat_file_links + ADD CONSTRAINT chat_file_links_file_id_key UNIQUE (file_id); + +-- The unique index replaces the purge lookup index. +DROP INDEX idx_chat_file_links_file_id; diff --git a/coderd/database/querier.go b/coderd/database/querier.go index c29f1b7b05e41..1d6216f3b8041 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -1260,8 +1260,11 @@ type sqlcQuerier interface { // time. chatstate calls this in a single query so the staleness check // is atomic and does not depend on the caller's local clock. IsChatHeartbeatStale(ctx context.Context, arg IsChatHeartbeatStaleParams) (bool, error) - // LinkChatFilesAfterLock requires the chat row lock. - // The lock serializes cap checks. The result counts rejected new links. + // LinkChatFilesAfterLock requires the chat row lock. When the batch would + // exceed the cap, the oldest files on the chat are deleted to make room; the + // cascade removes their links. A file links to at most one chat, so no other + // chat can lose a file here. The batch is rejected only when the batch itself + // exceeds the cap. LinkChatFilesAfterLock(ctx context.Context, arg LinkChatFilesAfterLockParams) (int32, error) ListAIBridgeClients(ctx context.Context, arg ListAIBridgeClientsParams) ([]string, error) // Finds all unique AI Bridge interception telemetry summaries combinations diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 637fbdc728a74..ba217c09523ed 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -2116,6 +2116,115 @@ func TestLinkChatFilesDeduplicatesInput(t *testing.T) { require.Equal(t, file.ID, files[0].ID) } +func TestLinkChatFilesEvictsOldest(t *testing.T) { + t.Parallel() + if testing.Short() { + t.SkipNow() + } + + ctx := testutil.Context(t, testutil.WaitMedium) + sqlDB := testSQLDB(t) + err := migrations.Up(sqlDB) + require.NoError(t, err) + db := database.New(sqlDB) + + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + newChat := func() database.Chat { + return dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + } + const maxLinks = 3 + base := dbtime.Now().Add(-time.Hour) + newFiles := func(n int) []uuid.UUID { + ids := make([]uuid.UUID, 0, n) + for i := 0; i < n; i++ { + file, err := db.InsertChatFile(ctx, database.InsertChatFileParams{ + OwnerID: user.ID, + OrganizationID: org.ID, + Name: fmt.Sprintf("file-%d.txt", i), + Mimetype: "text/plain", + Data: []byte("data"), + }) + require.NoError(t, err) + _, err = sqlDB.ExecContext(ctx, + "UPDATE chat_files SET created_at = $1 WHERE id = $2", + base.Add(time.Duration(i)*time.Second), file.ID) + require.NoError(t, err) + ids = append(ids, file.ID) + } + return ids + } + linkedIDs := func(chatID uuid.UUID) []uuid.UUID { + files, err := db.GetChatFileMetadataByChatID(ctx, chatID) + require.NoError(t, err) + ids := make([]uuid.UUID, 0, len(files)) + for _, f := range files { + ids = append(ids, f.ID) + } + return ids + } + + // Linking one file past the cap deletes the oldest file. + chat := newChat() + files := newFiles(maxLinks + 1) + rejected, err := db.LinkChatFiles(ctx, database.LinkChatFilesParams{ + ChatID: chat.ID, + FileIds: files[:maxLinks], + MaxFileLinks: maxLinks, + }) + require.NoError(t, err) + require.Zero(t, rejected) + rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ + ChatID: chat.ID, + FileIds: files[maxLinks:], + MaxFileLinks: maxLinks, + }) + require.NoError(t, err) + require.Zero(t, rejected) + require.Equal(t, files[1:], linkedIDs(chat.ID)) + _, err = db.GetChatFileByID(ctx, files[0]) + require.ErrorIs(t, err, sql.ErrNoRows) + + // A single batch over the cap is rejected and deletes nothing. + chat = newChat() + files = newFiles(maxLinks + 1) + rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ + ChatID: chat.ID, + FileIds: files, + MaxFileLinks: maxLinks, + }) + require.NoError(t, err) + require.EqualValues(t, maxLinks+1, rejected) + require.Empty(t, linkedIDs(chat.ID)) + for _, id := range files { + _, err = db.GetChatFileByID(ctx, id) + require.NoError(t, err) + } + + // A file links to one chat only. + chat = newChat() + files = newFiles(1) + rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ + ChatID: chat.ID, + FileIds: files, + MaxFileLinks: maxLinks, + }) + require.NoError(t, err) + require.Zero(t, rejected) + _, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ + ChatID: newChat().ID, + FileIds: files, + MaxFileLinks: maxLinks, + }) + require.True(t, database.IsUniqueViolation(err, database.UniqueChatFileLinksFileIDKey)) + require.Equal(t, files, linkedIDs(chat.ID)) +} + func TestGetChatFileDataPrefixesByIDs(t *testing.T) { t.Parallel() if testing.Short() { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 700b0fe036643..f277d76621958 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -11106,45 +11106,64 @@ func (q *sqlQuerier) IsChatHeartbeatStale(ctx context.Context, arg IsChatHeartbe } const linkChatFilesAfterLock = `-- name: LinkChatFilesAfterLock :one -WITH current AS ( - SELECT COUNT(*) AS cnt - FROM chat_file_links - WHERE chat_id = $1::uuid +WITH new_links AS ( + SELECT DISTINCT unnest($1::uuid[]) AS file_id ), -new_links AS ( - SELECT DISTINCT $1::uuid AS chat_id, unnest($2::uuid[]) AS file_id +fits AS ( + SELECT (SELECT COUNT(*) FROM new_links) <= $2::int AS ok ), genuinely_new AS ( - SELECT nl.chat_id, nl.file_id - FROM new_links nl + SELECT nl.file_id FROM new_links nl WHERE NOT EXISTS ( SELECT 1 FROM chat_file_links cfl - WHERE cfl.chat_id = nl.chat_id AND cfl.file_id = nl.file_id + WHERE cfl.chat_id = $3::uuid AND cfl.file_id = nl.file_id ) ), +needed AS ( + SELECT GREATEST( + (SELECT COUNT(*) FROM chat_file_links WHERE chat_id = $3::uuid) + + (SELECT COUNT(*) FROM genuinely_new) + - $2::int, 0)::int AS n +), +candidates AS ( + SELECT cf.id + FROM chat_file_links cfl + JOIN chat_files cf ON cf.id = cfl.file_id + WHERE cfl.chat_id = $3::uuid + AND NOT EXISTS (SELECT 1 FROM new_links nl WHERE nl.file_id = cf.id) + ORDER BY cf.created_at ASC, cf.id ASC + LIMIT (SELECT n FROM needed) +), +evicted AS ( + DELETE FROM chat_files cf + USING candidates c + WHERE cf.id = c.id AND (SELECT ok FROM fits) + RETURNING cf.id +), inserted AS ( INSERT INTO chat_file_links (chat_id, file_id) - SELECT gn.chat_id, gn.file_id - FROM genuinely_new gn, current c - WHERE c.cnt + (SELECT COUNT(*) FROM genuinely_new) <= $3::int + SELECT $3::uuid, gn.file_id FROM genuinely_new gn + WHERE (SELECT ok FROM fits) ON CONFLICT (chat_id, file_id) DO NOTHING RETURNING file_id ) -SELECT - (SELECT COUNT(*)::int FROM genuinely_new) - - (SELECT COUNT(*)::int FROM inserted) AS rejected_new_files +SELECT (SELECT COUNT(*)::int FROM genuinely_new) + - (SELECT COUNT(*)::int FROM inserted) AS rejected_new_files ` type LinkChatFilesAfterLockParams struct { - ChatID uuid.UUID `db:"chat_id" json:"chat_id"` FileIds []uuid.UUID `db:"file_ids" json:"file_ids"` MaxFileLinks int32 `db:"max_file_links" json:"max_file_links"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` } -// LinkChatFilesAfterLock requires the chat row lock. -// The lock serializes cap checks. The result counts rejected new links. +// LinkChatFilesAfterLock requires the chat row lock. When the batch would +// exceed the cap, the oldest files on the chat are deleted to make room; the +// cascade removes their links. A file links to at most one chat, so no other +// chat can lose a file here. The batch is rejected only when the batch itself +// exceeds the cap. func (q *sqlQuerier) LinkChatFilesAfterLock(ctx context.Context, arg LinkChatFilesAfterLockParams) (int32, error) { - row := q.db.QueryRowContext(ctx, linkChatFilesAfterLock, arg.ChatID, pq.Array(arg.FileIds), arg.MaxFileLinks) + row := q.db.QueryRowContext(ctx, linkChatFilesAfterLock, pq.Array(arg.FileIds), arg.MaxFileLinks, arg.ChatID) var rejected_new_files int32 err := row.Scan(&rejected_new_files) return rejected_new_files, err diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 616936f3c661f..ed0212a408a40 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -1755,35 +1755,54 @@ WHERE chat_id = @chat_id::uuid ORDER BY source ASC; -- name: LinkChatFilesAfterLock :one --- LinkChatFilesAfterLock requires the chat row lock. --- The lock serializes cap checks. The result counts rejected new links. -WITH current AS ( - SELECT COUNT(*) AS cnt - FROM chat_file_links - WHERE chat_id = @chat_id::uuid +-- LinkChatFilesAfterLock requires the chat row lock. When the batch would +-- exceed the cap, the oldest files on the chat are deleted to make room; the +-- cascade removes their links. A file links to at most one chat, so no other +-- chat can lose a file here. The batch is rejected only when the batch itself +-- exceeds the cap. +WITH new_links AS ( + SELECT DISTINCT unnest(@file_ids::uuid[]) AS file_id ), -new_links AS ( - SELECT DISTINCT @chat_id::uuid AS chat_id, unnest(@file_ids::uuid[]) AS file_id +fits AS ( + SELECT (SELECT COUNT(*) FROM new_links) <= @max_file_links::int AS ok ), genuinely_new AS ( - SELECT nl.chat_id, nl.file_id - FROM new_links nl + SELECT nl.file_id FROM new_links nl WHERE NOT EXISTS ( SELECT 1 FROM chat_file_links cfl - WHERE cfl.chat_id = nl.chat_id AND cfl.file_id = nl.file_id + WHERE cfl.chat_id = @chat_id::uuid AND cfl.file_id = nl.file_id ) ), +needed AS ( + SELECT GREATEST( + (SELECT COUNT(*) FROM chat_file_links WHERE chat_id = @chat_id::uuid) + + (SELECT COUNT(*) FROM genuinely_new) + - @max_file_links::int, 0)::int AS n +), +candidates AS ( + SELECT cf.id + FROM chat_file_links cfl + JOIN chat_files cf ON cf.id = cfl.file_id + WHERE cfl.chat_id = @chat_id::uuid + AND NOT EXISTS (SELECT 1 FROM new_links nl WHERE nl.file_id = cf.id) + ORDER BY cf.created_at ASC, cf.id ASC + LIMIT (SELECT n FROM needed) +), +evicted AS ( + DELETE FROM chat_files cf + USING candidates c + WHERE cf.id = c.id AND (SELECT ok FROM fits) + RETURNING cf.id +), inserted AS ( INSERT INTO chat_file_links (chat_id, file_id) - SELECT gn.chat_id, gn.file_id - FROM genuinely_new gn, current c - WHERE c.cnt + (SELECT COUNT(*) FROM genuinely_new) <= @max_file_links::int + SELECT @chat_id::uuid, gn.file_id FROM genuinely_new gn + WHERE (SELECT ok FROM fits) ON CONFLICT (chat_id, file_id) DO NOTHING RETURNING file_id ) -SELECT - (SELECT COUNT(*)::int FROM genuinely_new) - - (SELECT COUNT(*)::int FROM inserted) AS rejected_new_files; +SELECT (SELECT COUNT(*)::int FROM genuinely_new) + - (SELECT COUNT(*)::int FROM inserted) AS rejected_new_files; -- name: UpdateChatStatus :one WITH updated_chat AS ( diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 9d2390a6468ba..fe3f137552ec4 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -27,6 +27,7 @@ const ( UniqueChatDebugStepsPkey UniqueConstraint = "chat_debug_steps_pkey" // ALTER TABLE ONLY chat_debug_steps ADD CONSTRAINT chat_debug_steps_pkey PRIMARY KEY (id); UniqueChatDiffStatusesPkey UniqueConstraint = "chat_diff_statuses_pkey" // ALTER TABLE ONLY chat_diff_statuses ADD CONSTRAINT chat_diff_statuses_pkey PRIMARY KEY (chat_id); UniqueChatFileLinksChatIDFileIDKey UniqueConstraint = "chat_file_links_chat_id_file_id_key" // ALTER TABLE ONLY chat_file_links ADD CONSTRAINT chat_file_links_chat_id_file_id_key UNIQUE (chat_id, file_id); + UniqueChatFileLinksFileIDKey UniqueConstraint = "chat_file_links_file_id_key" // ALTER TABLE ONLY chat_file_links ADD CONSTRAINT chat_file_links_file_id_key UNIQUE (file_id); UniqueChatFilesPkey UniqueConstraint = "chat_files_pkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_pkey PRIMARY KEY (id); UniqueChatHeartbeatsPkey UniqueConstraint = "chat_heartbeats_pkey" // ALTER TABLE ONLY chat_heartbeats ADD CONSTRAINT chat_heartbeats_pkey PRIMARY KEY (chat_id, runner_id); UniqueChatMessagesPkey UniqueConstraint = "chat_messages_pkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_pkey PRIMARY KEY (id); diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 9f21f1c0adbac..963a70793c893 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -6626,7 +6626,7 @@ func writeChatFileError(ctx context.Context, rw http.ResponseWriter, err error) case errors.Is(err, chatstate.ErrChatFileCapExceeded): httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Chat attachment limit reached.", - Detail: fmt.Sprintf("A chat can reference at most %d attachments. Remove some attachments or start a new chat.", codersdk.MaxChatFileIDs), + Detail: fmt.Sprintf("A message can include at most %d attachments. Remove some attachments and retry.", codersdk.MaxChatFileIDs), }) case errors.Is(err, chatstate.ErrChatFileUnavailable): httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 0dc16fcf35f01..c38e2a5e44382 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -6497,7 +6497,7 @@ func TestGetChat(t *testing.T) { require.Equal(t, "text/markdown", f.MimeType) // Fill up to the cap by inserting more files via the - // chatd DB path, then verify the cap is enforced. + // chatd DB path, then verify the oldest file is evicted. for i := 1; i < codersdk.MaxChatFileIDs; i++ { extra, err := store.InsertChatFile(chatdCtx, database.InsertChatFileParams{ OwnerID: firstUser.UserID, @@ -6520,7 +6520,7 @@ func TestGetChat(t *testing.T) { require.NoError(t, err) require.Len(t, chatResult.Files, codersdk.MaxChatFileIDs) - // Attempt to add one more file — should be rejected (0 rows). + // Adding one more file evicts the oldest one. overflow, err := store.InsertChatFile(chatdCtx, database.InsertChatFileParams{ OwnerID: firstUser.UserID, OrganizationID: firstUser.OrganizationID, @@ -6535,18 +6535,20 @@ func TestGetChat(t *testing.T) { FileIds: []uuid.UUID{overflow.ID}, }) require.NoError(t, err) - require.Equal(t, int32(1), rejected, "cap should reject the 21st file") + require.Equal(t, int32(0), rejected, "linking past the cap should evict, not reject") + chatResult, err = client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.Len(t, chatResult.Files, codersdk.MaxChatFileIDs) + require.NotEqual(t, fileRow.ID, chatResult.Files[0].ID, "the oldest file should be evicted") + require.Equal(t, overflow.ID, chatResult.Files[len(chatResult.Files)-1].ID) - // Re-appending an already-linked ID at cap should succeed - // (dedup means no array growth). + // Re-appending an already-linked ID at cap is a no-op. rejected, err = store.LinkChatFiles(chatdCtx, database.LinkChatFilesParams{ ChatID: chat.ID, MaxFileLinks: int32(codersdk.MaxChatFileIDs), - FileIds: []uuid.UUID{fileRow.ID}, + FileIds: []uuid.UUID{overflow.ID}, }) require.NoError(t, err) - // ON CONFLICT DO NOTHING returns 0 rows when the link - // already exists, which is fine — the file is still linked. require.Equal(t, int32(0), rejected, "dedup of existing ID should be a no-op") // Count should still be exactly MaxChatFileIDs. @@ -9742,7 +9744,7 @@ func TestChatMessageWithFiles(t *testing.T) { require.Equal(t, uploadResp.ID, chatResult.Files[0].ID) }) - t.Run("FileCapExceeded", func(t *testing.T) { + t.Run("FileCapEvictsOldest", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) @@ -9781,45 +9783,30 @@ func TestChatMessageWithFiles(t *testing.T) { {Type: codersdk.ChatInputPartTypeFile, FileID: extraResp.ID}, }, }) - require.Error(t, err) - var sdkErr *codersdk.Error - require.ErrorAs(t, err, &sdkErr) - require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) - require.Contains(t, sdkErr.Message, "attachment limit") - - // getChatMessages reads history before queued messages, so a promotion - // can make one response miss the message in both places. Wait for the - // queue to empty, then read history again because promotion inserts a - // history row. - require.Eventually(t, func() bool { - m, err := client.GetChatMessages(ctx, chat.ID, nil) - return err == nil && len(m.QueuedMessages) == 0 - }, testutil.WaitLong, testutil.IntervalMedium) + require.NoError(t, err, "linking past the cap should evict the oldest file") - messages, err := client.GetChatMessages(ctx, chat.ID, nil) - require.NoError(t, err) - for _, msg := range messages.Messages { - for _, part := range msg.Content { - require.NotContains(t, part.Text, "one too many", "rejected send should not persist a message") - } - } - for _, queued := range messages.QueuedMessages { - for _, part := range queued.Content { - require.NotContains(t, part.Text, "one too many", "rejected send should not queue a message") + chatFileIDs := func() []uuid.UUID { + chatResult, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + ids := make([]uuid.UUID, 0, len(chatResult.Files)) + for _, f := range chatResult.Files { + ids = append(ids, f.ID) } + return ids } - chatResult, err := client.GetChat(ctx, chat.ID) - require.NoError(t, err) - require.Len(t, chatResult.Files, codersdk.MaxChatFileIDs, - "file count should not exceed the cap") + linked := chatFileIDs() + require.Len(t, linked, codersdk.MaxChatFileIDs, "file count should not exceed the cap") + require.Contains(t, linked, extraResp.ID) + require.NotContains(t, linked, fileIDs[0], "the oldest file should be evicted") _, err = client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{ {Type: codersdk.ChatInputPartTypeText, Text: "re-reference existing"}, - {Type: codersdk.ChatInputPartTypeFile, FileID: fileIDs[0]}, + {Type: codersdk.ChatInputPartTypeFile, FileID: fileIDs[1]}, }, }) require.NoError(t, err, "re-referencing an already-linked file must not count against the cap") + require.Equal(t, linked, chatFileIDs(), "re-referencing an already-linked file must not evict anything") }) t.Run("FileCapOnCreate", func(t *testing.T) { @@ -10466,7 +10453,7 @@ func TestPatchChatMessage(t *testing.T) { require.Equal(t, "image/png", f.MimeType) }) - t.Run("CapExceededOnEdit", func(t *testing.T) { + t.Run("CapEvictsOldestOnEdit", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) @@ -10479,9 +10466,11 @@ func TestPatchChatMessage(t *testing.T) { {Type: codersdk.ChatInputPartTypeText, Text: "fill to cap"}, } pngData := append([]byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}, make([]byte, 64)...) + fileIDs := make([]uuid.UUID, 0, codersdk.MaxChatFileIDs) for i := range codersdk.MaxChatFileIDs { up, err := client.UploadChatFile(ctx, firstUser.OrganizationID, "image/png", fmt.Sprintf("cap-%d.png", i), bytes.NewReader(pngData)) require.NoError(t, err) + fileIDs = append(fileIDs, up.ID) parts = append(parts, codersdk.ChatInputPart{Type: codersdk.ChatInputPartTypeFile, FileID: up.ID}) } chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{OrganizationID: firstUser.OrganizationID, Content: parts}) @@ -10500,7 +10489,7 @@ func TestPatchChatMessage(t *testing.T) { } require.NotZero(t, userMessageID) - // Upload one more file and try to link via edit. + // Upload one more file and link it via edit. extra, err := client.UploadChatFile(ctx, firstUser.OrganizationID, "image/png", "one-too-many.png", bytes.NewReader(pngData)) require.NoError(t, err) _, err = client.EditChatMessage(ctx, chat.ID, userMessageID, codersdk.EditChatMessageRequest{ @@ -10509,26 +10498,18 @@ func TestPatchChatMessage(t *testing.T) { {Type: codersdk.ChatInputPartTypeFile, FileID: extra.ID}, }, }) - require.Error(t, err, "edit over the cap should fail") - var sdkErr *codersdk.Error - require.ErrorAs(t, err, &sdkErr) - require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) - require.Contains(t, sdkErr.Message, "attachment limit") + require.NoError(t, err, "edit past the cap should evict the oldest file") - messagesResult, err = client.GetChatMessages(ctx, chat.ID, nil) - require.NoError(t, err) - var found bool - for _, msg := range messagesResult.Messages { - if msg.ID == userMessageID { - found = true - break - } - } - require.True(t, found, "original user message should survive a rejected edit") chatResult, err := client.GetChat(ctx, chat.ID) require.NoError(t, err) require.Len(t, chatResult.Files, codersdk.MaxChatFileIDs, "file count should not exceed the cap") + linked := make([]uuid.UUID, 0, len(chatResult.Files)) + for _, f := range chatResult.Files { + linked = append(linked, f.ID) + } + require.Contains(t, linked, extra.ID) + require.NotContains(t, linked, fileIDs[0], "the oldest file should be evicted") }) t.Run("ArchivedChat", func(t *testing.T) { diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 2a04bbaf6f578..0435b2e8d2bb8 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -47,7 +47,9 @@ There is other data that is held in the database and is associated with a chat, We call it **metadata**. The core state machine concerns itself with **execution state**. As a general guideline, a piece of data is execution state if the core state machine needs it to decide what the next state transition may be, or if it's directly modified by a state transition. For example, a queued message is part of the execution state because it impacts what the next action of the agent loop can be. If the agent loop finishes processing a user message and would otherwise stop, but there's a queued message, the agent loop will start processing the queued message instead. On the other hand, a chat's title does not impact the agent loop at all - it's just a label that helps the user identify the chat. -File links are metadata, but one invariant is enforced at transition time: if a transition persists message content that references uploaded files (chat create, message send, queued send, or message edit), it records the file links in the same transaction. If linking would exceed the per-chat attachment cap, the whole transition is rejected. File retention skips files that are still linked to existing chats, so a persisted message must never reference a file without a link. +File links are metadata, but they are written inside transitions: if a transition persists message content that references uploaded files (chat create, message send, queued send, or message edit), it records the file links in the same transaction. Two invariants are enforced when links are written. A file belongs to at most one chat: attaching a file that another chat already holds is refused in the same way as attaching a file that no longer exists. There is an upper bound on the number of files a chat holds. When a message's files would push the chat over the cap, the oldest files on the chat are deleted to make room, and their links go with them. Files in the same message are never evicted by that message, so only a message that is on its own larger than the cap is rejected. Files created by tools during a run take the same path, so a tool's attachment can evict a user's upload and vice versa. + +Eviction means that a persisted message may reference a file that no longer exists. That's expected: the UI shows the attachment as expired, and when the history is sent to the model, an evicted user upload is replaced with a short placeholder saying the content has expired, while evicted assistant and tool files are dropped. Editing a message that still references an evicted file is refused until the attachment is removed from the edit. If the distinction isn't completely clear to you at this point, don't worry. It should become clearer as you learn more about the core state machine. diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 1025a1cbf9cbc..488f0f6e79330 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -1739,8 +1739,9 @@ func TestMessageFileLinkingCapRollsBack(t *testing.T) { }) require.NoError(t, err) - capFileIDs := make([]uuid.UUID, 0, codersdk.MaxChatFileIDs) - for i := range codersdk.MaxChatFileIDs { + // A single batch over the cap is rejected. + tooMany := []codersdk.ChatMessagePart{codersdk.ChatMessageText("one too many")} + for i := range codersdk.MaxChatFileIDs + 1 { row, err := db.InsertChatFile(ctx, database.InsertChatFileParams{ OwnerID: user.ID, OrganizationID: org.ID, @@ -1749,24 +1750,8 @@ func TestMessageFileLinkingCapRollsBack(t *testing.T) { Data: []byte("png-bytes"), }) require.NoError(t, err) - capFileIDs = append(capFileIDs, row.ID) + tooMany = append(tooMany, codersdk.ChatMessageFile(row.ID, "image/png", row.Name)) } - rejected, err := db.LinkChatFiles(ctx, database.LinkChatFilesParams{ - ChatID: chat.ID, - MaxFileLinks: int32(codersdk.MaxChatFileIDs), - FileIds: capFileIDs, - }) - require.NoError(t, err) - require.Zero(t, rejected) - - extra, err := db.InsertChatFile(ctx, database.InsertChatFileParams{ - OwnerID: user.ID, - OrganizationID: org.ID, - Name: "extra.png", - Mimetype: "image/png", - Data: []byte("png-bytes"), - }) - require.NoError(t, err) chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat.ID, @@ -1780,11 +1765,8 @@ func TestMessageFileLinkingCapRollsBack(t *testing.T) { require.NoError(t, err) _, err = replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{ - codersdk.ChatMessageText("one too many"), - codersdk.ChatMessageFile(extra.ID, "image/png", "extra.png"), - }, + ChatID: chat.ID, + Content: tooMany, }) require.ErrorIs(t, err, chatstate.ErrChatFileCapExceeded) @@ -1796,14 +1778,11 @@ func TestMessageFileLinkingCapRollsBack(t *testing.T) { require.Len(t, messagesAfter, len(messagesBefore), "rejected send must not persist a message") files, err := db.GetChatFileMetadataByChatID(ctx, chat.ID) require.NoError(t, err) - require.Len(t, files, codersdk.MaxChatFileIDs) + require.Empty(t, files, "rejected send must not link files") sendResult, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{ - codersdk.ChatMessageText("re-reference"), - codersdk.ChatMessageFile(capFileIDs[0], "image/png", "cap-0.png"), - }, + ChatID: chat.ID, + Content: tooMany[:2], }) require.NoError(t, err) require.False(t, sendResult.Queued) diff --git a/coderd/x/chatd/chatstate/errors.go b/coderd/x/chatd/chatstate/errors.go index af37b12a0baa6..be1e99087446b 100644 --- a/coderd/x/chatd/chatstate/errors.go +++ b/coderd/x/chatd/chatstate/errors.go @@ -56,7 +56,8 @@ var ( // ErrChatFileCapExceeded reports a [LinkFiles] cap rejection. ErrChatFileCapExceeded = xerrors.New("chat attachment cap exceeded") - // ErrChatFileUnavailable reports a missing file passed to [LinkFiles]. + // ErrChatFileUnavailable reports a file passed to [LinkFiles] that is + // missing or already attached to another chat. ErrChatFileUnavailable = xerrors.New("chat attachment unavailable") // ErrToolResultDuplicate is returned by [Tx.CompleteRequiresAction] diff --git a/coderd/x/chatd/chatstate/files.go b/coderd/x/chatd/chatstate/files.go index e46d7458b6f31..5c7fff4baf0c8 100644 --- a/coderd/x/chatd/chatstate/files.go +++ b/coderd/x/chatd/chatstate/files.go @@ -12,8 +12,9 @@ import ( ) // LinkFiles links files, returning [ErrChatFileCapExceeded] for cap rejections -// and [ErrChatFileUnavailable] for missing files. Use the caller's transaction -// so failures roll back related writes; existing links use no additional slots. +// and [ErrChatFileUnavailable] for files that are missing or attached to +// another chat. Use the caller's transaction so failures roll back related +// writes; existing links use no additional slots. func LinkFiles(ctx context.Context, store database.Store, chatID uuid.UUID, fileIDs []uuid.UUID) error { if len(fileIDs) == 0 { return nil @@ -25,7 +26,8 @@ func LinkFiles(ctx context.Context, store database.Store, chatID uuid.UUID, file }) if err != nil { wrapped := xerrors.Errorf("link chat files: %w", err) - if database.IsForeignKeyViolation(err, database.ForeignKeyChatFileLinksFileID) { + if database.IsForeignKeyViolation(err, database.ForeignKeyChatFileLinksFileID) || + database.IsUniqueViolation(err, database.UniqueChatFileLinksFileIDKey) { return errors.Join(ErrChatFileUnavailable, wrapped) } return wrapped diff --git a/coderd/x/chatd/chatstate/files_test.go b/coderd/x/chatd/chatstate/files_test.go index ea571d9509b6b..b06d98a11f456 100644 --- a/coderd/x/chatd/chatstate/files_test.go +++ b/coderd/x/chatd/chatstate/files_test.go @@ -18,21 +18,32 @@ import ( func TestLinkFilesUnavailable(t *testing.T) { t.Parallel() - ctrl := gomock.NewController(t) - store := dbmock.NewMockStore(ctrl) - chatID := uuid.New() - fileID := uuid.New() - foreignKeyErr := &pq.Error{ - Code: pq.ErrorCode("23503"), - Constraint: string(database.ForeignKeyChatFileLinksFileID), - } - store.EXPECT().LinkChatFiles(gomock.Any(), database.LinkChatFilesParams{ - ChatID: chatID, - MaxFileLinks: int32(codersdk.MaxChatFileIDs), - FileIds: []uuid.UUID{fileID}, - }).Return(int32(0), foreignKeyErr) + for name, dbErr := range map[string]*pq.Error{ + "missing": { + Code: pq.ErrorCode("23503"), + Constraint: string(database.ForeignKeyChatFileLinksFileID), + }, + "linked to another chat": { + Code: pq.ErrorCode("23505"), + Constraint: string(database.UniqueChatFileLinksFileIDKey), + }, + } { + t.Run(name, func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + store := dbmock.NewMockStore(ctrl) + chatID := uuid.New() + fileID := uuid.New() + store.EXPECT().LinkChatFiles(gomock.Any(), database.LinkChatFilesParams{ + ChatID: chatID, + MaxFileLinks: int32(codersdk.MaxChatFileIDs), + FileIds: []uuid.UUID{fileID}, + }).Return(int32(0), dbErr) - err := chatstate.LinkFiles(context.Background(), store, chatID, []uuid.UUID{fileID}) - require.ErrorIs(t, err, chatstate.ErrChatFileUnavailable) - require.ErrorIs(t, err, foreignKeyErr) + err := chatstate.LinkFiles(context.Background(), store, chatID, []uuid.UUID{fileID}) + require.ErrorIs(t, err, chatstate.ErrChatFileUnavailable) + require.ErrorIs(t, err, dbErr) + }) + } } diff --git a/coderd/x/chatd/recording_internal_test.go b/coderd/x/chatd/recording_internal_test.go index 8d8b4d9b259c1..cb6e60c887d3c 100644 --- a/coderd/x/chatd/recording_internal_test.go +++ b/coderd/x/chatd/recording_internal_test.go @@ -3,6 +3,7 @@ package chatd import ( "bytes" "context" + "database/sql" "encoding/json" "fmt" "io" @@ -785,9 +786,9 @@ func TestStopAndStoreRecording_Empty(t *testing.T) { assert.Empty(t, result.recordingFileID, "empty recording should not be stored") } -// TestStopAndStoreRecording_LinkFailureRollsBackInsert verifies that a -// chat-file cap rejection does not leave behind an unlinked recording row. -func TestStopAndStoreRecording_LinkFailureRollsBackInsert(t *testing.T) { +// TestStopAndStoreRecording_EvictsOldestAtCap verifies that storing a +// recording on a chat at the file cap evicts the oldest file. +func TestStopAndStoreRecording_EvictsOldestAtCap(t *testing.T) { t.Parallel() db, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t) @@ -802,8 +803,9 @@ func TestStopAndStoreRecording_LinkFailureRollsBackInsert(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) parent, _ := createParentChildChats(ctx, t, server, user, org, model) + var oldest uuid.UUID for i := range codersdk.MaxChatFileIDs { - insertLinkedChatFile( + id := insertLinkedChatFile( ctx, t, db, @@ -814,6 +816,9 @@ func TestStopAndStoreRecording_LinkFailureRollsBackInsert(t *testing.T) { "text/plain", []byte("existing"), ) + if i == 0 { + oldest = id + } } var beforeCount int @@ -831,12 +836,14 @@ func TestStopAndStoreRecording_LinkFailureRollsBackInsert(t *testing.T) { uuid.NullUUID{UUID: workspace.ID, Valid: true}, ) - assert.Empty(t, result.recordingFileID) + require.NotEmpty(t, result.recordingFileID) assert.Empty(t, result.thumbnailFileID) var afterCount int require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM chat_files").Scan(&afterCount)) - assert.Equal(t, beforeCount, afterCount) + assert.Equal(t, beforeCount, afterCount, "the recording replaces the evicted file") + _, err := db.GetChatFileByID(ctx, oldest) + require.ErrorIs(t, err, sql.ErrNoRows) } // TestStopAndStoreRecording_WithThumbnail verifies that a multipart diff --git a/coderd/x/chatd/store_chat_attachment.go b/coderd/x/chatd/store_chat_attachment.go index 0702e110e2928..2e8bd1f0acb01 100644 --- a/coderd/x/chatd/store_chat_attachment.go +++ b/coderd/x/chatd/store_chat_attachment.go @@ -2,7 +2,6 @@ package chatd import ( "context" - "errors" "github.com/google/uuid" "golang.org/x/xerrors" @@ -11,7 +10,6 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/coderd/x/chatd/chattool" "github.com/coder/coder/v2/coderd/x/chatfiles" - "github.com/coder/coder/v2/codersdk" ) func (p *Server) newStoreChatAttachmentFunc(workspaceCtx *turnWorkspaceContext) chattool.StoreFileFunc { @@ -94,9 +92,6 @@ func storeLinkedChatFileTx( } if err := chatstate.LinkFiles(ctx, tx, chatID, []uuid.UUID{row.ID}); err != nil { - if errors.Is(err, chatstate.ErrChatFileCapExceeded) { - return chattool.AttachmentMetadata{}, xerrors.Errorf("chat already has the maximum of %d linked files", codersdk.MaxChatFileIDs) - } return chattool.AttachmentMetadata{}, err } diff --git a/coderd/x/chatd/store_chat_attachment_internal_test.go b/coderd/x/chatd/store_chat_attachment_internal_test.go index 327b78dacfa2a..34c5d027d4759 100644 --- a/coderd/x/chatd/store_chat_attachment_internal_test.go +++ b/coderd/x/chatd/store_chat_attachment_internal_test.go @@ -212,39 +212,6 @@ func TestStoreChatAttachment_InsertError(t *testing.T) { 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}, - } - - expectStoreChatAttachmentInTx(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() @@ -372,26 +339,21 @@ WHERE datname = current_database() barrierReleased = true require.NoError(t, err) - var successes, capRejections int for range 2 { select { case err := <-attachmentResults: - if err == nil { - successes++ - continue - } - require.ErrorContains(t, err, fmt.Sprintf("chat already has the maximum of %d linked files", codersdk.MaxChatFileIDs)) - capRejections++ + require.NoError(t, err) case <-ctx.Done(): require.Failf(t, "attachment store did not finish", "context ended: %v", ctx.Err()) } } - require.Equal(t, 1, successes) - require.Equal(t, 1, capRejections) + // The second attachment sees the first one under the chat lock and + // evicts the oldest file instead of exceeding the cap. files, err := db.GetChatFileMetadataByChatID(ctx, chat.ID) require.NoError(t, err) require.Len(t, files, codersdk.MaxChatFileIDs) + require.NotEqual(t, "existing-00.txt", files[0].Name) var fileCount int require.NoError(t, rawDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM chat_files").Scan(&fileCount)) diff --git a/codersdk/chats.go b/codersdk/chats.go index 52571b537543d..aab042a946f46 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -27,10 +27,9 @@ import ( // threshold settings. const ChatCompactionThresholdKeyPrefix = "chat_compaction_threshold_pct:" -// MaxChatFileIDs is the maximum number of file IDs that can be -// associated with a single chat. This limit prevents unbounded -// growth in the chat_file_links table. It is easier to raise -// this limit than to lower it. +// MaxChatFileIDs is the number of most recent attachments a chat +// keeps. Linking a new file past this cap deletes the oldest files +// on the chat. A single batch larger than the cap is rejected. const MaxChatFileIDs = 50 // MaxChatFileSizeBytes is the upload-endpoint cap for chat diff --git a/docs/ai-coder/agents/index.md b/docs/ai-coder/agents/index.md index 55ed08bb08f93..3d4acef7b8b19 100644 --- a/docs/ai-coder/agents/index.md +++ b/docs/ai-coder/agents/index.md @@ -133,7 +133,7 @@ direction. Users can attach files to chat messages by pasting from the clipboard, dragging files into the input area, or using the attachment button. Supported types are PNG, JPEG, GIF, and WebP images, plus plain text, Markdown, CSV, JSON, and PDF files. -Each upload can be up to 10 MiB, and a single conversation can reference at most 50 attachments. +Each upload can be up to 10 MiB. A conversation keeps its 50 most recent attachments, and older attachments are removed automatically. Attachments are sent to the model as multimodal content alongside the text prompt. This is useful for sharing screenshots of errors, UI mockups, terminal output, logs, or other context that helps the agent understand the task. diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 586d98366de54..070ac0442eef0 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -6120,10 +6120,9 @@ export const MaxAISpendLimitMicros = 1000000000000; // From codersdk/chats.go /** - * MaxChatFileIDs is the maximum number of file IDs that can be - * associated with a single chat. This limit prevents unbounded - * growth in the chat_file_links table. It is easier to raise - * this limit than to lower it. + * MaxChatFileIDs is the number of most recent attachments a chat + * keeps. Linking a new file past this cap deletes the oldest files + * on the chat. A single batch larger than the cap is rejected. */ export const MaxChatFileIDs = 50; diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index 29575a30c72d3..b3472cd3cb148 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -915,6 +915,7 @@ export const AgentChatPageView: FC = ({ key={agentId} organizationId={organizationId} store={store} + chatFiles={chat.files} initialActiveTurnMaxMessageId={initialActiveTurnMaxMessageId} persistedError={persistedError} hasMoreMessages={hasMoreMessages} diff --git a/site/src/pages/AgentsPage/components/ChatConversation/AttachmentBlocks.tsx b/site/src/pages/AgentsPage/components/ChatConversation/AttachmentBlocks.tsx index f33244ef7bb79..784444e4da407 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/AttachmentBlocks.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/AttachmentBlocks.tsx @@ -6,6 +6,7 @@ import { FileTextIcon, } from "lucide-react"; import { type FC, type ReactNode, useState } from "react"; +import { MaxChatFileIDs } from "#/api/typesGenerated"; import { Spinner } from "#/components/Spinner/Spinner"; import { Tooltip, @@ -220,11 +221,13 @@ const imageAttachmentFailureLabels: AttachmentFailureLabels = { failed: "Image failed to load", }; -const textAttachmentFailureLabels: AttachmentFailureLabels = { +const fileAttachmentFailureLabels: AttachmentFailureLabels = { expired: "Attachment expired", failed: "Attachment failed to load", }; +const expiredAttachmentExplanation = `A chat keeps its ${MaxChatFileIDs} most recent attachments, and older attachments are removed. Attachments that no chat references are deleted after this deployment's retention window.`; + const AttachmentFallbackTile: FC<{ state: AttachmentFailure; labels: AttachmentFailureLabels; @@ -250,14 +253,12 @@ const AttachmentFallbackTile: FC<{ ); // Only surface a tooltip when we have something to add: - // - "expired" explains the retention policy. + // - "expired" explains why the file is gone. // - "failed" with a detail surfaces the API error or network reason. // A bare "failed" (e.g. an inline base64 decode failure, where the // browser exposes nothing useful) stays a plain tile. const tooltipBody = - state.kind === "expired" - ? "Attachments are kept while any chat references them. After all references are removed, they are deleted once they are older than this deployment's retention window." - : state.detail; + state.kind === "expired" ? expiredAttachmentExplanation : state.detail; if (!tooltipBody) { return tile; } @@ -341,7 +342,7 @@ const RemoteTextAttachmentButton: FC<{ return ( ); @@ -350,7 +351,7 @@ const RemoteTextAttachmentButton: FC<{ return ( ); @@ -558,10 +559,21 @@ const FileCard: FC<{ block: FileAttachmentBlock; href: string; }> = ({ block, href }) => { + const { hasExpired } = useFileProbes(); const displayName = getAttachmentDisplayName(block); const downloadName = getAttachmentDownloadName(block); const badgeLabel = getAttachmentBadgeLabel(block); + if (block.file_id !== undefined && hasExpired(block.file_id)) { + return ( + + ); + } + return ( ([ }, ], ["storybook-expired-text", { status: 404, body: "" }], + ["storybook-expired-file", { status: 404, body: "" }], [ "storybook-failed-text", { @@ -268,6 +270,9 @@ const buildStoryArgs = (...messages: TypesGen.ChatMessage[]) => ({ parsedMessages: buildMessages(messages), }); +const buildChatFiles = (...fileIds: string[]): TypesGen.ChatFileMetadata[] => + fileIds.map((id) => ({ ...MockChatFileMetadata, id })); + const buildParsedReadFileEntry = ({ messageId, toolId, @@ -908,11 +913,11 @@ export const UserMessageWithExpiredImage: Story = { ).not.toBeInTheDocument(); expectNoCopyMessageButtonForElement(expiredTile); - // The tooltip explains the retention policy generically so the - // copy survives any operator-chosen retention window. + // The tooltip names the attachment cap and describes retention + // generically so the copy survives any operator-chosen window. await hoverAndExpectTooltip( expiredTile, - /kept while any chat references them/i, + /keeps its 50 most recent attachments/i, ); }, }; @@ -1201,35 +1206,75 @@ export const UserMessageWithTextAttachmentOnly: Story = { }, }; +/** + * A text attachment the chat record no longer lists, referenced before a + * remaining one, renders the placeholder on load without fetching the file. + */ export const UserMessageWithExpiredTextAttachment: Story = { - args: buildStoryArgs( - buildUserMessage({ - text: "This pasted context has expired", - files: [buildTextAttachmentPart("storybook-expired-text")], - }), - ), + args: { + ...buildStoryArgs( + buildUserMessage({ + id: 1, + text: "This pasted context has expired", + files: [buildTextAttachmentPart("storybook-expired-text")], + }), + buildUserMessage({ + id: 2, + text: "This newer context is still available", + files: [buildTextAttachmentPart("storybook-test-text")], + }), + ), + chatFiles: buildChatFiles("storybook-test-text"), + }, play: async ({ canvasElement }) => { const canvas = within(canvasElement); - const textButton = await canvas.findByRole("button", { - name: "View text attachment", - }); - expectNoCopyMessageButtonForElement(textButton); - await userEvent.click(textButton); const expiredTile = await findAttachmentTile(canvas, "Attachment expired"); expect( canvas.getByText("This pasted context has expired"), ).toBeInTheDocument(); expect( - canvas.queryByRole("button", { name: "View text attachment" }), - ).not.toBeInTheDocument(); + canvas.getAllByRole("button", { name: "View text attachment" }), + ).toHaveLength(1); + expectNoCopyMessageButtonForElement(expiredTile); await hoverAndExpectTooltip( expiredTile, - /kept while any chat references them/i, + /keeps its 50 most recent attachments/i, ); }, }; +/** An evicted downloadable file renders the placeholder instead of a dead link. */ +export const UserMessageWithExpiredDownloadableFile: Story = { + args: { + ...buildStoryArgs( + buildUserMessage({ + id: 1, + text: "The attached report has expired.", + files: [ + buildFilePart({ + media_type: "application/pdf", + file_id: "storybook-expired-file", + name: "old-report.pdf", + }), + ], + }), + buildUserMessage({ + id: 2, + text: "The newer report is still available.", + files: [ + buildFilePart({ + media_type: "application/pdf", + file_id: "storybook-current-file", + name: "new-report.pdf", + }), + ], + }), + ), + chatFiles: buildChatFiles("storybook-current-file"), + }, +}; + export const UserMessageWithFailedTextAttachment: Story = { args: buildStoryArgs( buildUserMessage({ diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx index dc18f5d782e7a..5b18210241eb7 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx @@ -36,6 +36,7 @@ import { } from "./liveStatusModel"; import { buildDisplayMessages, + deriveEvictedFileIds, deriveMessageDisplayState, } from "./messageHelpers"; import { getEditableUserMessagePayload } from "./messageParsing"; @@ -399,6 +400,7 @@ const ChatMessageItem = memo<{ interface ConversationTimelineProps { organizationId: string | undefined; parsedMessages: readonly ParsedMessageEntry[]; + chatFiles?: readonly TypesGen.ChatFileMetadata[]; initialActiveTurnMaxMessageId?: number; streamState?: StreamState | null; streamTools?: readonly MergedTool[]; @@ -426,6 +428,7 @@ export const ConversationTimeline = memo( ({ organizationId, parsedMessages, + chatFiles, initialActiveTurnMaxMessageId, streamState, streamTools = [], @@ -450,6 +453,7 @@ export const ConversationTimeline = memo( }; const displayMessages = buildDisplayMessages(parsedMessages); + const evictedFileIds = deriveEvictedFileIds(parsedMessages, chatFiles); const renderRows = assignTimelineRows( displayMessages, Boolean(liveStatus && shouldRenderLiveAssistant(liveStatus)), @@ -551,7 +555,7 @@ export const ConversationTimeline = memo( : undefined; return ( - + {renderRows.map((row) => { if (row.type === "live") { // This row only exists when liveStatus is set. diff --git a/site/src/pages/AgentsPage/components/ChatConversation/FileProbeContext.tsx b/site/src/pages/AgentsPage/components/ChatConversation/FileProbeContext.tsx index 560fd31bd9563..99e230ca6294a 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/FileProbeContext.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/FileProbeContext.tsx @@ -28,7 +28,9 @@ const FileProbeContext = createContext({ setProbeResult: () => {}, }); -export const FileProbeProvider: FC = ({ children }) => { +export const FileProbeProvider: FC< + PropsWithChildren<{ evictedFileIds: ReadonlySet }> +> = ({ evictedFileIds, children }) => { const [expiredFileIds, setExpiredFileIds] = useState>( () => new Set(), ); @@ -42,7 +44,8 @@ export const FileProbeProvider: FC = ({ children }) => { return ( expiredFileIds.has(fileId), + hasExpired: (fileId) => + evictedFileIds.has(fileId) || expiredFileIds.has(fileId), markExpired: (fileId) => { setExpiredFileIds((previous) => { if (previous.has(fileId)) { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/messageHelpers.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/messageHelpers.test.ts index d13bd65ccae0f..afe8ceada58b3 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/messageHelpers.test.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/messageHelpers.test.ts @@ -1,8 +1,12 @@ import { describe, expect, it } from "vitest"; import type * as TypesGen from "#/api/typesGenerated"; -import { MockChatMessage } from "#/testHelpers/chatEntities"; +import { + MockChatFileMetadata, + MockChatMessage, +} from "#/testHelpers/chatEntities"; import { buildDisplayMessages, + deriveEvictedFileIds, deriveMessageDisplayState, } from "./messageHelpers"; import { @@ -555,3 +559,95 @@ describe("buildDisplayMessages", () => { expect(result.map((entry) => entry.message.id)).toEqual([1, 2, 3]); }); }); + +describe("deriveEvictedFileIds", () => { + const fileMessage = (messageID: number, ...fileIds: string[]) => + entry({ + messageID, + role: "user", + content: fileIds.map((fileId) => ({ + type: "file", + media_type: "text/plain", + file_id: fileId, + })), + parsedOverrides: {}, + }); + const recordingMessage = ( + messageID: number, + recordingFileId: string, + thumbnailFileId: string, + ) => + entry({ + messageID, + content: [ + { + type: "tool-result", + tool_name: "wait_agent", + result: { + thumbnail_file_id: thumbnailFileId, + recording_file_id: recordingFileId, + }, + }, + ], + parsedOverrides: {}, + }); + const chatFiles = (...fileIds: string[]) => + fileIds.map((id) => ({ ...MockChatFileMetadata, id })); + + it("reports files referenced before the newest linked file as evicted", () => { + const evicted = deriveEvictedFileIds( + [fileMessage(1, "a"), fileMessage(2, "b", "c"), fileMessage(3, "d")], + chatFiles("c", "d"), + ); + + expect([...evicted]).toEqual(["a", "b"]); + }); + + it("presumes files referenced after the newest linked file are present", () => { + const evicted = deriveEvictedFileIds( + [fileMessage(1, "a"), fileMessage(2, "b"), fileMessage(3, "c")], + chatFiles("a", "b"), + ); + + expect(evicted.size).toBe(0); + }); + + it("reports nothing when the chat record has no files", () => { + expect(deriveEvictedFileIds([fileMessage(1, "a")], undefined).size).toBe(0); + expect(deriveEvictedFileIds([fileMessage(1, "a")], []).size).toBe(0); + }); + + it("ignores inline attachments without a file id", () => { + const inline = entry({ + messageID: 1, + role: "user", + content: [{ type: "file", media_type: "text/plain", data: "aGk=" }], + parsedOverrides: {}, + }); + + const evicted = deriveEvictedFileIds( + [inline, fileMessage(2, "a"), fileMessage(3, "b")], + chatFiles("b"), + ); + + expect([...evicted]).toEqual(["a"]); + }); + + it("orders a recording before its thumbnail", () => { + const evicted = deriveEvictedFileIds( + [recordingMessage(1, "rec-1", "thumb-1"), fileMessage(2, "a")], + chatFiles("thumb-1", "a"), + ); + + expect([...evicted]).toEqual(["rec-1"]); + }); + + it("reports an evicted recording and thumbnail", () => { + const evicted = deriveEvictedFileIds( + [recordingMessage(1, "rec-1", "thumb-1"), fileMessage(2, "a")], + chatFiles("a"), + ); + + expect([...evicted]).toEqual(["rec-1", "thumb-1"]); + }); +}); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/messageHelpers.ts b/site/src/pages/AgentsPage/components/ChatConversation/messageHelpers.ts index d45ee594060e9..8b776630f9902 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/messageHelpers.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/messageHelpers.ts @@ -264,3 +264,51 @@ export const buildDisplayMessages = ( flushReadFileEntries(); return grouped; }; + +const NO_FILE_IDS: ReadonlySet = new Set(); + +/** + * Chat files a message part references, in the order the server linked + * them. A wait_agent result stores its recording before its thumbnail. + */ +const partFileIds = (part: TypesGen.ChatMessagePart): string[] => { + switch (part.type) { + case "file": + return part.file_id ? [part.file_id] : []; + case "tool-result": + return [ + part.result?.recording_file_id, + part.result?.thumbnail_file_id, + ].filter((fileId): fileId is string => Boolean(fileId)); + default: + return []; + } +}; + +/** + * Eviction removes a chat's oldest attachments first, and the chat record + * lists the attachments that remain. An attachment referenced before the + * newest remaining one but absent from the record has been evicted. Later + * references are newer than the record, so a message that lands before the + * next chat refetch is never mistaken for an evicted one. + */ +export const deriveEvictedFileIds = ( + entries: readonly ParsedMessageEntry[], + chatFiles: readonly TypesGen.ChatFileMetadata[] | undefined, +): ReadonlySet => { + if (!chatFiles) { + return NO_FILE_IDS; + } + const linkedFileIds = new Set(chatFiles.map((file) => file.id)); + const referencedFileIds = entries.flatMap(({ message }) => + (message.content ?? []).flatMap(partFileIds), + ); + const newestLinkedIndex = referencedFileIds.findLastIndex((fileId) => + linkedFileIds.has(fileId), + ); + return new Set( + referencedFileIds + .slice(0, Math.max(newestLinkedIndex, 0)) + .filter((fileId) => !linkedFileIds.has(fileId)), + ); +}; diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/RecordingPreview.stories.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/RecordingPreview.stories.tsx index 60f9d64bb6b23..21e1ac635a638 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/RecordingPreview.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/RecordingPreview.stories.tsx @@ -1,5 +1,6 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; import { expect, fireEvent, userEvent, waitFor, within } from "storybook/test"; +import { FileProbeProvider } from "../../ChatConversation/FileProbeContext"; import { RecordingPreview } from "./RecordingPreview"; // Static assets stored in site/.storybook/static/. @@ -90,3 +91,17 @@ export const WithoutThumbnail: Story = { recordingFileId: "rec-id", }, }; + +export const Expired: Story = { + args: { + recordingFileId: "evicted-rec-id", + thumbnailFileId: "evicted-thumb-id", + }, + decorators: [ + (Story) => ( + + + + ), + ], +}; diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/RecordingPreview.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/RecordingPreview.tsx index e9d1c15906c82..94dd8d1711db5 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/RecordingPreview.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/RecordingPreview.tsx @@ -1,9 +1,33 @@ -import { ImageOffIcon, PlayIcon } from "lucide-react"; +import { + ImageOffIcon, + type LucideIcon, + PlayIcon, + VideoOffIcon, +} from "lucide-react"; import type React from "react"; import { useState } from "react"; import { getChatFileURL } from "../../../utils/chatAttachments"; +import { useFileProbes } from "../../ChatConversation/FileProbeContext"; import { VideoLightbox } from "../../VideoLightbox"; -import { DEFAULT_ASPECT, PREVIEW_HEIGHT } from "./previewConstants"; +import { + DEFAULT_ASPECT, + PREVIEW_HEIGHT, + RECORDING_EXPIRED_TEXT, +} from "./previewConstants"; + +const frameClassName = + "relative overflow-hidden rounded-lg border border-solid border-border-default"; +const frameStyle = { aspectRatio: DEFAULT_ASPECT, height: PREVIEW_HEIGHT }; + +const PreviewNotice: React.FC<{ icon: LucideIcon; children: string }> = ({ + icon: Icon, + children, +}) => ( +
+ + {children} +
+); interface RecordingPreviewProps { /** The chat file ID for the MP4 recording. */ @@ -22,7 +46,8 @@ interface RecordingPreviewProps { * Inline recording thumbnail with a play icon overlay. Clicking the * preview opens a full-screen VideoLightbox with native playback * controls. If the thumbnail fails to load, a "Thumbnail unavailable" - * message is shown but the video remains playable. + * message is shown but the video remains playable. A recording the chat + * has evicted renders as an expired notice with no playback control. */ export const RecordingPreview: React.FC = ({ recordingFileId, @@ -30,24 +55,29 @@ export const RecordingPreview: React.FC = ({ src: srcOverride, thumbnailSrc: thumbnailSrcOverride, }) => { + const { hasExpired } = useFileProbes(); const [showLightbox, setShowLightbox] = useState(false); const [thumbnailError, setThumbnailError] = useState(false); // Incremented each time the lightbox opens so the VideoLightbox // component remounts and resets its internal error state. const [lightboxKey, setLightboxKey] = useState(0); + if (hasExpired(recordingFileId)) { + return ( +
+ + {RECORDING_EXPIRED_TEXT} + +
+ ); + } + const videoSrc = srcOverride ?? getChatFileURL(recordingFileId); return ( -
+
{thumbnailError ? ( -
- - Thumbnail unavailable -
+ Thumbnail unavailable ) : thumbnailFileId ? ( = ({ organizationId, store, + chatFiles, persistedError, initialActiveTurnMaxMessageId, hasMoreMessages, @@ -208,6 +210,7 @@ export const ChatPageTimeline: FC = ({