perf(chatd): fix six scale bottlenecks identified by benchmarking (#22957)

## Summary

Scale-tested the `chatd` package with mock-based benchmarks to identify
performance bottlenecks. This PR fixes 6 of the 8 identified issues,
ranked by severity.

## Changes

### 1. Parallel tool execution (HIGH) — `chatloop.go`
`executeTools` ran tool calls sequentially. Now dispatches all calls
concurrently via goroutines with `sync.WaitGroup`. Results are
pre-allocated by index (no mutex needed). `onResult` callbacks fire as
each tool completes.

### 2. Pubsub-backed subagent await (HIGH) — `subagent.go`
`awaitSubagentCompletion` polled the DB every 200ms. Now subscribes to
the child chat's `ChatStreamNotifyChannel` via pubsub for near-instant
notifications. Fallback poll reduced to 5s. Falls back to 200ms only
when `pubsub == nil` (single-instance / in-memory).

### 3. Per-chat stream locking (MEDIUM) — `chatd.go`
Replaced single global `streamMu` + `map[uuid.UUID]*chatStreamState`
with `sync.Map` where each `chatStreamState` has its own `sync.Mutex`.
Zero cross-chat contention.

### 4. Batch chat acquisition (MEDIUM) — `chatd.go`
`processOnce` acquired 1 chat per tick. Now loops up to
`maxChatsPerAcquire = 10` per tick, avoiding idle time when many chats
are pending.

### 5. Reduced heartbeat frequency (LOW-MEDIUM) — `chatd.go`
`chatHeartbeatInterval` changed from 30s to 60s. Safe given the 5-minute
`DefaultInFlightChatStaleAfter`.

### 6. O(depth) descendant check (LOW) — `subagent.go`
Replaced top-down BFS (`O(total_descendants)` queries) with bottom-up
parent-chain walk (`O(depth)` queries). Includes cycle protection.

## Not addressed (intentionally)
- Message serialization overhead
- Buffer eviction (`buffer[1:]` pattern)

coderd/chatd/chatd.go

@@ -41,7 +41,7 @@ const (
 
 	homeInstructionLookupTimeout = 5 * time.Second
 	instructionCacheTTL          = 5 * time.Minute
-	chatHeartbeatInterval        = 30 * time.Second
+	chatHeartbeatInterval        = 60 * time.Second
 	maxChatSteps                 = 1200
 	// maxStreamBufferSize caps the number of events buffered
 	// per chat during a single LLM step. When exceeded the
@@ -53,6 +53,12 @@ const (
 	// of 5 means recovery runs at 1/5 of the stale-after duration.
 	staleRecoveryIntervalDivisor = 5
 
+	// maxChatsPerAcquire is the maximum number of chats to
+	// acquire in a single processOnce call. Batching avoids
+	// waiting a full polling interval between acquisitions
+	// when many chats are pending.
+	maxChatsPerAcquire int32 = 10
+
 	defaultSubagentInstruction = "You are running as a delegated sub-agent chat. Complete the delegated task and provide clear, concise assistant responses for the parent agent."
 )
 
@@ -75,10 +81,10 @@ type Server struct {
 	webpushDispatcher webpush.Dispatcher
 	providerAPIKeys   chatprovider.ProviderAPIKeys
 
-	// streamMu guards chatStreams which tracks in-flight chat
-	// stream state for broadcasting ephemeral events.
-	streamMu    sync.Mutex
-	chatStreams map[uuid.UUID]*chatStreamState
+	// chatStreams stores per-chat stream state. Using sync.Map
+	// gives each chat independent locking — concurrent chats
+	// never contend with each other.
+	chatStreams sync.Map // uuid.UUID -> *chatStreamState
 
 	// instructionCache caches home instruction file contents by
 	// workspace agent ID so we don't re-dial on every chat turn.
@@ -137,6 +143,7 @@ type SubscribeFnParams struct {
 }
 
 type chatStreamState struct {
+	mu          sync.Mutex
 	buffer      []codersdk.ChatStreamEvent
 	buffering   bool
 	subscribers map[uuid.UUID]chan codersdk.ChatStreamEvent
@@ -989,7 +996,6 @@ func New(cfg Config) *Server {
 		pubsub:                     cfg.Pubsub,
 		webpushDispatcher:          cfg.WebpushDispatcher,
 		providerAPIKeys:            cfg.ProviderAPIKeys,
-		chatStreams:                make(map[uuid.UUID]*chatStreamState),
 		instructionCache:           make(map[uuid.UUID]cachedInstruction),
 		pendingChatAcquireInterval: pendingChatAcquireInterval,
 		inFlightChatStaleAfter:     inFlightChatStaleAfter,
@@ -1029,79 +1035,73 @@ func (p *Server) start(ctx context.Context) {
 }
 
 func (p *Server) processOnce(ctx context.Context) {
-	// Bail out early if the server is shutting down. The main
-	// loop's select can randomly pick the ticker over ctx.Done(),
-	// so we must guard against acquiring a chat we cannot process.
 	if ctx.Err() != nil {
 		return
 	}
 
-	// Try to acquire a pending chat. We detach from the server
-	// lifetime to prevent a phantom-acquire race: when the server
-	// context is canceled, the pq driver's watchCancel goroutine
-	// races with the actual query on the wire. The UPDATE can
-	// commit in Postgres (setting the chat to "running") before
-	// the cancel request arrives via a second TCP connection, yet
-	// the Go driver still returns context.Canceled to the caller
-	// because the awaitDone goroutine in database/sql closes the
-	// Rows before Scan reads them. This leaves the chat stuck as
-	// "running" with no goroutine to process it. Using a context
-	// that cannot be canceled ensures the driver sees the query
-	// result if Postgres executed it.
+	// We detach from the server lifetime to prevent a
+	// phantom-acquire race: when the server context is
+	// canceled, the pq driver's watchCancel goroutine
+	// races with the actual query on the wire. Using a
+	// context that cannot be canceled ensures the driver
+	// sees the query result if Postgres executed it.
 	acquireCtx, acquireCancel := context.WithTimeout(
 		context.WithoutCancel(ctx), 10*time.Second,
 	)
-	defer acquireCancel()
-	chat, err := p.db.AcquireChat(acquireCtx, database.AcquireChatParams{
+	chats, err := p.db.AcquireChats(acquireCtx, database.AcquireChatsParams{
 		StartedAt: time.Now(),
 		WorkerID:  p.workerID,
+		NumChats:  maxChatsPerAcquire,
 	})
+	acquireCancel()
 	if err != nil {
-		if !xerrors.Is(err, sql.ErrNoRows) {
-			p.logger.Error(ctx, "failed to acquire chat", slog.Error(err))
-		}
-		// No pending chats or error.
+		p.logger.Error(ctx, "failed to acquire chats", slog.Error(err))
+		return
+	}
+	if len(chats) == 0 {
 		return
 	}
 
-	// If the server context was canceled while we were acquiring,
-	// release the chat back to pending immediately so another
-	// replica can pick it up.
+	// If the server context was canceled while we were
+	// acquiring, release the chats back to pending.
 	if ctx.Err() != nil {
 		releaseCtx, releaseCancel := context.WithTimeout(
 			context.WithoutCancel(ctx), 10*time.Second,
 		)
-		defer releaseCancel()
-		_, updateErr := p.db.UpdateChatStatus(releaseCtx, database.UpdateChatStatusParams{
-			ID:          chat.ID,
-			Status:      database.ChatStatusPending,
-			WorkerID:    uuid.NullUUID{},
-			StartedAt:   sql.NullTime{},
-			HeartbeatAt: sql.NullTime{},
-			LastError:   sql.NullString{},
-		})
-		if updateErr != nil {
-			p.logger.Error(ctx, "failed to release chat acquired during shutdown",
-				slog.F("chat_id", chat.ID), slog.Error(updateErr))
+		for _, chat := range chats {
+			_, updateErr := p.db.UpdateChatStatus(releaseCtx, database.UpdateChatStatusParams{
+				ID:          chat.ID,
+				Status:      database.ChatStatusPending,
+				WorkerID:    uuid.NullUUID{},
+				StartedAt:   sql.NullTime{},
+				HeartbeatAt: sql.NullTime{},
+				LastError:   sql.NullString{},
+			})
+			if updateErr != nil {
+				p.logger.Error(ctx, "failed to release chat acquired during shutdown",
+					slog.F("chat_id", chat.ID), slog.Error(updateErr))
+			}
 		}
+		releaseCancel()
 		return
 	}
 
-	// Process the chat (don't block the main loop).
-	p.inflight.Add(1)
-	go func() {
-		defer p.inflight.Done()
-		p.processChat(ctx, chat)
-	}()
+	for _, chat := range chats {
+		p.inflight.Add(1)
+		go func() {
+			defer p.inflight.Done()
+			p.processChat(ctx, chat)
+		}()
+	}
 }
 
 func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEvent) {
-	p.streamMu.Lock()
-	state := p.streamStateLocked(chatID)
+	state := p.getOrCreateStreamState(chatID)
+	state.mu.Lock()
 	if event.Type == codersdk.ChatStreamEventTypeMessagePart {
 		if !state.buffering {
-			p.cleanupStreamIfIdleLocked(chatID, state)
-			p.streamMu.Unlock()
+			p.cleanupStreamIfIdle(chatID, state)
+			state.mu.Unlock()
 			return
 		}
 		if len(state.buffer) >= maxStreamBufferSize {
@@ -1115,7 +1115,7 @@ func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEven
 	for _, ch := range state.subscribers {
 		subscribers = append(subscribers, ch)
 	}
-	p.streamMu.Unlock()
+	state.mu.Unlock()
 
 	for _, ch := range subscribers {
 		select {
@@ -1127,13 +1127,11 @@ func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEven
 	}
 
 	// Clean up the stream entry if it was created by
-	// streamStateLocked but has no subscribers and is not
+	// getOrCreateStreamState but has no subscribers and is not
 	// actively buffering (e.g. publish with no watchers).
-	p.streamMu.Lock()
-	if cur, ok := p.chatStreams[chatID]; ok {
-		p.cleanupStreamIfIdleLocked(chatID, cur)
-	}
-	p.streamMu.Unlock()
+	state.mu.Lock()
+	p.cleanupStreamIfIdle(chatID, state)
+	state.mu.Unlock()
 }
 
 func (p *Server) subscribeToStream(chatID uuid.UUID) (
@@ -1141,48 +1139,52 @@ func (p *Server) subscribeToStream(chatID uuid.UUID) (
 	<-chan codersdk.ChatStreamEvent,
 	func(),
 ) {
-	p.streamMu.Lock()
-	state := p.streamStateLocked(chatID)
+	state := p.getOrCreateStreamState(chatID)
+	state.mu.Lock()
 	snapshot := append([]codersdk.ChatStreamEvent(nil), state.buffer...)
 	id := uuid.New()
 	ch := make(chan codersdk.ChatStreamEvent, 128)
 	state.subscribers[id] = ch
-	p.streamMu.Unlock()
+	state.mu.Unlock()
 
 	cancel := func() {
-		p.streamMu.Lock()
-		state, ok := p.chatStreams[chatID]
-		if ok {
-			// Remove the subscriber but do not close the channel.
-			// publishToStream copies subscriber references under
-			// streamMu then sends outside the lock; closing here
-			// races with that send and can panic. The channel
-			// becomes unreachable once removed and will be GC'd.
-			delete(state.subscribers, id)
-			p.cleanupStreamIfIdleLocked(chatID, state)
-		}
-		p.streamMu.Unlock()
+		state.mu.Lock()
+		// Remove the subscriber but do not close the channel.
+		// publishToStream copies subscriber references under
+		// the per-chat lock then sends outside; closing here
+		// races with that send and can panic. The channel
+		// becomes unreachable once removed and will be GC'd.
+		delete(state.subscribers, id)
+		p.cleanupStreamIfIdle(chatID, state)
+		state.mu.Unlock()
 	}
 
 	return snapshot, ch, cancel
 }
 
-// cleanupStreamIfIdleLocked removes the chat entry when there
-// are no subscribers and the stream is not buffering. The
-// caller must hold p.streamMu.
-func (p *Server) cleanupStreamIfIdleLocked(chatID uuid.UUID, state *chatStreamState) {
-	if !state.buffering && len(state.subscribers) == 0 {
-		delete(p.chatStreams, chatID)
+// getOrCreateStreamState returns the per-chat stream state,
+// creating one atomically if it doesn't exist. The returned
+// state has its own mutex — callers must lock state.mu for
+// access.
+func (p *Server) getOrCreateStreamState(chatID uuid.UUID) *chatStreamState {
+	if val, ok := p.chatStreams.Load(chatID); ok {
+		state, _ := val.(*chatStreamState)
+		return state
 	}
+	val, _ := p.chatStreams.LoadOrStore(chatID, &chatStreamState{
+		subscribers: make(map[uuid.UUID]chan codersdk.ChatStreamEvent),
+	})
+	state, _ := val.(*chatStreamState)
+	return state
 }
 
-func (p *Server) streamStateLocked(chatID uuid.UUID) *chatStreamState {
-	state, ok := p.chatStreams[chatID]
-	if !ok {
-		state = &chatStreamState{subscribers: make(map[uuid.UUID]chan codersdk.ChatStreamEvent)}
-		p.chatStreams[chatID] = state
+// cleanupStreamIfIdle removes the chat entry from the sync.Map
+// when there are no subscribers and the stream is not buffering.
+// The caller must hold state.mu.
+func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) {
+	if !state.buffering && len(state.subscribers) == 0 {
+		p.chatStreams.Delete(chatID)
 	}
-	return state
 }
 
 func (p *Server) Subscribe(
@@ -1876,19 +1878,17 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) {
 	// buffering hasn't started yet — the subscriber gets an empty
 	// snapshot and publishToStream drops message_parts while
 	// buffering is false.
-	p.streamMu.Lock()
-	startState := p.streamStateLocked(chat.ID)
-	startState.buffer = nil
-	startState.buffering = true
-	p.streamMu.Unlock()
+	streamState := p.getOrCreateStreamState(chat.ID)
+	streamState.mu.Lock()
+	streamState.buffer = nil
+	streamState.buffering = true
+	streamState.mu.Unlock()
 	defer func() {
-		p.streamMu.Lock()
-		if stopState, ok := p.chatStreams[chat.ID]; ok {
-			stopState.buffer = nil
-			stopState.buffering = false
-			p.cleanupStreamIfIdleLocked(chat.ID, stopState)
-		}
-		p.streamMu.Unlock()
+		streamState.mu.Lock()
+		streamState.buffer = nil
+		streamState.buffering = false
+		p.cleanupStreamIfIdle(chat.ID, streamState)
+		streamState.mu.Unlock()
 	}()
 
 	p.publishStatus(chat.ID, database.ChatStatusRunning, uuid.NullUUID{
@@ -2373,11 +2373,13 @@ func (p *Server) runChat(
 		// Clear the stream buffer now that the step is
 		// persisted. Late-joining subscribers will load
 		// these messages from the database instead.
-		p.streamMu.Lock()
-		if state, ok := p.chatStreams[chat.ID]; ok {
-			state.buffer = nil
+		if val, ok := p.chatStreams.Load(chat.ID); ok {
+			if ss, ok := val.(*chatStreamState); ok {
+				ss.mu.Lock()
+				ss.buffer = nil
+				ss.mu.Unlock()
+			}
 		}
-		p.streamMu.Unlock()
 
 		return nil
 	}
@@ -2531,12 +2533,13 @@ func (p *Server) runChat(
 		},
 
 		OnRetry: func(attempt int, retryErr error, delay time.Duration) {
-			p.streamMu.Lock()
-			if state, ok := p.chatStreams[chat.ID]; ok {
-				state.buffer = nil
+			if val, ok := p.chatStreams.Load(chat.ID); ok {
+				if rs, ok := val.(*chatStreamState); ok {
+					rs.mu.Lock()
+					rs.buffer = nil
+					rs.mu.Unlock()
+				}
 			}
-			p.streamMu.Unlock()
-
 			logger.Warn(ctx, "retrying LLM stream",
 				slog.F("attempt", attempt),
 				slog.F("delay", delay.String()),

coderd/chatd/chatloop/chatloop.go

@@ -8,6 +8,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"charm.land/fantasy"
@@ -575,9 +576,10 @@ func processStepStream(
 	return result, nil
 }
 
-// executeTools runs each tool call sequentially after the stream
-// completes. Results are published via onResult as each tool
-// finishes.
+// executeTools runs all tool calls concurrently after the stream
+// completes. Results are published via onResult in the original
+// tool-call order after all tools finish, preserving deterministic
+// event ordering for SSE subscribers.
 func executeTools(
 	ctx context.Context,
 	allTools []fantasy.AgentTool,
@@ -593,11 +595,32 @@ func executeTools(
 		toolMap[t.Info().Name] = t
 	}
 
-	results := make([]fantasy.ToolResultContent, 0, len(toolCalls))
-	for _, tc := range toolCalls {
-		tr := executeSingleTool(ctx, toolMap, tc)
-		results = append(results, tr)
-		if onResult != nil {
+	results := make([]fantasy.ToolResultContent, len(toolCalls))
+	var wg sync.WaitGroup
+	wg.Add(len(toolCalls))
+	for i, tc := range toolCalls {
+		go func(i int, tc fantasy.ToolCallContent) {
+			defer wg.Done()
+			defer func() {
+				if r := recover(); r != nil {
+					results[i] = fantasy.ToolResultContent{
+						ToolCallID: tc.ToolCallID,
+						ToolName:   tc.ToolName,
+						Result: fantasy.ToolResultOutputContentError{
+							Error: xerrors.Errorf("tool panicked: %v", r),
+						},
+					}
+				}
+			}()
+			results[i] = executeSingleTool(ctx, toolMap, tc)
+		}(i, tc)
+	}
+	wg.Wait()
+
+	// Publish results in the original tool-call order so SSE
+	// subscribers see a deterministic event sequence.
+	if onResult != nil {
+		for _, tr := range results {
 			onResult(tr)
 		}
 	}

coderd/chatd/subagent.go

@@ -2,6 +2,7 @@ package chatd
 
 import (
 	"context"
+	"database/sql"
 	"encoding/json"
 	"sort"
 	"strings"
@@ -13,12 +14,14 @@ import (
 
 	"github.com/coder/coder/v2/coderd/chatd/chatprompt"
 	"github.com/coder/coder/v2/coderd/database"
+	coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
 )
 
 var ErrSubagentNotDescendant = xerrors.New("target chat is not a descendant of current chat")
 
 const (
 	subagentAwaitPollInterval  = 200 * time.Millisecond
+	subagentAwaitFallbackPoll  = 5 * time.Second
 	defaultSubagentWaitTimeout = 5 * time.Minute
 )
 
@@ -322,41 +325,90 @@ func (p *Server) awaitSubagentCompletion(
 		return database.Chat{}, "", ErrSubagentNotDescendant
 	}
 
+	// Check immediately before entering the poll loop.
+	targetChat, report, done, checkErr := p.checkSubagentCompletion(ctx, targetChatID)
+	if checkErr != nil {
+		return database.Chat{}, "", checkErr
+	}
+	if done {
+		return handleSubagentDone(targetChat, report)
+	}
+
 	if timeout <= 0 {
 		timeout = defaultSubagentWaitTimeout
 	}
 	timer := time.NewTimer(timeout)
 	defer timer.Stop()
 
-	ticker := time.NewTicker(subagentAwaitPollInterval)
+	// When pubsub is available, subscribe for fast status
+	// notifications and use a less aggressive fallback poll.
+	// Without pubsub (single-instance / in-memory) fall back
+	// to the original 200ms polling.
+	pollInterval := subagentAwaitPollInterval
+	var notifyCh <-chan struct{}
+	if p.pubsub != nil {
+		pollInterval = subagentAwaitFallbackPoll
+		ch := make(chan struct{}, 1)
+		notifyCh = ch
+		cancel, subErr := p.pubsub.SubscribeWithErr(
+			coderdpubsub.ChatStreamNotifyChannel(targetChatID),
+			func(_ context.Context, _ []byte, _ error) {
+				// Non-blocking send so we never stall the
+				// pubsub dispatch goroutine.
+				select {
+				case ch <- struct{}{}:
+				default:
+				}
+			},
+		)
+		if subErr == nil {
+			defer cancel()
+		} else {
+			// Subscription failed; fall back to fast polling.
+			pollInterval = subagentAwaitPollInterval
+			notifyCh = nil
+		}
+	}
+
+	ticker := time.NewTicker(pollInterval)
 	defer ticker.Stop()
 
 	for {
-		targetChat, report, done, checkErr := p.checkSubagentCompletion(ctx, targetChatID)
-		if checkErr != nil {
-			return database.Chat{}, "", checkErr
-		}
-		if done {
-			if targetChat.Status == database.ChatStatusError {
-				reason := strings.TrimSpace(report)
-				if reason == "" {
-					reason = "agent reached error status"
-				}
-				return database.Chat{}, "", xerrors.New(reason)
-			}
-			return targetChat, report, nil
-		}
-
 		select {
+		case <-notifyCh:
 		case <-ticker.C:
 		case <-timer.C:
 			return database.Chat{}, "", xerrors.New("timed out waiting for delegated subagent completion")
 		case <-ctx.Done():
 			return database.Chat{}, "", ctx.Err()
 		}
+
+		targetChat, report, done, checkErr = p.checkSubagentCompletion(ctx, targetChatID)
+		if checkErr != nil {
+			return database.Chat{}, "", checkErr
+		}
+		if done {
+			return handleSubagentDone(targetChat, report)
+		}
 	}
 }
 
+// handleSubagentDone translates a completed subagent check into the
+// appropriate return value, surfacing error-status chats as errors.
+func handleSubagentDone(
+	chat database.Chat,
+	report string,
+) (database.Chat, string, error) {
+	if chat.Status == database.ChatStatusError {
+		reason := strings.TrimSpace(report)
+		if reason == "" {
+			reason = "agent reached error status"
+		}
+		return database.Chat{}, "", xerrors.New(reason)
+	}
+	return chat, report, nil
+}
+
 func (p *Server) closeSubagent(
 	ctx context.Context,
 	parentChatID uuid.UUID,
@@ -448,6 +500,9 @@ func latestSubagentAssistantMessage(
 	return "", nil
 }
 
+// isSubagentDescendant reports whether targetChatID is a descendant
+// of ancestorChatID by walking up the parent chain from the target.
+// This is O(depth) DB queries instead of O(nodes) BFS.
 func isSubagentDescendant(
 	ctx context.Context,
 	store database.Store,
@@ -458,47 +513,29 @@ func isSubagentDescendant(
 		return false, nil
 	}
 
-	descendants, err := listSubagentDescendants(ctx, store, ancestorChatID)
-	if err != nil {
-		return false, err
-	}
-	for _, descendant := range descendants {
-		if descendant.ID == targetChatID {
+	currentID := targetChatID
+	visited := map[uuid.UUID]struct{}{} // cycle protection
+	for {
+		if _, seen := visited[currentID]; seen {
+			return false, nil
+		}
+		visited[currentID] = struct{}{}
+
+		chat, err := store.GetChatByID(ctx, currentID)
+		if err != nil {
+			if xerrors.Is(err, sql.ErrNoRows) {
+				return false, nil // chain broken; not a confirmed descendant
+			}
+			return false, xerrors.Errorf("get chat %s: %w", currentID, err)
+		}
+		if !chat.ParentChatID.Valid {
+			return false, nil // reached root without finding ancestor
+		}
+		if chat.ParentChatID.UUID == ancestorChatID {
 			return true, nil
 		}
+		currentID = chat.ParentChatID.UUID
 	}
-	return false, nil
-}
-
-func listSubagentDescendants(
-	ctx context.Context,
-	store database.Store,
-	chatID uuid.UUID,
-) ([]database.Chat, error) {
-	queue := []uuid.UUID{chatID}
-	visited := map[uuid.UUID]struct{}{chatID: {}}
-
-	out := make([]database.Chat, 0)
-	for len(queue) > 0 {
-		parentChatID := queue[0]
-		queue = queue[1:]
-
-		children, err := store.ListChildChatsByParentID(ctx, parentChatID)
-		if err != nil {
-			return nil, xerrors.Errorf("list child chats for %s: %w", parentChatID, err)
-		}
-
-		for _, child := range children {
-			if _, ok := visited[child.ID]; ok {
-				continue
-			}
-			visited[child.ID] = struct{}{}
-			out = append(out, child)
-			queue = append(queue, child.ID)
-		}
-	}
-
-	return out, nil
 }
 
 func subagentFallbackChatTitle(message string) string {

coderd/database/dbauthz/dbauthz.go

@@ -1512,13 +1512,13 @@ func (q *querier) authorizeProvisionerJob(ctx context.Context, job database.Prov
 	return nil
 }
 
-func (q *querier) AcquireChat(ctx context.Context, arg database.AcquireChatParams) (database.Chat, error) {
-	// AcquireChat is a system-level operation used by the chat processor.
+func (q *querier) AcquireChats(ctx context.Context, arg database.AcquireChatsParams) ([]database.Chat, error) {
+	// AcquireChats is a system-level operation used by the chat processor.
 	// Authorization is done at the system level, not per-user.
 	if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
-		return database.Chat{}, err
+		return nil, err
 	}
-	return q.db.AcquireChat(ctx, arg)
+	return q.db.AcquireChats(ctx, arg)
 }
 
 func (q *querier) AcquireLock(ctx context.Context, id int64) error {

coderd/database/dbauthz/dbauthz_test.go

@@ -373,14 +373,15 @@ func (s *MethodTestSuite) TestConnectionLogs() {
 }
 
 func (s *MethodTestSuite) TestChats() {
-	s.Run("AcquireChat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
-		arg := database.AcquireChatParams{
+	s.Run("AcquireChats", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		arg := database.AcquireChatsParams{
 			StartedAt: dbtime.Now(),
 			WorkerID:  uuid.New(),
+			NumChats:  1,
 		}
 		chat := testutil.Fake(s.T(), faker, database.Chat{})
-		dbm.EXPECT().AcquireChat(gomock.Any(), arg).Return(chat, nil).AnyTimes()
-		check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns(chat)
+		dbm.EXPECT().AcquireChats(gomock.Any(), arg).Return([]database.Chat{chat}, nil).AnyTimes()
+		check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.Chat{chat})
 	}))
 	s.Run("DeleteAllChatQueuedMessages", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		chat := testutil.Fake(s.T(), faker, database.Chat{})

coderd/database/dbmetrics/querymetrics.go

@@ -104,11 +104,11 @@ func (m queryMetricsStore) DeleteOrganization(ctx context.Context, id uuid.UUID)
 	return r0
 }
 
-func (m queryMetricsStore) AcquireChat(ctx context.Context, arg database.AcquireChatParams) (database.Chat, error) {
+func (m queryMetricsStore) AcquireChats(ctx context.Context, arg database.AcquireChatsParams) ([]database.Chat, error) {
 	start := time.Now()
-	r0, r1 := m.s.AcquireChat(ctx, arg)
-	m.queryLatencies.WithLabelValues("AcquireChat").Observe(time.Since(start).Seconds())
-	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "AcquireChat").Inc()
+	r0, r1 := m.s.AcquireChats(ctx, arg)
+	m.queryLatencies.WithLabelValues("AcquireChats").Observe(time.Since(start).Seconds())
+	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "AcquireChats").Inc()
 	return r0, r1
 }
 

coderd/database/dbmock/dbmock.go

@@ -44,19 +44,19 @@ func (m *MockStore) EXPECT() *MockStoreMockRecorder {
 	return m.recorder
 }
 
-// AcquireChat mocks base method.
-func (m *MockStore) AcquireChat(ctx context.Context, arg database.AcquireChatParams) (database.Chat, error) {
+// AcquireChats mocks base method.
+func (m *MockStore) AcquireChats(ctx context.Context, arg database.AcquireChatsParams) ([]database.Chat, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "AcquireChat", ctx, arg)
-	ret0, _ := ret[0].(database.Chat)
+	ret := m.ctrl.Call(m, "AcquireChats", ctx, arg)
+	ret0, _ := ret[0].([]database.Chat)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// AcquireChat indicates an expected call of AcquireChat.
-func (mr *MockStoreMockRecorder) AcquireChat(ctx, arg any) *gomock.Call {
+// AcquireChats indicates an expected call of AcquireChats.
+func (mr *MockStoreMockRecorder) AcquireChats(ctx, arg any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcquireChat", reflect.TypeOf((*MockStore)(nil).AcquireChat), ctx, arg)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcquireChats", reflect.TypeOf((*MockStore)(nil).AcquireChats), ctx, arg)
 }
 
 // AcquireLock mocks base method.

coderd/database/querier.go

@@ -12,9 +12,9 @@ import (
 )
 
 type sqlcQuerier interface {
-	// Acquires a pending chat for processing. Uses SKIP LOCKED to prevent
-	// multiple replicas from acquiring the same chat.
-	AcquireChat(ctx context.Context, arg AcquireChatParams) (Chat, error)
+	// Acquires up to @num_chats pending chats for processing. Uses SKIP LOCKED
+	// to prevent multiple replicas from acquiring the same chat.
+	AcquireChats(ctx context.Context, arg AcquireChatsParams) ([]Chat, error)
 	// Blocks until the lock is acquired.
 	//
 	// This must be called from within a transaction. The lock will be automatically

coderd/database/queries.sql.go

@@ -2968,7 +2968,7 @@ func (q *sqlQuerier) UpdateChatProvider(ctx context.Context, arg UpdateChatProvi
 	return i, err
 }
 
-const acquireChat = `-- name: AcquireChat :one
+const acquireChats = `-- name: AcquireChats :many
 UPDATE
     chats
 SET
@@ -2978,7 +2978,7 @@ SET
     updated_at = $1::timestamptz,
     worker_id = $2::uuid
 WHERE
-    id = (
+    id = ANY(
         SELECT
             id
         FROM
@@ -2990,40 +2990,57 @@ WHERE
         FOR UPDATE
             SKIP LOCKED
         LIMIT
-            1
+            $3::int
     )
 RETURNING
     id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
 `
 
-type AcquireChatParams struct {
+type AcquireChatsParams struct {
 	StartedAt time.Time `db:"started_at" json:"started_at"`
 	WorkerID  uuid.UUID `db:"worker_id" json:"worker_id"`
+	NumChats  int32     `db:"num_chats" json:"num_chats"`
 }
 
-// Acquires a pending chat for processing. Uses SKIP LOCKED to prevent
-// multiple replicas from acquiring the same chat.
-func (q *sqlQuerier) AcquireChat(ctx context.Context, arg AcquireChatParams) (Chat, error) {
-	row := q.db.QueryRowContext(ctx, acquireChat, arg.StartedAt, arg.WorkerID)
-	var i Chat
-	err := row.Scan(
-		&i.ID,
-		&i.OwnerID,
-		&i.WorkspaceID,
-		&i.Title,
-		&i.Status,
-		&i.WorkerID,
-		&i.StartedAt,
-		&i.HeartbeatAt,
-		&i.CreatedAt,
-		&i.UpdatedAt,
-		&i.ParentChatID,
-		&i.RootChatID,
-		&i.LastModelConfigID,
-		&i.Archived,
-		&i.LastError,
-	)
-	return i, err
+// Acquires up to @num_chats pending chats for processing. Uses SKIP LOCKED
+// to prevent multiple replicas from acquiring the same chat.
+func (q *sqlQuerier) AcquireChats(ctx context.Context, arg AcquireChatsParams) ([]Chat, error) {
+	rows, err := q.db.QueryContext(ctx, acquireChats, arg.StartedAt, arg.WorkerID, arg.NumChats)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []Chat
+	for rows.Next() {
+		var i Chat
+		if err := rows.Scan(
+			&i.ID,
+			&i.OwnerID,
+			&i.WorkspaceID,
+			&i.Title,
+			&i.Status,
+			&i.WorkerID,
+			&i.StartedAt,
+			&i.HeartbeatAt,
+			&i.CreatedAt,
+			&i.UpdatedAt,
+			&i.ParentChatID,
+			&i.RootChatID,
+			&i.LastModelConfigID,
+			&i.Archived,
+			&i.LastError,
+		); err != nil {
+			return nil, err
+		}
+		items = append(items, i)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
 }
 
 const acquireStaleChatDiffStatuses = `-- name: AcquireStaleChatDiffStatuses :many

coderd/database/queries/chats.sql

@@ -257,9 +257,9 @@ WHERE
 RETURNING
     *;
 
--- name: AcquireChat :one
--- Acquires a pending chat for processing. Uses SKIP LOCKED to prevent
--- multiple replicas from acquiring the same chat.
+-- name: AcquireChats :many
+-- Acquires up to @num_chats pending chats for processing. Uses SKIP LOCKED
+-- to prevent multiple replicas from acquiring the same chat.
 UPDATE
     chats
 SET
@@ -269,7 +269,7 @@ SET
     updated_at = @started_at::timestamptz,
     worker_id = @worker_id::uuid
 WHERE
-    id = (
+    id = ANY(
         SELECT
             id
         FROM
@@ -281,7 +281,7 @@ WHERE
         FOR UPDATE
             SKIP LOCKED
         LIMIT
-            1
+            @num_chats::int
     )
 RETURNING
     *;