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 lifecycleFailure struct {
231 attempt int
232 max int
233 reason RequestFailureReason
234 terminal bool
235}
236
237type recordingRequestLifecycle struct {
238 attempts []lifecycleAttempt
239 retries []lifecycleRetry
240 failures []lifecycleFailure
241}
242
243func (r *recordingRequestLifecycle) RequestAttempt(attempt, maxAttempts int) {
244 r.attempts = append(r.attempts, lifecycleAttempt{attempt: attempt, max: maxAttempts})
245}
246
247func (r *recordingRequestLifecycle) RetryScheduled(attempt, maxAttempts int, delay time.Duration) {
248 r.retries = append(r.retries, lifecycleRetry{attempt: attempt, max: maxAttempts, delay: delay})
249}
250
251func (r *recordingRequestLifecycle) RequestFailed(attempt, maxAttempts int, reason RequestFailureReason, terminal bool) {
252 r.failures = append(r.failures, lifecycleFailure{attempt: attempt, max: maxAttempts, reason: reason, terminal: terminal})
253}
254
255func TestDoStreamWithRetryReportsLifecycle(t *testing.T) {
256 var attempts atomic.Int32
257 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
258 if attempts.Add(1) == 1 {
259 w.WriteHeader(http.StatusServiceUnavailable)
260 _, _ = w.Write([]byte("servers overloaded: private provider detail"))
261 return
262 }
263 w.WriteHeader(http.StatusOK)
264 }))
265 defer srv.Close()
266
267 lifecycle := &recordingRequestLifecycle{}
268 resp, err := doStreamWithRetry(context.Background(), srv.Client(), func() (*http.Request, error) {
269 return http.NewRequestWithContext(context.Background(), http.MethodPost, srv.URL, nil)
270 }, lifecycle)
271 if err != nil {
272 t.Fatalf("doStreamWithRetry: %v", err)
273 }
274 resp.Body.Close()
275
276 if got := lifecycle.attempts; len(got) != 2 || got[0] != (lifecycleAttempt{attempt: 1, max: 3}) || got[1] != (lifecycleAttempt{attempt: 2, max: 3}) {
277 t.Fatalf("attempts = %#v, want attempts 1 and 2 of 3", got)
278 }
279 if got := lifecycle.retries; len(got) != 1 || got[0].attempt != 2 || got[0].max != 3 || got[0].delay != 250*time.Millisecond {
280 t.Fatalf("retries = %#v, want attempt 2 of 3 after 250ms", got)
281 }
282 if got := lifecycle.failures; len(got) != 1 || got[0] != (lifecycleFailure{attempt: 1, max: 3, reason: RequestFailureOverload}) {
283 t.Fatalf("failures = %#v, want one non-terminal overload", got)
284 }
285}
286
287func TestDoStreamWithRetryReportsTerminalClientFailure(t *testing.T) {
288 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
289 w.WriteHeader(http.StatusUnauthorized)
290 _, _ = w.Write([]byte("servers overloaded: private provider detail"))
291 }))
292 defer srv.Close()
293
294 lifecycle := &recordingRequestLifecycle{}
295 resp, err := doStreamWithRetry(context.Background(), srv.Client(), func() (*http.Request, error) {
296 return http.NewRequestWithContext(context.Background(), http.MethodPost, srv.URL, nil)
297 }, lifecycle)
298 if err != nil {
299 t.Fatalf("doStreamWithRetry: %v", err)
300 }
301 defer resp.Body.Close()
302 if resp.StatusCode != http.StatusUnauthorized {
303 t.Fatalf("status = %d, want 401", resp.StatusCode)
304 }
305 if got := lifecycle.failures; len(got) != 1 || got[0] != (lifecycleFailure{attempt: 1, max: 3, reason: RequestFailureClient, terminal: true}) {
306 t.Fatalf("failures = %#v, want one terminal client failure", got)
307 }
308 if len(lifecycle.retries) != 0 {
309 t.Fatalf("retries = %#v, want none", lifecycle.retries)
310 }
311}
312
313func TestDoStreamWithRetryReportsTerminalQuota(t *testing.T) {
314 const privateBody = `{"error":{"type":"insufficient_quota","message":"private provider detail"}}`
315 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
316 w.WriteHeader(http.StatusTooManyRequests)
317 _, _ = w.Write([]byte(privateBody))
318 }))
319 defer srv.Close()
320
321 lifecycle := &recordingRequestLifecycle{}
322 resp, err := doStreamWithRetry(context.Background(), srv.Client(), func() (*http.Request, error) {
323 return http.NewRequestWithContext(context.Background(), http.MethodPost, srv.URL, nil)
324 }, lifecycle)
325 if err != nil {
326 t.Fatalf("doStreamWithRetry: %v", err)
327 }
328 defer resp.Body.Close()
329 if resp.StatusCode != http.StatusTooManyRequests {
330 t.Fatalf("status = %d, want 429", resp.StatusCode)
331 }
332 if got := lifecycle.failures; len(got) != 1 || got[0] != (lifecycleFailure{attempt: 1, max: 3, reason: RequestFailureQuota, terminal: true}) {
333 t.Fatalf("failures = %#v, want one terminal quota failure", got)
334 }
335}
336
337func TestDoStreamWithRetryStopsLifecycleAfterCancellation(t *testing.T) {
338 ctx, cancel := context.WithCancel(context.Background())
339 cancel()
340 lifecycle := &recordingRequestLifecycle{}
341 _, err := doStreamWithRetry(ctx, http.DefaultClient, func() (*http.Request, error) {
342 t.Fatal("new request called after cancellation")
343 return nil, nil
344 }, lifecycle)
345 if !errors.Is(err, context.Canceled) {
346 t.Fatalf("error = %v, want context.Canceled", err)
347 }
348 if len(lifecycle.attempts) != 0 || len(lifecycle.retries) != 0 {
349 t.Fatalf("lifecycle after cancellation = %#v %#v", lifecycle.attempts, lifecycle.retries)
350 }
351}
352
353func TestRetryAfterDelay(t *testing.T) {
354 mk := func(kv ...string) http.Header {
355 h := http.Header{}
356 for i := 0; i+1 < len(kv); i += 2 {
357 h.Set(kv[i], kv[i+1])
358 }
359 return h
360 }
361 if d, ok := retryAfterDelay(mk("retry-after-ms", "1500")); !ok || d != 1500*time.Millisecond {
362 t.Errorf("retry-after-ms: got %v %v", d, ok)
363 }
364 if d, ok := retryAfterDelay(mk("Retry-After", "2")); !ok || d != 2*time.Second {
365 t.Errorf("seconds: got %v %v", d, ok)
366 }
367 // retry-after-ms takes precedence over Retry-After.
368 if d, ok := retryAfterDelay(mk("retry-after-ms", "100", "Retry-After", "30")); !ok || d != 100*time.Millisecond {
369 t.Errorf("precedence: got %v %v", d, ok)
370 }
371 // HTTP-date in the past clamps to zero.
372 if d, ok := retryAfterDelay(mk("Retry-After", "Mon, 02 Jan 2006 15:04:05 GMT")); !ok || d != 0 {
373 t.Errorf("past date: got %v %v", d, ok)
374 }
375 if _, ok := retryAfterDelay(mk()); ok {
376 t.Error("empty headers should not produce a delay")
377 }
378 if _, ok := retryAfterDelay(mk("Retry-After", "soon")); ok {
379 t.Error("garbage Retry-After should not produce a delay")
380 }
381}
382
383func TestIsTerminalRateLimitBody(t *testing.T) {
384 terminal := []string{
385 `{"error":{"type":"insufficient_quota"}}`,
386 "Monthly usage limit reached",
387 "you have hit your usage limit",
388 "quota exceeded",
389 "billing hard limit",
390 }
391 for _, b := range terminal {
392 if !isTerminalRateLimitBody(b) {
393 t.Errorf("%q should be terminal", b)
394 }
395 }
396 transient := []string{
397 "Our servers are currently overloaded. Please try again later.",
398 "Rate limit reached, retry shortly",
399 "",
400 }
401 for _, b := range transient {
402 if isTerminalRateLimitBody(b) {
403 t.Errorf("%q should NOT be terminal", b)
404 }
405 }
406}
407
408func TestIsTransientConnectError(t *testing.T) {
409 good := []string{
410 "read tcp: connection reset by peer",
411 "dial tcp: connection refused",
412 "unexpected EOF",
413 "upstream connect error or disconnect/reset before headers",
414 "tls handshake error: bad cert",
415 "i/o timeout",
416 }
417 for _, m := range good {
418 if !isTransientConnectError(stringErr(m)) {
419 t.Errorf("%q should be transient", m)
420 }
421 }
422 bad := []string{
423 "context canceled",
424 "json: cannot unmarshal",
425 }
426 for _, m := range bad {
427 if isTransientConnectError(stringErr(m)) {
428 t.Errorf("%q should NOT be transient", m)
429 }
430 }
431}
432
433type stringErr string
434
435func (s stringErr) Error() string { return string(s) }
436
437func TestDoStreamWithRetryDoesNotReportAttemptWhenRequestCreationFails(t *testing.T) {
438 lifecycle := &recordingRequestLifecycle{}
439 _, err := doStreamWithRetry(context.Background(), http.DefaultClient, func() (*http.Request, error) {
440 return nil, errors.New("cannot build request")
441 }, lifecycle)
442 if err == nil || !strings.Contains(err.Error(), "cannot build request") {
443 t.Fatalf("error = %v, want request creation failure", err)
444 }
445 if len(lifecycle.attempts) != 0 {
446 t.Fatalf("request attempts = %#v, want none", lifecycle.attempts)
447 }
448}
449
450type retryRoundTripper func(*http.Request) (*http.Response, error)
451
452func (f retryRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
453 return f(req)
454}
455
456func TestDoStreamWithRetryDoesNotReportCancelledRetry(t *testing.T) {
457 ctx, cancel := context.WithCancel(context.Background())
458 defer cancel()
459 client := &http.Client{Transport: retryRoundTripper(func(_ *http.Request) (*http.Response, error) {
460 cancel()
461 return &http.Response{
462 StatusCode: http.StatusServiceUnavailable,
463 Header: make(http.Header),
464 Body: http.NoBody,
465 }, nil
466 })}
467 lifecycle := &recordingRequestLifecycle{}
468 _, err := doStreamWithRetry(ctx, client, func() (*http.Request, error) {
469 return http.NewRequestWithContext(ctx, http.MethodPost, "https://example.test", nil)
470 }, lifecycle)
471 if !errors.Is(err, context.Canceled) {
472 t.Fatalf("error = %v, want context.Canceled", err)
473 }
474 if len(lifecycle.retries) != 0 {
475 t.Fatalf("retry lifecycle = %#v, want none after cancellation", lifecycle.retries)
476 }
477}