1package modes
2
3import (
4 "context"
5 "errors"
6 "sync"
7 "testing"
8 "time"
9
10 "github.com/bnema/zut/packages/core"
11 "github.com/bnema/zut/packages/provider"
12 "github.com/bnema/zut/packages/tui"
13)
14
15type compactQueueClient struct {
16 compactionStarted chan struct{}
17 releaseCompaction chan struct{}
18 followUpRequest chan provider.Request
19
20 mu sync.Mutex
21 calls int
22}
23
24func (c *compactQueueClient) Name() string { return "compact-queue-test" }
25
26func (c *compactQueueClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
27 c.mu.Lock()
28 c.calls++
29 call := c.calls
30 c.mu.Unlock()
31
32 out := make(chan provider.Event, 2)
33 go func() {
34 defer close(out)
35 if call == 1 {
36 close(c.compactionStarted)
37 select {
38 case <-ctx.Done():
39 out <- provider.EventDone{Stop: provider.StopAborted, Err: ctx.Err()}
40 case <-c.releaseCompaction:
41 out <- provider.EventTextDelta{Delta: "summary"}
42 out <- provider.EventDone{Stop: provider.StopEnd}
43 }
44 return
45 }
46
47 c.followUpRequest <- req
48 out <- provider.EventDone{
49 Stop: provider.StopEnd,
50 Message: provider.Message{
51 Role: provider.RoleAssistant,
52 Content: []provider.Content{provider.TextBlock{Text: "done"}},
53 },
54 }
55 }()
56 return out, nil
57}
58
59type autoCompactPacerClient struct {
60 compactionStarted chan struct{}
61 releaseCompaction chan struct{}
62
63 mu sync.Mutex
64 calls int
65}
66
67func (c *autoCompactPacerClient) Name() string { return "auto-compact-pacer-test" }
68
69func (c *autoCompactPacerClient) Stream(ctx context.Context, _ provider.Request) (<-chan provider.Event, error) {
70 c.mu.Lock()
71 c.calls++
72 call := c.calls
73 c.mu.Unlock()
74
75 out := make(chan provider.Event, 3)
76 go func() {
77 defer close(out)
78 switch call {
79 case 1:
80 const response = "a final response that is still buffered for typewriter rendering"
81 out <- provider.EventTextDelta{Delta: response}
82 out <- provider.EventUsage{Usage: provider.Usage{InputTokens: 150000}}
83 out <- provider.EventDone{
84 Stop: provider.StopEnd,
85 Message: provider.Message{
86 Role: provider.RoleAssistant,
87 Content: []provider.Content{provider.TextBlock{Text: response}},
88 },
89 }
90 case 2:
91 close(c.compactionStarted)
92 select {
93 case <-ctx.Done():
94 out <- provider.EventDone{Stop: provider.StopAborted, Err: ctx.Err()}
95 case <-c.releaseCompaction:
96 out <- provider.EventTextDelta{Delta: "summary"}
97 out <- provider.EventDone{Stop: provider.StopEnd}
98 }
99 }
100 }()
101 return out, nil
102}
103
104func TestAutoCompactionSettlesFinalStreamingStateBeforeCompacting(t *testing.T) {
105 client := &autoCompactPacerClient{
106 compactionStarted: make(chan struct{}),
107 releaseCompaction: make(chan struct{}),
108 }
109 agent := core.NewAgent(client, "test-model", "", nil)
110 threshold := 70
111 interactive := NewInteractive(InteractiveConfig{
112 Agent: agent,
113 Provider: "anthropic",
114 Model: "claude-sonnet-4-5-20250929",
115 AutoCompactThreshold: &threshold,
116 })
117 interactive.runCtx = context.Background()
118
119 // Do not start the pacer: the test controls the unpainted final delta so
120 // it can prove compaction never races with a still-live streaming frame.
121 interactive.startTurn(context.Background(), "initial request")
122 select {
123 case <-client.compactionStarted:
124 case <-time.After(2 * time.Second):
125 t.Fatal("auto-compaction did not start")
126 }
127
128 interactive.mu.Lock()
129 streaming, pending := interactive.streaming.Len(), len(interactive.streamPending)
130 streamOn, flushPending := interactive.streamOn, interactive.streamFlushPending
131 interactive.mu.Unlock()
132 if streaming != 0 || pending != 0 || streamOn || flushPending {
133 t.Fatalf("auto-compaction started with live stream state: rendered=%d pending=%d on=%t flush=%t", streaming, pending, streamOn, flushPending)
134 }
135
136 close(client.releaseCompaction)
137}
138
139func TestPromptSubmittedDuringCompactionStartsFollowUpTurn(t *testing.T) {
140 client := &compactQueueClient{
141 compactionStarted: make(chan struct{}),
142 releaseCompaction: make(chan struct{}),
143 followUpRequest: make(chan provider.Request, 1),
144 }
145 agent := core.NewAgent(client, "test-model", "", nil)
146 agent.SetMessages([]provider.Message{
147 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
148 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
149 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
150 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
151 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
152 })
153 interactive := NewInteractive(InteractiveConfig{Agent: agent})
154 interactive.runCtx = context.Background()
155
156 interactive.runCompact(context.Background(), false)
157 select {
158 case <-client.compactionStarted:
159 case <-time.After(2 * time.Second):
160 t.Fatal("compaction did not start")
161 }
162
163 interactive.ed.SetValue("follow up")
164 interactive.handleKey(context.Background(), tui.Key{Kind: tui.KeyEnter})
165 close(client.releaseCompaction)
166
167 select {
168 case req := <-client.followUpRequest:
169 if !requestContainsUserText(req, "follow up") {
170 t.Fatalf("follow-up request does not contain queued prompt: %#v", req.Messages)
171 }
172 case <-time.After(2 * time.Second):
173 t.Fatal("queued prompt did not start a turn after compaction")
174 }
175}
176
177type thresholdAutoCompactClient struct {
178 compactionStarted chan struct{}
179 releaseCompaction chan struct{}
180 followUpRequest chan provider.Request
181
182 mu sync.Mutex
183 calls int
184}
185
186func (c *thresholdAutoCompactClient) Name() string { return "threshold-auto-compact-test" }
187
188func (c *thresholdAutoCompactClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
189 c.mu.Lock()
190 c.calls++
191 call := c.calls
192 c.mu.Unlock()
193
194 out := make(chan provider.Event, 3)
195 go func() {
196 defer close(out)
197 switch call {
198 case 1:
199 out <- provider.EventUsage{Usage: provider.Usage{InputTokens: 150000}}
200 out <- provider.EventDone{
201 Stop: provider.StopEnd,
202 Message: provider.Message{
203 Role: provider.RoleAssistant,
204 Content: []provider.Content{provider.ReasoningBlock{Summary: "still working"}},
205 },
206 }
207 case 2:
208 close(c.compactionStarted)
209 select {
210 case <-ctx.Done():
211 out <- provider.EventDone{Stop: provider.StopAborted, Err: ctx.Err()}
212 case <-c.releaseCompaction:
213 out <- provider.EventTextDelta{Delta: "summary"}
214 out <- provider.EventDone{Stop: provider.StopEnd}
215 }
216 default:
217 c.followUpRequest <- req
218 out <- provider.EventDone{
219 Stop: provider.StopEnd,
220 Message: provider.Message{
221 Role: provider.RoleAssistant,
222 Content: []provider.Content{provider.TextBlock{Text: "continued"}},
223 },
224 }
225 }
226 }()
227 return out, nil
228}
229
230func TestThresholdAutoCompactionContinuesMostRecentIntent(t *testing.T) {
231 client := &thresholdAutoCompactClient{
232 compactionStarted: make(chan struct{}),
233 releaseCompaction: make(chan struct{}),
234 followUpRequest: make(chan provider.Request, 1),
235 }
236 agent := core.NewAgent(client, "test-model", "", nil)
237 threshold := 70
238 interactive := NewInteractive(InteractiveConfig{
239 Agent: agent,
240 Provider: "anthropic",
241 Model: "claude-sonnet-4-5-20250929",
242 AutoCompactThreshold: &threshold,
243 })
244 interactive.runCtx = context.Background()
245
246 interactive.startTurn(context.Background(), "initial request")
247 select {
248 case <-client.compactionStarted:
249 case <-time.After(2 * time.Second):
250 t.Fatal("threshold auto-compaction did not start")
251 }
252 close(client.releaseCompaction)
253
254 select {
255 case req := <-client.followUpRequest:
256 if got := requestUserTextCount(req, autoCompactContinuationPrompt); got != 1 {
257 t.Fatalf("continuation request contains auto-continue prompt %d times, want 1: %#v", got, req.Messages)
258 }
259 case <-time.After(2 * time.Second):
260 t.Fatal("threshold auto-compaction did not continue the most recent intent")
261 }
262}
263
264func TestPreTurnCompactionPreservesPromptImages(t *testing.T) {
265 client := &compactQueueClient{
266 compactionStarted: make(chan struct{}),
267 releaseCompaction: make(chan struct{}),
268 followUpRequest: make(chan provider.Request, 1),
269 }
270 agent := core.NewAgent(client, "test-model", "", nil)
271 agent.SetMessages([]provider.Message{{
272 Role: provider.RoleUser,
273 Content: []provider.Content{provider.TextBlock{Text: "one"}},
274 }})
275 threshold := 70
276 interactive := NewInteractive(InteractiveConfig{
277 Agent: agent,
278 Provider: "anthropic",
279 Model: "claude-sonnet-4-5",
280 AutoCompactThreshold: &threshold,
281 })
282 interactive.runCtx = context.Background()
283 interactive.lastCtxInput = 150000
284 image := provider.ImageBlock{MimeType: "image/png", Data: []byte("image-data")}
285
286 interactive.startTurnWithImages(context.Background(), "follow up", []provider.ImageBlock{image})
287 select {
288 case <-client.compactionStarted:
289 case <-time.After(2 * time.Second):
290 t.Fatal("pre-turn compaction did not start")
291 }
292 close(client.releaseCompaction)
293
294 select {
295 case req := <-client.followUpRequest:
296 if got := requestUserTextCount(req, "follow up"); got != 1 {
297 t.Fatalf("follow-up request contains prompt %d times, want 1: %#v", got, req.Messages)
298 }
299 if got := requestImageCount(req, image); got != 1 {
300 t.Fatalf("follow-up request contains image %d times, want 1: %#v", got, req.Messages)
301 }
302 case <-time.After(2 * time.Second):
303 t.Fatal("pre-turn compaction did not start the pending image request")
304 }
305}
306
307type contextRecoveryClient struct {
308 mu sync.Mutex
309 calls int
310 retried chan provider.Request
311 overflowRetry bool
312}
313
314func (c *contextRecoveryClient) Name() string { return "context-recovery-test" }
315
316func (c *contextRecoveryClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
317 c.mu.Lock()
318 c.calls++
319 call := c.calls
320 c.mu.Unlock()
321
322 out := make(chan provider.Event, 2)
323 go func() {
324 defer close(out)
325 switch call {
326 case 1:
327 out <- provider.EventDone{
328 Stop: provider.StopError,
329 Err: errors.New("provider error: Your input exceeds the context window of this model. Please adjust your input and try again."),
330 }
331 case 2:
332 out <- provider.EventTextDelta{Delta: "summary"}
333 out <- provider.EventDone{Stop: provider.StopEnd}
334 default:
335 c.retried <- req
336 if c.overflowRetry {
337 out <- provider.EventDone{
338 Stop: provider.StopError,
339 Err: errors.New("context window exceeded after compaction"),
340 }
341 return
342 }
343 out <- provider.EventDone{
344 Stop: provider.StopEnd,
345 Message: provider.Message{
346 Role: provider.RoleAssistant,
347 Content: []provider.Content{provider.TextBlock{Text: "done"}},
348 },
349 }
350 }
351 }()
352 return out, nil
353}
354
355func TestContextWindowErrorCompactsAndRetriesPromptOnce(t *testing.T) {
356 client := &contextRecoveryClient{retried: make(chan provider.Request, 1)}
357 agent := core.NewAgent(client, "test-model", "", nil)
358 agent.SetMessages([]provider.Message{
359 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
360 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
361 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
362 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
363 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
364 })
365 interactive := NewInteractive(InteractiveConfig{Agent: agent})
366 interactive.runCtx = context.Background()
367
368 image := provider.ImageBlock{MimeType: "image/png", Data: []byte("image-data")}
369 interactive.startTurnWithImages(context.Background(), "retry me", []provider.ImageBlock{image})
370
371 select {
372 case req := <-client.retried:
373 if got := requestUserTextCount(req, "retry me"); got != 1 {
374 t.Fatalf("retried request contains original prompt %d times, want 1: %#v", got, req.Messages)
375 }
376 if got := requestImageCount(req, image); got != 1 {
377 t.Fatalf("retried request contains original image %d times, want 1: %#v", got, req.Messages)
378 }
379 case <-time.After(2 * time.Second):
380 t.Fatal("context-window error did not compact and retry the prompt")
381 }
382}
383
384func TestContextWindowRecoveryStopsAfterOneRetry(t *testing.T) {
385 client := &contextRecoveryClient{
386 retried: make(chan provider.Request, 2),
387 overflowRetry: true,
388 }
389 agent := core.NewAgent(client, "test-model", "", nil)
390 agent.SetMessages([]provider.Message{
391 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
392 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
393 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
394 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
395 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
396 })
397 interactive := NewInteractive(InteractiveConfig{Agent: agent})
398 interactive.runCtx = context.Background()
399
400 interactive.startTurn(context.Background(), "still too large")
401 select {
402 case <-client.retried:
403 case <-time.After(2 * time.Second):
404 t.Fatal("context-window recovery did not retry")
405 }
406
407 deadline := time.NewTimer(2 * time.Second)
408 defer deadline.Stop()
409 poll := time.NewTicker(time.Millisecond)
410 defer poll.Stop()
411 for {
412 interactive.mu.Lock()
413 idleWithError := !interactive.busy && interactive.statusErr != ""
414 interactive.mu.Unlock()
415 if idleWithError {
416 break
417 }
418 select {
419 case <-deadline.C:
420 t.Fatal("second context-window error did not settle")
421 case <-poll.C:
422 }
423 }
424
425 client.mu.Lock()
426 calls := client.calls
427 client.mu.Unlock()
428 if calls != 3 {
429 t.Fatalf("provider called %d times, want initial request, compaction, and one retry", calls)
430 }
431}
432
433func requestContainsUserText(req provider.Request, want string) bool {
434 return requestUserTextCount(req, want) > 0
435}
436
437func requestUserTextCount(req provider.Request, want string) int {
438 count := 0
439 for _, message := range req.Messages {
440 if message.Role != provider.RoleUser {
441 continue
442 }
443 for _, content := range message.Content {
444 if text, ok := content.(provider.TextBlock); ok && text.Text == want {
445 count++
446 }
447 }
448 }
449 return count
450}
451
452func requestImageCount(req provider.Request, want provider.ImageBlock) int {
453 count := 0
454 for _, message := range req.Messages {
455 if message.Role != provider.RoleUser {
456 continue
457 }
458 for _, content := range message.Content {
459 if image, ok := content.(provider.ImageBlock); ok && image.MimeType == want.MimeType && string(image.Data) == string(want.Data) {
460 count++
461 }
462 }
463 }
464 return count
465}