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