Add retrying Provider wrapper with backoff and default timeout

llm.NewRetrying wraps a Provider and retries ErrRateLimited, ErrTimeout and 5xx ErrProvider at most twice. A rate limit waits RetryAfter, capped at 20 s; other retries use jittered exponential backoff. ErrAuth, 4xx and other errors are never retried, and no retry is made whose wait would not end before the caller's context deadline. Generate gets a 60 s per-attempt timeout when Request.Timeout is 0. Maestro task: 1a882923-fdb4-4576-ab82-9529e9974b26 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-04 10:48:07 -04:00
parent e105cdb9c5
commit cbf24e63e1
2 changed files with 378 additions and 0 deletions
@@ -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()
+	}
+}
@@ -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)
+	}
+}