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..2a403040e45 --- /dev/null +++ b/coderd/database/dbtestutil/chatwriteguard.go @@ -0,0 +1,236 @@ +package dbtestutil + +import ( + "context" + "database/sql" + "errors" + "slices" + "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() + return slices.Clone(r.rejections) +} + +// 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. mu guards it because a Store, + // including a transaction handle, may be used from several goroutines. + mu sync.Mutex + 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) +} + +// 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 { + g.allocated[chatID] = struct{}{} + } +} + +// 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() + if 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.markAllocated(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.markAllocated(id) + } + return chat, err +} + +func (g *chatWriteGuard) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.InsertChatMessagesRow, error) { + if err := g.requireSnapshot("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. 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) { + return xerrors.Errorf("resolve chat for message %d: %w", id, err) + } + if err == nil { + if err := g.requireSnapshot("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.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.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.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.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.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.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.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.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.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.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.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 d25f2508e40..ab8eaa7c9c6 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)) + recorder := &chatWriteRecorder{} + db = newChatWriteGuard(db, recorder) + t.Cleanup(func() { + 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) + } + }) 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..db1a2934df1 100644 --- a/coderd/database/dbtestutil/db_internal_test.go +++ b/coderd/database/dbtestutil/db_internal_test.go @@ -1,9 +1,22 @@ package dbtestutil import ( + "context" + "fmt" + "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 +43,287 @@ 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") } + +// 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. 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}) + 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, and record each rejection. +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{}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + 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}, + 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}, + } + + 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 := 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.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") + + 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.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()) + + queued, err := db.CountChatQueuedMessages(ctx, chat.ID) + require.NoError(t, err) + require.EqualValues(t, 1, queued, "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, 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|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]) { + 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, + } + 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() { + 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 + } + overrides = append(overrides, fn.Name.Name) + } + + guardedOrExempt := append([]string{"BackfillChatMessagesSearchTsv", "ReindexStaleChatMessagesSearchTsv"}, overrides...) + slices.Sort(writers) + 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 +// 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 _, cleanup := range slices.Backward(tb.cleanups) { + cleanup() + } + + require.Len(t, tb.errors, 1) + require.Contains(t, tb.errors[0], "chat write guard rejected InsertChatQueuedMessageWithCreator for chat "+chat.ID.String()) +} diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index babef6c49ec..c16d039a48a 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) { @@ -16411,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, @@ -16532,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, @@ -18315,25 +18296,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 +18528,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 +18604,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 +18777,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..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( - ctx context.Context, t testing.TB, db database.Store, chatID uuid.UUID, @@ -177,17 +175,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 +6629,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 +6704,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 +6920,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 +11268,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 @@ -12060,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, @@ -12141,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, @@ -12205,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, @@ -12281,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, @@ -12313,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) @@ -12371,25 +12341,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, @@ -12402,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, @@ -12495,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, @@ -18012,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) { @@ -18467,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, @@ -18502,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, @@ -18655,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/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..e69912b0ed3 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,7 @@ 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, t *testing.T, db database.Store, chatID uuid.UUID, @@ -124,18 +94,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(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 +124,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..1f0b12e8009 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) @@ -6529,8 +6528,8 @@ 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) - insertAssistantTextMessage(ctx, t, db, chat.ID, "working", model.ID) + insertSystemTextMessage(t, db, chat.ID, "sys-2", model.ID) + insertAssistantTextMessage(t, db, chat.ID, "working", model.ID) _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, @@ -7582,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, @@ -7689,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, @@ -7725,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, @@ -8607,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, @@ -8615,13 +8613,12 @@ 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), }) } func insertProviderToolPairMessageWithLocalTool( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -8662,8 +8659,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", @@ -8674,7 +8671,6 @@ func insertProviderToolPairMessageWithLocalTool( } func insertChatMessageParts( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -8686,17 +8682,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( @@ -8730,8 +8722,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 @@ -13330,7 +13322,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) @@ -13383,7 +13375,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) @@ -13425,7 +13417,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, @@ -13450,7 +13442,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, @@ -13480,7 +13472,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, @@ -13498,7 +13490,6 @@ func TestQueuedPromotionResolvesOrganizationModel(t *testing.T) { } func insertQueuedMessage( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -13509,14 +13500,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 +13526,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) @@ -14375,7 +14363,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, 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)