1package subagents
2
3import (
4 "context"
5 "errors"
6 "path/filepath"
7 "runtime"
8 "testing"
9 "time"
10)
11
12func TestResumeWithPromptQueuesFollowUpsForActiveWorker(t *testing.T) {
13 if runtime.GOOS == "windows" {
14 t.Skip("subagent inbox transport uses Unix-domain sockets")
15 }
16
17 root := shortSocketDir(t)
18 ready := make(chan struct{}, 1)
19 commands := make(chan Envelope, 2)
20 f := New(Config{
21 Root: root, RepoRoot: root,
22 NewRunner: func(a *Agent) Runner {
23 return RunnerFunc(func(ctx context.Context, _ Sink) error {
24 listener, err := Listen(a.InboxPath)
25 if err != nil {
26 return err
27 }
28 defer listener.Close()
29 a.setProcessState(ProcessAlive)
30 a.setTurnState(TurnRunning, "turn-1")
31 ready <- struct{}{}
32 for {
33 select {
34 case <-ctx.Done():
35 return ctx.Err()
36 case line, ok := <-listener.Lines():
37 if !ok {
38 return nil
39 }
40 command, err := ParseCommand(line)
41 if err != nil {
42 return err
43 }
44 if command.Type == CommandAgentShutdown {
45 return nil
46 }
47 commands <- command
48 }
49 }
50 })
51 },
52 })
53 defer f.StopAll()
54
55 a, err := f.Spawn(context.Background(), "review the implementation")
56 if err != nil {
57 t.Fatal(err)
58 }
59 select {
60 case <-ready:
61 case <-time.After(time.Second):
62 t.Fatal("active worker did not open its inbox")
63 }
64 canceled, cancel := context.WithCancel(context.Background())
65 cancel()
66 if _, err := f.ResumeWithPrompt(canceled, a.ID, "must not be accepted"); !errors.Is(err, context.Canceled) {
67 t.Fatalf("canceled active follow-up error = %v, want context canceled", err)
68 }
69
70 const (
71 firstFollowUp = "Check the parser behavior."
72 secondFollowUp = "Then check the error handling."
73 )
74 tracker := NewCompletionTracker()
75 if _, err := f.ResumeWithPromptBefore(context.Background(), a.ID, firstFollowUp, func(agent *Agent, prompt string) func() {
76 return tracker.TrackFutureTurn(agent, prompt, true)
77 }); err != nil {
78 t.Fatalf("queue first active follow-up: %v", err)
79 }
80 if a.OnTurnEnd == nil {
81 t.Fatal("active resume did not register its turn callback before queueing")
82 }
83 a.OnTurnEnd(2, "")
84 batch, err := tracker.WaitIdle(context.Background())
85 if err != nil || len(batch) != 1 || batch[0].Task != firstFollowUp {
86 t.Fatalf("active resume completion = (%+v, %v), want one %q completion", batch, err, firstFollowUp)
87 }
88 if _, err := f.ResumeWithPrompt(context.Background(), a.ID, secondFollowUp); err != nil {
89 t.Fatalf("queue second active follow-up: %v", err)
90 }
91 select {
92 case command := <-commands:
93 t.Fatalf("active follow-up was delivered before idle: %v", command.Type)
94 default:
95 }
96
97 persisted, err := readAgentMeta(filepath.Join(root, "agents", a.ID))
98 if err != nil {
99 t.Fatal(err)
100 }
101 if persisted.ResumePrompt != "" || len(persisted.ResumeQueue) != 2 {
102 t.Fatalf("active follow-ups persisted as prompt %q queue %d, want empty prompt and queue 2", persisted.ResumePrompt, len(persisted.ResumeQueue))
103 }
104 if firstID, secondID := persisted.ResumeQueue[0].CommandID, persisted.ResumeQueue[1].CommandID; firstID == "" || secondID == "" || firstID == secondID {
105 t.Fatalf("queued follow-up command IDs = %q, %q; want distinct non-empty IDs", firstID, secondID)
106 }
107
108 if err := updateAgentFromEvent(a, NewEvent(EventAgentIdle, map[string]any{"turn_id": "turn-1"})); err != nil {
109 t.Fatalf("apply first idle event: %v", err)
110 }
111 select {
112 case command := <-commands:
113 var payload TurnStartPayload
114 if command.Type != CommandTurnStart || command.DecodePayload(&payload) != nil || payload.Prompt != firstFollowUp {
115 t.Fatalf("first idle delivery = type %q prompt %q, want turn start %q", command.Type, payload.Prompt, firstFollowUp)
116 }
117 if !payload.NewRun {
118 t.Fatal("first queued follow-up did not request a fresh run budget")
119 }
120 case <-time.After(time.Second):
121 t.Fatal("first active follow-up was not delivered at idle")
122 }
123 if got := a.TurnState(); got != TurnQueued {
124 t.Fatalf("turn state after first queued delivery = %s, want %s", got, TurnQueued)
125 }
126
127 if err := updateAgentFromEvent(a, NewEvent(EventTurnStarted, map[string]any{
128 "turn_id": "turn-2", "step": float64(2), "lifetime_turns": 2, "current_run_turns": 1,
129 })); err != nil {
130 t.Fatalf("apply second turn started event: %v", err)
131 }
132 if err := updateAgentFromEvent(a, NewEvent(EventAgentIdle, map[string]any{"turn_id": "turn-2"})); err != nil {
133 t.Fatalf("apply second idle event: %v", err)
134 }
135 select {
136 case command := <-commands:
137 var payload TurnStartPayload
138 if command.Type != CommandTurnStart || command.DecodePayload(&payload) != nil || payload.Prompt != secondFollowUp {
139 t.Fatalf("second idle delivery = type %q prompt %q, want turn start %q", command.Type, payload.Prompt, secondFollowUp)
140 }
141 case <-time.After(time.Second):
142 t.Fatal("second active follow-up was not delivered at the next idle turn")
143 }
144
145 persisted, err = readAgentMeta(filepath.Join(root, "agents", a.ID))
146 if err != nil {
147 t.Fatal(err)
148 }
149 if persisted.ResumePrompt != secondFollowUp || len(persisted.ResumeQueue) != 0 {
150 t.Fatalf("second follow-up persistence = prompt %q queue %d, want prompt %q and empty queue", persisted.ResumePrompt, len(persisted.ResumeQueue), secondFollowUp)
151 }
152 if persisted.ResumePromptID == "" {
153 t.Fatal("promoted follow-up lost its command ID")
154 }
155 if err := f.Stop(a.ID); err != nil {
156 t.Fatal(err)
157 }
158 a.Wait()
159}
160
161func TestResumeWithPromptRestartTracksExistingPendingPromptBeforeNewPrompt(t *testing.T) {
162 root := t.TempDir()
163 started := make(chan *Agent, 2)
164 f := New(Config{
165 Root: root, RepoRoot: root,
166 NewRunner: func(a *Agent) Runner {
167 return RunnerFunc(func(ctx context.Context, _ Sink) error {
168 started <- a
169 <-ctx.Done()
170 return ctx.Err()
171 })
172 },
173 })
174 t.Cleanup(f.StopAll)
175
176 original, err := f.Spawn(context.Background(), "review the implementation")
177 if err != nil {
178 t.Fatal(err)
179 }
180 select {
181 case <-started:
182 case <-time.After(time.Second):
183 t.Fatal("original runner did not start")
184 }
185 if err := f.Stop(original.ID); err != nil {
186 t.Fatal(err)
187 }
188 original.Wait()
189
190 const (
191 pendingPrompt = "finish the queued review"
192 newPrompt = "then review the follow-up"
193 )
194 original.setResumePrompt(pendingPrompt, time.Now())
195 if err := f.persistAgent(original); err != nil {
196 t.Fatal(err)
197 }
198
199 tracker := NewCompletionTracker()
200 resumed, err := f.ResumeWithPromptBefore(context.Background(), original.ID, newPrompt, func(agent *Agent, prompt string) func() {
201 return tracker.TrackFutureTurn(agent, prompt, true)
202 })
203 if err != nil {
204 t.Fatal(err)
205 }
206 select {
207 case got := <-started:
208 if got != resumed {
209 t.Fatalf("started agent = %p, want resumed agent %p", got, resumed)
210 }
211 case <-time.After(time.Second):
212 t.Fatal("resumed runner did not start")
213 }
214 if got := resumed.resumePrompt(); got != pendingPrompt {
215 t.Fatalf("first scheduled prompt = %q, want pending prompt %q", got, pendingPrompt)
216 }
217 queue := resumed.resumePromptQueueSnapshot()
218 if len(queue) != 1 || queue[0].Prompt != newPrompt {
219 t.Fatalf("resume queue = %+v, want new prompt queued second", queue)
220 }
221 if resumed.OnTurnEnd == nil {
222 t.Fatal("restart did not install completion tracking before scheduling")
223 }
224 resumed.OnTurnEnd(1, "")
225 resumed.OnTurnEnd(2, "")
226 batch, err := tracker.WaitIdle(context.Background())
227 if err != nil {
228 t.Fatal(err)
229 }
230 if len(batch) != 2 || batch[0].Task != pendingPrompt || batch[1].Task != newPrompt {
231 t.Fatalf("completion order = %+v, want pending prompt then new prompt", batch)
232 }
233
234 // Recreate the durable queue-only state: no in-flight prompt, with the
235 // previously queued command still awaiting scheduler consumption.
236 if err := f.Stop(resumed.ID); err != nil {
237 t.Fatal(err)
238 }
239 resumed.Wait()
240 resumed.clearResumePrompt()
241 if err := f.persistAgent(resumed); err != nil {
242 t.Fatal(err)
243 }
244
245 const newestPrompt = "finally synthesize both reviews"
246 queueTracker := NewCompletionTracker()
247 restarted, err := f.ResumeWithPromptBefore(context.Background(), resumed.ID, newestPrompt, func(agent *Agent, prompt string) func() {
248 return queueTracker.TrackFutureTurn(agent, prompt, true)
249 })
250 if err != nil {
251 t.Fatal(err)
252 }
253 select {
254 case got := <-started:
255 if got != restarted {
256 t.Fatalf("queue-only restart agent = %p, want %p", got, restarted)
257 }
258 case <-time.After(time.Second):
259 t.Fatal("queue-only restarted runner did not start")
260 }
261 if got := restarted.resumePrompt(); got != newPrompt {
262 t.Fatalf("promoted queue prompt = %q, want %q", got, newPrompt)
263 }
264 queue = restarted.resumePromptQueueSnapshot()
265 if len(queue) != 1 || queue[0].Prompt != newestPrompt {
266 t.Fatalf("queue-only restart queue = %+v, want newest prompt second", queue)
267 }
268 restarted.OnTurnEnd(1, "")
269 restarted.OnTurnEnd(2, "")
270 batch, err = queueTracker.WaitIdle(context.Background())
271 if err != nil {
272 t.Fatal(err)
273 }
274 if len(batch) != 2 || batch[0].Task != newPrompt || batch[1].Task != newestPrompt {
275 t.Fatalf("queue-only completion order = %+v, want promoted prompt then newest prompt", batch)
276 }
277}
278
279func TestResumeWithPromptDeliversFollowUpToIdleWorker(t *testing.T) {
280 if runtime.GOOS == "windows" {
281 t.Skip("subagent inbox transport uses Unix-domain sockets")
282 }
283
284 root := shortSocketDir(t)
285 ready := make(chan struct{}, 1)
286 commands := make(chan Envelope, 2)
287 f := New(Config{
288 Root: root, RepoRoot: root,
289 NewRunner: func(a *Agent) Runner {
290 return RunnerFunc(func(ctx context.Context, _ Sink) error {
291 listener, err := Listen(a.InboxPath)
292 if err != nil {
293 return err
294 }
295 defer listener.Close()
296 a.setProcessState(ProcessAlive)
297 a.setTurnState(TurnIdle, "")
298 ready <- struct{}{}
299 for {
300 select {
301 case <-ctx.Done():
302 return ctx.Err()
303 case line, ok := <-listener.Lines():
304 if !ok {
305 return nil
306 }
307 command, err := ParseCommand(line)
308 if err != nil {
309 return err
310 }
311 if command.Type == CommandAgentShutdown {
312 return nil
313 }
314 commands <- command
315 }
316 }
317 })
318 },
319 })
320 defer f.StopAll()
321
322 a, err := f.Spawn(context.Background(), "review the implementation")
323 if err != nil {
324 t.Fatal(err)
325 }
326 select {
327 case <-ready:
328 case <-time.After(time.Second):
329 t.Fatal("idle worker did not open its inbox")
330 }
331 a.setTurnCounts(4, 2)
332
333 const followUp = "I applied your review. What do you think now?"
334 tracker := NewCompletionTracker()
335 continued, err := f.ResumeWithPromptBefore(context.Background(), a.ID, followUp, func(agent *Agent, prompt string) func() {
336 return tracker.TrackFutureTurn(agent, prompt, true)
337 })
338 if err != nil {
339 t.Fatalf("continue idle worker: %v", err)
340 }
341 if continued != a {
342 t.Fatalf("continued agent = %p, want existing worker %p", continued, a)
343 }
344 if got := a.TurnState(); got != TurnQueued {
345 t.Fatalf("turn state after sending follow-up = %s, want %s", got, TurnQueued)
346 }
347 if got := a.CurrentRunTurnsValue(); got != 0 {
348 t.Fatalf("current-run turns after explicit resume = %d, want 0 before turn starts", got)
349 }
350 if got := a.LifetimeTurnsValue(); got != 4 {
351 t.Fatalf("lifetime turns after explicit resume = %d, want 4", got)
352 }
353 if _, err := f.ResumeWithPrompt(context.Background(), a.ID, "duplicate follow-up"); err == nil {
354 t.Fatal("second follow-up succeeded while the first was queued")
355 }
356 select {
357 case command := <-commands:
358 if command.Type != CommandTurnStart {
359 t.Fatalf("command type = %q, want %q", command.Type, CommandTurnStart)
360 }
361 var payload TurnStartPayload
362 if err := command.DecodePayload(&payload); err != nil {
363 t.Fatal(err)
364 }
365 if payload.Prompt != followUp {
366 t.Fatalf("follow-up prompt = %q, want %q", payload.Prompt, followUp)
367 }
368 if !payload.NewRun {
369 t.Fatal("idle follow-up did not request a fresh run budget")
370 }
371 if a.OnTurnEnd == nil {
372 t.Fatal("idle resume callback was not installed before prompt delivery")
373 }
374 a.OnTurnEnd(5, "")
375 case <-time.After(time.Second):
376 t.Fatal("idle worker did not receive the follow-up")
377 }
378 batch, err := tracker.WaitIdle(context.Background())
379 if err != nil || len(batch) != 1 || batch[0].Task != followUp {
380 t.Fatalf("idle resume completion = (%+v, %v), want one %q completion", batch, err, followUp)
381 }
382 persisted, err := readAgentMeta(filepath.Join(root, "agents", a.ID))
383 if err != nil {
384 t.Fatal(err)
385 }
386 if persisted.ResumePrompt != followUp || persisted.ResumePromptAt.IsZero() {
387 t.Fatalf("queued follow-up was not persisted: prompt %q accepted %s", persisted.ResumePrompt, persisted.ResumePromptAt)
388 }
389 if err := f.Stop(a.ID); err != nil {
390 t.Fatal(err)
391 }
392 a.Wait()
393
394 restarted := make(chan *Agent, 1)
395 f2 := New(Config{
396 Root: root, RepoRoot: root,
397 NewRunner: func(agent *Agent) Runner {
398 return RunnerFunc(func(ctx context.Context, _ Sink) error {
399 restarted <- agent
400 <-ctx.Done()
401 return ctx.Err()
402 })
403 },
404 })
405 defer f2.StopAll()
406 if loaded, errs := f2.Reload(); loaded != 1 || len(errs) != 0 {
407 t.Fatalf("reload = (%d, %v), want (1, no errors)", loaded, errs)
408 }
409 resumed, err := f2.ResumeSession(context.Background(), a.ID)
410 if err != nil {
411 t.Fatal(err)
412 }
413 select {
414 case got := <-restarted:
415 if got != resumed || got.resumePrompt() != followUp {
416 t.Fatalf("unacknowledged follow-up was not replayed: agent %p prompt %q", got, got.resumePrompt())
417 }
418 case <-time.After(time.Second):
419 t.Fatal("restarted worker did not receive the follow-up")
420 }
421}