Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit d8e3f4f

Browse files
ethanndicksonethanndickson
authored
feat: evict the oldest attachments at the chat file cap (#29063)
Closes CODAGT-933 Long agent sessions were hitting the 50 attachment cap, after which every upload, screenshot and `attach_file` call was rejected, so people were abandoning otherwise healthy chats. Rather than raise the cap (which just moves the cliff and multiplies the worst-case per-chat `bytea` storage), a chat now keeps its 50 most recent attachments. `LinkChatFilesAfterLock` deletes the oldest `chat_files` rows in the same statement and the cascade drops their links. Files in the incoming batch are never evicted, so the only remaining rejection is a single message with more than 50 files. Nothing in the product ever linked one file to two chats, so rather than guard the eviction against it, this adds `UNIQUE (file_id)` on `chat_file_links`, with the migration dropping any existing duplicates in favour of the oldest chat. Evicted attachments render as the existing "Attachment expired" tile. A `pinned` flag is the obvious follow-up if someone needs to keep a specific file around.
1 parent 0f00f7a commit d8e3f4f

34 files changed

Lines changed: 627 additions & 272 deletions

‎coderd/database/dbpurge/dbpurge_test.go‎

Lines changed: 0 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2775,27 +2775,6 @@ func TestDeleteOldChatFiles(t *testing.T) {
27752775
now.Add(-10*24*time.Hour), recentArchivedChat.ID)
27762776
require.NoError(t, err)
27772777

2778-
// File F: 31 days old, in BOTH an active chat AND an old archived chat -> should be retained.
2779-
fileF := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now.Add(-31*24*time.Hour))
2780-
anotherOldArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour))
2781-
_, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
2782-
ChatID: anotherOldArchivedChat.ID,
2783-
MaxFileLinks: 100,
2784-
FileIds: []uuid.UUID{fileF},
2785-
})
2786-
require.NoError(t, err)
2787-
_, err = rawDB.ExecContext(ctx, "UPDATE chats SET updated_at = $1 WHERE id = $2",
2788-
now.Add(-31*24*time.Hour), anotherOldArchivedChat.ID)
2789-
require.NoError(t, err)
2790-
2791-
activeChatForF := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now)
2792-
_, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
2793-
ChatID: activeChatForF.ID,
2794-
MaxFileLinks: 100,
2795-
FileIds: []uuid.UUID{fileF},
2796-
})
2797-
require.NoError(t, err)
2798-
27992778
done := awaitDoTick(ctx, t, clk)
28002779
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
28012780
defer closer.Close()
@@ -2806,9 +2785,6 @@ func TestDeleteOldChatFiles(t *testing.T) {
28062785

28072786
_, err = db.GetChatFileByID(ctx, fileE)
28082787
require.NoError(t, err, "file E in recently archived chat should be retained")
2809-
2810-
_, err = db.GetChatFileByID(ctx, fileF)
2811-
require.NoError(t, err, "file F in active + old archived chat should be retained")
28122788
},
28132789
},
28142790
{

‎coderd/database/dump.sql‎

Lines changed: 3 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000591_chat_file_links_unique_file_id.down.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
CREATE INDEX idx_chat_file_links_file_id ON chat_file_links (file_id);
2+
3+
ALTER TABLE chat_file_links
4+
DROP CONSTRAINT chat_file_links_file_id_key;

‎coderd/database/migrations/000591_chat_file_links_unique_file_id.up.sql‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
-- A file belongs to one chat. Drop any duplicate links, keeping the link on
2+
-- the oldest chat, before enforcing it.
3+
DELETE FROM chat_file_links l
4+
USING chat_file_links o, chats lc, chats oc
5+
WHERE o.file_id = l.file_id
6+
AND lc.id = l.chat_id
7+
AND oc.id = o.chat_id
8+
AND (oc.created_at, oc.id) < (lc.created_at, lc.id);
9+
10+
ALTER TABLE chat_file_links
11+
ADD CONSTRAINT chat_file_links_file_id_key UNIQUE (file_id);
12+
13+
-- The unique index replaces the purge lookup index.
14+
DROP INDEX idx_chat_file_links_file_id;

‎coderd/database/querier.go‎

Lines changed: 5 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier_test.go‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2116,6 +2116,115 @@ func TestLinkChatFilesDeduplicatesInput(t *testing.T) {
21162116
require.Equal(t, file.ID, files[0].ID)
21172117
}
21182118

2119+
func TestLinkChatFilesEvictsOldest(t *testing.T) {
2120+
t.Parallel()
2121+
if testing.Short() {
2122+
t.SkipNow()
2123+
}
2124+
2125+
ctx := testutil.Context(t, testutil.WaitMedium)
2126+
sqlDB := testSQLDB(t)
2127+
err := migrations.Up(sqlDB)
2128+
require.NoError(t, err)
2129+
db := database.New(sqlDB)
2130+
2131+
user := dbgen.User(t, db, database.User{})
2132+
org := dbgen.Organization(t, db, database.Organization{})
2133+
model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{})
2134+
newChat := func() database.Chat {
2135+
return dbgen.Chat(t, db, database.Chat{
2136+
OrganizationID: org.ID,
2137+
OwnerID: user.ID,
2138+
LastModelConfigID: model.ID,
2139+
})
2140+
}
2141+
const maxLinks = 3
2142+
base := dbtime.Now().Add(-time.Hour)
2143+
newFiles := func(n int) []uuid.UUID {
2144+
ids := make([]uuid.UUID, 0, n)
2145+
for i := 0; i < n; i++ {
2146+
file, err := db.InsertChatFile(ctx, database.InsertChatFileParams{
2147+
OwnerID: user.ID,
2148+
OrganizationID: org.ID,
2149+
Name: fmt.Sprintf("file-%d.txt", i),
2150+
Mimetype: "text/plain",
2151+
Data: []byte("data"),
2152+
})
2153+
require.NoError(t, err)
2154+
_, err = sqlDB.ExecContext(ctx,
2155+
"UPDATE chat_files SET created_at = $1 WHERE id = $2",
2156+
base.Add(time.Duration(i)*time.Second), file.ID)
2157+
require.NoError(t, err)
2158+
ids = append(ids, file.ID)
2159+
}
2160+
return ids
2161+
}
2162+
linkedIDs := func(chatID uuid.UUID) []uuid.UUID {
2163+
files, err := db.GetChatFileMetadataByChatID(ctx, chatID)
2164+
require.NoError(t, err)
2165+
ids := make([]uuid.UUID, 0, len(files))
2166+
for _, f := range files {
2167+
ids = append(ids, f.ID)
2168+
}
2169+
return ids
2170+
}
2171+
2172+
// Linking one file past the cap deletes the oldest file.
2173+
chat := newChat()
2174+
files := newFiles(maxLinks + 1)
2175+
rejected, err := db.LinkChatFiles(ctx, database.LinkChatFilesParams{
2176+
ChatID: chat.ID,
2177+
FileIds: files[:maxLinks],
2178+
MaxFileLinks: maxLinks,
2179+
})
2180+
require.NoError(t, err)
2181+
require.Zero(t, rejected)
2182+
rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
2183+
ChatID: chat.ID,
2184+
FileIds: files[maxLinks:],
2185+
MaxFileLinks: maxLinks,
2186+
})
2187+
require.NoError(t, err)
2188+
require.Zero(t, rejected)
2189+
require.Equal(t, files[1:], linkedIDs(chat.ID))
2190+
_, err = db.GetChatFileByID(ctx, files[0])
2191+
require.ErrorIs(t, err, sql.ErrNoRows)
2192+
2193+
// A single batch over the cap is rejected and deletes nothing.
2194+
chat = newChat()
2195+
files = newFiles(maxLinks + 1)
2196+
rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
2197+
ChatID: chat.ID,
2198+
FileIds: files,
2199+
MaxFileLinks: maxLinks,
2200+
})
2201+
require.NoError(t, err)
2202+
require.EqualValues(t, maxLinks+1, rejected)
2203+
require.Empty(t, linkedIDs(chat.ID))
2204+
for _, id := range files {
2205+
_, err = db.GetChatFileByID(ctx, id)
2206+
require.NoError(t, err)
2207+
}
2208+
2209+
// A file links to one chat only.
2210+
chat = newChat()
2211+
files = newFiles(1)
2212+
rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
2213+
ChatID: chat.ID,
2214+
FileIds: files,
2215+
MaxFileLinks: maxLinks,
2216+
})
2217+
require.NoError(t, err)
2218+
require.Zero(t, rejected)
2219+
_, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
2220+
ChatID: newChat().ID,
2221+
FileIds: files,
2222+
MaxFileLinks: maxLinks,
2223+
})
2224+
require.True(t, database.IsUniqueViolation(err, database.UniqueChatFileLinksFileIDKey))
2225+
require.Equal(t, files, linkedIDs(chat.ID))
2226+
}
2227+
21192228
func TestGetChatFileDataPrefixesByIDs(t *testing.T) {
21202229
t.Parallel()
21212230
if testing.Short() {

‎coderd/database/queries.sql.go‎

Lines changed: 38 additions & 19 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/chats.sql‎

Lines changed: 36 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1755,35 +1755,54 @@ WHERE chat_id = @chat_id::uuid
17551755
ORDER BY source ASC;
17561756

17571757
-- name: LinkChatFilesAfterLock :one
1758-
-- LinkChatFilesAfterLock requires the chat row lock.
1759-
-- The lock serializes cap checks. The result counts rejected new links.
1760-
WITH current AS (
1761-
SELECT COUNT(*) AS cnt
1762-
FROM chat_file_links
1763-
WHERE chat_id = @chat_id::uuid
1758+
-- LinkChatFilesAfterLock requires the chat row lock. When the batch would
1759+
-- exceed the cap, the oldest files on the chat are deleted to make room; the
1760+
-- cascade removes their links. A file links to at most one chat, so no other
1761+
-- chat can lose a file here. The batch is rejected only when the batch itself
1762+
-- exceeds the cap.
1763+
WITH new_links AS (
1764+
SELECT DISTINCT unnest(@file_ids::uuid[]) AS file_id
17641765
),
1765-
new_links AS (
1766-
SELECT DISTINCT @chat_id::uuid AS chat_id, unnest(@file_ids::uuid[]) AS file_id
1766+
fits AS (
1767+
SELECT (SELECT COUNT(*) FROM new_links) <= @max_file_links::int AS ok
17671768
),
17681769
genuinely_new AS (
1769-
SELECT nl.chat_id, nl.file_id
1770-
FROM new_links nl
1770+
SELECT nl.file_id FROM new_links nl
17711771
WHERE NOT EXISTS (
17721772
SELECT 1 FROM chat_file_links cfl
1773-
WHERE cfl.chat_id = nl.chat_id AND cfl.file_id = nl.file_id
1773+
WHERE cfl.chat_id = @chat_id::uuid AND cfl.file_id = nl.file_id
17741774
)
17751775
),
1776+
needed AS (
1777+
SELECT GREATEST(
1778+
(SELECT COUNT(*) FROM chat_file_links WHERE chat_id = @chat_id::uuid)
1779+
+ (SELECT COUNT(*) FROM genuinely_new)
1780+
- @max_file_links::int, 0)::int AS n
1781+
),
1782+
candidates AS (
1783+
SELECT cf.id
1784+
FROM chat_file_links cfl
1785+
JOIN chat_files cf ON cf.id = cfl.file_id
1786+
WHERE cfl.chat_id = @chat_id::uuid
1787+
AND NOT EXISTS (SELECT 1 FROM new_links nl WHERE nl.file_id = cf.id)
1788+
ORDER BY cf.created_at ASC, cf.id ASC
1789+
LIMIT (SELECT n FROM needed)
1790+
),
1791+
evicted AS (
1792+
DELETE FROM chat_files cf
1793+
USING candidates c
1794+
WHERE cf.id = c.id AND (SELECT ok FROM fits)
1795+
RETURNING cf.id
1796+
),
17761797
inserted AS (
17771798
INSERT INTO chat_file_links (chat_id, file_id)
1778-
SELECT gn.chat_id, gn.file_id
1779-
FROM genuinely_new gn, current c
1780-
WHERE c.cnt + (SELECT COUNT(*) FROM genuinely_new) <= @max_file_links::int
1799+
SELECT @chat_id::uuid, gn.file_id FROM genuinely_new gn
1800+
WHERE (SELECT ok FROM fits)
17811801
ON CONFLICT (chat_id, file_id) DO NOTHING
17821802
RETURNING file_id
17831803
)
1784-
SELECT
1785-
(SELECT COUNT(*)::int FROM genuinely_new) -
1786-
(SELECT COUNT(*)::int FROM inserted) AS rejected_new_files;
1804+
SELECT (SELECT COUNT(*)::int FROM genuinely_new)
1805+
- (SELECT COUNT(*)::int FROM inserted) AS rejected_new_files;
17871806

17881807
-- name: UpdateChatStatus :one
17891808
WITH updated_chat AS (

‎coderd/database/unique_constraint.go‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/exp_chats.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6626,7 +6626,7 @@ func writeChatFileError(ctx context.Context, rw http.ResponseWriter, err error)
66266626
case errors.Is(err, chatstate.ErrChatFileCapExceeded):
66276627
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
66286628
Message: "Chat attachment limit reached.",
6629-
Detail: fmt.Sprintf("A chat can reference at most %d attachments. Remove some attachments or start a new chat.", codersdk.MaxChatFileIDs),
6629+
Detail: fmt.Sprintf("A message can include at most %d attachments. Remove some attachments and retry.", codersdk.MaxChatFileIDs),
66306630
})
66316631
case errors.Is(err, chatstate.ErrChatFileUnavailable):
66326632
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{

0 commit comments

Comments
 (0)