diff --git a/internal/ai/llm/anthropic/anthropic.go b/internal/ai/llm/anthropic/anthropic.go index 37b1239..c41c0f1 100644 --- a/internal/ai/llm/anthropic/anthropic.go +++ b/internal/ai/llm/anthropic/anthropic.go @@ -381,7 +381,7 @@ func (p *Provider) statusError(resp *http.Response, body []byte) error { case http.StatusTooManyRequests, statusOverloaded: return llm.NewRateLimitedError(retryAfter(resp.Header.Get("Retry-After"), time.Now()), p.apiKey) case http.StatusRequestTimeout: - return fmt.Errorf("HTTP %d: %w", resp.StatusCode, llm.ErrTimeout) + return llm.NewTimeoutStatusError(resp.StatusCode, p.apiKey) default: return llm.NewProviderError(resp.StatusCode, body, p.apiKey) } diff --git a/internal/ai/llm/errors.go b/internal/ai/llm/errors.go index e4f1bd2..b4c6480 100644 --- a/internal/ai/llm/errors.go +++ b/internal/ai/llm/errors.go @@ -137,3 +137,32 @@ func (e *ProviderError) GoString() string { return e.text } // Is reports whether target is ErrProvider. func (e *ProviderError) Is(target error) bool { return target == ErrProvider } + +// TimeoutStatusError is an HTTP status, such as 408, by which the vendor +// said the request timed out. It matches ErrTimeout, and Status keeps it +// apart from a transport timeout, so a caller can tell the vendor answered. +// Build it with NewTimeoutStatusError, which redacts its text. +type TimeoutStatusError struct { + status int + text string +} + +// NewTimeoutStatusError returns the error for a timeout HTTP status. +// secret is the credential sent with the request, or "" when the request +// sent none; every occurrence of a non-empty secret is removed from the +// error text, and the secret itself is not kept. +func NewTimeoutStatusError(status int, secret string) *TimeoutStatusError { + text := fmt.Sprintf("HTTP %d: %s", status, ErrTimeout) + return &TimeoutStatusError{status: status, text: Redact(text, secret)} +} + +// Status returns the HTTP status code. +func (e *TimeoutStatusError) Status() int { return e.status } + +func (e *TimeoutStatusError) Error() string { return e.text } + +// GoString keeps %#v to the redacted Error text instead of the fields. +func (e *TimeoutStatusError) GoString() string { return e.text } + +// Unwrap returns ErrTimeout. +func (e *TimeoutStatusError) Unwrap() error { return ErrTimeout } diff --git a/internal/ai/llm/errors_test.go b/internal/ai/llm/errors_test.go index 5ac5fa7..37406df 100644 --- a/internal/ai/llm/errors_test.go +++ b/internal/ai/llm/errors_test.go @@ -196,3 +196,17 @@ func TestTypedErrorsDoNotKeepSecret(t *testing.T) { } } } + +func TestTimeoutStatusError(t *testing.T) { + err := fmt.Errorf("openai: %w", NewTimeoutStatusError(408, "HTTP")) + if !errors.Is(err, ErrTimeout) || errors.Is(err, ErrProvider) { + t.Fatalf("%v: want ErrTimeout only", err) + } + var ts *TimeoutStatusError + if !errors.As(err, &ts) || ts.Status() != 408 { + t.Fatalf("TimeoutStatusError from %v", err) + } + if s := ts.Error() + fmt.Sprintf("%#v", ts); strings.Contains(s, "HTTP") { + t.Fatalf("text holds the secret: %s", s) + } +} diff --git a/internal/ai/llm/openai/openai.go b/internal/ai/llm/openai/openai.go index 70c94e9..d8bf66a 100644 --- a/internal/ai/llm/openai/openai.go +++ b/internal/ai/llm/openai/openai.go @@ -453,7 +453,7 @@ func (p *Provider) statusError(resp *http.Response, body []byte) error { case http.StatusTooManyRequests: return llm.NewRateLimitedError(retryAfter(resp.Header.Get("Retry-After"), time.Now()), p.apiKey) case http.StatusRequestTimeout: - return fmt.Errorf("HTTP %d: %w", resp.StatusCode, llm.ErrTimeout) + return llm.NewTimeoutStatusError(resp.StatusCode, p.apiKey) default: return llm.NewProviderError(resp.StatusCode, body, p.apiKey) } diff --git a/internal/ai/llm/presets/presets.go b/internal/ai/llm/presets/presets.go index 43b0d85..da86eda 100644 --- a/internal/ai/llm/presets/presets.go +++ b/internal/ai/llm/presets/presets.go @@ -178,13 +178,14 @@ var all = map[string]Preset{ // // The Provider sends s.Temperature with every request whose Temperature is // nil; a request that sets one keeps it. Neither is sent to a model the -// preset's OmitTemperature matches. +// preset's OmitTemperature matches. It is wrapped in llm.NewRetrying, so +// Generate retries transient vendor failures and has the default timeout. func FromSettings(s llm.Settings, key string) (llm.Provider, error) { cp, err := fromSettings(s, key) if err != nil { return nil, redact(err, key) } - return cp, nil + return llm.NewRetrying(cp), nil } func fromSettings(s llm.Settings, key string) (*configured, error) { diff --git a/internal/ai/llm/presets/presets_test.go b/internal/ai/llm/presets/presets_test.go index 142c564..41e4b3f 100644 --- a/internal/ai/llm/presets/presets_test.go +++ b/internal/ai/llm/presets/presets_test.go @@ -22,10 +22,8 @@ type captured struct { body map[string]json.RawMessage } -var errStop = errors.New("recorded") - -// recordingTransport records each request and fails it, so no request -// leaves the process. +// recordingTransport records each request and answers it 400, so no +// request leaves the process and the Provider never retries it. type recordingTransport struct{ reqs []captured } func (rt *recordingTransport) RoundTrip(r *http.Request) (*http.Response, error) { @@ -40,7 +38,12 @@ func (rt *recordingTransport) RoundTrip(r *http.Request) (*http.Response, error) } } rt.reqs = append(rt.reqs, c) - return nil, errStop + return &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{}`)), + Request: r, + }, nil } // generate builds the provider s describes and returns the request its @@ -327,7 +330,7 @@ func TestNegativeTemperatureRejectedForRejectingModels(t *testing.T) { } req := plainRequest() req.Temperature = ptr(-0.5) - if _, err := p.Generate(context.Background(), req); err == nil || errors.Is(err, errStop) || !strings.Contains(err.Error(), "temperature") { + if _, err := p.Generate(context.Background(), req); err == nil || errors.Is(err, llm.ErrProvider) || !strings.Contains(err.Error(), "temperature") { t.Errorf("%s: Generate = %v, want a temperature validation error", s.Model, err) } } @@ -465,7 +468,7 @@ func TestGenerateValidationErrorsNeverLeakKey(t *testing.T) { req := plainRequest() req.Temperature = ptr(-0.5) _, err = p.Generate(context.Background(), req) - if err == nil || errors.Is(err, errStop) { + if err == nil || errors.Is(err, llm.ErrProvider) { t.Fatalf("%s: Generate = %v, want a validation error", key, err) } for e := error(err); e != nil; e = errors.Unwrap(e) { @@ -518,3 +521,142 @@ func TestConfigErrorsNeverLeakKey(t *testing.T) { }) } } + +// scriptedTransport answers each request with the next scripted status and +// body, echoing the key in every failure body, and counts the requests. A +// scripted status of 0 is a transport timeout instead of a response. +type scriptedTransport struct { + statuses []int + key string + sent int +} + +func (st *scriptedTransport) RoundTrip(r *http.Request) (*http.Response, error) { + if st.sent == len(st.statuses) { + return nil, fmt.Errorf("unscripted request %d", st.sent+1) + } + status := st.statuses[st.sent] + st.sent++ + if status == 0 { + return nil, transportTimeout{} + } + header := http.Header{"Content-Type": {"application/json"}} + body := `{"error":{"message":"bad key ` + st.key + `"}}` + switch status { + case http.StatusOK: + body = `{"model":"m","choices":[{"message":{"content":"tacos"},"finish_reason":"stop"}]}` + case http.StatusTooManyRequests: + header.Set("Retry-After", "1") + } + return &http.Response{ + StatusCode: status, + Header: header, + Body: io.NopCloser(strings.NewReader(body)), + Request: r, + }, nil +} + +// transportTimeout is a net.Error that timed out. +type transportTimeout struct{} + +func (transportTimeout) Error() string { return "i/o timeout" } +func (transportTimeout) Timeout() bool { return true } +func (transportTimeout) Temporary() bool { return true } + +// useTransport makes the adapters send through rt for the rest of the test. +func useTransport(t *testing.T, rt http.RoundTripper) { + t.Helper() + orig := http.DefaultTransport + http.DefaultTransport = rt + t.Cleanup(func() { http.DefaultTransport = orig }) +} + +// The Provider FromSettings returns retries a rate-limited Generate. +func TestFromSettingsRetriesRateLimited(t *testing.T) { + st := &scriptedTransport{statuses: []int{http.StatusTooManyRequests, http.StatusOK}, key: testKey} + useTransport(t, st) + + p, err := FromSettings(llm.Settings{Provider: "openai", Model: "m"}, testKey) + if err != nil { + t.Fatalf("FromSettings: %v", err) + } + resp, err := p.Generate(context.Background(), plainRequest()) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if resp.Text != "tacos" { + t.Errorf("Text = %q, want %q", resp.Text, "tacos") + } + if st.sent != 2 { + t.Errorf("sent %d requests, want 2", st.sent) + } +} + +// When every retry fails, the error FromSettings's Provider returns keeps +// its sentinel and never holds the key the vendor echoed. +func TestFromSettingsRetryErrorNeverLeaksKey(t *testing.T) { + st := &scriptedTransport{statuses: []int{http.StatusTooManyRequests, http.StatusTooManyRequests, http.StatusServiceUnavailable}, key: testKey} + useTransport(t, st) + + p, err := FromSettings(llm.Settings{Provider: "openai", Model: "m"}, testKey) + if err != nil { + t.Fatalf("FromSettings: %v", err) + } + _, err = p.Generate(context.Background(), plainRequest()) + if !errors.Is(err, llm.ErrProvider) { + t.Fatalf("Generate = %v, want %v", err, llm.ErrProvider) + } + if st.sent != llm.MaxRetries+1 { + t.Errorf("sent %d requests, want %d", st.sent, llm.MaxRetries+1) + } + for e := error(err); e != nil; e = errors.Unwrap(e) { + if s := e.Error() + fmt.Sprintf("%#v", e); strings.Contains(s, testKey) { + t.Fatalf("error leaks the key: %s", s) + } + } +} + +// An HTTP 408 from the vendor is a 4xx answer and is not retried, though it +// matches llm.ErrTimeout. +func TestFromSettingsDoesNotRetryHTTP408(t *testing.T) { + for _, provider := range []string{"openai", "anthropic"} { + t.Run(provider, func(t *testing.T) { + st := &scriptedTransport{statuses: []int{http.StatusRequestTimeout}, key: testKey} + useTransport(t, st) + + p, err := FromSettings(llm.Settings{Provider: provider, Model: "m"}, testKey) + if err != nil { + t.Fatalf("FromSettings: %v", err) + } + _, err = p.Generate(context.Background(), plainRequest()) + var ts *llm.TimeoutStatusError + if !errors.Is(err, llm.ErrTimeout) || !errors.As(err, &ts) || ts.Status() != http.StatusRequestTimeout { + t.Fatalf("Generate = %v, want an HTTP 408 ErrTimeout", err) + } + if st.sent != 1 { + t.Errorf("sent %d requests, want 1", st.sent) + } + }) + } +} + +// A transport timeout carries no status and is retried. +func TestFromSettingsRetriesTransportTimeout(t *testing.T) { + st := &scriptedTransport{statuses: []int{0, http.StatusOK}, key: testKey} + useTransport(t, st) + + p, err := FromSettings(llm.Settings{Provider: "openai", Model: "m"}, testKey) + if err != nil { + t.Fatalf("FromSettings: %v", err) + } + resp, err := p.Generate(context.Background(), plainRequest()) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if resp.Text != "tacos" { + t.Errorf("Text = %q, want %q", resp.Text, "tacos") + } + if st.sent != 2 { + t.Errorf("sent %d requests, want 2", st.sent) + } +} diff --git a/internal/ai/llm/retry.go b/internal/ai/llm/retry.go new file mode 100644 index 0000000..ec50fa5 --- /dev/null +++ b/internal/ai/llm/retry.go @@ -0,0 +1,124 @@ +package llm + +import ( + "context" + "errors" + "math/rand/v2" + "time" +) + +const ( + // DefaultTimeout bounds a whole Generate call, retries and waits + // included, whose Request sets no Timeout. + DefaultTimeout = 60 * time.Second + // MaxRetries is how many times a failed Generate is retried, so a call + // makes at most MaxRetries+1 attempts. + MaxRetries = 2 + // MaxRetryAfter caps the wait a rate-limited vendor may ask for. + MaxRetryAfter = 20 * time.Second + // backoffBase is the first retry's backoff before jitter; each later + // retry doubles it. + backoffBase = time.Second +) + +// Retrying is a Provider that retries transient Generate failures of the +// Provider it wraps. It retries ErrRateLimited, a transport ErrTimeout and +// a 5xx ErrProvider at most MaxRetries times, and nothing else: an error +// carrying a 4xx HTTP status other than a rate limit, such as an HTTP 408 +// TimeoutStatusError, is not retried. A rate-limited retry waits for the +// vendor's RetryAfter, capped at MaxRetryAfter; every other retry, and a +// rate limit that names no wait, waits a jittered exponential backoff. The Request's Timeout, or DefaultTimeout, is one +// deadline for the whole call, retries and waits included; a caller's +// earlier context deadline wins. A retry whose wait would not end before +// that deadline is not made: the last error is returned at once. +// ListModels and Verify pass straight through. +type Retrying struct { + Provider + // now is the clock; tests replace it. + now func() time.Time + // sleep waits d or until ctx is done; tests replace it. + sleep func(ctx context.Context, d time.Duration) error + // jitter returns a random duration in [0, n) for n > 0; tests replace + // it. + jitter func(n time.Duration) time.Duration +} + +// NewRetrying returns inner wrapped with retries and the default Generate +// timeout. +func NewRetrying(inner Provider) *Retrying { + return &Retrying{ + Provider: inner, + now: time.Now, + sleep: sleep, + jitter: rand.N[time.Duration], + } +} + +// Generate sends req to the wrapped Provider, retrying transient failures. +// req.Timeout, or DefaultTimeout when it is 0, sets one deadline for the +// whole call, carried by the context each attempt receives; attempts get +// req with Timeout 0 so none restarts the clock. +func (r *Retrying) Generate(ctx context.Context, req Request) (Response, error) { + timeout := req.Timeout + if timeout == 0 { + timeout = DefaultTimeout + } + if timeout > 0 { + // A negative Timeout stays on req for the wrapped Provider's + // Validate to reject, with its key redaction. + req.Timeout = 0 + } + ctx, cancel := context.WithDeadline(ctx, r.now().Add(timeout)) + defer cancel() + deadline, _ := ctx.Deadline() + for retry := 0; ; retry++ { + resp, err := r.Provider.Generate(ctx, req) + if err == nil || retry == MaxRetries || ctx.Err() != nil { + return resp, err + } + wait, ok := r.wait(err, retry) + if !ok { + return resp, err + } + if !r.now().Add(wait).Before(deadline) { + return resp, err + } + if serr := r.sleep(ctx, wait); serr != nil { + return resp, err + } + } +} + +// wait returns how long to wait before retrying after err, the retry-th +// retry counting from 0, and false when err is not retried. +func (r *Retrying) wait(err error, retry int) (time.Duration, bool) { + var st interface{ Status() int } + if errors.As(err, &st) && st.Status() >= 400 && st.Status() <= 499 { + return 0, false + } + var rl *RateLimitedError + if errors.As(err, &rl) && rl.RetryAfter > 0 { + return min(rl.RetryAfter, MaxRetryAfter), true + } + var pe *ProviderError + switch { + case errors.Is(err, ErrRateLimited), errors.Is(err, ErrTimeout): + case errors.As(err, &pe) && pe.Status() >= 500 && pe.Status() <= 599: + default: + return 0, false + } + half := (backoffBase << retry) / 2 + return half + r.jitter(half), true +} + +// sleep waits d or until ctx is done, returning ctx's error in that case. +func sleep(ctx context.Context, d time.Duration) error { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-t.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} diff --git a/internal/ai/llm/retry_test.go b/internal/ai/llm/retry_test.go new file mode 100644 index 0000000..2095616 --- /dev/null +++ b/internal/ai/llm/retry_test.go @@ -0,0 +1,359 @@ +package llm + +import ( + "context" + "errors" + "fmt" + "slices" + "testing" + "time" +) + +// fakeClock is a clock whose sleeps return at once, advance it and are +// recorded. +type fakeClock struct { + t time.Time + sleeps []time.Duration +} + +func (c *fakeClock) now() time.Time { return c.t } + +func (c *fakeClock) sleep(ctx context.Context, d time.Duration) error { + if err := ctx.Err(); err != nil { + return err + } + c.sleeps = append(c.sleeps, d) + c.t = c.t.Add(d) + return nil +} + +// newTestRetrying returns a Retrying over a fresh Fake with a fake clock +// at the real current time, so real context deadlines compare sensibly, +// and jitter that always returns its maximum. +func newTestRetrying() (*Retrying, *Fake, *fakeClock) { + f := &Fake{} + c := &fakeClock{t: time.Now()} + r := NewRetrying(f) + r.now = c.now + r.sleep = c.sleep + r.jitter = func(n time.Duration) time.Duration { return n - 1 } + return r, f, c +} + +var retryReq = Request{Messages: []Message{{Role: RoleUser, Content: "plan dinners"}}} + +func ok(text string) GenerateResult { return GenerateResult{Response: Response{Text: text}} } + +func fail(err error) GenerateResult { return GenerateResult{Err: err} } + +func assertConsumed(t *testing.T, f *Fake, attempts int) { + t.Helper() + if g, _, _ := f.Remaining(); g != 0 { + t.Errorf("%d scripted Generate results unused", g) + } + if n := len(f.Requests()); n != attempts { + t.Errorf("made %d attempts, want %d", n, attempts) + } +} + +func TestRetryingRateLimitedThenSuccess(t *testing.T) { + r, f, c := newTestRetrying() + f.QueueGenerate(fail(NewRateLimitedError(3*time.Second, "")), ok("dinner")) + + resp, err := r.Generate(context.Background(), retryReq) + if err != nil || resp.Text != "dinner" { + t.Fatalf("Generate = %+v, %v", resp, err) + } + if want := []time.Duration{3 * time.Second}; !slices.Equal(c.sleeps, want) { + t.Errorf("sleeps = %v, want %v", c.sleeps, want) + } + assertConsumed(t, f, 2) +} + +func TestRetryingRateLimitedWaitIsCapped(t *testing.T) { + r, f, c := newTestRetrying() + f.QueueGenerate(fail(NewRateLimitedError(time.Hour, "")), ok("dinner")) + + if _, err := r.Generate(context.Background(), retryReq); err != nil { + t.Fatal(err) + } + if want := []time.Duration{20 * time.Second}; !slices.Equal(c.sleeps, want) { + t.Errorf("sleeps = %v, want %v", c.sleeps, want) + } + assertConsumed(t, f, 2) +} + +func TestRetryingRateLimitedWithoutRetryAfterBacksOff(t *testing.T) { + r, f, c := newTestRetrying() + f.QueueGenerate(fail(NewRateLimitedError(0, "")), ok("dinner")) + + if _, err := r.Generate(context.Background(), retryReq); err != nil { + t.Fatal(err) + } + if want := []time.Duration{time.Second - 1}; !slices.Equal(c.sleeps, want) { + t.Errorf("sleeps = %v, want %v", c.sleeps, want) + } +} + +func TestRetrying5xxTwiceThenSuccess(t *testing.T) { + r, f, c := newTestRetrying() + f.QueueGenerate( + fail(NewProviderError(503, []byte("busy"), "")), + // Adapters wrap the error they return. + fail(fmt.Errorf("openai: generate: %w", NewProviderError(503, []byte("busy"), ""))), + ok("dinner"), + ) + + resp, err := r.Generate(context.Background(), retryReq) + if err != nil || resp.Text != "dinner" { + t.Fatalf("Generate = %+v, %v", resp, err) + } + // Max jitter: [0.5 s, 1 s) then [1 s, 2 s). + if want := []time.Duration{time.Second - 1, 2*time.Second - 1}; !slices.Equal(c.sleeps, want) { + t.Errorf("sleeps = %v, want %v", c.sleeps, want) + } + assertConsumed(t, f, 3) +} + +func TestRetryingBackoffJitterLowerBound(t *testing.T) { + r, f, c := newTestRetrying() + r.jitter = func(time.Duration) time.Duration { return 0 } + f.QueueGenerate(fail(ErrTimeout), fail(ErrTimeout), ok("dinner")) + + if _, err := r.Generate(context.Background(), retryReq); err != nil { + t.Fatal(err) + } + if want := []time.Duration{500 * time.Millisecond, time.Second}; !slices.Equal(c.sleeps, want) { + t.Errorf("sleeps = %v, want %v", c.sleeps, want) + } +} + +func TestRetrying5xxThreeTimesFails(t *testing.T) { + r, f, c := newTestRetrying() + last := NewProviderError(503, []byte("third"), "") + f.QueueGenerate( + fail(NewProviderError(503, []byte("first"), "")), + fail(NewProviderError(503, []byte("second"), "")), + fail(last), + ) + + _, err := r.Generate(context.Background(), retryReq) + if err != last { + t.Fatalf("err = %v, want the last attempt's error", err) + } + if len(c.sleeps) != MaxRetries { + t.Errorf("slept %d times, want %d", len(c.sleeps), MaxRetries) + } + assertConsumed(t, f, 3) +} + +func TestRetryingDoesNotRetryPermanentErrors(t *testing.T) { + for name, err := range map[string]error{ + "401 auth": ErrAuth, + "400": NewProviderError(400, []byte("bad request"), ""), + "404": NewProviderError(404, []byte("no such model"), ""), + "408": NewTimeoutStatusError(408, ""), + "wrapped 408": fmt.Errorf("openai: generate: %w", NewTimeoutStatusError(408, "")), + "bad output": ErrBadOutput, + "other": errors.New("llm: request has no messages"), + } { + t.Run(name, func(t *testing.T) { + r, f, c := newTestRetrying() + f.QueueGenerate(fail(err)) + + if _, got := r.Generate(context.Background(), retryReq); got != err { + t.Fatalf("err = %v, want %v", got, err) + } + if len(c.sleeps) != 0 { + t.Errorf("slept %v", c.sleeps) + } + assertConsumed(t, f, 1) + }) + } +} + +// A transport timeout matches ErrTimeout like an HTTP 408, but carries no +// status, so it is retried. +func TestRetryingTransportTimeoutRetried(t *testing.T) { + r, f, _ := newTestRetrying() + f.QueueGenerate(fail(fmt.Errorf("openai: generate: %w", ErrTimeout)), ok("dinner")) + + if _, err := r.Generate(context.Background(), retryReq); err != nil { + t.Fatal(err) + } + assertConsumed(t, f, 2) +} + +func TestRetryingDeadlineShorterThanRetryAfter(t *testing.T) { + r, f, c := newTestRetrying() + rl := NewRateLimitedError(10*time.Second, "") + f.QueueGenerate(fail(rl)) + ctx, cancel := context.WithDeadline(context.Background(), c.t.Add(5*time.Second)) + defer cancel() + + if _, err := r.Generate(ctx, retryReq); err != rl { + t.Fatalf("err = %v, want the rate-limited error", err) + } + if len(c.sleeps) != 0 { + t.Errorf("slept %v", c.sleeps) + } + assertConsumed(t, f, 1) +} + +func TestRetryingDeadlineShorterThanBackoff(t *testing.T) { + r, f, c := newTestRetrying() + f.QueueGenerate(fail(ErrTimeout)) + // The first backoff with maximum jitter is just under 1 s. + ctx, cancel := context.WithDeadline(context.Background(), c.t.Add(time.Second-1)) + defer cancel() + + if _, err := r.Generate(ctx, retryReq); !errors.Is(err, ErrTimeout) { + t.Fatalf("err = %v, want ErrTimeout", err) + } + if len(c.sleeps) != 0 { + t.Errorf("slept %v", c.sleeps) + } + assertConsumed(t, f, 1) +} + +func TestRetryingRetriesWithinDeadline(t *testing.T) { + r, f, c := newTestRetrying() + f.QueueGenerate(fail(NewRateLimitedError(3*time.Second, "")), ok("dinner")) + ctx, cancel := context.WithDeadline(context.Background(), c.t.Add(time.Minute)) + defer cancel() + + if _, err := r.Generate(ctx, retryReq); err != nil { + t.Fatal(err) + } + assertConsumed(t, f, 2) +} + +func TestRetryingStopsWhenContextDone(t *testing.T) { + r, f, _ := newTestRetrying() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + f.QueueGenerate(fail(ErrTimeout)) + + if _, err := r.Generate(ctx, retryReq); !errors.Is(err, ErrTimeout) { + t.Fatalf("err = %v, want ErrTimeout", err) + } + assertConsumed(t, f, 1) +} + +// slowProvider is a Provider whose every attempt takes d on clock c, +// records the deadline of the context it received, and times out. +type slowProvider struct { + Provider + c *fakeClock + d time.Duration + reqs []Request + deadlines []time.Time +} + +func (p *slowProvider) Generate(ctx context.Context, req Request) (Response, error) { + p.reqs = append(p.reqs, req) + deadline, _ := ctx.Deadline() + p.deadlines = append(p.deadlines, deadline) + p.c.t = p.c.t.Add(p.d) + return Response{}, ErrTimeout +} + +func TestRetryingTimeoutBoundsWholeCall(t *testing.T) { + for _, tc := range []struct { + name string + timeout time.Duration + want time.Duration + }{ + {"default", 0, 60 * time.Second}, + {"custom", 10 * time.Second, 10 * time.Second}, + } { + t.Run(tc.name, func(t *testing.T) { + r, _, c := newTestRetrying() + // Each attempt uses most of the one deadline, so a fresh full + // timeout per attempt would allow every retry; one deadline + // leaves room for exactly one. + p := &slowProvider{c: c, d: tc.want * 5 / 6} + r.Provider = p + start := c.t + req := retryReq + req.Timeout = tc.timeout + + if _, err := r.Generate(context.Background(), req); !errors.Is(err, ErrTimeout) { + t.Fatalf("err = %v, want ErrTimeout", err) + } + if len(p.deadlines) != 2 { + t.Fatalf("made %d attempts, want 2", len(p.deadlines)) + } + for i, d := range p.deadlines { + if !d.Equal(start.Add(tc.want)) { + t.Errorf("attempt %d deadline = start+%v, want start+%v", i, d.Sub(start), tc.want) + } + if p.reqs[i].Timeout != 0 { + t.Errorf("attempt %d Timeout = %v, want 0", i, p.reqs[i].Timeout) + } + } + if want := []time.Duration{time.Second - 1}; !slices.Equal(c.sleeps, want) { + t.Errorf("sleeps = %v, want %v", c.sleeps, want) + } + }) + } +} + +func TestRetryingCallerDeadlineShorterThanTimeout(t *testing.T) { + r, _, c := newTestRetrying() + p := &slowProvider{c: c, d: time.Second} + r.Provider = p + ctx, cancel := context.WithDeadline(context.Background(), c.t.Add(5*time.Second)) + defer cancel() + + if _, err := r.Generate(ctx, retryReq); !errors.Is(err, ErrTimeout) { + t.Fatalf("err = %v, want ErrTimeout", err) + } + for i, d := range p.deadlines { + if want, _ := ctx.Deadline(); !d.Equal(want) { + t.Errorf("attempt %d deadline = %v, want caller's %v", i, d, want) + } + } +} + +func TestRetryingRejectsNegativeTimeout(t *testing.T) { + r, f, _ := newTestRetrying() + f.QueueGenerate(fail(errors.New("llm: timeout is negative"))) + req := retryReq + req.Timeout = -time.Second + + if _, err := r.Generate(context.Background(), req); err == nil { + t.Fatal("Generate accepted a negative timeout") + } + assertConsumed(t, f, 1) + if got := f.Requests()[0].Timeout; got != -time.Second { + t.Errorf("wrapped Provider got Timeout %v, want -1s to reject", got) + } +} + +func TestRetryingPassesThroughListModelsAndVerify(t *testing.T) { + r, f, _ := newTestRetrying() + f.QueueListModels(ListModelsResult{Err: ErrTimeout}) + f.QueueVerify(ErrTimeout) + + if _, err := r.ListModels(context.Background()); !errors.Is(err, ErrTimeout) { + t.Fatalf("ListModels err = %v", err) + } + if err := r.Verify(context.Background()); !errors.Is(err, ErrTimeout) { + t.Fatalf("Verify err = %v", err) + } + if _, l, v := f.Remaining(); l != 0 || v != 0 { + t.Fatalf("Remaining = %d, %d", l, v) + } +} + +func TestSleepStopsWhenContextDone(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := sleep(ctx, time.Hour); !errors.Is(err, context.Canceled) { + t.Fatalf("sleep err = %v", err) + } + if err := sleep(context.Background(), time.Millisecond); err != nil { + t.Fatalf("sleep err = %v", err) + } +}