Add provider-aware retry with backoff to the Generate path - #60
Open
hlubek wants to merge 2 commits into
Open
Conversation
Classify HTTP and network failures as retryable-transient, rate-limited, or permanent, and drive the Generate retry loop off that classification with full-jitter exponential backoff. - New llm/retry.go: classifyAPIError (408/409/5xx incl. 529 transient; 429 rate-limited honoring Retry-After / retry-after-ms / reset-*; insufficient_quota and other 4xx permanent), classifyNetworkError (network transient, ctx cancel/deadline not), header parsing, and the backoff/wait helpers. - LLMError gains Retryable / RetryAfter / StatusCode fields. - Config + LLMImpl gain MaxRetryDelay (LLM_MAX_RETRY_DELAY, default 60s) capping backoff; RetryDelay is now the backoff base rather than a fixed inter-attempt delay. - Generate returns the underlying error on exhaustion instead of masking it with an attempt count; per-attempt retries log at Debug, not Warn. - GenerateWithSchema is left on its existing retry-on-any loop (follow-up). Tests: table-driven classifier / header / delay coverage plus loop tests for transient-retry-then-success, permanent fail-fast, and insufficient_quota fail-fast. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
Reviewer's GuideIntroduce provider-aware error classification and exponential backoff with jitter into the LLM Generate path, adding structured retry metadata to LLM errors, new MaxRetryDelay configuration, and a retry loop that honors provider rate-limit headers while surfacing underlying failures. File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
hlubek
marked this pull request as ready for review
June 2, 2026 08:56
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- Consider updating LLMError.LoggableFields (or equivalent logging helpers) to include the new Retryable, RetryAfter, and StatusCode fields so the richer classification data is consistently exposed in structured logs.
- retryDelay currently uses the global math/rand without explicit seeding or isolation; consider using a per-LLMImpl *rand.Rand (or a deterministic source in tests) to avoid global contention and to make jitter behavior more predictable in controlled environments.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- Consider updating LLMError.LoggableFields (or equivalent logging helpers) to include the new Retryable, RetryAfter, and StatusCode fields so the richer classification data is consistently exposed in structured logs.
- retryDelay currently uses the global math/rand without explicit seeding or isolation; consider using a per-LLMImpl *rand.Rand (or a deterministic source in tests) to avoid global contention and to make jitter behavior more predictable in controlled environments.
## Individual Comments
### Comment 1
<location path="llm/retry.go" line_range="146-155" />
<code_context>
+// retryDelay computes the wait before the next attempt: full jitter over an
+// exponential backoff (RetryDelay × 2^attempt, capped at MaxRetryDelay),
+// floored by any server-provided delay. attempt is zero-based.
+func (l *LLMImpl) retryDelay(attempt int, serverDelay time.Duration) time.Duration {
+ base := l.RetryDelay
+ if base <= 0 {
+ base = 2 * time.Second
+ }
+
+ backoff := base
+ for i := 0; i < attempt; i++ {
+ backoff *= 2
+ if l.MaxRetryDelay > 0 && backoff > l.MaxRetryDelay {
+ backoff = l.MaxRetryDelay
+ break
+ }
+ }
+ if l.MaxRetryDelay > 0 && backoff > l.MaxRetryDelay {
+ backoff = l.MaxRetryDelay
+ }
</code_context>
<issue_to_address>
**suggestion:** Exponential backoff capping logic is duplicated and can be simplified.
The loop already enforces the `MaxRetryDelay` cap and breaks once it’s reached, so the final `if l.MaxRetryDelay > 0 && backoff > l.MaxRetryDelay` is redundant. You can drop the second check, or move the capping logic outside the loop and apply it once to make the behavior clearer.
Suggested implementation:
```golang
// retryDelay computes the wait before the next attempt: full jitter over an
// exponential backoff (RetryDelay × 2^attempt, capped at MaxRetryDelay),
// floored by any server-provided delay. attempt is zero-based.
+func (l *LLMImpl) retryDelay(attempt int, serverDelay time.Duration) time.Duration {
+ base := l.RetryDelay
+ if base <= 0 {
+ base = 2 * time.Second
+ }
+
+ backoff := base
+ for i := 0; i < attempt; i++ {
+ backoff *= 2
+ }
+ if l.MaxRetryDelay > 0 && backoff > l.MaxRetryDelay {
+ backoff = l.MaxRetryDelay
+ }
+
+ // NOTE: The rest of the function (jitter and serverDelay flooring) remains unchanged.
+ // It should follow here exactly as before.
+ // ...
+}
switch {
```
1. Ensure that the remainder of `retryDelay` (jitter application and `serverDelay` flooring) is kept exactly as it was after the removed snippet, immediately before the closing `}` I added in the replacement.
2. If the function previously did not end where I placed the `}`, adjust the position so that only the duplicated capping logic is replaced, keeping all existing behavior after the backoff calculation intact.
</issue_to_address>
### Comment 2
<location path="llm/retry_test.go" line_range="136-155" />
<code_context>
+ }
+}
+
+func TestRetryDelay(t *testing.T) {
+ l := &LLMImpl{RetryDelay: 2 * time.Second, MaxRetryDelay: 10 * time.Second}
+
+ // Backoff upper bound grows as base*2^attempt, capped at MaxRetryDelay.
+ // Full jitter keeps each sample within [0, bound].
+ bounds := map[int]time.Duration{0: 2 * time.Second, 1: 4 * time.Second, 2: 8 * time.Second, 3: 10 * time.Second, 4: 10 * time.Second}
+ for attempt, bound := range bounds {
+ for i := 0; i < 200; i++ {
+ got := l.retryDelay(attempt, 0)
+ if got < 0 || got > bound {
+ t.Fatalf("retryDelay(%d) = %v, want within [0, %v]", attempt, got, bound)
+ }
+ }
+ }
+
+ // A server-provided delay is an authoritative floor, even past the cap.
+ if got := l.retryDelay(0, 30*time.Second); got != 30*time.Second {
+ t.Errorf("server floor: retryDelay(0, 30s) = %v, want 30s", got)
+ }
</code_context>
<issue_to_address>
**suggestion (testing):** Add coverage for the default backoff behavior when `RetryDelay <= 0` and `MaxRetryDelay == 0`
The current test only exercises the path where `RetryDelay > 0` and `MaxRetryDelay > 0`, so the default-path logic isn’t covered. Please add a subtest (e.g. `RetryDelay: 0`, `MaxRetryDelay: 0`) to verify that:
- The backoff grows as `2s * 2^attempt` within the jitter range, and
- No cap is applied when `MaxRetryDelay` is unset.
This will guard against regressions in the default configuration used by most callers.
```suggestion
func TestRetryDelay(t *testing.T) {
l := &LLMImpl{RetryDelay: 2 * time.Second, MaxRetryDelay: 10 * time.Second}
// Backoff upper bound grows as base*2^attempt, capped at MaxRetryDelay.
// Full jitter keeps each sample within [0, bound].
bounds := map[int]time.Duration{
0: 2 * time.Second,
1: 4 * time.Second,
2: 8 * time.Second,
3: 10 * time.Second,
4: 10 * time.Second,
}
for attempt, bound := range bounds {
for i := 0; i < 200; i++ {
got := l.retryDelay(attempt, 0)
if got < 0 || got > bound {
t.Fatalf("retryDelay(%d) = %v, want within [0, %v]", attempt, got, bound)
}
}
}
// A server-provided delay is an authoritative floor, even past the cap.
if got := l.retryDelay(0, 30*time.Second); got != 30*time.Second {
t.Errorf("server floor: retryDelay(0, 30s) = %v, want 30s", got)
}
t.Run("defaultConfig", func(t *testing.T) {
// Default configuration: RetryDelay <= 0 and MaxRetryDelay == 0.
// Expect base backoff of 2s with no explicit cap, so the theoretical
// upper bound grows as 2s * 2^attempt, and samples stay within [0, bound].
l := &LLMImpl{}
defaultBounds := map[int]time.Duration{
0: 2 * time.Second, // 2s * 2^0
1: 4 * time.Second, // 2s * 2^1
2: 8 * time.Second, // 2s * 2^2
3: 16 * time.Second, // 2s * 2^3
4: 32 * time.Second, // 2s * 2^4
}
for attempt, bound := range defaultBounds {
for i := 0; i < 200; i++ {
got := l.retryDelay(attempt, 0)
if got < 0 || got > bound {
t.Fatalf("default retryDelay(%d) = %v, want within [0, %v]", attempt, got, bound)
}
}
}
})
}
```
</issue_to_address>
### Comment 3
<location path="llm/retry_test.go" line_range="210-228" />
<code_context>
+ }
+}
+
+func TestGenerateRetriesTransientThenSucceeds(t *testing.T) {
+ rt := &scriptedRT{steps: []func() (*http.Response, error){
+ func() (*http.Response, error) { return nil, errors.New("connection reset by peer") },
+ resp(503, `boom`),
+ resp(200, `ok`),
+ }}
+ l := newTestLLM(rt)
+
+ got, err := l.Generate(context.Background(), &Prompt{Input: "hi"})
+ if err != nil {
+ t.Fatalf("expected success after retries, got %v", err)
+ }
+ if got != "ok" {
+ t.Errorf("result = %q, want %q", got, "ok")
+ }
+ if rt.calls != 3 {
+ t.Errorf("calls = %d, want 3", rt.calls)
+ }
</code_context>
<issue_to_address>
**suggestion (testing):** Add a loop test that exercises rate-limit retries, including honoring `Retry-After`/`retry-after-ms` headers
The existing loop tests cover transient network/5xx -> success and permanent 4xx/`insufficient_quota` -> no retry, but they don’t validate the new rate-limit behavior. Please add a test where:
- A `429` with `Retry-After` or `retry-after-ms` is treated as retryable, and
- The retry loop uses the header-derived delay as a minimum (e.g., very small base backoff + larger header value, then assert `waitFor` is invoked with at least the header duration, via elapsed time or an instrumented `waitFor`).
This will exercise both the classification and backoff logic for provider rate limiting.
```suggestion
func TestGenerateRetriesTransientThenSucceeds(t *testing.T) {
rt := &scriptedRT{steps: []func() (*http.Response, error){
func() (*http.Response, error) { return nil, errors.New("connection reset by peer") },
resp(503, `boom`),
resp(200, `ok`),
}}
l := newTestLLM(rt)
got, err := l.Generate(context.Background(), &Prompt{Input: "hi"})
if err != nil {
t.Fatalf("expected success after retries, got %v", err)
}
if got != "ok" {
t.Errorf("result = %q, want %q", got, "ok")
}
if rt.calls != 3 {
t.Errorf("calls = %d, want 3", rt.calls)
}
}
func TestGenerateRetriesRateLimitedHonorsRetryAfter(t *testing.T) {
// First call: 429 with retry-after-ms, second: success.
rt := &scriptedRT{steps: []func() (*http.Response, error){
func() (*http.Response, error) {
return &http.Response{
StatusCode: http.StatusTooManyRequests,
Body: io.NopCloser(strings.NewReader(`rate limited`)),
Header: http.Header{
"retry-after-ms": []string{"25"},
},
}, nil
},
func() (*http.Response, error) {
return &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(`ok`)),
Header: make(http.Header),
}, nil
},
}}
l := newTestLLM(rt)
// Use a very small base backoff so the header value dominates.
l.RetryDelay = time.Millisecond
l.MaxRetryDelay = 10 * time.Second
var waits []time.Duration
l.waitFor = func(ctx context.Context, d time.Duration) error {
waits = append(waits, d)
return nil
}
got, err := l.Generate(context.Background(), &Prompt{Input: "hi"})
if err != nil {
t.Fatalf("expected success after rate-limit retries, got %v", err)
}
if got != "ok" {
t.Errorf("result = %q, want %q", got, "ok")
}
if rt.calls != 2 {
t.Fatalf("calls = %d, want 2", rt.calls)
}
if len(waits) == 0 {
t.Fatalf("expected waitFor to be called at least once")
}
var maxWait time.Duration
for _, d := range waits {
if d > maxWait {
maxWait = d
}
}
const headerWait = 25 * time.Millisecond
if maxWait < headerWait {
t.Errorf("max waitFor duration = %v, want at least %v derived from retry-after-ms header", maxWait, headerWait)
}
}
```
</issue_to_address>
### Comment 4
<location path="llm/retry_test.go" line_range="230-239" />
<code_context>
+ }
+}
+
+func TestGeneratePermanentFailsFast(t *testing.T) {
+ rt := &scriptedRT{steps: []func() (*http.Response, error){
+ resp(400, `{"error":{"type":"invalid_request_error"}}`),
+ }}
+ l := newTestLLM(rt)
+
+ _, err := l.Generate(context.Background(), &Prompt{Input: "hi"})
+ if err == nil {
+ t.Fatal("expected permanent error, got nil")
+ }
+ var le *LLMError
+ if !errors.As(err, &le) || le.StatusCode != 400 {
+ t.Errorf("want underlying 400 LLMError, got %v", err)
+ }
+ if rt.calls != 1 {
+ t.Errorf("calls = %d, want 1 (no retry on permanent)", rt.calls)
+ }
</code_context>
<issue_to_address>
**suggestion (testing):** Consider adding tests for context cancellation during the retry wait to validate `waitFor`/`shouldRetry` behavior
There are no tests asserting that this context-cancellable retry logic actually respects cancellation. Please add tests such as:
- Use a `context.WithCancel`, trigger a retryable error, cancel the context before/during the wait, and assert `Generate` stops retrying and returns the context error.
- A focused unit test for `waitFor` where `d > 0` and the context is cancelled shortly after it starts, asserting it returns promptly with the appropriate context error.
This helps ensure the loop won’t hang, leak goroutines, or continue retrying after the caller has cancelled.
Suggested implementation:
```golang
if rt.calls != 3 {
t.Errorf("calls = %d, want 3", rt.calls)
}
}
func TestGenerateContextCancelDuringRetry(t *testing.T) {
rt := &scriptedRT{steps: []func() (*http.Response, error){
// First call returns a retryable error so Generate will attempt to wait and retry.
resp(500, `{"error":{"type":"api_error"}}`),
}}
l := newTestLLM(rt)
// Use a short timeout so the context is cancelled during the retry wait.
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
start := time.Now()
_, err := l.Generate(ctx, &Prompt{Input: "hi"})
elapsed := time.Since(start)
if err == nil {
t.Fatal("expected error due to context cancellation, got nil")
}
if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) {
t.Fatalf("expected context cancellation error, got %v", err)
}
// Ensure we did not block for an excessively long time (i.e., waited out the full backoff).
if elapsed > 500*time.Millisecond {
t.Fatalf("Generate took %v, want <= 500ms (must respect context cancellation)", elapsed)
}
}
func TestWaitForContextCancel(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// Cancel shortly after waitFor starts.
go func() {
time.Sleep(10 * time.Millisecond)
cancel()
}()
start := time.Now()
err := waitFor(ctx, time.Hour)
elapsed := time.Since(start)
if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("expected context cancellation from waitFor, got %v", err)
}
if elapsed > 500*time.Millisecond {
t.Fatalf("waitFor took %v, want <= 500ms (must return promptly on cancellation)", elapsed)
}
}
for _, c := range cases {
```
1. Ensure `waitFor` is exported or in the same package (which it appears to be) so that `TestWaitForContextCancel` can call it directly.
2. Confirm that `llm/retry_test.go` already imports `context`, `errors`, and `time`. If any are missing, add them to the import block:
- `context`
- `errors`
- `time`
3. If your retry classification for a `500` response does *not* treat it as retryable, change the status code and body in `TestGenerateContextCancelDuringRetry` to whatever your implementation considers retryable (e.g., a different status or error type) so that the code path invoking `waitFor` is exercised.
4. Adjust the timing thresholds (`50ms` timeout, `500ms` upper bound) if your CI environment is particularly slow, while keeping them tight enough to assert that cancellation is respected and the full backoff is not waited out.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
… tests - retryDelay: cap the base up front and rely on the in-loop cap, dropping the redundant trailing cap while keeping overflow safety at high attempt counts. - LLMError.LoggableFields: include status_code / retryable / retry_after when meaningful, so classification data reaches structured logs without noising up non-HTTP errors. - Tests: default-config backoff path (RetryDelay<=0, MaxRetryDelay==0); a rate-limit (429 + Retry-After) retry loop asserting the server hint is honored over base backoff; and context-cancellation coverage for waitFor and the retry loop. Declined the suggestion to swap global math/rand for a per-instance *rand.Rand: the global source is concurrency-safe whereas a *rand.Rand is not (it would need its own mutex), and the tests assert ranges rather than exact values, so determinism is unnecessary. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Author
|
Thanks for the review — pushed 4119d4d addressing it. Summary of the two high-level notes: Logging: Done. Global The four inline suggestions are addressed in replies on each thread (the |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The
Generatepath currently makes a single attempt and collapses every non-200 response into a genericErrorTypeAPI, discarding the response headers. As a result transient failures (connection resets, 429 rate limits, 5xx) fail immediately, and a server'sRetry-Afteris never honored.This adds a classification + backoff layer:
llm/retry.go— classifies HTTP and network errors as:Retry-After(seconds), OpenAIretry-after-ms(milliseconds), andx-ratelimit-reset-*duration strings as a fallback;insufficient_quota.MaxRetriesand fully ctx-cancellable.MaxRetryDelayconfig (LLM_MAX_RETRY_DELAY, default60s) caps backoff growth.Generatenow returns the underlying error on exhaustion instead of masking it with"failed after N attempts", so callers can classify the failure. Per-attempt retries log at Debug rather than Warn.Behavior change to note:
RetryDelayshifts meaning from a fixed inter-attempt delay to the base of the exponential backoff curve. Existing callers that set it get exponential growth (floored at the same value) instead of a flat wait.Scope:
GenerateWithSchemais intentionally left on its existing retry-on-any-error loop. Routing it through the same policy is a sensible follow-up, but doing so here would change its re-ask-on-malformed-output semantics, so it's deferred.Tests are table-driven for the classifier, header parsing (including OpenAI
reset-*durations andinsufficient_quota), and delay/jitter bounds, plus loop tests covering transient-retry-then-success, permanent fail-fast, andinsufficient_quotafail-fast.Summary by Sourcery
Introduce provider-aware retry and backoff handling for LLM generation requests, including classification of transient, rate-limit, and permanent errors.
Enhancements:
Tests: