1package core
2
3import (
4 "context"
5 "fmt"
6 "strings"
7 "sync/atomic"
8 "testing"
9 "time"
10
11 "github.com/patriceckhart/zot/packages/provider"
12)
13
14type retryFakeClient struct {
15 calls int32
16}
17
18func (c *retryFakeClient) Name() string { return "retry-fake" }
19
20func (c *retryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
21 call := atomic.AddInt32(&c.calls, 1)
22 out := make(chan provider.Event, 4)
23 go func() {
24 defer close(out)
25 out <- provider.EventStart{Provider: "retry-fake", Model: req.Model}
26 if call == 1 {
27 out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("anthropic overloaded_error: Overloaded")}
28 return
29 }
30 out <- provider.EventTextDelta{Delta: "ok"}
31 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
32 Role: provider.RoleAssistant,
33 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
34 }}
35 }()
36 return out, nil
37}
38
39func TestAgentRetriesOverloadedStreamError(t *testing.T) {
40 client := &retryFakeClient{}
41 a := NewAgent(client, "fake-model", "system", Registry{})
42 a.RetryBaseDelay = time.Millisecond
43
44 var turnErrs []string
45 err := a.Prompt(context.Background(), "hello", nil, func(ev AgentEvent) {
46 if e, ok := ev.(EvTurnEnd); ok && e.Err != nil {
47 turnErrs = append(turnErrs, e.Err.Error())
48 }
49 })
50 if err != nil {
51 t.Fatalf("Prompt returned %v", err)
52 }
53 if got := atomic.LoadInt32(&client.calls); got != 2 {
54 t.Fatalf("Stream calls = %d; want 2", got)
55 }
56 if len(turnErrs) != 1 || !strings.Contains(turnErrs[0], "overloaded_error") {
57 t.Fatalf("turn errors = %v; want one overloaded error before retry", turnErrs)
58 }
59 msgs := a.Messages()
60 if len(msgs) != 2 {
61 t.Fatalf("message count = %d; want user + final assistant", len(msgs))
62 }
63 if got := extractText(msgs[1]); got != "ok" {
64 t.Fatalf("final assistant text = %q; want ok", got)
65 }
66}
67
68// codexRetryFakeClient reproduces the generic OpenAI Codex backend
69// failure ("An error occurred while processing your request. You can
70// retry your request, ...") on the first call and succeeds afterwards.
71type codexRetryFakeClient struct {
72 calls int32
73}
74
75func (c *codexRetryFakeClient) Name() string { return "openai-codex" }
76
77func (c *codexRetryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
78 call := atomic.AddInt32(&c.calls, 1)
79 out := make(chan provider.Event, 4)
80 go func() {
81 defer close(out)
82 out <- provider.EventStart{Provider: "openai-codex", Model: req.Model}
83 if call == 1 {
84 out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("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.")}
85 return
86 }
87 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
88 Role: provider.RoleAssistant,
89 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
90 }}
91 }()
92 return out, nil
93}
94
95func TestAgentRetriesCodexProcessingError(t *testing.T) {
96 client := &codexRetryFakeClient{}
97 a := NewAgent(client, "gpt-5.6-sol", "system", Registry{})
98 a.RetryBaseDelay = time.Millisecond
99
100 if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
101 t.Fatalf("Prompt returned %v; want retry to succeed", err)
102 }
103 if got := atomic.LoadInt32(&client.calls); got != 2 {
104 t.Fatalf("Stream calls = %d; want 2", got)
105 }
106 msgs := a.Messages()
107 if len(msgs) != 2 || extractText(msgs[1]) != "ok" {
108 t.Fatalf("messages = %v; want user + ok assistant", msgs)
109 }
110}
111
112type partialRetryFakeClient struct {
113 calls int32
114}
115
116func (c *partialRetryFakeClient) Name() string { return "partial-retry-fake" }
117
118func (c *partialRetryFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
119 call := atomic.AddInt32(&c.calls, 1)
120 out := make(chan provider.Event, 4)
121 go func() {
122 defer close(out)
123 out <- provider.EventStart{Provider: "partial-retry-fake", Model: req.Model}
124 if call == 1 {
125 out <- provider.EventTextDelta{Delta: "partial"}
126 out <- provider.EventDone{Stop: provider.StopError, Err: fmt.Errorf("provider returned error: 503"), Message: provider.Message{
127 Role: provider.RoleAssistant,
128 Content: []provider.Content{provider.TextBlock{Text: "partial"}},
129 }}
130 return
131 }
132 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
133 Role: provider.RoleAssistant,
134 Content: []provider.Content{provider.TextBlock{Text: "recovered"}},
135 }}
136 }()
137 return out, nil
138}
139
140func TestAgentDropsPartialAssistantBeforeRetry(t *testing.T) {
141 client := &partialRetryFakeClient{}
142 a := NewAgent(client, "fake-model", "system", Registry{})
143 a.RetryBaseDelay = time.Millisecond
144
145 if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
146 t.Fatalf("Prompt returned %v", err)
147 }
148 msgs := a.Messages()
149 if len(msgs) != 2 {
150 t.Fatalf("message count = %d; want user + recovered assistant", len(msgs))
151 }
152 if got := extractText(msgs[1]); got != "recovered" {
153 t.Fatalf("final assistant text = %q; want recovered", got)
154 }
155}
156
157// captureClient records the last Request it received so tests can
158// assert what the agent put on the wire.
159type captureClient struct {
160 lastReq provider.Request
161}
162
163func (c *captureClient) Name() string { return "capture" }
164
165func (c *captureClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
166 c.lastReq = req
167 out := make(chan provider.Event, 3)
168 go func() {
169 defer close(out)
170 out <- provider.EventStart{Provider: "capture", Model: req.Model}
171 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
172 Role: provider.RoleAssistant,
173 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
174 }}
175 }()
176 return out, nil
177}
178
179func TestAgentPropagatesMaxTokens(t *testing.T) {
180 client := &captureClient{}
181 a := NewAgent(client, "fake-model", "system", Registry{})
182 a.MaxTokens = 64000
183
184 if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
185 t.Fatalf("Prompt returned %v", err)
186 }
187 if client.lastReq.MaxTokens != 64000 {
188 t.Fatalf("request MaxTokens = %d; want 64000 (Agent.MaxTokens not propagated)", client.lastReq.MaxTokens)
189 }
190}
191
192func TestAgentPropagatesTemperature(t *testing.T) {
193 client := &captureClient{}
194 a := NewAgent(client, "fake-model", "system", Registry{})
195 temp := float32(0)
196 a.Temperature = &temp
197
198 if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
199 t.Fatalf("Prompt returned %v", err)
200 }
201 if client.lastReq.Temperature == nil || *client.lastReq.Temperature != temp {
202 t.Fatalf("request Temperature = %v; want %v", client.lastReq.Temperature, temp)
203 }
204}