1package subagents
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "os"
9 "path/filepath"
10 "testing"
11 "time"
12)
13
14func TestBatchIDsRemainUniqueWithFixedClock(t *testing.T) {
15 now := time.Unix(123, 0)
16 f := New(Config{
17 Root: t.TempDir(), RepoRoot: t.TempDir(), Now: func() time.Time { return now },
18 Policy: SubagentPolicy{MaxConcurrent: 2},
19 NewRunner: func(*Agent) Runner {
20 return RunnerFunc(func(context.Context, Sink) error { return nil })
21 },
22 })
23 t.Cleanup(f.StopAll)
24 first, err := f.SpawnBatch(context.Background(), BatchRequest{Tasks: []string{"one"}})
25 if err != nil {
26 t.Fatal(err)
27 }
28 second, err := f.SpawnBatch(context.Background(), BatchRequest{Tasks: []string{"two"}})
29 if err != nil {
30 t.Fatal(err)
31 }
32 first.Wait()
33 second.Wait()
34 if first.ID == second.ID {
35 t.Fatalf("duplicate batch ids: %q", first.ID)
36 }
37}
38
39func TestBatchReloadsDurableResult(t *testing.T) {
40 root := t.TempDir()
41 newSupervisor := func() *Supervisor {
42 return New(Config{
43 Root: root, RepoRoot: root,
44 NewRunner: func(*Agent) Runner {
45 return RunnerFunc(func(context.Context, Sink) error { return nil })
46 },
47 })
48 }
49 first := newSupervisor()
50 t.Cleanup(first.StopAll)
51 batch, err := first.SpawnBatch(context.Background(), BatchRequest{Context: "review", Tasks: []string{"one", "two"}})
52 if err != nil {
53 t.Fatal(err)
54 }
55 want, err := first.WaitBatch(batch.ID)
56 if err != nil {
57 t.Fatal(err)
58 }
59 if want.Status != BatchSucceeded {
60 t.Fatalf("first batch status = %s", want.Status)
61 }
62 // WaitBatch returns once the aggregate result is available, which can
63 // precede each child's final metadata write. Wait for that finalisation
64 // before another supervisor reloads the same durable state.
65 for _, childID := range batch.ChildIDs {
66 child := first.Get(childID)
67 if child == nil {
68 t.Fatalf("first supervisor is missing child %q", childID)
69 }
70 child.Wait()
71 }
72
73 second := newSupervisor()
74 t.Cleanup(second.StopAll)
75 loaded, errs := second.Reload()
76 if len(errs) != 0 {
77 t.Fatalf("reload errors: %v", errs)
78 }
79 if loaded != 2 {
80 t.Fatalf("reload loaded %d agents, want 2", loaded)
81 }
82 got, err := second.WaitBatch(batch.ID)
83 if err != nil {
84 t.Fatal(err)
85 }
86 if got.Status != BatchSucceeded || len(got.Results) != 2 {
87 t.Fatalf("reloaded batch = %#v", got)
88 }
89}
90
91func TestBatchReloadRejectsUnsafeIDs(t *testing.T) {
92 root := t.TempDir()
93 batchDir := filepath.Join(root, "batches")
94 if err := os.MkdirAll(batchDir, 0o700); err != nil {
95 t.Fatal(err)
96 }
97 data, err := json.Marshal(BatchResult{
98 Version: ProtocolVersion, BatchID: "../escape", TaskCount: 0, Status: BatchRunning,
99 })
100 if err != nil {
101 t.Fatal(err)
102 }
103 if err := os.WriteFile(filepath.Join(batchDir, "queued.json"), data, 0o600); err != nil {
104 t.Fatal(err)
105 }
106 f := New(Config{Root: root, RepoRoot: root})
107 if _, errs := f.Reload(); len(errs) == 0 {
108 t.Fatal("unsafe batch metadata was accepted")
109 }
110 if _, err := os.Stat(filepath.Join(root, "escape.json")); !os.IsNotExist(err) {
111 t.Fatalf("unsafe batch path was written: %v", err)
112 }
113}
114
115func TestBatchSpawnWaitAndCollect(t *testing.T) {
116 root := t.TempDir()
117 f := New(Config{
118 Root: root, RepoRoot: root,
119 Policy: SubagentPolicy{MaxConcurrent: 2, MaxConcurrentPerParent: 2, DefaultTimeout: 0, IdleTimeout: 0},
120 NewRunner: func(*Agent) Runner {
121 return RunnerFunc(func(context.Context, Sink) error { return nil })
122 },
123 })
124 t.Cleanup(f.StopAll)
125 batch, err := f.SpawnBatch(context.Background(), BatchRequest{Context: "review", Tasks: []string{"one", "two"}, MaxConcurrent: 1})
126 if err != nil {
127 t.Fatal(err)
128 }
129 result, err := f.WaitBatch(batch.ID)
130 if err != nil {
131 t.Fatal(err)
132 }
133 if result.Status != BatchSucceeded || len(result.ChildIDs) != 2 || len(result.Results) != 2 {
134 t.Fatalf("batch result = %#v", result)
135 }
136 if batch.Status() != BatchSucceeded {
137 t.Fatalf("batch status = %s", batch.Status())
138 }
139}
140
141func TestBatchResultWaitCancellationDoesNotCancelWorker(t *testing.T) {
142 started := make(chan struct{})
143 release := make(chan struct{})
144 f := New(Config{
145 Root: t.TempDir(), RepoRoot: t.TempDir(),
146 NewRunner: func(*Agent) Runner {
147 return RunnerFunc(func(context.Context, Sink) error {
148 close(started)
149 <-release
150 return nil
151 })
152 },
153 })
154 t.Cleanup(f.StopAll)
155 batch, err := f.SpawnBatch(context.Background(), BatchRequest{Tasks: []string{"one"}})
156 if err != nil {
157 t.Fatal(err)
158 }
159 <-started
160 ctx, cancel := context.WithCancel(context.Background())
161 cancel()
162
163 result, err := f.WaitBatchContext(ctx, batch.ID)
164 if !errors.Is(err, context.Canceled) {
165 t.Fatalf("wait error = %v, want context.Canceled", err)
166 }
167 if result != nil {
168 t.Fatalf("wait result = %#v, want nil", result)
169 }
170 child := f.Get(batch.ChildIDs[0])
171 if child == nil || child.Status() != StatusRunning {
172 t.Fatalf("worker after observer cancellation = %#v, want running", child)
173 }
174 close(release)
175 completed, err := f.WaitBatch(batch.ID)
176 if err != nil {
177 t.Fatal(err)
178 }
179 if completed.Status != BatchSucceeded || completed.Results[child.ID] == nil {
180 t.Fatalf("completed batch = %#v, want successful retained worker result", completed)
181 }
182}
183
184func TestBatchAllowsMoreThanFormerLifetimeLimit(t *testing.T) {
185 const workersBeyondFormerLifetimeLimit = 33
186 tasks := make([]string, workersBeyondFormerLifetimeLimit)
187 for i := range tasks {
188 tasks[i] = fmt.Sprintf("worker-%d", i)
189 }
190 f := New(Config{
191 Root: t.TempDir(), RepoRoot: t.TempDir(),
192 Policy: SubagentPolicy{MaxConcurrent: 1, MaxConcurrentPerParent: 1},
193 NewRunner: func(*Agent) Runner {
194 return RunnerFunc(func(context.Context, Sink) error { return nil })
195 },
196 })
197 t.Cleanup(f.StopAll)
198
199 batch, err := f.SpawnBatch(context.Background(), BatchRequest{Tasks: tasks, MaxConcurrent: 1})
200 if err != nil {
201 t.Fatal(err)
202 }
203 result, err := f.WaitBatch(batch.ID)
204 if err != nil {
205 t.Fatal(err)
206 }
207 if result.Status != BatchSucceeded || len(result.ChildIDs) != len(tasks) || len(result.Results) != len(tasks) {
208 t.Fatalf("batch result = %#v", result)
209 }
210}