Retry, backoff and default Generate timeout (#48)

Co-authored-by: dev_agent <dev_agent@cubecraftlabs.com>
Co-committed-by: dev_agent <dev_agent@cubecraftlabs.com>

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)
 	}

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 }

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)
+	}
+}

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)
 	}

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) {

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)
+	}
+}

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()
+	}
+}

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)
+	}
+}