diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index 64d088ee2b..3e33eacd33 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -43,6 +43,11 @@ const ( chatHeartbeatInterval = 30 * time.Second maxChatSteps = 1200 + // staleRecoveryIntervalDivisor determines how often the stale + // recovery loop runs relative to the stale threshold. A value + // of 5 means recovery runs at 1/5 of the stale-after duration. + staleRecoveryIntervalDivisor = 5 + defaultSubagentInstruction = "You are running as a delegated sub-agent chat. Complete the delegated task and provide clear, concise assistant responses for the parent agent." ) @@ -885,18 +890,25 @@ func New(cfg Config) *Server { func (p *Server) start(ctx context.Context) { defer close(p.closed) - // First, recover any stale chats from crashed workers. + // Recover stale chats on startup and periodically thereafter + // to handle chats orphaned by crashed or redeployed workers. p.recoverStaleChats(ctx) - ticker := time.NewTicker(p.pendingChatAcquireInterval) - defer ticker.Stop() + acquireTicker := time.NewTicker(p.pendingChatAcquireInterval) + defer acquireTicker.Stop() + + staleRecoveryInterval := p.inFlightChatStaleAfter / staleRecoveryIntervalDivisor + staleTicker := time.NewTicker(staleRecoveryInterval) + defer staleTicker.Stop() for { select { case <-ctx.Done(): return - case <-ticker.C: + case <-acquireTicker.C: p.processOnce(ctx) + case <-staleTicker.C: + p.recoverStaleChats(ctx) } } } @@ -1598,9 +1610,14 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { shouldPublishQueueUpdate := false defer func() { + // Use a context that is not canceled by Close() so we can + // reliably update the chat status in the database during + // graceful shutdown. + cleanupCtx := context.WithoutCancel(ctx) + // Handle panics gracefully. if r := recover(); r != nil { - logger.Error(ctx, "panic during chat processing", slog.F("panic", r)) + logger.Error(cleanupCtx, "panic during chat processing", slog.F("panic", r)) p.publishError(chat.ID, panicFailureReason(r)) status = database.ChatStatusError } @@ -1613,7 +1630,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { err := p.db.InTx(func(tx database.Store) error { // Re-read the chat status under lock — another caller // (e.g. promote) may have already set it to pending. - latestChat, lockErr := tx.GetChatByIDForUpdate(ctx, chat.ID) + latestChat, lockErr := tx.GetChatByIDForUpdate(cleanupCtx, chat.ID) if lockErr != nil { return xerrors.Errorf("lock chat for release: %w", lockErr) } @@ -1625,9 +1642,9 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { status = database.ChatStatusPending } else if status == database.ChatStatusWaiting { // Try to auto-promote the next queued message. - nextQueued, popErr := tx.PopNextQueuedMessage(ctx, chat.ID) + nextQueued, popErr := tx.PopNextQueuedMessage(cleanupCtx, chat.ID) if popErr == nil { - msg, insertErr := tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ + msg, insertErr := tx.InsertChatMessage(cleanupCtx, database.InsertChatMessageParams{ ChatID: chat.ID, ModelConfigID: uuid.NullUUID{UUID: latestChat.LastModelConfigID, Valid: true}, Role: "user", @@ -1646,7 +1663,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { Compressed: sql.NullBool{}, }) if insertErr != nil { - logger.Error(ctx, "failed to promote queued message", + logger.Error(cleanupCtx, "failed to promote queued message", slog.F("queued_message_id", nextQueued.ID), slog.Error(insertErr)) } else { status = database.ChatStatusPending @@ -1657,7 +1674,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { Message: &sdkMsg, }) - remaining, qErr := tx.GetChatQueuedMessages(ctx, chat.ID) + remaining, qErr := tx.GetChatQueuedMessages(cleanupCtx, chat.ID) if qErr == nil { remainingQueuedMessages = remaining shouldPublishQueueUpdate = true @@ -1666,7 +1683,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { } } - _, updateErr := tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + _, updateErr := tx.UpdateChatStatus(cleanupCtx, database.UpdateChatStatusParams{ ID: chat.ID, Status: status, WorkerID: uuid.NullUUID{}, @@ -1676,7 +1693,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { return updateErr }, nil) if err != nil { - logger.Error(ctx, "failed to release chat", slog.Error(err)) + logger.Error(cleanupCtx, "failed to release chat", slog.Error(err)) } if err == nil && shouldPublishQueueUpdate { p.publishEvent(chat.ID, codersdk.ChatStreamEvent{ diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index 8eb17db184..2f691df79b 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -404,6 +404,179 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { require.True(t, errors.Is(err, chatd.ErrEditedMessageNotUser)) } +func TestRecoverStaleChatsPeriodically(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + + ctx := testutil.Context(t, testutil.WaitLong) + user, model := seedChatDependencies(ctx, t, db) + + // Use a very short stale threshold so the periodic recovery + // kicks in quickly during the test. + staleAfter := 500 * time.Millisecond + + // Create a chat and simulate a dead worker by setting the chat + // to running with a heartbeat in the past. + deadWorkerID := uuid.New() + chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OwnerID: user.ID, + Title: "stale-recovery-periodic", + LastModelConfigID: model.ID, + }) + require.NoError(t, err) + + _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + ID: chat.ID, + Status: database.ChatStatusRunning, + WorkerID: uuid.NullUUID{UUID: deadWorkerID, Valid: true}, + StartedAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, + HeartbeatAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, + }) + require.NoError(t, err) + + // Start a new replica. Its startup recovery will reset the + // chat (since the heartbeat is old), but the key point is that + // the periodic loop also recovers newly-stale chats. + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + server := chatd.New(chatd.Config{ + Logger: logger, + Database: db, + ReplicaID: uuid.New(), + Pubsub: ps, + PendingChatAcquireInterval: testutil.WaitSuperLong, + InFlightChatStaleAfter: staleAfter, + }) + t.Cleanup(func() { + require.NoError(t, server.Close()) + }) + + // The startup recovery should have already reset our stale + // chat. + require.Eventually(t, func() bool { + fromDB, err := db.GetChatByID(ctx, chat.ID) + if err != nil { + return false + } + return fromDB.Status == database.ChatStatusPending + }, testutil.WaitMedium, testutil.IntervalFast) + + // Now simulate a second stale chat appearing AFTER startup. + // This tests the periodic recovery, not just the startup one. + deadWorkerID2 := uuid.New() + chat2, err := db.InsertChat(ctx, database.InsertChatParams{ + OwnerID: user.ID, + Title: "stale-recovery-periodic-2", + LastModelConfigID: model.ID, + }) + require.NoError(t, err) + + _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + ID: chat2.ID, + Status: database.ChatStatusRunning, + WorkerID: uuid.NullUUID{UUID: deadWorkerID2, Valid: true}, + StartedAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, + HeartbeatAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, + }) + require.NoError(t, err) + + // The periodic stale recovery loop (running at staleAfter/5 = + // 100ms intervals) should pick this up without a restart. + require.Eventually(t, func() bool { + fromDB, err := db.GetChatByID(ctx, chat2.ID) + if err != nil { + return false + } + return fromDB.Status == database.ChatStatusPending + }, testutil.WaitMedium, testutil.IntervalFast) +} + +func TestNewReplicaRecoversStaleChatFromDeadReplica(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + + ctx := testutil.Context(t, testutil.WaitLong) + user, model := seedChatDependencies(ctx, t, db) + + // Simulate a chat left running by a dead replica with a stale + // heartbeat (well beyond the stale threshold). + deadReplicaID := uuid.New() + chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OwnerID: user.ID, + Title: "orphaned-chat", + LastModelConfigID: model.ID, + }) + require.NoError(t, err) + + // Set the heartbeat far in the past so it's definitely stale. + _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + ID: chat.ID, + Status: database.ChatStatusRunning, + WorkerID: uuid.NullUUID{UUID: deadReplicaID, Valid: true}, + StartedAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, + HeartbeatAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, + }) + require.NoError(t, err) + + // Start a new replica — it should recover the stale chat on + // startup. + newReplica := newTestServer(t, db, ps, uuid.New()) + _ = newReplica + + require.Eventually(t, func() bool { + fromDB, err := db.GetChatByID(ctx, chat.ID) + if err != nil { + return false + } + return fromDB.Status == database.ChatStatusPending && + !fromDB.WorkerID.Valid + }, testutil.WaitMedium, testutil.IntervalFast) +} + +func TestWaitingChatsAreNotRecoveredAsStale(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + + ctx := testutil.Context(t, testutil.WaitLong) + user, model := seedChatDependencies(ctx, t, db) + + // Create a chat in waiting status — this should NOT be touched + // by stale recovery. + chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OwnerID: user.ID, + Title: "waiting-chat", + LastModelConfigID: model.ID, + }) + require.NoError(t, err) + + // Start a replica with a short stale threshold. + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + server := chatd.New(chatd.Config{ + Logger: logger, + Database: db, + ReplicaID: uuid.New(), + Pubsub: ps, + PendingChatAcquireInterval: testutil.WaitSuperLong, + InFlightChatStaleAfter: 500 * time.Millisecond, + }) + t.Cleanup(func() { + require.NoError(t, server.Close()) + }) + + // Wait long enough for multiple periodic recovery cycles to + // run (staleAfter/5 = 100ms intervals). + require.Never(t, func() bool { + fromDB, err := db.GetChatByID(ctx, chat.ID) + if err != nil { + return false + } + return fromDB.Status != database.ChatStatusWaiting + }, time.Second, testutil.IntervalFast, + "waiting chat should not be modified by stale recovery") +} + func newTestServer( t *testing.T, db database.Store,