1package subagents
2
3import (
4 "context"
5 "errors"
6 "strings"
7 "testing"
8 "time"
9)
10
11func newCompletionTestAgent(id, task string, status Status) *Agent {
12 return &Agent{
13 ID: id,
14 Task: task,
15 status: status,
16 activity: "testing",
17 done: make(chan struct{}),
18 transcript: []string{"tail line 1", "tail line 2"},
19 lastAssistant: "final response",
20 }
21}
22
23func TestCompletionTrackerTurnEndCompletesLongLivedWorker(t *testing.T) {
24 agent := newCompletionTestAgent("worker-1", "original task", StatusRunning)
25 t.Cleanup(func() { close(agent.done) })
26 tracker := NewCompletionTracker()
27 tracker.TrackTurn(agent, "ignored task", false)
28
29 if agent.OnTurnEnd == nil {
30 t.Fatal("TrackTurn did not install an event callback")
31 }
32 agent.OnTurnEnd(1, "")
33
34 batch, err := tracker.WaitIdle(context.Background())
35 if err != nil {
36 t.Fatal(err)
37 }
38 if len(batch) != 1 {
39 t.Fatalf("completion count = %d, want 1", len(batch))
40 }
41 if got := batch[0].Status; got != "completed" {
42 t.Fatalf("status = %q, want completed", got)
43 }
44 if got := batch[0].Task; got != agent.Task {
45 t.Fatalf("task = %q, want %q", got, agent.Task)
46 }
47 if tracker.Pending() != 0 {
48 t.Fatalf("pending = %d, want 0", tracker.Pending())
49 }
50}
51
52func TestCompletionTrackerStartupAndProcessFailuresUseExitPath(t *testing.T) {
53 tests := []struct {
54 name string
55 status Status
56 err error
57 want string
58 }{
59 {name: "startup failure", status: StatusFailed, err: errors.New("listener startup failed"), want: "failed"},
60 {name: "stopped process", status: StatusKilled, err: context.Canceled, want: "killed"},
61 }
62 for _, tt := range tests {
63 t.Run(tt.name, func(t *testing.T) {
64 agent := newCompletionTestAgent("worker-"+tt.name, "task", tt.status)
65 agent.lastErr = tt.err
66 tracker := NewCompletionTracker()
67 tracker.TrackExit(agent, agent.Task)
68 close(agent.done)
69
70 batch, err := tracker.WaitIdle(context.Background())
71 if err != nil {
72 t.Fatal(err)
73 }
74 if len(batch) != 1 || batch[0].Status != tt.want {
75 t.Fatalf("batch = %+v, want one %s completion", batch, tt.want)
76 }
77 if batch[0].Error != tt.err.Error() {
78 t.Fatalf("error = %q, want %q", batch[0].Error, tt.err)
79 }
80 })
81 }
82}
83
84func TestCompletionTrackerQueuesSequentialTurnsForOneWorker(t *testing.T) {
85 agent := newCompletionTestAgent("worker-1", "original task", StatusRunning)
86 tracker := NewCompletionTracker()
87 tracker.TrackTurn(agent, agent.Task, false)
88 tracker.TrackTurn(agent, "follow-up task", true)
89 tracker.TrackTurn(agent, "follow-up task", true)
90
91 if got := tracker.Pending(); got != 3 {
92 t.Fatalf("pending after queued turns = %d, want 3", got)
93 }
94 // Callback goroutines may be scheduled out of order. Completion still has
95 // to match the registered turn identity and flush in registration order.
96 agent.OnTurnEnd(2, "")
97 if got := tracker.Pending(); got != 2 {
98 t.Fatalf("pending after second turn completed first = %d, want 2", got)
99 }
100 // A duplicate lifecycle event for the same prompt-level turn must not
101 // consume the next accepted follow-up, even when its prompt is identical.
102 agent.OnTurnEnd(2, "")
103 if got := tracker.Pending(); got != 2 {
104 t.Fatalf("pending after duplicate second turn = %d, want 2", got)
105 }
106 agent.OnTurnEnd(1, "")
107 if got := tracker.Pending(); got != 1 {
108 t.Fatalf("pending after first turn = %d, want 1", got)
109 }
110 agent.OnTurnEnd(3, "")
111
112 batch, err := tracker.WaitIdle(context.Background())
113 if err != nil || len(batch) != 3 {
114 t.Fatalf("wait = (%+v, %v), want three turn completions", batch, err)
115 }
116 if batch[0].Task != agent.Task || batch[1].Task != "follow-up task" || batch[2].Task != "follow-up task" {
117 t.Fatalf("queued tasks = %q, %q, %q; want original then both follow-ups", batch[0].Task, batch[1].Task, batch[2].Task)
118 }
119 close(agent.done)
120}
121
122func TestCompletionTrackerTurnEndReleasesProcessWaiterWithoutDuplicate(t *testing.T) {
123 agent := newCompletionTestAgent("worker-1", "task", StatusRunning)
124 agent.lastErr = errors.New("late process failure")
125 tracker := NewCompletionTracker()
126 tracker.TrackTurn(agent, agent.Task, false)
127
128 tracker.mu.Lock()
129 state := tracker.active[agent]
130 waitDone := state.waitDone
131 tracker.mu.Unlock()
132 if waitDone == nil {
133 t.Fatal("turn tracking did not start a process fallback waiter")
134 }
135 agent.OnTurnEnd(1, "")
136 select {
137 case <-waitDone:
138 case <-time.After(time.Second):
139 t.Fatal("process fallback waiter was not released after turn_end")
140 }
141 close(agent.done)
142
143 batch, err := tracker.WaitIdle(context.Background())
144 if err != nil || len(batch) != 1 {
145 t.Fatalf("wait = (%+v, %v), want one completion", batch, err)
146 }
147 if batch[0].Status != "completed" {
148 t.Fatalf("status = %q, want completed", batch[0].Status)
149 }
150}
151
152func TestCompletionTrackerSingleWinnerDoesNotDuplicate(t *testing.T) {
153 agent := newCompletionTestAgent("worker-1", "task", StatusRunning)
154 tracker := NewCompletionTracker()
155 tracker.TrackTurn(agent, agent.Task, false)
156
157 agent.OnTurnEnd(1, "")
158 close(agent.done)
159
160 batch, err := tracker.WaitIdle(context.Background())
161 if err != nil {
162 t.Fatal(err)
163 }
164 if len(batch) != 1 {
165 t.Fatalf("completion count = %d, want exactly 1", len(batch))
166 }
167 if second, err := tracker.WaitIdle(context.Background()); err != nil || len(second) != 0 {
168 t.Fatalf("second wait = (%+v, %v), want no duplicate", second, err)
169 }
170}
171
172func TestCompletionTrackerBatchesUntilAllActiveEntriesComplete(t *testing.T) {
173 first := newCompletionTestAgent("worker-1", "first", StatusRunning)
174 second := newCompletionTestAgent("worker-2", "second", StatusRunning)
175 tracker := NewCompletionTracker()
176 tracker.TrackTurn(first, first.Task, false)
177 tracker.TrackTurn(second, second.Task, false)
178
179 waited := make(chan []Completion, 1)
180 go func() {
181 batch, _ := tracker.WaitIdle(context.Background())
182 waited <- batch
183 }()
184 // Complete the workers in reverse order. The batch must still follow the
185 // registration sequence rather than callback arrival order.
186 second.OnTurnEnd(1, "")
187 if got := tracker.Pending(); got != 1 {
188 t.Fatalf("pending after second completion = %d, want 1", got)
189 }
190 first.OnTurnEnd(1, "")
191
192 batch := <-waited
193 if len(batch) != 2 {
194 t.Fatalf("batch count = %d, want 2", len(batch))
195 }
196 if batch[0].AgentID != first.ID || batch[1].AgentID != second.ID {
197 t.Fatalf("batch order = %q, %q; want %q, %q", batch[0].AgentID, batch[1].AgentID, first.ID, second.ID)
198 }
199 close(first.done)
200 close(second.done)
201}
202
203func TestCompletionTrackerReplaysTurnEndReceivedBeforeTracking(t *testing.T) {
204 agent := newCompletionTestAgent("worker-1", "task", StatusRunning)
205 notifyPromptTurnEnd(agent, NewEvent("turn_end", map[string]any{"step": float64(1)}))
206 tracker := NewCompletionTracker()
207 tracker.TrackTurn(agent, agent.Task, false)
208
209 batch, err := tracker.WaitIdle(context.Background())
210 if err != nil || len(batch) != 1 {
211 t.Fatalf("wait = (%+v, %v), want replayed completion", batch, err)
212 }
213 if batch[0].Status != "completed" {
214 t.Fatalf("replayed status = %q, want completed", batch[0].Status)
215 }
216 close(agent.done)
217}
218
219func TestCompletionTrackerFutureTurnIgnoresPriorPendingTurnEnd(t *testing.T) {
220 agent := newCompletionTestAgent("worker-1", "original task", StatusRunning)
221 agent.setTurnCounts(1, 1)
222 agent.setTurnState(TurnIdle, "turn-1")
223 notifyPromptTurnEnd(agent, NewEvent("turn_end", map[string]any{"step": float64(1)}))
224
225 tracker := NewCompletionTracker()
226 tracker.TrackFutureTurn(agent, "follow-up task", true)
227 if got := tracker.Pending(); got != 1 {
228 t.Fatalf("pending after stale turn_end replay = %d, want 1", got)
229 }
230 agent.OnTurnEnd(2, "")
231
232 batch, err := tracker.WaitIdle(context.Background())
233 if err != nil || len(batch) != 1 || batch[0].Task != "follow-up task" {
234 t.Fatalf("future turn completion = (%+v, %v), want follow-up turn", batch, err)
235 }
236 close(agent.done)
237}
238
239func TestCompletionTrackerCurrentTurnReplaysFastCompletion(t *testing.T) {
240 agent := newCompletionTestAgent("worker-1", "original task", StatusRunning)
241 agent.setTurnCounts(2, 1)
242 agent.setTurnState(TurnSucceeded, "turn-2")
243 notifyPromptTurnEnd(agent, NewEvent("turn_end", map[string]any{"step": float64(2)}))
244
245 tracker := NewCompletionTracker()
246 tracker.TrackTurn(agent, "follow-up task", true)
247 batch, err := tracker.WaitIdle(context.Background())
248 if err != nil || len(batch) != 1 || batch[0].Task != "follow-up task" {
249 t.Fatalf("current turn replay = (%+v, %v), want follow-up turn", batch, err)
250 }
251 close(agent.done)
252}
253
254func TestCompletionTrackerPostSuccessRegistrationAcceptsNextStartedTurn(t *testing.T) {
255 agent := newCompletionTestAgent("worker-1", "original task", StatusRunning)
256 agent.setTurnCounts(1, 1)
257 agent.setTurnState(TurnQueued, "turn-1")
258
259 tracker := NewCompletionTracker()
260 tracker.TrackTurn(agent, "follow-up task", true)
261 agent.OnTurnEnd(2, "")
262 batch, err := tracker.WaitIdle(context.Background())
263 if err != nil || len(batch) != 1 || batch[0].Task != "follow-up task" {
264 t.Fatalf("post-success next turn = (%+v, %v), want follow-up turn", batch, err)
265 }
266 close(agent.done)
267}
268
269func TestCompletionTrackerKeepsCompletionsBeforeWaiting(t *testing.T) {
270 agent := newCompletionTestAgent("worker-1", "task", StatusRunning)
271 tracker := NewCompletionTracker()
272 tracker.TrackTurn(agent, agent.Task, false)
273 agent.OnTurnEnd(1, "")
274
275 batch, err := tracker.WaitIdle(context.Background())
276 if err != nil {
277 t.Fatal(err)
278 }
279 if len(batch) != 1 || batch[0].AgentID != agent.ID {
280 t.Fatalf("batch = %+v, want buffered worker completion", batch)
281 }
282 close(agent.done)
283}
284
285func TestCompletionTrackerResetDropsLateAndBufferedCompletions(t *testing.T) {
286 agent := newCompletionTestAgent("worker-1", "task", StatusRunning)
287 t.Cleanup(func() { close(agent.done) })
288 tracker := NewCompletionTracker()
289 tracker.TrackTurn(agent, agent.Task, false)
290 agent.OnTurnEnd(1, "")
291
292 tracker.Reset()
293 batch, err := tracker.WaitIdle(context.Background())
294 if err != nil {
295 t.Fatal(err)
296 }
297 if len(batch) != 0 {
298 t.Fatalf("reset batch = %+v, want no stale completion", batch)
299 }
300 if tracker.Pending() != 0 {
301 t.Fatalf("pending after reset = %d, want 0", tracker.Pending())
302 }
303
304 tracker.TrackTurn(agent, "new task", true)
305 agent.OnTurnEnd(2, "")
306 batch, err = tracker.WaitIdle(context.Background())
307 if err != nil || len(batch) != 1 || batch[0].Task != "new task" {
308 t.Fatalf("new registration completion = (%+v, %v), want one fresh completion", batch, err)
309 }
310}
311
312func TestCompletionTrackerWaitIdleCancellationLeavesActiveWork(t *testing.T) {
313 agent := newCompletionTestAgent("worker-1", "task", StatusRunning)
314 tracker := NewCompletionTracker()
315 tracker.TrackTurn(agent, agent.Task, false)
316 ctx, cancel := context.WithCancel(context.Background())
317 cancel()
318
319 batch, err := tracker.WaitIdle(ctx)
320 if !errors.Is(err, context.Canceled) {
321 t.Fatalf("error = %v, want context.Canceled", err)
322 }
323 if batch != nil {
324 t.Fatalf("cancelled batch = %+v, want nil", batch)
325 }
326 if tracker.Pending() != 1 {
327 t.Fatalf("pending after cancellation = %d, want 1", tracker.Pending())
328 }
329
330 agent.OnTurnEnd(1, "")
331 close(agent.done)
332 batch, err = tracker.WaitIdle(context.Background())
333 if err != nil || len(batch) != 1 {
334 t.Fatalf("post-cancellation wait = (%+v, %v), want one completion", batch, err)
335 }
336}
337
338func TestCompletionTrackerStoresRawFieldsAndResumedTask(t *testing.T) {
339 longTask := strings.Repeat("task ", 80)
340 longError := errors.New(strings.Repeat("error ", 80))
341 agent := newCompletionTestAgent("worker-1", "original task", StatusRunning)
342 agent.Task = longTask
343 agent.lastErr = longError
344 agent.lastAssistant = ""
345 tracker := NewCompletionTracker()
346 const followUp = "follow-up task"
347 tracker.TrackTurn(agent, followUp, true)
348 agent.OnTurnEnd(1, "")
349
350 batch, err := tracker.WaitIdle(context.Background())
351 if err != nil || len(batch) != 1 {
352 t.Fatalf("wait = (%+v, %v), want one completion", batch, err)
353 }
354 completion := batch[0]
355 if completion.Task != followUp {
356 t.Fatalf("resumed task = %q, want %q", completion.Task, followUp)
357 }
358 if completion.Error != longError.Error() {
359 t.Fatalf("stored error was truncated: got %d chars, want %d", len(completion.Error), len(longError.Error()))
360 }
361 if completion.Tail == "" {
362 t.Fatal("stored tail is empty")
363 }
364 close(agent.done)
365}
366
367func TestFormatCompletionUpdateBoundsFieldsAndPreservesEvidence(t *testing.T) {
368 longTask := strings.Repeat("t", 300)
369 longError := strings.Repeat("e", 300)
370 longTail := strings.Repeat("x", 700)
371 instruction := "Summarise the collective outcome without spawning new workers."
372 update := FormatCompletionUpdate([]Completion{{
373 AgentID: "worker-1", Status: "failed", Task: longTask, Error: longError, Tail: longTail,
374 }}, instruction)
375
376 if !strings.HasPrefix(update, "[auto-subagents update] 1 sub-agent(s) finished:\n\n") {
377 t.Fatalf("update has unexpected header: %q", update)
378 }
379 if !strings.Contains(update, "1. agent worker-1 - status: failed\n task: ") {
380 t.Fatalf("update omitted stable labels: %q", update)
381 }
382 if !strings.Contains(update, " error: ") || !strings.Contains(update, " tail: ") {
383 t.Fatalf("update omitted error or tail: %q", update)
384 }
385 if !strings.HasSuffix(update, instruction) {
386 t.Fatalf("update omitted instruction: %q", update)
387 }
388 if strings.Contains(strings.ToLower(update), "poll") {
389 t.Fatalf("update instructs the parent to poll: %q", update)
390 }
391
392 taskLine := ""
393 errorLine := ""
394 for _, line := range strings.Split(update, "\n") {
395 if strings.HasPrefix(line, " task: ") {
396 taskLine = strings.TrimPrefix(line, " task: ")
397 }
398 if strings.HasPrefix(line, " error: ") {
399 errorLine = strings.TrimPrefix(line, " error: ")
400 }
401 }
402 if got := len([]rune(taskLine)); got != 240 {
403 t.Fatalf("formatted task length = %d, want 240", got)
404 }
405 if got := len([]rune(errorLine)); got != 240 {
406 t.Fatalf("formatted error length = %d, want 240", got)
407 }
408 tailLine := ""
409 for _, line := range strings.Split(update, "\n") {
410 if strings.HasPrefix(line, " tail: ") {
411 tailLine = strings.TrimPrefix(line, " tail: ")
412 }
413 }
414 if got := len([]rune(tailLine)); got != 600 {
415 t.Fatalf("formatted tail length = %d, want 600", got)
416 }
417}
418
419func TestFormatCompletionUpdateUsesFinalResponseBeforeTailFallback(t *testing.T) {
420 final := "complete final response\nwith multiple lines"
421 withFinal := FormatCompletionUpdate([]Completion{{
422 AgentID: "worker-1", Status: "completed", FinalResponse: final, Tail: "old tail",
423 }}, "instruction")
424 if !strings.Contains(withFinal, " final response:\n"+final+"\n") {
425 t.Fatalf("final response missing: %q", withFinal)
426 }
427 if strings.Contains(withFinal, "tail:") {
428 t.Fatalf("tail should not be used when final response is present: %q", withFinal)
429 }
430
431 withTail := FormatCompletionUpdate([]Completion{{
432 AgentID: "worker-2", Status: "completed", Tail: " fallback tail ",
433 }}, "instruction")
434 if !strings.Contains(withTail, " tail: fallback tail\n") {
435 t.Fatalf("tail fallback missing or not trimmed: %q", withTail)
436 }
437}