1package modes
2
3import (
4 "bufio"
5 "bytes"
6 "context"
7 "encoding/json"
8 "errors"
9 "reflect"
10 "strings"
11 "sync"
12 "testing"
13
14 "github.com/bnema/zut/packages/core"
15 "github.com/bnema/zut/packages/provider"
16)
17
18type hostRecoveryClient struct {
19 mu sync.Mutex
20 calls int
21 retried chan provider.Request
22 retryOverflow bool
23}
24
25func (c *hostRecoveryClient) Name() string { return "host-context-recovery-test" }
26
27type proactiveContextClient struct {
28 mu sync.Mutex
29 requests []provider.Request
30 summaryErr error
31 onSummary func()
32 normalOverflow bool
33}
34
35func (c *proactiveContextClient) Name() string { return "proactive-context-test" }
36
37func (c *proactiveContextClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
38 c.mu.Lock()
39 c.requests = append(c.requests, req)
40 c.mu.Unlock()
41
42 events := make(chan provider.Event, 2)
43 go func() {
44 defer close(events)
45 if strings.Contains(req.System, "context summarization assistant") {
46 if c.onSummary != nil {
47 c.onSummary()
48 }
49 if c.summaryErr != nil {
50 events <- provider.EventDone{Stop: provider.StopError, Err: c.summaryErr}
51 return
52 }
53 events <- provider.EventTextDelta{Delta: "condensed history"}
54 events <- provider.EventDone{Stop: provider.StopEnd}
55 return
56 }
57 if c.normalOverflow {
58 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("context window exceeded")}
59 return
60 }
61 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
62 Role: provider.RoleAssistant,
63 Content: []provider.Content{provider.TextBlock{Text: "completed"}},
64 }}
65 }()
66 return events, nil
67}
68
69func TestPromptWithContextRecoveryProactivelyCompactsCacheOnlyUsage(t *testing.T) {
70 client := &proactiveContextClient{}
71 agent := core.NewAgent(client, "model", "system", nil)
72 history := []provider.Message{
73 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
74 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
75 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
76 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
77 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
78 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "six"}}},
79 }
80 agent.SetMessages(history)
81 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 85})
82
83 var persisted [][]provider.Message
84 result, err := PromptWithContextRecovery(context.Background(), agent, "continue", nil, nil, ContextRecoveryOptions{
85 ContextWindow: 100,
86 PersistCompaction: func(messages []provider.Message) error {
87 persisted = append(persisted, messages)
88 return nil
89 },
90 })
91 if err != nil {
92 t.Fatalf("PromptWithContextRecovery: %v", err)
93 }
94 if !result.Compacted || len(persisted) != 1 {
95 t.Fatalf("compacted/persisted = %t/%d, want true/1", result.Compacted, len(persisted))
96 }
97 if got := recoveredAssistantText(agent.Messages()[result.OutputStart:]); got != "completed" {
98 t.Fatalf("output after compacted boundary = %q, want completed", got)
99 }
100
101 client.mu.Lock()
102 requests := append([]provider.Request(nil), client.requests...)
103 client.mu.Unlock()
104 if len(requests) != 2 {
105 t.Fatalf("provider requests = %d, want compaction then prompt", len(requests))
106 }
107 if !recoveryRequestContainsText(requests[0], "<conversation>") || !recoveryRequestContainsText(requests[0], "one") {
108 t.Fatalf("first request was not compaction over old history: %#v", requests[0])
109 }
110 if got := requestUserTextCount(requests[1], "continue"); got != 1 {
111 t.Fatalf("prompt count after proactive compaction = %d, want 1", got)
112 }
113}
114
115func TestPromptWithContextRecoveryHonorsDisabledThreshold(t *testing.T) {
116 client := &proactiveContextClient{}
117 agent := core.NewAgent(client, "model", "system", nil)
118 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "history"}}}})
119 agent.SeedLastTurnUsage(provider.Usage{InputTokens: 100, CacheReadTokens: 100})
120 disabled := 0
121
122 result, err := PromptWithContextRecovery(context.Background(), agent, "continue", nil, nil, ContextRecoveryOptions{
123 ContextWindow: 100,
124 AutoCompactThreshold: &disabled,
125 })
126 if err != nil {
127 t.Fatalf("PromptWithContextRecovery: %v", err)
128 }
129 if result.Compacted {
130 t.Fatal("disabled threshold unexpectedly compacted")
131 }
132 client.mu.Lock()
133 requests := append([]provider.Request(nil), client.requests...)
134 client.mu.Unlock()
135 if len(requests) != 1 || requestUserTextCount(requests[0], "continue") != 1 {
136 t.Fatalf("requests with disabled threshold = %#v, want one direct prompt", requests)
137 }
138}
139
140func TestPromptWithContextRecoveryRestoresProactiveCompactionOnPersistenceFailure(t *testing.T) {
141 client := &proactiveContextClient{}
142 agent := core.NewAgent(client, "model", "system", nil)
143 history := []provider.Message{
144 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "original"}}},
145 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},
146 }
147 agent.SetMessages(history)
148 agent.SeedLastTurnUsage(provider.Usage{CacheWriteTokens: 90})
149
150 _, err := PromptWithContextRecovery(context.Background(), agent, "must not send", nil, nil, ContextRecoveryOptions{
151 ContextWindow: 100,
152 PersistCompaction: func([]provider.Message) error {
153 return errors.New("disk full")
154 },
155 })
156 if err == nil || !strings.Contains(err.Error(), "persist compacted transcript: disk full") {
157 t.Fatalf("error = %v, want persistence failure", err)
158 }
159 if got := agent.Messages(); !reflect.DeepEqual(got, history) {
160 t.Fatalf("transcript after persistence failure = %#v, want %#v", got, history)
161 }
162 client.mu.Lock()
163 requestCount := len(client.requests)
164 client.mu.Unlock()
165 if requestCount != 1 {
166 t.Fatalf("provider requests = %d, want compaction only", requestCount)
167 }
168}
169
170func TestPromptWithContextRecoveryDoesNotProactivelyCompactSoleTask(t *testing.T) {
171 client := &proactiveContextClient{}
172 agent := core.NewAgent(client, "model", "system", nil)
173 task := provider.Message{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "exact sole task"}}}
174 agent.SetMessages([]provider.Message{task})
175 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 95})
176
177 result, err := PromptWithContextRecovery(context.Background(), agent, "continue", nil, nil, ContextRecoveryOptions{ContextWindow: 100})
178 if err != nil {
179 t.Fatalf("PromptWithContextRecovery: %v", err)
180 }
181 if result.Compacted {
182 t.Fatal("sole task was proactively compacted")
183 }
184 client.mu.Lock()
185 requests := append([]provider.Request(nil), client.requests...)
186 client.mu.Unlock()
187 if len(requests) != 1 || !recoveryRequestContainsText(requests[0], "exact sole task") {
188 t.Fatalf("requests = %#v, want one direct request retaining the sole task", requests)
189 }
190}
191
192func TestPromptWithContextRecoveryFallsBackWhenProactiveSummaryOverflows(t *testing.T) {
193 client := &proactiveContextClient{summaryErr: errors.New("context window exceeded while summarizing")}
194 agent := core.NewAgent(client, "model", "system", nil)
195 history := []provider.Message{
196 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "history"}}},
197 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},
198 }
199 agent.SetMessages(history)
200 agent.SeedLastTurnUsage(provider.Usage{InputTokens: 90})
201
202 result, err := PromptWithContextRecovery(context.Background(), agent, "same task", nil, nil, ContextRecoveryOptions{ContextWindow: 100})
203 if err != nil {
204 t.Fatalf("PromptWithContextRecovery: %v", err)
205 }
206 if result.Compacted {
207 t.Fatal("overflowing proactive summary reported compaction")
208 }
209 client.mu.Lock()
210 requests := append([]provider.Request(nil), client.requests...)
211 client.mu.Unlock()
212 if len(requests) != 2 {
213 t.Fatalf("provider requests = %d, want one summary and one normal prompt", len(requests))
214 }
215 if got := requestUserTextCount(requests[1], "same task"); got != 1 {
216 t.Fatalf("normal recovery prompt count = %d, want 1", got)
217 }
218}
219
220func TestPromptWithContextRecoveryStopsBeforeContinuationWhenPersistenceCancels(t *testing.T) {
221 ctx, cancel := context.WithCancel(context.Background())
222 client := &proactiveContextClient{}
223 agent := core.NewAgent(client, "model", "system", nil)
224 history := []provider.Message{
225 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "history"}}},
226 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},
227 }
228 agent.SetMessages(history)
229 agent.SeedLastTurnUsage(provider.Usage{CacheWriteTokens: 90})
230 persisted := 0
231
232 result, err := PromptWithContextRecovery(ctx, agent, "must not run", nil, nil, ContextRecoveryOptions{
233 ContextWindow: 100,
234 PersistCompaction: func([]provider.Message) error {
235 persisted++
236 cancel()
237 return nil
238 },
239 })
240 if !errors.Is(err, context.Canceled) {
241 t.Fatalf("error = %v, want context cancellation", err)
242 }
243 if !result.Compacted || persisted != 1 {
244 t.Fatalf("compaction result/persistence = %v/%d, want true/1", result.Compacted, persisted)
245 }
246 client.mu.Lock()
247 requestCount := len(client.requests)
248 client.mu.Unlock()
249 if requestCount != 1 {
250 t.Fatalf("provider requests = %d, want summary only", requestCount)
251 }
252}
253
254func TestPromptWithContextRecoveryRollsBackWhenCanceledAfterSummary(t *testing.T) {
255 ctx, cancel := context.WithCancel(context.Background())
256 client := &proactiveContextClient{onSummary: cancel}
257 agent := core.NewAgent(client, "model", "system", nil)
258 history := []provider.Message{
259 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "history"}}},
260 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},
261 }
262 agent.SetMessages(history)
263 agent.SeedLastTurnUsage(provider.Usage{CacheWriteTokens: 90})
264 persisted := 0
265
266 _, err := PromptWithContextRecovery(ctx, agent, "must not run", nil, nil, ContextRecoveryOptions{
267 ContextWindow: 100,
268 PersistCompaction: func([]provider.Message) error {
269 persisted++
270 return nil
271 },
272 })
273 if !errors.Is(err, context.Canceled) {
274 t.Fatalf("error = %v, want context cancellation", err)
275 }
276 if persisted != 0 {
277 t.Fatalf("persistence calls = %d, want 0", persisted)
278 }
279 if got := agent.Messages(); !reflect.DeepEqual(got, history) {
280 t.Fatalf("transcript after cancellation = %#v, want %#v", got, history)
281 }
282 client.mu.Lock()
283 requestCount := len(client.requests)
284 client.mu.Unlock()
285 if requestCount != 1 {
286 t.Fatalf("provider requests = %d, want summary only", requestCount)
287 }
288}
289
290func TestPromptWithContextRecoveryRepairsToolPairAtCompactionBoundary(t *testing.T) {
291 client := &proactiveContextClient{}
292 agent := core.NewAgent(client, "model", "system", nil)
293 agent.SetMessages([]provider.Message{
294 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "old"}}},
295 {Role: provider.RoleAssistant, Content: []provider.Content{provider.ToolCallBlock{ID: "split-call", Name: "read"}}},
296 {Role: provider.RoleTool, Content: []provider.Content{provider.ToolResultBlock{CallID: "split-call", Content: []provider.Content{provider.TextBlock{Text: "result"}}}}},
297 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "after tool"}}},
298 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "next"}}},
299 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "ready"}}},
300 })
301 agent.SeedLastTurnUsage(provider.Usage{InputTokens: 90})
302
303 if _, err := PromptWithContextRecovery(context.Background(), agent, "continue", nil, nil, ContextRecoveryOptions{ContextWindow: 100}); err != nil {
304 t.Fatalf("PromptWithContextRecovery: %v", err)
305 }
306 client.mu.Lock()
307 requests := append([]provider.Request(nil), client.requests...)
308 client.mu.Unlock()
309 if len(requests) != 2 {
310 t.Fatalf("provider requests = %d, want summary then continuation", len(requests))
311 }
312 for _, message := range requests[1].Messages {
313 for _, content := range message.Content {
314 if result, ok := content.(provider.ToolResultBlock); ok && result.CallID == "split-call" {
315 t.Fatalf("continuation contains orphaned tool result: %#v", requests[1].Messages)
316 }
317 }
318 }
319}
320
321func recoveryRequestContainsText(req provider.Request, want string) bool {
322 for _, message := range req.Messages {
323 for _, content := range message.Content {
324 if text, ok := content.(provider.TextBlock); ok && strings.Contains(text.Text, want) {
325 return true
326 }
327 }
328 }
329 return false
330}
331
332func (c *hostRecoveryClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
333 c.mu.Lock()
334 c.calls++
335 call := c.calls
336 c.mu.Unlock()
337
338 events := make(chan provider.Event, 2)
339 go func() {
340 defer close(events)
341 switch call {
342 case 1:
343 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("provider error: input exceeds the context window")}
344 case 2:
345 events <- provider.EventTextDelta{Delta: "summary"}
346 events <- provider.EventDone{Stop: provider.StopEnd}
347 default:
348 c.retried <- req
349 if c.retryOverflow {
350 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("context window exceeded after compaction")}
351 return
352 }
353 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
354 Role: provider.RoleAssistant,
355 Content: []provider.Content{provider.TextBlock{Text: "recovered"}},
356 }}
357 }
358 }()
359 return events, nil
360}
361
362func TestPromptWithContextRecoveryCompactsOnceAndSuppressesInitialOverflow(t *testing.T) {
363 client := &hostRecoveryClient{retried: make(chan provider.Request, 1)}
364 agent := core.NewAgent(client, "test-model", "", nil)
365 agent.SetMessages([]provider.Message{
366 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
367 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
368 })
369
370 var sawOverflowTurnEnd bool
371 compacted := false
372 result, err := PromptWithContextRecovery(context.Background(), agent, "finish the task", nil, func(event core.AgentEvent) {
373 if turnEnd, ok := event.(core.EvTurnEnd); ok && provider.IsContextOverflowError(turnEnd.Err) {
374 sawOverflowTurnEnd = true
375 }
376 }, ContextRecoveryOptions{
377 PersistCompaction: func([]provider.Message) error {
378 compacted = true
379 return nil
380 },
381 SuppressInitialOverflowEvent: true,
382 })
383 if err != nil {
384 t.Fatalf("PromptWithContextRecovery returned %v", err)
385 }
386 if !result.Compacted || !compacted {
387 t.Fatalf("compaction result/persistence = %v/%v, want true/true", result.Compacted, compacted)
388 }
389 if sawOverflowTurnEnd {
390 t.Fatal("recoverable initial overflow turn_end reached the sink")
391 }
392 if got := recoveredAssistantText(agent.Messages()[result.OutputStart:]); got != "recovered" {
393 t.Fatalf("recovery output = %q, want recovered", got)
394 }
395 request := <-client.retried
396 if got := requestUserTextCount(request, "finish the task"); got != 1 {
397 t.Fatalf("retried request prompt count = %d, want 1: %#v", got, request.Messages)
398 }
399}
400
401func recoveredAssistantText(messages []provider.Message) string {
402 for _, message := range messages {
403 if message.Role != provider.RoleAssistant {
404 continue
405 }
406 for _, content := range message.Content {
407 if text, ok := content.(provider.TextBlock); ok {
408 return text.Text
409 }
410 }
411 }
412 return ""
413}
414
415func TestPromptWithContextRecoveryReportsRetryOverflowToSink(t *testing.T) {
416 client := &hostRecoveryClient{retried: make(chan provider.Request, 1), retryOverflow: true}
417 agent := core.NewAgent(client, "test-model", "", nil)
418 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
419
420 overflowTurnEnds := 0
421 _, err := PromptWithContextRecovery(context.Background(), agent, "finish the task", nil, func(event core.AgentEvent) {
422 if turnEnd, ok := event.(core.EvTurnEnd); ok && provider.IsContextOverflowError(turnEnd.Err) {
423 overflowTurnEnds++
424 }
425 }, ContextRecoveryOptions{SuppressInitialOverflowEvent: true})
426 if !provider.IsContextOverflowError(err) {
427 t.Fatalf("PromptWithContextRecovery error = %v, want retry context overflow", err)
428 }
429 if overflowTurnEnds != 1 {
430 t.Fatalf("overflow turn_end count = %d, want retry overflow only", overflowTurnEnds)
431 }
432 client.mu.Lock()
433 calls := client.calls
434 client.mu.Unlock()
435 if calls != 3 {
436 t.Fatalf("provider calls = %d, want prompt, compaction, and one continuation", calls)
437 }
438}
439
440func newRecoveredModeAgent() *core.Agent {
441 return core.NewAgent(&hostRecoveryClient{retried: make(chan provider.Request, 1)}, "test-model", "", nil)
442}
443
444func TestRunPrintRecoversContextOverflowWithoutLeakingIntermediateOutput(t *testing.T) {
445 agent := newRecoveredModeAgent()
446 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
447 var output bytes.Buffer
448
449 if _, err := RunPrint(context.Background(), agent, "finish the task", nil, &output); err != nil {
450 t.Fatalf("RunPrint returned %v", err)
451 }
452 if got := output.String(); got != "recovered\n" {
453 t.Fatalf("stdout = %q, want recovered only", got)
454 }
455}
456
457func TestRunStreamRecoversContextOverflowWithoutDiagnosticsLeak(t *testing.T) {
458 agent := newRecoveredModeAgent()
459 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
460 var output, diagnostics bytes.Buffer
461
462 if err := RunStreamWithDiag(context.Background(), agent, "finish the task", nil, &output, &diagnostics); err != nil {
463 t.Fatalf("RunStreamWithDiag returned %v", err)
464 }
465 if got := output.String(); got != "recovered\n" {
466 t.Fatalf("stdout = %q, want recovered assistant text once", got)
467 }
468 if got := diagnostics.String(); got != "" {
469 t.Fatalf("stderr diagnostics = %q, want no context-overflow diagnostic", got)
470 }
471}
472
473func TestRunJSONRecoversContextOverflowWithJSONLOnly(t *testing.T) {
474 agent := newRecoveredModeAgent()
475 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
476 var output bytes.Buffer
477
478 if err := RunJSON(context.Background(), agent, "finish the task", nil, &output); err != nil {
479 t.Fatalf("RunJSON returned %v", err)
480 }
481
482 var (
483 turnEnds int
484 done int
485 )
486 scanner := bufio.NewScanner(&output)
487 for scanner.Scan() {
488 var event map[string]any
489 if err := json.Unmarshal(scanner.Bytes(), &event); err != nil {
490 t.Fatalf("stdout line is not JSONL: %q: %v", scanner.Text(), err)
491 }
492 switch event["type"] {
493 case "error":
494 t.Fatalf("recoverable overflow wrote error frame: %#v", event)
495 case "turn_end":
496 turnEnds++
497 if _, hasError := event["error"]; hasError {
498 t.Fatalf("recoverable overflow wrote terminal turn_end: %#v", event)
499 }
500 case "done":
501 done++
502 }
503 }
504 if err := scanner.Err(); err != nil {
505 t.Fatalf("scan JSONL: %v", err)
506 }
507 if turnEnds != 1 || done != 1 {
508 t.Fatalf("turn_end/done count = %d/%d, want one recovered terminal sequence", turnEnds, done)
509 }
510}