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 TestPreTurnCompactionPreservesPromptImages(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 threshold := 70
112 interactive := NewInteractive(InteractiveConfig{
113 Agent: agent,
114 Provider: "anthropic",
115 Model: "claude-sonnet-4-5",
116 AutoCompactThreshold: &threshold,
117 })
118 interactive.runCtx = context.Background()
119 interactive.lastCtxInput = 150000
120 image := provider.ImageBlock{MimeType: "image/png", Data: []byte("image-data")}
121
122 interactive.startTurnWithImages(context.Background(), "follow up", []provider.ImageBlock{image})
123 select {
124 case <-client.compactionStarted:
125 case <-time.After(2 * time.Second):
126 t.Fatal("pre-turn compaction did not start")
127 }
128 close(client.releaseCompaction)
129
130 select {
131 case req := <-client.followUpRequest:
132 if got := requestUserTextCount(req, "follow up"); got != 1 {
133 t.Fatalf("follow-up request contains prompt %d times, want 1: %#v", got, req.Messages)
134 }
135 if got := requestImageCount(req, image); got != 1 {
136 t.Fatalf("follow-up request contains image %d times, want 1: %#v", got, req.Messages)
137 }
138 case <-time.After(2 * time.Second):
139 t.Fatal("pre-turn compaction did not start the pending image request")
140 }
141}
142
143type contextRecoveryClient struct {
144 mu sync.Mutex
145 calls int
146 retried chan provider.Request
147 overflowRetry bool
148}
149
150func (c *contextRecoveryClient) Name() string { return "context-recovery-test" }
151
152func (c *contextRecoveryClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
153 c.mu.Lock()
154 c.calls++
155 call := c.calls
156 c.mu.Unlock()
157
158 out := make(chan provider.Event, 2)
159 go func() {
160 defer close(out)
161 switch call {
162 case 1:
163 out <- provider.EventDone{
164 Stop: provider.StopError,
165 Err: errors.New("provider error: Your input exceeds the context window of this model. Please adjust your input and try again."),
166 }
167 case 2:
168 out <- provider.EventTextDelta{Delta: "summary"}
169 out <- provider.EventDone{Stop: provider.StopEnd}
170 default:
171 c.retried <- req
172 if c.overflowRetry {
173 out <- provider.EventDone{
174 Stop: provider.StopError,
175 Err: errors.New("context window exceeded after compaction"),
176 }
177 return
178 }
179 out <- provider.EventDone{
180 Stop: provider.StopEnd,
181 Message: provider.Message{
182 Role: provider.RoleAssistant,
183 Content: []provider.Content{provider.TextBlock{Text: "done"}},
184 },
185 }
186 }
187 }()
188 return out, nil
189}
190
191func TestContextWindowErrorCompactsAndRetriesPromptOnce(t *testing.T) {
192 client := &contextRecoveryClient{retried: make(chan provider.Request, 1)}
193 agent := core.NewAgent(client, "test-model", "", nil)
194 agent.SetMessages([]provider.Message{
195 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
196 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
197 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
198 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
199 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
200 })
201 interactive := NewInteractive(InteractiveConfig{Agent: agent})
202 interactive.runCtx = context.Background()
203
204 image := provider.ImageBlock{MimeType: "image/png", Data: []byte("image-data")}
205 interactive.startTurnWithImages(context.Background(), "retry me", []provider.ImageBlock{image})
206
207 select {
208 case req := <-client.retried:
209 if got := requestUserTextCount(req, "retry me"); got != 1 {
210 t.Fatalf("retried request contains original prompt %d times, want 1: %#v", got, req.Messages)
211 }
212 if got := requestImageCount(req, image); got != 1 {
213 t.Fatalf("retried request contains original image %d times, want 1: %#v", got, req.Messages)
214 }
215 case <-time.After(2 * time.Second):
216 t.Fatal("context-window error did not compact and retry the prompt")
217 }
218}
219
220func TestContextWindowRecoveryStopsAfterOneRetry(t *testing.T) {
221 client := &contextRecoveryClient{
222 retried: make(chan provider.Request, 2),
223 overflowRetry: true,
224 }
225 agent := core.NewAgent(client, "test-model", "", nil)
226 agent.SetMessages([]provider.Message{
227 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
228 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
229 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
230 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
231 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
232 })
233 interactive := NewInteractive(InteractiveConfig{Agent: agent})
234 interactive.runCtx = context.Background()
235
236 interactive.startTurn(context.Background(), "still too large")
237 select {
238 case <-client.retried:
239 case <-time.After(2 * time.Second):
240 t.Fatal("context-window recovery did not retry")
241 }
242
243 deadline := time.NewTimer(2 * time.Second)
244 defer deadline.Stop()
245 poll := time.NewTicker(time.Millisecond)
246 defer poll.Stop()
247 for {
248 interactive.mu.Lock()
249 idleWithError := !interactive.busy && interactive.statusErr != ""
250 interactive.mu.Unlock()
251 if idleWithError {
252 break
253 }
254 select {
255 case <-deadline.C:
256 t.Fatal("second context-window error did not settle")
257 case <-poll.C:
258 }
259 }
260
261 client.mu.Lock()
262 calls := client.calls
263 client.mu.Unlock()
264 if calls != 3 {
265 t.Fatalf("provider called %d times, want initial request, compaction, and one retry", calls)
266 }
267}
268
269func TestContextOverflowErrorRecognizesProviderMessages(t *testing.T) {
270 tests := []struct {
271 name string
272 err error
273 want bool
274 }{
275 {name: "http 413", err: errors.New("provider http 413: request rejected"), want: true},
276 {name: "context error code", err: errors.New("context_length_exceeded"), want: true},
277 {name: "maximum context length", err: errors.New("This model's maximum context length is 128000 tokens"), want: true},
278 {name: "context window exceeded", err: errors.New("context window exceeded"), want: true},
279 {name: "maximum input token count", err: errors.New("input token count exceeds the maximum number of tokens allowed"), want: true},
280 {name: "maximum output token count", err: errors.New("max_tokens exceeds the maximum number of tokens allowed"), want: false},
281 {name: "unrelated error", err: errors.New("rate limit exceeded"), want: false},
282 }
283
284 for _, tt := range tests {
285 t.Run(tt.name, func(t *testing.T) {
286 if got := isContextOverflowError(tt.err); got != tt.want {
287 t.Fatalf("isContextOverflowError(%q) = %t, want %t", tt.err, got, tt.want)
288 }
289 })
290 }
291}
292
293func requestContainsUserText(req provider.Request, want string) bool {
294 return requestUserTextCount(req, want) > 0
295}
296
297func requestUserTextCount(req provider.Request, want string) int {
298 count := 0
299 for _, message := range req.Messages {
300 if message.Role != provider.RoleUser {
301 continue
302 }
303 for _, content := range message.Content {
304 if text, ok := content.(provider.TextBlock); ok && text.Text == want {
305 count++
306 }
307 }
308 }
309 return count
310}
311
312func requestImageCount(req provider.Request, want provider.ImageBlock) int {
313 count := 0
314 for _, message := range req.Messages {
315 if message.Role != provider.RoleUser {
316 continue
317 }
318 for _, content := range message.Content {
319 if image, ok := content.(provider.ImageBlock); ok && image.MimeType == want.MimeType && string(image.Data) == string(want.Data) {
320 count++
321 }
322 }
323 }
324 return count
325}