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 ||