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