From 96130e2bc535a969671afd4ab4da6cc742196c73 Mon Sep 17 00:00:00 2001 From: Hugo Dutka Date: Mon, 6 Jul 2026 18:37:45 +0200 Subject: [PATCH] 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. --- coderd/x/chatd/auto_archive.go | 26 +- coderd/x/chatd/auto_archive_internal_test.go | 10 +- coderd/x/chatd/generation.go | 274 +++++++++---------- coderd/x/chatd/generation_internal_test.go | 49 ++++ coderd/x/chatd/helpers_test.go | 30 +- coderd/x/chatd/tasks.go | 97 +++---- coderd/x/chatd/tasks_test.go | 37 +-- coderd/x/chatd/worker.go | 3 + coderd/x/chatd/worker_internal_test.go | 13 +- 9 files changed, 294 insertions(+), 245 deletions(-) create mode 100644 coderd/x/chatd/generation_internal_test.go diff --git a/coderd/x/chatd/auto_archive.go b/coderd/x/chatd/auto_archive.go index 6283aa3afe..ac460f74a3 100644 --- a/coderd/x/chatd/auto_archive.go +++ b/coderd/x/chatd/auto_archive.go @@ -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 diff --git a/coderd/x/chatd/auto_archive_internal_test.go b/coderd/x/chatd/auto_archive_internal_test.go index 0d50bfb7f0..e08709ae15 100644 --- a/coderd/x/chatd/auto_archive_internal_test.go +++ b/coderd/x/chatd/auto_archive_internal_test.go @@ -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 } diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index da3b1df6ba..ac3bcd9a12 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -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, diff --git a/coderd/x/chatd/generation_internal_test.go b/coderd/x/chatd/generation_internal_test.go new file mode 100644 index 0000000000..9090506b67 --- /dev/null +++ b/coderd/x/chatd/generation_internal_test.go @@ -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) + } + }) + } +} diff --git a/coderd/x/chatd/helpers_test.go b/coderd/x/chatd/helpers_test.go index bb295728b0..2a5f0a42dd 100644 --- a/coderd/x/chatd/helpers_test.go +++ b/coderd/x/chatd/helpers_test.go @@ -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()) }) diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index 82a4ada1e3..c1c0c840e4 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -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, diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 4f549554e2..31882f5991 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -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(), diff --git a/coderd/x/chatd/worker.go b/coderd/x/chatd/worker.go index 7b3e8d5666..5aed8e5713 100644 --- a/coderd/x/chatd/worker.go +++ b/coderd/x/chatd/worker.go @@ -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 diff --git a/coderd/x/chatd/worker_internal_test.go b/coderd/x/chatd/worker_internal_test.go index f01bb0d69c..a7bae0d7dd 100644 --- a/coderd/x/chatd/worker_internal_test.go +++ b/coderd/x/chatd/worker_internal_test.go @@ -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()))