1package modes
2
3import (
4 "context"
5 "errors"
6 "strings"
7 "testing"
8 "time"
9
10 "github.com/patriceckhart/zot/packages/agent/swarm"
11 "github.com/patriceckhart/zot/packages/core"
12)
13
14func TestAutoSwarmSummaryIncludesCompleteFinalResponseAfterLongTask(t *testing.T) {
15 response := "first response line\n" + strings.Repeat("result ", 150) + "final marker"
16 mgr := swarm.New(swarm.Config{
17 Root: t.TempDir(),
18 RepoRoot: t.TempDir(),
19 NewRunner: func(*swarm.Agent) swarm.Runner {
20 return swarm.RunnerFunc(func(_ context.Context, sink swarm.Sink) error {
21 sink.Transcript(response)
22 return nil
23 })
24 },
25 })
26 a, err := mgr.Spawn(context.Background(), strings.Repeat("long task ", 150))
27 if err != nil {
28 t.Fatal(err)
29 }
30
31 iv := newQueuedAutoSwarmInteractive()
32 iv.TrackSwarmAgent(a, a.Task)
33 update := waitForQueuedPrompt(t, iv)
34 if !strings.Contains(update, "final response:\n"+response) {
35 t.Fatalf("update missing complete final response: %q", update)
36 }
37}
38
39func TestTrackSwarmAgentReportsStartupFailure(t *testing.T) {
40 mgr := swarm.New(swarm.Config{
41 Root: t.TempDir(),
42 RepoRoot: t.TempDir(),
43 NewRunner: func(*swarm.Agent) swarm.Runner {
44 return swarm.RunnerFunc(func(context.Context, swarm.Sink) error {
45 return errors.New("listener startup failed")
46 })
47 },
48 })
49 a, err := mgr.Spawn(context.Background(), "report the date")
50 if err != nil {
51 t.Fatal(err)
52 }
53
54 iv := newQueuedAutoSwarmInteractive()
55 iv.TrackSwarmAgent(a, a.Task)
56
57 update := waitForQueuedPrompt(t, iv)
58 if !strings.Contains(update, "status: failed") {
59 t.Fatalf("update missing failed status: %q", update)
60 }
61 if !strings.Contains(update, "listener startup failed") {
62 t.Fatalf("update missing startup error: %q", update)
63 }
64}
65
66func TestCompleteSwarmWatchReportsTurnOutcomeOnce(t *testing.T) {
67 started := make(chan struct{})
68 mgr := swarm.New(swarm.Config{
69 Root: t.TempDir(),
70 RepoRoot: t.TempDir(),
71 NewRunner: func(*swarm.Agent) swarm.Runner {
72 return swarm.RunnerFunc(func(ctx context.Context, _ swarm.Sink) error {
73 close(started)
74 <-ctx.Done()
75 return ctx.Err()
76 })
77 },
78 })
79 a, err := mgr.Spawn(context.Background(), "report the date")
80 if err != nil {
81 t.Fatal(err)
82 }
83 <-started
84 t.Cleanup(func() {
85 _ = mgr.Stop(a.ID)
86 a.Wait()
87 })
88
89 iv := newQueuedAutoSwarmInteractive()
90 entry := &swarmWatchEntry{agent: a, task: a.Task}
91 iv.swarmWatch = []*swarmWatchEntry{entry}
92 iv.completeSwarmWatchEntry(entry, "completed", "")
93 iv.completeSwarmWatchEntry(entry, "failed", "late error")
94
95 update := waitForQueuedPrompt(t, iv)
96 if !strings.Contains(update, "status: completed") {
97 t.Fatalf("update uses daemon status instead of turn outcome: %q", update)
98 }
99 iv.mu.Lock()
100 queued := len(iv.queued)
101 iv.mu.Unlock()
102 if queued != 1 {
103 t.Fatalf("queued updates = %d; want exactly one", queued)
104 }
105}
106
107func newQueuedAutoSwarmInteractive() *Interactive {
108 return &Interactive{
109 agent: &core.Agent{},
110 busy: true,
111 compacting: true,
112 dirty: make(chan struct{}, 1),
113 }
114}
115
116func waitForQueuedPrompt(t *testing.T, iv *Interactive) string {
117 t.Helper()
118 deadline := time.NewTimer(time.Second)
119 defer deadline.Stop()
120 ticker := time.NewTicker(time.Millisecond)
121 defer ticker.Stop()
122 for {
123 iv.mu.Lock()
124 if len(iv.queued) > 0 {
125 prompt := iv.queued[0]
126 iv.mu.Unlock()
127 return prompt
128 }
129 iv.mu.Unlock()
130 select {
131 case <-deadline.C:
132 t.Fatal("timed out waiting for auto-swarm update")
133 case <-ticker.C:
134 }
135 }
136}