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