1package core
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "strings"
8 "sync/atomic"
9 "testing"
10 "time"
11
12 "github.com/patriceckhart/zot/packages/provider"
13)
14
15type retryFakeClient struct {
16 calls int32
17}
18
19func (c *retryFakeClient) Name() string { return "retry-fake" }
20
21func (c *retryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
22 call := atomic.AddInt32(&c.calls, 1)
23 out := make(chan provider.Event, 4)
24 go func() {
25 defer close(out)
26 out <- provider.EventStart{Provider: "retry-fake", Model: req.Model}
27 if call == 1 {
28 out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("anthropic overloaded_error: Overloaded")}
29 return
30 }
31 out <- provider.EventTextDelta{Delta: "ok"}
32 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
33 Role: provider.RoleAssistant,
34 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
35 }}
36 }()
37 return out, nil
38}
39
40func TestAgentRetriesOverloadedStreamError(t *testing.T) {
41 client := &retryFakeClient{}
42 a := NewAgent(client, "fake-model", "system", Registry{})
43 a.RetryBaseDelay = time.Millisecond
44
45 var turnErrs []string
46 err := a.Prompt(context.Background(), "hello", nil, func(ev AgentEvent) {
47 if e, ok := ev.(EvTurnEnd); ok && e.Err != nil {
48 turnErrs = append(turnErrs, e.Err.Error())
49 }
50 })
51 if err != nil {
52 t.Fatalf("Prompt returned %v", err)
53 }
54 if got := atomic.LoadInt32(&client.calls); got != 2 {
55 t.Fatalf("Stream calls = %d; want 2", got)
56 }
57 if len(turnErrs) != 1 || !strings.Contains(turnErrs[0], "overloaded_error") {
58 t.Fatalf("turn errors = %v; want one overloaded error before retry", turnErrs)
59 }
60 msgs := a.Messages()
61 if len(msgs) != 2 {
62 t.Fatalf("message count = %d; want user + final assistant", len(msgs))
63 }
64 if got := extractText(msgs[1]); got != "ok" {
65 t.Fatalf("final assistant text = %q; want ok", got)
66 }
67}
68
69// codexRetryFakeClient reproduces a transient OpenAI Codex backend
70// failure on the first call and succeeds afterwards.
71type codexRetryFakeClient struct {
72 calls int32
73 firstErr string
74}
75
76func (c *codexRetryFakeClient) Name() string { return "openai-codex" }
77
78func (c *codexRetryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
79 call := atomic.AddInt32(&c.calls, 1)
80 out := make(chan provider.Event, 4)
81 go func() {
82 defer close(out)
83 out <- provider.EventStart{Provider: "openai-codex", Model: req.Model}
84 if call == 1 {
85 out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("%s", c.firstErr)}
86 return
87 }
88 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
89 Role: provider.RoleAssistant,
90 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
91 }}
92 }()
93 return out, nil
94}
95
96func TestAgentRetriesCodexProcessingError(t *testing.T) {
97 cases := []struct {
98 name string
99 err string
100 }{
101 {
102 name: "processing error",
103 err: "codex error: An error occurred while processing your request. You can retry your request, or contact us through our help center at help.openai.com if the error persists. Please include the request ID 60c8ebbd-20bd-42e4-b756-6e844041cfc0 in your message.",
104 },
105 {
106 name: "servers overloaded",
107 err: "codex error: Our servers are currently overloaded. Please try again later.",
108 },
109 {
110 name: "try again later only",
111 err: "codex error: Something went wrong. Please try again later.",
112 },
113 }
114 for _, tc := range cases {
115 t.Run(tc.name, func(t *testing.T) {
116 client := &codexRetryFakeClient{firstErr: tc.err}
117 a := NewAgent(client, "gpt-5.6-sol", "system", Registry{})
118 a.RetryBaseDelay = time.Millisecond
119
120 if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
121 t.Fatalf("Prompt returned %v; want retry to succeed", err)
122 }
123 if got := atomic.LoadInt32(&client.calls); got != 2 {
124 t.Fatalf("Stream calls = %d; want 2", got)
125 }
126 msgs := a.Messages()
127 if len(msgs) != 2 || extractText(msgs[1]) != "ok" {
128 t.Fatalf("messages = %v; want user + ok assistant", msgs)
129 }
130 })
131 }
132}
133
134// TestCanRetryErrorCapacityMessages pins the classification of Codex
135// capacity wording and makes sure quota/usage limits stay terminal even
136// when they carry "try again later" style advice.
137func TestCanRetryErrorCapacityMessages(t *testing.T) {
138 a := NewAgent(nil, "gpt-5.6-sol", "system", Registry{})
139 cases := []struct {
140 msg string
141 want bool
142 }{
143 {"codex error: Our servers are currently overloaded. Please try again later.", true},
144 {"codex error: Our servers are busy right now.", true},
145 {"codex error: Please try again later.", true},
146 {"codex error: You have hit your monthly usage limit. Try again later.", false},
147 {"codex error: quota exceeded, try again later", false},
148 {"codex error: unsupported parameter: reasoning", false},
149 }
150 for _, tc := range cases {
151 if got := a.canRetryError(errors.New(tc.msg), 0); got != tc.want {
152 t.Errorf("canRetryError(%q) = %v; want %v", tc.msg, got, tc.want)
153 }
154 }
155}
156
157type partialRetryFakeClient struct {
158 calls int32
159}
160
161func (c *partialRetryFakeClient) Name() string { return "partial-retry-fake" }
162
163func (c *partialRetryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
164 call := atomic.AddInt32(&c.calls, 1)
165 out := make(chan provider.Event, 4)
166 go func() {
167 defer close(out)
168 out <- provider.EventStart{Provider: "partial-retry-fake", Model: req.Model}
169 if call == 1 {
170 out <- provider.EventTextDelta{Delta: "partial"}
171 out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("provider returned error: 503"), Message: provider.Message{
172 Role: provider.RoleAssistant,
173 Content: []provider.Content{provider.TextBlock{Text: "partial"}},
174 }}
175 return
176 }
177 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
178 Role: provider.RoleAssistant,
179 Content: []provider.Content{provider.TextBlock{Text: "recovered"}},
180 }}
181 }()
182 return out, nil
183}
184
185func TestAgentDropsPartialAssistantBeforeRetry(t *testing.T) {
186 client := &partialRetryFakeClient{}
187 a := NewAgent(client, "fake-model", "system", Registry{})
188 a.RetryBaseDelay = time.Millisecond
189
190 if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
191 t.Fatalf("Prompt returned %v", err)
192 }
193 msgs := a.Messages()
194 if len(msgs) != 2 {
195 t.Fatalf("message count = %d; want user + recovered assistant", len(msgs))
196 }
197 if got := extractText(msgs[1]); got != "recovered" {
198 t.Fatalf("final assistant text = %q; want recovered", got)
199 }
200}
201
202// captureClient records the last Request it received so tests can
203// assert what the agent put on the wire.
204type captureClient struct {
205 lastReq provider.Request
206}
207
208func (c *captureClient) Name() string { return "capture" }
209
210func (c *captureClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
211 c.lastReq = req
212 out := make(chan provider.Event, 3)
213 go func() {
214 defer close(out)
215 out <- provider.EventStart{Provider: "capture", Model: req.Model}
216 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
217 Role: provider.RoleAssistant,
218 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
219 }}
220 }()
221 return out, nil
222}
223
224func TestAgentPropagatesMaxTokens(t *testing.T) {
225 client := &captureClient{}
226 a := NewAgent(client, "fake-model", "system", Registry{})
227 a.MaxTokens = 64000
228
229 if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
230 t.Fatalf("Prompt returned %v", err)
231 }
232 if client.lastReq.MaxTokens != 64000 {
233 t.Fatalf("request MaxTokens = %d; want 64000 (Agent.MaxTokens not propagated)", client.lastReq.MaxTokens)
234 }
235}
236
237func TestAgentPropagatesTemperature(t *testing.T) {
238 client := &captureClient{}
239 a := NewAgent(client, "fake-model", "system", Registry{})
240 temp := float32(0)
241 a.Temperature = &temp
242
243 if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
244 t.Fatalf("Prompt returned %v", err)
245 }
246 if client.lastReq.Temperature == nil || *client.lastReq.Temperature != temp {
247 t.Fatalf("request Temperature = %v; want %v", client.lastReq.Temperature, temp)
248 }
249}