fix(coderd/x/chatd): bind goInflight contexts to server lifetime (#26811)

fix(coderd/x/chatd): bind goInflight contexts to server lifetime
- Add Server.inflightContext: WithoutCancel(reqCtx) bound to p.ctx via
  context.AfterFunc, so Close cancels in-flight work instead of blocking
  on the caller's timeout while a provider is unreachable.
- Apply at GenerateChatTitleAsync, finalizeSuccessfulTurnStatusLabelWithAfterFunc,
  setLastTurnSummaryAsync, clearLastTurnSummaryAsync, and scheduleDebugCleanup.
- Honor cleanupCtx in the debug retry-delay timer so cancellation lands
  promptly between attempts.

> 🤖

coderd/x/chatd/chatd.go

@@ -4713,8 +4713,9 @@ func (p *Server) finalizeSuccessfulTurnStatusLabelWithAfterFunc(
 	logger slog.Logger,
 	afterFinalize func(context.Context, string),
 ) {
+	finalizeCtx, stopFinalizeCtx := p.inflightContext(ctx)
 	if err := p.goInflight(func() {
-		finalizeCtx := context.WithoutCancel(ctx)
+		defer stopFinalizeCtx()
 		statusLabel := p.generateFinalTurnStatusLabel(finalizeCtx, chat, status, runResult, logger)
 		logger.Debug(finalizeCtx, "generated chat turn status label",
 			slog.F("chat_id", chat.ID),
@@ -4726,6 +4727,7 @@ func (p *Server) finalizeSuccessfulTurnStatusLabelWithAfterFunc(
 
 		afterFinalize(finalizeCtx, statusLabel)
 	}); err != nil {
+		stopFinalizeCtx()
 		logger.Error(context.WithoutCancel(ctx), "failed to schedule chat turn status finalization",
 			slog.F("chat_id", chat.ID),
 			slog.F("status", status),
@@ -4813,9 +4815,12 @@ func (p *Server) setLastTurnSummaryAsync(
 	if chat.LastTurnSummary.Valid && strings.TrimSpace(chat.LastTurnSummary.String) == summary {
 		return
 	}
+	updateCtx, stopUpdateCtx := p.inflightContext(ctx)
 	if err := p.goInflight(func() {
-		p.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.HistoryVersion, summary, logger)
+		defer stopUpdateCtx()
+		p.updateLastTurnSummary(updateCtx, chat, chat.HistoryVersion, summary, logger)
 	}); err != nil {
+		stopUpdateCtx()
 		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),
@@ -4830,9 +4835,12 @@ func (p *Server) clearLastTurnSummaryAsync(
 	chat database.Chat,
 	logger slog.Logger,
 ) {
+	clearCtx, stopClearCtx := p.inflightContext(ctx)
 	if err := p.goInflight(func() {
-		p.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.HistoryVersion, "", logger)
+		defer stopClearCtx()
+		p.updateLastTurnSummary(clearCtx, chat, chat.HistoryVersion, "", logger)
 	}); err != nil {
+		stopClearCtx()
 		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),
@@ -4942,6 +4950,22 @@ func (p *Server) Close() error {
 	return nil
 }
 
+// inflightContext returns a context for an in-flight goroutine launched
+// via goInflight. It is detached from reqCtx's cancellation so the work
+// can outlive the originating request, while preserving its values for
+// auth, routing, and tracing. The context is bound to the server
+// lifetime via p.ctx so Close cancels it promptly. The returned stop
+// must be called once the work completes to release the shutdown hook.
+// The caller is responsible for providing their own timeout.
+func (p *Server) inflightContext(reqCtx context.Context) (context.Context, func()) {
+	ctx, cancel := context.WithCancel(context.WithoutCancel(reqCtx))
+	stop := context.AfterFunc(p.ctx, cancel)
+	return ctx, func() {
+		stop()
+		cancel()
+	}
+}
+
 func (p *Server) goInflight(f func()) error {
 	if p.inflightClosed.Load() {
 		return errInflightClosed

coderd/x/chatd/chatd_debug.go

@@ -76,12 +76,19 @@ func (p *Server) scheduleDebugCleanup(
 		return
 	}
 
+	cleanupCtx, stopCleanupCtx := p.inflightContext(ctx)
 	if err := p.goInflight(func() {
-		cleanupCtx := context.WithoutCancel(ctx)
+		defer stopCleanupCtx()
 		for attempt := 0; attempt < debugCleanupAttempts; attempt++ {
 			if attempt > 0 {
 				timer := p.clock.NewTimer(debugCleanupRetryDelay, "chatd", "debug_cleanup")
-				<-timer.C
+				defer timer.Stop()
+				select {
+				case <-timer.C:
+				case <-cleanupCtx.Done():
+					timer.Stop()
+					return
+				}
 			}
 
 			passCtx, cancel := context.WithTimeout(cleanupCtx, debugCleanupTimeout)
@@ -99,6 +106,7 @@ func (p *Server) scheduleDebugCleanup(
 			p.logger.Warn(cleanupCtx, logMessage, logFields...)
 		}
 	}); err != nil {
+		stopCleanupCtx()
 		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...)

coderd/x/chatd/chatd_internal_test.go

@@ -3485,3 +3485,40 @@ func TestGetWorkspaceConnBumpsWorkspaceUsage(t *testing.T) {
 	require.True(t, updatedBuild.Deadline.After(now.Add(time.Hour-2*time.Minute)))
 	require.True(t, updatedBuild.Deadline.Before(now.Add(time.Hour+2*time.Minute)))
 }
+
+func TestServer_inflightContext(t *testing.T) {
+	t.Parallel()
+
+	serverCtx, serverCancel := context.WithCancel(context.Background())
+	t.Cleanup(serverCancel)
+	server := &Server{ctx: serverCtx}
+
+	type ctxKey string
+	const key ctxKey = "inflight-test"
+	reqCtx, reqCancel := context.WithCancel(context.WithValue(context.Background(), key, "value"))
+	t.Cleanup(reqCancel)
+
+	inflightCtx, stop := server.inflightContext(reqCtx)
+	t.Cleanup(stop)
+
+	// Auth and routing values must carry over from the request.
+	require.Equal(t, "value", inflightCtx.Value(key))
+
+	// Request cancellation must not cancel in-flight work: it has to outlive
+	// the originating request.
+	reqCancel()
+	select {
+	case <-inflightCtx.Done():
+		t.Fatal("inflight context canceled by request cancellation")
+	case <-time.After(testutil.IntervalFast):
+	}
+
+	// Server shutdown must cancel in-flight work so Close does not block
+	// on long-running callees while a provider is unreachable.
+	serverCancel()
+	select {
+	case <-inflightCtx.Done():
+	case <-time.After(testutil.WaitShort):
+		t.Fatal("inflight context not canceled on server shutdown")
+	}
+}

coderd/x/chatd/quickgen.go

@@ -141,9 +141,11 @@ type generatedTurnStatusLabel struct {
 // chat-creation endpoint right after the chat and its initial user
 // message are persisted.
 //
-// The work runs in a tracked goroutine with a detached context so it
-// neither blocks the HTTP response nor is canceled when the request
-// completes. It resolves the chat's model and provider keys, then
+// The work runs in a tracked goroutine with a context detached from the
+// request but bound to the server: it neither blocks the HTTP response
+// nor is canceled when the request completes, and Close cancels it
+// instead of blocking on the title timeout while a provider is
+// unreachable. It resolves the chat's model and provider keys, then
 // delegates to maybeGenerateChatTitle, which only acts on the first user
 // turn (see titleInput) and is otherwise a no-op. Errors are logged and
 // swallowed.
@@ -166,21 +168,21 @@ func (p *Server) GenerateChatTitleAsync(ctx context.Context, chat database.Chat)
 	if _, ok := titleInput(chat, messages); !ok {
 		return
 	}
-	// Detach from the request lifetime so title generation can finish
-	// even after the create response is written.
-	titleCtx := context.WithoutCancel(ctx)
+	// Detach from request; bind to server so Close cancels it.
+	titleCtx, stopTitleCtx := p.inflightContext(ctx)
 	if err := p.goInflight(func() {
+		defer stopTitleCtx()
 		modelOpts := modelBuildOptionsFromMessages(messages)
-		titleCtx = withActiveTurnAPIKeyID(titleCtx, modelOpts)
-		model, modelConfig, keys, route, _, _, _, err := p.resolveChatModel(titleCtx, chat, modelOpts)
+		turnCtx := withActiveTurnAPIKeyID(titleCtx, modelOpts)
+		model, modelConfig, keys, route, _, _, _, err := p.resolveChatModel(turnCtx, chat, modelOpts)
 		if err != nil {
-			logger.Debug(titleCtx, "failed to resolve model for automatic title generation",
+			logger.Debug(turnCtx, "failed to resolve model for automatic title generation",
 				slog.Error(err),
 			)
 			return
 		}
 		p.maybeGenerateChatTitle(
-			titleCtx,
+			turnCtx,
 			chat,
 			messages,
 			modelConfig.Provider,
@@ -194,7 +196,8 @@ func (p *Server) GenerateChatTitleAsync(ctx context.Context, chat database.Chat)
 			p.existingDebugService(),
 		)
 	}); err != nil {
-		logger.Error(titleCtx, "failed to schedule automatic chat title generation",
+		stopTitleCtx()
+		logger.Error(context.WithoutCancel(ctx), "failed to schedule automatic chat title generation",
 			slog.F("chat_id", chat.ID),
 			slog.F("owner_id", chat.OwnerID),
 			slog.Error(err),

coderd/x/chatd/turn_summary_internal_test.go

@@ -178,7 +178,7 @@ func TestPendingChatPersistsSummaryButSkipsWebPush(t *testing.T) {
 
 	dispatcher := &recordingWebpushDispatcher{}
 	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
-	server := &Server{db: db, pubsub: ps, webpushDispatcher: dispatcher}
+	server := &Server{ctx: t.Context(), db: db, pubsub: ps, webpushDispatcher: dispatcher}
 	server.maybeFinalizeTurnStatusLabelAndPush(
 		context.WithoutCancel(ctx),
 		chat,
@@ -260,6 +260,7 @@ func TestSuccessfulChildChatOutcomeSkipsSummaryAndWebPush(t *testing.T) {
 
 	dispatcher := &recordingWebpushDispatcher{}
 	server := &Server{
+		ctx:               t.Context(),
 		db:                db,
 		pubsub:            ps,
 		logger:            slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),