chore(coderd/x/chatd): address generation review items (#26517)

Addresses the deferred `coderd/x/chatd/generation.go` review comments from PR #26109: [required generation dependencies](https://github.com/coder/coder/pull/26109#discussion_r3380311853), [scoped chat variables](https://github.com/coder/coder/pull/26109#discussion_r3387161874), [generation state error handling](https://github.com/coder/coder/pull/26109#discussion_r3387191468), [generation attempt return values](https://github.com/coder/coder/pull/26109#discussion_r3387251382), [generation fence verification](https://github.com/coder/coder/pull/26109#discussion_r3387288234), and [chatdebug outcome recording](https://github.com/coder/coder/pull/26109#discussion_r3387544273). This makes generation task dependencies explicit, packages generation attempt episode state into a struct, and centralizes generation task fence checks for generation transitions. Generated by Coder Agents, closely reviewed by Hugo.
2026-09-24 15:04:27 +08:00 · 2026-07-06 16:37:45 +00:00
parent 2de8b3c74a
commit 96130e2bc5
9 changed files with 294 additions and 245 deletions
@@ -147,8 +147,8 @@ func (w *chatWorker) archiveCandidate(
 	if len(familyChats) == 0 {
 		return nil, nil
 	}
-	w.scheduleArchiveDebugCleanup(ctx, familyChats)
-	w.publishArchiveWatchEvents(familyChats)
+	w.server.scheduleArchiveDebugCleanup(ctx, familyChats)
+	w.server.publishChatPubsubEvents(familyChats, codersdk.ChatWatchEventKindDeleted)

 	archived := make([]autoArchivedChat, 0, len(familyChats))
 	for _, chat := range familyChats {
@@ -172,28 +172,6 @@ func isExpectedAutoArchiveError(err error) bool {
 		errors.Is(err, chatstate.ErrTransitionNotAllowed)
 }

-func (w *chatWorker) publishArchiveWatchEvents(familyChats []database.Chat) {
-	if w.server != nil {
-		w.server.publishChatPubsubEvents(familyChats, codersdk.ChatWatchEventKindDeleted)
-		return
-	}
-	for _, chat := range familyChats {
-		if err := publishChatWatchEvent(w.opts.Pubsub, chat, codersdk.ChatWatchEventKindDeleted); err != nil {
-			w.opts.Logger.Warn(context.Background(), "chatworker auto-archive watch publish failed",
-				slog.F("chat_id", chat.ID),
-				slog.Error(err),
-			)
-		}
-	}
-}
-
-func (w *chatWorker) scheduleArchiveDebugCleanup(ctx context.Context, familyChats []database.Chat) {
-	if w.server == nil || len(familyChats) == 0 {
-		return
-	}
-	w.server.scheduleArchiveDebugCleanup(ctx, familyChats)
-}
-
 func (p *Server) scheduleArchiveDebugCleanup(ctx context.Context, familyChats []database.Chat) {
 	if len(familyChats) == 0 {
 		return
@@ -293,7 +293,15 @@ func (f *workerTestFixture) newArchiveWorkerWithOptions(t *testing.T, opts chatW
 	if opts.NotificationsEnqueuer == nil {
 		opts.NotificationsEnqueuer = notificationstest.NewFakeEnqueuer()
 	}
-	worker, err := newChatWorker(nil, opts)
+	// The archive tick dereferences the worker's server for debug cleanup
+	// and pubsub events, so give it a real (unstarted) Server. Route the
+	// server through the recording pubsub when one is in use so tests can
+	// observe server-published watch events.
+	serverPS := f.pubsub
+	if recording, ok := opts.Pubsub.(*recordingPubsub); ok {
+		serverPS = recording
+	}
+	worker, err := newChatWorker(newUnstartedServer(t, serverPS, f.db), opts)
 	require.NoError(t, err)
 	return worker
 }
@@ -2,7 +2,6 @@ package chatd

 import (
 	"context"
-	"database/sql"
 	"encoding/json"
 	"errors"
 	"strings"
@@ -294,17 +293,14 @@ func hasExclusiveToolCall(toolCalls []fantasy.ToolCallContent, exclusiveToolName
 }

 func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskStartInput) error {
-	if s.server == nil {
-		return xerrors.New("chatworker: server is required")
-	}
 	machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID)
 	for {
-		locked, messages, err := loadGenerationState(ctx, machine, input)
+		chat, messages, err := loadGenerationState(ctx, machine, input)
 		if err != nil {
 			return xerrors.Errorf("load generation state: %w", err)
 		}
 		prepareInput := generationPrepareInput{
-			Chat:     locked,
+			Chat:     chat,
 			Messages: messages,
 		}
 		prepared, err := retryGenerationPhase(ctx, s, "prepare", func() (generationPrepared, error) {
@@ -314,7 +310,7 @@ func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskS
 			if errors.Is(err, errTaskExpectedExit) || errors.Is(err, errTaskRetryable) {
 				return xerrors.Errorf("prepare generation: %w", err)
 			}
-			return s.finishGenerationError(ctx, machine, input, 0, err, generationAttemptNotRequired)
+			return s.finishGenerationError(ctx, machine, input, err, generationAttemptNotRequired)
 		}
 		cleanup := prepared.Cleanup
 		decision, err := retryGenerationPhase(ctx, s, "decide", func() (generationDecision, error) {
@@ -344,7 +340,7 @@ func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskS
 					errCompactionStillOverLimit,
 				)
 			}
-			return s.finishGenerationError(ctx, machine, input, 0, err, generationAttemptNotRequired)
+			return s.finishGenerationError(ctx, machine, input, err, generationAttemptNotRequired)
 		}

 		var actionErr error
@@ -354,7 +350,7 @@ func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskS
 			return s.enterRequiresAction(ctx, machine, input)
 		case generationActionFinishTurn:
 			cleanup()
-			return s.finishGenerationTurn(ctx, machine, input, 0, decision, generationAttemptNotRequired)
+			return s.finishGenerationTurn(ctx, machine, input, decision, generationAttemptNotRequired)
 		case generationActionGenerateAssistant:
 			actionErr = s.generateAssistant(ctx, machine, input, prepared)
 		case generationActionExecuteLocalTools:
@@ -362,7 +358,7 @@ func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskS
 		case generationActionCompact:
 			actionErr = s.generateCompaction(ctx, machine, input, prepared)
 		default:
-			return s.finishGenerationError(ctx, machine, input, 0, xerrors.Errorf("unknown generation action %q", decision.kind), generationAttemptNotRequired)
+			return s.finishGenerationError(ctx, machine, input, xerrors.Errorf("unknown generation action %q", decision.kind), generationAttemptNotRequired)
 		}
 		cleanup()
 		if actionErr == nil {
@@ -402,9 +398,9 @@ func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskS
 				}
 				continue
 			}
-			return s.finishGenerationError(ctx, machine, input, decision.generationAttempt, actionErr, generationAttemptRequired)
+			return s.finishGenerationError(ctx, machine, input, actionErr, requireGenerationAttempt(decision.generationAttempt))
 		}
-		return s.finishGenerationError(ctx, machine, input, 0, actionErr, generationAttemptNotRequired)
+		return s.finishGenerationError(ctx, machine, input, actionErr, generationAttemptNotRequired)
 	}
 }

@@ -413,34 +409,28 @@ func loadGenerationState(
 	machine *chatstate.ChatMachine,
 	input chatWorkerTaskStartInput,
 ) (database.Chat, []database.ChatMessage, error) {
-	var locked database.Chat
+	var chat database.Chat
 	var messages []database.ChatMessage
 	err := machine.ReadLock(ctx, func(store database.Store) error {
-		chat, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load locked chat: %w", err))
-		}
+		loadedChat, err := loadChatForTask(ctx, store, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true})
 		if err != nil {
-			return xerrors.Errorf("load locked chat: %w", err)
+			return xerrors.Errorf("load chat for task: %w", err)
 		}
-		if err := verifyTaskFence(chat, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
-		}
-		loaded, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
+		loadedMessages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
 			ChatID:  input.ChatID,
 			AfterID: 0,
 		})
 		if err != nil {
 			return xerrors.Errorf("load chat messages: %w", err)
 		}
-		locked = chat
-		messages = loaded
+		chat = loadedChat
+		messages = loadedMessages
 		return nil
 	})
 	if err != nil {
 		return database.Chat{}, nil, normalizeTaskInfrastructureError(err, "lock chat for generation")
 	}
-	return locked, messages, nil
+	return chat, messages, nil
 }

 func (*taskStarter) recordGenerationRetry(
@@ -452,23 +442,17 @@ func (*taskStarter) recordGenerationRetry(
 	var decision generationRetryDecision
 	var payload *codersdk.ChatStreamRetry
 	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-		}
+		chat, err := loadChatForTask(ctx, store, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true})
 		if err != nil {
-			return xerrors.Errorf("load chat: %w", err)
+			return xerrors.Errorf("load chat for task: %w", err)
 		}
-		if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
-		}
-		decision.generationAttempt = locked.GenerationAttempt
-		if locked.GenerationAttempt <= 0 || locked.GenerationAttempt >= int64(chatretry.MaxAttempts) {
+		decision.generationAttempt = chat.GenerationAttempt
+		if chat.GenerationAttempt <= 0 || chat.GenerationAttempt >= int64(chatretry.MaxAttempts) {
 			decision.retry = false
 			return errRetryStateDecisionOnly
 		}

-		attempt := int(locked.GenerationAttempt)
+		attempt := int(chat.GenerationAttempt)
 		delay := chatretry.Delay(attempt - 1)
 		if classified.RetryAfter > delay {
 			delay = classified.RetryAfter
@@ -587,11 +571,11 @@ func (s *taskStarter) generateAssistant(
 	input chatWorkerTaskStartInput,
 	prepared generationPrepared,
 ) error {
-	attempt, _, publish, closeEpisode, err := s.beginGenerationAttempt(ctx, machine, input)
+	attempt, err := s.beginGenerationAttempt(ctx, machine, input)
 	if err != nil {
 		return xerrors.Errorf("begin generation attempt: %w", err)
 	}
-	defer closeEpisode()
+	defer attempt.closeEpisode()
 	runCtx := input.DebugTurn.Ensure(ctx, prepared.Chat, prepared.Debug)
 	outcome, err := chatloop.GenerateAssistant(runCtx, chatloop.GenerateAssistantOptions{
 		Model:                prepared.Model,
@@ -603,7 +587,7 @@ func (s *taskStarter) generateAssistant(
 		ContextLimitFallback: prepared.ContextLimitFallback,
 		ModelConfig:          prepared.ModelConfig,
 		ProviderOptions:      prepared.ProviderOptions,
-		PublishMessagePart:   publish,
+		PublishMessagePart:   attempt.publish,
 		Logger:               s.opts.Logger,
 		Clock:                s.opts.Clock,
 		Metrics:              s.server.metrics,
@@ -612,7 +596,7 @@ func (s *taskStarter) generateAssistant(
 		return xerrors.Errorf("generate assistant: %w", err)
 	}
 	if len(outcome.Step.Content) == 0 {
-		return s.finishGenerationTurn(ctx, machine, input, attempt, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, generationAttemptRequired)
+		return s.finishGenerationTurn(ctx, machine, input, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, requireGenerationAttempt(attempt.number))
 	}
 	messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{
 		modelConfigID:      prepared.ModelConfigID,
@@ -623,9 +607,9 @@ func (s *taskStarter) generateAssistant(
 		contentVersion:     chatprompt.CurrentContentVersion,
 	})
 	if err != nil {
-		return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired)
+		return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number))
 	}
-	return s.commitGenerationStep(ctx, machine, input, attempt, generationActionGenerateAssistant, messages)
+	return s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionGenerateAssistant, messages)
 }

 func (s *taskStarter) executeLocalTools(
@@ -635,11 +619,11 @@ func (s *taskStarter) executeLocalTools(
 	prepared generationPrepared,
 	decision generationDecision,
 ) error {
-	attempt, _, publish, closeEpisode, err := s.beginGenerationAttempt(ctx, machine, input)
+	attempt, err := s.beginGenerationAttempt(ctx, machine, input)
 	if err != nil {
 		return xerrors.Errorf("beginGenerationAttempt: %w", err)
 	}
-	defer closeEpisode()
+	defer attempt.closeEpisode()
 	provider := ""
 	modelName := ""
 	if prepared.Model != nil {
@@ -662,7 +646,7 @@ func (s *taskStarter) executeLocalTools(
 		ModelName:          modelName,
 		ContextLimit:       prepared.ContextLimitFallback,
 		ToolNameAliases:    subagentToolNameAliases,
-		PublishMessagePart: publish,
+		PublishMessagePart: attempt.publish,
 		Logger:             s.opts.Logger,
 		Metrics:            s.server.metrics,
 		Clock:              s.opts.Clock,
@@ -679,9 +663,9 @@ func (s *taskStarter) executeLocalTools(
 		contentVersion:     chatprompt.CurrentContentVersion,
 	})
 	if err != nil {
-		return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired)
+		return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number))
 	}
-	return s.commitGenerationStep(ctx, machine, input, attempt, generationActionExecuteLocalTools, messages)
+	return s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionExecuteLocalTools, messages)
 }

 func (s *taskStarter) generateCompaction(
@@ -690,16 +674,16 @@ func (s *taskStarter) generateCompaction(
 	input chatWorkerTaskStartInput,
 	prepared generationPrepared,
 ) error {
-	attempt, _, publish, closeEpisode, err := s.beginGenerationAttempt(ctx, machine, input)
+	attempt, err := s.beginGenerationAttempt(ctx, machine, input)
 	if err != nil {
 		return xerrors.Errorf("beginGenerationAttempt: %w", err)
 	}
-	defer closeEpisode()
+	defer attempt.closeEpisode()
 	if prepared.Compaction == nil {
-		return s.finishGenerationError(ctx, machine, input, attempt, xerrors.New("compaction action missing options"), generationAttemptRequired)
+		return s.finishGenerationError(ctx, machine, input, xerrors.New("compaction action missing options"), requireGenerationAttempt(attempt.number))
 	}
 	compactionOpts := prepared.Compaction.Options
-	compactionOpts.PublishMessagePart = publish
+	compactionOpts.PublishMessagePart = attempt.publish
 	outcome, err := chatloop.GenerateCompaction(ctx, compactionOpts)
 	if err != nil {
 		s.server.metrics.RecordCompaction(compactionProvider(compactionOpts), compactionModel(compactionOpts), false, err)
@@ -708,7 +692,7 @@ func (s *taskStarter) generateCompaction(
 	if strings.TrimSpace(outcome.SystemSummary) == "" || strings.TrimSpace(outcome.SummaryReport) == "" {
 		err := xerrors.New("compaction produced no summary")
 		s.server.metrics.RecordCompaction(compactionProvider(compactionOpts), compactionModel(compactionOpts), false, err)
-		return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired)
+		return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number))
 	}
 	messages, err := buildCompactionMessages(buildCompactionMessagesInput{
 		modelConfigID:  prepared.ModelConfigID,
@@ -720,9 +704,9 @@ func (s *taskStarter) generateCompaction(
 	})
 	if err != nil {
 		s.server.metrics.RecordCompaction(compactionProvider(compactionOpts), compactionModel(compactionOpts), false, err)
-		return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired)
+		return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number))
 	}
-	err = s.commitGenerationStep(ctx, machine, input, attempt, generationActionCompact, stepMessagesForCommit{
+	err = s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionCompact, stepMessagesForCommit{
 		Messages:       messages.Messages,
 		VisibleIndexes: visibleMessageIndexes(messages.Messages),
 	})
@@ -747,23 +731,27 @@ func compactionModel(opts chatloop.GenerateCompactionOptions) string {
 	return opts.Model.Model()
 }

+// generationAttempt groups the state a generation action needs after
+// recording a new attempt.
+type generationAttempt struct {
+	number int64
+	// publish streams a message part into the attempt's buffer episode.
+	publish func(codersdk.ChatMessageRole, codersdk.ChatMessagePart)
+	// closeEpisode closes the attempt's buffer episode. It is always
+	// non-nil when beginGenerationAttempt succeeds.
+	closeEpisode func()
+}
+
 func (s *taskStarter) beginGenerationAttempt(
 	ctx context.Context,
 	machine *chatstate.ChatMachine,
 	input chatWorkerTaskStartInput,
-) (int64, messagepartbuffer.Key, func(codersdk.ChatMessageRole, codersdk.ChatMessagePart), func(), error) {
+) (generationAttempt, error) {
 	var attempt int64
 	var committed database.Chat
 	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-		}
-		if err != nil {
-			return xerrors.Errorf("load chat: %w", err)
-		}
-		if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
+		if _, err := loadChatForTask(ctx, store, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
+			return xerrors.Errorf("load chat for task: %w", err)
 		}
 		result, err := tx.RecordGenerationAttempt(chatstate.RecordGenerationAttemptInput{})
 		if err != nil {
@@ -777,7 +765,7 @@ func (s *taskStarter) beginGenerationAttempt(
 		return nil
 	})
 	if err != nil {
-		return 0, messagepartbuffer.Key{}, nil, nil, normalizeTaskTransitionError(err, "record generation attempt")
+		return generationAttempt{}, normalizeTaskTransitionError(err, "record generation attempt")
 	}
 	key := messagepartbuffer.Key{
 		ChatID:            input.ChatID,
@@ -785,15 +773,17 @@ func (s *taskStarter) beginGenerationAttempt(
 		GenerationAttempt: attempt,
 	}
 	if err := s.opts.MessagePartBuffer.CreateEpisode(key); err != nil && ctx.Err() == nil {
-		return 0, messagepartbuffer.Key{}, nil, nil, taskRetryableError{err: xerrors.Errorf("create message part episode: %w", err)}
+		return generationAttempt{}, taskRetryableError{err: xerrors.Errorf("create message part episode: %w", err)}
 	}
-	publish := func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) {
-		_ = s.opts.MessagePartBuffer.AddPart(key, role, part)
-	}
-	closeEpisode := func() {
-		_ = s.opts.MessagePartBuffer.CloseEpisode(key)
-	}
-	return attempt, key, publish, closeEpisode, nil
+	return generationAttempt{
+		number: attempt,
+		publish: func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) {
+			_ = s.opts.MessagePartBuffer.AddPart(key, role, part)
+		},
+		closeEpisode: func() {
+			_ = s.opts.MessagePartBuffer.CloseEpisode(key)
+		},
+	}, nil
 }

 func (s *taskStarter) commitGenerationStep(
@@ -805,20 +795,13 @@ func (s *taskStarter) commitGenerationStep(
 	messages stepMessagesForCommit,
 ) error {
 	if len(messages.Messages) == 0 {
-		return s.finishGenerationTurn(ctx, machine, input, attempt, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, generationAttemptRequired)
+		return s.finishGenerationTurn(ctx, machine, input, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, requireGenerationAttempt(attempt))
 	}
 	var committed database.Chat
 	insertedMessages := []runnerActionMessage{}
 	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-		}
-		if err != nil {
-			return xerrors.Errorf("load chat: %w", err)
-		}
-		if err := verifyGenerationFence(locked, input, attempt); err != nil {
-			return xerrors.Errorf("verifyGenerationFence: %w", err)
+		if _, err := loadChatForGeneration(ctx, store, input, requireGenerationAttempt(attempt)); err != nil {
+			return xerrors.Errorf("load chat for generation: %w", err)
 		}
 		commitResult, err := tx.CommitStep(chatstate.CommitStepInput{Messages: messages.Messages})
 		if err != nil {
@@ -852,23 +835,17 @@ func (s *taskStarter) enterRequiresAction(
 ) error {
 	var committed database.Chat
 	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-		}
-		if err != nil {
-			return xerrors.Errorf("load chat: %w", err)
-		}
-		if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
+		if _, err := loadChatForTask(ctx, store, 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)
 		}
-		committed, err = store.GetChatByID(ctx, input.ChatID)
+		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 {
@@ -884,36 +861,65 @@ func (s *taskStarter) enterRequiresAction(
 	})
 }

-type generationAttemptFence int
+// generationAttemptFence controls whether a terminal generation
+// transition also verifies that the chat's GenerationAttempt counter
+// matches an expected value.
+type generationAttemptFence struct {
+	required bool
+	attempt  int64
+}

-const (
-	generationAttemptNotRequired generationAttemptFence = iota
-	generationAttemptRequired
-)
+// generationAttemptNotRequired skips the generation attempt fence; only the
+// running-task fence is verified.
+var generationAttemptNotRequired = generationAttemptFence{}
+
+// requireGenerationAttempt returns a fence that also verifies the chat's
+// generation attempt matches the given value.
+func requireGenerationAttempt(attempt int64) generationAttemptFence {
+	return generationAttemptFence{required: true, attempt: attempt}
+}
+
+// loadChatForGeneration loads the chat and verifies the running-task fence,
+// additionally verifying the generation attempt fence when required.
+func loadChatForGeneration(
+	ctx context.Context,
+	store database.Store,
+	input chatWorkerTaskStartInput,
+	fence generationAttemptFence,
+) (database.Chat, error) {
+	chat, err := loadChatForTask(ctx, store, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true})
+	if err != nil {
+		return database.Chat{}, err
+	}
+	if fence.required && chat.GenerationAttempt != fence.attempt {
+		return database.Chat{}, errors.Join(errTaskExpectedExit, xerrors.Errorf("generation fence mismatch: %d != %d", chat.GenerationAttempt, fence.attempt))
+	}
+	return chat, nil
+}
+
+// recordGenerationFinishFailure records an error outcome on the debug turn
+// when a terminal generation transition fails, so the debug run is not
+// finalized as interrupted when work was actually done. It skips expected
+// exits (fence lost, chat deleted) where another task owns the turn, and
+// retryable errors where a task retry will record the real outcome.
+func recordGenerationFinishFailure(turn *runnerDebugTurn, err error) {
+	if errors.Is(err, errTaskExpectedExit) || errors.Is(err, errTaskRetryable) {
+		return
+	}
+	turn.RecordOutcome(chatdebug.StatusError)
+}

 func (s *taskStarter) finishGenerationTurn(
 	ctx context.Context,
 	machine *chatstate.ChatMachine,
 	input chatWorkerTaskStartInput,
-	attempt int64,
 	decision generationDecision,
-	attemptFence generationAttemptFence,
+	fence generationAttemptFence,
 ) error {
 	var committed database.Chat
 	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-		}
-		if err != nil {
-			return xerrors.Errorf("load chat: %w", err)
-		}
-		if attemptFence == generationAttemptRequired {
-			if err := verifyGenerationFence(locked, input, attempt); err != nil {
-				return xerrors.Errorf("verifyGenerationFence: %w", err)
-			}
-		} else if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
+		if _, err := loadChatForGeneration(ctx, store, input, fence); err != nil {
+			return xerrors.Errorf("load chat for generation: %w", err)
 		}
 		finishResult, err := tx.FinishTurn(chatstate.FinishTurnInput{})
 		if err != nil {
@@ -926,7 +932,9 @@ func (s *taskStarter) finishGenerationTurn(
 		return nil
 	})
 	if err != nil {
-		return normalizeTaskTransitionError(err, "finish generation turn")
+		err := normalizeTaskTransitionError(err, "finish generation turn")
+		recordGenerationFinishFailure(input.DebugTurn, err)
+		return err
 	}
 	input.DebugTurn.RecordOutcome(chatdebug.StatusCompleted)
 	watchCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), postCommitWatchPublishTimeout)
@@ -950,9 +958,8 @@ func (s *taskStarter) finishGenerationError(
 	ctx context.Context,
 	machine *chatstate.ChatMachine,
 	input chatWorkerTaskStartInput,
-	attempt int64,
 	cause error,
-	attemptFence generationAttemptFence,
+	fence generationAttemptFence,
 ) error {
 	classified := chaterror.Classify(cause)
 	// Log the unsanitized cause before persisting so administrators can
@@ -971,31 +978,23 @@ func (s *taskStarter) finishGenerationError(
 	lastError, message := generationLastError(cause)
 	var committed database.Chat
 	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-		}
-		if err != nil {
-			return xerrors.Errorf("load chat: %w", err)
-		}
-		if attemptFence == generationAttemptRequired {
-			if err := verifyGenerationFence(locked, input, attempt); err != nil {
-				return xerrors.Errorf("verifyGenerationFence: %w", err)
-			}
-		} else if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
+		if _, err := loadChatForGeneration(ctx, store, 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)
 		}
-		committed, err = store.GetChatByID(ctx, input.ChatID)
+		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 {
-		return normalizeTaskTransitionError(err, "finish generation error")
+		err := normalizeTaskTransitionError(err, "finish generation error")
+		recordGenerationFinishFailure(input.DebugTurn, err)
+		return err
 	}
 	input.DebugTurn.RecordOutcome(chatdebug.StatusError)
 	if err := s.publishWatchAndRoute(ctx, committed, codersdk.ChatWatchEventKindStatusChange); err != nil {
@@ -1026,25 +1025,12 @@ func generationLastError(err error) (pqtype.NullRawMessage, string) {
 }

 func (s *taskStarter) afterGenerationOutcome(ctx context.Context, outcome generationOutcome) error {
-	if s.server == nil {
-		return nil
-	}
 	if err := s.server.afterGenerationOutcome(ctx, outcome); err != nil {
 		return taskRetryableError{err: xerrors.Errorf("generation post-outcome side effects: %w", err)}
 	}
 	return nil
 }

-func verifyGenerationFence(chat database.Chat, input chatWorkerTaskStartInput, attempt int64) error {
-	if err := verifyTaskFence(chat, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil {
-		return xerrors.Errorf("verifyTaskFence: %w", err)
-	}
-	if chat.GenerationAttempt != attempt {
-		return errors.Join(errTaskExpectedExit, xerrors.Errorf("generation fence mismatch: %d != %d", chat.GenerationAttempt, attempt))
-	}
-	return nil
-}
-
 func stepDataFromPersisted(step chatloop.PersistedStep) stepData {
 	return stepData{
 		Content:              step.Content,
@@ -0,0 +1,49 @@
+package chatd //nolint:testpackage // Exercises unexported generation helpers.
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/coderd/x/chatd/chatdebug"
+	"github.com/coder/coder/v2/coderd/x/chatd/chatstate"
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestRecordGenerationFinishFailure(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name         string
+		err          error
+		wantRecorded bool
+	}{
+		{
+			name:         "TerminalFailureRecordsError",
+			err:          normalizeTaskTransitionError(chatstate.ErrTransitionNotAllowed, "finish generation error"),
+			wantRecorded: true,
+		},
+		{
+			name:         "ExpectedExitSkips",
+			err:          normalizeTaskTransitionError(errTaskExpectedExit, "finish generation error"),
+			wantRecorded: false,
+		},
+		{
+			name:         "RetryableSkips",
+			err:          normalizeTaskTransitionError(xerrors.New("transient infrastructure failure"), "finish generation error"),
+			wantRecorded: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			turn := newRunnerDebugTurn(testutil.Context(t, testutil.WaitShort), testutil.Logger(t))
+			recordGenerationFinishFailure(turn, tt.err)
+			require.Equal(t, tt.wantRecorded, turn.statusSet)
+			if tt.wantRecorded {
+				require.Equal(t, chatdebug.StatusError, turn.status)
+			}
+		})
+	}
+}
@@ -45,13 +45,15 @@ type publishedEvent struct {
 }

 type recordingPubsub struct {
-	inner  dbpubsub.Pubsub
+	// Embed the full pubsub so a recordingPubsub can also stand in for a
+	// Server's pubsub, recording every published event.
+	dbpubsub.Pubsub
 	mu     sync.Mutex
 	events []publishedEvent
 }

 func newRecordingPubsub(inner dbpubsub.Pubsub) *recordingPubsub {
-	return &recordingPubsub{inner: inner}
+	return &recordingPubsub{Pubsub: inner}
 }

 func (p *recordingPubsub) Publish(channel string, payload []byte) error {
@@ -61,11 +63,7 @@ func (p *recordingPubsub) Publish(channel string, payload []byte) error {
 		payload: append([]byte(nil), payload...),
 	})
 	p.mu.Unlock()
-	return p.inner.Publish(channel, payload)
-}
-
-func (p *recordingPubsub) SubscribeWithErr(channel string, listener dbpubsub.ListenerWithErr) (func(), error) {
-	return p.inner.SubscribeWithErr(channel, listener)
+	return p.Pubsub.Publish(channel, payload)
 }

 func (p *recordingPubsub) ownershipMessages(t *testing.T) []coderdpubsub.ChatStateOwnershipMessage {
@@ -268,9 +266,25 @@ func testOptions(t *testing.T, f *workerTestFixture, starter chatWorkerTaskStart
 	}
 }

+// newUnstartedServer builds a real Server backed by the given pubsub and
+// store. The server is never started; it only provides the dependencies
+// that workers and task starters dereference.
+func newUnstartedServer(t *testing.T, ps dbpubsub.Pubsub, db database.Store) *Server {
+	t.Helper()
+	server := New(ps, Config{
+		Logger:    testutil.Logger(t),
+		Database:  db,
+		ReplicaID: uuid.New(),
+	})
+	t.Cleanup(func() { _ = server.Close() })
+	return server
+}
+
 func startWorker(t *testing.T, opts chatWorkerOptions) *chatWorker {
 	t.Helper()
-	worker, err := newChatWorker(nil, opts)
+	ps, ok := opts.Pubsub.(dbpubsub.Pubsub)
+	require.True(t, ok, "worker pubsub must implement the full pubsub interface")
+	worker, err := newChatWorker(newUnstartedServer(t, ps, opts.Store), opts)
 	require.NoError(t, err)
 	require.NoError(t, worker.Start(context.Background()))
 	t.Cleanup(func() { require.NoError(t, worker.Close()) })
@@ -193,6 +193,9 @@ func newTaskStarter(
 	routeStateHint func(context.Context, runnerStateUpdate),
 	requestCleanup func(context.Context, runnerKey),
 ) (*taskStarter, error) {
+	if server == nil {
+		return nil, xerrors.New("chatworker: server is required")
+	}
 	if opts.Store == nil {
 		return nil, xerrors.New("chatworker: task store is required")
 	}
@@ -241,17 +244,11 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt
 	machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID)
 	var chat database.Chat
 	err := machine.ReadLock(ctx, func(store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load locked chat: %w", err))
-		}
+		loadedChat, err := loadChatForTask(ctx, store, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true})
 		if err != nil {
-			return xerrors.Errorf("load locked chat: %w", err)
+			return xerrors.Errorf("load chat for task: %w", err)
 		}
-		if err := verifyTaskFence(locked, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
-		}
-		chat = locked
+		chat = loadedChat
 		return nil
 	})
 	if err != nil {
@@ -293,18 +290,12 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt

 	var committed database.Chat
 	err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-		}
+		chat, err := loadChatForTask(ctx, store, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true})
 		if err != nil {
-			return xerrors.Errorf("load chat: %w", err)
-		}
-		if err := verifyTaskFence(locked, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
+			return xerrors.Errorf("load chat for task: %w", err)
 		}
 		messages := partialMessages
-		committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, locked, s.opts.Clock.Now("chatworker", "interrupt"))
+		committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, s.opts.Clock.Now("chatworker", "interrupt"))
 		if err != nil {
 			return xerrors.Errorf("committed pending local tool cancellation messages: %w", err)
 		}
@@ -339,7 +330,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt

 func (s *taskStarter) runAfterInterruptionOutcome(ctx context.Context, outcome interruptionOutcome) error {
 	afterOutcome := s.afterInterruptionOutcome
-	if afterOutcome == nil && s.server != nil {
+	if afterOutcome == nil {
 		afterOutcome = s.server.afterInterruptionOutcome
 	}
 	if afterOutcome == nil {
@@ -383,17 +374,11 @@ func decideRequiresActionTimeout(
 ) (requiresActionTimeoutDecision, error) {
 	var decision requiresActionTimeoutDecision
 	err := machine.ReadLock(ctx, func(store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load locked chat: %w", err))
-		}
+		chat, err := loadChatForTask(ctx, store, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true})
 		if err != nil {
-			return xerrors.Errorf("load locked chat: %w", err)
+			return xerrors.Errorf("load chat for task: %w", err)
 		}
-		if err := verifyTaskFence(locked, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
-		}
-		if !locked.RequiresActionDeadlineAt.Valid {
+		if !chat.RequiresActionDeadlineAt.Valid {
 			decision.cancel = true
 			decision.reason = "Tool execution canceled because the action deadline was missing"
 			return nil
@@ -402,8 +387,8 @@ func decideRequiresActionTimeout(
 		if err != nil {
 			return xerrors.Errorf("get database time: %w", err)
 		}
-		if now.Before(locked.RequiresActionDeadlineAt.Time) {
-			decision.waitUntil = locked.RequiresActionDeadlineAt
+		if now.Before(chat.RequiresActionDeadlineAt.Time) {
+			decision.waitUntil = chat.RequiresActionDeadlineAt
 			return nil
 		}
 		decision.cancel = true
@@ -439,22 +424,16 @@ func (s *taskStarter) cancelRequiresAction(
 ) error {
 	var committed database.Chat
 	err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load locked chat: %w", err))
-		}
+		chat, err := loadChatForTask(ctx, store, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true})
 		if err != nil {
-			return xerrors.Errorf("load chat: %w", err)
+			return xerrors.Errorf("load chat for task: %w", err)
 		}
-		if err := verifyTaskFence(locked, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true}); err != nil {
-			return xerrors.Errorf("verifyTaskFence: %w", err)
-		}
-		if locked.RequiresActionDeadlineAt.Valid {
+		if chat.RequiresActionDeadlineAt.Valid {
 			now, err := store.GetDatabaseNow(ctx)
 			if err != nil {
 				return xerrors.Errorf("get database time: %w", err)
 			}
-			if now.Before(locked.RequiresActionDeadlineAt.Time) {
+			if now.Before(chat.RequiresActionDeadlineAt.Time) {
 				return errors.Join(errTaskExpectedExit, xerrors.Errorf("requires action deadline is in the future"))
 			}
 		}
@@ -480,19 +459,19 @@ func (s *taskStarter) StartAbandon(ctx context.Context, input chatWorkerTaskStar
 	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 {
-		locked, err := store.GetChatByID(ctx, input.ChatID)
-		if errors.Is(err, sql.ErrNoRows) {
-			mismatch = true
-			return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
-		}
+		chat, err := store.GetChatByID(ctx, input.ChatID)
 		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(locked, input) {
+		if !ownedByTask(chat, input) {
 			mismatch = true
 			return errors.Join(errTaskExpectedExit, xerrors.Errorf("chat not owned by task"))
 		}
-		if err := verifyTaskFence(locked, input, input.Status, taskFenceOptions{requireHistory: true, allowArchived: true}); err != nil {
+		if err := verifyTaskFence(chat, input, input.Status, taskFenceOptions{requireHistory: true, allowArchived: true}); err != nil {
 			return xerrors.Errorf("task fence mismatch: %w", err)
 		}
 		if _, err := tx.Abandon(chatstate.AbandonInput{}); err != nil {
@@ -595,6 +574,30 @@ type taskFenceOptions struct {
 	allowArchived  bool
 }

+// loadChatForTask loads the chat row and verifies the task fence in one
+// step so call sites cannot skip the fence check. It returns an error
+// wrapping errTaskExpectedExit when the chat no longer exists or the fence
+// no longer matches.
+func loadChatForTask(
+	ctx context.Context,
+	store database.Store,
+	input chatWorkerTaskStartInput,
+	status database.ChatStatus,
+	opts taskFenceOptions,
+) (database.Chat, error) {
+	chat, err := store.GetChatByID(ctx, input.ChatID)
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return database.Chat{}, errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
+		}
+		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,
@@ -578,7 +578,7 @@ func TestGenerationTask_RecordRetryState(t *testing.T) {
 	recorder := newTaskSideEffectRecorder()
 	starter := newTestTaskStarter(t, f, recorder)

-	attempt, _, _, closeEpisode, err := starter.beginGenerationAttempt(
+	attempt, err := starter.beginGenerationAttempt(
 		testutil.Context(t, testutil.WaitLong),
 		chatstate.NewChatMachine(f.db, f.pubsub, chat.ID),
 		chatWorkerTaskStartInput{
@@ -590,8 +590,8 @@ func TestGenerationTask_RecordRetryState(t *testing.T) {
 		},
 	)
 	require.NoError(t, err)
-	closeEpisode()
-	require.Equal(t, int64(1), attempt)
+	attempt.closeEpisode()
+	require.Equal(t, int64(1), attempt.number)
 	before, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID)
 	require.NoError(t, err)
 	require.False(t, before.RetryState.Valid)
@@ -650,7 +650,7 @@ func TestGenerationTask_RecordRetryStateUsesDurableGenerationAttempt(t *testing.
 	machine := chatstate.NewChatMachine(f.db, f.pubsub, chat.ID)

 	for range 3 {
-		attempt, _, _, closeEpisode, err := starter.beginGenerationAttempt(
+		attempt, err := starter.beginGenerationAttempt(
 			testutil.Context(t, testutil.WaitLong),
 			machine,
 			chatWorkerTaskStartInput{
@@ -662,8 +662,8 @@ func TestGenerationTask_RecordRetryStateUsesDurableGenerationAttempt(t *testing.
 			},
 		)
 		require.NoError(t, err)
-		closeEpisode()
-		require.Positive(t, attempt)
+		attempt.closeEpisode()
+		require.Positive(t, attempt.number)
 	}

 	decision, err := starter.recordGenerationRetry(
@@ -714,10 +714,10 @@ func TestGenerationTask_RecordRetryStateClearedByNextAttempt(t *testing.T) {
 		Status:         database.ChatStatusRunning,
 	}

-	attempt, _, _, closeEpisode, err := starter.beginGenerationAttempt(testutil.Context(t, testutil.WaitLong), machine, input)
+	attempt, err := starter.beginGenerationAttempt(testutil.Context(t, testutil.WaitLong), machine, input)
 	require.NoError(t, err)
-	closeEpisode()
-	require.Equal(t, int64(1), attempt)
+	attempt.closeEpisode()
+	require.Equal(t, int64(1), attempt.number)
 	_, err = starter.recordGenerationRetry(
 		testutil.Context(t, testutil.WaitLong),
 		machine,
@@ -734,10 +734,10 @@ func TestGenerationTask_RecordRetryStateClearedByNextAttempt(t *testing.T) {
 	require.NoError(t, err)
 	require.True(t, withRetry.RetryState.Valid)

-	attempt, _, _, closeEpisode, err = starter.beginGenerationAttempt(testutil.Context(t, testutil.WaitLong), machine, input)
+	attempt, err = starter.beginGenerationAttempt(testutil.Context(t, testutil.WaitLong), machine, input)
 	require.NoError(t, err)
-	closeEpisode()
-	require.Equal(t, int64(2), attempt)
+	attempt.closeEpisode()
+	require.Equal(t, int64(2), attempt.number)
 	after, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID)
 	require.NoError(t, err)
 	require.False(t, after.RetryState.Valid)
@@ -755,7 +755,7 @@ func TestGenerationTask_RecordRetryStateStaleFenceExits(t *testing.T) {
 	acquired := f.acquireChat(t, chat.ID, workerID, runnerID)
 	starter := newTestTaskStarter(t, f, newTaskSideEffectRecorder())
 	machine := chatstate.NewChatMachine(f.db, f.pubsub, chat.ID)
-	attempt, _, _, closeEpisode, err := starter.beginGenerationAttempt(
+	attempt, err := starter.beginGenerationAttempt(
 		testutil.Context(t, testutil.WaitLong),
 		machine,
 		chatWorkerTaskStartInput{
@@ -767,8 +767,8 @@ func TestGenerationTask_RecordRetryStateStaleFenceExits(t *testing.T) {
 		},
 	)
 	require.NoError(t, err)
-	closeEpisode()
-	require.Equal(t, int64(1), attempt)
+	attempt.closeEpisode()
+	require.Equal(t, int64(1), attempt.number)

 	otherWorkerID := uuid.New()
 	otherRunnerID := uuid.New()
@@ -855,6 +855,7 @@ func TestRunner_StartsRealAbandonTask(t *testing.T) {
 type taskTestFixture struct {
 	db     database.Store
 	pubsub *taskRecordingPubsub
+	rawPS  dbpubsub.Pubsub
 	sqlDB  *sql.DB
 	user   database.User
 	org    database.Organization
@@ -875,7 +876,7 @@ func newTaskTestFixture(t *testing.T) *taskTestFixture {
 	})
 	model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{IsDefault: true})
 	apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID})
-	return &taskTestFixture{db: db, pubsub: newTaskRecordingPubsub(ps), sqlDB: sqlDB, user: user, org: org, model: model, apiKey: apiKey}
+	return &taskTestFixture{db: db, pubsub: newTaskRecordingPubsub(ps), rawPS: ps, sqlDB: sqlDB, user: user, org: org, model: model, apiKey: apiKey}
 }

 func (f *taskTestFixture) createRunningChat(t *testing.T) database.Chat {
@@ -1124,7 +1125,7 @@ func startRealTaskWorker(t *testing.T, f *taskTestFixture) *chatWorker {
 	t.Helper()
 	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
 	t.Cleanup(buffer.Close)
-	worker, err := newChatWorker(nil, chatWorkerOptions{
+	worker, err := newChatWorker(newUnstartedServer(t, f.rawPS, f.db), chatWorkerOptions{
 		WorkerID:                   uuid.New(),
 		Store:                      f.db,
 		Pubsub:                     f.pubsub,
@@ -1246,7 +1247,7 @@ func newTestTaskStarter(t *testing.T, f *taskTestFixture, recorder *taskSideEffe
 	t.Helper()
 	buffer := messagepartbuffer.New(messagepartbuffer.Options{})
 	t.Cleanup(buffer.Close)
-	starter, err := newTaskStarter(nil, chatWorkerOptions{
+	starter, err := newTaskStarter(newUnstartedServer(t, f.rawPS, f.db), chatWorkerOptions{
 		Store:                   f.db,
 		Pubsub:                  f.pubsub,
 		Logger:                  slog.Make(),
@@ -33,6 +33,9 @@ type chatWorker struct {
 // newChatWorker constructs a chat worker. The worker is idle until Start is
 // called.
 func newChatWorker(server *Server, opts chatWorkerOptions) (*chatWorker, error) {
+	if server == nil {
+		return nil, xerrors.New("chatworker: server is required")
+	}
 	withDefaults, err := opts.withDefaults()
 	if err != nil {
 		return nil, err
@@ -18,16 +18,23 @@ import (
 func TestWorker_NewRequiresTaskStarterOrMessagePartBuffer(t *testing.T) {
 	t.Parallel()
 	f := newWorkerTestFixture(t)
-	_, err := newChatWorker(nil, chatWorkerOptions{WorkerID: uuid.New(), Store: f.db, Pubsub: f.pubsub})
+	_, err := newChatWorker(newUnstartedServer(t, f.pubsub, f.db), chatWorkerOptions{WorkerID: uuid.New(), Store: f.db, Pubsub: f.pubsub})
 	require.ErrorContains(t, err, "task starter or message part buffer is required")
 }

+func TestWorker_NewRequiresServer(t *testing.T) {
+	t.Parallel()
+	f := newWorkerTestFixture(t)
+	_, err := newChatWorker(nil, testOptions(t, f, newRecordingTaskStarter()))
+	require.ErrorContains(t, err, "server is required")
+}
+
 func TestWorker_NewRequiresWorkerID(t *testing.T) {
 	t.Parallel()
 	f := newWorkerTestFixture(t)
 	opts := testOptions(t, f, newRecordingTaskStarter())
 	opts.WorkerID = uuid.Nil
-	_, err := newChatWorker(nil, opts)
+	_, err := newChatWorker(newUnstartedServer(t, f.pubsub, f.db), opts)
 	require.ErrorContains(t, err, "worker ID is required")
 }

@@ -37,7 +44,7 @@ func TestWorker_UsesConfiguredWorkerID(t *testing.T) {
 	starter := newRecordingTaskStarter()
 	opts := testOptions(t, f, starter)
 	workerID := opts.WorkerID
-	worker, err := newChatWorker(nil, opts)
+	worker, err := newChatWorker(newUnstartedServer(t, f.pubsub, f.db), opts)
 	require.NoError(t, err)
 	require.Equal(t, workerID, worker.chatWorkerID())
 	require.NoError(t, worker.Start(context.Background()))