1package provider
2
3import (
4 "context"
5 "errors"
6 "net/http"
7 "net/http/httptest"
8 "strings"
9 "sync/atomic"
10 "testing"
11 "time"
12)
13
14// TestOpenAIStreamRetriesOn503 verifies the OpenAI/Kimi stream client
15// silently retries when the upstream returns 503 a couple of times,
16// then succeeds. The user-visible error path should NOT fire because
17// a later attempt landed on a 200.
18func TestOpenAIStreamRetriesOn503(t *testing.T) {
19 var attempts atomic.Int32
20 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
21 n := attempts.Add(1)
22 if n < 3 {
23 w.WriteHeader(http.StatusServiceUnavailable)
24 _, _ = w.Write([]byte("upstream connect error"))
25 return
26 }
27 w.Header().Set("content-type", "text/event-stream")
28 w.WriteHeader(http.StatusOK)
29 fl, _ := w.(http.Flusher)
30 write := func(s string) {
31 _, _ = w.Write([]byte(s))
32 if fl != nil {
33 fl.Flush()
34 }
35 }
36 write("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\n")
37 write("data: [DONE]\n\n")
38 }))
39 defer srv.Close()
40
41 c := NewOpenAI("x", srv.URL)
42 lifecycle := &recordingRequestLifecycle{}
43 evs, err := c.Stream(context.Background(), Request{Model: "gpt-5", Lifecycle: lifecycle})
44 if err != nil {
45 t.Fatalf("stream returned error: %v", err)
46 }
47 gotText := ""
48 for ev := range evs {
49 if td, ok := ev.(EventTextDelta); ok {
50 gotText += td.Delta
51 }
52 }
53 if gotText != "ok" {
54 t.Fatalf("text=%q", gotText)
55 }
56 if got := attempts.Load(); got != 3 {
57 t.Fatalf("attempts=%d want 3", got)
58 }
59 if got := lifecycle.attempts; len(got) != 3 || got[0] != (lifecycleAttempt{attempt: 1, max: 3}) || got[2] != (lifecycleAttempt{attempt: 3, max: 3}) {
60 t.Fatalf("lifecycle attempts = %#v, want 1 through 3", got)
61 }
62 if got := lifecycle.retries; len(got) != 2 || got[0].attempt != 2 || got[1].attempt != 3 {
63 t.Fatalf("lifecycle retries = %#v, want retries before attempts 2 and 3", got)
64 }
65}
66
67// TestOpenAIStreamSurfaces503AfterRetriesExhausted verifies that when
68// every retry attempt returns 503 the caller sees a normal "http 503"
69// error, not a retry-internal placeholder. The body bytes must be
70// preserved so the rescue picker / red banner can show them.
71func TestOpenAIStreamSurfaces503AfterRetriesExhausted(t *testing.T) {
72 var attempts atomic.Int32
73 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
74 attempts.Add(1)
75 w.WriteHeader(http.StatusServiceUnavailable)
76 _, _ = w.Write([]byte("upstream connect error"))
77 }))
78 defer srv.Close()
79
80 c := NewOpenAI("x", srv.URL)
81 _, err := c.Stream(context.Background(), Request{Model: "gpt-5"})
82 if err == nil {
83 t.Fatalf("expected error after exhausted retries")
84 }
85 if !strings.Contains(err.Error(), "503") {
86 t.Fatalf("error %q should mention 503", err)
87 }
88 if got := attempts.Load(); got != streamRetryAttempts+1 {
89 t.Fatalf("attempts=%d want %d", got, streamRetryAttempts+1)
90 }
91}
92
93// TestOpenAIStreamDoesNotRetryOn400 ensures non-transient errors are
94// returned immediately without retrying. We don't want to amplify
95// load on a provider that's actively rejecting our request shape.
96func TestOpenAIStreamDoesNotRetryOn400(t *testing.T) {
97 var attempts atomic.Int32
98 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
99 attempts.Add(1)
100 w.WriteHeader(http.StatusBadRequest)
101 _, _ = w.Write([]byte(`{"error":{"message":"nope"}}`))
102 }))
103 defer srv.Close()
104
105 c := NewOpenAI("x", srv.URL)
106 _, err := c.Stream(context.Background(), Request{Model: "gpt-5"})
107 if err == nil || !strings.Contains(err.Error(), "400") {
108 t.Fatalf("want 400 err got %v", err)
109 }
110 if got := attempts.Load(); got != 1 {
111 t.Fatalf("attempts=%d want 1", got)
112 }
113}
114
115func TestIsTransientHTTPStatus(t *testing.T) {
116 for _, code := range []int{429, 500, 502, 503, 504, 524, 529} {
117 if !isTransientHTTPStatus(code) {
118 t.Errorf("%d should be transient", code)
119 }
120 }
121 for _, code := range []int{200, 400, 401, 403, 404, 501} {
122 if isTransientHTTPStatus(code) {
123 t.Errorf("%d should NOT be transient", code)
124 }
125 }
126}
127
128// TestOpenAIStreamRetriesOn429WithRetryAfter verifies that a transient
129// 429 is retried and that a server-requested Retry-After delay is
130// honored instead of the fixed backoff.
131func TestOpenAIStreamRetriesOn429WithRetryAfter(t *testing.T) {
132 var attempts atomic.Int32
133 var firstRetryAt atomic.Int64
134 start := time.Now()
135 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
136 n := attempts.Add(1)
137 if n == 1 {
138 w.Header().Set("retry-after-ms", "50")
139 w.WriteHeader(http.StatusTooManyRequests)
140 _, _ = w.Write([]byte("Our servers are currently overloaded. Please try again later."))
141 return
142 }
143 firstRetryAt.Store(int64(time.Since(start)))
144 w.Header().Set("content-type", "text/event-stream")
145 w.WriteHeader(http.StatusOK)
146 fl, _ := w.(http.Flusher)
147 write := func(s string) {
148 _, _ = w.Write([]byte(s))
149 if fl != nil {
150 fl.Flush()
151 }
152 }
153 write("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\n")
154 write("data: [DONE]\n\n")
155 }))
156 defer srv.Close()
157
158 c := NewOpenAI("x", srv.URL)
159 evs, err := c.Stream(context.Background(), Request{Model: "gpt-5"})
160 if err != nil {
161 t.Fatalf("stream returned error: %v", err)
162 }
163 for range evs {
164 }
165 if got := attempts.Load(); got != 2 {
166 t.Fatalf("attempts=%d want 2", got)
167 }
168 if got := time.Duration(firstRetryAt.Load()); got < 50*time.Millisecond {
169 t.Fatalf("retry fired after %v; want >= 50ms (server-requested delay)", got)
170 }
171}
172
173// TestOpenAIStreamDoesNotRetryTerminal429 ensures quota/billing 429s
174// are surfaced immediately without retrying: a usage-limit block will
175// not clear within the retry window, and hammering the endpoint can
176// extend the block on some backends.
177func TestOpenAIStreamDoesNotRetryTerminal429(t *testing.T) {
178 var attempts atomic.Int32
179 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
180 attempts.Add(1)
181 w.WriteHeader(http.StatusTooManyRequests)
182 _, _ = w.Write([]byte(`{"error":{"type":"insufficient_quota","message":"You exceeded your current quota"}}`))
183 }))
184 defer srv.Close()
185
186 c := NewOpenAI("x", srv.URL)
187 _, err := c.Stream(context.Background(), Request{Model: "gpt-5"})
188 if err == nil || !strings.Contains(err.Error(), "429") {
189 t.Fatalf("want 429 err got %v", err)
190 }
191 if got := attempts.Load(); got != 1 {
192 t.Fatalf("attempts=%d want 1", got)
193 }
194}
195
196// TestOpenAIStreamRejectsExcessiveRetryAfter verifies that a
197// server-requested delay beyond maxServerRetryDelay fails immediately
198// with an informative error instead of blocking silently.
199func TestOpenAIStreamRejectsExcessiveRetryAfter(t *testing.T) {
200 var attempts atomic.Int32
201 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
202 attempts.Add(1)
203 w.Header().Set("Retry-After", "18000") // 5 hours
204 w.WriteHeader(http.StatusTooManyRequests)
205 _, _ = w.Write([]byte("rate limited"))
206 }))
207 defer srv.Close()
208
209 c := NewOpenAI("x", srv.URL)
210 _, err := c.Stream(context.Background(), Request{Model: "gpt-5"})
211 if err == nil || !strings.Contains(err.Error(), "retry delay") {
212 t.Fatalf("want retry-delay err got %v", err)
213 }
214 if got := attempts.Load(); got != 1 {
215 t.Fatalf("attempts=%d want 1", got)
216 }
217}
218
219type lifecycleAttempt struct {
220 attempt int
221 max int
222}
223
224type lifecycleRetry struct {
225 attempt int
226 max int
227 delay time.Duration
228}
229
230type recordingRequestLifecycle struct {
231 attempts []lifecycleAttempt
232 retries []lifecycleRetry
233}
234
235func (r *recordingRequestLifecycle) RequestAttempt(attempt, maxAttempts int) {
236 r.attempts = append(r.attempts, lifecycleAttempt{attempt: attempt, max: maxAttempts})
237}
238
239func (r *recordingRequestLifecycle) RetryScheduled(attempt, maxAttempts int, delay time.Duration) {
240 r.retries = append(r.retries, lifecycleRetry{attempt: attempt, max: maxAttempts, delay: delay})
241}
242
243func TestDoStreamWithRetryReportsLifecycle(t *testing.T) {
244 var attempts atomic.Int32
245 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
246 if attempts.Add(1) == 1 {
247 w.WriteHeader(http.StatusServiceUnavailable)
248 return
249 }
250 w.WriteHeader(http.StatusOK)
251 }))
252 defer srv.Close()
253
254 lifecycle := &recordingRequestLifecycle{}
255 resp, err := doStreamWithRetry(context.Background(), srv.Client(), func() (*http.Request, error) {
256 return http.NewRequestWithContext(context.Background(), http.MethodPost, srv.URL, nil)
257 }, lifecycle)
258 if err != nil {
259 t.Fatalf("doStreamWithRetry: %v", err)
260 }
261 resp.Body.Close()
262
263 if got := lifecycle.attempts; len(got) != 2 || got[0] != (lifecycleAttempt{attempt: 1, max: 3}) || got[1] != (lifecycleAttempt{attempt: 2, max: 3}) {
264 t.Fatalf("attempts = %#v, want attempts 1 and 2 of 3", got)
265 }
266 if got := lifecycle.retries; len(got) != 1 || got[0].attempt != 2 || got[0].max != 3 || got[0].delay != 250*time.Millisecond {
267 t.Fatalf("retries = %#v, want attempt 2 of 3 after 250ms", got)
268 }
269}
270
271func TestDoStreamWithRetryStopsLifecycleAfterCancellation(t *testing.T) {
272 ctx, cancel := context.WithCancel(context.Background())
273 cancel()
274 lifecycle := &recordingRequestLifecycle{}
275 _, err := doStreamWithRetry(ctx, http.DefaultClient, func() (*http.Request, error) {
276 t.Fatal("new request called after cancellation")
277 return nil, nil
278 }, lifecycle)
279 if !errors.Is(err, context.Canceled) {
280 t.Fatalf("error = %v, want context.Canceled", err)
281 }
282 if len(lifecycle.attempts) != 0 || len(lifecycle.retries) != 0 {
283 t.Fatalf("lifecycle after cancellation = %#v %#v", lifecycle.attempts, lifecycle.retries)
284 }
285}
286
287func TestRetryAfterDelay(t *testing.T) {
288 mk := func(kv ...string) http.Header {
289 h := http.Header{}
290 for i := 0; i+1 < len(kv); i += 2 {
291 h.Set(kv[i], kv[i+1])
292 }
293 return h
294 }
295 if d, ok := retryAfterDelay(mk("retry-after-ms", "1500")); !ok || d != 1500*time.Millisecond {
296 t.Errorf("retry-after-ms: got %v %v", d, ok)
297 }
298 if d, ok := retryAfterDelay(mk("Retry-After", "2")); !ok || d != 2*time.Second {
299 t.Errorf("seconds: got %v %v", d, ok)
300 }
301 // retry-after-ms takes precedence over Retry-After.
302 if d, ok := retryAfterDelay(mk("retry-after-ms", "100", "Retry-After", "30")); !ok || d != 100*time.Millisecond {
303 t.Errorf("precedence: got %v %v", d, ok)
304 }
305 // HTTP-date in the past clamps to zero.
306 if d, ok := retryAfterDelay(mk("Retry-After", "Mon, 02 Jan 2006 15:04:05 GMT")); !ok || d != 0 {
307 t.Errorf("past date: got %v %v", d, ok)
308 }
309 if _, ok := retryAfterDelay(mk()); ok {
310 t.Error("empty headers should not produce a delay")
311 }
312 if _, ok := retryAfterDelay(mk("Retry-After", "soon")); ok {
313 t.Error("garbage Retry-After should not produce a delay")
314 }
315}
316
317func TestIsTerminalRateLimitBody(t *testing.T) {
318 terminal := []string{
319 `{"error":{"type":"insufficient_quota"}}`,
320 "Monthly usage limit reached",
321 "you have hit your usage limit",
322 "quota exceeded",
323 "billing hard limit",
324 }
325 for _, b := range terminal {
326 if !isTerminalRateLimitBody(b) {
327 t.Errorf("%q should be terminal", b)
328 }
329 }
330 transient := []string{
331 "Our servers are currently overloaded. Please try again later.",
332 "Rate limit reached, retry shortly",
333 "",
334 }
335 for _, b := range transient {
336 if isTerminalRateLimitBody(b) {
337 t.Errorf("%q should NOT be terminal", b)
338 }
339 }
340}
341
342func TestIsTransientConnectError(t *testing.T) {
343 good := []string{
344 "read tcp: connection reset by peer",
345 "dial tcp: connection refused",
346 "unexpected EOF",
347 "upstream connect error or disconnect/reset before headers",
348 "tls handshake error: bad cert",
349 "i/o timeout",
350 }
351 for _, m := range good {
352 if !isTransientConnectError(stringErr(m)) {
353 t.Errorf("%q should be transient", m)
354 }
355 }
356 bad := []string{
357 "context canceled",
358 "json: cannot unmarshal",
359 }
360 for _, m := range bad {
361 if isTransientConnectError(stringErr(m)) {
362 t.Errorf("%q should NOT be transient", m)
363 }
364 }
365}
366
367type stringErr string
368
369func (s stringErr) Error() string { return string(s) }
370
371func TestDoStreamWithRetryDoesNotReportAttemptWhenRequestCreationFails(t *testing.T) {
372 lifecycle := &recordingRequestLifecycle{}
373 _, err := doStreamWithRetry(context.Background(), http.DefaultClient, func() (*http.Request, error) {
374 return nil, errors.New("cannot build request")
375 }, lifecycle)
376 if err == nil || !strings.Contains(err.Error(), "cannot build request") {
377 t.Fatalf("error = %v, want request creation failure", err)
378 }
379 if len(lifecycle.attempts) != 0 {
380 t.Fatalf("request attempts = %#v, want none", lifecycle.attempts)
381 }
382}
383
384type retryRoundTripper func(*http.Request) (*http.Response, error)
385
386func (f retryRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
387 return f(req)
388}
389
390func TestDoStreamWithRetryDoesNotReportCancelledRetry(t *testing.T) {
391 ctx, cancel := context.WithCancel(context.Background())
392 defer cancel()
393 client := &http.Client{Transport: retryRoundTripper(func(_ *http.Request) (*http.Response, error) {
394 cancel()
395 return &http.Response{
396 StatusCode: http.StatusServiceUnavailable,
397 Header: make(http.Header),
398 Body: http.NoBody,
399 }, nil
400 })}
401 lifecycle := &recordingRequestLifecycle{}
402 _, err := doStreamWithRetry(ctx, client, func() (*http.Request, error) {
403 return http.NewRequestWithContext(ctx, http.MethodPost, "https://example.test", nil)
404 }, lifecycle)
405 if !errors.Is(err, context.Canceled) {
406 t.Fatalf("error = %v, want context.Canceled", err)
407 }
408 if len(lifecycle.retries) != 0 {
409 t.Fatalf("retry lifecycle = %#v, want none after cancellation", lifecycle.retries)
410 }
411}