diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 327a98638a2ad..fb7ab0b8e4079 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 0000000000000..f0bb3a2f3c9ff --- /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 0000000000000..914bea3397ce1 --- /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'; diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 0eb2f1d78d924..eefcf9eadcf36 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -469,6 +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) + // 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 c5511ebda3c1f..46a25dd10434d 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). @@ -12426,14 +12426,49 @@ 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 := 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}, chatMessageIDs(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}, chatMessageIDs(got), + "the boundary is compared with id, so it must also be selected by id") + require.NotContains(t, chatMessageIDs(got), staleSummary.ID) + }) + t.Run("NoCompaction", func(t *testing.T) { t.Parallel() chat := newChat(t) @@ -12444,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) { @@ -12463,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) { @@ -12490,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") @@ -12529,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) { @@ -12549,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 30bf07a624ae6..97c31e6b891be 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,11 @@ WHERE ) ) ORDER BY - created_at ASC, id ASC ` +// 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 fec4a006aecbc..b672ec5068a1b 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -484,6 +484,8 @@ LIMIT COALESCE(NULLIF(@limit_val::int, 0), 500); -- name: GetChatMessagesForPromptByChatID :many +-- 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 @@ -495,7 +497,6 @@ WITH latest_compressed_summary AS ( AND deleted = false AND visibility = 'model' ORDER BY - created_at DESC, id DESC LIMIT 1 @@ -538,7 +539,6 @@ WHERE ) ) ORDER BY - created_at ASC, id ASC; -- name: GetChats :many