From d90886ddc22e5dad6cf1b3d848d00d4dd04ce9f4 Mon Sep 17 00:00:00 2001 From: dev_agent Date: Sun, 4 Oct 2026 10:53:53 -0400 Subject: [PATCH] Bound whole retrying Generate call by one deadline Request.Timeout, or the 60-second default, now sets a single context deadline covering every attempt and wait, instead of restarting for each attempt. A shorter caller deadline still wins; attempts receive Timeout 0 so adapters do not restart the clock. A negative Timeout is left for the wrapped Provider to reject with its key redaction. Co-Authored-By: Claude Opus 5.5 --- internal/ai/llm/retry.go | 32 ++++++++---- internal/ai/llm/retry_test.go | 96 +++++++++++++++++++++++++++++++---- 2 files changed, 107 insertions(+), 21 deletions(-) diff --git a/internal/ai/llm/retry.go b/internal/ai/llm/retry.go index 72eb534..bb724ab 100644 --- a/internal/ai/llm/retry.go +++ b/internal/ai/llm/retry.go @@ -8,8 +8,8 @@ import ( ) const ( - // DefaultTimeout bounds each Generate attempt whose Request sets no - // Timeout. + // 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. @@ -26,9 +26,11 @@ const ( // 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. +// 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. @@ -52,12 +54,22 @@ func NewRetrying(inner Provider) *Retrying { } // 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. +// 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) { - if req.Timeout == 0 { - req.Timeout = DefaultTimeout + 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 { @@ -67,7 +79,7 @@ func (r *Retrying) Generate(ctx context.Context, req Request) (Response, error) if !ok { return resp, err } - if deadline, has := ctx.Deadline(); has && !r.now().Add(wait).Before(deadline) { + if !r.now().Add(wait).Before(deadline) { return resp, err } if serr := r.sleep(ctx, wait); serr != nil { diff --git a/internal/ai/llm/retry_test.go b/internal/ai/llm/retry_test.go index 6000f9a..1c0a7b7 100644 --- a/internal/ai/llm/retry_test.go +++ b/internal/ai/llm/retry_test.go @@ -226,20 +226,94 @@ func TestRetryingStopsWhenContextDone(t *testing.T) { 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 +// 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 +} - for _, req := range []Request{retryReq, custom} { - if _, err := r.Generate(context.Background(), req); err != nil { - t.Fatal(err) +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, DefaultTimeout}, + {"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) } } - 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 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) } }