From 6697f8ae2febb2474bc6863631a5706fc98889cd Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 11 Sep 2026 14:48:00 +0000 Subject: [PATCH 1/6] test(coderd): guard chat history writes in test databases Every chat_messages and chat_queued_messages write must run in the transaction that allocated a snapshot for the chat, which production does through chatstate.CreateChat and ChatMachine.Update. Nothing enforced this for tests, so fixtures wrote history rows the runner could not attribute to a snapshot. dbtestutil.NewDB now wraps the store in a guard that rejects the thirteen writer methods unless InsertChat or LockChatAndBumpSnapshotVersion ran earlier in the same transaction, and fails the test at cleanup for every rejection even when the caller dropped the error. dbgen.ChatMessage and the new dbgen.ChatQueuedMessage allocate inside their own transaction, and the direct seeds in chatd, chatstate, coderd, database and telemetry tests move to them or to an allocating InTx around the writer under test. A completeness test parses the generated queries so a new writer cannot be added without the guard learning about it. Raw-SQL fixtures are unchanged and unguarded by design. --- coderd/database/dbgen/dbgen.go | 73 ++++-- coderd/database/dbtestutil/chatwriteguard.go | 224 ++++++++++++++++ coderd/database/dbtestutil/db.go | 16 ++ .../database/dbtestutil/db_internal_test.go | 143 +++++++++++ coderd/database/querier_test.go | 242 ++++++++---------- coderd/exp_chats_chatstate_test.go | 12 +- coderd/exp_chats_test.go | 177 +++++-------- coderd/telemetry/telemetry_test.go | 7 +- coderd/x/chatd/chatd_helpers_test.go | 63 +---- coderd/x/chatd/chatd_test.go | 44 ++-- coderd/x/chatd/chatstate/family_test.go | 8 +- .../chatstate/transitions_helpers_test.go | 4 +- coderd/x/chatd/chatstate/transitions_test.go | 3 +- coderd/x/chatd/chatstate/trigger_test.go | 73 ++++-- coderd/x/chatd/hooks_test.go | 31 ++- coderd/x/chatd/synthetickey_internal_test.go | 3 +- 16 files changed, 728 insertions(+), 395 deletions(-) create mode 100644 coderd/database/dbtestutil/chatwriteguard.go diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index b6a8f82867c..9dec91b6f23 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -115,6 +115,8 @@ func Chat(t testing.TB, db database.Store, seed database.Chat) database.Chat { return chat } +// ChatMessage inserts one chat message. It allocates a snapshot for the chat +// in the same transaction, as every chat history write must. func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) database.ChatMessage { t.Helper() @@ -124,30 +126,63 @@ func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) dat } role := takeFirst(seed.Role, database.ChatMessageRoleUser) - msgs, err := db.InsertChatMessages(genCtx, database.InsertChatMessagesParams{ - ChatID: seed.ChatID, - CreatedBy: []uuid.UUID{seed.CreatedBy.UUID}, - ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID}, - ReasoningEffort: []string{string(seed.ReasoningEffort.ChatReasoningEffort)}, - Role: []database.ChatMessageRole{role}, - Content: []string{content}, - ContentVersion: []int16{takeFirst(seed.ContentVersion, chatprompt.CurrentContentVersion)}, - Visibility: []database.ChatMessageVisibility{takeFirst(seed.Visibility, database.ChatMessageVisibilityBoth)}, - InputTokens: []int64{seed.InputTokens.Int64}, - OutputTokens: []int64{seed.OutputTokens.Int64}, - TotalTokens: []int64{seed.TotalTokens.Int64}, - ReasoningTokens: []int64{seed.ReasoningTokens.Int64}, - CacheCreationTokens: []int64{seed.CacheCreationTokens.Int64}, - CacheReadTokens: []int64{seed.CacheReadTokens.Int64}, - ContextLimit: []int64{seed.ContextLimit.Int64}, - Compressed: []bool{seed.Compressed}, - RuntimeMs: []int64{seed.RuntimeMs.Int64}, - }) + var msgs []database.InsertChatMessagesRow + err := db.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(genCtx, seed.ChatID); err != nil { + return xerrors.Errorf("allocate chat snapshot: %w", err) + } + var err error + msgs, err = tx.InsertChatMessages(genCtx, database.InsertChatMessagesParams{ + ChatID: seed.ChatID, + CreatedBy: []uuid.UUID{seed.CreatedBy.UUID}, + ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID}, + ReasoningEffort: []string{string(seed.ReasoningEffort.ChatReasoningEffort)}, + Role: []database.ChatMessageRole{role}, + Content: []string{content}, + ContentVersion: []int16{takeFirst(seed.ContentVersion, chatprompt.CurrentContentVersion)}, + Visibility: []database.ChatMessageVisibility{takeFirst(seed.Visibility, database.ChatMessageVisibilityBoth)}, + InputTokens: []int64{seed.InputTokens.Int64}, + OutputTokens: []int64{seed.OutputTokens.Int64}, + TotalTokens: []int64{seed.TotalTokens.Int64}, + ReasoningTokens: []int64{seed.ReasoningTokens.Int64}, + CacheCreationTokens: []int64{seed.CacheCreationTokens.Int64}, + CacheReadTokens: []int64{seed.CacheReadTokens.Int64}, + ContextLimit: []int64{seed.ContextLimit.Int64}, + Compressed: []bool{seed.Compressed}, + RuntimeMs: []int64{seed.RuntimeMs.Int64}, + }) + return err + }, nil) require.NoError(t, err, "insert chat message") require.Len(t, msgs, 1) return database.ChatMessage(msgs[0]) } +// ChatQueuedMessage inserts one queued chat message. It allocates a snapshot +// for the chat in the same transaction, as every chat queue write must. +// CreatedBy defaults to the chat owner. +func ChatQueuedMessage(t testing.TB, db database.Store, seed database.ChatQueuedMessage) database.ChatQueuedMessage { + t.Helper() + + var queued database.ChatQueuedMessage + err := db.InTx(func(tx database.Store) error { + chat, err := tx.LockChatAndBumpSnapshotVersion(genCtx, seed.ChatID) + if err != nil { + return xerrors.Errorf("allocate chat snapshot: %w", err) + } + queued, err = tx.InsertChatQueuedMessageWithCreator(genCtx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: seed.ChatID, + Content: takeFirstSlice(seed.Content, json.RawMessage("[]")), + ModelConfigID: seed.ModelConfigID, + ReasoningEffort: seed.ReasoningEffort, + CreatedBy: takeFirst(seed.CreatedBy, chat.OwnerID), + }) + return err + }, nil) + require.NoError(t, err, "insert chat queued message") + return queued +} + const ( // Match the default OpenAI test model's effective context settings. defaultChatModelContextLimit int64 = 128000 diff --git a/coderd/database/dbtestutil/chatwriteguard.go b/coderd/database/dbtestutil/chatwriteguard.go new file mode 100644 index 00000000000..4bf773637a2 --- /dev/null +++ b/coderd/database/dbtestutil/chatwriteguard.go @@ -0,0 +1,224 @@ +package dbtestutil + +import ( + "context" + "database/sql" + "errors" + "sync" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" +) + +// chatWriteRejection records one guarded write that ran without a snapshot +// allocation for its chat. +type chatWriteRejection struct { + method string + chatID uuid.UUID +} + +// chatWriteRecorder collects every rejection made by the guards that share +// it. NewDB owns one recorder per test database and fails the test at +// cleanup for each recorded rejection, so a rejection fails the test even +// when the caller drops the returned error. +type chatWriteRecorder struct { + mu sync.Mutex + rejections []chatWriteRejection +} + +// add records a rejection and returns the error the guarded method hands +// back to its caller. +func (r *chatWriteRecorder) add(method string, chatID uuid.UUID) error { + r.mu.Lock() + r.rejections = append(r.rejections, chatWriteRejection{method: method, chatID: chatID}) + r.mu.Unlock() + return xerrors.Errorf("%s for chat %s outside a chat state transition (no snapshot allocated in this transaction); "+ + "write history through chatstate.ChatMachine.Update or CreateChat, or dbgen in tests", method, chatID) +} + +// list returns a copy of the recorded rejections. +func (r *chatWriteRecorder) list() []chatWriteRejection { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]chatWriteRejection, len(r.rejections)) + copy(out, r.rejections) + return out +} + +// reset discards the recorded rejections. +func (r *chatWriteRecorder) reset() { + r.mu.Lock() + defer r.mu.Unlock() + r.rejections = nil +} + +// chatWriteGuard is a database.Store that rejects writes to chat_messages +// and chat_queued_messages unless the same transaction allocated a snapshot +// for the chat first, through InsertChat or LockChatAndBumpSnapshotVersion. +// Production satisfies this by construction: chatstate.CreateChat inserts +// the chat row and chatstate.ChatMachine.Update bumps the snapshot version +// before any history write. Test fixtures satisfy it through dbgen, which +// allocates inside its own transaction. +// +// The root handle never allocates, so every guarded write on it is +// rejected. Each transaction started through InTx gets its own allocation +// set; a nested InTx that reuses the outer transaction shares the outer set. +// The guard wraps the store returned by database.New directly, so the +// nested check compares the transaction store by pointer identity. +type chatWriteGuard struct { + database.Store + rec *chatWriteRecorder + // allocated is nil on the root handle and holds the chats whose snapshot + // this transaction allocated otherwise. + allocated map[uuid.UUID]struct{} +} + +func newChatWriteGuard(store database.Store, rec *chatWriteRecorder) *chatWriteGuard { + return &chatWriteGuard{Store: store, rec: rec} +} + +func (g *chatWriteGuard) Wrappers() []string { + return append(g.Store.Wrappers(), "dbtestutil.chatWriteGuard") +} + +func (g *chatWriteGuard) InTx(fn func(database.Store) error, opts *database.TxOptions) error { + return g.Store.InTx(func(tx database.Store) error { + if tx == g.Store { + return fn(g) + } + return fn(&chatWriteGuard{Store: tx, rec: g.rec, allocated: map[uuid.UUID]struct{}{}}) + }, opts) +} + +func (g *chatWriteGuard) mark(chatID uuid.UUID) { + if g.allocated != nil { + g.allocated[chatID] = struct{}{} + } +} + +// require returns the rejection error unless this transaction allocated a +// snapshot for chatID. +func (g *chatWriteGuard) require(method string, chatID uuid.UUID) error { + if _, ok := g.allocated[chatID]; ok { + return nil + } + return g.rec.add(method, chatID) +} + +func (g *chatWriteGuard) InsertChat(ctx context.Context, arg database.InsertChatParams) (database.Chat, error) { + chat, err := g.Store.InsertChat(ctx, arg) + if err == nil { + g.mark(chat.ID) + } + return chat, err +} + +func (g *chatWriteGuard) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) { + chat, err := g.Store.LockChatAndBumpSnapshotVersion(ctx, id) + if err == nil { + g.mark(id) + } + return chat, err +} + +func (g *chatWriteGuard) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.InsertChatMessagesRow, error) { + if err := g.require("InsertChatMessages", arg.ChatID); err != nil { + return nil, err + } + return g.Store.InsertChatMessages(ctx, arg) +} + +// SoftDeleteChatMessageByID resolves the chat from the message because the +// parameters carry no chat id. A missing message passes through unchecked +// because the underlying update affects no row. +func (g *chatWriteGuard) SoftDeleteChatMessageByID(ctx context.Context, id int64) error { + msg, err := g.GetChatMessageByID(ctx, id) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("resolve chat for message %d: %w", id, err) + } + if err == nil { + if err := g.require("SoftDeleteChatMessageByID", msg.ChatID); err != nil { + return err + } + } + return g.Store.SoftDeleteChatMessageByID(ctx, id) +} + +func (g *chatWriteGuard) SoftDeleteChatMessagesAfterID(ctx context.Context, arg database.SoftDeleteChatMessagesAfterIDParams) error { + if err := g.require("SoftDeleteChatMessagesAfterID", arg.ChatID); err != nil { + return err + } + return g.Store.SoftDeleteChatMessagesAfterID(ctx, arg) +} + +func (g *chatWriteGuard) SoftDeleteContextFileMessages(ctx context.Context, chatID uuid.UUID) error { + if err := g.require("SoftDeleteContextFileMessages", chatID); err != nil { + return err + } + return g.Store.SoftDeleteContextFileMessages(ctx, chatID) +} + +func (g *chatWriteGuard) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.ChatQueuedMessage, error) { + if err := g.require("InsertChatQueuedMessage", arg.ChatID); err != nil { + return database.ChatQueuedMessage{}, err + } + return g.Store.InsertChatQueuedMessage(ctx, arg) +} + +func (g *chatWriteGuard) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) { + if err := g.require("InsertChatQueuedMessageWithCreator", arg.ChatID); err != nil { + return database.ChatQueuedMessage{}, err + } + return g.Store.InsertChatQueuedMessageWithCreator(ctx, arg) +} + +func (g *chatWriteGuard) DeleteChatQueuedMessage(ctx context.Context, arg database.DeleteChatQueuedMessageParams) error { + if err := g.require("DeleteChatQueuedMessage", arg.ChatID); err != nil { + return err + } + return g.Store.DeleteChatQueuedMessage(ctx, arg) +} + +func (g *chatWriteGuard) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg database.DeleteChatQueuedMessageReturningCountParams) (int64, error) { + if err := g.require("DeleteChatQueuedMessageReturningCount", arg.ChatID); err != nil { + return 0, err + } + return g.Store.DeleteChatQueuedMessageReturningCount(ctx, arg) +} + +func (g *chatWriteGuard) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error { + if err := g.require("DeleteAllChatQueuedMessages", chatID); err != nil { + return err + } + return g.Store.DeleteAllChatQueuedMessages(ctx, chatID) +} + +func (g *chatWriteGuard) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) { + if err := g.require("DeleteAllChatQueuedMessagesReturningCount", chatID); err != nil { + return 0, err + } + return g.Store.DeleteAllChatQueuedMessagesReturningCount(ctx, chatID) +} + +func (g *chatWriteGuard) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { + if err := g.require("PopNextQueuedMessage", chatID); err != nil { + return database.ChatQueuedMessage{}, err + } + return g.Store.PopNextQueuedMessage(ctx, chatID) +} + +func (g *chatWriteGuard) ReorderChatQueuedMessageToFront(ctx context.Context, arg database.ReorderChatQueuedMessageToFrontParams) (int64, error) { + if err := g.require("ReorderChatQueuedMessageToFront", arg.ChatID); err != nil { + return 0, err + } + return g.Store.ReorderChatQueuedMessageToFront(ctx, arg) +} + +func (g *chatWriteGuard) ReorderChatQueuedMessageToHead(ctx context.Context, arg database.ReorderChatQueuedMessageToHeadParams) (int64, error) { + if err := g.require("ReorderChatQueuedMessageToHead", arg.ChatID); err != nil { + return 0, err + } + return g.Store.ReorderChatQueuedMessageToHead(ctx, arg) +} diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go index d25f2508e40..4e799588bc5 100644 --- a/coderd/database/dbtestutil/db.go +++ b/coderd/database/dbtestutil/db.go @@ -89,6 +89,13 @@ func NowInDefaultTimezone() time.Time { return time.Now().In(loc).Round(time.Microsecond) } +// NewDB opens a PostgreSQL database for the test and returns its store and +// pubsub. The store rejects writes to chat_messages and chat_queued_messages +// that run without a snapshot allocation in the same transaction and fails +// the test at cleanup for every rejection, whether or not the caller +// propagated the error. Seed chat history and queue rows through +// dbgen.ChatMessage and dbgen.ChatQueuedMessage, which allocate a snapshot +// inside their own transaction. func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) { t.Helper() @@ -135,6 +142,15 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) { } // Unit tests should not retry serial transaction failures. db = database.New(sqlDB, database.WithSerialRetryCount(1)) + rejections := &chatWriteRecorder{} + db = newChatWriteGuard(db, rejections) + t.Cleanup(func() { + for _, r := range rejections.list() { + t.Errorf("chat write guard rejected %s for chat %s outside a chat state transition; "+ + "the call site is the assertion that received this error, otherwise search for callers of %s", + r.method, r.chatID, r.method) + } + }) ps, err = pubsub.New(context.Background(), o.logger, sqlDB, connectionURL) require.NoError(t, err) diff --git a/coderd/database/dbtestutil/db_internal_test.go b/coderd/database/dbtestutil/db_internal_test.go index fb4d71b5652..8a1b44a4be5 100644 --- a/coderd/database/dbtestutil/db_internal_test.go +++ b/coderd/database/dbtestutil/db_internal_test.go @@ -1,9 +1,20 @@ package dbtestutil import ( + "go/ast" + "go/parser" + "go/token" + "os" + "regexp" + "slices" "testing" + "github.com/google/uuid" "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/testutil" ) // Recent pg_dump versions (13.22+ / 14.19+ / 15.14+ / 16.10+ / 17.6+) emit @@ -30,3 +41,135 @@ CREATE TABLE foo; require.NotContains(t, out, `\unrestrict`, `normalizeDump must strip \unrestrict psql meta-command`) require.Contains(t, out, "CREATE TABLE foo;", "normalizeDump must preserve real SQL between the meta-commands") } + +// The guard installed by NewDB must reject chat history and queue writes on +// the root handle and inside a transaction that has not allocated a snapshot +// for the chat, record each rejection, and let the same write through once +// the transaction has allocated. +func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { + t.Parallel() + + db, _ := NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + require.Contains(t, db.Wrappers(), "dbtestutil.chatWriteGuard") + + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + messageParams := database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{user.ID}, + ModelConfigID: []uuid.UUID{model.ID}, + ReasoningEffort: []string{""}, + Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, + Content: []string{`[{"type":"text","text":"hello"}]`}, + ContentVersion: []int16{1}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + RuntimeMs: []int64{0}, + } + + _, err := db.InsertChatMessages(ctx, messageParams) + require.ErrorContains(t, err, "InsertChatMessages for chat "+chat.ID.String()+" outside a chat state transition") + + _, err = db.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: chat.ID, + Content: []byte(`[]`), + CreatedBy: user.ID, + }) + require.ErrorContains(t, err, "InsertChatQueuedMessageWithCreator for chat "+chat.ID.String()+" outside a chat state transition") + + guard, ok := db.(*chatWriteGuard) + require.True(t, ok, "NewDB must return the chat write guard") + require.Equal(t, []chatWriteRejection{ + {method: "InsertChatMessages", chatID: chat.ID}, + {method: "InsertChatQueuedMessageWithCreator", chatID: chat.ID}, + }, guard.rec.list()) + // The rejections above are the expected outcome of this test, not + // failures to report at cleanup. + guard.rec.reset() + + err = db.InTx(func(tx database.Store) error { + _, err := tx.InsertChatMessages(ctx, messageParams) + return err + }, nil) + require.ErrorContains(t, err, "InsertChatMessages for chat "+chat.ID.String()+" outside a chat state transition") + require.Equal(t, []chatWriteRejection{{method: "InsertChatMessages", chatID: chat.ID}}, guard.rec.list()) + guard.rec.reset() + + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.Empty(t, messages, "rejected writes must not reach the database") + + err = db.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, chat.ID); err != nil { + return err + } + _, err := tx.InsertChatMessages(ctx, messageParams) + return err + }, nil) + require.NoError(t, err) + require.Empty(t, guard.rec.list()) +} + +// Every generated query that inserts, updates or deletes rows of +// chat_messages or chat_queued_messages must be overridden by the guard, +// except the two search_tsv maintenance queries, which change only the +// search columns. A new writer query fails this test until the guard learns +// about it. +func TestChatWriteGuardCoversEveryChatHistoryWriter(t *testing.T) { + t.Parallel() + + queries, err := os.ReadFile("../queries.sql.go") + require.NoError(t, err) + queryConst := regexp.MustCompile("(?s)\nconst \\w+ = `-- name: (\\w+) :\\w+\n([^`]*)`") + writesChatTables := regexp.MustCompile(`(?i)\b(?:INSERT\s+INTO|UPDATE|DELETE\s+FROM)\s+(?:chat_messages|chat_queued_messages)\b`) + var writers []string + for _, match := range queryConst.FindAllSubmatch(queries, -1) { + if writesChatTables.Match(match[2]) { + writers = append(writers, string(match[1])) + } + } + require.NotEmpty(t, writers, "no writer queries found; the query constant pattern no longer matches queries.sql.go") + + file, err := parser.ParseFile(token.NewFileSet(), "chatwriteguard.go", nil, parser.SkipObjectResolution) + require.NoError(t, err) + notWriters := map[string]bool{ + "InTx": true, + "Wrappers": true, + "InsertChat": true, + "LockChatAndBumpSnapshotVersion": true, + } + covered := []string{"BackfillChatMessagesSearchTsv", "ReindexStaleChatMessagesSearchTsv"} + for _, decl := range file.Decls { + fn, ok := decl.(*ast.FuncDecl) + if !ok || fn.Recv == nil || len(fn.Recv.List) != 1 || !fn.Name.IsExported() { + continue + } + star, ok := fn.Recv.List[0].Type.(*ast.StarExpr) + if !ok { + continue + } + if recv, ok := star.X.(*ast.Ident); !ok || recv.Name != "chatWriteGuard" || notWriters[fn.Name.Name] { + continue + } + covered = append(covered, fn.Name.Name) + } + + slices.Sort(writers) + slices.Sort(covered) + require.Equal(t, writers, covered, + "queries writing chat_messages or chat_queued_messages must be guarded by chatWriteGuard or listed here as search_tsv maintenance") +} diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index babef6c49ec..bbe238fb81d 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -13072,24 +13072,30 @@ func TestInsertChatMessages(t *testing.T) { insertMessage := func(t *testing.T, store database.Store, ctx context.Context, chatID, userID, modelConfigID uuid.UUID, content string) { t.Helper() - _, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{userID}, - ModelConfigID: []uuid.UUID{modelConfigID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - Content: []string{fmt.Sprintf("%q", content)}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, - }) + err := store.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, chatID); err != nil { + return err + } + _, err := tx.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + ChatID: chatID, + CreatedBy: []uuid.UUID{userID}, + ModelConfigID: []uuid.UUID{modelConfigID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + Content: []string{fmt.Sprintf("%q", content)}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + RuntimeMs: []int64{0}, + }) + return err + }, nil) require.NoError(t, err) } @@ -13132,24 +13138,32 @@ func TestInsertChatMessages(t *testing.T) { t.Parallel() store, ctx, user, chat, _, modelConfigA := setupChat(t) - msgs, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{user.ID, uuid.Nil, uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfigA.ID, modelConfigA.ID, modelConfigA.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser, database.ChatMessageRoleAssistant, database.ChatMessageRoleTool}, - ContentVersion: []int16{chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth}, - Content: []string{`"hello"`, `"response"`, `"tool result"`}, - InputTokens: []int64{10, 0, 0}, - OutputTokens: []int64{0, 20, 0}, - TotalTokens: []int64{10, 20, 0}, - ReasoningTokens: []int64{0, 5, 0}, - CacheCreationTokens: []int64{0, 0, 0}, - CacheReadTokens: []int64{0, 0, 0}, - ContextLimit: []int64{0, 0, 0}, - Compressed: []bool{false, false, false}, - RuntimeMs: []int64{0, 500, 0}, - }) + var msgs []database.InsertChatMessagesRow + err := store.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, chat.ID); err != nil { + return err + } + var err error + msgs, err = tx.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{user.ID, uuid.Nil, uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfigA.ID, modelConfigA.ID, modelConfigA.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleUser, database.ChatMessageRoleAssistant, database.ChatMessageRoleTool}, + ContentVersion: []int16{chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth}, + Content: []string{`"hello"`, `"response"`, `"tool result"`}, + InputTokens: []int64{10, 0, 0}, + OutputTokens: []int64{0, 20, 0}, + TotalTokens: []int64{10, 20, 0}, + ReasoningTokens: []int64{0, 5, 0}, + CacheCreationTokens: []int64{0, 0, 0}, + CacheReadTokens: []int64{0, 0, 0}, + ContextLimit: []int64{0, 0, 0}, + Compressed: []bool{false, false, false}, + RuntimeMs: []int64{0, 500, 0}, + }) + return err + }, nil) require.NoError(t, err) require.Len(t, msgs, 3) @@ -13204,24 +13218,32 @@ func insertChatMessagesInvertedTimestamps(t *testing.T, db database.Store, sqlDB }) 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), - RuntimeMs: make([]int64, count), - }) + var inserted []database.InsertChatMessagesRow + err := db.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, chat.ID); err != nil { + return err + } + var err error + inserted, err = tx.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), + RuntimeMs: make([]int64, count), + }) + return err + }, nil) require.NoError(t, err) require.Len(t, inserted, count) @@ -13375,26 +13397,13 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { content string, ) database.ChatMessage { t.Helper() - results, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{uuid.Nil}, - Role: []database.ChatMessageRole{role}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{vis}, - Compressed: []bool{compressed}, - Content: []string{`"` + content + `"`}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - RuntimeMs: []int64{0}, - }) - require.NoError(t, err) - return database.ChatMessage(results[0]) + return dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + Role: role, + Visibility: vis, + Compressed: compressed, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`"` + content + `"`), Valid: true}, + }) } invertCreatedAt := func(t *testing.T, chatID uuid.UUID) { @@ -18315,25 +18324,13 @@ func TestGetChatsFilter(t *testing.T) { makeUnread := func(chatID uuid.UUID) { t.Helper() - _, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{user.ID}, - ModelConfigID: []uuid.UUID{modelCfg.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - Content: []string{`[{"type":"text","text":"hello"}]`}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, + dbgen.ChatMessage(t, store, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"hello"}]`), Valid: true}, }) - require.NoError(t, err) } markRead := func(chatID uuid.UUID) { @@ -18559,27 +18556,15 @@ func TestGetChatsSearch(t *testing.T) { insertMsg := func(chatID uuid.UUID, role database.ChatMessageRole, visibility database.ChatMessageVisibility, text string) database.ChatMessage { t.Helper() - msgs, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{user.ID}, - ModelConfigID: []uuid.UUID{modelCfg.ID}, - Role: []database.ChatMessageRole{role}, - Content: []string{`[{"type":"text","text":` + strconv.Quote(text) + `}]`}, - ContentVersion: []int16{1}, - Visibility: []database.ChatMessageVisibility{visibility}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, - }) - require.NoError(t, err) - require.Len(t, msgs, 1) - return database.ChatMessage(msgs[0]) + return dbgen.ChatMessage(t, store, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: role, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":` + strconv.Quote(text) + `}]`), Valid: true}, + ContentVersion: 1, + Visibility: visibility, + }) } linkPR := func(chatID uuid.UUID, url, state, prTitle string, prNumber int32, gitRemoteOrigin string) { @@ -18647,7 +18632,12 @@ func TestGetChatsSearch(t *testing.T) { // Soft-deleted rows stay excluded even though search_tsv remains // populated. - err = store.SoftDeleteChatMessageByID(ctx, deletedMsg.ID) + err = store.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, deletedMsgChat.ID); err != nil { + return err + } + return tx.SoftDeleteChatMessageByID(ctx, deletedMsg.ID) + }, nil) require.NoError(t, err) // Inserted after backfill: search_tsv IS NULL, must match nothing. @@ -18815,25 +18805,13 @@ func TestChatHasUnread(t *testing.T) { // Helper to insert a single chat message. insertMsg := func(role database.ChatMessageRole, text string) { t.Helper() - _, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{user.ID}, - ModelConfigID: []uuid.UUID{modelCfg.ID}, - Role: []database.ChatMessageRole{role}, - Content: []string{fmt.Sprintf(`[{"type":"text","text":%q}]`, text)}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, + dbgen.ChatMessage(t, store, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: role, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(fmt.Sprintf(`[{"type":"text","text":%q}]`, text)), Valid: true}, }) - require.NoError(t, err) } // Insert an assistant message: becomes unread. diff --git a/coderd/exp_chats_chatstate_test.go b/coderd/exp_chats_chatstate_test.go index 10746f57e8f..5f44b79bc47 100644 --- a/coderd/exp_chats_chatstate_test.go +++ b/coderd/exp_chats_chatstate_test.go @@ -754,18 +754,16 @@ func driveChatToInvalidWaitingWithQueue( // Seed the queue with one row attributed to the chat owner. The // content is a minimal valid JSON payload; only the row's - // presence matters for ClassifyExecutionState. The owner_id is - // filled from the chat row by the SQL. + // presence matters for ClassifyExecutionState. dbgen fills + // created_by from the chat row. rawContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - _, err = api.Database.InsertChatQueuedMessage(sysCtx, database.InsertChatQueuedMessageParams{ - ChatID: chatID, - Content: rawContent.RawMessage, - ModelConfigID: uuid.NullUUID{}, + dbgen.ChatQueuedMessage(t, api.Database, database.ChatQueuedMessage{ + ChatID: chatID, + Content: rawContent.RawMessage, }) - require.NoError(t, err) // Flip the chat's status to waiting via a raw execution-state // update. This bypasses the transition matrix to produce the diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index c38e2a5e443..6ee402b5136 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -167,7 +167,7 @@ func insertTestChatQueuedMessage( } func insertTestChatQueuedMessageWithReasoningEffort( - ctx context.Context, + _ context.Context, t testing.TB, db database.Store, chatID uuid.UUID, @@ -177,17 +177,12 @@ func insertTestChatQueuedMessageWithReasoningEffort( ) database.ChatQueuedMessage { t.Helper() - queued, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chatID, - Content: content, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, - ReasoningEffort: database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffort(reasoningEffort), Valid: reasoningEffort != ""}, - }, - ) - require.NoError(t, err) - return queued + return dbgen.ChatQueuedMessage(t, db, database.ChatQueuedMessage{ + ChatID: chatID, + Content: content, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, + ReasoningEffort: database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffort(reasoningEffort), Valid: reasoningEffort != ""}, + }) } // findUserMessage returns the first user-role message from a slice of chat @@ -6636,30 +6631,24 @@ func TestGetChatUserPrompts(t *testing.T) { t.Helper() content, err := chatprompt.MarshalParts(parts) require.NoError(t, err) - msgs, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{userID}, - ModelConfigID: []uuid.UUID{modelConfigID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(content.RawMessage)}, - Visibility: []database.ChatMessageVisibility{visibility}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, - }) - require.NoError(t, err) - require.Len(t, msgs, 1) + msg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: visibility, + }) if deleted { - require.NoError(t, db.SoftDeleteChatMessageByID(dbauthz.AsSystemRestricted(ctx), msgs[0].ID)) + err := db.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(dbauthz.AsSystemRestricted(ctx), chatID); err != nil { + return err + } + return tx.SoftDeleteChatMessageByID(dbauthz.AsSystemRestricted(ctx), msg.ID) + }, nil) + require.NoError(t, err) } - return database.ChatMessage(msgs[0]) + return msg } t.Run("NewestFirstFiltering", func(t *testing.T) { @@ -6717,50 +6706,46 @@ func TestGetChatUserPrompts(t *testing.T) { {Type: codersdk.ChatMessagePartTypeText, Text: "assistant reply"}, }) require.NoError(t, err) - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{user.UserID}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.UserID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, }) - require.NoError(t, err) // Legacy V0 user message stored as a scalar JSON string // (predates migration 000434). The jsonb_typeof guard in // GetChatUserPromptsByChatID must silently exclude this row; // without the guard, jsonb_array_elements would raise // "cannot extract elements from a scalar" and the request - // would 500. - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{user.UserID}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - ContentVersion: []int16{chatprompt.ContentVersionV0}, - Content: []string{`"plain text from V0"`}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, - }) + // would 500. dbgen.ChatMessage defaults content_version to the + // current version, so the V0 row is inserted directly inside + // an allocating transaction. + err = db.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(dbauthz.AsSystemRestricted(ctx), chat.ID); err != nil { + return err + } + _, err := tx.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{user.UserID}, + ModelConfigID: []uuid.UUID{modelConfig.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, + ContentVersion: []int16{chatprompt.ContentVersionV0}, + Content: []string{`"plain text from V0"`}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + RuntimeMs: []int64{0}, + }) + return err + }, nil) require.NoError(t, err) // Soft-deleted prompt; must not appear. @@ -6937,25 +6922,13 @@ func TestGetChatUserPrompts(t *testing.T) { {Type: codersdk.ChatMessagePartTypeText, Text: "assistant reply"}, }) require.NoError(t, err) - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: assistantOnlyChat.ID, - CreatedBy: []uuid.UUID{user.UserID}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: assistantOnlyChat.ID, + CreatedBy: uuid.NullUUID{UUID: user.UserID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, }) - require.NoError(t, err) resp, err = client.GetChatPrompts(ctx, assistantOnlyChat.ID, nil) require.NoError(t, err) @@ -11297,12 +11270,11 @@ func TestClearChat(t *testing.T) { codersdk.ChatMessageText("queued follow-up"), }) require.NoError(t, err) - _, err = db.InsertChatQueuedMessageWithCreator(dbauthz.AsSystemRestricted(ctx), database.InsertChatQueuedMessageWithCreatorParams{ + dbgen.ChatQueuedMessage(t, db, database.ChatQueuedMessage{ ChatID: chat.ID, Content: queuedContent.RawMessage, CreatedBy: user.UserID, }) - require.NoError(t, err) // No waiting-with-queue state exists, so a synchronous clear // from E1 is rejected; the user deletes or promotes the queue @@ -12371,25 +12343,12 @@ func TestPromoteChatQueuedMessage(t *testing.T) { }}) require.NoError(t, err) - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, }) - require.NoError(t, err) _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, diff --git a/coderd/telemetry/telemetry_test.go b/coderd/telemetry/telemetry_test.go index d37f54bd787..8a9291d0a31 100644 --- a/coderd/telemetry/telemetry_test.go +++ b/coderd/telemetry/telemetry_test.go @@ -1164,7 +1164,12 @@ func TestChatsTelemetry(t *testing.T) { ContextLimit: sql.NullInt64{Int64: 200000, Valid: true}, RuntimeMs: sql.NullInt64{Int64: 999999, Valid: true}, }) - err = db.SoftDeleteChatMessageByID(ctx, poisonMsg.ID) + err = db.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, rootChat.ID); err != nil { + return err + } + return tx.SoftDeleteChatMessageByID(ctx, poisonMsg.ID) + }, nil) require.NoError(t, err) _, snapshot := collectSnapshot(ctx, t, db, nil) diff --git a/coderd/x/chatd/chatd_helpers_test.go b/coderd/x/chatd/chatd_helpers_test.go index 70df2a5a72c..05c34db238a 100644 --- a/coderd/x/chatd/chatd_helpers_test.go +++ b/coderd/x/chatd/chatd_helpers_test.go @@ -12,7 +12,6 @@ import ( "charm.land/fantasy" fantasyanthropic "charm.land/fantasy/providers/anthropic" "github.com/google/uuid" - "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/database" @@ -85,36 +84,8 @@ func anthropicRequestBody(t *testing.T, req chattest.AnthropicRequest) string { return string(data) } -func singleChatMessageInsertParams( - chatID uuid.UUID, - role database.ChatMessageRole, - content pqtype.NullRawMessage, - modelConfigID uuid.UUID, - createdBy uuid.UUID, -) database.InsertChatMessagesParams { - return database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{createdBy}, - ModelConfigID: []uuid.UUID{modelConfigID}, - ReasoningEffort: []string{""}, - Role: []database.ChatMessageRole{role}, - Content: []string{string(content.RawMessage)}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, - } -} - func insertSystemTextMessage( - ctx context.Context, + _ context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -124,18 +95,15 @@ func insertSystemTextMessage( t.Helper() content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) require.NoError(t, err) - params := singleChatMessageInsertParams( - chatID, - database.ChatMessageRoleSystem, - content, - modelID, - uuid.Nil, - ) - _, err = db.InsertChatMessages(ctx, params) - require.NoError(t, err) + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + Role: database.ChatMessageRoleSystem, + Content: content, + }) } -func insertOrphanProviderToolCall(ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, modelID uuid.UUID) { +func insertOrphanProviderToolCall(_ context.Context, t *testing.T, db database.Store, chatID uuid.UUID, modelID uuid.UUID) { t.Helper() reasoningMetadata, err := json.Marshal(fantasy.ProviderMetadata{ fantasyanthropic.Name: &fantasyanthropic.ReasoningOptionMetadata{RedactedData: "redacted-payload"}, @@ -157,15 +125,12 @@ func insertOrphanProviderToolCall(ctx context.Context, t *testing.T, db database } content, err := chatprompt.MarshalParts(parts) require.NoError(t, err) - params := singleChatMessageInsertParams( - chatID, - database.ChatMessageRoleAssistant, - content, - modelID, - uuid.Nil, - ) - _, err = db.InsertChatMessages(ctx, params) - require.NoError(t, err) + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: content, + }) } func createChatThroughServer( diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 488f0f6e793..8545b0eb6f8 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -5242,14 +5242,13 @@ func TestActiveServer_RoutingPreservesAPIKeyAfterCompaction(t *testing.T) { ContextFileDirectory: "/home/coder/project", }}) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, singleChatMessageInsertParams( - chat.ID, - database.ChatMessageRoleUser, - contextContent, - model.ID, - user.ID, - )) - require.NoError(t, err) + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: contextContent, + }) _ = newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(factory) @@ -8674,7 +8673,7 @@ func insertProviderToolPairMessageWithLocalTool( } func insertChatMessageParts( - ctx context.Context, + _ context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -8686,17 +8685,13 @@ func insertChatMessageParts( t.Helper() content, err := chatprompt.MarshalParts(parts) require.NoError(t, err) - params := singleChatMessageInsertParams( - chatID, - role, - content, - modelID, - createdBy, - ) - messages, err := db.InsertChatMessages(ctx, params) - require.NoError(t, err) - require.Len(t, messages, 1) - return database.ChatMessage(messages[0]) + return dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: createdBy != uuid.Nil}, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + Role: role, + Content: content, + }) } func createPlanSubagentChatWithHistory( @@ -13498,7 +13493,7 @@ func TestQueuedPromotionResolvesOrganizationModel(t *testing.T) { } func insertQueuedMessage( - ctx context.Context, + _ context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -13509,14 +13504,12 @@ func insertQueuedMessage( t.Helper() content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) require.NoError(t, err) - queued, err := db.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + return dbgen.ChatQueuedMessage(t, db, database.ChatQueuedMessage{ ChatID: chatID, Content: content.RawMessage, ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, CreatedBy: createdBy, }) - require.NoError(t, err) - return queued } func TestPromoteQueuedPreservesReasoningEffort(t *testing.T) { @@ -13537,14 +13530,13 @@ func TestPromoteQueuedPreservesReasoningEffort(t *testing.T) { }) content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}) require.NoError(t, err) - queued, err := db.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + queued := dbgen.ChatQueuedMessage(t, db, database.ChatQueuedMessage{ ChatID: chat.ID, Content: content.RawMessage, ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, ReasoningEffort: database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffortHigh, Valid: true}, CreatedBy: user.ID, }) - require.NoError(t, err) require.True(t, queued.ReasoningEffort.Valid) require.Equal(t, database.ChatReasoningEffortHigh, queued.ReasoningEffort.ChatReasoningEffort) diff --git a/coderd/x/chatd/chatstate/family_test.go b/coderd/x/chatd/chatstate/family_test.go index 7fcbd18310f..395d9f618c2 100644 --- a/coderd/x/chatd/chatstate/family_test.go +++ b/coderd/x/chatd/chatstate/family_test.go @@ -131,12 +131,10 @@ func TestSetFamilyArchivedRejectsInvalidStateEvenWhenAlreadyDesired(t *testing.T codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: child.ID, - Content: rawContent.RawMessage, - ModelConfigID: uuid.NullUUID{}, + dbgen.ChatQueuedMessage(t, db, database.ChatQueuedMessage{ + ChatID: child.ID, + Content: rawContent.RawMessage, }) - require.NoError(t, err) pub := newRecordingPubsub() _, err = chatstate.SetFamilyArchived(ctx, db, pub, chatstate.SetFamilyArchivedInput{ diff --git a/coderd/x/chatd/chatstate/transitions_helpers_test.go b/coderd/x/chatd/chatstate/transitions_helpers_test.go index c44e5af1647..57b1828c6cd 100644 --- a/coderd/x/chatd/chatstate/transitions_helpers_test.go +++ b/coderd/x/chatd/chatstate/transitions_helpers_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatstate" @@ -509,13 +510,12 @@ func seedA1WithMixedOutstandingToolCalls(t *testing.T, f *testFixture, queuedExt for i := range queuedExtras { body := fmt.Sprintf("queued-%s-%d", namePrefix, i) createdBy := uuid.New() - queued, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + queued := dbgen.ChatQueuedMessage(t, f.DB, database.ChatQueuedMessage{ ChatID: created.Chat.ID, Content: userMessageContent(t, body), ModelConfigID: uuid.NullUUID{UUID: f.Model.ID, Valid: true}, CreatedBy: createdBy, }) - require.NoError(t, err) queuedIDs = append(queuedIDs, queued.ID) queuedBodies = append(queuedBodies, body) queuedCreatedBy = append(queuedCreatedBy, createdBy) diff --git a/coderd/x/chatd/chatstate/transitions_test.go b/coderd/x/chatd/chatstate/transitions_test.go index 31dbe23459a..6336fc3e00d 100644 --- a/coderd/x/chatd/chatstate/transitions_test.go +++ b/coderd/x/chatd/chatstate/transitions_test.go @@ -229,12 +229,11 @@ func promoteQueuedMessageWithModel( ctx := testutil.Context(t, testutil.WaitShort) created := createTestChat(t, f) message := userTextMessage("queued model resolution", f.User.ID, modelConfigID) - queued, err := f.DB.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ + queued := dbgen.ChatQueuedMessage(t, f.DB, database.ChatQueuedMessage{ ChatID: created.Chat.ID, Content: message.Content.RawMessage, ModelConfigID: message.ModelConfigID, }) - require.NoError(t, err) machine := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { _, err := tx.FinishError(chatstate.FinishErrorInput{ diff --git a/coderd/x/chatd/chatstate/trigger_test.go b/coderd/x/chatd/chatstate/trigger_test.go index 004e1704648..9635af9432b 100644 --- a/coderd/x/chatd/chatstate/trigger_test.go +++ b/coderd/x/chatd/chatstate/trigger_test.go @@ -350,15 +350,21 @@ func TestQueueInsertUpdatesQueueVersion(t *testing.T) { require.NoError(t, err) require.Equal(t, int64(0), before.QueueVersion) - bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) - require.NoError(t, err) - content := userMessageContent(t, "queued") - _, err = f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ - ChatID: created.Chat.ID, - Content: content, - CreatedBy: f.User.ID, - }) + var bumped database.Chat + err = f.DB.InTx(func(tx database.Store) error { + var err error + bumped, err = tx.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + if err != nil { + return err + } + _, err = tx.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: created.Chat.ID, + Content: content, + CreatedBy: f.User.ID, + }) + return err + }, nil) require.NoError(t, err) after, err := f.DB.GetChatByID(ctx, created.Chat.ID) @@ -392,10 +398,18 @@ func TestLegacyQueuedMessageInsertUsesChatOwnerAsCreator(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) created := createTestChat(t, f) - queued, err := f.DB.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: created.Chat.ID, - Content: userMessageContent(t, "legacy-queued"), - }) + var queued database.ChatQueuedMessage + err := f.DB.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID); err != nil { + return err + } + var err error + queued, err = tx.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ + ChatID: created.Chat.ID, + Content: userMessageContent(t, "legacy-queued"), + }) + return err + }, nil) require.NoError(t, err) require.Equal(t, created.Chat.OwnerID, queued.CreatedBy) } @@ -410,12 +424,11 @@ func TestQueueUpdateContentUpdatesQueueVersion(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) created := createTestChat(t, f) - queued, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + queued := dbgen.ChatQueuedMessage(t, f.DB, database.ChatQueuedMessage{ ChatID: created.Chat.ID, Content: userMessageContent(t, "initial"), CreatedBy: f.User.ID, }) - require.NoError(t, err) before, err := f.DB.GetChatByID(ctx, created.Chat.ID) require.NoError(t, err) @@ -445,18 +458,16 @@ func TestQueueUpdatePositionUpdatesQueueVersion(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) created := createTestChat(t, f) - q1, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + q1 := dbgen.ChatQueuedMessage(t, f.DB, database.ChatQueuedMessage{ ChatID: created.Chat.ID, Content: userMessageContent(t, "first"), CreatedBy: f.User.ID, }) - require.NoError(t, err) - q2, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + q2 := dbgen.ChatQueuedMessage(t, f.DB, database.ChatQueuedMessage{ ChatID: created.Chat.ID, Content: userMessageContent(t, "second"), CreatedBy: f.User.ID, }) - require.NoError(t, err) require.NotEqual(t, q1.ID, q2.ID) bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) @@ -483,20 +494,28 @@ func TestQueueDeleteUpdatesQueueVersion(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) created := createTestChat(t, f) - queued, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + queued := dbgen.ChatQueuedMessage(t, f.DB, database.ChatQueuedMessage{ ChatID: created.Chat.ID, Content: userMessageContent(t, "to delete"), CreatedBy: f.User.ID, }) - require.NoError(t, err) - - bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) - require.NoError(t, err) - rows, err := f.DB.DeleteChatQueuedMessageReturningCount(ctx, database.DeleteChatQueuedMessageReturningCountParams{ - ID: queued.ID, - ChatID: created.Chat.ID, - }) + var ( + bumped database.Chat + rows int64 + ) + err := f.DB.InTx(func(tx database.Store) error { + var err error + bumped, err = tx.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + if err != nil { + return err + } + rows, err = tx.DeleteChatQueuedMessageReturningCount(ctx, database.DeleteChatQueuedMessageReturningCountParams{ + ID: queued.ID, + ChatID: created.Chat.ID, + }) + return err + }, nil) require.NoError(t, err) require.Equal(t, int64(1), rows) diff --git a/coderd/x/chatd/hooks_test.go b/coderd/x/chatd/hooks_test.go index 52410998f42..98c4bd38e18 100644 --- a/coderd/x/chatd/hooks_test.go +++ b/coderd/x/chatd/hooks_test.go @@ -448,11 +448,13 @@ func TestEditMessageUserPromptSubmitHook(t *testing.T) { }) content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}) require.NoError(t, err) - inserted, err := db.InsertChatMessages(ctx, singleChatMessageInsertParams( - chat.ID, database.ChatMessageRoleUser, content, model.ID, user.ID, - )) - require.NoError(t, err) - require.Len(t, inserted, 1) + inserted := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: content, + }) type receivedHook struct { request agenthooks.Request claims agenthooks.Claims @@ -490,7 +492,7 @@ func TestEditMessageUserPromptSubmitHook(t *testing.T) { result, err := server.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, - EditedMessageID: inserted[0].ID, + EditedMessageID: inserted.ID, Content: []codersdk.ChatMessagePart{ reference, codersdk.ChatMessageText("edited original"), @@ -627,15 +629,17 @@ func TestPromptHooksAdmissionPreflight(t *testing.T) { content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}) require.NoError(t, err) - inserted, err := db.InsertChatMessages(ctx, singleChatMessageInsertParams( - chat.ID, database.ChatMessageRoleUser, content, model.ID, user.ID, - )) - require.NoError(t, err) - require.Len(t, inserted, 1) + inserted := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: content, + }) _, err = server.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, - EditedMessageID: inserted[0].ID, + EditedMessageID: inserted.ID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("bad model edit")}, ModelConfigID: uuid.New(), }) @@ -650,13 +654,12 @@ func TestPromptHooksAdmissionPreflight(t *testing.T) { queuedContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}) require.NoError(t, err) for range chatstate.MaxQueueSize { - _, err = db.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + dbgen.ChatQueuedMessage(t, db, database.ChatQueuedMessage{ ChatID: busy.ID, Content: queuedContent.RawMessage, ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, CreatedBy: user.ID, }) - require.NoError(t, err) } _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: busy.ID, diff --git a/coderd/x/chatd/synthetickey_internal_test.go b/coderd/x/chatd/synthetickey_internal_test.go index 5c18d783aaf..41336b7f9c8 100644 --- a/coderd/x/chatd/synthetickey_internal_test.go +++ b/coderd/x/chatd/synthetickey_internal_test.go @@ -223,12 +223,11 @@ func TestSyntheticAPIKeyDeletionDoesNotMutateChatState(t *testing.T) { ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, Role: database.ChatMessageRoleUser, }) - _, err = db.InsertChatQueuedMessage(t.Context(), database.InsertChatQueuedMessageParams{ + dbgen.ChatQueuedMessage(t, db, database.ChatQueuedMessage{ ChatID: chat.ID, Content: json.RawMessage(`[]`), ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, }) - require.NoError(t, err) before, err := db.GetChatByID(t.Context(), chat.ID) require.NoError(t, err) From 5387935db4849bf52c1e667fbc43a77626b6557d Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 11 Sep 2026 15:12:12 +0000 Subject: [PATCH 2/6] test(coderd): seed every chat fixture through dbgen The seed helpers stopped using their context once they moved to dbgen, so the parameter now misstates what they do; drop it and its callers' argument. TestUpdateChatLastTurnSummary and TestUpdateChatSummary build their store without NewDB and were still bumping the snapshot and inserting the message in two autocommit statements, the shape the guard exists to reject; they seed through dbgen.ChatMessage like every other fixture. The canary no longer re-proves the allocating path that every dbgen seed in the suite already exercises. --- .../database/dbtestutil/db_internal_test.go | 13 +---- coderd/database/querier_test.go | 52 +++++-------------- coderd/exp_chats_test.go | 26 +++++----- coderd/x/chatd/chatd_helpers_test.go | 3 +- coderd/x/chatd/chatd_test.go | 28 +++++----- 5 files changed, 39 insertions(+), 83 deletions(-) diff --git a/coderd/database/dbtestutil/db_internal_test.go b/coderd/database/dbtestutil/db_internal_test.go index 8a1b44a4be5..ce2a10609ae 100644 --- a/coderd/database/dbtestutil/db_internal_test.go +++ b/coderd/database/dbtestutil/db_internal_test.go @@ -44,8 +44,7 @@ CREATE TABLE foo; // The guard installed by NewDB must reject chat history and queue writes on // the root handle and inside a transaction that has not allocated a snapshot -// for the chat, record each rejection, and let the same write through once -// the transaction has allocated. +// for the chat, and record each rejection. func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { t.Parallel() @@ -112,16 +111,6 @@ func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) require.NoError(t, err) require.Empty(t, messages, "rejected writes must not reach the database") - - err = db.InTx(func(tx database.Store) error { - if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, chat.ID); err != nil { - return err - } - _, err := tx.InsertChatMessages(ctx, messageParams) - return err - }, nil) - require.NoError(t, err) - require.Empty(t, guard.rec.list()) } // Every generated query that inserts, updates or deletes rows of diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index bbe238fb81d..c16d039a48a 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -16420,27 +16420,13 @@ func TestUpdateChatLastTurnSummary(t *testing.T) { require.Equal(t, sql.NullString{String: "still fresh summary", Valid: true}, fetched.LastTurnSummary) require.Equal(t, advanced.UpdatedAt, fetched.UpdatedAt) - _, err = db.LockChatAndBumpSnapshotVersion(ctx, chat.ID) - require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{owner.ID}, - ModelConfigID: []uuid.UUID{modelCfg.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - Content: []string{`[{"type":"text","text":"new request"}]`}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"new request"}]`), Valid: true}, }) - require.NoError(t, err) affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{ ID: chat.ID, @@ -16541,27 +16527,13 @@ func TestUpdateChatSummary(t *testing.T) { require.NoError(t, err) require.EqualValues(t, 1, affected) - _, err = db.LockChatAndBumpSnapshotVersion(ctx, chat.ID) - require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{owner.ID}, - ModelConfigID: []uuid.UUID{modelCfg.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - Content: []string{`[{"type":"text","text":"new request"}]`}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - RuntimeMs: []int64{0}, + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"new request"}]`), Valid: true}, }) - require.NoError(t, err) affected, err = db.UpdateChatSummary(ctx, database.UpdateChatSummaryParams{ ID: chat.ID, diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 6ee402b5136..f966c7e3bed 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -155,7 +155,6 @@ func newChatClientWithoutAIBridge(t testing.TB, overrides ...func(*coderdtest.Op } func insertTestChatQueuedMessage( - ctx context.Context, t testing.TB, db database.Store, chatID uuid.UUID, @@ -163,11 +162,10 @@ func insertTestChatQueuedMessage( modelConfigID uuid.UUID, ) database.ChatQueuedMessage { t.Helper() - return insertTestChatQueuedMessageWithReasoningEffort(ctx, t, db, chatID, content, modelConfigID, "") + return insertTestChatQueuedMessageWithReasoningEffort(t, db, chatID, content, modelConfigID, "") } func insertTestChatQueuedMessageWithReasoningEffort( - _ context.Context, t testing.TB, db database.Store, chatID uuid.UUID, @@ -12032,7 +12030,7 @@ func TestDeleteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("queued message for delete route"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, deleteContent, modelConfig.ID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, deleteContent, modelConfig.ID) res, err := client.Request( ctx, @@ -12113,7 +12111,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText(queuedText), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, chat.LastModelConfigID) promoteRes, err := client.Request( ctx, @@ -12177,7 +12175,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("require a local model for promotion"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, foreignConfig.ID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, foreignConfig.ID) promoteRes, err := memberClient.Request( ctx, @@ -12253,7 +12251,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("queued message promoted by member"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, chat.LastModelConfigID) promoteRes, err := memberClient.Request( ctx, @@ -12285,7 +12283,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, chat.LastModelConfigID) // Archive the chat. _, err = db.ArchiveChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) @@ -12361,7 +12359,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText(queuedText), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, chat.LastModelConfigID) promoteRes, err := client.Request( ctx, @@ -12454,7 +12452,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("running-promote"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, chat.LastModelConfigID) promoteRes, err := client.Request( ctx, @@ -17971,7 +17969,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - _ = insertTestChatQueuedMessage(ctx, t, db, chatID, content, modelConfigID) + _ = insertTestChatQueuedMessage(t, db, chatID, content, modelConfigID) } t.Run("NoCursorReturnsAllDESCPlusQueued", func(t *testing.T) { @@ -18426,7 +18424,7 @@ func TestChatReadOnlySharedWriteHandlers(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, chat.LastModelConfigID) res, err := sharedClient.Request( ctx, @@ -18461,7 +18459,7 @@ func TestChatReadOnlySharedWriteHandlers(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, chat.LastModelConfigID) res, err := sharedClient.Request( ctx, @@ -18614,7 +18612,7 @@ func TestChatOwnerOnlyWriteHandlers(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID) + queuedMessage := insertTestChatQueuedMessage(t, db, chat.ID, queuedContent, chat.LastModelConfigID) // Org admin tries to promote. promoteRes, err := adminClient.Request( diff --git a/coderd/x/chatd/chatd_helpers_test.go b/coderd/x/chatd/chatd_helpers_test.go index 05c34db238a..e69912b0ed3 100644 --- a/coderd/x/chatd/chatd_helpers_test.go +++ b/coderd/x/chatd/chatd_helpers_test.go @@ -85,7 +85,6 @@ func anthropicRequestBody(t *testing.T, req chattest.AnthropicRequest) string { } func insertSystemTextMessage( - _ context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -103,7 +102,7 @@ func insertSystemTextMessage( }) } -func insertOrphanProviderToolCall(_ context.Context, t *testing.T, db database.Store, chatID uuid.UUID, modelID uuid.UUID) { +func insertOrphanProviderToolCall(t *testing.T, db database.Store, chatID uuid.UUID, modelID uuid.UUID) { t.Helper() reasoningMetadata, err := json.Marshal(fantasy.ProviderMetadata{ fantasyanthropic.Name: &fantasyanthropic.ReasoningOptionMetadata{RedactedData: "redacted-payload"}, diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 8545b0eb6f8..557db593dfd 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6528,7 +6528,7 @@ func TestActiveServer_BasicAssistantGenerationAndPromptPreparation(t *testing.T) }) chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) - insertSystemTextMessage(ctx, t, db, chat.ID, "sys-2", model.ID) + insertSystemTextMessage(t, db, chat.ID, "sys-2", model.ID) insertAssistantTextMessage(ctx, t, db, chat.ID, "working", model.ID) _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, @@ -7581,7 +7581,7 @@ func TestActiveServer_AnthropicSanitizesProviderToolBeforeRequest(t *testing.T) }) chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search for coder") waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) - insertOrphanProviderToolCall(ctx, t, db, chat.ID, model.ID) + insertOrphanProviderToolCall(t, db, chat.ID, model.ID) _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, @@ -8614,7 +8614,7 @@ func insertAssistantTextMessage( modelID uuid.UUID, ) { t.Helper() - insertChatMessageParts(ctx, t, db, chatID, database.ChatMessageRoleAssistant, modelID, uuid.Nil, []codersdk.ChatMessagePart{ + insertChatMessageParts(t, db, chatID, database.ChatMessageRoleAssistant, modelID, uuid.Nil, []codersdk.ChatMessagePart{ codersdk.ChatMessageText(text), }) } @@ -8661,8 +8661,8 @@ func insertProviderToolPairMessageWithLocalTool( ToolName: "read_file", Args: json.RawMessage(`{"path":"main.go"}`), }) - insertChatMessageParts(ctx, t, db, chatID, database.ChatMessageRoleAssistant, modelID, uuid.Nil, parts) - insertChatMessageParts(ctx, t, db, chatID, database.ChatMessageRoleTool, modelID, uuid.Nil, []codersdk.ChatMessagePart{ + insertChatMessageParts(t, db, chatID, database.ChatMessageRoleAssistant, modelID, uuid.Nil, parts) + insertChatMessageParts(t, db, chatID, database.ChatMessageRoleTool, modelID, uuid.Nil, []codersdk.ChatMessagePart{ { Type: codersdk.ChatMessagePartTypeToolResult, ToolCallID: "tc-1", @@ -8673,7 +8673,6 @@ func insertProviderToolPairMessageWithLocalTool( } func insertChatMessageParts( - _ context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -8725,8 +8724,8 @@ func createPlanSubagentChatWithHistory( MCPServerIDs: []uuid.UUID{}, ClientType: database.ChatClientTypeApi, }) - insertSystemTextMessage(ctx, t, db, chat.ID, "You are not currently connected to a workspace.", modelID) - insertChatMessageParts(ctx, t, db, chat.ID, database.ChatMessageRoleUser, modelID, userID, []codersdk.ChatMessagePart{ + insertSystemTextMessage(t, db, chat.ID, "You are not currently connected to a workspace.", modelID) + insertChatMessageParts(t, db, chat.ID, database.ChatMessageRoleUser, modelID, userID, []codersdk.ChatMessagePart{ codersdk.ChatMessageText("hello"), }) return chat @@ -13325,7 +13324,7 @@ func TestQueuedCompletionResolvesOrganizationModel(t *testing.T) { foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, localDefault.ID, "first") testutil.TryReceive(ctx, t, firstRunStarted) - insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + insertQueuedMessage(t, db, chat.ID, user.ID, foreignModel.ID, "queued") close(allowFirstRunFinish) waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) @@ -13378,7 +13377,7 @@ func TestQueuedCompletionResolvesOrganizationModel(t *testing.T) { foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, localDefault.ID, "first") testutil.TryReceive(ctx, t, streamStarted) - insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + insertQueuedMessage(t, db, chat.ID, user.ID, foreignModel.ID, "queued") current, err := db.GetChatByID(ctx, chat.ID) require.NoError(t, err) _, err = server.InterruptChat(ctx, current) @@ -13420,7 +13419,7 @@ func TestQueuedPromotionResolvesOrganizationModel(t *testing.T) { Title: "resolve explicit queued promotion", Status: database.ChatStatusError, }) - queued := insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + queued := insertQueuedMessage(t, db, chat.ID, user.ID, foreignModel.ID, "queued") result, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ ChatID: chat.ID, @@ -13445,7 +13444,7 @@ func TestQueuedPromotionResolvesOrganizationModel(t *testing.T) { Title: "resolve error send promotion", Status: database.ChatStatusError, }) - insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + insertQueuedMessage(t, db, chat.ID, user.ID, foreignModel.ID, "queued") result, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, @@ -13475,7 +13474,7 @@ func TestQueuedPromotionResolvesOrganizationModel(t *testing.T) { Title: "keep unresolved queued row", Status: database.ChatStatusError, }) - queued := insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + queued := insertQueuedMessage(t, db, chat.ID, user.ID, foreignModel.ID, "queued") _, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ ChatID: chat.ID, @@ -13493,7 +13492,6 @@ func TestQueuedPromotionResolvesOrganizationModel(t *testing.T) { } func insertQueuedMessage( - _ context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -14367,7 +14365,7 @@ func TestProviderSwitchSanitizesAndRestoresPEToolHistory(t *testing.T) { }) require.NoError(t, err) waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) - insertChatMessageParts(ctx, t, db, chat.ID, database.ChatMessageRoleAssistant, mA.ID, uuid.Nil, + insertChatMessageParts(t, db, chat.ID, database.ChatMessageRoleAssistant, mA.ID, uuid.Nil, []codersdk.ChatMessagePart{ { Type: codersdk.ChatMessagePartTypeToolCall, From ef0434a50dd144464ecefddb16fefb83558217c6 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 11 Sep 2026 15:55:23 +0000 Subject: [PATCH 3/6] test(coderd): lock the chat guard's allocation set A database.Store handle, including a transaction handle, may be used from more than one goroutine, so the per-transaction allocation set needs a mutex like the recorder already has. Two more chatd seed helpers had kept a context parameter they no longer use; it goes the same way as the others. --- coderd/database/dbtestutil/chatwriteguard.go | 11 +++++++++-- coderd/x/chatd/chatd_test.go | 8 +++----- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/coderd/database/dbtestutil/chatwriteguard.go b/coderd/database/dbtestutil/chatwriteguard.go index 4bf773637a2..355713b7c3a 100644 --- a/coderd/database/dbtestutil/chatwriteguard.go +++ b/coderd/database/dbtestutil/chatwriteguard.go @@ -71,7 +71,9 @@ type chatWriteGuard struct { database.Store rec *chatWriteRecorder // allocated is nil on the root handle and holds the chats whose snapshot - // this transaction allocated otherwise. + // this transaction allocated otherwise. mu guards it because a Store, + // including a transaction handle, may be used from several goroutines. + mu sync.Mutex allocated map[uuid.UUID]struct{} } @@ -93,6 +95,8 @@ func (g *chatWriteGuard) InTx(fn func(database.Store) error, opts *database.TxOp } func (g *chatWriteGuard) mark(chatID uuid.UUID) { + g.mu.Lock() + defer g.mu.Unlock() if g.allocated != nil { g.allocated[chatID] = struct{}{} } @@ -101,7 +105,10 @@ func (g *chatWriteGuard) mark(chatID uuid.UUID) { // require returns the rejection error unless this transaction allocated a // snapshot for chatID. func (g *chatWriteGuard) require(method string, chatID uuid.UUID) error { - if _, ok := g.allocated[chatID]; ok { + g.mu.Lock() + _, ok := g.allocated[chatID] + g.mu.Unlock() + if ok { return nil } return g.rec.add(method, chatID) diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 557db593dfd..1f0b12e8009 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6529,7 +6529,7 @@ func TestActiveServer_BasicAssistantGenerationAndPromptPreparation(t *testing.T) chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) insertSystemTextMessage(t, db, chat.ID, "sys-2", model.ID) - insertAssistantTextMessage(ctx, t, db, chat.ID, "working", model.ID) + insertAssistantTextMessage(t, db, chat.ID, "working", model.ID) _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, @@ -7688,7 +7688,7 @@ func TestActiveServer_AnthropicProviderToolPreRequestGuard(t *testing.T) { }) chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search") waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) - insertProviderToolPairMessageWithLocalTool(ctx, t, db, chat.ID, model.ID, "ws-allowed") + insertProviderToolPairMessageWithLocalTool(t, db, chat.ID, model.ID, "ws-allowed") _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, @@ -7724,7 +7724,7 @@ func TestActiveServer_AnthropicProviderToolPreRequestGuard(t *testing.T) { }) chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search and read") waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) - insertProviderToolPairMessageWithLocalTool(ctx, t, db, chat.ID, model.ID, "ws-disabled") + insertProviderToolPairMessageWithLocalTool(t, db, chat.ID, model.ID, "ws-disabled") _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, @@ -8606,7 +8606,6 @@ func updateChatModelCallConfig(t *testing.T, db database.Store, model database.C } func insertAssistantTextMessage( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -8620,7 +8619,6 @@ func insertAssistantTextMessage( } func insertProviderToolPairMessageWithLocalTool( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, From e784d9b24e6954084828974253d96b7d82e16e13 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 11 Sep 2026 16:15:11 +0000 Subject: [PATCH 4/6] test(coderd): prove the chat guard reports at cleanup The cleanup report is what fails a test whose caller dropped the rejection error, and no test exercised it: deleting the block left the guard tests green. A fake testing.TB captures NewDB's cleanups and Errorf calls so the report can be asserted directly. The soft-delete pass-through comment now names the real mechanism (GetChatMessageByID excludes deleted rows; the trigger treats an identical rewrite as a no-op), and the completeness failure message describes both directions of a mismatch. --- coderd/database/dbtestutil/chatwriteguard.go | 7 ++- .../database/dbtestutil/db_internal_test.go | 56 ++++++++++++++++++- 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/coderd/database/dbtestutil/chatwriteguard.go b/coderd/database/dbtestutil/chatwriteguard.go index 355713b7c3a..59f3c7eed30 100644 --- a/coderd/database/dbtestutil/chatwriteguard.go +++ b/coderd/database/dbtestutil/chatwriteguard.go @@ -138,8 +138,11 @@ func (g *chatWriteGuard) InsertChatMessages(ctx context.Context, arg database.In } // SoftDeleteChatMessageByID resolves the chat from the message because the -// parameters carry no chat id. A missing message passes through unchecked -// because the underlying update affects no row. +// parameters carry no chat id. GetChatMessageByID excludes deleted rows, so +// a missing or already deleted id passes through unchecked; the update then +// matches no row or rewrites a row with identical values, which the BEFORE +// UPDATE trigger on chat_messages treats as a no-op without assigning a +// revision or bumping history_version. func (g *chatWriteGuard) SoftDeleteChatMessageByID(ctx context.Context, id int64) error { msg, err := g.GetChatMessageByID(ctx, id) if err != nil && !errors.Is(err, sql.ErrNoRows) { diff --git a/coderd/database/dbtestutil/db_internal_test.go b/coderd/database/dbtestutil/db_internal_test.go index ce2a10609ae..bf40c360644 100644 --- a/coderd/database/dbtestutil/db_internal_test.go +++ b/coderd/database/dbtestutil/db_internal_test.go @@ -1,6 +1,7 @@ package dbtestutil import ( + "fmt" "go/ast" "go/parser" "go/token" @@ -160,5 +161,58 @@ func TestChatWriteGuardCoversEveryChatHistoryWriter(t *testing.T) { slices.Sort(writers) slices.Sort(covered) require.Equal(t, writers, covered, - "queries writing chat_messages or chat_queued_messages must be guarded by chatWriteGuard or listed here as search_tsv maintenance") + "every query writing chat_messages or chat_queued_messages must have a chatWriteGuard override "+ + "or be listed here as search_tsv maintenance, and every chatWriteGuard override must be one of "+ + "those writers or be listed in notWriters") +} + +// rejectionReportingTB records the Cleanup functions and Errorf calls that +// NewDB makes so a test can run the cleanup itself and inspect what it +// reported. Every other testing.TB method goes to the embedded test. +type rejectionReportingTB struct { + *testing.T + cleanups []func() + errors []string +} + +func (tb *rejectionReportingTB) Cleanup(f func()) { + tb.cleanups = append(tb.cleanups, f) +} + +func (tb *rejectionReportingTB) Errorf(format string, args ...any) { + tb.errors = append(tb.errors, fmt.Sprintf(format, args...)) +} + +// A rejection must fail the test at cleanup even when the caller drops the +// returned error. +func TestChatWriteGuardReportsRejectionsAtCleanup(t *testing.T) { + t.Parallel() + + tb := &rejectionReportingTB{T: t} + db, _ := NewDB(tb) + ctx := testutil.Context(t, testutil.WaitShort) + + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + + _, _ = db.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: chat.ID, + Content: []byte(`[]`), + CreatedBy: user.ID, + }) + + // NewDB registers its cleanups on tb, so they run here instead of at the + // end of the test, in the order testing would use. + for i := len(tb.cleanups) - 1; i >= 0; i-- { + tb.cleanups[i]() + } + + require.Len(t, tb.errors, 1) + require.Contains(t, tb.errors[0], "chat write guard rejected InsertChatQueuedMessageWithCreator for chat "+chat.ID.String()) } From 6326a0aade8b8eeeb79a462d38663eaca4190236 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 11 Sep 2026 17:16:20 +0000 Subject: [PATCH 5/6] test(coderd/database/dbtestutil): drive every guarded writer to a rejection The completeness test proved an override exists per writer query but not that it rejects; only two of the thirteen were ever driven to the guard error. A table now calls each writer on the root handle and the completeness test requires the table to name every override, so a future override that drops the check fails a test instead of passing silently. A nested InTx case covers the one branch that shares the outer allocation set. Also widen the writer query pattern to ONLY and MERGE INTO forms, and rename require, mark, covered and rejections to say what they hold. --- coderd/database/dbtestutil/chatwriteguard.go | 46 +++--- coderd/database/dbtestutil/db.go | 6 +- .../database/dbtestutil/db_internal_test.go | 154 ++++++++++++++---- 3 files changed, 148 insertions(+), 58 deletions(-) diff --git a/coderd/database/dbtestutil/chatwriteguard.go b/coderd/database/dbtestutil/chatwriteguard.go index 59f3c7eed30..2a403040e45 100644 --- a/coderd/database/dbtestutil/chatwriteguard.go +++ b/coderd/database/dbtestutil/chatwriteguard.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "errors" + "slices" "sync" "github.com/google/uuid" @@ -42,9 +43,7 @@ func (r *chatWriteRecorder) add(method string, chatID uuid.UUID) error { func (r *chatWriteRecorder) list() []chatWriteRejection { r.mu.Lock() defer r.mu.Unlock() - out := make([]chatWriteRejection, len(r.rejections)) - copy(out, r.rejections) - return out + return slices.Clone(r.rejections) } // reset discards the recorded rejections. @@ -94,7 +93,10 @@ func (g *chatWriteGuard) InTx(fn func(database.Store) error, opts *database.TxOp }, opts) } -func (g *chatWriteGuard) mark(chatID uuid.UUID) { +// markAllocated records that this transaction allocated a snapshot for +// chatID. The root handle has no allocation set, so the call is a no-op +// there and its guarded writes stay rejected. +func (g *chatWriteGuard) markAllocated(chatID uuid.UUID) { g.mu.Lock() defer g.mu.Unlock() if g.allocated != nil { @@ -102,9 +104,9 @@ func (g *chatWriteGuard) mark(chatID uuid.UUID) { } } -// require returns the rejection error unless this transaction allocated a -// snapshot for chatID. -func (g *chatWriteGuard) require(method string, chatID uuid.UUID) error { +// requireSnapshot returns the rejection error unless this transaction +// allocated a snapshot for chatID. +func (g *chatWriteGuard) requireSnapshot(method string, chatID uuid.UUID) error { g.mu.Lock() _, ok := g.allocated[chatID] g.mu.Unlock() @@ -117,7 +119,7 @@ func (g *chatWriteGuard) require(method string, chatID uuid.UUID) error { func (g *chatWriteGuard) InsertChat(ctx context.Context, arg database.InsertChatParams) (database.Chat, error) { chat, err := g.Store.InsertChat(ctx, arg) if err == nil { - g.mark(chat.ID) + g.markAllocated(chat.ID) } return chat, err } @@ -125,13 +127,13 @@ func (g *chatWriteGuard) InsertChat(ctx context.Context, arg database.InsertChat func (g *chatWriteGuard) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) { chat, err := g.Store.LockChatAndBumpSnapshotVersion(ctx, id) if err == nil { - g.mark(id) + g.markAllocated(id) } return chat, err } func (g *chatWriteGuard) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.InsertChatMessagesRow, error) { - if err := g.require("InsertChatMessages", arg.ChatID); err != nil { + if err := g.requireSnapshot("InsertChatMessages", arg.ChatID); err != nil { return nil, err } return g.Store.InsertChatMessages(ctx, arg) @@ -149,7 +151,7 @@ func (g *chatWriteGuard) SoftDeleteChatMessageByID(ctx context.Context, id int64 return xerrors.Errorf("resolve chat for message %d: %w", id, err) } if err == nil { - if err := g.require("SoftDeleteChatMessageByID", msg.ChatID); err != nil { + if err := g.requireSnapshot("SoftDeleteChatMessageByID", msg.ChatID); err != nil { return err } } @@ -157,77 +159,77 @@ func (g *chatWriteGuard) SoftDeleteChatMessageByID(ctx context.Context, id int64 } func (g *chatWriteGuard) SoftDeleteChatMessagesAfterID(ctx context.Context, arg database.SoftDeleteChatMessagesAfterIDParams) error { - if err := g.require("SoftDeleteChatMessagesAfterID", arg.ChatID); err != nil { + if err := g.requireSnapshot("SoftDeleteChatMessagesAfterID", arg.ChatID); err != nil { return err } return g.Store.SoftDeleteChatMessagesAfterID(ctx, arg) } func (g *chatWriteGuard) SoftDeleteContextFileMessages(ctx context.Context, chatID uuid.UUID) error { - if err := g.require("SoftDeleteContextFileMessages", chatID); err != nil { + if err := g.requireSnapshot("SoftDeleteContextFileMessages", chatID); err != nil { return err } return g.Store.SoftDeleteContextFileMessages(ctx, chatID) } func (g *chatWriteGuard) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.ChatQueuedMessage, error) { - if err := g.require("InsertChatQueuedMessage", arg.ChatID); err != nil { + if err := g.requireSnapshot("InsertChatQueuedMessage", arg.ChatID); err != nil { return database.ChatQueuedMessage{}, err } return g.Store.InsertChatQueuedMessage(ctx, arg) } func (g *chatWriteGuard) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) { - if err := g.require("InsertChatQueuedMessageWithCreator", arg.ChatID); err != nil { + if err := g.requireSnapshot("InsertChatQueuedMessageWithCreator", arg.ChatID); err != nil { return database.ChatQueuedMessage{}, err } return g.Store.InsertChatQueuedMessageWithCreator(ctx, arg) } func (g *chatWriteGuard) DeleteChatQueuedMessage(ctx context.Context, arg database.DeleteChatQueuedMessageParams) error { - if err := g.require("DeleteChatQueuedMessage", arg.ChatID); err != nil { + if err := g.requireSnapshot("DeleteChatQueuedMessage", arg.ChatID); err != nil { return err } return g.Store.DeleteChatQueuedMessage(ctx, arg) } func (g *chatWriteGuard) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg database.DeleteChatQueuedMessageReturningCountParams) (int64, error) { - if err := g.require("DeleteChatQueuedMessageReturningCount", arg.ChatID); err != nil { + if err := g.requireSnapshot("DeleteChatQueuedMessageReturningCount", arg.ChatID); err != nil { return 0, err } return g.Store.DeleteChatQueuedMessageReturningCount(ctx, arg) } func (g *chatWriteGuard) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error { - if err := g.require("DeleteAllChatQueuedMessages", chatID); err != nil { + if err := g.requireSnapshot("DeleteAllChatQueuedMessages", chatID); err != nil { return err } return g.Store.DeleteAllChatQueuedMessages(ctx, chatID) } func (g *chatWriteGuard) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) { - if err := g.require("DeleteAllChatQueuedMessagesReturningCount", chatID); err != nil { + if err := g.requireSnapshot("DeleteAllChatQueuedMessagesReturningCount", chatID); err != nil { return 0, err } return g.Store.DeleteAllChatQueuedMessagesReturningCount(ctx, chatID) } func (g *chatWriteGuard) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { - if err := g.require("PopNextQueuedMessage", chatID); err != nil { + if err := g.requireSnapshot("PopNextQueuedMessage", chatID); err != nil { return database.ChatQueuedMessage{}, err } return g.Store.PopNextQueuedMessage(ctx, chatID) } func (g *chatWriteGuard) ReorderChatQueuedMessageToFront(ctx context.Context, arg database.ReorderChatQueuedMessageToFrontParams) (int64, error) { - if err := g.require("ReorderChatQueuedMessageToFront", arg.ChatID); err != nil { + if err := g.requireSnapshot("ReorderChatQueuedMessageToFront", arg.ChatID); err != nil { return 0, err } return g.Store.ReorderChatQueuedMessageToFront(ctx, arg) } func (g *chatWriteGuard) ReorderChatQueuedMessageToHead(ctx context.Context, arg database.ReorderChatQueuedMessageToHeadParams) (int64, error) { - if err := g.require("ReorderChatQueuedMessageToHead", arg.ChatID); err != nil { + if err := g.requireSnapshot("ReorderChatQueuedMessageToHead", arg.ChatID); err != nil { return 0, err } return g.Store.ReorderChatQueuedMessageToHead(ctx, arg) diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go index 4e799588bc5..ab8eaa7c9c6 100644 --- a/coderd/database/dbtestutil/db.go +++ b/coderd/database/dbtestutil/db.go @@ -142,10 +142,10 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) { } // Unit tests should not retry serial transaction failures. db = database.New(sqlDB, database.WithSerialRetryCount(1)) - rejections := &chatWriteRecorder{} - db = newChatWriteGuard(db, rejections) + recorder := &chatWriteRecorder{} + db = newChatWriteGuard(db, recorder) t.Cleanup(func() { - for _, r := range rejections.list() { + for _, r := range recorder.list() { t.Errorf("chat write guard rejected %s for chat %s outside a chat state transition; "+ "the call site is the assertion that received this error, otherwise search for callers of %s", r.method, r.chatID, r.method) diff --git a/coderd/database/dbtestutil/db_internal_test.go b/coderd/database/dbtestutil/db_internal_test.go index bf40c360644..f96481a38ab 100644 --- a/coderd/database/dbtestutil/db_internal_test.go +++ b/coderd/database/dbtestutil/db_internal_test.go @@ -1,6 +1,7 @@ package dbtestutil import ( + "context" "fmt" "go/ast" "go/parser" @@ -43,15 +44,80 @@ CREATE TABLE foo; require.Contains(t, out, "CREATE TABLE foo;", "normalizeDump must preserve real SQL between the meta-commands") } -// The guard installed by NewDB must reject chat history and queue writes on -// the root handle and inside a transaction that has not allocated a snapshot -// for the chat, and record each rejection. +// guardedWriteCall invokes one guarded writer against chatID. messageID is +// a live message of that chat, for the writer that resolves the chat from a +// message id. +type guardedWriteCall struct { + method string + call func(ctx context.Context, store database.Store, chatID uuid.UUID, messageID int64) error +} + +// guardedWriteCalls has one entry per chatWriteGuard writer override; the +// completeness test enforces that equality so no override ships without a +// rejection case. Only the chat id is set because the guard rejects before +// the query runs. +var guardedWriteCalls = []guardedWriteCall{ + {method: "InsertChatMessages", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + _, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ChatID: chatID}) + return err + }}, + {method: "SoftDeleteChatMessageByID", call: func(ctx context.Context, store database.Store, _ uuid.UUID, messageID int64) error { + return store.SoftDeleteChatMessageByID(ctx, messageID) + }}, + {method: "SoftDeleteChatMessagesAfterID", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + return store.SoftDeleteChatMessagesAfterID(ctx, database.SoftDeleteChatMessagesAfterIDParams{ChatID: chatID}) + }}, + {method: "SoftDeleteContextFileMessages", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + return store.SoftDeleteContextFileMessages(ctx, chatID) + }}, + {method: "InsertChatQueuedMessage", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + _, err := store.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ChatID: chatID}) + return err + }}, + {method: "InsertChatQueuedMessageWithCreator", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + _, err := store.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ChatID: chatID}) + return err + }}, + {method: "DeleteChatQueuedMessage", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + return store.DeleteChatQueuedMessage(ctx, database.DeleteChatQueuedMessageParams{ChatID: chatID}) + }}, + {method: "DeleteChatQueuedMessageReturningCount", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + _, err := store.DeleteChatQueuedMessageReturningCount(ctx, database.DeleteChatQueuedMessageReturningCountParams{ChatID: chatID}) + return err + }}, + {method: "DeleteAllChatQueuedMessages", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + return store.DeleteAllChatQueuedMessages(ctx, chatID) + }}, + {method: "DeleteAllChatQueuedMessagesReturningCount", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + _, err := store.DeleteAllChatQueuedMessagesReturningCount(ctx, chatID) + return err + }}, + {method: "PopNextQueuedMessage", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + _, err := store.PopNextQueuedMessage(ctx, chatID) + return err + }}, + {method: "ReorderChatQueuedMessageToFront", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + _, err := store.ReorderChatQueuedMessageToFront(ctx, database.ReorderChatQueuedMessageToFrontParams{ChatID: chatID}) + return err + }}, + {method: "ReorderChatQueuedMessageToHead", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { + _, err := store.ReorderChatQueuedMessageToHead(ctx, database.ReorderChatQueuedMessageToHeadParams{ChatID: chatID}) + return err + }}, +} + +// The guard installed by NewDB must reject every guarded writer on the root +// handle and inside a transaction that has not allocated a snapshot for the +// chat, record each rejection, and let a write through once the same +// transaction has allocated, including from a nested InTx. func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { t.Parallel() db, _ := NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) require.Contains(t, db.Wrappers(), "dbtestutil.chatWriteGuard") + guard, ok := db.(*chatWriteGuard) + require.True(t, ok, "NewDB must return the chat write guard") user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) @@ -61,6 +127,11 @@ func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { OwnerID: user.ID, LastModelConfigID: model.ID, }) + seeded := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }) messageParams := database.InsertChatMessagesParams{ ChatID: chat.ID, CreatedBy: []uuid.UUID{user.ID}, @@ -81,27 +152,16 @@ func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { RuntimeMs: []int64{0}, } - _, err := db.InsertChatMessages(ctx, messageParams) - require.ErrorContains(t, err, "InsertChatMessages for chat "+chat.ID.String()+" outside a chat state transition") - - _, err = db.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ - ChatID: chat.ID, - Content: []byte(`[]`), - CreatedBy: user.ID, - }) - require.ErrorContains(t, err, "InsertChatQueuedMessageWithCreator for chat "+chat.ID.String()+" outside a chat state transition") - - guard, ok := db.(*chatWriteGuard) - require.True(t, ok, "NewDB must return the chat write guard") - require.Equal(t, []chatWriteRejection{ - {method: "InsertChatMessages", chatID: chat.ID}, - {method: "InsertChatQueuedMessageWithCreator", chatID: chat.ID}, - }, guard.rec.list()) - // The rejections above are the expected outcome of this test, not - // failures to report at cleanup. - guard.rec.reset() + for _, tc := range guardedWriteCalls { + err := tc.call(ctx, db, chat.ID, seeded.ID) + require.ErrorContains(t, err, tc.method+" for chat "+chat.ID.String()+" outside a chat state transition") + require.Equal(t, []chatWriteRejection{{method: tc.method, chatID: chat.ID}}, guard.rec.list()) + // The rejection is the expected outcome of this test, not a failure + // to report at cleanup. + guard.rec.reset() + } - err = db.InTx(func(tx database.Store) error { + err := db.InTx(func(tx database.Store) error { _, err := tx.InsertChatMessages(ctx, messageParams) return err }, nil) @@ -111,21 +171,39 @@ func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) require.NoError(t, err) - require.Empty(t, messages, "rejected writes must not reach the database") + require.Len(t, messages, 1, "rejected writes must not reach the database") + + // A nested InTx runs on the outer transaction handle, so it must see the + // outer allocation instead of starting an empty set. + err = db.InTx(func(tx database.Store) error { + if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, chat.ID); err != nil { + return err + } + return tx.InTx(func(inner database.Store) error { + _, err := inner.InsertChatMessages(ctx, messageParams) + return err + }, nil) + }, nil) + require.NoError(t, err) + require.Empty(t, guard.rec.list()) + messages, err = db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.Len(t, messages, 2, "the nested allocated write must reach the database") } // Every generated query that inserts, updates or deletes rows of // chat_messages or chat_queued_messages must be overridden by the guard, // except the two search_tsv maintenance queries, which change only the -// search columns. A new writer query fails this test until the guard learns -// about it. +// search columns, and every override must have a rejection case in +// guardedWriteCalls. A new writer query fails this test until the guard +// learns about it and the rejection test covers it. func TestChatWriteGuardCoversEveryChatHistoryWriter(t *testing.T) { t.Parallel() queries, err := os.ReadFile("../queries.sql.go") require.NoError(t, err) queryConst := regexp.MustCompile("(?s)\nconst \\w+ = `-- name: (\\w+) :\\w+\n([^`]*)`") - writesChatTables := regexp.MustCompile(`(?i)\b(?:INSERT\s+INTO|UPDATE|DELETE\s+FROM)\s+(?:chat_messages|chat_queued_messages)\b`) + writesChatTables := regexp.MustCompile(`(?i)\b(?:INSERT\s+INTO|MERGE\s+INTO|UPDATE|DELETE\s+FROM)\s+(?:ONLY\s+)?(?:chat_messages|chat_queued_messages)\b`) var writers []string for _, match := range queryConst.FindAllSubmatch(queries, -1) { if writesChatTables.Match(match[2]) { @@ -142,7 +220,7 @@ func TestChatWriteGuardCoversEveryChatHistoryWriter(t *testing.T) { "InsertChat": true, "LockChatAndBumpSnapshotVersion": true, } - covered := []string{"BackfillChatMessagesSearchTsv", "ReindexStaleChatMessagesSearchTsv"} + var overrides []string for _, decl := range file.Decls { fn, ok := decl.(*ast.FuncDecl) if !ok || fn.Recv == nil || len(fn.Recv.List) != 1 || !fn.Name.IsExported() { @@ -155,15 +233,25 @@ func TestChatWriteGuardCoversEveryChatHistoryWriter(t *testing.T) { if recv, ok := star.X.(*ast.Ident); !ok || recv.Name != "chatWriteGuard" || notWriters[fn.Name.Name] { continue } - covered = append(covered, fn.Name.Name) + overrides = append(overrides, fn.Name.Name) } + guardedOrExempt := append([]string{"BackfillChatMessagesSearchTsv", "ReindexStaleChatMessagesSearchTsv"}, overrides...) slices.Sort(writers) - slices.Sort(covered) - require.Equal(t, writers, covered, + slices.Sort(guardedOrExempt) + require.Equal(t, writers, guardedOrExempt, "every query writing chat_messages or chat_queued_messages must have a chatWriteGuard override "+ "or be listed here as search_tsv maintenance, and every chatWriteGuard override must be one of "+ "those writers or be listed in notWriters") + + tested := make([]string, 0, len(guardedWriteCalls)) + for _, tc := range guardedWriteCalls { + tested = append(tested, tc.method) + } + slices.Sort(overrides) + slices.Sort(tested) + require.Equal(t, overrides, tested, + "every chatWriteGuard override must have a rejection case in guardedWriteCalls, and every case must name an override") } // rejectionReportingTB records the Cleanup functions and Errorf calls that @@ -209,8 +297,8 @@ func TestChatWriteGuardReportsRejectionsAtCleanup(t *testing.T) { // NewDB registers its cleanups on tb, so they run here instead of at the // end of the test, in the order testing would use. - for i := len(tb.cleanups) - 1; i >= 0; i-- { - tb.cleanups[i]() + for _, cleanup := range slices.Backward(tb.cleanups) { + cleanup() } require.Len(t, tb.errors, 1) From fb2006d4b3a5d592ab86e119453a7b1b97fbd5ba Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 11 Sep 2026 17:31:19 +0000 Subject: [PATCH 6/6] test(coderd/database/dbtestutil): split the nested InTx case into its own test The rejection test also proved the positive path, so a failure there would report under a name that says "rejects". The nested allocation case now has a test named for what it asserts. The table comment also claimed only the chat id is set, which the SoftDeleteChatMessageByID entry contradicts; it now describes the id the guard reads. --- .../database/dbtestutil/db_internal_test.go | 43 ++++++++++++++----- 1 file changed, 33 insertions(+), 10 deletions(-) diff --git a/coderd/database/dbtestutil/db_internal_test.go b/coderd/database/dbtestutil/db_internal_test.go index f96481a38ab..db1a2934df1 100644 --- a/coderd/database/dbtestutil/db_internal_test.go +++ b/coderd/database/dbtestutil/db_internal_test.go @@ -54,8 +54,8 @@ type guardedWriteCall struct { // guardedWriteCalls has one entry per chatWriteGuard writer override; the // completeness test enforces that equality so no override ships without a -// rejection case. Only the chat id is set because the guard rejects before -// the query runs. +// rejection case. Each call sets only the id the guard reads to find the +// chat, so the guarded query itself never runs. var guardedWriteCalls = []guardedWriteCall{ {method: "InsertChatMessages", call: func(ctx context.Context, store database.Store, chatID uuid.UUID, _ int64) error { _, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ChatID: chatID}) @@ -108,8 +108,7 @@ var guardedWriteCalls = []guardedWriteCall{ // The guard installed by NewDB must reject every guarded writer on the root // handle and inside a transaction that has not allocated a snapshot for the -// chat, record each rejection, and let a write through once the same -// transaction has allocated, including from a nested InTx. +// chat, and record each rejection. func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { t.Parallel() @@ -172,23 +171,47 @@ func TestChatWriteGuardRejectsWritesWithoutSnapshot(t *testing.T) { messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) require.NoError(t, err) require.Len(t, messages, 1, "rejected writes must not reach the database") +} + +// A nested InTx runs on the outer transaction handle, so a write inside it +// must see the outer transaction's allocation: the write lands and nothing +// is recorded. +func TestChatWriteGuardSharesAllocationAcrossNestedInTx(t *testing.T) { + t.Parallel() + + db, _ := NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + guard, ok := db.(*chatWriteGuard) + require.True(t, ok, "NewDB must return the chat write guard") - // A nested InTx runs on the outer transaction handle, so it must see the - // outer allocation instead of starting an empty set. - err = db.InTx(func(tx database.Store) error { + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + + err := db.InTx(func(tx database.Store) error { if _, err := tx.LockChatAndBumpSnapshotVersion(ctx, chat.ID); err != nil { return err } return tx.InTx(func(inner database.Store) error { - _, err := inner.InsertChatMessages(ctx, messageParams) + _, err := inner.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: chat.ID, + Content: []byte(`[]`), + CreatedBy: user.ID, + }) return err }, nil) }, nil) require.NoError(t, err) require.Empty(t, guard.rec.list()) - messages, err = db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + + queued, err := db.CountChatQueuedMessages(ctx, chat.ID) require.NoError(t, err) - require.Len(t, messages, 2, "the nested allocated write must reach the database") + require.EqualValues(t, 1, queued, "the nested allocated write must reach the database") } // Every generated query that inserts, updates or deletes rows of