fix(coderd/x/chatd): resolve inflight race (#26460)

Using `WaitGroup.Go` must be synchronized with `WaitGroup.Wait`
according to [go docs](https://pkg.go.dev/sync#WaitGroup.Go):

> If the WaitGroup is empty, Go must happen before a
[WaitGroup.Wait](https://pkg.go.dev/sync#WaitGroup.Wait).

There were a couple of places in chatd that violated this principle.
This was caught as a data race in
https://github.com/coder/internal/issues/1599. This PR ensures that all
functions that spawn inflight goroutines synchronize with each other.

I also noticed that inflight goroutines may be spawned after the server
is closed, which was surprising and looked like a bug. This PR therefore
also introduces a mechanism that disallows spawning inflight goroutines
after the server is closed, and ensures that any code that tries doing
it logs an error.

Closes https://github.com/coder/internal/issues/1599.

coderd/x/chatd/chatd.go

@@ -140,6 +140,7 @@ var (
 			"The agent may still be reachable on the next attempt.",
 	)
 	errChatExternalAgentUnavailable = xerrors.New("external workspace agent unavailable")
+	errInflightClosed               = xerrors.New("chatd server inflight closed")
 )
 
 type chatExternalAgentUnavailableError struct {
@@ -162,11 +163,12 @@ func newChatExternalAgentUnavailableError(agent database.WorkspaceAgent) error {
 
 // Server handles background processing of pending chats.
 type Server struct {
-	cancel     context.CancelFunc
-	ctx        context.Context
-	wg         sync.WaitGroup
-	inflight   sync.WaitGroup
-	inflightMu sync.Mutex
+	cancel         context.CancelFunc
+	ctx            context.Context
+	wg             sync.WaitGroup
+	inflight       sync.WaitGroup
+	inflightMu     sync.Mutex
+	inflightClosed atomic.Bool
 
 	db       database.Store
 	workerID uuid.UUID
@@ -4181,9 +4183,16 @@ func (p *Server) appendRootChatTools(
 		// burns the full budget for nothing.
 		snapshot := opts.workspaceCtx.currentChatSnapshot()
 		if snapshot.WorkspaceID.Valid && snapshot.AgentID.Valid {
-			p.inflight.Go(func() {
+			if err := p.goInflight(func() {
 				p.primeWorkspaceMCPCache(opts.primerCtx, p.logger, snapshot.ID, opts.workspaceCtx)
-			})
+			}); err != nil {
+				p.logger.Error(context.WithoutCancel(ctx), "failed to schedule workspace MCP cache primer",
+					slog.F("chat_id", snapshot.ID),
+					slog.F("workspace_id", snapshot.WorkspaceID.UUID),
+					slog.F("agent_id", snapshot.AgentID.UUID),
+					slog.Error(err),
+				)
+			}
 		}
 	}
 
@@ -5128,10 +5137,7 @@ func (p *Server) finalizeSuccessfulTurnStatusLabelWithAfterFunc(
 	logger slog.Logger,
 	afterFinalize func(context.Context, string),
 ) {
-	// This helper runs during processChat cleanup, while processChat is
-	// still counted in p.inflight. Do not take inflightMu here because
-	// drainInflight holds it while waiting.
-	p.inflight.Go(func() {
+	if err := p.goInflight(func() {
 		finalizeCtx := context.WithoutCancel(ctx)
 		statusLabel := p.generateFinalTurnStatusLabel(finalizeCtx, chat, status, runResult, logger)
 		logger.Debug(finalizeCtx, "generated chat turn status label",
@@ -5143,7 +5149,13 @@ func (p *Server) finalizeSuccessfulTurnStatusLabelWithAfterFunc(
 		p.updateLastTurnSummary(finalizeCtx, chat, chat.HistoryVersion, statusLabel, logger)
 
 		afterFinalize(finalizeCtx, statusLabel)
-	})
+	}); err != nil {
+		logger.Error(context.WithoutCancel(ctx), "failed to schedule chat turn status finalization",
+			slog.F("chat_id", chat.ID),
+			slog.F("status", status),
+			slog.Error(err),
+		)
+	}
 }
 
 func (p *Server) generateFinalTurnStatusLabel(
@@ -5225,12 +5237,16 @@ func (p *Server) setLastTurnSummaryAsync(
 	if chat.LastTurnSummary.Valid && strings.TrimSpace(chat.LastTurnSummary.String) == summary {
 		return
 	}
-	// This helper runs during processChat cleanup, while processChat is
-	// still counted in p.inflight. Do not take inflightMu here because
-	// drainInflight holds it while waiting.
-	p.inflight.Go(func() {
+	if err := p.goInflight(func() {
 		p.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.HistoryVersion, summary, logger)
-	})
+	}); err != nil {
+		logger.Error(context.WithoutCancel(ctx), "failed to schedule chat turn summary update",
+			slog.F("chat_id", chat.ID),
+			slog.F("expected_history_version", chat.HistoryVersion),
+			slog.F("summary_length", len(summary)),
+			slog.Error(err),
+		)
+	}
 }
 
 func (p *Server) clearLastTurnSummaryAsync(
@@ -5238,12 +5254,15 @@ func (p *Server) clearLastTurnSummaryAsync(
 	chat database.Chat,
 	logger slog.Logger,
 ) {
-	// This helper runs during processChat cleanup, while processChat is
-	// still counted in p.inflight. Do not take inflightMu here because
-	// drainInflight holds it while waiting.
-	p.inflight.Go(func() {
+	if err := p.goInflight(func() {
 		p.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.HistoryVersion, "", logger)
-	})
+	}); err != nil {
+		logger.Error(context.WithoutCancel(ctx), "failed to schedule chat turn summary clear",
+			slog.F("chat_id", chat.ID),
+			slog.F("expected_history_version", chat.HistoryVersion),
+			slog.Error(err),
+		)
+	}
 }
 
 // updateLastTurnSummary writes the cached sidebar summary for a chat.
@@ -5325,6 +5344,7 @@ func (p *Server) dispatchPush(
 
 // Close stops the processor and waits for it to finish.
 func (p *Server) Close() error {
+	p.closeInflightAdmission()
 	if unsub := p.configCacheUnsubscribe; unsub != nil {
 		p.configCacheUnsubscribe = nil
 		unsub()
@@ -5346,10 +5366,30 @@ func (p *Server) Close() error {
 	return nil
 }
 
-// drainInflight waits for all in-flight operations to complete.
-// It acquires inflightMu to prevent processOnce from spawning
-// new goroutines (via inflight.Add) concurrently with Wait,
-// which would violate sync.WaitGroup's contract.
+func (p *Server) goInflight(f func()) error {
+	if p.inflightClosed.Load() {
+		return errInflightClosed
+	}
+
+	// Acquire inflightMu around the inflight.Go so Close() cannot
+	// call drainInflight concurrently when the counter is at zero.
+	// See drainInflight for the WaitGroup contract this preserves.
+	p.inflightMu.Lock()
+	defer p.inflightMu.Unlock()
+	if p.inflightClosed.Load() {
+		return errInflightClosed
+	}
+	p.inflight.Go(f)
+	return nil
+}
+
+func (p *Server) closeInflightAdmission() {
+	p.inflightClosed.Store(true)
+}
+
+// drainInflight waits for already-admitted in-flight operations to complete.
+// It acquires inflightMu so Wait cannot race with a positive Add from
+// goInflight when the WaitGroup counter is zero.
 //
 // https://pkg.go.dev/sync#WaitGroup.Add
 // > Note that calls with a positive delta that occur when the counter is zero must happen before a Wait.

coderd/x/chatd/chatd_debug.go

@@ -76,15 +76,7 @@ func (p *Server) scheduleDebugCleanup(
 		return
 	}
 
-	// Acquire inflightMu around the positive Add so Close() cannot
-	// call drainInflight concurrently when the counter is at zero.
-	// See drainInflight for the WaitGroup contract this preserves.
-	p.inflightMu.Lock()
-	p.inflight.Add(1)
-	p.inflightMu.Unlock()
-	go func() {
-		defer p.inflight.Done()
-
+	if err := p.goInflight(func() {
 		cleanupCtx := context.WithoutCancel(ctx)
 		for attempt := 0; attempt < debugCleanupAttempts; attempt++ {
 			if attempt > 0 {
@@ -106,7 +98,11 @@ func (p *Server) scheduleDebugCleanup(
 			logFields = append(logFields, slog.Error(err))
 			p.logger.Warn(cleanupCtx, logMessage, logFields...)
 		}
-	}()
+	}); err != nil {
+		logFields := append([]slog.Field{slog.F("cleanup", logMessage)}, fields...)
+		logFields = append(logFields, slog.Error(err))
+		p.logger.Error(context.WithoutCancel(ctx), "failed to schedule chat debug cleanup", logFields...)
+	}
 }
 
 func (p *Server) newDebugAwareModel(

coderd/x/chatd/quickgen.go

@@ -169,7 +169,7 @@ func (p *Server) GenerateChatTitleAsync(ctx context.Context, chat database.Chat)
 	// Detach from the request lifetime so title generation can finish
 	// even after the create response is written.
 	titleCtx := context.WithoutCancel(ctx)
-	p.inflight.Go(func() {
+	if err := p.goInflight(func() {
 		modelOpts := modelBuildOptionsFromMessages(messages)
 		titleCtx = withActiveTurnAPIKeyID(titleCtx, modelOpts)
 		model, modelConfig, keys, route, _, _, _, err := p.resolveChatModel(titleCtx, chat, modelOpts)
@@ -193,7 +193,13 @@ func (p *Server) GenerateChatTitleAsync(ctx context.Context, chat database.Chat)
 			logger,
 			p.existingDebugService(),
 		)
-	})
+	}); err != nil {
+		logger.Error(titleCtx, "failed to schedule automatic chat title generation",
+			slog.F("chat_id", chat.ID),
+			slog.F("owner_id", chat.OwnerID),
+			slog.Error(err),
+		)
+	}
 }
 
 // maybeGenerateChatTitle generates an AI title for the chat when

coderd/x/chatd/recording_internal_test.go

@@ -188,7 +188,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) {
 
 	// Wait for background processing triggered by CreateChat to
 	// settle before setting up the mock agent connection.
-	server.drainInflight()
+	WaitUntilIdleForTest(server)
 
 	// Now wire up the mock agent connection.
 	server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
@@ -274,7 +274,7 @@ func TestWaitAgentComputerUseRecordingWithThumbnail(t *testing.T) {
 		"parent-recording-thumb", "computer-use-child-thumb",
 	)
 
-	server.drainInflight()
+	WaitUntilIdleForTest(server)
 
 	server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
 		require.Equal(t, agent.ID, agentID)
@@ -368,7 +368,7 @@ func TestWaitAgentNonComputerUseNoRecording(t *testing.T) {
 
 	// Wait for background processing triggered by CreateChat to
 	// settle before setting up the mock agent connection.
-	server.drainInflight()
+	WaitUntilIdleForTest(server)
 
 	// Wire up the mock agent connection. The mock has zero
 	// expectations — gomock will fail if StartDesktopRecording

coderd/x/chatd/subagent_internal_test.go

@@ -3315,7 +3315,7 @@ func TestWaitAgentDoesNotRelayRegularSubagentAttachments(t *testing.T) {
 	server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
 
 	parent, child := createParentChildChats(ctx, t, server, user, org, model)
-	server.drainInflight()
+	WaitUntilIdleForTest(server)
 
 	insertedFile := insertLinkedChatFile(
 		ctx,
@@ -3501,6 +3501,9 @@ func TestAwaitSubagentCompletion(t *testing.T) {
 
 		parent, child := createParentChildChats(ctx, t, server, user, org, model)
 
+		// signalWake from CreateChat may trigger immediate processing.
+		// Wait for it to settle, then reset chats to the state we need.
+		WaitUntilIdleForTest(server)
 		setChatStatus(ctx, t, db, parent.ID, database.ChatStatusRunning, "")
 		setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
 
@@ -3578,7 +3581,8 @@ func TestAwaitSubagentCompletion(t *testing.T) {
 		parent, child := createParentChildChats(ctx, t, server, user, org, model)
 
 		// This case should return immediately, so use the shared
-		// real-clock passive server instead of a mock clock.
+		// real-clock server instead of a mock clock.
+		WaitUntilIdleForTest(server)
 		setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
 
 		gotChat, report, err := server.awaitSubagentCompletion(

coderd/x/chatd/testhooks.go

@@ -10,11 +10,10 @@ import (
 // database state only after asynchronous chat processing has completed.
 // Close waits for the same tracked work, but also stops the server.
 func WaitUntilIdleForTest(server *Server) {
-	server.drainInflight()
-	if server.chatWorker == nil {
-		return
+	if server.chatWorker != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		_ = server.chatWorker.WaitIdle(ctx)
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
-	_ = server.chatWorker.WaitIdle(ctx)
+	server.drainInflight()
 }