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