1package agent
2
3import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "errors"
8 "io"
9 "os"
10 "path/filepath"
11 "strings"
12 "sync"
13 "testing"
14
15 "github.com/bnema/zut/packages/agent/subagents"
16 "github.com/bnema/zut/packages/core"
17 "github.com/bnema/zut/packages/provider"
18)
19
20// TestWorkerOutputLimitsUseSupervisorPolicy verifies that the child reads
21// the effective output caps propagated by the supervisor and retains safe
22// defaults for standalone workers.
23func TestWorkerOutputLimitsUseSupervisorPolicy(t *testing.T) {
24 t.Setenv("ZUT_SUBAGENT_MAX_OUTPUT_BYTES", "123")
25 t.Setenv("ZUT_SUBAGENT_MAX_OUTPUT_LINES", "7")
26 bytesLimit, linesLimit := workerOutputLimits()
27 if bytesLimit != 123 || linesLimit != 7 {
28 t.Fatalf("worker output limits = (%d, %d), want (123, 7)", bytesLimit, linesLimit)
29 }
30 t.Setenv("ZUT_SUBAGENT_MAX_OUTPUT_BYTES", "invalid")
31 t.Setenv("ZUT_SUBAGENT_MAX_OUTPUT_LINES", "0")
32 bytesLimit, linesLimit = workerOutputLimits()
33 if bytesLimit != 500_000 || linesLimit != 5_000 {
34 t.Fatalf("invalid worker output limits = (%d, %d), want defaults", bytesLimit, linesLimit)
35 }
36}
37
38type subagentContextRecoveryClient struct {
39 mu sync.Mutex
40 calls int
41 requests []provider.Request
42 retried chan provider.Request
43 compactionErr error
44 retryOverflow bool
45}
46
47func (c *subagentContextRecoveryClient) Name() string { return "subagent-context-recovery-test" }
48
49func (c *subagentContextRecoveryClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
50 c.mu.Lock()
51 c.calls++
52 call := c.calls
53 c.requests = append(c.requests, req)
54 c.mu.Unlock()
55
56 events := make(chan provider.Event, 2)
57 go func() {
58 defer close(events)
59 switch call {
60 case 1:
61 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("provider error: input exceeds the context window")}
62 case 2:
63 if c.compactionErr != nil {
64 events <- provider.EventDone{Stop: provider.StopError, Err: c.compactionErr}
65 return
66 }
67 events <- provider.EventTextDelta{Delta: "summary"}
68 events <- provider.EventDone{Stop: provider.StopEnd}
69 default:
70 c.retried <- req
71 if c.retryOverflow {
72 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("context window exceeded after compaction")}
73 return
74 }
75 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
76 Role: provider.RoleAssistant,
77 Content: []provider.Content{provider.TextBlock{Text: "recovered"}},
78 }}
79 }
80 }()
81 return events, nil
82}
83
84func TestPromptWithContextRecoveryCompactsAndContinues(t *testing.T) {
85 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
86 agent := core.NewAgent(client, "test-model", "", nil)
87 agent.SetMessages([]provider.Message{
88 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
89 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
90 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
91 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
92 })
93
94 outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, nil)
95 if err != nil {
96 t.Fatalf("promptWithContextRecovery returned %v", err)
97 }
98 if got := finalAssistantText(agent.Messages()[outputStart:]); got != "recovered" {
99 t.Fatalf("final assistant text = %q, want recovered", got)
100 }
101
102 request := <-client.retried
103 if got := workerRequestUserTextCount(request, "finish the task"); got != 1 {
104 t.Fatalf("retried request contains prompt %d times, want 1: %#v", got, request.Messages)
105 }
106 client.mu.Lock()
107 calls := client.calls
108 requests := append([]provider.Request(nil), client.requests...)
109 client.mu.Unlock()
110 if calls != 3 || len(requests) != 3 {
111 t.Fatalf("provider calls/requests = %d/%d, want 3/3", calls, len(requests))
112 }
113 if !workerRequestContainsText(requests[1], "<conversation>") || !workerRequestContainsText(requests[1], "one") {
114 t.Fatalf("compaction request does not include the older transcript: %#v", requests[1].Messages)
115 }
116}
117
118func TestPromptWithContextRecoveryStopsAfterOneRetry(t *testing.T) {
119 client := &subagentContextRecoveryClient{
120 retried: make(chan provider.Request, 1),
121 retryOverflow: true,
122 }
123 agent := core.NewAgent(client, "test-model", "", nil)
124 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
125
126 _, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, nil)
127 if !provider.IsContextOverflowError(err) {
128 t.Fatalf("promptWithContextRecovery error = %v, want context overflow", err)
129 }
130 <-client.retried
131 client.mu.Lock()
132 calls := client.calls
133 client.mu.Unlock()
134 if calls != 3 {
135 t.Fatalf("provider calls = %d, want one initial prompt, compaction, and one continuation", calls)
136 }
137}
138
139func TestPromptWithContextRecoveryDoesNotCompactAfterCancellation(t *testing.T) {
140 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
141 agent := core.NewAgent(client, "test-model", "", nil)
142 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
143 ctx, cancel := context.WithCancel(context.Background())
144 defer cancel()
145
146 _, err := promptWithContextRecovery(ctx, agent, "finish the task", func(event core.AgentEvent) {
147 if turnEnd, ok := event.(core.EvTurnEnd); ok && provider.IsContextOverflowError(turnEnd.Err) {
148 cancel()
149 }
150 }, nil)
151 if !provider.IsContextOverflowError(err) {
152 t.Fatalf("promptWithContextRecovery error = %v, want context overflow", err)
153 }
154 if err := ctx.Err(); !errors.Is(err, context.Canceled) {
155 t.Fatalf("context error = %v, want %v", err, context.Canceled)
156 }
157 client.mu.Lock()
158 calls := client.calls
159 client.mu.Unlock()
160 if calls != 1 {
161 t.Fatalf("provider calls = %d, want only the initial prompt after cancellation", calls)
162 }
163}
164
165func TestPromptWithContextRecoveryPropagatesCompactionFailure(t *testing.T) {
166 client := &subagentContextRecoveryClient{
167 retried: make(chan provider.Request, 1),
168 compactionErr: errors.New("compaction provider unavailable"),
169 }
170 agent := core.NewAgent(client, "test-model", "", nil)
171 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
172
173 _, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, nil)
174 if err == nil || !strings.Contains(err.Error(), "compact transcript after initial overflow") {
175 t.Fatalf("promptWithContextRecovery error = %v, want compaction failure", err)
176 }
177 payload := resultErrorPayload(err)
178 if payload["code"] != "turn_failed" {
179 t.Fatalf("result error payload = %#v, want turn_failed", payload)
180 }
181}
182
183func TestPromptWithContextRecoveryKeepsSinglePromptIntact(t *testing.T) {
184 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
185 agent := core.NewAgent(client, "test-model", "", nil)
186
187 _, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, nil)
188 if !provider.IsContextOverflowError(err) {
189 t.Fatalf("promptWithContextRecovery error = %v, want context overflow", err)
190 }
191 client.mu.Lock()
192 calls := client.calls
193 requests := append([]provider.Request(nil), client.requests...)
194 client.mu.Unlock()
195 if len(requests) != 1 || workerRequestUserTextCount(requests[0], "finish the task") != 1 {
196 t.Fatalf("initial request does not retain the prompt: %#v", requests)
197 }
198 if len(agent.Messages()) != 1 || workerRequestMessageTextCount(agent.Messages(), "finish the task") != 1 {
199 t.Fatalf("single prompt was not retained: %#v", agent.Messages())
200 }
201 if calls != 1 {
202 t.Fatalf("provider calls = %d, want 1 without compaction", calls)
203 }
204}
205
206func TestPromptWithContextRecoveryRetainsPromptWhenPersistenceFails(t *testing.T) {
207 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
208 history := []provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}}
209 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
210 if err != nil {
211 t.Fatalf("create session: %v", err)
212 }
213 if err := session.AppendMessage(history[0]); err != nil {
214 t.Fatalf("append history: %v", err)
215 }
216 agent := core.NewAgent(client, "test-model", "", nil)
217 agent.SetMessages(history)
218
219 outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, func([]provider.Message) error {
220 return errors.New("persist failed")
221 })
222 if err == nil || !strings.Contains(err.Error(), "persist compacted transcript") {
223 t.Fatalf("promptWithContextRecovery error = %v, want persistence failure", err)
224 }
225 if got := workerRequestMessageTextCount(agent.Messages(), "prior context"); got != 1 || workerRequestMessageTextCount(agent.Messages(), "finish the task") != 1 {
226 t.Fatalf("agent transcript lost the pending task after persistence failure: %#v", agent.Messages())
227 }
228 WriteNewTranscript(agent, session, outputStart)
229 if err := session.Close(); err != nil {
230 t.Fatalf("close session: %v", err)
231 }
232 reopened, messages, err := core.OpenSession(session.Path)
233 if err != nil {
234 t.Fatalf("reopen session: %v", err)
235 }
236 t.Cleanup(func() {
237 if err := reopened.Close(); err != nil {
238 t.Errorf("close reopened session: %v", err)
239 }
240 })
241 if got := workerRequestMessageTextCount(messages, "finish the task"); got != 1 {
242 t.Fatalf("persisted prompt count = %d, want 1: %#v", got, messages)
243 }
244 client.mu.Lock()
245 calls := client.calls
246 client.mu.Unlock()
247 if calls != 2 {
248 t.Fatalf("provider calls = %d, want initial prompt and compaction only", calls)
249 }
250}
251
252func TestPromptWithContextRecoveryPersistsCompaction(t *testing.T) {
253 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
254 history := []provider.Message{
255 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
256 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
257 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
258 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
259 }
260 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
261 if err != nil {
262 t.Fatalf("create session: %v", err)
263 }
264 for _, message := range history {
265 if err := session.AppendMessage(message); err != nil {
266 t.Fatalf("append history: %v", err)
267 }
268 }
269
270 agent := core.NewAgent(client, "test-model", "", nil)
271 agent.SetMessages(history)
272 outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, session.AppendCompaction)
273 if err != nil {
274 t.Fatalf("promptWithContextRecovery returned %v", err)
275 }
276 WriteNewTranscript(agent, session, outputStart)
277 if err := session.Close(); err != nil {
278 t.Fatalf("close session: %v", err)
279 }
280
281 records, err := os.ReadFile(session.Path)
282 if err != nil {
283 t.Fatalf("read session: %v", err)
284 }
285 compactions := 0
286 for _, record := range bytes.Split(bytes.TrimSpace(records), []byte("\n")) {
287 var entry struct {
288 Type string `json:"type"`
289 }
290 if err := json.Unmarshal(record, &entry); err != nil {
291 t.Fatalf("decode session record: %v", err)
292 }
293 if entry.Type == "compaction" {
294 compactions++
295 }
296 }
297 if compactions != 1 {
298 t.Fatalf("compaction records = %d, want 1", compactions)
299 }
300
301 reopened, messages, err := core.OpenSession(session.Path)
302 if err != nil {
303 t.Fatalf("reopen session: %v", err)
304 }
305 t.Cleanup(func() {
306 if err := reopened.Close(); err != nil {
307 t.Errorf("close reopened session: %v", err)
308 }
309 })
310 if got := workerRequestMessageTextCount(messages, "finish the task"); got != 1 {
311 t.Fatalf("persisted prompt count = %d, want 1: %#v", got, messages)
312 }
313 if got := finalAssistantText(messages); got != "two\nfour\nrecovered" {
314 t.Fatalf("persisted assistant text = %q, want prior tail plus recovered result", got)
315 }
316}
317
318func workerRequestMessageTextCount(messages []provider.Message, want string) int {
319 count := 0
320 for _, message := range messages {
321 for _, block := range message.Content {
322 if text, ok := block.(provider.TextBlock); ok && text.Text == want {
323 count++
324 }
325 }
326 }
327 return count
328}
329
330func workerRequestContainsText(req provider.Request, want string) bool {
331 for _, message := range req.Messages {
332 for _, block := range message.Content {
333 if text, ok := block.(provider.TextBlock); ok && strings.Contains(text.Text, want) {
334 return true
335 }
336 }
337 }
338 return false
339}
340
341func workerRequestUserTextCount(req provider.Request, want string) int {
342 count := 0
343 for _, message := range req.Messages {
344 if message.Role != provider.RoleUser {
345 continue
346 }
347 for _, block := range message.Content {
348 if text, ok := block.(provider.TextBlock); ok && text.Text == want {
349 count++
350 }
351 }
352 }
353 return count
354}
355
356func TestResultErrorPayloadSanitizesContextOverflow(t *testing.T) {
357 payload := resultErrorPayload(errors.New("provider error: input exceeds the context window; echoed task text"))
358 if payload["code"] != "context_limit" {
359 t.Fatalf("error code = %q, want context_limit", payload["code"])
360 }
361 if payload["message"] != subagentContextLimitMessage {
362 t.Fatalf("error message = %q, want %q", payload["message"], subagentContextLimitMessage)
363 }
364}
365
366func TestSubagentEventDataSanitizesContextOverflow(t *testing.T) {
367 payload := subagentEventData(core.EvTurnEnd{Err: errors.New("provider error: input exceeds the context window; echoed task text")})
368 if payload["error"] != subagentContextLimitMessage {
369 t.Fatalf("event error = %q, want %q", payload["error"], subagentContextLimitMessage)
370 }
371}
372
373// TestSupervisorEmitterMirrorDormantUntilStdoutBreaks regresses the
374// "everything is doubled after reopening a subagent agent" bug.
375//
376// Symptom: events.jsonl held two copies of every event because the
377// child mirrored each event to disk AND the supervisor parsed the
378// child's stdout and appended each event to disk too. On next zut
379// launch the replay produced two transcript lines per real one.
380//
381// Fix invariant: while stdout writes succeed (i.e. the supervisor is
382// alive on the other end of the pipe), the child's mirror writes
383// NOTHING. Only when a stdout write fails (broken pipe → orphan)
384// does the mirror take over so events still get persisted.
385func TestSupervisorEmitterMirrorDormantUntilStdoutBreaks(t *testing.T) {
386 // Real *os.File for the emitter's stdout-equivalent so the
387 // emitter's write() path (which expects *os.File) actually runs.
388 stdoutPath := filepath.Join(t.TempDir(), "stdout.fifo")
389 stdoutFile, err := os.Create(stdoutPath)
390 if err != nil {
391 t.Fatalf("create stdout file: %v", err)
392 }
393 defer stdoutFile.Close()
394
395 // Mirror writes go to a separate events.jsonl that we can read
396 // at the end to assert how many events the mirror emitted.
397 mirrorPath := filepath.Join(t.TempDir(), "events.jsonl")
398 mirror, err := subagents.OpenEventLog(mirrorPath)
399 if err != nil {
400 t.Fatalf("open mirror: %v", err)
401 }
402 defer mirror.Close()
403
404 em := newSubagentEmitter(stdoutFile, mirror)
405
406 // Healthy stdout: emit three events. Mirror must stay empty.
407 em.emit("turn_start", map[string]any{"step": 1})
408 em.emit("assistant_message", map[string]any{"text": "hi"})
409 em.emit("turn_end", map[string]any{"step": 1})
410
411 got, err := subagents.ReadEventLog(mirrorPath)
412 if err != nil {
413 t.Fatalf("read mirror: %v", err)
414 }
415 if len(got) != 0 {
416 t.Fatalf("mirror wrote %d events while supervisor was alive; want 0 (every event would otherwise double on the next reload)\n%+v",
417 len(got), got)
418 }
419
420 // Simulate supervisor death: close stdout so the next Write
421 // returns EBADF / broken pipe. The emitter must flip into
422 // orphan mode and start writing through the mirror.
423 if err := stdoutFile.Close(); err != nil {
424 t.Fatalf("close stdout: %v", err)
425 }
426
427 em.emit("assistant_message", map[string]any{"text": "after orphan"})
428 em.emit("turn_end", map[string]any{"step": 2})
429
430 got, err = subagents.ReadEventLog(mirrorPath)
431 if err != nil {
432 t.Fatalf("read mirror post-orphan: %v", err)
433 }
434 if len(got) < 2 {
435 t.Fatalf("mirror failed to take over after stdout died: got %d events", len(got))
436 }
437 if got[len(got)-1].Type != "turn_end" {
438 t.Errorf("last mirrored event type = %q; want turn_end", got[len(got)-1].Type)
439 }
440}
441
442// TestSupervisorEmitterLargeResultIsNotDuplicatedOnWire regresses the
443// large-result wire duplication bug.
444func TestSupervisorEmitterLargeResultIsNotDuplicatedOnWire(t *testing.T) {
445 path := filepath.Join(t.TempDir(), "events.jsonl")
446 file, err := os.Create(path)
447 if err != nil {
448 t.Fatal(err)
449 }
450 defer file.Close()
451
452 em := newSubagentEmitter(file, nil)
453 em.setProtocolIdentity("agent-1")
454 em.emit("turn.result", map[string]any{
455 "status": "succeeded",
456 "turn_id": "turn-1",
457 "output": strings.Repeat("x", 500_000),
458 })
459 if err := file.Close(); err != nil {
460 t.Fatal(err)
461 }
462 wire, err := os.ReadFile(path)
463 if err != nil {
464 t.Fatal(err)
465 }
466 if len(wire) >= 900_000 {
467 t.Fatalf("large result appears duplicated on wire: %d bytes", len(wire))
468 }
469 var object map[string]any
470 if err := json.Unmarshal(bytes.TrimSpace(wire), &object); err != nil {
471 t.Fatal(err)
472 }
473 if _, exists := object["output"]; exists {
474 t.Fatal("large output was flattened a second time")
475 }
476 payload, ok := object["payload"].(map[string]any)
477 if !ok || len(payload["output"].(string)) != 500_000 {
478 t.Fatal("canonical payload lost the complete bounded output")
479 }
480 events, err := subagents.ReadEventLog(path)
481 if err != nil || len(events) != 1 {
482 t.Fatalf("supervisor event parser could not recover the large result: events=%d err=%v", len(events), err)
483 }
484 if output, ok := events[0].Data["output"].(string); !ok || len(output) != 500_000 {
485 t.Fatal("supervisor parser lost the large result payload")
486 }
487}
488
489// TestSupervisorEmitterStdoutShapeMatchesSupervisorParser pins the
490// wire-format contract: each emitted event lands on stdout as one
491// JSON object per line with type+time at top level alongside the
492// data fields. The supervisor's runner parses this exact shape.
493func TestSupervisorEmitterStdoutShapeMatchesSupervisorParser(t *testing.T) {
494 // Pipe so we can read what the emitter wrote.
495 r, w, err := os.Pipe()
496 if err != nil {
497 t.Fatalf("pipe: %v", err)
498 }
499 defer r.Close()
500
501 em := newSubagentEmitter(w, nil)
502 em.emit("turn_start", map[string]any{"step": 1})
503 _ = w.Close()
504
505 body, err := io.ReadAll(r)
506 if err != nil {
507 t.Fatalf("read pipe: %v", err)
508 }
509 // One trailing newline => one event line.
510 lines := bytes.Split(bytes.TrimRight(body, "\n"), []byte("\n"))
511 if len(lines) != 1 {
512 t.Fatalf("expected 1 event line, got %d: %q", len(lines), body)
513 }
514 var object map[string]any
515 if err := json.Unmarshal(lines[0], &object); err != nil {
516 t.Fatalf("not valid json: %v\n%s", err, lines[0])
517 }
518 if object["type"] != "turn.started" {
519 t.Errorf("type field missing or wrong: %v", object["type"])
520 }
521 if _, ok := object["timestamp"].(string); !ok {
522 t.Errorf("timestamp field missing: %v", object["timestamp"])
523 }
524 payload, ok := object["payload"].(map[string]any)
525 if !ok || payload["step"] != float64(1) {
526 t.Errorf("payload step field missing or wrong: %v", object["payload"])
527 }
528}