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
132// Reset discards all active and ready completions. Hosts use it when the
133// owning turn is cancelled so a late result cannot be attributed to a later
134// registration for the same worker.
135func (t *CompletionTracker) Reset() {
136 if t == nil {
137 return
138 }
139 var cancels []context.CancelFunc
140 t.mu.Lock()
141 t.initLocked()
142 for _, state := range t.active {
143 for _, entry := range state.entries {
144 entry.done = true
145 }
146 if state.waitCancel != nil {
147 cancels = append(cancels, state.waitCancel)
148 state.waitCancel = nil
149 }
150 state.waitingForProcess = false
151 }
152 t.active = make(map[*Agent]*completionAgent)
153 t.ready = nil
154 t.signalLocked()
155 t.mu.Unlock()
156 for _, cancel := range cancels {
157 cancel()
158 }
159}
160
161func (t *CompletionTracker) addTurn(a *Agent, task string, followUp, future bool) (*completionAgent, *completionEntry, bool) {
162 entry := &completionEntry{task: task, followUp: followUp, expectedStep: expectedTurnStep(a, future), strictStep: future}
163 t.mu.Lock()
164 defer t.mu.Unlock()
165 t.initLocked()
166 state := t.active[a]
167 if state == nil {
168 state = &completionAgent{agent: a, seenTurnEnds: make(map[int]struct{})}
169 t.active[a] = state
170 }
171 if state.seenTurnEnds == nil {
172 state.seenTurnEnds = make(map[int]struct{})
173 }
174 if count := len(state.entries); count != 0 {
175 previousStep := state.entries[count-1].expectedStep
176 if entry.expectedStep <= previousStep {
177 entry.expectedStep = previousStep + 1
178 }
179 }
180 t.nextRegistration++
181 entry.registrationSeq = t.nextRegistration
182 entry.state = state
183 state.entries = append(state.entries, entry)
184 installCallback := !state.callbackInstalled
185 state.callbackInstalled = true
186 return state, entry, installCallback
187}
188
189func (t *CompletionTracker) addExit(a *Agent, task string) (*completionAgent, *completionEntry) {
190 t.mu.Lock()
191 defer t.mu.Unlock()
192 t.initLocked()
193 if state := t.active[a]; state != nil && len(state.entries) != 0 {
194 // A turn watcher already owns this worker. Its process fallback is
195 // sufficient and must not create a duplicate terminal completion.
196 return nil, nil
197 }
198 state := &completionAgent{agent: a}
199 t.nextRegistration++
200 entry := &completionEntry{state: state, task: task, registrationSeq: t.nextRegistration}
201 state.entries = append(state.entries, entry)
202 t.active[a] = state
203 return state, entry
204}
205
206func (t *CompletionTracker) startProcessWaiter(state *completionAgent) {
207 if state == nil || state.agent == nil {
208 return
209 }
210 t.mu.Lock()
211 if t.active[state.agent] != state || state.waitingForProcess {
212 t.mu.Unlock()
213 return
214 }
215 waitCtx, cancel := context.WithCancel(context.Background())
216 state.waitingForProcess = true
217 state.waitCancel = cancel
218 state.waitDone = make(chan struct{})
219 waitDone := state.waitDone
220 t.mu.Unlock()
221
222 go func() {
223 defer close(waitDone)
224 err := state.agent.WaitContext(waitCtx)
225 if err == nil {
226 t.completeProcess(state)
227 }
228 }()
229}
230
231func expectedTurnStep(a *Agent, future bool) int {
232 step := a.LifetimeTurnsValue()
233 if turnID := strings.TrimPrefix(a.CurrentTurnID(), "turn-"); turnID != "" {
234 if current, err := strconv.Atoi(turnID); err == nil && current > step {
235 step = current
236 }
237 }
238 if future || step == 0 {
239 step++
240 }
241 return step
242}
243
244func (t *CompletionTracker) completeNextTurn(state *completionAgent, step int, turnErr string) {
245 if state == nil {
246 return
247 }
248 t.mu.Lock()
249 if t.active[state.agent] != state || len(state.entries) == 0 {
250 t.mu.Unlock()
251 return
252 }
253 if _, duplicate := state.seenTurnEnds[step]; duplicate {
254 t.mu.Unlock()
255 return
256 }
257 state.seenTurnEnds[step] = struct{}{}
258 var entry *completionEntry
259 for _, candidate := range state.entries {
260 if !candidate.done && candidate.expectedStep == step {
261 entry = candidate
262 break
263 }
264 }
265 if entry == nil {
266 // Post-success callers can register after a resumed worker was queued
267 // but before its next turn.start updated metadata. Accept the first
268 // later step for that compatibility path; pre-delivery registrations
269 // use strict identity and never consume a stale or unrelated notice.
270 for _, candidate := range state.entries {
271 if !candidate.done && !candidate.strictStep && candidate.expectedStep < step {
272 entry = candidate
273 candidate.expectedStep = step
274 break
275 }
276 }
277 }
278 t.mu.Unlock()
279 if entry != nil {
280 t.completeTurn(entry, turnErr)
281 }
282}
283
284func (t *CompletionTracker) cancelEntry(entry *completionEntry) {
285 if entry == nil || entry.state == nil {
286 return
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 t.flushReadyLocked(state)
297 if len(state.entries) == 0 {
298 delete(t.active, state.agent)
299 cancel = state.waitCancel
300 state.waitCancel = nil
301 state.waitingForProcess = false
302 }
303 t.signalLocked()
304 t.mu.Unlock()
305 if cancel != nil {
306 cancel()
307 }
308}
309
310func (t *CompletionTracker) completeTurn(entry *completionEntry, turnErr string) {
311 if entry == nil || entry.state == nil || entry.state.agent == nil {
312 return
313 }
314 snap := entry.state.agent.Snapshot()
315 completion := makeCompletion(entry, snap, turnErr, true)
316
317 var cancel context.CancelFunc
318 t.mu.Lock()
319 state := entry.state
320 if entry.done || t.active[state.agent] != state {
321 t.mu.Unlock()
322 return
323 }
324 entry.done = true
325 entry.completion = &completion
326 t.flushReadyLocked(state)
327 if len(state.entries) == 0 {
328 delete(t.active, state.agent)
329 cancel = state.waitCancel
330 state.waitCancel = nil
331 state.waitingForProcess = false
332 }
333 t.signalLocked()
334 t.mu.Unlock()
335 if cancel != nil {
336 cancel()
337 }
338}
339
340func (t *CompletionTracker) completeProcess(state *completionAgent) {
341 if state == nil || state.agent == nil {
342 return
343 }
344 snap := state.agent.Snapshot()
345
346 t.mu.Lock()
347 if t.active[state.agent] != state {
348 t.mu.Unlock()
349 return
350 }
351 for _, entry := range state.entries {
352 if entry.done {
353 continue
354 }
355 entry.done = true
356 completion := makeCompletion(entry, snap, "", false)
357 entry.completion = &completion
358 }
359 t.flushReadyLocked(state)
360 delete(t.active, state.agent)
361 state.waitCancel = nil
362 state.waitingForProcess = false
363 t.signalLocked()
364 t.mu.Unlock()
365}
366
367func (t *CompletionTracker) flushReadyLocked(state *completionAgent) {
368 for len(state.entries) != 0 && state.entries[0].done {
369 entry := state.entries[0]
370 state.entries[0] = nil
371 state.entries = state.entries[1:]
372 if entry.completion != nil {
373 t.ready = append(t.ready, *entry.completion)
374 }
375 }
376}
377
378func makeCompletion(entry *completionEntry, snap AgentSnapshot, turnErr string, turnEnd bool) Completion {
379 status := string(snap.Status)
380 if turnEnd {
381 status = "completed"
382 if turnErr != "" {
383 status = "failed"
384 }
385 } else if snap.Status == StatusDone {
386 status = "completed"
387 }
388
389 task := entry.task
390 if !entry.followUp && snap.Task != "" {
391 task = snap.Task
392 }
393 if task == "" {
394 task = snap.Task
395 }
396
397 errText := ""
398 cancellation := false
399 if snap.Result != nil && snap.Result.ShutdownOrigin != "" {
400 if snap.Result.Status == ResultCanceled {
401 status = "cancelled"
402 }
403 if snap.Result.Error != nil {
404 errText = snap.Result.Error.Message
405 cancellation = true
406 }
407 }
408 turnError := false
409 if errText == "" {
410 errText = snap.Err
411 }
412 if errText == "" && turnErr != "" {
413 errText = turnErr
414 turnError = true
415 }
416 return Completion{
417 AgentID: snap.ID,
418 Status: status,
419 Task: task,
420 Error: errText,
421 FinalResponse: snap.LastAssistant,
422 Tail: snap.Tail,
423 turnError: turnError,
424 cancellation: cancellation,
425 registrationSeq: entry.registrationSeq,
426 }
427}
428
429func (t *CompletionTracker) signalLocked() {
430 close(t.changed)
431 t.changed = make(chan struct{})
432}
433
434func (t *CompletionTracker) takeReadyLocked() []Completion {
435 if len(t.ready) == 0 {
436 return nil
437 }
438 batch := append([]Completion(nil), t.ready...)
439 t.ready = nil
440 sort.SliceStable(batch, func(i, j int) bool {
441 return batch[i].registrationSeq < batch[j].registrationSeq
442 })
443 return batch
444}
445
446// Pending reports the number of tracked entries that have not completed.
447func (t *CompletionTracker) Pending() int {
448 if t == nil {
449 return 0
450 }
451 t.mu.Lock()
452 defer t.mu.Unlock()
453 pending := 0
454 for _, state := range t.active {
455 for _, entry := range state.entries {
456 if !entry.done {
457 pending++
458 }
459 }
460 }
461 return pending
462}
463
464// WaitIdle returns all completions buffered so far and all completions that
465// arrive until no tracked entry remains. A cancelled wait leaves active and
466// buffered entries intact for a later caller.
467func (t *CompletionTracker) WaitIdle(ctx context.Context) ([]Completion, error) {
468 if t == nil {
469 return nil, nil
470 }
471 if ctx == nil {
472 ctx = context.Background()
473 }
474
475 for {
476 t.mu.Lock()
477 t.initLocked()
478 if len(t.active) == 0 {
479 batch := t.takeReadyLocked()
480 t.mu.Unlock()
481 return batch, nil
482 }
483 changed := t.changed
484 t.mu.Unlock()
485
486 select {
487 case <-changed:
488 case <-ctx.Done():
489 // Prefer an already-finished batch if completion won the race
490 // with cancellation; otherwise preserve the active work.
491 t.mu.Lock()
492 if len(t.active) == 0 {
493 batch := t.takeReadyLocked()
494 t.mu.Unlock()
495 return batch, nil
496 }
497 t.mu.Unlock()
498 return nil, ctx.Err()
499 }
500 }
501}
502
503// FormatCompletionUpdate formats the stable evidence payload injected into a
504// parent agent. Task and error fields are bounded to 240 characters and tail
505// evidence is bounded to 600 characters; final responses retain their full
506// inline text as they did in interactive mode.
507func FormatCompletionUpdate(batch []Completion, instruction string) string {
508 if len(batch) == 0 {
509 return ""
510 }
511
512 var sb strings.Builder
513 fmt.Fprintf(&sb, "[auto-subagents update] %d sub-agent(s) finished:\n\n", len(batch))
514 for idx, completion := range batch {
515 fmt.Fprintf(&sb, "%d. agent %s - status: %s\n", idx+1, completion.AgentID, completion.Status)
516 fmt.Fprintf(&sb, " task: %s\n", truncateCompletionField(completion.Task, 240))
517 if completion.Error != "" {
518 label := "error"
519 if completion.cancellation {
520 label = "cancellation"
521 } else if completion.turnError {
522 label = "turn error"
523 }
524 fmt.Fprintf(&sb, " %s: %s\n", label, truncateCompletionField(completion.Error, 240))
525 }
526 if completion.FinalResponse != "" {
527 sb.WriteString(" final response:\n")
528 sb.WriteString(completion.FinalResponse)
529 sb.WriteString("\n")
530 } else if tail := strings.TrimSpace(completion.Tail); tail != "" {
531 fmt.Fprintf(&sb, " tail: %s\n", truncateCompletionField(tail, 600))
532 }
533 sb.WriteString("\n")
534 }
535 if instruction != "" {
536 sb.WriteString(instruction)
537 }
538 return sb.String()
539}
540
541func truncateCompletionField(value string, limit int) string {
542 value = strings.TrimSpace(value)
543 if limit <= 0 {
544 return ""
545 }
546 runes := []rune(value)
547 if len(runes) <= limit {
548 return value
549 }
550 if limit <= 3 {
551 return string(runes[:limit])
552 }
553 return string(runes[:limit-3]) + "..."
554}