fix(coderd/x/chatd): inflight wait group data race (#24007)

Addresses https://github.com/coder/internal/issues/1450

coderd/x/chatd/chatd.go

@@ -104,9 +104,10 @@ var errChatHasNoWorkspaceAgent = xerrors.New("workspace has no running agent: th
 
 // Server handles background processing of pending chats.
 type Server struct {
-	cancel   context.CancelFunc
-	closed   chan struct{}
-	inflight sync.WaitGroup
+	cancel     context.CancelFunc
+	closed     chan struct{}
+	inflight   sync.WaitGroup
+	inflightMu sync.Mutex
 
 	db       database.Store
 	workerID uuid.UUID
@@ -2513,6 +2514,7 @@ func (p *Server) processOnce(ctx context.Context) {
 		return
 	}
 
+	p.inflightMu.Lock()
 	for _, chat := range chats {
 		p.inflight.Add(1)
 		go func() {
@@ -2520,6 +2522,7 @@ func (p *Server) processOnce(ctx context.Context) {
 			p.processChat(ctx, chat)
 		}()
 	}
+	p.inflightMu.Unlock()
 }
 
 func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEvent) {
@@ -5454,10 +5457,23 @@ func (p *Server) Close() error {
 	}
 	p.cancel()
 	<-p.closed
-	p.inflight.Wait()
+	p.drainInflight()
 	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.
+//
+// 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.
+func (p *Server) drainInflight() {
+	p.inflightMu.Lock()
+	p.inflight.Wait()
+	p.inflightMu.Unlock()
+}
+
 // refreshExpiredMCPTokens checks each MCP OAuth2 token and refreshes
 // any that are expired (or about to expire). Tokens without a
 // refresh_token or that fail to refresh are returned unchanged so the

coderd/x/chatd/export_test.go

@@ -5,5 +5,5 @@ package chatd
 // 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.inflight.Wait()
+	server.drainInflight()
 }

coderd/x/chatd/recording_internal_test.go

@@ -143,7 +143,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) {
 
 	// Wait for background processing triggered by CreateChat to
 	// settle before setting up the mock agent connection.
-	server.inflight.Wait()
+	server.drainInflight()
 
 	// Now wire up the mock agent connection.
 	server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
@@ -222,7 +222,7 @@ func TestWaitAgentNonComputerUseNoRecording(t *testing.T) {
 
 	// Wait for background processing triggered by CreateChat to
 	// settle before setting up the mock agent connection.
-	server.inflight.Wait()
+	server.drainInflight()
 
 	// Wire up the mock agent connection. The mock has zero
 	// expectations — gomock will fail if StartDesktopRecording

coderd/x/chatd/subagent_internal_test.go

@@ -1008,7 +1008,7 @@ func TestAwaitSubagentCompletion(t *testing.T) {
 
 		// signalWake from CreateChat may trigger immediate processing.
 		// Wait for it to settle, then reset chats to the state we need.
-		server.inflight.Wait()
+		server.drainInflight()
 		setChatStatus(ctx, t, db, parent.ID, database.ChatStatusRunning, "")
 		setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
 
@@ -1090,7 +1090,7 @@ func TestAwaitSubagentCompletion(t *testing.T) {
 		// Wait for it to settle, then set the terminal state we need.
 		// This case should return immediately, so use the shared
 		// real-clock server instead of a mock clock.
-		server.inflight.Wait()
+		server.drainInflight()
 		setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
 
 		gotChat, report, err := server.awaitSubagentCompletion(