perf(coderd): batch chat heartbeat queries into single UPDATE per interval (#24037)

## Summary

Replaces N per-chat heartbeat goroutines with a single centralized
heartbeat loop that issues one `UPDATE` per 30s interval for all running
chats on a worker.

## Problem

Each running chat spawned a dedicated goroutine that issued an
individual `UPDATE chats SET heartbeat_at = NOW() WHERE id = $1 AND
worker_id = $2 AND status = 'running'` query every 30 seconds. At 10,000
concurrent chats this produces **~333 DB queries/second** just for
heartbeats, plus ~333 `ActivityBumpWorkspace` CTE queries/second from
`trackWorkspaceUsage`.

## Solution

New `UpdateChatHeartbeats` (plural) SQL query replaces the old singular
`UpdateChatHeartbeat`:

```sql
UPDATE chats
SET    heartbeat_at = @now::timestamptz
WHERE  worker_id = @worker_id::uuid
  AND  status = 'running'::chat_status
RETURNING id;
```

A single `heartbeatLoop` goroutine on the `Server`:
1. Ticks every `chatHeartbeatInterval` (30s)
2. Issues one batch UPDATE for all registered chats
3. Detects stolen/completed chats via set-difference (equivalent of old
`rows == 0`)
4. Calls `trackWorkspaceUsage` for surviving chats

`processChat` registers an entry in the heartbeat registry instead of
spawning a goroutine.

## Impact

| Metric | Before (10K chats) | After (10K chats) |
|---|---|---|
| Heartbeat queries/sec | ~333 | ~0.03 (1 per 30s per replica) |
| Heartbeat goroutines | 10,000 | 1 |
| Self-interrupt detection | Per-chat `rows==0` | Batch set-difference |

---

> 🤖 Generated by Coder Agents

<details><summary>Implementation notes</summary>

- Uses `@now` parameter instead of `NOW()` so tests with `quartz.Mock`
can control timestamps.
- `heartbeatEntry` stores `context.CancelCauseFunc` + workspace state
for the centralized loop.
- `recoverStaleChats` is unaffected — it reads `heartbeat_at` which is
still updated.
- The old singular `UpdateChatHeartbeat` is removed entirely.
- `dbauthz` wrapper uses system-level `rbac.ResourceChat` authorization
(same pattern as `AcquireChats`).

</details>

coderd/database/dbauthz/dbauthz.go

@@ -5782,15 +5782,15 @@ func (q *querier) UpdateChatByID(ctx context.Context, arg database.UpdateChatByI
 	return q.db.UpdateChatByID(ctx, arg)
 }
 
-func (q *querier) UpdateChatHeartbeat(ctx context.Context, arg database.UpdateChatHeartbeatParams) (int64, error) {
-	chat, err := q.db.GetChatByID(ctx, arg.ID)
-	if err != nil {
-		return 0, err
+func (q *querier) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) {
+	// The batch heartbeat is a system-level operation filtered by
+	// worker_id. Authorization is enforced by the AsChatd context
+	// at the call site rather than per-row, because checking each
+	// row individually would defeat the purpose of batching.
+	if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
+		return nil, err
 	}
-	if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
-		return 0, err
-	}
-	return q.db.UpdateChatHeartbeat(ctx, arg)
+	return q.db.UpdateChatHeartbeats(ctx, arg)
 }
 
 func (q *querier) UpdateChatLabelsByID(ctx context.Context, arg database.UpdateChatLabelsByIDParams) (database.Chat, error) {

coderd/database/dbauthz/dbauthz_test.go

@@ -842,15 +842,15 @@ func (s *MethodTestSuite) TestChats() {
 		dbm.EXPECT().UpdateChatStatusPreserveUpdatedAt(gomock.Any(), arg).Return(chat, nil).AnyTimes()
 		check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(chat)
 	}))
-	s.Run("UpdateChatHeartbeat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
-		chat := testutil.Fake(s.T(), faker, database.Chat{})
-		arg := database.UpdateChatHeartbeatParams{
-			ID:       chat.ID,
+	s.Run("UpdateChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		resultID := uuid.New()
+		arg := database.UpdateChatHeartbeatsParams{
+			IDs:      []uuid.UUID{resultID},
 			WorkerID: uuid.New(),
+			Now:      time.Now(),
 		}
-		dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
-		dbm.EXPECT().UpdateChatHeartbeat(gomock.Any(), arg).Return(int64(1), nil).AnyTimes()
-		check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(int64(1))
+		dbm.EXPECT().UpdateChatHeartbeats(gomock.Any(), arg).Return([]uuid.UUID{resultID}, nil).AnyTimes()
+		check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]uuid.UUID{resultID})
 	}))
 	s.Run("UpdateChatMessageByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		chat := testutil.Fake(s.T(), faker, database.Chat{})

coderd/database/dbmetrics/querymetrics.go

@@ -4136,11 +4136,11 @@ func (m queryMetricsStore) UpdateChatByID(ctx context.Context, arg database.Upda
 	return r0, r1
 }
 
-func (m queryMetricsStore) UpdateChatHeartbeat(ctx context.Context, arg database.UpdateChatHeartbeatParams) (int64, error) {
+func (m queryMetricsStore) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) {
 	start := time.Now()
-	r0, r1 := m.s.UpdateChatHeartbeat(ctx, arg)
-	m.queryLatencies.WithLabelValues("UpdateChatHeartbeat").Observe(time.Since(start).Seconds())
-	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatHeartbeat").Inc()
+	r0, r1 := m.s.UpdateChatHeartbeats(ctx, arg)
+	m.queryLatencies.WithLabelValues("UpdateChatHeartbeats").Observe(time.Since(start).Seconds())
+	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatHeartbeats").Inc()
 	return r0, r1
 }
 

coderd/database/dbmock/dbmock.go

@@ -7835,19 +7835,19 @@ func (mr *MockStoreMockRecorder) UpdateChatByID(ctx, arg any) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatByID", reflect.TypeOf((*MockStore)(nil).UpdateChatByID), ctx, arg)
 }
 
-// UpdateChatHeartbeat mocks base method.
-func (m *MockStore) UpdateChatHeartbeat(ctx context.Context, arg database.UpdateChatHeartbeatParams) (int64, error) {
+// UpdateChatHeartbeats mocks base method.
+func (m *MockStore) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "UpdateChatHeartbeat", ctx, arg)
-	ret0, _ := ret[0].(int64)
+	ret := m.ctrl.Call(m, "UpdateChatHeartbeats", ctx, arg)
+	ret0, _ := ret[0].([]uuid.UUID)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// UpdateChatHeartbeat indicates an expected call of UpdateChatHeartbeat.
-func (mr *MockStoreMockRecorder) UpdateChatHeartbeat(ctx, arg any) *gomock.Call {
+// UpdateChatHeartbeats indicates an expected call of UpdateChatHeartbeats.
+func (mr *MockStoreMockRecorder) UpdateChatHeartbeats(ctx, arg any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatHeartbeat", reflect.TypeOf((*MockStore)(nil).UpdateChatHeartbeat), ctx, arg)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatHeartbeats", reflect.TypeOf((*MockStore)(nil).UpdateChatHeartbeats), ctx, arg)
 }
 
 // UpdateChatLabelsByID mocks base method.

coderd/database/querier.go

@@ -870,9 +870,11 @@ type sqlcQuerier interface {
 	UpdateAPIKeyByID(ctx context.Context, arg UpdateAPIKeyByIDParams) error
 	UpdateChatBuildAgentBinding(ctx context.Context, arg UpdateChatBuildAgentBindingParams) (Chat, error)
 	UpdateChatByID(ctx context.Context, arg UpdateChatByIDParams) (Chat, error)
-	// Bumps the heartbeat timestamp for a running chat so that other
-	// replicas know the worker is still alive.
-	UpdateChatHeartbeat(ctx context.Context, arg UpdateChatHeartbeatParams) (int64, error)
+	// Bumps the heartbeat timestamp for the given set of chat IDs,
+	// provided they are still running and owned by the specified
+	// worker. Returns the IDs that were actually updated so the
+	// caller can detect stolen or completed chats via set-difference.
+	UpdateChatHeartbeats(ctx context.Context, arg UpdateChatHeartbeatsParams) ([]uuid.UUID, error)
 	UpdateChatLabelsByID(ctx context.Context, arg UpdateChatLabelsByIDParams) (Chat, error)
 	// Updates the cached injected context parts (AGENTS.md +
 	// skills) on the chat row. Called only when context changes

coderd/database/queries.sql.go

@@ -6615,30 +6615,49 @@ func (q *sqlQuerier) UpdateChatByID(ctx context.Context, arg UpdateChatByIDParam
 	return i, err
 }
 
-const updateChatHeartbeat = `-- name: UpdateChatHeartbeat :execrows
+const updateChatHeartbeats = `-- name: UpdateChatHeartbeats :many
 UPDATE
     chats
 SET
-    heartbeat_at = NOW()
+    heartbeat_at = $1::timestamptz
 WHERE
-    id = $1::uuid
-    AND worker_id = $2::uuid
+    id = ANY($2::uuid[])
+    AND worker_id = $3::uuid
     AND status = 'running'::chat_status
+RETURNING id
 `
 
-type UpdateChatHeartbeatParams struct {
-	ID       uuid.UUID `db:"id" json:"id"`
-	WorkerID uuid.UUID `db:"worker_id" json:"worker_id"`
+type UpdateChatHeartbeatsParams struct {
+	Now      time.Time   `db:"now" json:"now"`
+	IDs      []uuid.UUID `db:"ids" json:"ids"`
+	WorkerID uuid.UUID   `db:"worker_id" json:"worker_id"`
 }
 
-// Bumps the heartbeat timestamp for a running chat so that other
-// replicas know the worker is still alive.
-func (q *sqlQuerier) UpdateChatHeartbeat(ctx context.Context, arg UpdateChatHeartbeatParams) (int64, error) {
-	result, err := q.db.ExecContext(ctx, updateChatHeartbeat, arg.ID, arg.WorkerID)
+// Bumps the heartbeat timestamp for the given set of chat IDs,
+// provided they are still running and owned by the specified
+// worker. Returns the IDs that were actually updated so the
+// caller can detect stolen or completed chats via set-difference.
+func (q *sqlQuerier) UpdateChatHeartbeats(ctx context.Context, arg UpdateChatHeartbeatsParams) ([]uuid.UUID, error) {
+	rows, err := q.db.QueryContext(ctx, updateChatHeartbeats, arg.Now, pq.Array(arg.IDs), arg.WorkerID)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
-	return result.RowsAffected()
+	defer rows.Close()
+	var items []uuid.UUID
+	for rows.Next() {
+		var id uuid.UUID
+		if err := rows.Scan(&id); err != nil {
+			return nil, err
+		}
+		items = append(items, id)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
 }
 
 const updateChatLabelsByID = `-- name: UpdateChatLabelsByID :one

coderd/database/queries/chats.sql

@@ -674,17 +674,20 @@ WHERE
     status = 'running'::chat_status
     AND heartbeat_at < @stale_threshold::timestamptz;
 
--- name: UpdateChatHeartbeat :execrows
--- Bumps the heartbeat timestamp for a running chat so that other
--- replicas know the worker is still alive.
+-- name: UpdateChatHeartbeats :many
+-- Bumps the heartbeat timestamp for the given set of chat IDs,
+-- provided they are still running and owned by the specified
+-- worker. Returns the IDs that were actually updated so the
+-- caller can detect stolen or completed chats via set-difference.
 UPDATE
     chats
 SET
-    heartbeat_at = NOW()
+    heartbeat_at = @now::timestamptz
 WHERE
-    id = @id::uuid
+    id = ANY(@ids::uuid[])
     AND worker_id = @worker_id::uuid
-    AND status = 'running'::chat_status;
+    AND status = 'running'::chat_status
+RETURNING id;
 
 -- name: GetChatDiffStatusByChatID :one
 SELECT

coderd/x/chatd/chatd.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"maps"
 	"net/http"
 	"slices"
 	"strconv"
@@ -151,6 +152,12 @@ type Server struct {
 	inFlightChatStaleAfter     time.Duration
 	chatHeartbeatInterval      time.Duration
 
+	// heartbeatMu guards heartbeatRegistry.
+	heartbeatMu sync.Mutex
+	// heartbeatRegistry maps chat IDs to their cancel functions
+	// and workspace state for the centralized heartbeat loop.
+	heartbeatRegistry map[uuid.UUID]*heartbeatEntry
+
 	// wakeCh is signaled by SendMessage, EditMessage, CreateChat,
 	// and PromoteQueued so the run loop calls processOnce
 	// immediately instead of waiting for the next ticker.
@@ -706,6 +713,17 @@ type chatStreamState struct {
 	bufferRetainedAt time.Time
 }
 
+// heartbeatEntry tracks a single chat's cancel function and workspace
+// state for the centralized heartbeat loop. Instead of spawning a
+// per-chat goroutine, processChat registers an entry here and the
+// single heartbeatLoop goroutine handles all chats.
+type heartbeatEntry struct {
+	cancelWithCause context.CancelCauseFunc
+	chatID          uuid.UUID
+	workspaceID     uuid.NullUUID
+	logger          slog.Logger
+}
+
 // resetDropCounters zeroes the rate-limiting state for both buffer
 // and subscriber drop warnings. The caller must hold s.mu.
 func (s *chatStreamState) resetDropCounters() {
@@ -2420,8 +2438,8 @@ func New(cfg Config) *Server {
 		clock:                          clk,
 		recordingSem:                   make(chan struct{}, maxConcurrentRecordingUploads),
 		wakeCh:                         make(chan struct{}, 1),
+		heartbeatRegistry:              make(map[uuid.UUID]*heartbeatEntry),
 	}
-
 	//nolint:gocritic // The chat processor uses a scoped chatd context.
 	ctx = dbauthz.AsChatd(ctx)
 
@@ -2461,6 +2479,9 @@ func (p *Server) start(ctx context.Context) {
 	// to handle chats orphaned by crashed or redeployed workers.
 	p.recoverStaleChats(ctx)
 
+	// Single heartbeat loop for all chats on this replica.
+	go p.heartbeatLoop(ctx)
+
 	acquireTicker := p.clock.NewTicker(
 		p.pendingChatAcquireInterval,
 		"chatd",
@@ -2730,6 +2751,97 @@ func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) {
 	p.workspaceMCPToolsCache.Delete(chatID)
 }
 
+// registerHeartbeat enrolls a chat in the centralized batch
+// heartbeat loop. Must be called after chatCtx is created.
+func (p *Server) registerHeartbeat(entry *heartbeatEntry) {
+	p.heartbeatMu.Lock()
+	defer p.heartbeatMu.Unlock()
+	if _, exists := p.heartbeatRegistry[entry.chatID]; exists {
+		p.logger.Warn(context.Background(),
+			"duplicate heartbeat registration, skipping",
+			slog.F("chat_id", entry.chatID))
+		return
+	}
+	p.heartbeatRegistry[entry.chatID] = entry
+}
+
+// unregisterHeartbeat removes a chat from the centralized
+// heartbeat loop when chat processing finishes.
+func (p *Server) unregisterHeartbeat(chatID uuid.UUID) {
+	p.heartbeatMu.Lock()
+	defer p.heartbeatMu.Unlock()
+	delete(p.heartbeatRegistry, chatID)
+}
+
+// heartbeatLoop runs in a single goroutine, issuing one batch
+// heartbeat query per interval for all registered chats.
+func (p *Server) heartbeatLoop(ctx context.Context) {
+	ticker := p.clock.NewTicker(p.chatHeartbeatInterval, "chatd", "batch-heartbeat")
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			p.heartbeatTick(ctx)
+		}
+	}
+}
+
+// heartbeatTick issues a single batch UPDATE for all running chats
+// owned by this worker. Chats missing from the result set are
+// interrupted (stolen by another replica or already completed).
+func (p *Server) heartbeatTick(ctx context.Context) {
+	// Snapshot the registry under the lock.
+	p.heartbeatMu.Lock()
+	snapshot := maps.Clone(p.heartbeatRegistry)
+	p.heartbeatMu.Unlock()
+
+	if len(snapshot) == 0 {
+		return
+	}
+
+	// Collect the IDs we believe we own.
+	ids := slices.Collect(maps.Keys(snapshot))
+
+	//nolint:gocritic // AsChatd provides narrowly-scoped daemon
+	// access for batch-updating heartbeats.
+	chatdCtx := dbauthz.AsChatd(ctx)
+	updatedIDs, err := p.db.UpdateChatHeartbeats(chatdCtx, database.UpdateChatHeartbeatsParams{
+		IDs:      ids,
+		WorkerID: p.workerID,
+		Now:      p.clock.Now(),
+	})
+	if err != nil {
+		p.logger.Error(ctx, "batch heartbeat failed", slog.Error(err))
+		return
+	}
+
+	// Build a set of IDs that were successfully updated.
+	updated := make(map[uuid.UUID]struct{}, len(updatedIDs))
+	for _, id := range updatedIDs {
+		updated[id] = struct{}{}
+	}
+
+	// Interrupt registered chats that were not in the result
+	// (stolen by another replica or already completed).
+	for id, entry := range snapshot {
+		if _, ok := updated[id]; !ok {
+			entry.logger.Warn(ctx, "chat not in batch heartbeat result, interrupting")
+			entry.cancelWithCause(chatloop.ErrInterrupted)
+			continue
+		}
+		// Bump workspace usage for surviving chats.
+		newWsID := p.trackWorkspaceUsage(ctx, entry.chatID, entry.workspaceID, entry.logger)
+		// Update workspace ID in the registry for next tick.
+		p.heartbeatMu.Lock()
+		if current, exists := p.heartbeatRegistry[id]; exists {
+			current.workspaceID = newWsID
+		}
+		p.heartbeatMu.Unlock()
+	}
+}
+
 func (p *Server) Subscribe(
 	ctx context.Context,
 	chatID uuid.UUID,
@@ -3575,33 +3687,17 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) {
 		}
 	}()
 
-	// Periodically update the heartbeat so other replicas know this
-	// worker is still alive. The goroutine stops when chatCtx is
-	// canceled (either by completion or interruption).
-	go func() {
-		ticker := p.clock.NewTicker(p.chatHeartbeatInterval, "chatd", "heartbeat")
-		defer ticker.Stop()
-		for {
-			select {
-			case <-chatCtx.Done():
-				return
-			case <-ticker.C:
-				rows, err := p.db.UpdateChatHeartbeat(chatCtx, database.UpdateChatHeartbeatParams{
-					ID:       chat.ID,
-					WorkerID: p.workerID,
-				})
-				if err != nil {
-					logger.Warn(chatCtx, "failed to update chat heartbeat", slog.Error(err))
-					continue
-				}
-				if rows == 0 {
-					cancel(chatloop.ErrInterrupted)
-					return
-				}
-				chat.WorkspaceID = p.trackWorkspaceUsage(chatCtx, chat.ID, chat.WorkspaceID, logger)
-			}
-		}
-	}()
+	// Register with the centralized heartbeat loop instead of
+	// running a per-chat goroutine. The loop issues a single batch
+	// UPDATE for all chats on this worker and detects stolen chats
+	// via set-difference.
+	p.registerHeartbeat(&heartbeatEntry{
+		cancelWithCause: cancel,
+		chatID:          chat.ID,
+		workspaceID:     chat.WorkspaceID,
+		logger:          logger,
+	})
+	defer p.unregisterHeartbeat(chat.ID)
 
 	// Start buffering stream events BEFORE publishing the running
 	// status. This closes a race where a subscriber sees

coderd/x/chatd/chatd_internal_test.go

@@ -21,6 +21,7 @@ import (
 	dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub"
 	coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
 	"github.com/coder/coder/v2/coderd/x/chatd/chaterror"
+	"github.com/coder/coder/v2/coderd/x/chatd/chatloop"
 	"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
 	"github.com/coder/coder/v2/coderd/x/chatd/chattest"
 	"github.com/coder/coder/v2/coderd/x/chatd/chattool"
@@ -2071,6 +2072,7 @@ func TestProcessChat_IgnoresStaleControlNotification(t *testing.T) {
 		workerID:              workerID,
 		chatHeartbeatInterval: time.Minute,
 		configCache:           newChatConfigCache(ctx, db, clock),
+		heartbeatRegistry:     make(map[uuid.UUID]*heartbeatEntry),
 	}
 
 	// Publish a stale "pending" notification on the control channel
@@ -2133,3 +2135,130 @@ func TestProcessChat_IgnoresStaleControlNotification(t *testing.T) {
 	require.Equal(t, database.ChatStatusError, finalStatus,
 		"processChat should have reached runChat (error), not been interrupted (waiting)")
 }
+
+// TestHeartbeatTick_StolenChatIsInterrupted verifies that when the
+// batch heartbeat UPDATE does not return a registered chat's ID
+// (because another replica stole it or it was completed), the
+// heartbeat tick cancels that chat's context with ErrInterrupted
+// while leaving surviving chats untouched.
+func TestHeartbeatTick_StolenChatIsInterrupted(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+	ctrl := gomock.NewController(t)
+	db := dbmock.NewMockStore(ctrl)
+	clock := quartz.NewMock(t)
+
+	workerID := uuid.New()
+
+	server := &Server{
+		db:                    db,
+		logger:                logger,
+		clock:                 clock,
+		workerID:              workerID,
+		chatHeartbeatInterval: time.Minute,
+		heartbeatRegistry:     make(map[uuid.UUID]*heartbeatEntry),
+	}
+
+	// Create three chats with independent cancel functions.
+	chat1 := uuid.New()
+	chat2 := uuid.New()
+	chat3 := uuid.New()
+
+	_, cancel1 := context.WithCancelCause(ctx)
+	_, cancel2 := context.WithCancelCause(ctx)
+	ctx3, cancel3 := context.WithCancelCause(ctx)
+
+	server.registerHeartbeat(&heartbeatEntry{
+		cancelWithCause: cancel1,
+		chatID:          chat1,
+		logger:          logger,
+	})
+	server.registerHeartbeat(&heartbeatEntry{
+		cancelWithCause: cancel2,
+		chatID:          chat2,
+		logger:          logger,
+	})
+	server.registerHeartbeat(&heartbeatEntry{
+		cancelWithCause: cancel3,
+		chatID:          chat3,
+		logger:          logger,
+	})
+
+	// The batch UPDATE returns only chat1 and chat2 —
+	// chat3 was "stolen" by another replica.
+	db.EXPECT().UpdateChatHeartbeats(gomock.Any(), gomock.Any()).DoAndReturn(
+		func(_ context.Context, params database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) {
+			require.Equal(t, workerID, params.WorkerID)
+			require.Len(t, params.IDs, 3)
+			// Return only chat1 and chat2 as surviving.
+			return []uuid.UUID{chat1, chat2}, nil
+		},
+	)
+
+	server.heartbeatTick(ctx)
+
+	// chat3's context should be canceled with ErrInterrupted.
+	require.ErrorIs(t, context.Cause(ctx3), chatloop.ErrInterrupted,
+		"stolen chat should be interrupted")
+
+	// chat3 should have been removed from the registry by
+	// unregister (in production this happens via defer in
+	// processChat). The heartbeat tick itself does not
+	// unregister — it only cancels. Verify the entry is
+	// still present (processChat's defer would clean it up).
+	server.heartbeatMu.Lock()
+	_, chat1Exists := server.heartbeatRegistry[chat1]
+	_, chat2Exists := server.heartbeatRegistry[chat2]
+	_, chat3Exists := server.heartbeatRegistry[chat3]
+	server.heartbeatMu.Unlock()
+
+	require.True(t, chat1Exists, "surviving chat1 should remain registered")
+	require.True(t, chat2Exists, "surviving chat2 should remain registered")
+	require.True(t, chat3Exists,
+		"stolen chat3 should still be in registry (processChat defer removes it)")
+}
+
+// TestHeartbeatTick_DBErrorDoesNotInterruptChats verifies that a
+// transient database failure causes the tick to log and return
+// without canceling any registered chats.
+func TestHeartbeatTick_DBErrorDoesNotInterruptChats(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+	ctrl := gomock.NewController(t)
+	db := dbmock.NewMockStore(ctrl)
+	clock := quartz.NewMock(t)
+
+	server := &Server{
+		db:                    db,
+		logger:                logger,
+		clock:                 clock,
+		workerID:              uuid.New(),
+		chatHeartbeatInterval: time.Minute,
+		heartbeatRegistry:     make(map[uuid.UUID]*heartbeatEntry),
+	}
+
+	chatID := uuid.New()
+	chatCtx, cancel := context.WithCancelCause(ctx)
+
+	server.registerHeartbeat(&heartbeatEntry{
+		cancelWithCause: cancel,
+		chatID:          chatID,
+		logger:          logger,
+	})
+
+	// Simulate a transient DB error.
+	db.EXPECT().UpdateChatHeartbeats(gomock.Any(), gomock.Any()).Return(
+		nil, xerrors.New("connection reset"),
+	)
+
+	server.heartbeatTick(ctx)
+
+	// Chat should NOT be interrupted — the tick logged and
+	// returned early.
+	require.NoError(t, chatCtx.Err(),
+		"chat context should not be canceled on transient DB error")
+}

coderd/x/chatd/chatd_test.go

@@ -474,7 +474,7 @@ func TestArchiveChatInterruptsActiveProcessing(t *testing.T) {
 	require.Equal(t, 1, userMessages, "expected queued message to stay queued after archive")
 }
 
-func TestUpdateChatHeartbeatRequiresOwnership(t *testing.T) {
+func TestUpdateChatHeartbeatsRequiresOwnership(t *testing.T) {
 	t.Parallel()
 
 	db, ps := dbtestutil.NewDB(t)
@@ -501,19 +501,24 @@ func TestUpdateChatHeartbeatRequiresOwnership(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	rows, err := db.UpdateChatHeartbeat(ctx, database.UpdateChatHeartbeatParams{
-		ID:       chat.ID,
+	// Wrong worker_id should return no IDs.
+	ids, err := db.UpdateChatHeartbeats(ctx, database.UpdateChatHeartbeatsParams{
+		IDs:      []uuid.UUID{chat.ID},
 		WorkerID: uuid.New(),
+		Now:      time.Now(),
 	})
 	require.NoError(t, err)
-	require.Equal(t, int64(0), rows)
+	require.Empty(t, ids)
 
-	rows, err = db.UpdateChatHeartbeat(ctx, database.UpdateChatHeartbeatParams{
-		ID:       chat.ID,
+	// Correct worker_id should return the chat's ID.
+	ids, err = db.UpdateChatHeartbeats(ctx, database.UpdateChatHeartbeatsParams{
+		IDs:      []uuid.UUID{chat.ID},
 		WorkerID: workerID,
+		Now:      time.Now(),
 	})
 	require.NoError(t, err)
-	require.Equal(t, int64(1), rows)
+	require.Len(t, ids, 1)
+	require.Equal(t, chat.ID, ids[0])
 }
 
 func TestSendMessageQueueBehaviorQueuesWhenBusy(t *testing.T) {