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 <noreply@anthropic.com>

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 {

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