| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "fmt" |
| 8 | "path/filepath" |
| 9 | "strings" |
| 10 | "sync/atomic" |
| 11 | "testing" |
| 12 | "time" |
| 13 | |
| 14 | "reasonix/internal/checkpoint" |
| 15 | "reasonix/internal/event" |
| 16 | "reasonix/internal/jobs" |
| 17 | "reasonix/internal/provider" |
| 18 | "reasonix/internal/tool" |
| 19 | ) |
| 20 | |
| 21 | func TestBackgroundFleetRegistersEveryWriterUntilCompletion(t *testing.T) { |
| 22 | root := t.TempDir() |
| 23 | prov := &fleetHoldProvider{started: make(chan struct{}, 2), release: make(chan struct{})} |
| 24 | store := checkpoint.New("", root) |
| 25 | observer := checkpoint.NewMutationObserver(checkpoint.ObserverOptions{Store: store}) |
| 26 | task := NewTaskTool(prov, nil, tool.NewRegistry(), 20, 0, 0, 0, 0, 0, 0, 0.0, "", "sys", nil, 0, "", "", nil). |
| 27 | WithTranscripts(mustSubagentStore(t), root, "base", "high"). |
| 28 | WithScheduler(NewSubagentScheduler(2, 2)). |
| 29 | WithMutationObserver(observer) |
| 30 | fleet := NewFleetTool(task) |
| 31 | manager := jobs.NewManager(event.Discard) |
| 32 | defer manager.Close() |
| 33 | ctx := withCallContext(context.Background(), "fleet-call", event.Discard, nil, false) |
| 34 | ctx = jobs.WithManager(ctx, manager) |
| 35 | ctx = jobs.WithSession(ctx, "parent-session") |
| 36 | args := json.RawMessage(`{ |
| 37 | "run_in_background":true, |
| 38 | "tasks":[ |
| 39 | {"prompt":"first","write_paths":["first.md"]}, |
| 40 | {"prompt":"second","write_paths":["second.md"]} |
| 41 | ] |
| 42 | }`) |
| 43 | if _, err := fleet.Execute(ctx, args); err != nil { |
| 44 | t.Fatal(err) |
| 45 | } |
| 46 | for range 2 { |
| 47 | select { |
| 48 | case <-prov.started: |
| 49 | case <-time.After(2 * time.Second): |
| 50 | t.Fatal("timed out waiting for background fleet writer") |
| 51 | } |
| 52 | } |
| 53 | if writers := observer.ActiveWriters(); len(writers) != 3 { |
| 54 | t.Fatalf("active fleet writers = %+v, want two item writers plus one fleet reservation", writers) |
| 55 | } |
| 56 | running := manager.RunningForSession("parent-session") |
| 57 | if len(running) != 1 { |
| 58 | t.Fatalf("running fleet jobs = %+v, want 1", running) |
| 59 | } |
| 60 | close(prov.release) |
| 61 | result := manager.WaitForSession(context.Background(), "parent-session", []string{running[0].ID}, 5) |
| 62 | if len(result) != 1 || result[0].Status != jobs.Done { |
| 63 | t.Fatalf("background fleet result = %+v", result) |
| 64 | } |
| 65 | if writers := observer.ActiveWriters(); len(writers) != 0 { |
| 66 | t.Fatalf("fleet writers still registered after completion: %+v", writers) |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | // TestBackgroundFleetProgressLifecycleUsesStableIDs guards both sides of the |
| 71 | // background handoff: Execute must leave the shared merger alive for the job, |
| 72 | // and group/child progress must be emitted through the raw parent sink so IDs |
| 73 | // are namespaced exactly once and match the cards already dispatched. |
| 74 | func TestBackgroundFleetProgressLifecycleUsesStableIDs(t *testing.T) { |
| 75 | root := t.TempDir() |
| 76 | rec := &recordSink{} |
| 77 | prov := &fleetHoldProvider{started: make(chan struct{}, 2), release: make(chan struct{})} |
| 78 | task := NewTaskTool(prov, nil, tool.NewRegistry(), 20, 0, 0, 0, 0, 0, 0, 0.0, "", "sys", nil, 0, "", "", nil). |
| 79 | WithTranscripts(mustSubagentStore(t), root, "base", "high"). |
| 80 | WithScheduler(NewSubagentScheduler(2, 2)) |
| 81 | fleet := NewFleetTool(task) |
| 82 | manager := jobs.NewManager(event.Discard) |
| 83 | defer manager.Close() |
| 84 | ctx := withCallContext(context.Background(), "fleet-call", rec, nil, false) |
| 85 | ctx = jobs.WithManager(ctx, manager) |
| 86 | ctx = jobs.WithSession(ctx, "progress-session") |
| 87 | args := json.RawMessage(`{ |
| 88 | "run_in_background":true, |
| 89 | "tasks":[ |
| 90 | {"prompt":"first","write_paths":["first.md"]}, |
| 91 | {"prompt":"second","write_paths":["second.md"]} |
| 92 | ] |
| 93 | }`) |
| 94 | if _, err := fleet.Execute(ctx, args); err != nil { |
| 95 | t.Fatal(err) |
| 96 | } |
| 97 | for range 2 { |
| 98 | select { |
| 99 | case <-prov.started: |
| 100 | case <-time.After(2 * time.Second): |
| 101 | t.Fatal("timed out waiting for background fleet child") |
| 102 | } |
| 103 | } |
| 104 | close(prov.release) |
| 105 | running := manager.RunningForSession("progress-session") |
| 106 | if len(running) != 1 { |
| 107 | t.Fatalf("running fleet jobs = %+v, want 1", running) |
| 108 | } |
| 109 | result := manager.WaitForSession(context.Background(), "progress-session", []string{running[0].ID}, 5) |
| 110 | if len(result) != 1 || result[0].Status != jobs.Done { |
| 111 | t.Fatalf("background fleet result = %+v, want one completed job", result) |
| 112 | } |
| 113 | |
| 114 | groupStatuses := []string{} |
| 115 | childStatuses := map[string][]string{} |
| 116 | childPreviews := map[string]bool{} |
| 117 | for _, e := range rec.kinds(event.ToolProgress) { |
| 118 | if strings.Contains(e.Tool.ID, "fleet-call/fleet-call") { |
| 119 | t.Fatalf("progress ID was namespaced twice: %+v", e.Tool) |
| 120 | } |
| 121 | switch { |
| 122 | case e.Tool.ID == "fleet-call" && progressName(e) == event.SubagentProgressStatusName: |
| 123 | if e.Tool.ParentID != "" { |
| 124 | t.Fatalf("group progress ParentID = %q, want empty", e.Tool.ParentID) |
| 125 | } |
| 126 | groupStatuses = append(groupStatuses, progressOutput(e)) |
| 127 | case strings.HasPrefix(e.Tool.ID, "fleet-call/fleet-"): |
| 128 | if e.Tool.ParentID != "fleet-call" { |
| 129 | t.Fatalf("child progress ParentID = %q, want fleet-call", e.Tool.ParentID) |
| 130 | } |
| 131 | if progressName(e) == event.SubagentProgressStatusName { |
| 132 | childStatuses[e.Tool.ID] = append(childStatuses[e.Tool.ID], progressOutput(e)) |
| 133 | } |
| 134 | if progressName(e) == event.SubagentProgressTextName && progressOutput(e) != "" { |
| 135 | childPreviews[e.Tool.ID] = true |
| 136 | } |
| 137 | } |
| 138 | } |
| 139 | if len(groupStatuses) != 2 || groupStatuses[0] != string(subagentPhaseRunning) || groupStatuses[1] != string(subagentPhaseCompleted) { |
| 140 | t.Fatalf("group lifecycle = %v, want running → completed", groupStatuses) |
| 141 | } |
| 142 | for _, id := range []string{"fleet-call/fleet-1", "fleet-call/fleet-2"} { |
| 143 | statuses := childStatuses[id] |
| 144 | if len(statuses) < 2 || statuses[0] != string(subagentPhaseRunning) || statuses[len(statuses)-1] != string(subagentPhaseCompleted) { |
| 145 | t.Fatalf("child %s lifecycle = %v, want running → … → completed", id, statuses) |
| 146 | } |
| 147 | terminals := 0 |
| 148 | for _, status := range statuses { |
| 149 | if isTerminalStatusOutput(status) { |
| 150 | terminals++ |
| 151 | } |
| 152 | } |
| 153 | if terminals != 1 { |
| 154 | t.Fatalf("child %s terminals = %d, want exactly one", id, terminals) |
| 155 | } |
| 156 | if !childPreviews[id] { |
| 157 | t.Fatalf("child %s never emitted its text preview", id) |
| 158 | } |
| 159 | } |
| 160 | } |
| 161 | |
| 162 | func TestBackgroundFleetRegistersReservationWhileItemsAreQueued(t *testing.T) { |
| 163 | root := t.TempDir() |
| 164 | store := checkpoint.New("", root) |
| 165 | observer := checkpoint.NewMutationObserver(checkpoint.ObserverOptions{Store: store}) |
| 166 | scheduler := NewSubagentScheduler(1, 1) |
| 167 | releaseSlot, err := scheduler.Acquire(context.Background(), AcquireRequest{Writer: false}) |
| 168 | if err != nil { |
| 169 | t.Fatal(err) |
| 170 | } |
| 171 | task := NewTaskTool(&mockProvider{name: "sub"}, nil, tool.NewRegistry(), 20, 0, 0, 0, 0, 0, 0, 0.0, "", "sys", nil, 0, "", "", nil). |
| 172 | WithTranscripts(mustSubagentStore(t), root, "base", "high"). |
| 173 | WithScheduler(scheduler). |
| 174 | WithMutationObserver(observer) |
| 175 | fleet := NewFleetTool(task) |
| 176 | manager := jobs.NewManager(event.Discard) |
| 177 | defer manager.Close() |
| 178 | ctx := withCallContext(context.Background(), "queued-fleet", event.Discard, nil, false) |
| 179 | ctx = jobs.WithManager(ctx, manager) |
| 180 | ctx = jobs.WithSession(ctx, "queued-session") |
| 181 | args := json.RawMessage(`{ |
| 182 | "run_in_background":true, |
| 183 | "tasks":[ |
| 184 | {"prompt":"first","write_paths":["first.md"]}, |
| 185 | {"prompt":"second","write_paths":["second.md"]} |
| 186 | ] |
| 187 | }`) |
| 188 | if _, err := fleet.Execute(ctx, args); err != nil { |
| 189 | t.Fatal(err) |
| 190 | } |
| 191 | writers := observer.ActiveWriters() |
| 192 | if len(writers) != 1 || writers[0].Kind != "background_fleet" { |
| 193 | t.Fatalf("queued fleet reservation = %+v, want one rewind exclusion", writers) |
| 194 | } |
| 195 | releaseSlot() |
| 196 | running := manager.RunningForSession("queued-session") |
| 197 | if len(running) != 1 { |
| 198 | t.Fatalf("running fleet jobs = %+v, want 1", running) |
| 199 | } |
| 200 | result := manager.WaitForSession(context.Background(), "queued-session", []string{running[0].ID}, 5) |
| 201 | if len(result) != 1 || result[0].Status != jobs.Done { |
| 202 | t.Fatalf("background fleet result = %+v, want one completed job", result) |
| 203 | } |
| 204 | if writers := observer.ActiveWriters(); len(writers) != 0 { |
| 205 | t.Fatalf("completed background fleet still registered: %+v", writers) |
| 206 | } |
| 207 | } |
| 208 | |
| 209 | func TestFleetSchemaStableAndBounds(t *testing.T) { |
| 210 | f := NewFleetTool(&TaskTool{}) |
| 211 | schema := string(f.Schema()) |
| 212 | for _, want := range []string{`"profile"`, `"write_paths"`, `"read_only"`, `"run_in_background"`} { |
| 213 | if !strings.Contains(schema, want) { |
| 214 | t.Fatalf("schema missing %s: %s", want, schema) |
| 215 | } |
| 216 | } |
| 217 | // Profile names must not be enumerated in schema (cache stability). |
| 218 | if strings.Contains(schema, "doc-rewriter") || strings.Contains(schema, "enum") { |
| 219 | t.Fatalf("schema must not embed profile names: %s", schema) |
| 220 | } |
| 221 | if f.Name() != "fleet" { |
| 222 | t.Fatalf("name = %q", f.Name()) |
| 223 | } |
| 224 | } |
| 225 | |
| 226 | func TestFleetRejectsSingleTaskAndPathConflict(t *testing.T) { |
| 227 | root := t.TempDir() |
| 228 | task := newTestTaskTool(t, &mockProvider{name: "sub"}, tool.NewRegistry(), "sys", "", "", nil). |
| 229 | WithTranscripts(mustSubagentStore(t), root, "base", "high"). |
| 230 | WithScheduler(NewSubagentScheduler(6, 3)) |
| 231 | f := NewFleetTool(task) |
| 232 | |
| 233 | _, err := f.Execute(context.Background(), json.RawMessage(`{"tasks":[{"prompt":"only one"}]}`)) |
| 234 | if err == nil || !strings.Contains(err.Error(), "between") { |
| 235 | t.Fatalf("single task error = %v", err) |
| 236 | } |
| 237 | |
| 238 | args, _ := json.Marshal(map[string]any{ |
| 239 | "tasks": []map[string]any{ |
| 240 | {"prompt": "a", "write_paths": []string{"same.md"}}, |
| 241 | {"prompt": "b", "write_paths": []string{"same.md"}}, |
| 242 | }, |
| 243 | }) |
| 244 | _, err = f.Execute(withCallContext(context.Background(), "fleet-call", event.Discard, nil, false), args) |
| 245 | if err == nil || !strings.Contains(err.Error(), "conflict") { |
| 246 | t.Fatalf("path conflict error = %v", err) |
| 247 | } |
| 248 | |
| 249 | // Read-only items must not shift the caller-visible task numbers in the |
| 250 | // preflight diagnostic. |
| 251 | args, _ = json.Marshal(map[string]any{ |
| 252 | "tasks": []map[string]any{ |
| 253 | {"prompt": "inspect", "read_only": true}, |
| 254 | {"prompt": "writer a", "write_paths": []string{"same.md"}}, |
| 255 | {"prompt": "writer b", "write_paths": []string{"same.md"}}, |
| 256 | }, |
| 257 | }) |
| 258 | _, err = f.Execute(withCallContext(context.Background(), "fleet-call", event.Discard, nil, false), args) |
| 259 | if err == nil || !strings.Contains(err.Error(), "task 2 and task 3") { |
| 260 | t.Fatalf("mixed-task conflict error = %v, want original task numbers 2 and 3", err) |
| 261 | } |
| 262 | } |
| 263 | |
| 264 | func TestFleetCancellationPreservesStartedItemStatus(t *testing.T) { |
| 265 | root := t.TempDir() |
| 266 | prov := &fleetCancelProvider{ |
| 267 | started: make(chan struct{}, 2), |
| 268 | observed: make(chan struct{}, 2), |
| 269 | release: make(chan struct{}), |
| 270 | } |
| 271 | reg := tool.NewRegistry() |
| 272 | task := NewTaskTool(prov, nil, reg, 20, 0, 0, 0, 0, 0, 0, 0.0, "", "sys", nil, 0, "", "", nil). |
| 273 | WithTranscripts(mustSubagentStore(t), root, "base", "high"). |
| 274 | WithScheduler(NewSubagentScheduler(2, 2)) |
| 275 | f := NewFleetTool(task) |
| 276 | |
| 277 | ctx, cancel := context.WithCancel(withCallContext(context.Background(), "fleet-call", event.Discard, nil, false)) |
| 278 | done := make(chan struct { |
| 279 | out string |
| 280 | err error |
| 281 | }, 1) |
| 282 | go func() { |
| 283 | out, err := f.Execute(ctx, json.RawMessage(`{ |
| 284 | "tasks":[ |
| 285 | {"prompt":"first","write_paths":["first.md"]}, |
| 286 | {"prompt":"second","write_paths":["second.md"]} |
| 287 | ] |
| 288 | }`)) |
| 289 | done <- struct { |
| 290 | out string |
| 291 | err error |
| 292 | }{out: out, err: err} |
| 293 | }() |
| 294 | |
| 295 | // Both workers are inside the provider before cancellation. Hold their |
| 296 | // terminal results until the fleet has observed ctx.Done, then release them. |
| 297 | waitSignal := func(name string, ch <-chan struct{}) { |
| 298 | t.Helper() |
| 299 | select { |
| 300 | case <-ch: |
| 301 | case <-time.After(2 * time.Second): |
| 302 | t.Fatalf("timed out waiting for %s", name) |
| 303 | } |
| 304 | } |
| 305 | for range 2 { |
| 306 | waitSignal("provider start", prov.started) |
| 307 | } |
| 308 | cancel() |
| 309 | for range 2 { |
| 310 | waitSignal("provider cancellation", prov.observed) |
| 311 | } |
| 312 | close(prov.release) |
| 313 | |
| 314 | var got struct { |
| 315 | out string |
| 316 | err error |
| 317 | } |
| 318 | select { |
| 319 | case got = <-done: |
| 320 | case <-time.After(2 * time.Second): |
| 321 | t.Fatal("timed out waiting for fleet cancellation result") |
| 322 | } |
| 323 | if !errors.Is(got.err, context.Canceled) { |
| 324 | t.Fatalf("fleet error = %v, want context.Canceled", got.err) |
| 325 | } |
| 326 | if strings.Contains(got.out, "status: skipped") { |
| 327 | t.Fatalf("started tasks must not be reported skipped after cancellation:\n%s", got.out) |
| 328 | } |
| 329 | if count := strings.Count(got.out, "status: cancelled"); count != 2 { |
| 330 | t.Fatalf("cancelled status count = %d, want 2:\n%s", count, got.out) |
| 331 | } |
| 332 | } |
| 333 | |
| 334 | func TestFleetParallelDisjointWriters(t *testing.T) { |
| 335 | root := t.TempDir() |
| 336 | var concurrent atomic.Int32 |
| 337 | var maxConcurrent atomic.Int32 |
| 338 | prov := &fleetBarrierProvider{ |
| 339 | onPrompt: func() { |
| 340 | cur := concurrent.Add(1) |
| 341 | for { |
| 342 | old := maxConcurrent.Load() |
| 343 | if cur <= old || maxConcurrent.CompareAndSwap(old, cur) { |
| 344 | break |
| 345 | } |
| 346 | } |
| 347 | time.Sleep(30 * time.Millisecond) |
| 348 | concurrent.Add(-1) |
| 349 | }, |
| 350 | } |
| 351 | reg := tool.NewRegistry() |
| 352 | // No writer tools needed — provider finishes without tools. |
| 353 | task := NewTaskTool(prov, nil, reg, 20, 0, 0, 0, 0, 0, 0, 0.0, "", "sys", nil, 0, "", "", nil). |
| 354 | WithTranscripts(mustSubagentStore(t), root, "base", "high"). |
| 355 | WithScheduler(NewSubagentScheduler(10, 10)) |
| 356 | f := NewFleetTool(task) |
| 357 | |
| 358 | tasks := make([]map[string]any, 0, 4) |
| 359 | for i := 0; i < 4; i++ { |
| 360 | path := filepath.Join("docs", "f"+string(rune('0'+i))+".md") |
| 361 | tasks = append(tasks, map[string]any{ |
| 362 | "prompt": "handle " + path, |
| 363 | "write_paths": []string{path}, |
| 364 | "description": path, |
| 365 | }) |
| 366 | } |
| 367 | args, _ := json.Marshal(map[string]any{"tasks": tasks}) |
| 368 | ctx := withCallContext(context.Background(), "fleet-call", event.Discard, nil, false) |
| 369 | out, err := f.Execute(ctx, args) |
| 370 | if err != nil { |
| 371 | t.Fatalf("fleet: %v", err) |
| 372 | } |
| 373 | if !strings.Contains(out, "Completed fleet of 4") { |
| 374 | t.Fatalf("output = %s", out) |
| 375 | } |
| 376 | if maxConcurrent.Load() < 2 { |
| 377 | t.Fatalf("expected concurrent starts, max=%d", maxConcurrent.Load()) |
| 378 | } |
| 379 | } |
| 380 | |
| 381 | func TestFleetAggregatePreservesEveryReferenceUnderToolLimit(t *testing.T) { |
| 382 | results := make([]fleetItemResult, 3) |
| 383 | for i := range results { |
| 384 | results[i] = fleetItemResult{ |
| 385 | index: i, |
| 386 | status: fleetItemCompleted, |
| 387 | output: fmt.Sprintf("BEGIN-%d\n%s\nEND-%d", i+1, strings.Repeat(string(rune('a'+i)), 20*1024), i+1), |
| 388 | ref: fmt.Sprintf("sa_result_%d", i+1), |
| 389 | } |
| 390 | } |
| 391 | out := formatFleetAggregate(results, false) |
| 392 | if len(out) > subagentAggregateBudgetBytes { |
| 393 | t.Fatalf("aggregate bytes = %d, want <= %d", len(out), subagentAggregateBudgetBytes) |
| 394 | } |
| 395 | if _, notice := truncateToolOutput(out); notice != "" { |
| 396 | t.Fatalf("bounded fleet aggregate still hit generic truncation: %s", notice) |
| 397 | } |
| 398 | for i := range results { |
| 399 | if !strings.Contains(out, results[i].ref) { |
| 400 | t.Fatalf("aggregate lost ref %q", results[i].ref) |
| 401 | } |
| 402 | } |
| 403 | } |
| 404 | |
| 405 | type fleetBarrierProvider struct { |
| 406 | onPrompt func() |
| 407 | } |
| 408 | |
| 409 | type fleetCancelProvider struct { |
| 410 | started chan struct{} |
| 411 | observed chan struct{} |
| 412 | release chan struct{} |
| 413 | } |
| 414 | |
| 415 | type fleetHoldProvider struct { |
| 416 | started chan struct{} |
| 417 | release chan struct{} |
| 418 | } |
| 419 | |
| 420 | func (p *fleetHoldProvider) Name() string { return "fleet-hold" } |
| 421 | |
| 422 | func (p *fleetHoldProvider) Stream(_ context.Context, _ provider.Request) (<-chan provider.Chunk, error) { |
| 423 | p.started <- struct{}{} |
| 424 | <-p.release |
| 425 | ch := make(chan provider.Chunk, 1) |
| 426 | ch <- provider.Chunk{Type: provider.ChunkText, Text: "done"} |
| 427 | close(ch) |
| 428 | return ch, nil |
| 429 | } |
| 430 | |
| 431 | func (p *fleetCancelProvider) Name() string { return "fleet-cancel" } |
| 432 | |
| 433 | func (p *fleetCancelProvider) Stream(ctx context.Context, _ provider.Request) (<-chan provider.Chunk, error) { |
| 434 | p.started <- struct{}{} |
| 435 | <-ctx.Done() |
| 436 | p.observed <- struct{}{} |
| 437 | <-p.release |
| 438 | return nil, ctx.Err() |
| 439 | } |
| 440 | |
| 441 | func (p *fleetBarrierProvider) Name() string { return "fleet-barrier" } |
| 442 | |
| 443 | func (p *fleetBarrierProvider) Stream(_ context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 444 | if p.onPrompt != nil { |
| 445 | p.onPrompt() |
| 446 | } |
| 447 | ch := make(chan provider.Chunk, 2) |
| 448 | ch <- provider.Chunk{Type: provider.ChunkText, Text: "done"} |
| 449 | close(ch) |
| 450 | return ch, nil |
| 451 | } |
| 452 | |
| 453 | func mustSubagentStore(t *testing.T) *SubagentStore { |
| 454 | t.Helper() |
| 455 | return NewSubagentStore(t.TempDir()) |
| 456 | } |
| 457 |