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