From 955e2ad611e425573b3550ba76fbdc9fb591ec9e Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 28 Jul 2026 21:17:39 +0000 Subject: [PATCH 1/3] fix(coderd/database): order the chat prompt query and its boundary by id The prompt query led with created_at while selecting its compaction boundary by created_at and then applying that boundary with an id comparison. Because created_at is the transaction start time and is shared across an insert batch, the prompt could present a tool result before the assistant message that requested it, and could retain a stale compressed summary as the boundary. --- coderd/database/querier.go | 3 ++ coderd/database/querier_test.go | 47 ++++++++++++++++++++++++++++++- coderd/database/queries.sql.go | 5 ++-- coderd/database/queries/chats.sql | 5 ++-- 4 files changed, 55 insertions(+), 5 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 0eb2f1d78d9..694c2918ae3 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -469,6 +469,9 @@ type sqlcQuerier interface { GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedParams) ([]ChatMessage, error) // Stream deltas and reset snapshots must use the same message order. GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) + // Ordered by id throughout: the boundary row below is compared with id, and the + // prompt must present roles in append order so tool results follow the assistant + // message that requested them. GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error) GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index c5511ebda3c..68d2a2e8ce6 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -12346,7 +12346,7 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { // This test exercises a complex CTE query for prompt // reconstruction after compaction. It requires Postgres. - db, _ := dbtestutil.NewDB(t) + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) ctx := context.Background() // Helper: create a chat model config (required FK for chats). @@ -12434,6 +12434,51 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { return ids } + // invertCreatedAt makes created_at descend as id ascends, so a reader that + // leads with created_at sees the chat in reverse append order. + invertCreatedAt := func(t *testing.T, chatID uuid.UUID) { + t.Helper() + _, err := sqlDB.ExecContext(ctx, + "UPDATE chat_messages SET created_at = now() - (id || ' seconds')::interval WHERE chat_id = $1", + chatID) + require.NoError(t, err) + } + + t.Run("OrdersByIDWhenTimestampsDisagree", func(t *testing.T) { + t.Parallel() + chat := newChat(t) + + sys := insertMsg(t, chat.ID, database.ChatMessageRoleSystem, database.ChatMessageVisibilityModel, false, "system prompt") + usr := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "question") + ast := insertMsg(t, chat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, false, "tool call") + tool := insertMsg(t, chat.ID, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, false, "tool result") + invertCreatedAt(t, chat.ID) + + got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, []int64{sys.ID, usr.ID, ast.ID, tool.ID}, msgIDs(got), + "the prompt must keep append order so a tool result follows its assistant call") + }) + + t.Run("CompactionBoundaryUsesID", func(t *testing.T) { + t.Parallel() + chat := newChat(t) + + sys := insertMsg(t, chat.ID, database.ChatMessageRoleSystem, database.ChatMessageVisibilityModel, false, "system prompt") + insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "before first summary") + staleSummary := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, true, "first summary") + insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "between summaries") + latestSummary := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, true, "second summary") + afterLatest := insertMsg(t, chat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, false, "after second summary") + invertCreatedAt(t, chat.ID) + + got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, []int64{sys.ID, latestSummary.ID, afterLatest.ID}, msgIDs(got), + "the boundary is compared with id, so it must also be selected by id") + require.NotContains(t, msgIDs(got), staleSummary.ID) + }) + t.Run("NoCompaction", func(t *testing.T) { t.Parallel() chat := newChat(t) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 30bf07a624a..915c486574d 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -8417,7 +8417,6 @@ WITH latest_compressed_summary AS ( AND deleted = false AND visibility = 'model' ORDER BY - created_at DESC, id DESC LIMIT 1 @@ -8460,10 +8459,12 @@ WHERE ) ) ORDER BY - created_at ASC, id ASC ` +// Ordered by id throughout: the boundary row below is compared with id, and the +// prompt must present roles in append order so tool results follow the assistant +// message that requested them. func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) { rows, err := q.db.QueryContext(ctx, getChatMessagesForPromptByChatID, chatID) if err != nil { diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index fec4a006aec..a99558ce20c 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -484,6 +484,9 @@ LIMIT COALESCE(NULLIF(@limit_val::int, 0), 500); -- name: GetChatMessagesForPromptByChatID :many +-- Ordered by id throughout: the boundary row below is compared with id, and the +-- prompt must present roles in append order so tool results follow the assistant +-- message that requested them. WITH latest_compressed_summary AS ( SELECT id @@ -495,7 +498,6 @@ WITH latest_compressed_summary AS ( AND deleted = false AND visibility = 'model' ORDER BY - created_at DESC, id DESC LIMIT 1 @@ -538,7 +540,6 @@ WHERE ) ) ORDER BY - created_at ASC, id ASC; -- name: GetChats :many From d974c9eee521a3ffcaed150adafdf1dbe0c974bd Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 28 Jul 2026 21:57:29 +0000 Subject: [PATCH 2/3] test(coderd/database): tighten prompt query comment and reuse the id helper --- coderd/database/querier.go | 5 ++--- coderd/database/querier_test.go | 28 +++++++++------------------- coderd/database/queries.sql.go | 5 ++--- coderd/database/queries/chats.sql | 5 ++--- 4 files changed, 15 insertions(+), 28 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 694c2918ae3..eefcf9eadcf 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -469,9 +469,8 @@ type sqlcQuerier interface { GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedParams) ([]ChatMessage, error) // Stream deltas and reset snapshots must use the same message order. GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) - // Ordered by id throughout: the boundary row below is compared with id, and the - // prompt must present roles in append order so tool results follow the assistant - // message that requested them. + // The compaction boundary and final ordering must use the same key so tool + // results remain after their assistant calls. GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error) GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 68d2a2e8ce6..46a25dd1043 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -12426,16 +12426,6 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { return database.ChatMessage(results[0]) } - msgIDs := func(msgs []database.ChatMessage) []int64 { - ids := make([]int64, len(msgs)) - for i, m := range msgs { - ids[i] = m.ID - } - return ids - } - - // invertCreatedAt makes created_at descend as id ascends, so a reader that - // leads with created_at sees the chat in reverse append order. invertCreatedAt := func(t *testing.T, chatID uuid.UUID) { t.Helper() _, err := sqlDB.ExecContext(ctx, @@ -12456,7 +12446,7 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) - require.Equal(t, []int64{sys.ID, usr.ID, ast.ID, tool.ID}, msgIDs(got), + require.Equal(t, []int64{sys.ID, usr.ID, ast.ID, tool.ID}, chatMessageIDs(got), "the prompt must keep append order so a tool result follows its assistant call") }) @@ -12474,9 +12464,9 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) - require.Equal(t, []int64{sys.ID, latestSummary.ID, afterLatest.ID}, msgIDs(got), + require.Equal(t, []int64{sys.ID, latestSummary.ID, afterLatest.ID}, chatMessageIDs(got), "the boundary is compared with id, so it must also be selected by id") - require.NotContains(t, msgIDs(got), staleSummary.ID) + require.NotContains(t, chatMessageIDs(got), staleSummary.ID) }) t.Run("NoCompaction", func(t *testing.T) { @@ -12489,7 +12479,7 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) - require.Equal(t, []int64{sys.ID, usr.ID, ast.ID}, msgIDs(got)) + require.Equal(t, []int64{sys.ID, usr.ID, ast.ID}, chatMessageIDs(got)) }) t.Run("UserOnlyVisibilityExcluded", func(t *testing.T) { @@ -12508,7 +12498,7 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { require.NotEqual(t, database.ChatMessageVisibilityUser, m.Visibility, "visibility=user messages should not appear in the prompt") } - require.Contains(t, msgIDs(got), usr.ID) + require.Contains(t, chatMessageIDs(got), usr.ID) }) t.Run("AfterCompaction", func(t *testing.T) { @@ -12535,7 +12525,7 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) - gotIDs := msgIDs(got) + gotIDs := chatMessageIDs(got) // Must include: system prompt, summary, post-compaction. require.Contains(t, gotIDs, sys.ID, "system prompt must be included") @@ -12574,8 +12564,8 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { } require.True(t, hasNonSystem, "prompt must contain at least one non-system message after compaction") - require.Contains(t, msgIDs(got), summary.ID) - require.Contains(t, msgIDs(got), newUsr.ID) + require.Contains(t, chatMessageIDs(got), summary.ID) + require.Contains(t, chatMessageIDs(got), newUsr.ID) }) t.Run("CompressedToolResultNotPickedAsSummary", func(t *testing.T) { @@ -12594,7 +12584,7 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { got, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) - gotIDs := msgIDs(got) + gotIDs := chatMessageIDs(got) require.Contains(t, gotIDs, summary.ID, "real summary must be included") require.NotContains(t, gotIDs, compressedTool.ID, "compressed tool result must not be included") diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 915c486574d..97c31e6b891 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -8462,9 +8462,8 @@ ORDER BY id ASC ` -// Ordered by id throughout: the boundary row below is compared with id, and the -// prompt must present roles in append order so tool results follow the assistant -// message that requested them. +// The compaction boundary and final ordering must use the same key so tool +// results remain after their assistant calls. func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) { rows, err := q.db.QueryContext(ctx, getChatMessagesForPromptByChatID, chatID) if err != nil { diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index a99558ce20c..b672ec5068a 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -484,9 +484,8 @@ LIMIT COALESCE(NULLIF(@limit_val::int, 0), 500); -- name: GetChatMessagesForPromptByChatID :many --- Ordered by id throughout: the boundary row below is compared with id, and the --- prompt must present roles in append order so tool results follow the assistant --- message that requested them. +-- The compaction boundary and final ordering must use the same key so tool +-- results remain after their assistant calls. WITH latest_compressed_summary AS ( SELECT id From 387640a10016b80304b8750fea48e8272b845e7a Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Wed, 29 Jul 2026 08:06:31 +0000 Subject: [PATCH 3/3] fix(coderd/database): rebuild the chat summary boundary index by id idx_chat_messages_compressed_summary_boundary required role = 'system', but compaction writes its summary with the user role, so the index never matched a row. It also led with created_at, which no longer matches the boundary lookup now that it orders by id. Rebuilt as (chat_id, id DESC) WHERE compressed AND NOT deleted AND visibility = 'model'. The boundary lookup becomes an index-only scan: 2 buffers instead of 267 on a 20k-message chat. --- coderd/database/dump.sql | 2 +- .../000560_chat_summary_boundary_index_by_id.down.sql | 7 +++++++ .../000560_chat_summary_boundary_index_by_id.up.sql | 10 ++++++++++ 3 files changed, 18 insertions(+), 1 deletion(-) create mode 100644 coderd/database/migrations/000560_chat_summary_boundary_index_by_id.down.sql create mode 100644 coderd/database/migrations/000560_chat_summary_boundary_index_by_id.up.sql diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 327a98638a2..fb7ab0b8e40 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -4770,7 +4770,7 @@ CREATE INDEX idx_chat_messages_chat_created ON chat_messages USING btree (chat_i CREATE INDEX idx_chat_messages_chat_role_id ON chat_messages USING btree (chat_id, role, id DESC) WHERE (deleted = false); -CREATE INDEX idx_chat_messages_compressed_summary_boundary ON chat_messages USING btree (chat_id, created_at DESC, id DESC) WHERE ((compressed = true) AND (role = 'system'::chat_message_role) AND (visibility = ANY (ARRAY['model'::chat_message_visibility, 'both'::chat_message_visibility]))); +CREATE INDEX idx_chat_messages_compressed_summary_boundary ON chat_messages USING btree (chat_id, id DESC) WHERE ((compressed = true) AND (deleted = false) AND (visibility = 'model'::chat_message_visibility)); CREATE INDEX idx_chat_messages_created_at ON chat_messages USING btree (created_at); diff --git a/coderd/database/migrations/000560_chat_summary_boundary_index_by_id.down.sql b/coderd/database/migrations/000560_chat_summary_boundary_index_by_id.down.sql new file mode 100644 index 00000000000..f0bb3a2f3c9 --- /dev/null +++ b/coderd/database/migrations/000560_chat_summary_boundary_index_by_id.down.sql @@ -0,0 +1,7 @@ +DROP INDEX idx_chat_messages_compressed_summary_boundary; + +CREATE INDEX idx_chat_messages_compressed_summary_boundary + ON chat_messages(chat_id, created_at DESC, id DESC) + WHERE compressed = TRUE + AND role = 'system' + AND visibility IN ('model', 'both'); diff --git a/coderd/database/migrations/000560_chat_summary_boundary_index_by_id.up.sql b/coderd/database/migrations/000560_chat_summary_boundary_index_by_id.up.sql new file mode 100644 index 00000000000..914bea3397c --- /dev/null +++ b/coderd/database/migrations/000560_chat_summary_boundary_index_by_id.up.sql @@ -0,0 +1,10 @@ +-- The predicate required role = 'system', but compaction writes its summary +-- with the user role, so this index has never matched a row. Rebuild it to +-- match GetChatMessagesForPromptByChatID's boundary lookup, which orders by id. +DROP INDEX idx_chat_messages_compressed_summary_boundary; + +CREATE INDEX idx_chat_messages_compressed_summary_boundary + ON chat_messages(chat_id, id DESC) + WHERE compressed = TRUE + AND deleted = false + AND visibility = 'model';