From e96d8646e2574434108c82f64f7c338eb01ac634 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:17:42 +0200 Subject: [PATCH] fix(coderd): give chat message ids an append-order guarantee (#27495) Chat message ordering was derived from `created_at`, which is `now()` and therefore the transaction start time. That makes it unusable as an append-order column for two independent reasons: every row in one `InsertChatMessages` batch shares a single timestamp, and two concurrent transactions can commit in the opposite order to the one they started in. This PR gives `chat_messages.id` a real append-order guarantee and moves the history reads onto it. ## Changes **`InsertChatMessages` had no input-order guarantee.** Callers index the returned slice by input position. That only worked because PostgreSQL happens to evaluate the `BIGSERIAL` default in row order. Ids are now allocated up front and the k-th smallest is assigned to input index k, so the pairing does not depend on where the column default is evaluated. Returned rows are explicitly `ORDER BY id`. **Three history reads now order by `id`.** | Query | Was | Now | |---|---|---| | `GetChatMessagesByChatID` | `created_at ASC` | `id ASC` | | `GetChatMessagesByRevisionForStream` | `created_at ASC, id ASC` | `id ASC` | | `GetLastChatMessageByRole` | `created_at DESC, id DESC` | `id DESC` | `GetChatMessagesByChatID` paginated by `id` while ordering by `created_at`, which is incoherent on its own terms. The other two matter because of who consumes them. The stream query supplies incremental updates on the same socket that emits a full `GetChatMessagesByChatID` snapshot on history reset, so once that snapshot moved to `id` the two disagreed under timestamp skew. `GetLastChatMessageByRole` returns an id that is then used as an id cursor, both as `AfterID` when synthesizing tool cancellations and as `chats.last_read_message_id`, where a stale anchor leaves later assistant messages permanently unread. A tie-breaker would not have fixed either one. It only resolves equal timestamps; leading with `created_at` is the actual defect. **`GetLastChatMessageByRole` loses its index, so this adds one.** `ORDER BY created_at DESC, id DESC` could take an ordered scan of `idx_chat_messages_chat_created`. Nothing in the schema can supply `ORDER BY id DESC LIMIT 1` for a given `chat_id` and `role`, so the planner switches to a backward scan of the primary key and filters every newer row in the table, scanning all of it when the chat has no message in that role, which is the routine case for a fresh chat. Migration `000559` adds `(chat_id, role, id DESC) WHERE deleted = false`, the same shape as the existing `idx_chat_messages_user_prompts`. This matters because the query is hot: it runs on every stream connect and disconnect, and once per turn when synthesizing tool cancellations. `GetChatMessagesForPromptByChatID` has the same defect and is fixed in the stacked PR, because its compaction boundary change is semantic and deserves a separate review. Auto-archive stays timestamp-based deliberately: it measures activity, not order. Wrapping the insert in a CTE (needed because `INSERT` cannot take `ORDER BY`) makes sqlc synthesize `InsertChatMessagesRow`. It is structurally identical to `ChatMessage`, so the call sites use a direct struct conversion that stops compiling if the two ever diverge. ## Testing Behavior tests write `created_at` values inverted against id order, so a reader that leads with `created_at` returns the batch backwards. All three queries were verified red by reverting the `ORDER BY` and regenerating: the stream query returned `[3,2,1]` for `[1,2,3]`, and `GetLastChatMessageByRole` picked id 1 instead of id 3. `TestInsertChatMessagesOrderContract` asserts against the generated SQL, covering what a behavior test cannot: PostgreSQL evaluates the id default in row order anyway, so a batch still looks ordered once the guarantee is removed. `TestChatMessagesSequenceCacheIsOne` guards the cross-batch half of the invariant. Ids follow chat row lock order only while the sequence hands out one value at a time; sequence cache blocks are per session, so with a cache above one a backend holding stale cached values can lock second and still commit lower ids. Bumping a sequence cache is an ordinary throughput tweak, and it would silently corrupt history order. The index was checked on a 200k row fixture. Without it, the zero-match lookup filters all 200,000 rows over 2763 buffers; with it, the plan is an index scan with both `chat_id` and `role` in the index condition, no sort node, and 3 buffers. Note that the within-batch mapping does not depend on the cache size. It is established by `ROW_NUMBER() OVER (ORDER BY id)` over the allocated ids, so it holds regardless of `nextval` evaluation order. ## Note on the deleted subagent hand-sort The subagent history reader's hand-sort stays deleted, but calling it redundant was imprecise. It sorted by `created_at` then `id`, so it is only equivalent to `id` ordering when the two agree. When they disagree the old code selected a different "latest assistant". This is a deliberate behavior change to match the new invariant, not dead-code removal. > Opened by Mux on behalf of Mike. --- coderd/database/dbauthz/dbauthz.go | 2 +- coderd/database/dbauthz/dbauthz_test.go | 2 +- coderd/database/dbgen/dbgen.go | 2 +- coderd/database/dbmetrics/querymetrics.go | 2 +- coderd/database/dbmock/dbmock.go | 4 +- coderd/database/dump.sql | 2 + ..._chat_messages_last_by_role_index.down.sql | 1 + ...59_chat_messages_last_by_role_index.up.sql | 5 + coderd/database/modelqueries_internal_test.go | 22 +++ coderd/database/querier.go | 11 +- coderd/database/querier_test.go | 133 +++++++++++++- coderd/database/queries.sql.go | 163 ++++++++++++------ coderd/database/queries/chats.sql | 116 ++++++++----- coderd/exp_chats_test.go | 2 +- coderd/x/chatd/chatd_test.go | 2 +- coderd/x/chatd/chatstate/messages.go | 11 ++ coderd/x/chatd/chatstate/transitions.go | 4 +- coderd/x/chatd/subagent.go | 8 - 18 files changed, 373 insertions(+), 119 deletions(-) create mode 100644 coderd/database/migrations/000559_chat_messages_last_by_role_index.down.sql create mode 100644 coderd/database/migrations/000559_chat_messages_last_by_role_index.up.sql diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 03f8ce5cac..de153b24c2 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -6086,7 +6086,7 @@ func (q *querier) InsertChatFile(ctx context.Context, arg database.InsertChatFil return insert(q.log, q.auth, rbac.ResourceChat.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), q.db.InsertChatFile)(ctx, arg) } -func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { +func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.InsertChatMessagesRow, error) { // Authorize create on the parent chat (using update permission). chat, err := q.db.GetChatByID(ctx, arg.ChatID) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 25d2622de1..ecc58968c2 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -1281,7 +1281,7 @@ func (s *MethodTestSuite) TestChats() { s.Run("InsertChatMessages", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) arg := testutil.Fake(s.T(), faker, database.InsertChatMessagesParams{ChatID: chat.ID}) - msgs := []database.ChatMessage{testutil.Fake(s.T(), faker, database.ChatMessage{ChatID: chat.ID})} + msgs := []database.InsertChatMessagesRow{testutil.Fake(s.T(), faker, database.InsertChatMessagesRow{ChatID: chat.ID})} dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().InsertChatMessages(gomock.Any(), arg).Return(msgs, nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(msgs) diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 98c8718697..a85e500986 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -145,7 +145,7 @@ func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) dat }) require.NoError(t, err, "insert chat message") require.Len(t, msgs, 1) - return msgs[0] + return database.ChatMessage(msgs[0]) } const ( diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index c56b38190b..26f47ff51d 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -4169,7 +4169,7 @@ func (m queryMetricsStore) InsertChatFile(ctx context.Context, arg database.Inse return r0, r1 } -func (m queryMetricsStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { +func (m queryMetricsStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.InsertChatMessagesRow, error) { start := time.Now() r0, r1 := m.s.InsertChatMessages(ctx, arg) m.queryLatencies.WithLabelValues("InsertChatMessages").Observe(time.Since(start).Seconds()) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index fff45208fc..ac848e0f99 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -7812,10 +7812,10 @@ func (mr *MockStoreMockRecorder) InsertChatFile(ctx, arg any) *gomock.Call { } // InsertChatMessages mocks base method. -func (m *MockStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { +func (m *MockStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.InsertChatMessagesRow, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "InsertChatMessages", ctx, arg) - ret0, _ := ret[0].([]database.ChatMessage) + ret0, _ := ret[0].([]database.InsertChatMessagesRow) ret1, _ := ret[1].(error) return ret0, ret1 } diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index c167dd3beb..327a98638a 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -4768,6 +4768,8 @@ CREATE INDEX idx_chat_messages_chat ON chat_messages USING btree (chat_id); CREATE INDEX idx_chat_messages_chat_created ON chat_messages USING btree (chat_id, created_at); +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_created_at ON chat_messages USING btree (created_at); diff --git a/coderd/database/migrations/000559_chat_messages_last_by_role_index.down.sql b/coderd/database/migrations/000559_chat_messages_last_by_role_index.down.sql new file mode 100644 index 0000000000..e9c83f584e --- /dev/null +++ b/coderd/database/migrations/000559_chat_messages_last_by_role_index.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS idx_chat_messages_chat_role_id; diff --git a/coderd/database/migrations/000559_chat_messages_last_by_role_index.up.sql b/coderd/database/migrations/000559_chat_messages_last_by_role_index.up.sql new file mode 100644 index 0000000000..91c46b8134 --- /dev/null +++ b/coderd/database/migrations/000559_chat_messages_last_by_role_index.up.sql @@ -0,0 +1,5 @@ +-- Serves GetLastChatMessageByRole. It orders by id, so the existing +-- (chat_id, created_at) index cannot supply the LIMIT 1 row in index order. +CREATE INDEX idx_chat_messages_chat_role_id + ON chat_messages (chat_id, role, id DESC) + WHERE deleted = false; diff --git a/coderd/database/modelqueries_internal_test.go b/coderd/database/modelqueries_internal_test.go index 698954e39b..ed5dc9a8d2 100644 --- a/coderd/database/modelqueries_internal_test.go +++ b/coderd/database/modelqueries_internal_test.go @@ -1,6 +1,7 @@ package database import ( + "reflect" "regexp" "slices" "strings" @@ -168,6 +169,27 @@ func TestFinalizeStaleChatDebugRows_TerminalStatusAlignment(t *testing.T) { } } +// TestInsertChatMessagesOrderContract guards the input-order guarantee that +// callers rely on when indexing the returned slice. A behavior test cannot: +// Postgres evaluates the id default in row order anyway, so a batch still looks +// ordered once the guarantee is removed. +func TestInsertChatMessagesOrderContract(t *testing.T) { + t.Parallel() + + require.Contains(t, insertChatMessages, "nextval('chat_messages_id_seq')", + "ids must be allocated explicitly so they can be correlated to input array position") + require.Contains(t, insertChatMessages, "ROW_NUMBER() OVER (ORDER BY id)", + "the k-th smallest allocated id must be assigned to input index k") + require.Regexp(t, `(?s)ORDER BY id\s*\z`, strings.TrimSpace(insertChatMessages), + "returned rows must be explicitly ordered by id rather than relying on RETURNING order") + + // Every parallel input array must be read at the allocated ordinal. A column + // left on UNNEST would be positioned by the executor instead. + subscripted := regexp.MustCompile(`\)\[allocated\.ord\]`).FindAllString(insertChatMessages, -1) + require.Len(t, subscripted, reflect.TypeOf(InsertChatMessagesParams{}).NumField()-1, + "each InsertChatMessagesParams array field, all but ChatID, must be subscripted by allocated.ord") +} + // extractWhereClause extracts the WHERE clause from a SQL query string func extractWhereClause(query string) string { // Find WHERE and get everything after it diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 703cd40755..0eb2f1d78d 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -461,9 +461,13 @@ type sqlcQuerier interface { // after the given timestamp. Uses message created_at so that // ongoing activity in long-running chats is captured each window. GetChatMessageSummariesPerChat(ctx context.Context, createdAfter time.Time) ([]GetChatMessageSummariesPerChatRow, error) + // Ordered by id to match the @after_id cursor. created_at is the transaction + // start time, so it can disagree with append order when a transaction takes the + // chat row lock later than one that started after it. GetChatMessagesByChatID(ctx context.Context, arg GetChatMessagesByChatIDParams) ([]ChatMessage, error) GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedParams) ([]ChatMessage, error) 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) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error) @@ -637,6 +641,8 @@ type sqlcQuerier interface { // param created_at_opt: The created_at timestamp to filter by. This parameter is usd for pagination - it fetches notifications created before the specified timestamp if it is not the zero value // param limit_opt: The limit of notifications to fetch. If the limit is not specified, it defaults to 25 GetInboxNotificationsByUserID(ctx context.Context, arg GetInboxNotificationsByUserIDParams) ([]InboxNotification, error) + // The returned id becomes both an AfterID cursor and last_read_message_id, so + // "last" must use id order. GetLastChatMessageByRole(ctx context.Context, arg GetLastChatMessageByRoleParams) (ChatMessage, error) GetLastUpdateCheck(ctx context.Context) (string, error) GetLatestCryptoKeyByFeature(ctx context.Context, feature CryptoKeyFeature) (CryptoKey, error) @@ -1105,7 +1111,10 @@ type sqlcQuerier interface { // with concurrent FinalizeStale under READ COMMITTED isolation. InsertChatDebugStep(ctx context.Context, arg InsertChatDebugStepParams) (ChatDebugStep, error) InsertChatFile(ctx context.Context, arg InsertChatFileParams) (InsertChatFileRow, error) - InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) + // Returns the inserted rows in input array order. Ids are allocated before the + // insert and the k-th smallest is assigned to input index k, so callers may + // index the result positionally. + InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]InsertChatMessagesRow, error) InsertChatModelConfig(ctx context.Context, arg InsertChatModelConfigParams) (ChatModelConfig, error) // Legacy queue insertion path. When no caller-supplied creator exists, // preserve the created_by invariant by attributing the queued row to the diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index a94729b9c8..c5511ebda3 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -12212,6 +12212,135 @@ func TestInsertChatMessages(t *testing.T) { }) } +// The returned ids are in insert order, which the inverted created_at values +// deliberately contradict. +func insertChatMessagesInvertedTimestamps(t *testing.T, db database.Store, sqlDB *sql.DB, roles []database.ChatMessageRole) (database.Chat, []int64) { + t.Helper() + + ctx := context.Background() + org := dbgen.Organization(t, db, database.Organization{}) + owner := dbgen.User(t, db, database.User{}) + modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, + }) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: owner.ID, + LastModelConfigID: modelCfg.ID, + }) + + count := len(roles) + inserted, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: slices.Repeat([]uuid.UUID{owner.ID}, count), + ModelConfigID: slices.Repeat([]uuid.UUID{modelCfg.ID}, count), + Role: roles, + ContentVersion: slices.Repeat([]int16{chatprompt.CurrentContentVersion}, count), + Visibility: slices.Repeat([]database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, count), + Content: slices.Repeat([]string{`"message"`}, count), + InputTokens: make([]int64, count), + OutputTokens: make([]int64, count), + TotalTokens: make([]int64, count), + ReasoningTokens: make([]int64, count), + CacheCreationTokens: make([]int64, count), + CacheReadTokens: make([]int64, count), + ContextLimit: make([]int64, count), + Compressed: make([]bool, count), + TotalCostMicros: make([]int64, count), + RuntimeMs: make([]int64, count), + }) + require.NoError(t, err) + require.Len(t, inserted, count) + + insertedIDs := make([]int64, count) + for i, message := range inserted { + insertedIDs[i] = message.ID + _, err := sqlDB.ExecContext(ctx, + "UPDATE chat_messages SET created_at = $1 WHERE id = $2", + message.CreatedAt.Add(time.Duration(count-i)*time.Minute), message.ID) + require.NoError(t, err) + } + + return chat, insertedIDs +} + +func chatMessageIDs(messages []database.ChatMessage) []int64 { + ids := make([]int64, len(messages)) + for i, message := range messages { + ids[i] = message.ID + } + return ids +} + +func TestGetChatMessagesByChatIDOrdersByID(t *testing.T) { + t.Parallel() + + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + ctx := context.Background() + + chat, insertedIDs := insertChatMessagesInvertedTimestamps(t, db, sqlDB, + slices.Repeat([]database.ChatMessageRole{database.ChatMessageRoleUser}, 3)) + + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + require.NoError(t, err) + require.Equal(t, insertedIDs, chatMessageIDs(messages)) +} + +func TestGetChatMessagesByRevisionForStreamOrdersByID(t *testing.T) { + t.Parallel() + + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + ctx := context.Background() + + chat, insertedIDs := insertChatMessagesInvertedTimestamps(t, db, sqlDB, + slices.Repeat([]database.ChatMessageRole{database.ChatMessageRoleUser}, 3)) + + messages, err := db.GetChatMessagesByRevisionForStream(ctx, database.GetChatMessagesByRevisionForStreamParams{ + ChatID: chat.ID, + AfterRevision: 0, + }) + require.NoError(t, err) + require.Equal(t, insertedIDs, chatMessageIDs(messages)) +} + +func TestGetLastChatMessageByRoleOrdersByID(t *testing.T) { + t.Parallel() + + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + ctx := context.Background() + + chat, insertedIDs := insertChatMessagesInvertedTimestamps(t, db, sqlDB, + slices.Repeat([]database.ChatMessageRole{database.ChatMessageRoleAssistant}, 3)) + + last, err := db.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{ + ChatID: chat.ID, + Role: database.ChatMessageRoleAssistant, + }) + require.NoError(t, err) + require.Equal(t, insertedIDs[len(insertedIDs)-1], last.ID) +} + +// Sequence cache blocks are handed out per session, so above cache 1 a backend +// holding stale cached values can take the chat row lock second and still commit +// lower ids. Bumping a sequence cache is an ordinary throughput tweak. +func TestChatMessagesSequenceCacheIsOne(t *testing.T) { + t.Parallel() + + _, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + ctx := context.Background() + + var cacheSize int64 + err := sqlDB.QueryRowContext(ctx, + "SELECT cache_size FROM pg_sequences WHERE sequencename = 'chat_messages_id_seq'"). + Scan(&cacheSize) + require.NoError(t, err) + require.Equal(t, int64(1), cacheSize, "chat_messages_id_seq must use cache 1") +} + func TestGetChatMessagesForPromptByChatID(t *testing.T) { t.Parallel() @@ -12294,7 +12423,7 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { RuntimeMs: []int64{0}, }) require.NoError(t, err) - return results[0] + return database.ChatMessage(results[0]) } msgIDs := func(msgs []database.ChatMessage) []int64 { @@ -17173,7 +17302,7 @@ func TestGetChatsSearch(t *testing.T) { }) require.NoError(t, err) require.Len(t, msgs, 1) - return msgs[0] + return database.ChatMessage(msgs[0]) } linkPR := func(chatID uuid.UUID, url, state, prTitle string, prNumber int32, gitRemoteOrigin string) { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index b86e7cb57c..30bf07a624 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -8126,7 +8126,7 @@ WHERE AND visibility IN ('user', 'both') AND deleted = false ORDER BY - created_at ASC + id ASC ` type GetChatMessagesByChatIDParams struct { @@ -8134,6 +8134,9 @@ type GetChatMessagesByChatIDParams struct { AfterID int64 `db:"after_id" json:"after_id"` } +// Ordered by id to match the @after_id cursor. created_at is the transaction +// start time, so it can disagree with append order when a transaction takes the +// chat row lock later than one that started after it. func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMessagesByChatIDParams) ([]ChatMessage, error) { rows, err := q.db.QueryContext(ctx, getChatMessagesByChatID, arg.ChatID, arg.AfterID) if err != nil { @@ -8345,7 +8348,7 @@ WHERE AND revision > $2::bigint AND visibility IN ('user', 'both') ORDER BY - created_at ASC, id ASC + id ASC ` type GetChatMessagesByRevisionForStreamParams struct { @@ -8353,6 +8356,7 @@ type GetChatMessagesByRevisionForStreamParams struct { AfterRevision int64 `db:"after_revision" json:"after_revision"` } +// Stream deltas and reset snapshots must use the same message order. func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) { rows, err := q.db.QueryContext(ctx, getChatMessagesByRevisionForStream, arg.ChatID, arg.AfterRevision) if err != nil { @@ -9867,7 +9871,7 @@ WHERE AND role = $2::chat_message_role AND deleted = false ORDER BY - created_at DESC, id DESC + id DESC LIMIT 1 ` @@ -9877,6 +9881,8 @@ type GetLastChatMessageByRoleParams struct { Role ChatMessageRole `db:"role" json:"role"` } +// The returned id becomes both an AfterID cursor and last_read_message_id, so +// "last" must use id order. func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastChatMessageByRoleParams) (ChatMessage, error) { row := q.db.QueryRowContext(ctx, getLastChatMessageByRole, arg.ChatID, arg.Role) var i ChatMessage @@ -10375,7 +10381,7 @@ WITH batch AS ( SELECT ( SELECT val - FROM UNNEST($3::uuid[]) + FROM UNNEST($1::uuid[]) WITH ORDINALITY AS t(val, ord) WHERE val != '00000000-0000-0000-0000-000000000000'::uuid ORDER BY ord DESC @@ -10383,7 +10389,7 @@ WITH batch AS ( ) AS last_model_config_id, ( SELECT NULLIF(val, '')::chat_reasoning_effort - FROM UNNEST($4::text[]) + FROM UNNEST($2::text[]) WITH ORDINALITY AS t(val, ord) WHERE val != '' ORDER BY ord DESC @@ -10398,61 +10404,80 @@ updated_chat AS ( last_reasoning_effort = COALESCE(batch.last_reasoning_effort, chats.last_reasoning_effort) FROM batch WHERE - chats.id = $1::uuid + chats.id = $3::uuid AND ( chats.last_model_config_id IS DISTINCT FROM COALESCE(batch.last_model_config_id, chats.last_model_config_id) OR chats.last_reasoning_effort IS DISTINCT FROM COALESCE(batch.last_reasoning_effort, chats.last_reasoning_effort) ) +), +allocated AS MATERIALIZED ( + -- Numbering the ids by value, rather than by the order nextval produced + -- them, is what makes ordinal k always the k-th smallest id. MATERIALIZED + -- is redundant while nextval is volatile, and pins that if it changes. + SELECT + id, + (ROW_NUMBER() OVER (ORDER BY id))::int AS ord + FROM ( + SELECT nextval('chat_messages_id_seq') AS id + FROM generate_series(1, cardinality($4::chat_message_role[])) + ) s +), +inserted AS ( + INSERT INTO chat_messages ( + id, + chat_id, + created_by, + model_config_id, + reasoning_effort, + role, + content, + content_version, + visibility, + input_tokens, + output_tokens, + total_tokens, + reasoning_tokens, + cache_creation_tokens, + cache_read_tokens, + context_limit, + compressed, + total_cost_micros, + runtime_ms + ) + SELECT + allocated.id, + $3::uuid, + NULLIF(($5::uuid[])[allocated.ord], '00000000-0000-0000-0000-000000000000'::uuid), + NULLIF(($1::uuid[])[allocated.ord], '00000000-0000-0000-0000-000000000000'::uuid), + NULLIF(($2::text[])[allocated.ord], '')::chat_reasoning_effort, + ($4::chat_message_role[])[allocated.ord], + ($6::text[])[allocated.ord]::jsonb, + ($7::smallint[])[allocated.ord], + ($8::chat_message_visibility[])[allocated.ord], + NULLIF(($9::bigint[])[allocated.ord], 0), + NULLIF(($10::bigint[])[allocated.ord], 0), + NULLIF(($11::bigint[])[allocated.ord], 0), + NULLIF(($12::bigint[])[allocated.ord], 0), + NULLIF(($13::bigint[])[allocated.ord], 0), + NULLIF(($14::bigint[])[allocated.ord], 0), + NULLIF(($15::bigint[])[allocated.ord], 0), + ($16::boolean[])[allocated.ord], + NULLIF(($17::bigint[])[allocated.ord], 0), + NULLIF(($18::bigint[])[allocated.ord], 0) + FROM allocated + RETURNING id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, revision, reasoning_effort, search_tsv ) -INSERT INTO chat_messages ( - chat_id, - created_by, - model_config_id, - reasoning_effort, - role, - content, - content_version, - visibility, - input_tokens, - output_tokens, - total_tokens, - reasoning_tokens, - cache_creation_tokens, - cache_read_tokens, - context_limit, - compressed, - total_cost_micros, - runtime_ms -) -SELECT - $1::uuid, - NULLIF(UNNEST($2::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), - NULLIF(UNNEST($3::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), - NULLIF(UNNEST($4::text[]), '')::chat_reasoning_effort, - UNNEST($5::chat_message_role[]), - UNNEST($6::text[])::jsonb, - UNNEST($7::smallint[]), - UNNEST($8::chat_message_visibility[]), - NULLIF(UNNEST($9::bigint[]), 0), - NULLIF(UNNEST($10::bigint[]), 0), - NULLIF(UNNEST($11::bigint[]), 0), - NULLIF(UNNEST($12::bigint[]), 0), - NULLIF(UNNEST($13::bigint[]), 0), - NULLIF(UNNEST($14::bigint[]), 0), - NULLIF(UNNEST($15::bigint[]), 0), - UNNEST($16::boolean[]), - NULLIF(UNNEST($17::bigint[]), 0), - NULLIF(UNNEST($18::bigint[]), 0) -RETURNING - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, revision, reasoning_effort, search_tsv +SELECT id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, revision, reasoning_effort, search_tsv +FROM inserted +ORDER BY id ` type InsertChatMessagesParams struct { - ChatID uuid.UUID `db:"chat_id" json:"chat_id"` - CreatedBy []uuid.UUID `db:"created_by" json:"created_by"` ModelConfigID []uuid.UUID `db:"model_config_id" json:"model_config_id"` ReasoningEffort []string `db:"reasoning_effort" json:"reasoning_effort"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` Role []ChatMessageRole `db:"role" json:"role"` + CreatedBy []uuid.UUID `db:"created_by" json:"created_by"` Content []string `db:"content" json:"content"` ContentVersion []int16 `db:"content_version" json:"content_version"` Visibility []ChatMessageVisibility `db:"visibility" json:"visibility"` @@ -10468,13 +10493,43 @@ type InsertChatMessagesParams struct { RuntimeMs []int64 `db:"runtime_ms" json:"runtime_ms"` } -func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) { +type InsertChatMessagesRow struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + Role ChatMessageRole `db:"role" json:"role"` + Content pqtype.NullRawMessage `db:"content" json:"content"` + Visibility ChatMessageVisibility `db:"visibility" json:"visibility"` + InputTokens sql.NullInt64 `db:"input_tokens" json:"input_tokens"` + OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` + TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` + ContextLimit sql.NullInt64 `db:"context_limit" json:"context_limit"` + Compressed bool `db:"compressed" json:"compressed"` + CreatedBy uuid.NullUUID `db:"created_by" json:"created_by"` + ContentVersion int16 `db:"content_version" json:"content_version"` + TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` + SearchTsv interface{} `db:"search_tsv" json:"search_tsv"` +} + +// Returns the inserted rows in input array order. Ids are allocated before the +// insert and the k-th smallest is assigned to input index k, so callers may +// index the result positionally. +func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]InsertChatMessagesRow, error) { rows, err := q.db.QueryContext(ctx, insertChatMessages, - arg.ChatID, - pq.Array(arg.CreatedBy), pq.Array(arg.ModelConfigID), pq.Array(arg.ReasoningEffort), + arg.ChatID, pq.Array(arg.Role), + pq.Array(arg.CreatedBy), pq.Array(arg.Content), pq.Array(arg.ContentVersion), pq.Array(arg.Visibility), @@ -10493,9 +10548,9 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa return nil, err } defer rows.Close() - var items []ChatMessage + var items []InsertChatMessagesRow for rows.Next() { - var i ChatMessage + var i InsertChatMessagesRow if err := rows.Scan( &i.ID, &i.ChatID, diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 836d1200b3..fec4a006ae 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -386,6 +386,9 @@ WHERE AND deleted = false; -- name: GetChatMessagesByChatID :many +-- Ordered by id to match the @after_id cursor. created_at is the transaction +-- start time, so it can disagree with append order when a transaction takes the +-- chat row lock later than one that started after it. SELECT * FROM @@ -396,9 +399,10 @@ WHERE AND visibility IN ('user', 'both') AND deleted = false ORDER BY - created_at ASC; + id ASC; -- name: GetChatMessagesByRevisionForStream :many +-- Stream deltas and reset snapshots must use the same message order. SELECT * FROM @@ -408,7 +412,7 @@ WHERE AND revision > @after_revision::bigint AND visibility IN ('user', 'both') ORDER BY - created_at ASC, id ASC; + id ASC; -- name: GetChatMessagesByChatIDAscPaginated :many SELECT @@ -868,6 +872,9 @@ SELECT * FROM chats_expanded; -- name: InsertChatMessages :many +-- Returns the inserted rows in input array order. Ids are allocated before the +-- insert and the k-th smallest is assigned to input index k, so callers may +-- index the result positionally. WITH batch AS ( SELECT ( @@ -900,48 +907,67 @@ updated_chat AS ( chats.last_model_config_id IS DISTINCT FROM COALESCE(batch.last_model_config_id, chats.last_model_config_id) OR chats.last_reasoning_effort IS DISTINCT FROM COALESCE(batch.last_reasoning_effort, chats.last_reasoning_effort) ) +), +allocated AS MATERIALIZED ( + -- Numbering the ids by value, rather than by the order nextval produced + -- them, is what makes ordinal k always the k-th smallest id. MATERIALIZED + -- is redundant while nextval is volatile, and pins that if it changes. + SELECT + id, + (ROW_NUMBER() OVER (ORDER BY id))::int AS ord + FROM ( + SELECT nextval('chat_messages_id_seq') AS id + FROM generate_series(1, cardinality(@role::chat_message_role[])) + ) s +), +inserted AS ( + INSERT INTO chat_messages ( + id, + chat_id, + created_by, + model_config_id, + reasoning_effort, + role, + content, + content_version, + visibility, + input_tokens, + output_tokens, + total_tokens, + reasoning_tokens, + cache_creation_tokens, + cache_read_tokens, + context_limit, + compressed, + total_cost_micros, + runtime_ms + ) + SELECT + allocated.id, + @chat_id::uuid, + NULLIF((@created_by::uuid[])[allocated.ord], '00000000-0000-0000-0000-000000000000'::uuid), + NULLIF((@model_config_id::uuid[])[allocated.ord], '00000000-0000-0000-0000-000000000000'::uuid), + NULLIF((@reasoning_effort::text[])[allocated.ord], '')::chat_reasoning_effort, + (@role::chat_message_role[])[allocated.ord], + (@content::text[])[allocated.ord]::jsonb, + (@content_version::smallint[])[allocated.ord], + (@visibility::chat_message_visibility[])[allocated.ord], + NULLIF((@input_tokens::bigint[])[allocated.ord], 0), + NULLIF((@output_tokens::bigint[])[allocated.ord], 0), + NULLIF((@total_tokens::bigint[])[allocated.ord], 0), + NULLIF((@reasoning_tokens::bigint[])[allocated.ord], 0), + NULLIF((@cache_creation_tokens::bigint[])[allocated.ord], 0), + NULLIF((@cache_read_tokens::bigint[])[allocated.ord], 0), + NULLIF((@context_limit::bigint[])[allocated.ord], 0), + (@compressed::boolean[])[allocated.ord], + NULLIF((@total_cost_micros::bigint[])[allocated.ord], 0), + NULLIF((@runtime_ms::bigint[])[allocated.ord], 0) + FROM allocated + RETURNING * ) -INSERT INTO chat_messages ( - chat_id, - created_by, - model_config_id, - reasoning_effort, - role, - content, - content_version, - visibility, - input_tokens, - output_tokens, - total_tokens, - reasoning_tokens, - cache_creation_tokens, - cache_read_tokens, - context_limit, - compressed, - total_cost_micros, - runtime_ms -) -SELECT - @chat_id::uuid, - NULLIF(UNNEST(@created_by::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), - NULLIF(UNNEST(@model_config_id::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), - NULLIF(UNNEST(@reasoning_effort::text[]), '')::chat_reasoning_effort, - UNNEST(@role::chat_message_role[]), - UNNEST(@content::text[])::jsonb, - UNNEST(@content_version::smallint[]), - UNNEST(@visibility::chat_message_visibility[]), - NULLIF(UNNEST(@input_tokens::bigint[]), 0), - NULLIF(UNNEST(@output_tokens::bigint[]), 0), - NULLIF(UNNEST(@total_tokens::bigint[]), 0), - NULLIF(UNNEST(@reasoning_tokens::bigint[]), 0), - NULLIF(UNNEST(@cache_creation_tokens::bigint[]), 0), - NULLIF(UNNEST(@cache_read_tokens::bigint[]), 0), - NULLIF(UNNEST(@context_limit::bigint[]), 0), - UNNEST(@compressed::boolean[]), - NULLIF(UNNEST(@total_cost_micros::bigint[]), 0), - NULLIF(UNNEST(@runtime_ms::bigint[]), 0) -RETURNING - *; +SELECT * +FROM inserted +ORDER BY id; -- name: UpdateChatByID :one WITH updated_chat AS ( @@ -1948,6 +1974,8 @@ SET created_at = ( WHERE target.id = @target_id AND target.chat_id = @chat_id; -- name: GetLastChatMessageByRole :one +-- The returned id becomes both an AfterID cursor and last_read_message_id, so +-- "last" must use id order. SELECT * FROM @@ -1957,7 +1985,7 @@ WHERE AND role = @role::chat_message_role AND deleted = false ORDER BY - created_at DESC, id DESC + id DESC LIMIT 1; diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index f37f082cca..8627dd4f69 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -5333,7 +5333,7 @@ func TestGetChatUserPrompts(t *testing.T) { if deleted { require.NoError(t, db.SoftDeleteChatMessageByID(dbauthz.AsSystemRestricted(ctx), msgs[0].ID)) } - return msgs[0] + return database.ChatMessage(msgs[0]) } t.Run("NewestFirstFiltering", func(t *testing.T) { diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index baf9753462..6fcd08b9ce 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -8182,7 +8182,7 @@ func insertChatMessageParts( messages, err := db.InsertChatMessages(ctx, params) require.NoError(t, err) require.Len(t, messages, 1) - return messages[0] + return database.ChatMessage(messages[0]) } func createPlanSubagentChatWithHistory( diff --git a/coderd/x/chatd/chatstate/messages.go b/coderd/x/chatd/chatstate/messages.go index b867c1f05e..13e9e122ca 100644 --- a/coderd/x/chatd/chatstate/messages.go +++ b/coderd/x/chatd/chatstate/messages.go @@ -95,6 +95,17 @@ func toInsertParams(chatID uuid.UUID, messages []Message) database.InsertChatMes return params } +// fromInsertedRows converts the rows returned by `InsertChatMessages`, which +// sqlc types separately because the query wraps the insert in a CTE. The +// conversion stops compiling if the row ever stops matching ChatMessage. +func fromInsertedRows(rows []database.InsertChatMessagesRow) []database.ChatMessage { + messages := make([]database.ChatMessage, len(rows)) + for i, row := range rows { + messages[i] = database.ChatMessage(row) + } + return messages +} + func nullUUIDOrNil(u uuid.NullUUID) uuid.UUID { if u.Valid { return u.UUID diff --git a/coderd/x/chatd/chatstate/transitions.go b/coderd/x/chatd/chatstate/transitions.go index 6b8593eff8..5fbb4eeba8 100644 --- a/coderd/x/chatd/chatstate/transitions.go +++ b/coderd/x/chatd/chatstate/transitions.go @@ -111,7 +111,7 @@ func CreateChat( } result = CreateChatResult{ Chat: refreshed, - InitialMessages: inserted, + InitialMessages: fromInsertedRows(inserted), } if err := buffer.Publish( coderdpubsub.ChatStateUpdateChannel(refreshed.ID), @@ -182,7 +182,7 @@ func (tx *Tx) insertMessages(messages []Message) ([]database.ChatMessage, error) if err != nil { return nil, xerrors.Errorf("insert messages: %w", err) } - return inserted, nil + return fromInsertedRows(inserted), nil } // clearQueue deletes all queued messages on the chat and returns the diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index 98df7f1f0c..0a93275ee2 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -7,7 +7,6 @@ import ( "errors" "fmt" "slices" - "sort" "strings" "time" @@ -1654,13 +1653,6 @@ func latestSubagentAssistantMessage( return "", xerrors.Errorf("get chat messages: %w", err) } - sort.Slice(messages, func(i, j int) bool { - if messages[i].CreatedAt.Equal(messages[j].CreatedAt) { - return messages[i].ID < messages[j].ID - } - return messages[i].CreatedAt.Before(messages[j].CreatedAt) - }) - for i := len(messages) - 1; i >= 0; i-- { message := messages[i] if message.Role != database.ChatMessageRoleAssistant ||