1package subagents
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "io"
9 "os"
10 "path/filepath"
11 "strings"
12 "sync"
13 "time"
14)
15
16// BatchRequest starts independent child tasks under one aggregate id. Batch
17// scheduling is deliberately shallow; dependency graphs belong in workflows.
18type BatchRequest struct {
19 Context string
20 Tasks []string
21 MaxConcurrent int
22 Spawn SpawnRequest
23}
24
25type BatchStatus string
26
27const (
28 BatchQueued BatchStatus = "queued"
29 BatchRunning BatchStatus = "running"
30 BatchSucceeded BatchStatus = "succeeded"
31 BatchFailed BatchStatus = "failed"
32 BatchCanceled BatchStatus = "canceled"
33)
34
35type BatchResult struct {
36 Version int `json:"version"`
37 BatchID string `json:"batch_id"`
38 TaskCount int `json:"task_count,omitempty"`
39 Status BatchStatus `json:"status"`
40 ChildIDs []string `json:"child_ids"`
41 Results map[string]*TurnResult `json:"results,omitempty"`
42 Error string `json:"error,omitempty"`
43 CreatedAt time.Time `json:"created_at"`
44 UpdatedAt time.Time `json:"updated_at"`
45}
46
47const maxBatchFileBytes = 32 * 1024 * 1024
48
49type Batch struct {
50 ID string
51 Context string
52 ChildIDs []string
53 MaxConcurrent int
54 TaskCount int
55 CreatedAt time.Time
56
57 mu sync.Mutex
58 status BatchStatus
59 result *BatchResult
60 done chan struct{}
61 once sync.Once
62}
63
64func (b *Batch) Status() BatchStatus {
65 b.mu.Lock()
66 defer b.mu.Unlock()
67 return b.status
68}
69
70func (b *Batch) Result() *BatchResult {
71 b.mu.Lock()
72 defer b.mu.Unlock()
73 return cloneBatchResult(b.result)
74}
75
76func (b *Batch) Wait() { <-b.done }
77func (b *Batch) closeDone() { b.once.Do(func() { close(b.done) }) }
78
79func (f *Supervisor) SpawnBatch(ctx context.Context, req BatchRequest) (*Batch, error) {
80 if len(req.Tasks) == 0 {
81 return nil, fmt.Errorf("subagents: batch requires at least one task")
82 }
83 if ctx == nil {
84 ctx = context.Background()
85 }
86 now := f.cfg.Now()
87 baseID := newAgentID("batch "+now.Format("150405.000"), now)
88 f.mu.Lock()
89 if f.batches == nil {
90 f.batches = make(map[string]*Batch)
91 }
92 batchID := baseID
93 for suffix := 2; ; suffix++ {
94 if _, exists := f.batches[batchID]; !exists {
95 break
96 }
97 batchID = fmt.Sprintf("%s-%d", baseID, suffix)
98 }
99 batch := &Batch{
100 ID: batchID,
101 Context: strings.TrimSpace(req.Context),
102 MaxConcurrent: req.MaxConcurrent,
103 TaskCount: len(req.Tasks),
104 CreatedAt: now,
105 status: BatchQueued,
106 result: &BatchResult{
107 Version: ProtocolVersion, BatchID: batchID, TaskCount: len(req.Tasks), Status: BatchQueued,
108 CreatedAt: now, UpdatedAt: now,
109 },
110 done: make(chan struct{}),
111 }
112 f.batches[batch.ID] = batch
113 f.mu.Unlock()
114 // Write a queued manifest before admitting children so a crash during
115 // batch admission still leaves an aggregate id to reload and reconcile.
116 f.persistBatch(batch)
117
118 for _, task := range req.Tasks {
119 spawnReq := req.Spawn
120 spawnReq.Task = task
121 spawnReq.ParentID = batch.ID
122 spawnReq.BatchID = batch.ID
123 a, err := f.SpawnReq(ctx, spawnReq)
124 if err != nil {
125 childIDs := append([]string(nil), batch.ChildIDs...)
126 for _, childID := range childIDs {
127 _ = f.Stop(childID)
128 }
129 for _, childID := range childIDs {
130 if child := f.Get(childID); child != nil {
131 child.Wait()
132 }
133 }
134 batch.mu.Lock()
135 batch.status = BatchFailed
136 batch.result = &BatchResult{
137 Version: ProtocolVersion, BatchID: batch.ID, TaskCount: batch.TaskCount, Status: BatchFailed,
138 ChildIDs: childIDs, Error: err.Error(), CreatedAt: now, UpdatedAt: f.cfg.Now(),
139 }
140 batch.mu.Unlock()
141 f.persistBatch(batch)
142 batch.closeDone()
143 return batch, err
144 }
145 batch.ChildIDs = append(batch.ChildIDs, a.ID)
146 batch.mu.Lock()
147 batch.result = &BatchResult{
148 Version: ProtocolVersion, BatchID: batch.ID, TaskCount: batch.TaskCount, Status: BatchQueued,
149 ChildIDs: append([]string(nil), batch.ChildIDs...), CreatedAt: batch.CreatedAt,
150 UpdatedAt: f.cfg.Now(),
151 }
152 batch.mu.Unlock()
153 f.persistBatch(batch)
154 }
155 batch.mu.Lock()
156 batch.status = BatchRunning
157 batch.mu.Unlock()
158 f.persistBatch(batch)
159 go f.collectBatch(batch)
160 return batch, nil
161}
162
163func (f *Supervisor) collectBatch(batch *Batch) {
164 if batch.TaskCount > 0 && len(batch.ChildIDs) < batch.TaskCount {
165 batch.mu.Lock()
166 batch.status = BatchFailed
167 batch.result = &BatchResult{
168 Version: ProtocolVersion, BatchID: batch.ID, TaskCount: batch.TaskCount,
169 Status: BatchFailed, ChildIDs: append([]string(nil), batch.ChildIDs...),
170 Error: "batch admission was incomplete", CreatedAt: batch.CreatedAt, UpdatedAt: f.cfg.Now(),
171 }
172 batch.mu.Unlock()
173 f.persistBatch(batch)
174 batch.closeDone()
175 return
176 }
177 results := make(map[string]*TurnResult, len(batch.ChildIDs))
178 status := BatchSucceeded
179 var firstErr string
180 for _, id := range batch.ChildIDs {
181 a := f.Get(id)
182 if a == nil {
183 status = BatchFailed
184 if firstErr == "" {
185 firstErr = "missing child " + id
186 }
187 continue
188 }
189 result, err := f.waitForBatchResult(a)
190 if err == nil {
191 results[id] = result
192 switch result.Status {
193 case ResultFailed:
194 if status == BatchSucceeded {
195 status = BatchFailed
196 if result.Error != nil {
197 firstErr = result.Error.Message
198 }
199 }
200 case ResultCanceled:
201 if status == BatchSucceeded {
202 status = BatchCanceled
203 firstErr = "one or more batch children were canceled"
204 }
205 }
206 } else {
207 status = BatchFailed
208 if firstErr == "" {
209 firstErr = err.Error()
210 }
211 }
212 }
213 batch.mu.Lock()
214 batch.status = status
215 batch.result = &BatchResult{
216 Version: ProtocolVersion, BatchID: batch.ID, TaskCount: batch.TaskCount, Status: status,
217 ChildIDs: append([]string(nil), batch.ChildIDs...), Results: results,
218 Error: firstErr, CreatedAt: batch.CreatedAt, UpdatedAt: f.cfg.Now(),
219 }
220 batch.mu.Unlock()
221 f.persistBatch(batch)
222 batch.closeDone()
223}
224
225func (f *Supervisor) waitForBatchResult(a *Agent) (*TurnResult, error) {
226 if a == nil {
227 return nil, errors.New("subagents: missing batch child")
228 }
229 if result := a.Result(); result != nil {
230 return result, nil
231 }
232 if a.Status() != StatusDetached {
233 return a.waitForTurnResult(context.Background())
234 }
235
236 // A reloaded worker has a closed supervisor-side done channel even when
237 // its orphan process is still alive. Reconcile durable result.json/events
238 // before treating that closed channel as a failed child.
239 readDurable := func() *TurnResult {
240 if result, err := readTurnResult(a.stateDirectory(f.cfg.Root)); err == nil && validateTurnResultAgent(result, a.ID) == nil {
241 result = result.Bounded(f.cfg.Policy.MaxOutputBytes, f.cfg.Policy.MaxOutputLines)
242 a.setResult(result)
243 return result
244 }
245 events, err := ReadEventLog(a.EventLogPath)
246 if err != nil {
247 return nil
248 }
249 for _, event := range events {
250 if event.Type != EventTurnResult && event.Type != "turn_result" {
251 continue
252 }
253 result, decodeErr := decodeTurnResultEvent(event, a.ID, f.cfg.Policy.MaxOutputBytes, f.cfg.Policy.MaxOutputLines)
254 if decodeErr != nil {
255 continue
256 }
257 a.setResult(result)
258 return result
259 }
260 return nil
261 }
262
263 if result := readDurable(); result != nil {
264 return result, nil
265 }
266 if !inboxLive(a.InboxPath) {
267 return nil, fmt.Errorf("subagents: agent %s exited without a turn result", a.ID)
268 }
269
270 interval := 100 * time.Millisecond
271 deadline := time.NewTimer(f.cfg.Policy.DefaultTimeout)
272 defer deadline.Stop()
273 ticker := time.NewTicker(interval)
274 defer ticker.Stop()
275 for {
276 select {
277 case <-ticker.C:
278 if result := readDurable(); result != nil {
279 return result, nil
280 }
281 if !inboxLive(a.InboxPath) {
282 return nil, fmt.Errorf("subagents: agent %s exited without a turn result", a.ID)
283 }
284 case <-deadline.C:
285 return nil, fmt.Errorf("subagents: agent %s batch result timed out", a.ID)
286 }
287 }
288}
289
290func (f *Supervisor) WaitBatch(id string) (*BatchResult, error) {
291 f.mu.Lock()
292 batch := f.batches[id]
293 f.mu.Unlock()
294 if batch == nil {
295 return nil, fmt.Errorf("subagents: no such batch %q", id)
296 }
297 batch.Wait()
298 return batch.Result(), nil
299}
300
301func (f *Supervisor) GetBatch(id string) *Batch {
302 f.mu.Lock()
303 defer f.mu.Unlock()
304 return f.batches[id]
305}
306
307func cloneBatchResult(result *BatchResult) *BatchResult {
308 if result == nil {
309 return nil
310 }
311 copy := *result
312 copy.ChildIDs = append([]string(nil), result.ChildIDs...)
313 if result.Results != nil {
314 copy.Results = make(map[string]*TurnResult, len(result.Results))
315 for id, child := range result.Results {
316 copy.Results[id] = cloneTurnResult(child)
317 }
318 }
319 return ©
320}
321
322func (f *Supervisor) reloadBatches(root string) []error {
323 entries, err := os.ReadDir(filepath.Join(root, "batches"))
324 if err != nil {
325 if errors.Is(err, os.ErrNotExist) {
326 return nil
327 }
328 return []error{fmt.Errorf("subagents reload batches %s: %w", root, err)}
329 }
330 var errs []error
331 for _, entry := range entries {
332 if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
333 continue
334 }
335 fileID := strings.TrimSuffix(entry.Name(), ".json")
336 if !safeAgentID(fileID) {
337 errs = append(errs, fmt.Errorf("subagents reload batch %s: unsafe batch id", entry.Name()))
338 continue
339 }
340 path := filepath.Join(root, "batches", entry.Name())
341 file, openErr := os.Open(path)
342 if openErr != nil {
343 errs = append(errs, fmt.Errorf("subagents reload batch %s: %w", path, openErr))
344 continue
345 }
346 data, readErr := io.ReadAll(io.LimitReader(file, maxBatchFileBytes+1))
347 _ = file.Close()
348 if readErr != nil {
349 errs = append(errs, fmt.Errorf("subagents reload batch %s: %w", path, readErr))
350 continue
351 }
352 if len(data) > maxBatchFileBytes {
353 errs = append(errs, fmt.Errorf("subagents reload batch %s: exceeds %d bytes", path, maxBatchFileBytes))
354 continue
355 }
356 var result BatchResult
357 if unmarshalErr := json.Unmarshal(data, &result); unmarshalErr != nil {
358 errs = append(errs, fmt.Errorf("subagents reload batch %s: %w", path, unmarshalErr))
359 continue
360 }
361 if result.Version == 0 {
362 result.Version = ProtocolVersion
363 }
364 if result.Version != ProtocolVersion || result.BatchID != fileID || !safeAgentID(result.BatchID) || result.TaskCount < 0 || len(result.ChildIDs) > result.TaskCount {
365 errs = append(errs, fmt.Errorf("subagents reload batch %s: invalid result metadata", path))
366 continue
367 }
368 batch := &Batch{
369 ID: result.BatchID,
370 ChildIDs: append([]string(nil), result.ChildIDs...),
371 TaskCount: result.TaskCount,
372 CreatedAt: result.CreatedAt,
373 status: result.Status,
374 result: cloneBatchResult(&result),
375 done: make(chan struct{}),
376 }
377 if batch.CreatedAt.IsZero() {
378 batch.CreatedAt = result.UpdatedAt
379 }
380 f.mu.Lock()
381 _, exists := f.batches[batch.ID]
382 if !exists {
383 f.batches[batch.ID] = batch
384 }
385 f.mu.Unlock()
386 if exists {
387 continue
388 }
389 switch batch.status {
390 case BatchSucceeded, BatchFailed, BatchCanceled:
391 batch.closeDone()
392 default:
393 go f.collectBatch(batch)
394 }
395 }
396 return errs
397}
398
399func (f *Supervisor) persistBatch(batch *Batch) {
400 if batch == nil || !safeAgentID(batch.ID) {
401 return
402 }
403 result := batch.Result()
404 if result == nil {
405 return
406 }
407 data, err := json.MarshalIndent(result, "", " ")
408 if err != nil || len(data) > maxBatchFileBytes {
409 return
410 }
411 dir := filepath.Join(f.cfg.Root, "batches")
412 if err := os.MkdirAll(dir, 0o700); err != nil {
413 return
414 }
415 path := filepath.Join(dir, batch.ID+".json")
416 tmp := path + ".tmp"
417 if err := os.WriteFile(tmp, append(data, '\n'), 0o600); err != nil {
418 return
419 }
420 if err := os.Rename(tmp, path); err != nil {
421 _ = os.Remove(tmp)
422 }
423}