diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 3f739f33dd..63e3c423d6 100644 --- a/coderd/x/chatd/chatd.go +++ b/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 diff --git a/coderd/x/chatd/export_test.go b/coderd/x/chatd/export_test.go index 33fff4d6be..7c7177b88b 100644 --- a/coderd/x/chatd/export_test.go +++ b/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() } diff --git a/coderd/x/chatd/recording_internal_test.go b/coderd/x/chatd/recording_internal_test.go index 1b0094a236..b511d18168 100644 --- a/coderd/x/chatd/recording_internal_test.go +++ b/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 diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 83248800b7..b61276d478 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/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(