diff --git a/internal/ai/llm/retry.go b/internal/ai/llm/retry.go new file mode 100644 index 0000000..72eb534 --- /dev/null +++ b/internal/ai/llm/retry.go @@ -0,0 +1,107 @@ +package llm + +import ( + "context" + "errors" + "math/rand/v2" + "time" +) + +const ( + // DefaultTimeout bounds each Generate attempt 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, ErrTimeout and a 5xx +// ErrProvider at most MaxRetries times, and nothing else. 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. A retry whose wait would not end before the +// caller's context 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. +// A req.Timeout of 0 becomes DefaultTimeout; either bounds each attempt, +// not the whole call, which ctx bounds. +func (r *Retrying) Generate(ctx context.Context, req Request) (Response, error) { + if req.Timeout == 0 { + req.Timeout = DefaultTimeout + } + 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 deadline, has := ctx.Deadline(); has && !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 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..6000f9a --- /dev/null +++ b/internal/ai/llm/retry_test.go @@ -0,0 +1,271 @@ +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{MaxRetryAfter}; !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"), ""), + "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) + }) + } +} + +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) +} + +func TestRetryingDefaultTimeout(t *testing.T) { + r, f, _ := newTestRetrying() + f.QueueGenerate(ok("a"), ok("b")) + custom := retryReq + custom.Timeout = 5 * time.Second + + for _, req := range []Request{retryReq, custom} { + if _, err := r.Generate(context.Background(), req); err != nil { + t.Fatal(err) + } + } + got := f.Requests() + if got[0].Timeout != DefaultTimeout || got[1].Timeout != 5*time.Second { + t.Fatalf("timeouts = %v, %v; want %v, 5s", got[0].Timeout, got[1].Timeout, DefaultTimeout) + } +} + +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) + } +}