From cb6a75717cfb8da217dd1d9f56a6d6315ee347c2 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 30 Jun 2026 10:51:59 +0100 Subject: [PATCH] fix(coderd/x/chatd): bind goInflight contexts to server lifetime (#26811) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 30 ++++++++++++++-- coderd/x/chatd/chatd_debug.go | 12 +++++-- coderd/x/chatd/chatd_internal_test.go | 37 ++++++++++++++++++++ coderd/x/chatd/quickgen.go | 25 +++++++------ coderd/x/chatd/turn_summary_internal_test.go | 3 +- 5 files changed, 90 insertions(+), 17 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index d2db52f031..e983330e89 100644 --- a/coderd/x/chatd/chatd.go +++ b/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 diff --git a/coderd/x/chatd/chatd_debug.go b/coderd/x/chatd/chatd_debug.go index 70f8c88e9d..ccae7c4a27 100644 --- a/coderd/x/chatd/chatd_debug.go +++ b/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...) diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 3c15609a3c..6f44e2ef4c 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/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") + } +} diff --git a/coderd/x/chatd/quickgen.go b/coderd/x/chatd/quickgen.go index 88642d67de..96fad11f3f 100644 --- a/coderd/x/chatd/quickgen.go +++ b/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), diff --git a/coderd/x/chatd/turn_summary_internal_test.go b/coderd/x/chatd/turn_summary_internal_test.go index 9d0abd7004..9547f9a759 100644 --- a/coderd/x/chatd/turn_summary_internal_test.go +++ b/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}),