opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(coderd/database): order the chat prompt query and its boundary by id (#27619)

Browse Source 
## Stack context

Follows #27495 (merged), which gives `chat_messages.id` an append-order
guarantee and moves the history reads onto it. This PR applies the same
fix to the query that builds the model prompt.

## Why?

`GetChatMessagesForPromptByChatID` mixed two orderings. It selected the
compaction boundary with `created_at DESC, id DESC`, then applied that
boundary with an `id >` comparison, and returned rows with `created_at
ASC, id ASC`.

`created_at` is `now()`, so it is the transaction start time. Every row
in one insert batch shares it, and concurrent transactions can commit in
the opposite order to the one they started in. Two consequences, both
reaching the provider:

- **Malformed prompts.** A tool result could be ordered ahead of the
assistant message that requested it.
`chatprompt.injectMissingToolResults` does not repair this: it only
handles tool rows already contiguous after an assistant row, and adds
missing results. It never moves a tool row that precedes its assistant,
and nothing re-sorts the rows in Go.
- **Wrong compaction boundary.** The boundary is picked by timestamp but
compared by id, so a stale compressed summary could be retained while
the actual latest one was dropped.

## Changes

Both the boundary CTE and the outer query order by `id`. The `id >`
predicate is unchanged, which is the point: the ordering now matches the
comparison that was always being made.

**The boundary index was dead, so it is rebuilt to match.**
`idx_chat_messages_compressed_summary_boundary` was created for exactly
this lookup, but its predicate requires `role = 'system'` while
compaction writes its summary with the user role
(`message_conversion.go:334`, the only writer of `compressed = true`).
It matched zero rows, and no other query can use it. Migration `000560`
rebuilds it as `(chat_id, id DESC) WHERE compressed AND NOT deleted AND
visibility = 'model'`, which also matches the new order key.

Measured on PostgreSQL 13 with a 20k-message chat, 11 summaries, and 14
sibling chats so `chat_id` is selective:

| boundary lookup | plan | buffers |
|---|---|---|
| old predicate | Index Scan `idx_chat_messages_chat`, 19,989 rows
filtered | 267 |
| rebuilt index | Index Only Scan | 2 |

Not in scope: the outer `SELECT` still inspects every row of the chat,
because its `role = 'system' AND compressed = FALSE` disjunct has no
lower `id` bound. That predates this PR and needs a query rewrite rather
than an index.

## Testing

Two subtests, both verified red by reverting the `ORDER BY` and
regenerating:

- `OrdersByIDWhenTimestampsDisagree` returned `[4,3,2,1]` instead of
`[1,2,3,4]`, placing the tool result before the assistant call.
- `CompactionBoundaryUsesID` selected the stale summary and leaked the
messages between the two summaries into the prompt.

Existing subtests pass unchanged. Migration up/down tests pass, and the
rebuilt index was verified red-green: restoring the old predicate
returns the plan to a 267-buffer scan, and the old predicate matches 0
rows in the fixture.

> Opened by Mux on behalf of Mike.
This commit is contained in:
Michael Suchacz
2026-07-29 14:04:59 +02:00
committed by GitHub
parent b371262e5c
commit 1c722ff969
7 changed files with 72 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dump.sql
Generated
+1 -1
View File
@@ -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);
coderd/database/migrations/000560_chat_summary_boundary_index_by_id.down.sql
+7
View File
@@ -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');
coderd/database/migrations/000560_chat_summary_boundary_index_by_id.up.sql
+10
View File
@@ -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';
coderd/database/querier.go
Generated
+2
View File
@@ -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)
coderd/database/querier_test.go
+48 -13
View File
@@ -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")
coderd/database/queries.sql.go
Generated
+2 -2
View File
@@ -8496,7 +8496,6 @@ WITH latest_compressed_summary AS (
AND deleted = false
AND visibility = 'model'
ORDER BY
created_at DESC,
id DESC
LIMIT
1
@@ -8539,10 +8538,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 {
coderd/database/queries/chats.sql
+2 -2
View File
@@ -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