perf(coderd/x/chatd): validate chat transitions from the locked row

ChatMachine.Update seeds the transaction with the row returned by
LockChatAndBumpSnapshotVersion plus a queue count taken right after it.
Callbacks read it with tx.Current() and the first transition validates
against it, so neither re-reads the chat while the row lock is held.

ChatMachine.UpdateReturning returns the row Update already reads after
the callback for publication. updateChat (send, edit, promote, tool
results, interrupt, compact, clear, reconcile), hook failure, family
archive, and the worker's generation and task transitions use it
instead of reloading the chat under the lock.

Generated by Coder Agents on behalf of @cstyan.

coderd/x/chatd/ARCHITECTURE.md

@@ -112,6 +112,10 @@ We will not define the SQL queries that correspond to each transition - it'd tak
 
 Each transaction that applies one or more transitions advances the `snapshot_version` field on the `chats` table by 1 immediately after locking the chat row and before mutating any tables. This lets us version the chat's execution state. The chat worker and the stream loop rely on it to ensure they do not process outdated or out of order notifications.
 
+TODO (tx.Current): `ChatMachine.Update` counts queued messages right after the lock and seeds the transaction with the locked row. Callbacks read it with `tx.Current()` and the first transition validates against it, so neither reads the chat again under the lock. Later transitions in the same transaction and the publication read still go to the database. Describe this here.
+
+TODO (UpdateReturning): `ChatMachine.UpdateReturning` returns the row that `Update` reads after the callback for publication. Callers that need the post-transition chat use it instead of reading the chat again at the end of the callback. Describe this here.
+
 Chat-message changes update `history_version` on the `chats` table and the `revision` fields on the `chat_messages` table automatically via Postgres triggers described in [Message revisions and history version](#message-revisions-and-history-version). `history_version` stores the latest `snapshot_version` in which chat message history changed. The chat runner and the stream loop rely on it to ensure they are fully aware of the chat's history changes. See [Event processing](#event-processing) for how the runner uses `history_version` differently from `snapshot_version`.
 
 Queue changes update `queue_version` automatically via Postgres triggers described in [Queue version](#queue-version).

coderd/x/chatd/chat_mutations.go

@@ -21,35 +21,27 @@ import (
 	"github.com/coder/coder/v2/codersdk/x/agenthooks"
 )
 
-// updateChat runs transition in one chatstate transaction, reloads the chat
-// inside it so the result and the sidebar watch event carry the transition's
-// snapshot bump, then publishes only that watch event (chat:update consumers
-// handle stream side effects). A transition that already holds the
-// post-transition row sets *updated to skip the reload.
+// updateChat runs transition in one chatstate transaction and returns the
+// chat Update reads after it, so the result and the sidebar watch event
+// carry the transition's snapshot bump, then publishes only that watch event
+// (chat:update consumers handle stream side effects). A transition that
+// already holds the post-transition row may set *updated to return it
+// instead.
 func (p *Server) updateChat(
 	ctx context.Context,
 	chatID uuid.UUID,
-	operation string,
 	transition func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error,
 ) (database.Chat, error) {
 	var updated database.Chat
-	err := p.newChatMachine(chatID).Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		if err := transition(tx, store, &updated); err != nil {
-			return err
-		}
-		if updated.ID != uuid.Nil {
-			return nil
-		}
-		chat, err := store.GetChatByID(ctx, chatID)
-		if err != nil {
-			return xerrors.Errorf("reload chat after %s: %w", operation, err)
-		}
-		updated = chat
-		return nil
+	published, err := p.newChatMachine(chatID).UpdateReturning(ctx, func(tx *chatstate.Tx, store database.Store) error {
+		return transition(tx, store, &updated)
 	})
 	if err != nil {
 		return database.Chat{}, err
 	}
+	if updated.ID == uuid.Nil {
+		updated = published
+	}
 	p.publishChatPubsubEvent(updated, codersdk.ChatWatchEventKindStatusChange, nil)
 	return updated, nil
 }
@@ -127,8 +119,8 @@ func (p *Server) SendMessage(
 		promotedQueuedAt     time.Time
 		finishedInterruption bool
 	)
-	refreshed, err := p.updateChat(ctx, opts.ChatID, "send", func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error {
-		lockedChat, err := store.GetChatByID(ctx, opts.ChatID)
+	refreshed, err := p.updateChat(ctx, opts.ChatID, func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error {
+		lockedChat, _, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat: %w", err)
 		}
@@ -332,8 +324,8 @@ func (p *Server) EditMessage(
 		result    EditMessageResult
 		editedMsg database.ChatMessage
 	)
-	refreshed, err := p.updateChat(ctx, opts.ChatID, "edit", func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error {
-		lockedChat, err := store.GetChatByID(ctx, opts.ChatID)
+	refreshed, err := p.updateChat(ctx, opts.ChatID, func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error {
+		lockedChat, _, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat: %w", err)
 		}
@@ -543,8 +535,8 @@ func (p *Server) PromoteQueued(
 		result           PromoteQueuedResult
 		promotedQueuedAt time.Time
 	)
-	refreshed, err := p.updateChat(ctx, opts.ChatID, "promote", func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error {
-		lockedChat, err := store.GetChatByID(ctx, opts.ChatID)
+	refreshed, err := p.updateChat(ctx, opts.ChatID, func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error {
+		lockedChat, _, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat: %w", err)
 		}
@@ -598,8 +590,8 @@ func (p *Server) SubmitToolResults(
 		}
 	}
 
-	_, err := p.updateChat(ctx, opts.ChatID, "tool results", func(tx *chatstate.Tx, store database.Store, _ *database.Chat) error {
-		locked, err := store.GetChatByID(ctx, opts.ChatID)
+	_, err := p.updateChat(ctx, opts.ChatID, func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error {
+		locked, _, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat: %w", err)
 		}
@@ -696,7 +688,7 @@ func (p *Server) InterruptChat(
 	}
 
 	var result chatstate.InterruptResult
-	refreshed, err := p.updateChat(ctx, chat.ID, "interrupt", func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error {
+	refreshed, err := p.updateChat(ctx, chat.ID, func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error {
 		var err error
 		result, err = tx.Interrupt(chatstate.InterruptInput{
 			Reason: "Tool execution interrupted by user",
@@ -750,8 +742,8 @@ func (p *Server) CompactChat(
 		return chat, xerrors.New("chat_id is required")
 	}
 
-	refreshed, err := p.updateChat(ctx, chat.ID, "compact", func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error {
-		lockedChat, err := store.GetChatByID(ctx, chat.ID)
+	refreshed, err := p.updateChat(ctx, chat.ID, func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error {
+		lockedChat, _, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat: %w", err)
 		}
@@ -803,8 +795,8 @@ func (p *Server) ClearChat(
 		return chat, xerrors.New("chat_id is required")
 	}
 
-	refreshed, err := p.updateChat(ctx, chat.ID, "clear", func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error {
-		lockedChat, err := store.GetChatByID(ctx, chat.ID)
+	refreshed, err := p.updateChat(ctx, chat.ID, func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error {
+		lockedChat, _, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat: %w", err)
 		}
@@ -866,7 +858,7 @@ func (p *Server) ReconcileInvalidStateChat(
 		return chat, xerrors.New("chat_id is required")
 	}
 
-	refreshed, err := p.updateChat(ctx, chat.ID, "reconcile", func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error {
+	refreshed, err := p.updateChat(ctx, chat.ID, func(tx *chatstate.Tx, _ database.Store, _ *database.Chat) error {
 		_, err := tx.ReconcileInvalidState(chatstate.ReconcileInvalidStateInput{})
 		return err
 	})

coderd/x/chatd/chat_mutations_internal_test.go

@@ -33,12 +33,12 @@ func TestUpdateChat(t *testing.T) {
 		return err
 	}
 
-	t.Run("ReloadsAndPublishesAfterStateUpdate", func(t *testing.T) {
+	t.Run("ReturnsAndPublishesPostTransitionChat", func(t *testing.T) {
 		t.Parallel()
 		ctx := testutil.Context(t, testutil.WaitLong)
 		fixture, chat, recorder, server := setup(t)
 
-		updated, err := server.updateChat(ctx, chat.ID, "test", requestCompaction)
+		updated, err := server.updateChat(ctx, chat.ID, requestCompaction)
 		require.NoError(t, err)
 		stored, err := fixture.db.GetChatByID(ctx, chat.ID)
 		require.NoError(t, err)
@@ -66,13 +66,13 @@ func TestUpdateChat(t *testing.T) {
 		require.Greater(t, watchIndex, stateIndex)
 	})
 
-	t.Run("TransitionRowSkipsReload", func(t *testing.T) {
+	t.Run("TransitionRowOverridesResult", func(t *testing.T) {
 		t.Parallel()
 		ctx := testutil.Context(t, testutil.WaitLong)
 		_, chat, recorder, server := setup(t)
 		sentinel := database.Chat{ID: chat.ID, OwnerID: chat.OwnerID, Title: "from transition"}
 
-		updated, err := server.updateChat(ctx, chat.ID, "test", func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error {
+		updated, err := server.updateChat(ctx, chat.ID, func(tx *chatstate.Tx, store database.Store, updated *database.Chat) error {
 			*updated = sentinel
 			return requestCompaction(tx, store, nil)
 		})
@@ -89,7 +89,7 @@ func TestUpdateChat(t *testing.T) {
 		fixture, chat, recorder, server := setup(t)
 		wantErr := xerrors.New("transition failed")
 
-		_, err := server.updateChat(ctx, chat.ID, "test", func(*chatstate.Tx, database.Store, *database.Chat) error {
+		_, err := server.updateChat(ctx, chat.ID, func(*chatstate.Tx, database.Store, *database.Chat) error {
 			return wantErr
 		})
 		require.ErrorIs(t, err, wantErr)

coderd/x/chatd/chatstate/family.go

@@ -109,14 +109,13 @@ func SetFamilyArchived(
 		}
 		familyChats = make([]database.Chat, 0, len(ids))
 		for _, id := range ids {
-			var chat database.Chat
 			machine := NewChatMachine(tx, buffer, id)
-			err := machine.Update(ctx, func(state *Tx, _ database.Store) error {
+			chat, err := machine.UpdateReturning(ctx, func(state *Tx, _ database.Store) error {
 				// Classify each member so any invalid execution state
 				// aborts and rolls back the whole family update, even
 				// when that member already has the requested archived
 				// value.
-				current, from, err := state.loadState()
+				current, from, err := state.Current()
 				if err != nil {
 					return err
 				}
@@ -124,17 +123,10 @@ func SetFamilyArchived(
 					return ErrInvalidState
 				}
 				if current.Archived == input.Archived {
-					chat = current
 					return nil
 				}
-				if _, err := state.SetArchived(SetArchivedInput{Archived: input.Archived}); err != nil {
-					return err
-				}
-				chat, err = state.Store().GetChatByID(state.Ctx(), state.ChatID())
-				if err != nil {
-					return xerrors.Errorf("reload archived chat: %w", err)
-				}
-				return nil
+				_, err = state.SetArchived(SetArchivedInput{Archived: input.Archived})
+				return err
 			})
 			if err != nil {
 				return err

coderd/x/chatd/chatstate/machine.go

@@ -52,14 +52,26 @@ func NewChatMachine(
 
 // Tx is the per-transaction handle passed to [ChatMachine.Update]
 // callbacks. It carries the active context, the transactional store,
-// and the chat ID. Tx does not cache mutable chat state across calls:
-// every transition method reads the chat row and queue cardinality
-// from the database on entry, so a bundle of transitions inside one
-// Update callback always validates against the latest committed state.
+// the chat ID, and the row returned by the transition lock.
+//
+// The locked row is safe to reuse because the lock blocks every other
+// writer of the chat until commit. Transitions write the chat, so the
+// first loadState consumes the locked row and later reads go to the
+// database. Writes a callback makes through the raw store are not
+// reflected in the locked row.
 type Tx struct {
 	ctx    context.Context
 	store  database.Store
 	chatID uuid.UUID
+
+	locked *lockedRow
+}
+
+// lockedRow is the chat row returned by the transition lock plus the
+// queue state read right after it.
+type lockedRow struct {
+	chat      database.Chat
+	hasQueued bool
 }
 
 // Ctx returns the context the surrounding [ChatMachine.Update] call
@@ -81,11 +93,31 @@ func (tx *Tx) ChatID() uuid.UUID { return tx.chatID }
 // matrix.
 func (tx *Tx) Store() database.Store { return tx.store }
 
-// loadState reads the current chat row and queue cardinality from the
-// active transaction, classifies the execution state, and returns the
-// inputs every transition method needs. Returns ErrChatNotFound if
-// the chat row was deleted in this transaction (or never existed).
+// loadState returns the chat row and execution state for a transition
+// to validate against. The locked row is returned at most once because
+// the transition that reads it is about to write the chat.
 func (tx *Tx) loadState() (database.Chat, ExecutionState, error) {
+	if l := tx.locked; l != nil {
+		tx.locked = nil
+		return l.chat, ClassifyExecutionState(l.chat, l.hasQueued, true), nil
+	}
+	return tx.readState()
+}
+
+// Current returns the chat row and execution state, reusing the locked
+// row without consuming it when no transition has run yet.
+func (tx *Tx) Current() (database.Chat, ExecutionState, error) {
+	if l := tx.locked; l != nil {
+		return l.chat, ClassifyExecutionState(l.chat, l.hasQueued, true), nil
+	}
+	return tx.readState()
+}
+
+// readState reads the current chat row and queue cardinality from the
+// active transaction and classifies the execution state. Returns
+// ErrChatNotFound if the chat row was deleted in this transaction (or
+// never existed).
+func (tx *Tx) readState() (database.Chat, ExecutionState, error) {
 	chat, err := tx.store.GetChatByID(tx.ctx, tx.chatID)
 	if err != nil {
 		if errors.Is(err, sql.ErrNoRows) {
@@ -148,35 +180,57 @@ func (m *ChatMachine) Update(
 	ctx context.Context,
 	fn func(*Tx, database.Store) error,
 ) error {
+	_, err := m.UpdateReturning(ctx, fn)
+	return err
+}
+
+// UpdateReturning is [ChatMachine.Update] that also returns the chat row
+// Update reads after fn for publication, so callers do not read it again
+// under the lock. The row is zero if fn or that read failed. It is still
+// set when the commit or publication fails afterwards, so treat it as
+// committed only when err is nil.
+func (m *ChatMachine) UpdateReturning(
+	ctx context.Context,
+	fn func(*Tx, database.Store) error,
+) (database.Chat, error) {
 	if m.store == nil {
-		return xerrors.New("chatstate: ChatMachine has nil store")
+		return database.Chat{}, xerrors.New("chatstate: ChatMachine has nil store")
 	}
 	if m.publisher == nil {
-		return xerrors.New("chatstate: ChatMachine has nil publisher")
+		return database.Chat{}, xerrors.New("chatstate: ChatMachine has nil publisher")
 	}
 
 	buffer := NewPublishBuffer(m.publisher)
 	defer buffer.Discard()
 
+	var updated database.Chat
 	err := m.store.InTx(func(store database.Store) error {
-		if _, err := store.LockChatAndBumpSnapshotVersion(ctx, m.chatID); err != nil {
+		locked, err := store.LockChatAndBumpSnapshotVersion(ctx, m.chatID)
+		if err != nil {
 			if errors.Is(err, sql.ErrNoRows) {
 				return ErrChatNotFound
 			}
 			return xerrors.Errorf("lock chat and bump snapshot: %w", err)
 		}
+		queued, err := store.CountChatQueuedMessages(ctx, m.chatID)
+		if err != nil {
+			return xerrors.Errorf("count queued messages: %w", err)
+		}
 		tx := &Tx{
 			ctx:    ctx,
 			store:  store,
 			chatID: m.chatID,
+			locked: &lockedRow{chat: locked, hasQueued: queued > 0},
 		}
 		if err := fn(tx, store); err != nil {
 			return err
 		}
-		chat, state, err := tx.loadState()
+		// The callback may have written the chat, so publish a fresh read.
+		chat, state, err := tx.readState()
 		if err != nil {
 			return err
 		}
+		updated = chat
 		if err := buffer.Publish(
 			coderdpubsub.ChatStateUpdateChannel(chat.ID),
 			buildChatUpdateMessage(chat),
@@ -200,9 +254,9 @@ func (m *ChatMachine) Update(
 		return nil
 	}, nil)
 	if err != nil {
-		return err
+		return updated, err
 	}
-	return buffer.Flush()
+	return updated, buffer.Flush()
 }
 
 // Lock locks the chat row with FOR UPDATE and runs fn in a

coderd/x/chatd/chatstate/machine_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbgen"
@@ -173,6 +174,136 @@ func TestChatMachine_Update_RejectsMissingChat(t *testing.T) {
 	require.Empty(t, f.Pub.channels)
 }
 
+// chatReadCounts counts chat reads issued on transactional handles.
+type chatReadCounts struct {
+	mu          sync.Mutex
+	getChatByID int
+	countQueued int
+}
+
+type chatReadCountingStore struct {
+	database.Store
+	counts *chatReadCounts
+	inTx   bool
+}
+
+func (s *chatReadCountingStore) InTx(fn func(database.Store) error, opts *database.TxOptions) error {
+	return s.Store.InTx(func(tx database.Store) error {
+		return fn(&chatReadCountingStore{Store: tx, counts: s.counts, inTx: true})
+	}, opts)
+}
+
+func (s *chatReadCountingStore) GetChatByID(ctx context.Context, id uuid.UUID) (database.Chat, error) {
+	if s.inTx {
+		s.counts.mu.Lock()
+		s.counts.getChatByID++
+		s.counts.mu.Unlock()
+	}
+	return s.Store.GetChatByID(ctx, id)
+}
+
+func (s *chatReadCountingStore) CountChatQueuedMessages(ctx context.Context, id uuid.UUID) (int64, error) {
+	if s.inTx {
+		s.counts.mu.Lock()
+		s.counts.countQueued++
+		s.counts.mu.Unlock()
+	}
+	return s.Store.CountChatQueuedMessages(ctx, id)
+}
+
+// Reads under the transition lock directly lengthen lock hold time, so
+// validation must use the row the lock returned.
+func TestChatMachine_Update_ValidatesFromLockedRow(t *testing.T) {
+	t.Parallel()
+	f := newTestFixture(t)
+	ctx := testutil.Context(t, testutil.WaitShort)
+	created := createTestChat(t, f)
+
+	counts := &chatReadCounts{}
+	m := chatstate.NewChatMachine(&chatReadCountingStore{Store: f.DB, counts: counts}, f.Pub, created.Chat.ID)
+	require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, _, err := tx.Current(); err != nil {
+			return err
+		}
+		_, err := tx.FinishTurn(chatstate.FinishTurnInput{})
+		return err
+	}))
+
+	require.Equal(t, 1, counts.getChatByID, "only the publication read")
+	require.Equal(t, 2, counts.countQueued, "the post-lock count and the publication read")
+}
+
+func TestChatMachine_Update_CurrentMatchesDatabase(t *testing.T) {
+	t.Parallel()
+	f := newTestFixture(t)
+	ctx := testutil.Context(t, testutil.WaitShort)
+	created := createTestChat(t, f)
+
+	m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
+	require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
+		locked, state, err := tx.Current()
+		require.NoError(t, err)
+		read, err := store.GetChatByID(ctx, created.Chat.ID)
+		require.NoError(t, err)
+		require.Equal(t, read, locked, "the locked row must match an in-transaction read")
+		require.Equal(t, created.Chat.SnapshotVersion+1, locked.SnapshotVersion)
+		require.Equal(t, chatstate.StateR0, state)
+
+		if _, err := tx.FinishTurn(chatstate.FinishTurnInput{}); err != nil {
+			return err
+		}
+		after, state, err := tx.Current()
+		require.NoError(t, err)
+		require.Equal(t, database.ChatStatusWaiting, after.Status, "Current must observe the transition's write")
+		require.Equal(t, chatstate.StateW, state)
+		return nil
+	}))
+}
+
+func TestChatMachine_UpdateReturning(t *testing.T) {
+	t.Parallel()
+
+	t.Run("ReturnsCommittedRow", func(t *testing.T) {
+		t.Parallel()
+		f := newTestFixture(t)
+		ctx := testutil.Context(t, testutil.WaitShort)
+		created := createTestChat(t, f)
+
+		counts := &chatReadCounts{}
+		m := chatstate.NewChatMachine(&chatReadCountingStore{Store: f.DB, counts: counts}, f.Pub, created.Chat.ID)
+		updated, err := m.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+			_, err := tx.FinishTurn(chatstate.FinishTurnInput{})
+			return err
+		})
+		require.NoError(t, err)
+		require.Equal(t, 1, counts.getChatByID, "the returned row is the publication read")
+
+		committed := f.readChat(ctx, t, created.Chat.ID)
+		require.Equal(t, committed, updated)
+		require.Equal(t, database.ChatStatusWaiting, updated.Status)
+		require.Equal(t, created.Chat.SnapshotVersion+1, updated.SnapshotVersion)
+	})
+
+	t.Run("CallbackErrorReturnsNoRow", func(t *testing.T) {
+		t.Parallel()
+		f := newTestFixture(t)
+		ctx := testutil.Context(t, testutil.WaitShort)
+		created := createTestChat(t, f)
+
+		m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
+		wantErr := xerrors.New("callback failed")
+		updated, err := m.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+			if _, err := tx.FinishTurn(chatstate.FinishTurnInput{}); err != nil {
+				return err
+			}
+			return wantErr
+		})
+		require.ErrorIs(t, err, wantErr)
+		require.Zero(t, updated)
+		require.Equal(t, created.Chat.SnapshotVersion, f.readChat(ctx, t, created.Chat.ID).SnapshotVersion)
+	})
+}
+
 func TestChatMachine_Lock_DoesNotBumpSnapshot(t *testing.T) {
 	t.Parallel()
 	f := newTestFixture(t)

coderd/x/chatd/generation.go

@@ -439,9 +439,8 @@ func applySessionStartResponse(
 		return sessionStartResult{}, err
 	}
 
-	var applied sessionStartResult
-	err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		if _, err := loadChatForGeneration(ctx, store, input, generationAttemptNotRequired); err != nil {
+	updated, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, err := loadLockedChatForGeneration(tx, input, generationAttemptNotRequired); err != nil {
 			return xerrors.Errorf("load chat for session_start response: %w", err)
 		}
 		if len(eventMessages) > 0 {
@@ -449,16 +448,12 @@ func applySessionStartResponse(
 				return xerrors.Errorf("insert session_start response messages: %w", err)
 			}
 		}
-		applied.Chat, err = store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("reload chat after session_start response: %w", err)
-		}
 		return nil
 	})
 	if err != nil {
 		return sessionStartResult{}, normalizeTaskTransitionError(err, "apply session_start response")
 	}
-	return applied, nil
+	return sessionStartResult{Chat: updated}, nil
 }
 
 // syncMemoryIndex commits the model-only memory index message the chat
@@ -479,24 +474,19 @@ func (s *taskStarter) syncMemoryIndex(
 	if !ok {
 		return sessionStartResult{}, nil
 	}
-	var applied sessionStartResult
-	err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		if _, err := loadChatForGeneration(ctx, store, input, generationAttemptNotRequired); err != nil {
+	updated, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, err := loadLockedChatForGeneration(tx, input, generationAttemptNotRequired); err != nil {
 			return xerrors.Errorf("load chat for memory index: %w", err)
 		}
 		if _, err := tx.CommitStep(chatstate.CommitStepInput{Messages: []chatstate.Message{message}}); err != nil {
 			return xerrors.Errorf("insert memory index message: %w", err)
 		}
-		applied.Chat, err = store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("reload chat after memory index: %w", err)
-		}
 		return nil
 	})
 	if err != nil {
 		return sessionStartResult{}, normalizeTaskTransitionError(err, "apply memory index")
 	}
-	return applied, nil
+	return sessionStartResult{Chat: updated}, nil
 }
 
 func (s *taskStarter) startGenerationSession(
@@ -805,8 +795,8 @@ func (s *taskStarter) recordGenerationRetry(
 ) (generationRetryDecision, error) {
 	var decision generationRetryDecision
 	var payload *codersdk.ChatStreamRetry
-	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		chat, err := loadChatForTask(ctx, store, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true})
+	err := machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		chat, err := loadLockedChatForTask(tx, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true})
 		if err != nil {
 			return xerrors.Errorf("load chat for task: %w", err)
 		}
@@ -1446,9 +1436,8 @@ func (s *taskStarter) beginGenerationAttempt(
 	input chatWorkerTaskStartInput,
 ) (generationAttempt, error) {
 	var attempt int64
-	var committed database.Chat
-	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		if _, err := loadChatForTask(ctx, store, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
+	committed, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, err := loadLockedChatForTask(tx, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
 			return xerrors.Errorf("load chat for task: %w", err)
 		}
 		result, err := tx.RecordGenerationAttempt(chatstate.RecordGenerationAttemptInput{})
@@ -1456,10 +1445,6 @@ func (s *taskStarter) beginGenerationAttempt(
 			return xerrors.Errorf("tx.RecordGenerationAttempt: %w", err)
 		}
 		attempt = result.GenerationAttempt
-		committed, err = store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("load committed chat: %w", err)
-		}
 		return nil
 	})
 	if err != nil {
@@ -1529,13 +1514,12 @@ func (s *taskStarter) commitGenerationStep(
 		)
 		postCommitLastError, postCommitMessage = generationLastError(commitHooks.PostCommitError)
 	}
-	var committed database.Chat
 	commitCtx, commitSpan := s.server.stages.Start(ctx, chatloop.StageCommit,
 		attribute.String(chatloop.AttrGenerationAction, string(kind)),
 		attribute.Int64(chatloop.AttrGenerationAttempt, attempt),
 	)
-	err := machine.Update(commitCtx, func(tx *chatstate.Tx, store database.Store) error {
-		if _, err := loadChatForGeneration(ctx, store, input, requireGenerationAttempt(attempt)); err != nil {
+	committed, err := machine.UpdateReturning(commitCtx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, err := loadLockedChatForGeneration(tx, input, requireGenerationAttempt(attempt)); err != nil {
 			return xerrors.Errorf("load chat for generation: %w", err)
 		}
 		if _, err := tx.CommitStep(chatstate.CommitStepInput{
@@ -1552,11 +1536,6 @@ func (s *taskStarter) commitGenerationStep(
 				return xerrors.Errorf("tx.FinishError: %w", err)
 			}
 		}
-		loadedChat, err := store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("load committed chat: %w", err)
-		}
-		committed = loadedChat
 		return nil
 	})
 	commitSpan.End(err)
@@ -1589,19 +1568,13 @@ func (s *taskStarter) enterRequiresAction(
 	machine *chatstate.ChatMachine,
 	input chatWorkerTaskStartInput,
 ) error {
-	var committed database.Chat
-	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		if _, err := loadChatForTask(ctx, store, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
+	committed, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, err := loadLockedChatForTask(tx, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
 			return xerrors.Errorf("load chat for task: %w", err)
 		}
 		if _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}); err != nil {
 			return xerrors.Errorf("tx.EnterRequiresAction: %w", err)
 		}
-		chat, err := store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("load committed chat: %w", err)
-		}
-		committed = chat
 		return nil
 	})
 	if err != nil {
@@ -1648,6 +1621,20 @@ func loadChatForGeneration(
 	if err != nil {
 		return database.Chat{}, err
 	}
+	return verifyGenerationFence(chat, fence)
+}
+
+// loadLockedChatForGeneration is loadChatForGeneration for
+// [chatstate.ChatMachine.Update] callbacks; see loadLockedChatForTask.
+func loadLockedChatForGeneration(tx *chatstate.Tx, input chatWorkerTaskStartInput, fence generationAttemptFence) (database.Chat, error) {
+	chat, err := loadLockedChatForTask(tx, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true})
+	if err != nil {
+		return database.Chat{}, err
+	}
+	return verifyGenerationFence(chat, fence)
+}
+
+func verifyGenerationFence(chat database.Chat, fence generationAttemptFence) (database.Chat, error) {
 	if fence.required && chat.GenerationAttempt != fence.attempt {
 		return database.Chat{}, errors.Join(errTaskExpectedExit, xerrors.Errorf("generation fence mismatch: %d != %d", chat.GenerationAttempt, fence.attempt))
 	}
@@ -1696,8 +1683,8 @@ func (s *taskStarter) finishGenerationTurnWithoutHook(
 ) error {
 	var committed database.Chat
 	var promotedQueuedAt time.Time
-	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		if _, err := loadChatForGeneration(ctx, store, input, fence); err != nil {
+	err := machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, err := loadLockedChatForGeneration(tx, input, fence); err != nil {
 			return xerrors.Errorf("load chat for generation: %w", err)
 		}
 		finishResult, err := tx.FinishTurn(chatstate.FinishTurnInput{})
@@ -1763,8 +1750,8 @@ func (s *taskStarter) finishGenerationTurn(
 
 	var committed database.Chat
 	var promotedQueuedAt time.Time
-	err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		if _, err := loadChatForGeneration(ctx, store, input, fence); err != nil {
+	updated, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, err := loadLockedChatForGeneration(tx, input, fence); err != nil {
 			return xerrors.Errorf("load chat for generation: %w", err)
 		}
 		if len(stopMessages) > 0 {
@@ -1779,13 +1766,7 @@ func (s *taskStarter) finishGenerationTurn(
 			}
 			promotedQueuedAt = finishResult.PromotedQueuedAt
 			committed = finishResult.Chat
-			return nil
 		}
-		loadedChat, err := store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("load committed chat: %w", err)
-		}
-		committed = loadedChat
 		return nil
 	})
 	if err != nil {
@@ -1797,6 +1778,7 @@ func (s *taskStarter) finishGenerationTurn(
 		return err
 	}
 	if continueTurn {
+		committed = updated
 		s.routeStateHint(ctx, stateUpdateFromChat(committed))
 		return s.afterGenerationOutcome(ctx, generationOutcome{
 			Chat: committed,
@@ -1830,19 +1812,13 @@ func (s *taskStarter) finishGenerationError(
 		slog.Error(cause),
 	)
 	lastError, message := generationLastError(cause)
-	var committed database.Chat
-	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		if _, err := loadChatForGeneration(ctx, store, input, fence); err != nil {
+	committed, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		if _, err := loadLockedChatForGeneration(tx, input, fence); err != nil {
 			return xerrors.Errorf("load chat for generation: %w", err)
 		}
 		if _, err := tx.FinishError(chatstate.FinishErrorInput{LastError: lastError}); err != nil {
 			return xerrors.Errorf("tx.FinishError: %w", err)
 		}
-		chat, err := store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("load committed chat: %w", err)
-		}
-		committed = chat
 		return nil
 	})
 	if err != nil {

coderd/x/chatd/hook_server.go

@@ -79,10 +79,9 @@ func (p *Server) handleAPIDispatchError(ctx context.Context, chatID uuid.UUID, e
 	if marshalErr != nil {
 		return errors.Join(dispatchErr, xerrors.Errorf("encode hook dispatch error: %w", marshalErr))
 	}
-	var failedChat database.Chat
 	machine := p.newChatMachine(chatID)
-	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		current, err := store.GetChatByID(ctx, chatID)
+	failedChat, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		current, _, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat for hook failure: %w", err)
 		}
@@ -92,17 +91,10 @@ func (p *Server) handleAPIDispatchError(ctx context.Context, chatID uuid.UUID, e
 		if current.Status != database.ChatStatusWaiting {
 			return chatstate.ErrTransitionNotAllowed
 		}
-		if _, err := tx.FinishError(chatstate.FinishErrorInput{
+		_, err = tx.FinishError(chatstate.FinishErrorInput{
 			LastError: pqtype.NullRawMessage{RawMessage: encoded, Valid: true},
-		}); err != nil {
-			return err
-		}
-		chat, err := store.GetChatByID(ctx, chatID)
-		if err != nil {
-			return xerrors.Errorf("reload chat after hook failure: %w", err)
-		}
-		failedChat = chat
-		return nil
+		})
+		return err
 	})
 	if errors.Is(err, chatstate.ErrTransitionNotAllowed) {
 		return dispatchErr

coderd/x/chatd/tasks.go

@@ -325,12 +325,9 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt
 
 	canceled := s.cancelUnresolvedToolCalls(ctx, chat, cancelable, toolCallIDs)
 
-	var (
-		committed        database.Chat
-		promotedQueuedAt time.Time
-	)
-	err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		chat, err := loadChatForTask(ctx, store, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true})
+	var promotedQueuedAt time.Time
+	committed, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, store database.Store) error {
+		chat, err := loadLockedChatForTask(tx, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true})
 		if err != nil {
 			return xerrors.Errorf("load chat for task: %w", err)
 		}
@@ -351,10 +348,6 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt
 			return xerrors.Errorf("finish interruption: %w", err)
 		}
 		promotedQueuedAt = finishResult.PromotedQueuedAt
-		committed, err = store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("load committed chat: %w", err)
-		}
 		return nil
 	})
 	if err != nil {
@@ -474,9 +467,8 @@ func (s *taskStarter) cancelRequiresAction(
 	input chatWorkerTaskStartInput,
 	reason string,
 ) error {
-	var committed database.Chat
-	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		chat, err := loadChatForTask(ctx, store, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true})
+	committed, err := machine.UpdateReturning(ctx, func(tx *chatstate.Tx, store database.Store) error {
+		chat, err := loadLockedChatForTask(tx, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true})
 		if err != nil {
 			return xerrors.Errorf("load chat for task: %w", err)
 		}
@@ -492,10 +484,6 @@ func (s *taskStarter) cancelRequiresAction(
 		if _, err := tx.CancelRequiresAction(chatstate.CancelRequiresActionInput{Reason: reason}); err != nil {
 			return xerrors.Errorf("cancel requires action: %w", err)
 		}
-		committed, err = store.GetChatByID(ctx, input.ChatID)
-		if err != nil {
-			return xerrors.Errorf("load committed chat: %w", err)
-		}
 		return nil
 	})
 	if err != nil {
@@ -510,13 +498,9 @@ func (s *taskStarter) cancelRequiresAction(
 func (s *taskStarter) StartAbandon(ctx context.Context, input chatWorkerTaskStartInput) error {
 	machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID)
 	mismatch := false
-	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		chat, err := store.GetChatByID(ctx, input.ChatID)
+	err := machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error {
+		chat, _, err := tx.Current()
 		if err != nil {
-			if errors.Is(err, sql.ErrNoRows) {
-				mismatch = true
-				return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-			}
 			return xerrors.Errorf("load chat: %w", err)
 		}
 		if !ownedByTask(chat, input) {
@@ -650,6 +634,25 @@ func loadChatForTask(
 	return chat, nil
 }
 
+// loadLockedChatForTask is loadChatForTask for [chatstate.ChatMachine.Update]
+// callbacks: it fences against the row the transition lock returned instead
+// of reading the chat again while the lock is held.
+func loadLockedChatForTask(
+	tx *chatstate.Tx,
+	input chatWorkerTaskStartInput,
+	status database.ChatStatus,
+	opts taskFenceOptions,
+) (database.Chat, error) {
+	chat, _, err := tx.Current()
+	if err != nil {
+		return database.Chat{}, xerrors.Errorf("load chat: %w", err)
+	}
+	if err := verifyTaskFence(chat, input, status, opts); err != nil {
+		return database.Chat{}, xerrors.Errorf("verifyTaskFence: %w", err)
+	}
+	return chat, nil
+}
+
 func verifyTaskFence(
 	chat database.Chat,
 	input chatWorkerTaskStartInput,

coderd/x/chatd/worker.go

@@ -2,7 +2,6 @@ package chatd
 
 import (
 	"context"
-	"database/sql"
 	"errors"
 	"sync"
 	"time"
@@ -263,18 +262,11 @@ func (w *chatWorker) acquireCandidate(
 	var takenOver bool
 	machine := chatstate.NewChatMachine(w.opts.Store, w.opts.Pubsub, chatID)
 	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		chat, err := store.GetChatByID(ctx, chatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errSkipAcquire
-		}
+		chat, state, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat: %w", err)
 		}
-		queueCount, err := store.CountChatQueuedMessages(ctx, chatID)
-		if err != nil {
-			return xerrors.Errorf("count queue: %w", err)
-		}
-		if !chatstate.ClassifyExecutionState(chat, queueCount > 0, true).IsRunnable() || chat.Archived {
+		if !state.IsRunnable() || chat.Archived {
 			return errSkipAcquire
 		}
 		if chat.WorkerID.Valid && chat.RunnerID.Valid {
@@ -359,11 +351,8 @@ func (w *chatWorker) abandonAcquiredChat(ctx context.Context, workerID uuid.UUID
 	cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), shutdownCleanupTimeout)
 	defer cancel()
 	machine := chatstate.NewChatMachine(w.opts.Store, w.opts.Pubsub, chatID)
-	err := machine.Update(cleanupCtx, func(tx *chatstate.Tx, store database.Store) error {
-		chat, err := store.GetChatByID(cleanupCtx, chatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errSkipAcquire
-		}
+	err := machine.Update(cleanupCtx, func(tx *chatstate.Tx, _ database.Store) error {
+		chat, _, err := tx.Current()
 		if err != nil {
 			return xerrors.Errorf("load chat: %w", err)
 		}