1package subagents
2
3import (
4 "context"
5 "errors"
6 "os"
7 "os/exec"
8 "path/filepath"
9 "runtime"
10 "strings"
11 "testing"
12 "time"
13)
14
15type notifyingAgentSink struct {
16 agentSink
17 deltas chan<- struct{}
18}
19
20func (s notifyingAgentSink) assistantDelta(text string) {
21 s.agentSink.assistantDelta(text)
22 select {
23 case s.deltas <- struct{}{}:
24 default:
25 }
26}
27
28const runnerProcessTimeout = 10 * time.Second
29
30type controlledDeadlineContext struct{ done chan struct{} }
31
32func (c *controlledDeadlineContext) Deadline() (time.Time, bool) {
33 return time.Now().Add(time.Hour), true
34}
35func (c *controlledDeadlineContext) Done() <-chan struct{} { return c.done }
36func (c *controlledDeadlineContext) Err() error {
37 select {
38 case <-c.done:
39 return context.DeadlineExceeded
40 default:
41 return nil
42 }
43}
44func (c *controlledDeadlineContext) Value(any) any { return nil }
45func (c *controlledDeadlineContext) expire() {
46 select {
47 case <-c.done:
48 default:
49 close(c.done)
50 }
51}
52
53// TestExecRunnerDeadlineGracefullyStopsWorkerAndPreservesStreamedOutput
54// verifies that deadline expiry shuts a worker down through its inbox rather
55// than killing it before its terminal events can be recorded.
56func TestExecRunnerDeadlineGracefullyStopsWorkerAndPreservesStreamedOutput(t *testing.T) {
57 if runtime.GOOS == "windows" {
58 t.Skip("unix sockets not supported")
59 }
60 if testing.Short() {
61 t.Skip("skip end-to-end runner test in -short mode")
62 }
63
64 exe := buildStubChild(t)
65 t.Setenv("ZUT_STUB_PROTOCOL", "1")
66 t.Setenv("ZUT_STUB_BLOCK_INITIAL", "1")
67
68 root := t.TempDir()
69 inboxPath, err := inboxSocketPath(root, "deadline-test")
70 if err != nil {
71 t.Fatalf("inboxSocketPath: %v", err)
72 }
73 a := &Agent{
74 ID: "deadline-test",
75 Task: "long-running task",
76 Dir: root,
77 SessionPath: filepath.Join(root, "session.jsonl"),
78 InboxPath: inboxPath,
79 EventLogPath: filepath.Join(root, "events.jsonl"),
80 }
81 a.inbox = NewInbox(a.InboxPath)
82 t.Cleanup(func() { _ = a.inbox.Close() })
83
84 r := &execRunner{
85 agent: a,
86 Command: subagentWorkerArgs(subagentWorkerArgsOpts{Exe: exe, Dir: a.Dir, SessionPath: a.SessionPath, InboxPath: a.InboxPath, Task: a.Task}),
87 GracePeriod: time.Second,
88 }
89 ctx := &controlledDeadlineContext{done: make(chan struct{})}
90 defer ctx.expire()
91 deltas := make(chan struct{}, 1)
92 runDone := make(chan error, 1)
93 go func() { runDone <- r.Run(ctx, notifyingAgentSink{agentSink: agentSink{a: a}, deltas: deltas}) }()
94 select {
95 case <-deltas:
96 case <-time.After(runnerProcessTimeout):
97 t.Fatal("worker did not emit its initial delta")
98 }
99 ctx.expire()
100 select {
101 case err := <-runDone:
102 if !errors.Is(err, context.DeadlineExceeded) {
103 t.Fatalf("Run error = %v, want context deadline exceeded", err)
104 }
105 case <-time.After(runnerProcessTimeout):
106 t.Fatal("Run did not return after deadline expiry")
107 }
108
109 events, err := ReadEventLog(a.EventLogPath)
110 if err != nil {
111 t.Fatalf("ReadEventLog: %v", err)
112 }
113 var sawDelta, sawGracefulTurnEnd bool
114 for _, event := range events {
115 if event.Type == "message.delta" && event.Data["delta"] == "partial answer" {
116 sawDelta = true
117 }
118 if event.Type == "turn_end" && event.Data["stop"] == "cancelled" {
119 sawGracefulTurnEnd = true
120 }
121 }
122 if !sawDelta || !sawGracefulTurnEnd {
123 t.Fatalf("deadline did not preserve a graceful partial turn: %s", formatEvents(events))
124 }
125
126 snapshot := a.Snapshot()
127 if snapshot.LastAssistant != "partial answer" {
128 t.Fatalf("live partial output = %q, want %q", snapshot.LastAssistant, "partial answer")
129 }
130 replayed := &Agent{}
131 for _, event := range events {
132 replayEventTranscript(replayed, event)
133 }
134 if snapshot := replayed.Snapshot(); snapshot.LastAssistant != "partial answer" {
135 t.Fatalf("replayed partial output = %q, want %q", snapshot.LastAssistant, "partial answer")
136 }
137}
138
139func TestExecRunnerCancelsBlockedWorker(t *testing.T) {
140 if runtime.GOOS == "windows" {
141 t.Skip("unix sockets not supported")
142 }
143 if testing.Short() {
144 t.Skip("skip end-to-end runner test in -short mode")
145 }
146
147 exe := buildStubChild(t)
148 t.Setenv("ZUT_STUB_PROTOCOL", "1")
149 t.Setenv("ZUT_STUB_BLOCK_INITIAL", "1")
150
151 root := t.TempDir()
152 inboxPath, err := inboxSocketPath(root, "cancel-test")
153 if err != nil {
154 t.Fatalf("inboxSocketPath: %v", err)
155 }
156 a := &Agent{
157 ID: "cancel-test",
158 Task: "blocked task",
159 Dir: root,
160 SessionPath: filepath.Join(root, "session.jsonl"),
161 InboxPath: inboxPath,
162 EventLogPath: filepath.Join(root, "events.jsonl"),
163 }
164 a.inbox = NewInbox(a.InboxPath)
165 t.Cleanup(func() { _ = a.inbox.Close() })
166
167 r := &execRunner{
168 agent: a,
169 Command: subagentWorkerArgs(subagentWorkerArgsOpts{Exe: exe, Dir: a.Dir, SessionPath: a.SessionPath, InboxPath: a.InboxPath, Task: a.Task}),
170 }
171 ctx, cancel := context.WithCancel(context.Background())
172 defer cancel()
173 deltas := make(chan struct{}, 1)
174 runDone := make(chan error, 1)
175 go func() { runDone <- r.Run(ctx, notifyingAgentSink{agentSink: agentSink{a: a}, deltas: deltas}) }()
176
177 select {
178 case <-deltas:
179 case <-time.After(runnerProcessTimeout):
180 t.Fatal("worker did not emit its initial delta")
181 }
182 cancel()
183 select {
184 case err := <-runDone:
185 if !errors.Is(err, context.Canceled) {
186 t.Fatalf("Run error = %v, want context canceled", err)
187 }
188 case <-time.After(runnerProcessTimeout):
189 t.Fatal("Run did not return after context cancellation")
190 }
191}
192
193// TestRunnerEndToEndWithStubChild drives the daemon-mode runner through an
194// initial turn, a follow-up, and a graceful stop using the stub child. It is
195// skipped on platforms without Unix sockets.
196func TestRunnerEndToEndWithStubChild(t *testing.T) {
197 if runtime.GOOS == "windows" {
198 t.Skip("unix sockets not supported")
199 }
200 if testing.Short() {
201 t.Skip("skip end-to-end runner test in -short mode")
202 }
203
204 exe := buildStubChild(t)
205 t.Setenv("ZUT_STUB_PROTOCOL", "1")
206
207 root := t.TempDir()
208 repo := t.TempDir()
209 f := New(Config{
210 Root: root,
211 RepoRoot: repo,
212 NewRunner: func(a *Agent) Runner {
213 return &execRunner{
214 agent: a,
215 Command: subagentWorkerArgs(subagentWorkerArgsOpts{
216 Exe: exe,
217 Dir: a.Dir,
218 SessionPath: a.SessionPath,
219 InboxPath: a.InboxPath,
220 Task: a.Task,
221 Model: a.Model,
222 Provider: a.Provider,
223 }),
224 }
225 },
226 })
227 defer f.StopAll()
228
229 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
230 defer cancel()
231
232 a, err := f.Spawn(ctx, "first task")
233 if err != nil {
234 t.Fatalf("Spawn: %v", err)
235 }
236
237 // Wait until the durable log has at least one assistant_message
238 // from the initial task. That confirms stdout→log→follower.
239 waitFor := func(want string) {
240 t.Helper()
241 deadline := time.Now().Add(5 * time.Second)
242 for time.Now().Before(deadline) {
243 evs, _ := ReadEventLog(a.EventLogPath)
244 for _, ev := range evs {
245 if strings.Contains(eventText(ev), want) {
246 return
247 }
248 }
249 time.Sleep(20 * time.Millisecond)
250 }
251 evs, _ := ReadEventLog(a.EventLogPath)
252 t.Fatalf("timed out waiting for %q in event log; got %d events:\n%s\n%s",
253 want, len(evs), formatEvents(evs), dumpEventsVerbose(evs))
254 }
255 waitFor("echo: first task")
256
257 // Send a follow-up over the inbox. The stub echoes the text into
258 // an assistant_message we can poll for in the log.
259 if err := retrySend(f, a.ID, "follow up", time.Second); err != nil {
260 t.Fatalf("SendUserTurn: %v", err)
261 }
262 waitFor("echo: follow up")
263
264 // Shut the agent down gracefully via the inbox.
265 if err := a.inbox.SendCommand(NewCommand(CommandAgentShutdown, a.ID, a.CurrentTurnID(), AgentShutdownPayload{})); err != nil && !errors.Is(err, ErrNotReady) {
266 t.Fatalf("shutdown send: %v", err)
267 }
268 a.Wait()
269 if got := a.Status(); got != StatusDone && got != StatusKilled {
270 t.Fatalf("final status = %s; want done/killed", got)
271 }
272}
273
274// retrySend exists because the inbox dial races against the child
275// opening the socket. Production callers handle ErrNotReady with a
276// status message; tests retry within a small window.
277func retrySend(f *Supervisor, id, msg string, timeout time.Duration) error {
278 deadline := time.Now().Add(timeout)
279 var lastErr error
280 for time.Now().Before(deadline) {
281 err := f.SendUserTurn(id, msg)
282 if err == nil {
283 return nil
284 }
285 lastErr = err
286 if !errors.Is(err, ErrNotReady) {
287 return err
288 }
289 time.Sleep(30 * time.Millisecond)
290 }
291 return lastErr
292}
293
294func eventText(ev Event) string {
295 if ev.Type != "assistant_message" && ev.Type != "user_message" {
296 return ""
297 }
298 content, _ := ev.Data["content"].([]any)
299 var sb strings.Builder
300 for _, c := range content {
301 m, _ := c.(map[string]any)
302 if t, _ := m["type"].(string); t == "text" {
303 if txt, _ := m["text"].(string); txt != "" {
304 sb.WriteString(txt)
305 sb.WriteByte('\n')
306 }
307 }
308 }
309 return sb.String()
310}
311
312func dumpEventsVerbose(evs []Event) string {
313 var sb strings.Builder
314 for _, ev := range evs {
315 sb.WriteString(ev.Type)
316 sb.WriteString("\t")
317 for k, v := range ev.Data {
318 sb.WriteString(k)
319 sb.WriteString("=")
320 switch vv := v.(type) {
321 case string:
322 sb.WriteString(vv)
323 default:
324 sb.WriteString("<...>")
325 }
326 sb.WriteString(" ")
327 }
328 sb.WriteString("\n")
329 }
330 return sb.String()
331}
332
333func formatEvents(evs []Event) string {
334 var sb strings.Builder
335 for _, ev := range evs {
336 sb.WriteString(ev.Type)
337 sb.WriteString(" ")
338 sb.WriteString(ev.Time.Format(time.RFC3339Nano))
339 sb.WriteString("\n")
340 }
341 return sb.String()
342}
343
344func buildStubChild(t *testing.T) string {
345 t.Helper()
346 out := filepath.Join(t.TempDir(), "stubchild")
347 cmd := exec.Command("go", "build", "-o", out, "./testdata/cmd/stubchild")
348 // Pass through the test runner's env so `go build` can find
349 // HOME, PATH, GOCACHE, etc. CGO is disabled to keep the build
350 // hermetic across machines.
351 cmd.Env = append(os.Environ(), "CGO_ENABLED=0")
352 if b, err := cmd.CombinedOutput(); err != nil {
353 t.Fatalf("build stubchild: %v\n%s", err, b)
354 }
355 return out
356}