1package modes
2
3import (
4 "context"
5 "errors"
6 "sync"
7 "testing"
8 "time"
9
10 "github.com/patriceckhart/zot/packages/core"
11 "github.com/patriceckhart/zot/packages/provider"
12 "github.com/patriceckhart/zot/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
59func TestPromptSubmittedDuringCompactionStartsFollowUpTurn(t *testing.T) {
60 client := &compactQueueClient{
61 compactionStarted: make(chan struct{}),
62 releaseCompaction: make(chan struct{}),
63 followUpRequest: make(chan provider.Request, 1),
64 }
65 agent := core.NewAgent(client, "test-model", "", nil)
66 agent.SetMessages([]provider.Message{
67 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
68 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
69 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
70 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
71 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
72 })
73 interactive := NewInteractive(InteractiveConfig{Agent: agent})
74 interactive.runCtx = context.Background()
75
76 interactive.runCompact(context.Background(), false)
77 select {
78 case <-client.compactionStarted:
79 case <-time.After(2 * time.Second):
80 t.Fatal("compaction did not start")
81 }
82
83 interactive.ed.SetValue("follow up")
84 interactive.handleKey(context.Background(), tui.Key{Kind: tui.KeyEnter})
85 close(client.releaseCompaction)
86
87 select {
88 case req := <-client.followUpRequest:
89 if !requestContainsUserText(req, "follow up") {
90 t.Fatalf("follow-up request does not contain queued prompt: %#v", req.Messages)
91 }
92 case <-time.After(2 * time.Second):
93 t.Fatal("queued prompt did not start a turn after compaction")
94 }
95}
96
97func TestCompactionRetainsContextEstimate(t *testing.T) {
98 client := &compactQueueClient{
99 compactionStarted: make(chan struct{}),
100 releaseCompaction: make(chan struct{}),
101 followUpRequest: make(chan provider.Request, 1),
102 }
103 agent := core.NewAgent(client, "test-model", "", nil)
104 agent.SetMessages([]provider.Message{
105 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
106 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
107 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
108 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
109 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
110 })
111 interactive := NewInteractive(InteractiveConfig{Agent: agent})
112
113 interactive.runCompact(context.Background(), false)
114 select {
115 case <-client.compactionStarted:
116 case <-time.After(2 * time.Second):
117 t.Fatal("compaction did not start")
118 }
119 close(client.releaseCompaction)
120
121 deadline := time.NewTimer(2 * time.Second)
122 defer deadline.Stop()
123 poll := time.NewTicker(time.Millisecond)
124 defer poll.Stop()
125 for {
126 interactive.mu.Lock()
127 busy := interactive.busy
128 got := interactive.lastCtxInput
129 interactive.mu.Unlock()
130 if !busy {
131 want := estimateTimelineMessageTokens(agent.Messages())
132 if got != want || got == 0 {
133 t.Fatalf("context estimate after compaction = %d, want %d", got, want)
134 }
135 break
136 }
137 select {
138 case <-deadline.C:
139 t.Fatal("compaction did not finish")
140 case <-poll.C:
141 }
142 }
143}
144
145func TestPreTurnCompactionPreservesPromptImages(t *testing.T) {
146 client := &compactQueueClient{
147 compactionStarted: make(chan struct{}),
148 releaseCompaction: make(chan struct{}),
149 followUpRequest: make(chan provider.Request, 1),
150 }
151 agent := core.NewAgent(client, "test-model", "", nil)
152 agent.SetMessages([]provider.Message{
153 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
154 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
155 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
156 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
157 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
158 })
159 threshold := 70
160 interactive := NewInteractive(InteractiveConfig{
161 Agent: agent,
162 Provider: "anthropic",
163 Model: "claude-sonnet-4-5",
164 AutoCompactThreshold: &threshold,
165 })
166 interactive.runCtx = context.Background()
167 interactive.lastCtxInput = 150000
168 image := provider.ImageBlock{MimeType: "image/png", Data: []byte("image-data")}
169
170 interactive.startTurnWithImages(context.Background(), "follow up", []provider.ImageBlock{image})
171 select {
172 case <-client.compactionStarted:
173 case <-time.After(2 * time.Second):
174 t.Fatal("pre-turn compaction did not start")
175 }
176 close(client.releaseCompaction)
177
178 select {
179 case req := <-client.followUpRequest:
180 if got := requestUserTextCount(req, "follow up"); got != 1 {
181 t.Fatalf("follow-up request contains prompt %d times, want 1: %#v", got, req.Messages)
182 }
183 if got := requestImageCount(req, image); got != 1 {
184 t.Fatalf("follow-up request contains image %d times, want 1: %#v", got, req.Messages)
185 }
186 case <-time.After(2 * time.Second):
187 t.Fatal("pre-turn compaction did not start the pending image request")
188 }
189}
190
191type contextRecoveryClient struct {
192 mu sync.Mutex
193 calls int
194 retried chan provider.Request
195 overflowRetry bool
196}
197
198func (c *contextRecoveryClient) Name() string { return "context-recovery-test" }
199
200func (c *contextRecoveryClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
201 c.mu.Lock()
202 c.calls++
203 call := c.calls
204 c.mu.Unlock()
205
206 out := make(chan provider.Event, 2)
207 go func() {
208 defer close(out)
209 switch call {
210 case 1:
211 out <- provider.EventDone{
212 Stop: provider.StopError,
213 Err: errors.New("provider error: Your input exceeds the context window of this model. Please adjust your input and try again."),
214 }
215 case 2:
216 out <- provider.EventTextDelta{Delta: "summary"}
217 out <- provider.EventDone{Stop: provider.StopEnd}
218 default:
219 c.retried <- req
220 if c.overflowRetry {
221 out <- provider.EventDone{
222 Stop: provider.StopError,
223 Err: errors.New("context window exceeded after compaction"),
224 }
225 return
226 }
227 out <- provider.EventDone{
228 Stop: provider.StopEnd,
229 Message: provider.Message{
230 Role: provider.RoleAssistant,
231 Content: []provider.Content{provider.TextBlock{Text: "done"}},
232 },
233 }
234 }
235 }()
236 return out, nil
237}
238
239func TestContextWindowErrorCompactsAndRetriesPromptOnce(t *testing.T) {
240 client := &contextRecoveryClient{retried: make(chan provider.Request, 1)}
241 agent := core.NewAgent(client, "test-model", "", nil)
242 agent.SetMessages([]provider.Message{
243 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
244 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
245 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
246 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
247 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
248 })
249 interactive := NewInteractive(InteractiveConfig{Agent: agent})
250 interactive.runCtx = context.Background()
251
252 image := provider.ImageBlock{MimeType: "image/png", Data: []byte("image-data")}
253 interactive.startTurnWithImages(context.Background(), "retry me", []provider.ImageBlock{image})
254
255 select {
256 case req := <-client.retried:
257 if got := requestUserTextCount(req, "retry me"); got != 1 {
258 t.Fatalf("retried request contains original prompt %d times, want 1: %#v", got, req.Messages)
259 }
260 if got := requestImageCount(req, image); got != 1 {
261 t.Fatalf("retried request contains original image %d times, want 1: %#v", got, req.Messages)
262 }
263 case <-time.After(2 * time.Second):
264 t.Fatal("context-window error did not compact and retry the prompt")
265 }
266}
267
268func TestContextWindowRecoveryStopsAfterOneRetry(t *testing.T) {
269 client := &contextRecoveryClient{
270 retried: make(chan provider.Request, 2),
271 overflowRetry: true,
272 }
273 agent := core.NewAgent(client, "test-model", "", nil)
274 agent.SetMessages([]provider.Message{
275 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
276 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
277 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
278 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
279 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
280 })
281 interactive := NewInteractive(InteractiveConfig{Agent: agent})
282 interactive.runCtx = context.Background()
283
284 interactive.startTurn(context.Background(), "still too large")
285 select {
286 case <-client.retried:
287 case <-time.After(2 * time.Second):
288 t.Fatal("context-window recovery did not retry")
289 }
290
291 deadline := time.NewTimer(2 * time.Second)
292 defer deadline.Stop()
293 poll := time.NewTicker(time.Millisecond)
294 defer poll.Stop()
295 for {
296 interactive.mu.Lock()
297 idleWithError := !interactive.busy && interactive.statusErr != ""
298 interactive.mu.Unlock()
299 if idleWithError {
300 break
301 }
302 select {
303 case <-deadline.C:
304 t.Fatal("second context-window error did not settle")
305 case <-poll.C:
306 }
307 }
308
309 client.mu.Lock()
310 calls := client.calls
311 client.mu.Unlock()
312 if calls != 3 {
313 t.Fatalf("provider called %d times, want initial request, compaction, and one retry", calls)
314 }
315}
316
317func TestContextOverflowErrorRecognizesProviderMessages(t *testing.T) {
318 tests := []struct {
319 name string
320 err error
321 want bool
322 }{
323 {name: "http 413", err: errors.New("provider http 413: request rejected"), want: true},
324 {name: "context error code", err: errors.New("context_length_exceeded"), want: true},
325 {name: "maximum context length", err: errors.New("This model's maximum context length is 128000 tokens"), want: true},
326 {name: "context window exceeded", err: errors.New("context window exceeded"), want: true},
327 {name: "maximum input token count", err: errors.New("input token count exceeds the maximum number of tokens allowed"), want: true},
328 {name: "maximum output token count", err: errors.New("max_tokens exceeds the maximum number of tokens allowed"), want: false},
329 {name: "unrelated error", err: errors.New("rate limit exceeded"), want: false},
330 }
331
332 for _, tt := range tests {
333 t.Run(tt.name, func(t *testing.T) {
334 if got := isContextOverflowError(tt.err); got != tt.want {
335 t.Fatalf("isContextOverflowError(%q) = %t, want %t", tt.err, got, tt.want)
336 }
337 })
338 }
339}
340
341func requestContainsUserText(req provider.Request, want string) bool {
342 return requestUserTextCount(req, want) > 0
343}
344
345func requestUserTextCount(req provider.Request, want string) int {
346 count := 0
347 for _, message := range req.Messages {
348 if message.Role != provider.RoleUser {
349 continue
350 }
351 for _, content := range message.Content {
352 if text, ok := content.(provider.TextBlock); ok && text.Text == want {
353 count++
354 }
355 }
356 }
357 return count
358}
359
360func requestImageCount(req provider.Request, want provider.ImageBlock) int {
361 count := 0
362 for _, message := range req.Messages {
363 if message.Role != provider.RoleUser {
364 continue
365 }
366 for _, content := range message.Content {
367 if image, ok := content.(provider.ImageBlock); ok && image.MimeType == want.MimeType && string(image.Data) == string(want.Data) {
368 count++
369 }
370 }
371 }
372 return count
373}