1package subagents
2
3import (
4 "context"
5 "fmt"
6 "sort"
7 "strconv"
8 "strings"
9 "sync"
10)
11
12// Completion is the observed outcome of one delegated turn. The values are
13// kept unbounded so callers can decide how much evidence to retain; the
14// shared formatter applies the bounds used in parent-agent updates.
15type Completion struct {
16 AgentID string
17 Status string
18 Task string
19 Error string
20 FinalResponse string
21 Tail string
22
23 // turnError preserves the label used by the interactive update for an
24 // error reported by a prompt-level turn_end rather than by the worker
25 // process itself.
26 turnError bool
27 cancellation bool
28
29 registrationSeq uint64
30}
31
32type completionEntry struct {
33 state *completionAgent
34 task string
35 followUp bool
36 expectedStep int
37 strictStep bool
38 registrationSeq uint64
39 done bool
40 completion *Completion
41}
42
43type completionAgent struct {
44 agent *Agent
45 entries []*completionEntry
46 seenTurnEnds map[int]struct{}
47 turnMu sync.Mutex
48 callbackInstalled bool
49 waitCancel context.CancelFunc
50 waitingForProcess bool
51 waitDone chan struct{}
52}
53
54// CompletionTracker observes delegated turns through Agent's event and
55// lifecycle callbacks. Completed values are buffered so a worker can finish
56// before the parent begins waiting for the batch.
57type CompletionTracker struct {
58 mu sync.Mutex
59 active map[*Agent]*completionAgent
60 ready []Completion
61 nextRegistration uint64
62 changed chan struct{}
63}
64
65// NewCompletionTracker creates an empty completion tracker.
66func NewCompletionTracker() *CompletionTracker {
67 return &CompletionTracker{
68 active: make(map[*Agent]*completionAgent),
69 changed: make(chan struct{}),
70 }
71}
72
73func (t *CompletionTracker) initLocked() {
74 if t.active == nil {
75 t.active = make(map[*Agent]*completionAgent)
76 }
77 if t.changed == nil {
78 t.changed = make(chan struct{})
79 }
80}
81
82// TrackTurn tracks the worker's current prompt-level turn, or its first turn
83// when no turn has started yet. Multiple turns for the same long-lived worker
84// are retained and dispatched in registration order.
85func (t *CompletionTracker) TrackTurn(a *Agent, task string, followUp bool) func() {
86 return t.trackTurn(a, task, followUp, false)
87}
88
89// TrackFutureTurn tracks a prompt-level turn before its prompt is delivered.
90// The returned function removes the registration without producing a
91// completion; callers use it when prompt delivery is rejected.
92func (t *CompletionTracker) TrackFutureTurn(a *Agent, task string, followUp bool) func() {
93 return t.trackTurn(a, task, followUp, true)
94}
95
96func (t *CompletionTracker) trackTurn(a *Agent, task string, followUp, future bool) func() {
97 if t == nil || a == nil {
98 return nil
99 }
100
101 state, entry, installCallback := t.addTurn(a, task, followUp, future)
102 if state == nil {
103 return nil
104 }
105 if installCallback {
106 // SetOnTurnEnd is deliberately installed before the process waiter and
107 // before the caller sends a follow-up. It also replays notices that
108 // arrived while no callback was registered.
109 a.SetOnTurnEnd(func(step int, errMsg string) {
110 state.turnMu.Lock()
111 defer state.turnMu.Unlock()
112 t.completeNextTurn(state, step, errMsg)
113 })
114 }
115 t.startProcessWaiter(state)
116 return func() { t.cancelEntry(entry) }
117}
118
119// TrackExit tracks only process termination. It is used for explicit worker
120// stops when no prompt-level turn watcher owns the worker.
121func (t *CompletionTracker) TrackExit(a *Agent, task string) {
122 if t == nil || a == nil {
123 return
124 }
125 state, entry := t.addExit(a, task)
126 if state == nil || entry == nil {
127 return
128 }
129 t.startProcessWaiter(state)
130}
131
132func (t *CompletionTracker) addTurn(a *Agent, task string, followUp, future bool) (*completionAgent, *completionEntry, bool) {
133 entry := &completionEntry{task: task, followUp: followUp, expectedStep: expectedTurnStep(a, future), strictStep: future}
134 t.mu.Lock()
135 defer t.mu.Unlock()
136 t.initLocked()
137 state := t.active[a]
138 if state == nil {
139 state = &completionAgent{agent: a, seenTurnEnds: make(map[int]struct{})}
140 t.active[a] = state
141 }
142 if state.seenTurnEnds == nil {
143 state.seenTurnEnds = make(map[int]struct{})
144 }
145 if count := len(state.entries); count != 0 {
146 previousStep := state.entries[count-1].expectedStep
147 if entry.expectedStep <= previousStep {
148 entry.expectedStep = previousStep + 1
149 }
150 }
151 t.nextRegistration++
152 entry.registrationSeq = t.nextRegistration
153 entry.state = state
154 state.entries = append(state.entries, entry)
155 installCallback := !state.callbackInstalled
156 state.callbackInstalled = true
157 return state, entry, installCallback
158}
159
160func (t *CompletionTracker) addExit(a *Agent, task string) (*completionAgent, *completionEntry) {
161 t.mu.Lock()
162 defer t.mu.Unlock()
163 t.initLocked()
164 if state := t.active[a]; state != nil && len(state.entries) != 0 {
165 // A turn watcher already owns this worker. Its process fallback is
166 // sufficient and must not create a duplicate terminal completion.
167 return nil, nil
168 }
169 state := &completionAgent{agent: a}
170 t.nextRegistration++
171 entry := &completionEntry{state: state, task: task, registrationSeq: t.nextRegistration}
172 state.entries = append(state.entries, entry)
173 t.active[a] = state
174 return state, entry
175}
176
177func (t *CompletionTracker) startProcessWaiter(state *completionAgent) {
178 if state == nil || state.agent == nil {
179 return
180 }
181 t.mu.Lock()
182 if t.active[state.agent] != state || state.waitingForProcess {
183 t.mu.Unlock()
184 return
185 }
186 waitCtx, cancel := context.WithCancel(context.Background())
187 state.waitingForProcess = true
188 state.waitCancel = cancel
189 state.waitDone = make(chan struct{})
190 waitDone := state.waitDone
191 t.mu.Unlock()
192
193 go func() {
194 defer close(waitDone)
195 err := state.agent.WaitContext(waitCtx)
196 if err == nil {
197 t.completeProcess(state)
198 }
199 }()
200}
201
202func expectedTurnStep(a *Agent, future bool) int {
203 step := a.LifetimeTurnsValue()
204 if turnID := strings.TrimPrefix(a.CurrentTurnID(), "turn-"); turnID != "" {
205 if current, err := strconv.Atoi(turnID); err == nil && current > step {
206 step = current
207 }
208 }
209 if future || step == 0 {
210 step++
211 }
212 return step
213}
214
215func (t *CompletionTracker) completeNextTurn(state *completionAgent, step int, turnErr string) {
216 if state == nil {
217 return
218 }
219 t.mu.Lock()
220 if t.active[state.agent] != state || len(state.entries) == 0 {
221 t.mu.Unlock()
222 return
223 }
224 if _, duplicate := state.seenTurnEnds[step]; duplicate {
225 t.mu.Unlock()
226 return
227 }
228 state.seenTurnEnds[step] = struct{}{}
229 var entry *completionEntry
230 for _, candidate := range state.entries {
231 if !candidate.done && candidate.expectedStep == step {
232 entry = candidate
233 break
234 }
235 }
236 if entry == nil {
237 // Post-success callers can register after a resumed worker was queued
238 // but before its next turn.start updated metadata. Accept the first
239 // later step for that compatibility path; pre-delivery registrations
240 // use strict identity and never consume a stale or unrelated notice.
241 for _, candidate := range state.entries {
242 if !candidate.done && !candidate.strictStep && candidate.expectedStep < step {
243 entry = candidate
244 candidate.expectedStep = step
245 break
246 }
247 }
248 }
249 t.mu.Unlock()
250 if entry != nil {
251 t.completeTurn(entry, turnErr)
252 }
253}
254
255func (t *CompletionTracker) cancelEntry(entry *completionEntry) {
256 if entry == nil || entry.state == nil {
257 return
258 }
259 var cancel context.CancelFunc
260 t.mu.Lock()
261 state := entry.state
262 if entry.done || t.active[state.agent] != state {
263 t.mu.Unlock()
264 return
265 }
266 entry.done = true
267 t.flushReadyLocked(state)
268 if len(state.entries) == 0 {
269 delete(t.active, state.agent)
270 cancel = state.waitCancel
271 state.waitCancel = nil
272 state.waitingForProcess = false
273 }
274 t.signalLocked()
275 t.mu.Unlock()
276 if cancel != nil {
277 cancel()
278 }
279}
280
281func (t *CompletionTracker) completeTurn(entry *completionEntry, turnErr string) {
282 if entry == nil || entry.state == nil || entry.state.agent == nil {
283 return
284 }
285 snap := entry.state.agent.Snapshot()
286 completion := makeCompletion(entry, snap, turnErr, true)
287
288 var cancel context.CancelFunc
289 t.mu.Lock()
290 state := entry.state
291 if entry.done || t.active[state.agent] != state {
292 t.mu.Unlock()
293 return
294 }
295 entry.done = true
296 entry.completion = &completion
297 t.flushReadyLocked(state)
298 if len(state.entries) == 0 {
299 delete(t.active, state.agent)
300 cancel = state.waitCancel
301 state.waitCancel = nil
302 state.waitingForProcess = false
303 }
304 t.signalLocked()
305 t.mu.Unlock()
306 if cancel != nil {
307 cancel()
308 }
309}
310
311func (t *CompletionTracker) completeProcess(state *completionAgent) {
312 if state == nil || state.agent == nil {
313 return
314 }
315 snap := state.agent.Snapshot()
316
317 t.mu.Lock()
318 if t.active[state.agent] != state {
319 t.mu.Unlock()
320 return
321 }
322 for _, entry := range state.entries {
323 if entry.done {
324 continue
325 }
326 entry.done = true
327 completion := makeCompletion(entry, snap, "", false)
328 entry.completion = &completion
329 }
330 t.flushReadyLocked(state)
331 delete(t.active, state.agent)
332 state.waitCancel = nil
333 state.waitingForProcess = false
334 t.signalLocked()
335 t.mu.Unlock()
336}
337
338func (t *CompletionTracker) flushReadyLocked(state *completionAgent) {
339 for len(state.entries) != 0 && state.entries[0].done {
340 entry := state.entries[0]
341 state.entries[0] = nil
342 state.entries = state.entries[1:]
343 if entry.completion != nil {
344 t.ready = append(t.ready, *entry.completion)
345 }
346 }
347}
348
349func makeCompletion(entry *completionEntry, snap AgentSnapshot, turnErr string, turnEnd bool) Completion {
350 status := string(snap.Status)
351 if turnEnd {
352 status = "completed"
353 if turnErr != "" {
354 status = "failed"
355 }
356 } else if snap.Status == StatusDone {
357 status = "completed"
358 }
359
360 task := entry.task
361 if !entry.followUp && snap.Task != "" {
362 task = snap.Task
363 }
364 if task == "" {
365 task = snap.Task
366 }
367
368 errText := ""
369 cancellation := false
370 if snap.Result != nil && snap.Result.ShutdownOrigin != "" {
371 if snap.Result.Status == ResultCanceled {
372 status = "cancelled"
373 }
374 if snap.Result.Error != nil {
375 errText = snap.Result.Error.Message
376 cancellation = true
377 }
378 }
379 turnError := false
380 if errText == "" {
381 errText = snap.Err
382 }
383 if errText == "" && turnErr != "" {
384 errText = turnErr
385 turnError = true
386 }
387 return Completion{
388 AgentID: snap.ID,
389 Status: status,
390 Task: task,
391 Error: errText,
392 FinalResponse: snap.LastAssistant,
393 Tail: snap.Tail,
394 turnError: turnError,
395 cancellation: cancellation,
396 registrationSeq: entry.registrationSeq,
397 }
398}
399
400func (t *CompletionTracker) signalLocked() {
401 close(t.changed)
402 t.changed = make(chan struct{})
403}
404
405func (t *CompletionTracker) takeReadyLocked() []Completion {
406 if len(t.ready) == 0 {
407 return nil
408 }
409 batch := append([]Completion(nil), t.ready...)
410 t.ready = nil
411 sort.SliceStable(batch, func(i, j int) bool {
412 return batch[i].registrationSeq < batch[j].registrationSeq
413 })
414 return batch
415}
416
417// Pending reports the number of tracked entries that have not completed.
418func (t *CompletionTracker) Pending() int {
419 if t == nil {
420 return 0
421 }
422 t.mu.Lock()
423 defer t.mu.Unlock()
424 pending := 0
425 for _, state := range t.active {
426 for _, entry := range state.entries {
427 if !entry.done {
428 pending++
429 }
430 }
431 }
432 return pending
433}
434
435// WaitIdle returns all completions buffered so far and all completions that
436// arrive until no tracked entry remains. A cancelled wait leaves active and
437// buffered entries intact for a later caller.
438func (t *CompletionTracker) WaitIdle(ctx context.Context) ([]Completion, error) {
439 if t == nil {
440 return nil, nil
441 }
442 if ctx == nil {
443 ctx = context.Background()
444 }
445
446 for {
447 t.mu.Lock()
448 t.initLocked()
449 if len(t.active) == 0 {
450 batch := t.takeReadyLocked()
451 t.mu.Unlock()
452 return batch, nil
453 }
454 changed := t.changed
455 t.mu.Unlock()
456
457 select {
458 case <-changed:
459 case <-ctx.Done():
460 // Prefer an already-finished batch if completion won the race
461 // with cancellation; otherwise preserve the active work.
462 t.mu.Lock()
463 if len(t.active) == 0 {
464 batch := t.takeReadyLocked()
465 t.mu.Unlock()
466 return batch, nil
467 }
468 t.mu.Unlock()
469 return nil, ctx.Err()
470 }
471 }
472}
473
474// FormatCompletionUpdate formats the stable evidence payload injected into a
475// parent agent. Task and error fields are bounded to 240 characters and tail
476// evidence is bounded to 600 characters; final responses retain their full
477// inline text as they did in interactive mode.
478func FormatCompletionUpdate(batch []Completion, instruction string) string {
479 if len(batch) == 0 {
480 return ""
481 }
482
483 var sb strings.Builder
484 fmt.Fprintf(&sb, "[auto-subagents update] %d sub-agent(s) finished:\n\n", len(batch))
485 for idx, completion := range batch {
486 fmt.Fprintf(&sb, "%d. agent %s - status: %s\n", idx+1, completion.AgentID, completion.Status)
487 fmt.Fprintf(&sb, " task: %s\n", truncateCompletionField(completion.Task, 240))
488 if completion.Error != "" {
489 label := "error"
490 if completion.cancellation {
491 label = "cancellation"
492 } else if completion.turnError {
493 label = "turn error"
494 }
495 fmt.Fprintf(&sb, " %s: %s\n", label, truncateCompletionField(completion.Error, 240))
496 }
497 if completion.FinalResponse != "" {
498 sb.WriteString(" final response:\n")
499 sb.WriteString(completion.FinalResponse)
500 sb.WriteString("\n")
501 } else if tail := strings.TrimSpace(completion.Tail); tail != "" {
502 fmt.Fprintf(&sb, " tail: %s\n", truncateCompletionField(tail, 600))
503 }
504 sb.WriteString("\n")
505 }
506 if instruction != "" {
507 sb.WriteString(instruction)
508 }
509 return sb.String()
510}
511
512func truncateCompletionField(value string, limit int) string {
513 value = strings.TrimSpace(value)
514 if limit <= 0 {
515 return ""
516 }
517 runes := []rune(value)
518 if len(runes) <= limit {
519 return value
520 }
521 if limit <= 3 {
522 return string(runes[:limit])
523 }
524 return string(runes[:limit-3]) + "..."
525}