mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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)
This commit is contained in:
+110
-107
@@ -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()),
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
+91
-54
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user