diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 33fc87f33b..0473f554f1 100644 --- a/coderd/x/chatd/chatd.go +++ b/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. diff --git a/coderd/x/chatd/chatd_debug.go b/coderd/x/chatd/chatd_debug.go index bbb66cb82a..70f8c88e9d 100644 --- a/coderd/x/chatd/chatd_debug.go +++ b/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( diff --git a/coderd/x/chatd/quickgen.go b/coderd/x/chatd/quickgen.go index ec9dc01356..fb4ddf3f21 100644 --- a/coderd/x/chatd/quickgen.go +++ b/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 diff --git a/coderd/x/chatd/recording_internal_test.go b/coderd/x/chatd/recording_internal_test.go index f4b576a1ea..23d4f2ff2a 100644 --- a/coderd/x/chatd/recording_internal_test.go +++ b/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 diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 4f324380f8..bb02bf66cd 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/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( diff --git a/coderd/x/chatd/testhooks.go b/coderd/x/chatd/testhooks.go index c1356ee3d0..4598bbeec6 100644 --- a/coderd/x/chatd/testhooks.go +++ b/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() }