1package subagents
2
3import (
4 "context"
5 "strings"
6 "sync"
7 "testing"
8 "time"
9)
10
11// steerableRuntime is a ResidentRuntime whose turn blocks until released or
12// canceled, recording steered follow-ups it would deliver at the next boundary.
13type steerableRuntime struct {
14 started chan string
15 release chan struct{}
16
17 mu sync.Mutex
18 running bool
19 pending []string
20 delivered []string
21}
22
23func newSteerableRuntime() *steerableRuntime {
24 return &steerableRuntime{started: make(chan string, 4), release: make(chan struct{})}
25}
26
27func (r *steerableRuntime) Run(ctx context.Context, prompt string) error {
28 r.mu.Lock()
29 r.running = true
30 r.mu.Unlock()
31 defer func() {
32 r.mu.Lock()
33 r.running = false
34 r.mu.Unlock()
35 }()
36 r.started <- prompt
37 select {
38 case <-ctx.Done():
39 return ctx.Err()
40 case <-r.release:
41 }
42 r.mu.Lock()
43 r.delivered = append(r.delivered, r.pending...)
44 r.pending = nil
45 r.mu.Unlock()
46 return nil
47}
48
49func (r *steerableRuntime) Steer(prompt string) bool {
50 r.mu.Lock()
51 defer r.mu.Unlock()
52 if !r.running {
53 return false
54 }
55 r.pending = append(r.pending, prompt)
56 return true
57}
58
59func (r *steerableRuntime) PendingSteers() int {
60 r.mu.Lock()
61 defer r.mu.Unlock()
62 return len(r.pending)
63}
64
65func (r *steerableRuntime) DrainSteers() []string {
66 r.mu.Lock()
67 defer r.mu.Unlock()
68 out := r.pending
69 r.pending = nil
70 return out
71}
72
73func (r *steerableRuntime) awaitStart(t *testing.T) string {
74 t.Helper()
75 select {
76 case prompt := <-r.started:
77 return prompt
78 case <-time.After(5 * time.Second):
79 t.Fatal("turn did not start")
80 return ""
81 }
82}
83
84func newSteerTestManager(t *testing.T, runtime ResidentRuntime) (*ResidentManager, chan ResidentCompletion) {
85 t.Helper()
86 manager := NewResidentManager(t.TempDir(), func(ResidentChildSpec, *ResidentJournal) (ResidentRuntime, error) { return runtime, nil })
87 completions := make(chan ResidentCompletion, 8)
88 manager.SetCompletionObserver(func(c ResidentCompletion) bool { completions <- c; return true })
89 t.Cleanup(func() { _ = manager.Close(context.Background()) })
90 return manager, completions
91}
92
93func drainCompletions(t *testing.T, completions chan ResidentCompletion, want int) []ResidentCompletion {
94 t.Helper()
95 var got []ResidentCompletion
96 for len(got) < want {
97 select {
98 case c := <-completions:
99 got = append(got, c)
100 case <-time.After(5 * time.Second):
101 t.Fatalf("got %d completions, want %d: %#v", len(got), want, got)
102 }
103 }
104 select {
105 case extra := <-completions:
106 t.Fatalf("unexpected extra completion %#v", extra)
107 case <-time.After(50 * time.Millisecond):
108 }
109 return got
110}
111
112// Issue 226: stopping a child must not echo an undelivered queued follow-up
113// back to the parent as if it were an interrupted task.
114func TestResidentStopReportsQueuedFollowUpAsDroppedNotInterrupted(t *testing.T) {
115 runtime := ResidentTurnRunner(func(ctx context.Context, _ string) error { <-ctx.Done(); return ctx.Err() })
116 started := make(chan struct{}, 1)
117 manager, completions := newSteerTestManager(t, ResidentTurnRunner(func(ctx context.Context, prompt string) error {
118 started <- struct{}{}
119 return runtime(ctx, prompt)
120 }))
121 spec := ResidentChildSpec{ID: "stop-echo", InitialTurnID: "initial", SessionID: "session", Provider: "openai", Model: "test"}
122 if _, err := manager.Spawn(t.Context(), spec, "investigate"); err != nil {
123 t.Fatal(err)
124 }
125 <-started
126 if _, err := manager.ResumeFollowUp(t.Context(), spec.ID, "stop and report now", ResumeQueue, "follow-up", nil); err != nil {
127 t.Fatal(err)
128 }
129 if err := manager.Stop(t.Context(), spec.ID); err != nil {
130 t.Fatal(err)
131 }
132 got := drainCompletions(t, completions, 2)
133 byTurn := map[string]Completion{}
134 for _, c := range got {
135 byTurn[c.TurnID] = c.Completion()
136 }
137 if byTurn["initial"].Status != string(ResidentInterrupted) {
138 t.Fatalf("initial completion = %#v", byTurn["initial"])
139 }
140 if byTurn["follow-up"].Status != CompletionDropped {
141 t.Fatalf("follow-up completion = %#v, want dropped", byTurn["follow-up"])
142 }
143 update := FormatCompletionUpdate([]Completion{byTurn["initial"], byTurn["follow-up"]}, "")
144 if strings.Contains(update, "interrupted — stop and report now") || !strings.Contains(update, "dropped (queued follow-up never started: stop and report now)") {
145 t.Fatalf("update = %q", update)
146 }
147}
148
149// A steered follow-up joins the running turn: no new turn, no extra
150// completion, and the running runtime receives it.
151func TestResidentResumeSteersRunningTurn(t *testing.T) {
152 runtime := newSteerableRuntime()
153 manager, completions := newSteerTestManager(t, runtime)
154 spec := ResidentChildSpec{ID: "steered", InitialTurnID: "initial", SessionID: "session", Provider: "openai", Model: "test"}
155 if _, err := manager.Spawn(t.Context(), spec, "investigate"); err != nil {
156 t.Fatal(err)
157 }
158 runtime.awaitStart(t)
159 var watched string
160 outcome, err := manager.ResumeFollowUp(t.Context(), spec.ID, "wrap up now", ResumeSteer, "unused", func(turnID string) { watched = turnID })
161 if err != nil {
162 t.Fatal(err)
163 }
164 if !outcome.Steered || outcome.TurnID != "initial" || watched != "initial" {
165 t.Fatalf("outcome = %#v, watched = %q", outcome, watched)
166 }
167 if snapshot, _ := manager.SnapshotFor(spec.ID); snapshot.PendingFollowUps != 1 {
168 t.Fatalf("pending follow-ups = %d, want 1", snapshot.PendingFollowUps)
169 }
170 close(runtime.release)
171 got := drainCompletions(t, completions, 1)
172 if got[0].TurnID != "initial" || got[0].Err != nil || len(got[0].Undelivered) != 0 {
173 t.Fatalf("completion = %#v", got[0])
174 }
175 if len(runtime.delivered) != 1 || runtime.delivered[0] != "wrap up now" {
176 t.Fatalf("delivered = %#v", runtime.delivered)
177 }
178}
179
180// An idle child cannot be steered, so steer mode falls back to a new turn.
181func TestResidentResumeSteerFallsBackToNewTurnWhenIdle(t *testing.T) {
182 runtime := newSteerableRuntime()
183 close(runtime.release)
184 manager, completions := newSteerTestManager(t, runtime)
185 spec := ResidentChildSpec{ID: "idle", InitialTurnID: "initial", SessionID: "session", Provider: "openai", Model: "test"}
186 if _, err := manager.Spawn(t.Context(), spec, "task"); err != nil {
187 t.Fatal(err)
188 }
189 drainCompletions(t, completions, 1)
190 outcome, err := manager.ResumeFollowUp(t.Context(), spec.ID, "next", ResumeSteer, "second", nil)
191 if err != nil || outcome.Steered || outcome.TurnID != "second" {
192 t.Fatalf("outcome = %#v, err = %v", outcome, err)
193 }
194 if got := drainCompletions(t, completions, 1); got[0].TurnID != "second" {
195 t.Fatalf("completion = %#v", got[0])
196 }
197}
198
199// Stopping a child with a steered follow-up it never read reports that
200// follow-up as undelivered on the interrupted turn.
201func TestResidentStopReportsUndeliveredSteer(t *testing.T) {
202 runtime := newSteerableRuntime()
203 manager, completions := newSteerTestManager(t, runtime)
204 spec := ResidentChildSpec{ID: "undelivered", InitialTurnID: "initial", SessionID: "session", Provider: "openai", Model: "test"}
205 if _, err := manager.Spawn(t.Context(), spec, "investigate"); err != nil {
206 t.Fatal(err)
207 }
208 runtime.awaitStart(t)
209 if _, err := manager.ResumeFollowUp(t.Context(), spec.ID, "stop and report now", ResumeSteer, "", nil); err != nil {
210 t.Fatal(err)
211 }
212 if err := manager.Stop(t.Context(), spec.ID); err != nil {
213 t.Fatal(err)
214 }
215 got := drainCompletions(t, completions, 1)
216 completion := got[0].Completion()
217 if completion.Status != string(ResidentInterrupted) || len(completion.Undelivered) != 1 || completion.Undelivered[0] != "stop and report now" {
218 t.Fatalf("completion = %#v", completion)
219 }
220 if update := FormatCompletionUpdate([]Completion{completion}, ""); !strings.Contains(update, "undelivered: 1 follow-up(s) the child never read:\n - stop and report now") {
221 t.Fatalf("update = %q", update)
222 }
223}
224
225// Interrupt cancels only the running turn; the child stays live and a later
226// resume runs as a new turn with the same runtime.
227func TestResidentInterruptKeepsChildResumable(t *testing.T) {
228 runtime := newSteerableRuntime()
229 manager, completions := newSteerTestManager(t, runtime)
230 spec := ResidentChildSpec{ID: "interrupted", InitialTurnID: "initial", SessionID: "session", Provider: "openai", Model: "test", Required: true}
231 if _, err := manager.Spawn(t.Context(), spec, "investigate"); err != nil {
232 t.Fatal(err)
233 }
234 runtime.awaitStart(t)
235 if _, err := manager.ResumeFollowUp(t.Context(), spec.ID, "queued work", ResumeQueue, "queued", nil); err != nil {
236 t.Fatal(err)
237 }
238 interrupted, err := manager.Interrupt(t.Context(), spec.ID)
239 if err != nil || !interrupted {
240 t.Fatalf("Interrupt = %t, %v", interrupted, err)
241 }
242 got := drainCompletions(t, completions, 2)
243 byTurn := map[string]Completion{}
244 for _, c := range got {
245 byTurn[c.TurnID] = c.Completion()
246 }
247 if byTurn["initial"].Status != string(ResidentInterrupted) || byTurn["queued"].Status != CompletionDropped {
248 t.Fatalf("completions = %#v", byTurn)
249 }
250 if manager.Get(spec.ID) == nil {
251 t.Fatal("interrupt removed the live child")
252 }
253 if again, err := manager.Interrupt(t.Context(), spec.ID); err != nil || again {
254 t.Fatalf("second Interrupt = %t, %v, want no running turn", again, err)
255 }
256 close(runtime.release)
257 outcome, err := manager.ResumeFollowUp(t.Context(), spec.ID, "wrap up with what you have", ResumeSteer, "wrap-up", nil)
258 if err != nil || outcome.Steered {
259 t.Fatalf("resume after interrupt = %#v, %v", outcome, err)
260 }
261 if prompt := runtime.awaitStart(t); prompt != "wrap up with what you have" {
262 t.Fatalf("resumed prompt = %q", prompt)
263 }
264 if got := drainCompletions(t, completions, 1); got[0].TurnID != "wrap-up" || got[0].Err != nil {
265 t.Fatalf("wrap-up completion = %#v", got[0])
266 }
267 if len(manager.UnmetRequired()) != 0 {
268 t.Fatal("successful wrap-up left required work unmet")
269 }
270}
271
272// An interrupt can be accepted just before the turn finishes on its own. The
273// turn must still report interrupted exactly once, and a late interrupt on the
274// now-idle child must be refused.
275func TestResidentInterruptRacingNaturalCompletion(t *testing.T) {
276 started, release := make(chan struct{}, 1), make(chan struct{})
277 // This runtime ignores cancellation, so its turn ends successfully after
278 // the interrupt was already accepted.
279 runtime := ResidentTurnRunner(func(context.Context, string) error {
280 started <- struct{}{}
281 <-release
282 return nil
283 })
284 manager, completions := newSteerTestManager(t, runtime)
285 spec := ResidentChildSpec{ID: "racing", InitialTurnID: "initial", SessionID: "session", Provider: "openai", Model: "test"}
286 if _, err := manager.Spawn(t.Context(), spec, "work"); err != nil {
287 t.Fatal(err)
288 }
289 <-started
290 if ok, err := manager.Interrupt(t.Context(), spec.ID); err != nil || !ok {
291 t.Fatalf("Interrupt = %t, %v", ok, err)
292 }
293 close(release)
294 got := drainCompletions(t, completions, 1)
295 if got[0].TurnID != "initial" || got[0].Completion().Status != string(ResidentInterrupted) {
296 t.Fatalf("completion = %#v", got[0])
297 }
298 if ok, err := manager.Interrupt(t.Context(), spec.ID); err != nil || ok {
299 t.Fatalf("late Interrupt = %t, %v, want refused on idle child", ok, err)
300 }
301 if manager.Get(spec.ID) == nil {
302 t.Fatal("interrupt removed the live child")
303 }
304}
305
306// Steering must not jump ahead of a queued follow-up.
307func TestResidentSteerQueuesBehindPendingTurns(t *testing.T) {
308 runtime := newSteerableRuntime()
309 manager, _ := newSteerTestManager(t, runtime)
310 spec := ResidentChildSpec{ID: "ordered", InitialTurnID: "initial", SessionID: "session", Provider: "openai", Model: "test"}
311 if _, err := manager.Spawn(t.Context(), spec, "investigate"); err != nil {
312 t.Fatal(err)
313 }
314 runtime.awaitStart(t)
315 if _, err := manager.ResumeFollowUp(t.Context(), spec.ID, "first", ResumeQueue, "first", nil); err != nil {
316 t.Fatal(err)
317 }
318 outcome, err := manager.ResumeFollowUp(t.Context(), spec.ID, "second", ResumeSteer, "second", nil)
319 if err != nil || outcome.Steered || outcome.TurnID != "second" {
320 t.Fatalf("outcome = %#v, err = %v", outcome, err)
321 }
322 if snapshot, _ := manager.SnapshotFor(spec.ID); snapshot.PendingFollowUps != 2 {
323 t.Fatalf("pending follow-ups = %d, want 2 queued turns", snapshot.PendingFollowUps)
324 }
325}