1package agent
2
3import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "errors"
8 "io"
9 "os"
10 "path/filepath"
11 "reflect"
12 "strings"
13 "sync"
14 "testing"
15 "time"
16
17 "github.com/bnema/zut/packages/agent/modes"
18 "github.com/bnema/zut/packages/agent/subagents"
19 "github.com/bnema/zut/packages/core"
20 "github.com/bnema/zut/packages/provider"
21)
22
23func TestSubagentTurnContextResetsDeadlineAndHonorsShutdown(t *testing.T) {
24 expired, stopExpired := subagentTurnContext(context.Background(), time.Nanosecond)
25 defer stopExpired()
26 <-expired.Done()
27 if !errors.Is(expired.Err(), context.DeadlineExceeded) {
28 t.Fatalf("expired turn error = %v, want context deadline exceeded", expired.Err())
29 }
30
31 parent, stopParent := context.WithCancel(context.Background())
32 first, stopFirst := subagentTurnContext(parent, time.Hour)
33 if _, ok := first.Deadline(); !ok {
34 t.Fatal("active turn context has no deadline")
35 }
36 stopFirst()
37 if !errors.Is(first.Err(), context.Canceled) {
38 t.Fatalf("first turn error = %v, want context canceled", first.Err())
39 }
40
41 second, stopSecond := subagentTurnContext(parent, time.Hour)
42 defer stopSecond()
43 select {
44 case <-second.Done():
45 t.Fatalf("new turn inherited previous turn cancellation: %v", second.Err())
46 default:
47 }
48 stopParent()
49 <-second.Done()
50 if !errors.Is(second.Err(), context.Canceled) {
51 t.Fatalf("shutdown error = %v, want context canceled", second.Err())
52 }
53}
54
55func TestResultErrorPayloadClassifiesTurnDeadline(t *testing.T) {
56 payload := resultErrorPayload(context.DeadlineExceeded)
57 if payload["code"] != "deadline_exceeded" {
58 t.Fatalf("deadline error code = %v, want deadline_exceeded", payload["code"])
59 }
60}
61
62func TestShutdownExitPayloadIncludesOnlyKnownOrigin(t *testing.T) {
63 payload := shutdownExitPayload(subagents.ShutdownOriginSession)
64 if payload["reason"] != "shutdown" || payload["origin"] != subagents.ShutdownOriginSession {
65 t.Fatalf("known shutdown payload = %#v", payload)
66 }
67 payload = shutdownExitPayload(subagents.ShutdownOrigin("private detail"))
68 if _, ok := payload["origin"]; ok {
69 t.Fatalf("unknown shutdown origin was persisted: %#v", payload)
70 }
71}
72
73// TestWorkerOutputLimitsUseSupervisorPolicy verifies that the child reads
74// the effective output caps propagated by the supervisor and retains safe
75// defaults for standalone workers.
76func TestWorkerOutputLimitsUseSupervisorPolicy(t *testing.T) {
77 t.Setenv("ZUT_SUBAGENT_MAX_OUTPUT_BYTES", "123")
78 t.Setenv("ZUT_SUBAGENT_MAX_OUTPUT_LINES", "7")
79 bytesLimit, linesLimit := workerOutputLimits()
80 if bytesLimit != 123 || linesLimit != 7 {
81 t.Fatalf("worker output limits = (%d, %d), want (123, 7)", bytesLimit, linesLimit)
82 }
83 t.Setenv("ZUT_SUBAGENT_MAX_OUTPUT_BYTES", "invalid")
84 t.Setenv("ZUT_SUBAGENT_MAX_OUTPUT_LINES", "0")
85 bytesLimit, linesLimit = workerOutputLimits()
86 if bytesLimit != 500_000 || linesLimit != 5_000 {
87 t.Fatalf("invalid worker output limits = (%d, %d), want defaults", bytesLimit, linesLimit)
88 }
89}
90
91type subagentContextRecoveryClient struct {
92 mu sync.Mutex
93 calls int
94 requests []provider.Request
95 retried chan provider.Request
96 compactionErr error
97 retryOverflow bool
98}
99
100func (c *subagentContextRecoveryClient) Name() string { return "subagent-context-recovery-test" }
101
102func (c *subagentContextRecoveryClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
103 c.mu.Lock()
104 c.calls++
105 call := c.calls
106 c.requests = append(c.requests, req)
107 c.mu.Unlock()
108
109 events := make(chan provider.Event, 2)
110 go func() {
111 defer close(events)
112 switch call {
113 case 1:
114 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("provider error: input exceeds the context window")}
115 case 2:
116 if c.compactionErr != nil {
117 events <- provider.EventDone{Stop: provider.StopError, Err: c.compactionErr}
118 return
119 }
120 events <- provider.EventTextDelta{Delta: "summary"}
121 events <- provider.EventDone{Stop: provider.StopEnd}
122 default:
123 c.retried <- req
124 if c.retryOverflow {
125 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("context window exceeded after compaction")}
126 return
127 }
128 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
129 Role: provider.RoleAssistant,
130 Content: []provider.Content{provider.TextBlock{Text: "recovered"}},
131 }}
132 }
133 }()
134 return events, nil
135}
136
137type subagentProactiveClient struct {
138 mu sync.Mutex
139 requests []provider.Request
140}
141
142func (c *subagentProactiveClient) Name() string { return "subagent-proactive-test" }
143
144func (c *subagentProactiveClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
145 c.mu.Lock()
146 c.requests = append(c.requests, req)
147 c.mu.Unlock()
148 events := make(chan provider.Event, 2)
149 go func() {
150 defer close(events)
151 if strings.Contains(req.System, "context summarization assistant") {
152 events <- provider.EventTextDelta{Delta: "worker summary"}
153 events <- provider.EventDone{Stop: provider.StopEnd}
154 return
155 }
156 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
157 Role: provider.RoleAssistant,
158 Content: []provider.Content{provider.TextBlock{Text: "continued same task"}},
159 }}
160 }()
161 return events, nil
162}
163
164func TestPromptWithContextRecoveryCompactsAndContinues(t *testing.T) {
165 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
166 agent := core.NewAgent(client, "test-model", "", nil)
167 agent.SetMessages([]provider.Message{
168 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
169 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
170 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
171 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
172 })
173
174 recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{})
175 if err != nil {
176 t.Fatalf("promptWithContextRecovery returned %v", err)
177 }
178 if !recovery.Compacted {
179 t.Fatal("promptWithContextRecovery did not report compaction")
180 }
181 if got := finalAssistantText(agent.Messages()[recovery.OutputStart:]); got != "recovered" {
182 t.Fatalf("final assistant text = %q, want recovered", got)
183 }
184
185 request := <-client.retried
186 if got := workerRequestUserTextCount(request, "finish the task"); got != 1 {
187 t.Fatalf("retried request contains prompt %d times, want 1: %#v", got, request.Messages)
188 }
189 client.mu.Lock()
190 calls := client.calls
191 requests := append([]provider.Request(nil), client.requests...)
192 client.mu.Unlock()
193 if calls != 3 || len(requests) != 3 {
194 t.Fatalf("provider calls/requests = %d/%d, want 3/3", calls, len(requests))
195 }
196 if !workerRequestContainsText(requests[1], "<conversation>") || !workerRequestContainsText(requests[1], "one") {
197 t.Fatalf("compaction request does not include the older transcript: %#v", requests[1].Messages)
198 }
199}
200
201func TestPromptWithContextRecoveryStopsAfterOneRetry(t *testing.T) {
202 client := &subagentContextRecoveryClient{
203 retried: make(chan provider.Request, 1),
204 retryOverflow: true,
205 }
206 agent := core.NewAgent(client, "test-model", "", nil)
207 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
208
209 _, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{})
210 if !provider.IsContextOverflowError(err) {
211 t.Fatalf("promptWithContextRecovery error = %v, want context overflow", err)
212 }
213 <-client.retried
214 client.mu.Lock()
215 calls := client.calls
216 client.mu.Unlock()
217 if calls != 3 {
218 t.Fatalf("provider calls = %d, want one initial prompt, compaction, and one continuation", calls)
219 }
220}
221
222func TestPromptWithContextRecoveryDoesNotCompactAfterCancellation(t *testing.T) {
223 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
224 agent := core.NewAgent(client, "test-model", "", nil)
225 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
226 ctx, cancel := context.WithCancel(context.Background())
227 defer cancel()
228
229 _, err := promptWithContextRecovery(ctx, agent, "finish the task", func(event core.AgentEvent) {
230 if turnEnd, ok := event.(core.EvTurnEnd); ok && provider.IsContextOverflowError(turnEnd.Err) {
231 cancel()
232 }
233 }, modes.ContextRecoveryOptions{})
234 if !provider.IsContextOverflowError(err) {
235 t.Fatalf("promptWithContextRecovery error = %v, want context overflow", err)
236 }
237 if err := ctx.Err(); !errors.Is(err, context.Canceled) {
238 t.Fatalf("context error = %v, want %v", err, context.Canceled)
239 }
240 client.mu.Lock()
241 calls := client.calls
242 client.mu.Unlock()
243 if calls != 1 {
244 t.Fatalf("provider calls = %d, want only the initial prompt after cancellation", calls)
245 }
246}
247
248func TestPromptWithContextRecoveryPropagatesCompactionFailure(t *testing.T) {
249 client := &subagentContextRecoveryClient{
250 retried: make(chan provider.Request, 1),
251 compactionErr: errors.New("compaction provider unavailable"),
252 }
253 agent := core.NewAgent(client, "test-model", "", nil)
254 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
255
256 _, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{})
257 if err == nil || !strings.Contains(err.Error(), "compact transcript after initial overflow") {
258 t.Fatalf("promptWithContextRecovery error = %v, want compaction failure", err)
259 }
260 payload := resultErrorPayload(err)
261 if payload["code"] != "turn_failed" {
262 t.Fatalf("result error payload = %#v, want turn_failed", payload)
263 }
264}
265
266func TestPromptWithContextRecoveryKeepsSinglePromptIntact(t *testing.T) {
267 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
268 agent := core.NewAgent(client, "test-model", "", nil)
269
270 _, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{})
271 if !provider.IsContextOverflowError(err) {
272 t.Fatalf("promptWithContextRecovery error = %v, want context overflow", err)
273 }
274 client.mu.Lock()
275 calls := client.calls
276 requests := append([]provider.Request(nil), client.requests...)
277 client.mu.Unlock()
278 if len(requests) != 1 || workerRequestUserTextCount(requests[0], "finish the task") != 1 {
279 t.Fatalf("initial request does not retain the prompt: %#v", requests)
280 }
281 if len(agent.Messages()) != 1 || workerRequestMessageTextCount(agent.Messages(), "finish the task") != 1 {
282 t.Fatalf("single prompt was not retained: %#v", agent.Messages())
283 }
284 if calls != 1 {
285 t.Fatalf("provider calls = %d, want 1 without compaction", calls)
286 }
287}
288
289func TestPromptWithContextRecoveryRetainsPromptWhenPersistenceFails(t *testing.T) {
290 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
291 history := []provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}}
292 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
293 if err != nil {
294 t.Fatalf("create session: %v", err)
295 }
296 if err := session.AppendMessage(history[0]); err != nil {
297 t.Fatalf("append history: %v", err)
298 }
299 agent := core.NewAgent(client, "test-model", "", nil)
300 agent.SetMessages(history)
301
302 recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{
303 PersistCompaction: func([]provider.Message) error { return errors.New("persist failed") },
304 })
305 if err == nil || !strings.Contains(err.Error(), "persist compacted transcript") {
306 t.Fatalf("promptWithContextRecovery error = %v, want persistence failure", err)
307 }
308 if got := workerRequestMessageTextCount(agent.Messages(), "prior context"); got != 1 || workerRequestMessageTextCount(agent.Messages(), "finish the task") != 1 {
309 t.Fatalf("agent transcript lost the pending task after persistence failure: %#v", agent.Messages())
310 }
311 if err := WriteNewTranscript(agent, session, recovery.OutputStart); err != nil {
312 t.Fatalf("persist transcript: %v", err)
313 }
314 if err := session.Close(); err != nil {
315 t.Fatalf("close session: %v", err)
316 }
317 reopened, messages, err := core.OpenSession(session.Path)
318 if err != nil {
319 t.Fatalf("reopen session: %v", err)
320 }
321 t.Cleanup(func() {
322 if err := reopened.Close(); err != nil {
323 t.Errorf("close reopened session: %v", err)
324 }
325 })
326 if got := workerRequestMessageTextCount(messages, "finish the task"); got != 1 {
327 t.Fatalf("persisted prompt count = %d, want 1: %#v", got, messages)
328 }
329 client.mu.Lock()
330 calls := client.calls
331 client.mu.Unlock()
332 if calls != 2 {
333 t.Fatalf("provider calls = %d, want initial prompt and compaction only", calls)
334 }
335}
336
337func TestPromptWithContextRecoveryPersistsCompaction(t *testing.T) {
338 client := &subagentContextRecoveryClient{retried: make(chan provider.Request, 1)}
339 history := []provider.Message{
340 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
341 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
342 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
343 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
344 }
345 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
346 if err != nil {
347 t.Fatalf("create session: %v", err)
348 }
349 for _, message := range history {
350 if err := session.AppendMessage(message); err != nil {
351 t.Fatalf("append history: %v", err)
352 }
353 }
354
355 agent := core.NewAgent(client, "test-model", "", nil)
356 agent.SetMessages(history)
357 recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{
358 PersistCompaction: session.AppendCompaction,
359 })
360 if err != nil {
361 t.Fatalf("promptWithContextRecovery returned %v", err)
362 }
363 if err := WriteNewTranscript(agent, session, recovery.OutputStart); err != nil {
364 t.Fatalf("persist transcript: %v", err)
365 }
366 if err := session.Close(); err != nil {
367 t.Fatalf("close session: %v", err)
368 }
369
370 records, err := os.ReadFile(session.Path)
371 if err != nil {
372 t.Fatalf("read session: %v", err)
373 }
374 compactions := 0
375 for _, record := range bytes.Split(bytes.TrimSpace(records), []byte("\n")) {
376 var entry struct {
377 Type string `json:"type"`
378 }
379 if err := json.Unmarshal(record, &entry); err != nil {
380 t.Fatalf("decode session record: %v", err)
381 }
382 if entry.Type == "compaction" {
383 compactions++
384 }
385 }
386 if compactions != 1 {
387 t.Fatalf("compaction records = %d, want 1", compactions)
388 }
389
390 reopened, messages, err := core.OpenSession(session.Path)
391 if err != nil {
392 t.Fatalf("reopen session: %v", err)
393 }
394 t.Cleanup(func() {
395 if err := reopened.Close(); err != nil {
396 t.Errorf("close reopened session: %v", err)
397 }
398 })
399 if got := workerRequestMessageTextCount(messages, "finish the task"); got != 1 {
400 t.Fatalf("persisted prompt count = %d, want 1: %#v", got, messages)
401 }
402 if got := finalAssistantText(messages); got != "two\nfour\nrecovered" {
403 t.Fatalf("persisted assistant text = %q, want prior tail plus recovered result", got)
404 }
405}
406
407func TestWorkerUsesConfiguredAutoCompactionThreshold(t *testing.T) {
408 t.Setenv("ZUT_HOME", t.TempDir())
409 threshold := 70
410 if err := SaveConfig(Config{AutoCompactThreshold: &threshold}); err != nil {
411 t.Fatalf("save config: %v", err)
412 }
413 cfg, err := loadSubagentWorkerConfig()
414 if err != nil {
415 t.Fatalf("load worker config: %v", err)
416 }
417
418 client := &subagentProactiveClient{}
419 agent := core.NewAgent(client, "model", "", nil)
420 agent.SetMessages([]provider.Message{
421 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "history"}}},
422 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},
423 })
424 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 75})
425 recovery, err := promptWithContextRecovery(context.Background(), agent, "finish exact task", nil, modes.ContextRecoveryOptions{
426 ContextWindow: 100,
427 AutoCompactThreshold: cfg.AutoCompactThreshold,
428 })
429 if err != nil {
430 t.Fatalf("promptWithContextRecovery: %v", err)
431 }
432 if got := finalAssistantText(agent.Messages()[recovery.OutputStart:]); got != "continued same task" {
433 t.Fatalf("continued output = %q", got)
434 }
435 client.mu.Lock()
436 requests := append([]provider.Request(nil), client.requests...)
437 client.mu.Unlock()
438 if len(requests) != 2 || workerRequestUserTextCount(requests[1], "finish exact task") != 1 {
439 t.Fatalf("requests = %#v, want compaction then the same task once", requests)
440 }
441}
442
443func TestWorkerConfigLoadRejectsMalformedConfig(t *testing.T) {
444 t.Setenv("ZUT_HOME", t.TempDir())
445 if err := os.MkdirAll(filepath.Dir(ConfigPath()), 0o755); err != nil {
446 t.Fatal(err)
447 }
448 if err := os.WriteFile(ConfigPath(), []byte("{malformed"), 0o600); err != nil {
449 t.Fatal(err)
450 }
451 if _, err := loadSubagentWorkerConfig(); err == nil || !strings.Contains(err.Error(), "load subagent worker config") {
452 t.Fatalf("loadSubagentWorkerConfig error = %v, want parse failure", err)
453 }
454}
455
456type subagentSummaryUsageClient struct{}
457
458func (c *subagentSummaryUsageClient) Name() string { return "subagent-summary-usage-test" }
459
460func (c *subagentSummaryUsageClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
461 events := make(chan provider.Event, 3)
462 go func() {
463 defer close(events)
464 if strings.Contains(req.System, "context summarization assistant") {
465 events <- provider.EventUsage{Usage: provider.Usage{InputTokens: 3, OutputTokens: 1, CostUSD: 0.03}}
466 events <- provider.EventTextDelta{Delta: "worker summary"}
467 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
468 Role: provider.RoleAssistant,
469 Content: []provider.Content{provider.TextBlock{Text: "worker summary"}},
470 }}
471 return
472 }
473 events <- provider.EventUsage{Usage: provider.Usage{InputTokens: 5, OutputTokens: 2, CostUSD: 0.05}}
474 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
475 Role: provider.RoleAssistant,
476 Content: []provider.Content{provider.TextBlock{Text: "continued"}},
477 }}
478 }()
479 return events, nil
480}
481
482func TestPersistWorkerSummaryUsageReopensWithCumulativeUsage(t *testing.T) {
483 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
484 if err != nil {
485 t.Fatal(err)
486 }
487 history := []provider.Message{
488 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "old task"}}},
489 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "old answer"}}},
490 }
491 for _, message := range history {
492 if err := session.AppendMessage(message); err != nil {
493 t.Fatal(err)
494 }
495 }
496 baseline := provider.Usage{InputTokens: 10, OutputTokens: 2, CostUSD: 0.1}
497 if err := session.AppendUsage(baseline, baseline); err != nil {
498 t.Fatal(err)
499 }
500 agent := core.NewAgent(&subagentSummaryUsageClient{}, "model", "", nil)
501 agent.SetMessages(history)
502 agent.SeedCost(baseline)
503 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 90})
504 recovery, err := promptWithContextRecovery(context.Background(), agent, "continue task", nil, modes.ContextRecoveryOptions{
505 ContextWindow: 100,
506 })
507 if err != nil {
508 t.Fatalf("promptWithContextRecovery: %v", err)
509 }
510 if recovery.OutputStart != len(history) || !recovery.Compacted {
511 t.Fatalf("compaction recovery = %+v, want output start %d and compacted", recovery, len(history))
512 }
513 expected := provider.Usage{InputTokens: 18, OutputTokens: 5, CostUSD: 0.18}
514 if !reflect.DeepEqual(agent.Cost(), expected) {
515 t.Fatalf("worker cumulative usage = %#v, want %#v", agent.Cost(), expected)
516 }
517 if err := persistWorkerTranscript(context.Background(), agent, session, recovery.OutputStart, recovery.Compacted, workerTranscriptState{messages: history, cost: baseline}); err != nil {
518 t.Fatalf("persistWorkerTranscript: %v", err)
519 }
520 if err := session.Close(); err != nil {
521 t.Fatal(err)
522 }
523 cumulative, lastTurn, err := core.SessionUsageDetail(session.Path)
524 if err != nil {
525 t.Fatal(err)
526 }
527 if !reflect.DeepEqual(cumulative, expected) {
528 t.Fatalf("reopened cumulative usage = %#v, want %#v", cumulative, expected)
529 }
530 wantLast := provider.Usage{InputTokens: 8, OutputTokens: 3, CostUSD: expected.CostUSD - baseline.CostUSD}
531 if !reflect.DeepEqual(lastTurn, wantLast) {
532 t.Fatalf("reopened last-turn usage = %#v, want %#v", lastTurn, wantLast)
533 }
534}
535
536func TestPersistWorkerCompactedTranscriptReopensWithTurnAndUsage(t *testing.T) {
537 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
538 if err != nil {
539 t.Fatal(err)
540 }
541 baselineMessage := provider.Message{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "baseline"}}}
542 if err := session.AppendMessage(baselineMessage); err != nil {
543 t.Fatal(err)
544 }
545 baselineUsage := provider.Usage{InputTokens: 10, OutputTokens: 2}
546 if err := session.AppendUsage(baselineUsage, baselineUsage); err != nil {
547 t.Fatal(err)
548 }
549
550 agent := core.NewAgent(&subagentProactiveClient{}, "model", "", nil)
551 agent.SetMessages([]provider.Message{
552 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "summary"}}},
553 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "current task"}}},
554 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "current answer"}}},
555 })
556 cumulative := provider.Usage{InputTokens: 30, OutputTokens: 8, CacheReadTokens: 4}
557 agent.SeedCost(cumulative)
558 if err := persistWorkerTranscript(context.Background(), agent, session, 1, true, workerTranscriptState{
559 messages: []provider.Message{baselineMessage},
560 cost: baselineUsage,
561 }); err != nil {
562 t.Fatalf("persistWorkerTranscript: %v", err)
563 }
564 if err := session.Close(); err != nil {
565 t.Fatal(err)
566 }
567
568 snapshot, err := core.ReadSessionSnapshot(session.Path)
569 if err != nil {
570 t.Fatal(err)
571 }
572 if workerRequestMessageTextCount(snapshot.Messages, "current task") != 1 || finalAssistantText(snapshot.Messages) != "current answer" {
573 t.Fatalf("reopened messages = %#v, want full continued turn", snapshot.Messages)
574 }
575 if len(snapshot.UsageCheckpoints) == 0 || !reflect.DeepEqual(snapshot.UsageCheckpoints[len(snapshot.UsageCheckpoints)-1].Cumulative, cumulative) {
576 t.Fatalf("usage checkpoints = %#v, want cumulative %#v", snapshot.UsageCheckpoints, cumulative)
577 }
578}
579
580func TestPersistWorkerCompactionFailureRollsBackAndKeepsSessionReopenable(t *testing.T) {
581 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
582 if err != nil {
583 t.Fatal(err)
584 }
585 baselineMessage := provider.Message{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "baseline"}}}
586 if err := session.AppendMessage(baselineMessage); err != nil {
587 t.Fatal(err)
588 }
589 baselineUsage := provider.Usage{InputTokens: 10, OutputTokens: 2}
590 if err := session.AppendUsage(baselineUsage, baselineUsage); err != nil {
591 t.Fatal(err)
592 }
593 if err := session.Close(); err != nil {
594 t.Fatal(err)
595 }
596
597 agent := core.NewAgent(&subagentProactiveClient{}, "model", "", nil)
598 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "unsaved compacted turn"}}}})
599 agent.SeedCost(provider.Usage{InputTokens: 99})
600 baseline := workerTranscriptState{messages: []provider.Message{baselineMessage}, cost: baselineUsage, lastTurn: baselineUsage}
601 if err := persistWorkerTranscript(context.Background(), agent, session, 0, true, baseline); err == nil {
602 t.Fatal("persistWorkerTranscript succeeded on closed session")
603 }
604 if !reflect.DeepEqual(agent.Messages(), baseline.messages) || !reflect.DeepEqual(agent.Cost(), baselineUsage) {
605 t.Fatalf("agent was not rolled back: messages=%#v cost=%#v", agent.Messages(), agent.Cost())
606 }
607 snapshot, err := core.ReadSessionSnapshot(session.Path)
608 if err != nil {
609 t.Fatalf("reopen snapshot: %v", err)
610 }
611 if !reflect.DeepEqual(snapshot.Messages, baseline.messages) {
612 t.Fatalf("reopened transcript = %#v, want baseline %#v", snapshot.Messages, baseline.messages)
613 }
614 if len(snapshot.UsageCheckpoints) != 1 || !reflect.DeepEqual(snapshot.UsageCheckpoints[0].Cumulative, baselineUsage) {
615 t.Fatalf("reopened usage = %#v, want baseline", snapshot.UsageCheckpoints)
616 }
617}
618
619func TestPersistWorkerTranscriptSkipsCanceledNonCompactedTurn(t *testing.T) {
620 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
621 if err != nil {
622 t.Fatal(err)
623 }
624 baselineMessage := provider.Message{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "baseline"}}}
625 if err := session.AppendMessage(baselineMessage); err != nil {
626 t.Fatal(err)
627 }
628 agent := core.NewAgent(&subagentProactiveClient{}, "model", "", nil)
629 agent.SetMessages([]provider.Message{
630 baselineMessage,
631 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "canceled answer"}}},
632 })
633 ctx, cancel := context.WithCancel(context.Background())
634 cancel()
635 if err := persistWorkerTranscript(ctx, agent, session, 1, false, workerTranscriptState{messages: []provider.Message{baselineMessage}}); !errors.Is(err, context.Canceled) {
636 t.Fatalf("persistWorkerTranscript error = %v, want cancellation", err)
637 }
638 if err := session.Close(); err != nil {
639 t.Fatal(err)
640 }
641 snapshot, err := core.ReadSessionSnapshot(session.Path)
642 if err != nil {
643 t.Fatal(err)
644 }
645 if !reflect.DeepEqual(snapshot.Messages, []provider.Message{baselineMessage}) {
646 t.Fatalf("reopened messages = %#v, want baseline only", snapshot.Messages)
647 }
648}
649
650func TestPersistWorkerCompactionHonorsCancellation(t *testing.T) {
651 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
652 if err != nil {
653 t.Fatal(err)
654 }
655 baselineMessage := provider.Message{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "baseline"}}}
656 if err := session.AppendMessage(baselineMessage); err != nil {
657 t.Fatal(err)
658 }
659 baseline := workerTranscriptState{messages: []provider.Message{baselineMessage}}
660 agent := core.NewAgent(&subagentProactiveClient{}, "model", "", nil)
661 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "must not persist"}}}})
662 ctx, cancel := context.WithCancel(context.Background())
663 cancel()
664 if err := persistWorkerTranscript(ctx, agent, session, 0, true, baseline); !errors.Is(err, context.Canceled) {
665 t.Fatalf("persistWorkerTranscript error = %v, want cancellation", err)
666 }
667 if !reflect.DeepEqual(agent.Messages(), baseline.messages) {
668 t.Fatalf("agent messages = %#v, want rollback %#v", agent.Messages(), baseline.messages)
669 }
670 if err := session.Close(); err != nil {
671 t.Fatal(err)
672 }
673 snapshot, err := core.ReadSessionSnapshot(session.Path)
674 if err != nil {
675 t.Fatal(err)
676 }
677 if !reflect.DeepEqual(snapshot.Messages, baseline.messages) {
678 t.Fatalf("reopened messages = %#v, want baseline", snapshot.Messages)
679 }
680}
681
682func TestWriteNewTranscriptReturnsPersistenceErrors(t *testing.T) {
683 session, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
684 if err != nil {
685 t.Fatal(err)
686 }
687 if err := session.AppendMessage(provider.Message{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "baseline"}}}); err != nil {
688 t.Fatal(err)
689 }
690 if err := session.Close(); err != nil {
691 t.Fatal(err)
692 }
693 agent := core.NewAgent(&subagentProactiveClient{}, "model", "", nil)
694 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "new"}}}})
695 if err := WriteNewTranscript(agent, session, 0); err == nil || !strings.Contains(err.Error(), "append transcript message") {
696 t.Fatalf("WriteNewTranscript error = %v, want append failure", err)
697 }
698}
699
700func workerRequestMessageTextCount(messages []provider.Message, want string) int {
701 count := 0
702 for _, message := range messages {
703 for _, block := range message.Content {
704 if text, ok := block.(provider.TextBlock); ok && text.Text == want {
705 count++
706 }
707 }
708 }
709 return count
710}
711
712func workerRequestContainsText(req provider.Request, want string) bool {
713 for _, message := range req.Messages {
714 for _, block := range message.Content {
715 if text, ok := block.(provider.TextBlock); ok && strings.Contains(text.Text, want) {
716 return true
717 }
718 }
719 }
720 return false
721}
722
723func workerRequestUserTextCount(req provider.Request, want string) int {
724 count := 0
725 for _, message := range req.Messages {
726 if message.Role != provider.RoleUser {
727 continue
728 }
729 for _, block := range message.Content {
730 if text, ok := block.(provider.TextBlock); ok && text.Text == want {
731 count++
732 }
733 }
734 }
735 return count
736}
737
738func TestResultErrorPayloadSanitizesContextOverflow(t *testing.T) {
739 payload := resultErrorPayload(errors.New("provider error: input exceeds the context window; echoed task text"))
740 if payload["code"] != "context_limit" {
741 t.Fatalf("error code = %q, want context_limit", payload["code"])
742 }
743 if payload["message"] != subagentContextLimitMessage {
744 t.Fatalf("error message = %q, want %q", payload["message"], subagentContextLimitMessage)
745 }
746}
747
748func TestSubagentEventDataSanitizesContextOverflow(t *testing.T) {
749 payload := subagentEventData(core.EvTurnEnd{Err: errors.New("provider error: input exceeds the context window; echoed task text")})
750 if payload["error"] != subagentContextLimitMessage {
751 t.Fatalf("event error = %q, want %q", payload["error"], subagentContextLimitMessage)
752 }
753}
754
755// TestSupervisorEmitterMirrorDormantUntilStdoutBreaks regresses the
756// "everything is doubled after reopening a subagent agent" bug.
757//
758// Symptom: events.jsonl held two copies of every event because the
759// child mirrored each event to disk AND the supervisor parsed the
760// child's stdout and appended each event to disk too. On next zut
761// launch the replay produced two transcript lines per real one.
762//
763// Fix invariant: while stdout writes succeed (i.e. the supervisor is
764// alive on the other end of the pipe), the child's mirror writes
765// NOTHING. Only when a stdout write fails (broken pipe → orphan)
766// does the mirror take over so events still get persisted.
767func TestSupervisorEmitterMirrorDormantUntilStdoutBreaks(t *testing.T) {
768 // Real *os.File for the emitter's stdout-equivalent so the
769 // emitter's write() path (which expects *os.File) actually runs.
770 stdoutPath := filepath.Join(t.TempDir(), "stdout.fifo")
771 stdoutFile, err := os.Create(stdoutPath)
772 if err != nil {
773 t.Fatalf("create stdout file: %v", err)
774 }
775 defer stdoutFile.Close()
776
777 // Mirror writes go to a separate events.jsonl that we can read
778 // at the end to assert how many events the mirror emitted.
779 mirrorPath := filepath.Join(t.TempDir(), "events.jsonl")
780 mirror, err := subagents.OpenEventLog(mirrorPath)
781 if err != nil {
782 t.Fatalf("open mirror: %v", err)
783 }
784 defer mirror.Close()
785
786 em := newSubagentEmitter(stdoutFile, mirror)
787
788 // Healthy stdout: emit three events. Mirror must stay empty.
789 em.emit("turn_start", map[string]any{"step": 1})
790 em.emit("assistant_message", map[string]any{"text": "hi"})
791 em.emit("turn_end", map[string]any{"step": 1})
792
793 got, err := subagents.ReadEventLog(mirrorPath)
794 if err != nil {
795 t.Fatalf("read mirror: %v", err)
796 }
797 if len(got) != 0 {
798 t.Fatalf("mirror wrote %d events while supervisor was alive; want 0 (every event would otherwise double on the next reload)\n%+v",
799 len(got), got)
800 }
801
802 // Simulate supervisor death: close stdout so the next Write
803 // returns EBADF / broken pipe. The emitter must flip into
804 // orphan mode and start writing through the mirror.
805 if err := stdoutFile.Close(); err != nil {
806 t.Fatalf("close stdout: %v", err)
807 }
808
809 em.emit("assistant_message", map[string]any{"text": "after orphan"})
810 em.emit("turn_end", map[string]any{"step": 2})
811
812 got, err = subagents.ReadEventLog(mirrorPath)
813 if err != nil {
814 t.Fatalf("read mirror post-orphan: %v", err)
815 }
816 if len(got) < 2 {
817 t.Fatalf("mirror failed to take over after stdout died: got %d events", len(got))
818 }
819 if got[len(got)-1].Type != "turn_end" {
820 t.Errorf("last mirrored event type = %q; want turn_end", got[len(got)-1].Type)
821 }
822}
823
824// TestSupervisorEmitterLargeResultIsNotDuplicatedOnWire regresses the
825// large-result wire duplication bug.
826func TestSupervisorEmitterLargeResultIsNotDuplicatedOnWire(t *testing.T) {
827 path := filepath.Join(t.TempDir(), "events.jsonl")
828 file, err := os.Create(path)
829 if err != nil {
830 t.Fatal(err)
831 }
832 defer file.Close()
833
834 em := newSubagentEmitter(file, nil)
835 em.setProtocolIdentity("agent-1")
836 em.emit("turn.result", map[string]any{
837 "status": "succeeded",
838 "turn_id": "turn-1",
839 "output": strings.Repeat("x", 500_000),
840 })
841 if err := file.Close(); err != nil {
842 t.Fatal(err)
843 }
844 wire, err := os.ReadFile(path)
845 if err != nil {
846 t.Fatal(err)
847 }
848 if len(wire) >= 900_000 {
849 t.Fatalf("large result appears duplicated on wire: %d bytes", len(wire))
850 }
851 var object map[string]any
852 if err := json.Unmarshal(bytes.TrimSpace(wire), &object); err != nil {
853 t.Fatal(err)
854 }
855 if _, exists := object["output"]; exists {
856 t.Fatal("large output was flattened a second time")
857 }
858 payload, ok := object["payload"].(map[string]any)
859 if !ok || len(payload["output"].(string)) != 500_000 {
860 t.Fatal("canonical payload lost the complete bounded output")
861 }
862 events, err := subagents.ReadEventLog(path)
863 if err != nil || len(events) != 1 {
864 t.Fatalf("supervisor event parser could not recover the large result: events=%d err=%v", len(events), err)
865 }
866 if output, ok := events[0].Data["output"].(string); !ok || len(output) != 500_000 {
867 t.Fatal("supervisor parser lost the large result payload")
868 }
869}
870
871// TestSupervisorEmitterStdoutShapeMatchesSupervisorParser pins the
872// wire-format contract: each emitted event lands on stdout as one
873// JSON object per line with type+time at top level alongside the
874// data fields. The supervisor's runner parses this exact shape.
875func TestSupervisorEmitterStdoutShapeMatchesSupervisorParser(t *testing.T) {
876 // Pipe so we can read what the emitter wrote.
877 r, w, err := os.Pipe()
878 if err != nil {
879 t.Fatalf("pipe: %v", err)
880 }
881 defer r.Close()
882
883 em := newSubagentEmitter(w, nil)
884 em.emit("turn_start", map[string]any{"step": 1})
885 _ = w.Close()
886
887 body, err := io.ReadAll(r)
888 if err != nil {
889 t.Fatalf("read pipe: %v", err)
890 }
891 // One trailing newline => one event line.
892 lines := bytes.Split(bytes.TrimRight(body, "\n"), []byte("\n"))
893 if len(lines) != 1 {
894 t.Fatalf("expected 1 event line, got %d: %q", len(lines), body)
895 }
896 var object map[string]any
897 if err := json.Unmarshal(lines[0], &object); err != nil {
898 t.Fatalf("not valid json: %v\n%s", err, lines[0])
899 }
900 if object["type"] != "turn.started" {
901 t.Errorf("type field missing or wrong: %v", object["type"])
902 }
903 if _, ok := object["timestamp"].(string); !ok {
904 t.Errorf("timestamp field missing: %v", object["timestamp"])
905 }
906 payload, ok := object["payload"].(map[string]any)
907 if !ok || payload["step"] != float64(1) {
908 t.Errorf("payload step field missing or wrong: %v", object["payload"])
909 }
910}
911
912func TestInitialTurnNumberIgnoresNestedProviderTurns(t *testing.T) {
913 path := filepath.Join(t.TempDir(), "events.jsonl")
914 log, err := subagents.OpenEventLog(path)
915 if err != nil {
916 t.Fatal(err)
917 }
918 if err := log.Append(subagents.NewEvent(subagents.EventTurnStarted, map[string]any{
919 "step": float64(1), "lifetime_turns": 1, "current_run_turns": 1,
920 })); err != nil {
921 t.Fatal(err)
922 }
923 if err := log.Append(subagents.NewEvent(subagents.EventTurnStarted, map[string]any{
924 "step": float64(99), "nested_turn": true,
925 })); err != nil {
926 t.Fatal(err)
927 }
928 if err := log.Close(); err != nil {
929 t.Fatal(err)
930 }
931 if got := initialTurnNumber(path); got != 1 {
932 t.Fatalf("initial turn number = %d, want 1", got)
933 }
934}
935
936func TestWorkerTurnBudgetFreshAllowanceOnlyForExplicitRun(t *testing.T) {
937 var budget workerTurnBudget
938
939 if _, lifetime, current, admitted := budget.start(2, false); !admitted || lifetime != 1 || current != 1 {
940 t.Fatalf("first turn = (%d, %d, %v), want (1, 1, true)", lifetime, current, admitted)
941 }
942 // Provider retries, model loops, and compaction stay inside the admitted
943 // turn and therefore do not call start or consume another allowance.
944 if budget.lifetime != 1 || budget.current != 1 {
945 t.Fatalf("nested work changed budget to (lifetime=%d, current=%d)", budget.lifetime, budget.current)
946 }
947 if _, lifetime, current, admitted := budget.start(2, false); !admitted || lifetime != 2 || current != 2 {
948 t.Fatalf("second turn = (%d, %d, %v), want (2, 2, true)", lifetime, current, admitted)
949 }
950 if _, lifetime, current, admitted := budget.start(2, false); admitted || lifetime != 2 || current != 2 {
951 t.Fatalf("exhausted turn = (%d, %d, %v), want (2, 2, false)", lifetime, current, admitted)
952 }
953 if _, lifetime, current, admitted := budget.start(2, true); !admitted || lifetime != 3 || current != 1 {
954 t.Fatalf("explicit resumed turn = (%d, %d, %v), want (3, 1, true)", lifetime, current, admitted)
955 }
956}