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