1package modes
2
3import (
4 "context"
5 "errors"
6 "strings"
7 "sync"
8 "testing"
9 "time"
10
11 "github.com/bnema/zut/packages/agent/subagents"
12 "github.com/bnema/zut/packages/core"
13 "github.com/bnema/zut/packages/provider"
14)
15
16const autoSubagentsTestTimeout = 10 * time.Second
17
18func TestAutoSubagentsSystemPromptKeepsProfileManifestWhileTogglingDelegationGuidance(t *testing.T) {
19 const onDemand = "delegate only on an explicit user request"
20 iv := &Interactive{
21 agent: &core.Agent{System: "base system"},
22 cfg: InteractiveConfig{
23 AutoSubagentsSystemAddendum: "auto-subagents guidance",
24 OnDemandSubagentsSystemAddendum: onDemand,
25 SubagentsSystemAddendum: "[subagents_list]\n- reviewer\n[/subagents_list]",
26 Supervisor: &subagents.Supervisor{},
27 },
28 }
29
30 iv.applyAutoSubagentsSystemPrompt(true)
31 if !strings.Contains(iv.agent.System, "[subagents_list]") || !strings.Contains(iv.agent.System, "auto-subagents guidance") {
32 t.Fatalf("enabled system prompt missing subagent blocks: %q", iv.agent.System)
33 }
34 before := iv.agent.System
35 iv.applyAutoSubagentsSystemPrompt(true)
36 if iv.agent.System != before {
37 t.Fatalf("enabling twice changed system prompt: %q -> %q", before, iv.agent.System)
38 }
39
40 iv.applyAutoSubagentsSystemPrompt(false)
41 if !strings.Contains(iv.agent.System, "[subagents_list]") {
42 t.Fatalf("disabled system prompt omitted profile manifest: %q", iv.agent.System)
43 }
44 if strings.Contains(iv.agent.System, "auto-subagents guidance") {
45 t.Fatalf("disabled system prompt retained orchestrator guidance: %q", iv.agent.System)
46 }
47 if !strings.Contains(iv.agent.System, onDemand) {
48 t.Fatalf("disabled system prompt omitted on-demand delegation guidance: %q", iv.agent.System)
49 }
50}
51
52func TestAutoSubagentsTogglePreservesMatchingBasePromptBlocks(t *testing.T) {
53 manifest := "[subagents_list]\n- reviewer\n[/subagents_list]"
54 guidance := "auto-subagents guidance"
55 base := "base system\n\n" + manifest + "\n\n" + guidance
56 iv := &Interactive{
57 agent: &core.Agent{System: base},
58 cfg: InteractiveConfig{
59 AutoSubagentsSystemAddendum: guidance,
60 SubagentsSystemAddendum: manifest,
61 Supervisor: &subagents.Supervisor{},
62 },
63 }
64
65 iv.applyAutoSubagentsSystemPrompt(true)
66 if strings.Count(iv.agent.System, manifest) != 2 || strings.Count(iv.agent.System, guidance) != 2 {
67 t.Fatalf("enabling should append owned duplicate blocks: %q", iv.agent.System)
68 }
69 iv.applyAutoSubagentsSystemPrompt(false)
70 want := base + "\n\n" + manifest
71 if got := strings.TrimSpace(iv.agent.System); got != want {
72 t.Fatalf("disabling changed matching base prompt blocks: %q; want %q", got, want)
73 }
74}
75
76func TestAutoSubagentsSummaryIncludesCompleteFinalResponseAfterLongTask(t *testing.T) {
77 response := "first response line\n" + strings.Repeat("result ", 150) + "final marker"
78 mgr := subagents.New(subagents.Config{
79 Root: t.TempDir(),
80 RepoRoot: t.TempDir(),
81 NewRunner: func(*subagents.Agent) subagents.Runner {
82 return subagents.RunnerFunc(func(_ context.Context, sink subagents.Sink) error {
83 sink.Transcript(response)
84 return nil
85 })
86 },
87 })
88 a, err := mgr.Spawn(context.Background(), strings.Repeat("long task ", 150))
89 if err != nil {
90 t.Fatal(err)
91 }
92
93 iv := newQueuedAutoSubagentsInteractive()
94 iv.TrackSubagentWorker(a, a.Task, false)
95 update := waitForQueuedPrompt(t, iv)
96 if !strings.Contains(update, "final response:\n"+response) {
97 t.Fatalf("update missing complete final response: %q", update)
98 }
99}
100
101func TestTrackSubagentWorkerReportsStartupFailure(t *testing.T) {
102 mgr := subagents.New(subagents.Config{
103 Root: t.TempDir(),
104 RepoRoot: t.TempDir(),
105 NewRunner: func(*subagents.Agent) subagents.Runner {
106 return subagents.RunnerFunc(func(context.Context, subagents.Sink) error {
107 return errors.New("listener startup failed")
108 })
109 },
110 })
111 a, err := mgr.Spawn(context.Background(), "report the date")
112 if err != nil {
113 t.Fatal(err)
114 }
115
116 iv := newQueuedAutoSubagentsInteractive()
117 iv.TrackSubagentWorker(a, a.Task, false)
118
119 update := waitForQueuedPrompt(t, iv)
120 if !strings.Contains(update, "status: failed") {
121 t.Fatalf("update missing failed status: %q", update)
122 }
123 if !strings.Contains(update, "listener startup failed") {
124 t.Fatalf("update missing startup error: %q", update)
125 }
126}
127
128func TestTrackStoppedSubagentWorkerDeliversTerminalUpdate(t *testing.T) {
129 started := make(chan struct{})
130 mgr := subagents.New(subagents.Config{
131 Root: t.TempDir(),
132 RepoRoot: t.TempDir(),
133 NewRunner: func(*subagents.Agent) subagents.Runner {
134 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
135 close(started)
136 <-ctx.Done()
137 return ctx.Err()
138 })
139 },
140 })
141 a, err := mgr.Spawn(context.Background(), "investigate the stuck worker")
142 if err != nil {
143 t.Fatal(err)
144 }
145 <-started
146 t.Cleanup(func() {
147 _ = mgr.Stop(a.ID)
148 a.Wait()
149 })
150
151 iv := newQueuedAutoSubagentsInteractive()
152 iv.TrackStoppedSubagentWorker(a)
153 if err := mgr.Stop(a.ID); err != nil {
154 t.Fatal(err)
155 }
156
157 update := waitForQueuedPrompt(t, iv)
158 if !strings.Contains(update, "task: investigate the stuck worker") {
159 t.Fatalf("stop update missing task: %q", update)
160 }
161 if !strings.Contains(update, "status: cancelled") {
162 t.Fatalf("stop update missing cancellation status: %q", update)
163 }
164 if !strings.Contains(update, "subagent stopped by request") || strings.Contains(update, "context canceled") {
165 t.Fatalf("stop update missing attributed cancellation: %q", update)
166 }
167}
168
169func TestTrackDeadlineFailureUsesAttributedResultError(t *testing.T) {
170 mgr := subagents.New(subagents.Config{
171 Root: t.TempDir(),
172 RepoRoot: t.TempDir(),
173 NewRunner: func(*subagents.Agent) subagents.Runner {
174 return subagents.RunnerFunc(func(context.Context, subagents.Sink) error {
175 return context.DeadlineExceeded
176 })
177 },
178 })
179 a, err := mgr.Spawn(context.Background(), "run until deadline")
180 if err != nil {
181 t.Fatal(err)
182 }
183
184 iv := newQueuedAutoSubagentsInteractive()
185 iv.TrackSubagentWorker(a, a.Task, false)
186 update := waitForQueuedPrompt(t, iv)
187 if !strings.Contains(update, "status: failed") {
188 t.Fatalf("deadline update changed failed status: %q", update)
189 }
190 want := subagents.ShutdownResultError(subagents.ShutdownOriginDeadline).Message
191 if !strings.Contains(update, want) {
192 t.Fatalf("deadline update missing attributed result error: %q", update)
193 }
194}
195
196func TestTrackResumedSubagentWorkerDeliversFollowUp(t *testing.T) {
197 started := make(chan struct{})
198 mgr := subagents.New(subagents.Config{
199 Root: t.TempDir(),
200 RepoRoot: t.TempDir(),
201 NewRunner: func(*subagents.Agent) subagents.Runner {
202 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
203 close(started)
204 <-ctx.Done()
205 return ctx.Err()
206 })
207 },
208 })
209 a, err := mgr.Spawn(context.Background(), "review the implementation")
210 if err != nil {
211 t.Fatal(err)
212 }
213 <-started
214 t.Cleanup(func() {
215 _ = mgr.Stop(a.ID)
216 a.Wait()
217 })
218
219 iv := newQueuedAutoSubagentsInteractive()
220 const followUp = "I applied your review. What do you think now?"
221 iv.TrackResumedSubagentWorker(a, followUp, false)
222 if a.OnTurnEnd == nil {
223 t.Fatal("resumed worker has no completion watcher")
224 }
225 a.OnTurnEnd(1, "")
226
227 update := waitForQueuedPrompt(t, iv)
228 if !strings.Contains(update, "task: "+followUp) {
229 t.Fatalf("follow-up update uses the original task instead of the new prompt: %q", update)
230 }
231 if !strings.Contains(update, "status: completed") {
232 t.Fatalf("follow-up update missing completed status: %q", update)
233 }
234}
235
236func TestTrackResumedSubagentWorkerQueuesActiveFollowUpOnSameWorker(t *testing.T) {
237 started := make(chan struct{})
238 mgr := subagents.New(subagents.Config{
239 Root: t.TempDir(),
240 RepoRoot: t.TempDir(),
241 NewRunner: func(*subagents.Agent) subagents.Runner {
242 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
243 select {
244 case started <- struct{}{}:
245 default:
246 }
247 <-ctx.Done()
248 return ctx.Err()
249 })
250 },
251 })
252 first, err := mgr.Spawn(context.Background(), "original task")
253 if err != nil {
254 t.Fatal(err)
255 }
256 <-started
257 t.Cleanup(func() {
258 _ = mgr.Stop(first.ID)
259 first.Wait()
260 })
261
262 iv := newQueuedAutoSubagentsInteractive()
263 iv.TrackSubagentWorker(first, first.Task, false)
264 const followUp = "queued follow-up"
265 cancel := iv.PrepareResumedSubagentWorker(first, followUp, false)
266 if cancel == nil {
267 t.Fatal("resume registration returned nil cleanup")
268 }
269 if first.OnTurnEnd == nil {
270 t.Fatal("worker has no completion watcher")
271 }
272 first.OnTurnEnd(1, "")
273 first.OnTurnEnd(2, "")
274
275 update := waitForQueuedPrompt(t, iv)
276 if !strings.Contains(update, "[auto-subagents update] 2 sub-agent(s) finished:") {
277 t.Fatalf("active follow-up was not batched with the original turn: %q", update)
278 }
279 if !strings.Contains(update, "task: "+followUp) {
280 t.Fatalf("active follow-up task missing: %q", update)
281 }
282}
283
284func TestCompletionDeliveryHoldBatchesFastAndLateRegistrations(t *testing.T) {
285 mgr := subagents.New(subagents.Config{
286 Root: t.TempDir(),
287 RepoRoot: t.TempDir(),
288 NewRunner: func(*subagents.Agent) subagents.Runner {
289 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
290 <-ctx.Done()
291 return ctx.Err()
292 })
293 },
294 })
295 first, err := mgr.Spawn(context.Background(), "first task")
296 if err != nil {
297 t.Fatal(err)
298 }
299 second, err := mgr.Spawn(context.Background(), "second task")
300 if err != nil {
301 t.Fatal(err)
302 }
303 t.Cleanup(func() {
304 _ = mgr.Stop(first.ID)
305 _ = mgr.Stop(second.ID)
306 first.Wait()
307 second.Wait()
308 })
309
310 iv := newQueuedAutoSubagentsInteractive()
311 release := iv.beginCompletionDeliveryHold()
312 iv.TrackSubagentWorker(first, first.Task, false)
313 first.OnTurnEnd(1, "")
314 iv.TrackSubagentWorker(second, second.Task, false)
315 second.OnTurnEnd(1, "")
316 release()
317
318 update := waitForQueuedPrompt(t, iv)
319 if !strings.Contains(update, "[auto-subagents update] 2 sub-agent(s) finished:") {
320 t.Fatalf("fast and late registrations were not batched: %q", update)
321 }
322 if !strings.Contains(update, "task: first task") || !strings.Contains(update, "task: second task") {
323 t.Fatalf("batched update omitted one task: %q", update)
324 }
325 iv.mu.Lock()
326 queued := len(iv.queued)
327 iv.mu.Unlock()
328 if queued != 1 {
329 t.Fatalf("queued updates = %d; want exactly one batched update", queued)
330 }
331}
332
333func TestCompleteSupervisorWatchReportsTurnOutcomeOnce(t *testing.T) {
334 started := make(chan struct{})
335 mgr := subagents.New(subagents.Config{
336 Root: t.TempDir(),
337 RepoRoot: t.TempDir(),
338 NewRunner: func(*subagents.Agent) subagents.Runner {
339 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
340 close(started)
341 <-ctx.Done()
342 return ctx.Err()
343 })
344 },
345 })
346 a, err := mgr.Spawn(context.Background(), "report the date")
347 if err != nil {
348 t.Fatal(err)
349 }
350 <-started
351 t.Cleanup(func() {
352 _ = mgr.Stop(a.ID)
353 a.Wait()
354 })
355
356 iv := newQueuedAutoSubagentsInteractive()
357 iv.TrackSubagentWorker(a, a.Task, false)
358 if a.OnTurnEnd == nil {
359 t.Fatal("worker has no completion watcher")
360 }
361 a.OnTurnEnd(1, "")
362
363 update := waitForQueuedPrompt(t, iv)
364 if !strings.Contains(update, "status: completed") {
365 t.Fatalf("update uses daemon status instead of turn outcome: %q", update)
366 }
367 iv.mu.Lock()
368 queued := len(iv.queued)
369 iv.mu.Unlock()
370 if queued != 1 {
371 t.Fatalf("queued updates = %d; want exactly one", queued)
372 }
373}
374
375type autoSubagentsQueueClient struct {
376 requests chan provider.Request
377}
378
379func (c *autoSubagentsQueueClient) Name() string { return "auto-subagents-queue-test" }
380
381func (c *autoSubagentsQueueClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
382 c.requests <- req
383 out := make(chan provider.Event, 1)
384 go func() {
385 defer close(out)
386 out <- provider.EventDone{
387 Stop: provider.StopEnd,
388 Message: provider.Message{
389 Role: provider.RoleAssistant,
390 Content: []provider.Content{provider.TextBlock{Text: "done"}},
391 },
392 }
393 }()
394 return out, nil
395}
396
397func TestAutoSubagentsUpdateQueuedAtTurnEndStartsFollowUp(t *testing.T) {
398 client := &autoSubagentsQueueClient{requests: make(chan provider.Request, 2)}
399 agent := core.NewAgent(client, "test-model", "", nil)
400 interactive := NewInteractive(InteractiveConfig{Agent: agent})
401 interactive.runCtx = context.Background()
402
403 var once sync.Once
404 agent.OnEvent = func(ev core.AgentEvent) {
405 if _, ok := ev.(core.EvDone); !ok {
406 return
407 }
408 once.Do(func() {
409 instruction := "Briefly summarise the collective outcome for the user. Reference the agents by id. If any failed, suggest a follow-up; otherwise confirm completion. Do not spawn new sub-agents unless the user asks."
410 update := subagents.FormatCompletionUpdate([]subagents.Completion{{Task: "smoke test"}}, instruction)
411 interactive.SubmitOrQueue(update, nil)
412 })
413 }
414
415 interactive.startTurn(context.Background(), "initial")
416 select {
417 case <-client.requests:
418 case <-time.After(2 * time.Second):
419 t.Fatal("initial request did not start")
420 }
421
422 var followUp provider.Request
423 select {
424 case followUp = <-client.requests:
425 case <-time.After(2 * time.Second):
426 t.Fatal("queued auto-subagents update did not start a follow-up request")
427 }
428 foundUpdate := false
429 for _, message := range followUp.Messages {
430 for _, content := range message.Content {
431 if text, ok := content.(provider.TextBlock); ok && strings.Contains(text.Text, "[auto-subagents update]") {
432 foundUpdate = true
433 }
434 }
435 }
436 if !foundUpdate {
437 t.Fatalf("follow-up request did not contain the auto-subagents update: %#v", followUp.Messages)
438 }
439
440 deadline := time.NewTimer(2 * time.Second)
441 defer deadline.Stop()
442 for {
443 interactive.mu.Lock()
444 busy := interactive.busy
445 interactive.mu.Unlock()
446 if !busy && agent.QueuedMessageCount() == 0 {
447 return
448 }
449 select {
450 case <-deadline.C:
451 t.Fatal("interactive turn remained busy or retained a queued update")
452 case <-time.After(time.Millisecond):
453 }
454 }
455}
456
457func newQueuedAutoSubagentsInteractive() *Interactive {
458 return &Interactive{
459 agent: &core.Agent{},
460 busy: true,
461 compacting: true,
462 dirty: make(chan struct{}, 1),
463 }
464}
465
466func waitForQueuedPrompt(t *testing.T, iv *Interactive) string {
467 t.Helper()
468 deadline := time.NewTimer(autoSubagentsTestTimeout)
469 defer deadline.Stop()
470 ticker := time.NewTicker(time.Millisecond)
471 defer ticker.Stop()
472 for {
473 iv.mu.Lock()
474 if len(iv.queued) > 0 {
475 prompt := iv.queued[0]
476 iv.mu.Unlock()
477 return prompt
478 }
479 iv.mu.Unlock()
480 select {
481 case <-deadline.C:
482 t.Fatal("timed out waiting for auto-subagents update")
483 case <-ticker.C:
484 }
485 }
486}