From e6fd24db0a1d4ef71250ba713e9d01a96cacf3d9 Mon Sep 17 00:00:00 2001
From: dev_agent <dev_agent@cubecraftlabs.com>
Date: Mon, 5 Oct 2026 11:41:23 -0400
Subject: [PATCH] Return context error when a retry backoff wait is cancelled
 (#64)

Co-authored-by: dev_agent <dev_agent@cubecraftlabs.com>
---
 internal/ai/llm/retry.go      |  7 ++++---
 internal/ai/llm/retry_test.go | 37 +++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/internal/ai/llm/retry.go b/internal/ai/llm/retry.go
index ec50fa5..ec8cbc4 100644
--- a/internal/ai/llm/retry.go
+++ b/internal/ai/llm/retry.go
@@ -30,7 +30,8 @@ const (
 // 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.
+// that deadline is not made: the last error is returned at once. A wait
+// cut short by cancellation or the deadline returns the context's error.
 // ListModels and Verify pass straight through.
 type Retrying struct {
 	Provider
@@ -83,8 +84,8 @@ func (r *Retrying) Generate(ctx context.Context, req Request) (Response, error)
 		if !r.now().Add(wait).Before(deadline) {
 			return resp, err
 		}
-		if serr := r.sleep(ctx, wait); serr != nil {
-			return resp, err
+		if err := r.sleep(ctx, wait); err != nil {
+			return Response{}, err
 		}
 	}
 }
diff --git a/internal/ai/llm/retry_test.go b/internal/ai/llm/retry_test.go
index 2095616..8548283 100644
--- a/internal/ai/llm/retry_test.go
+++ b/internal/ai/llm/retry_test.go
@@ -240,6 +240,43 @@ func TestRetryingStopsWhenContextDone(t *testing.T) {
 	assertConsumed(t, f, 1)
 }
 
+func TestRetryingWaitInterruptedReturnsContextError(t *testing.T) {
+	t.Run("cancelled", func(t *testing.T) {
+		r, f, _ := newTestRetrying()
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		r.sleep = func(ctx context.Context, d time.Duration) error {
+			cancel()
+			return sleep(ctx, d)
+		}
+		f.QueueGenerate(fail(NewRateLimitedError(3*time.Second, "")))
+
+		_, err := r.Generate(ctx, retryReq)
+		if !errors.Is(err, context.Canceled) || errors.Is(err, ErrRateLimited) {
+			t.Fatalf("err = %v, want context.Canceled", err)
+		}
+		assertConsumed(t, f, 1)
+	})
+	t.Run("deadline", func(t *testing.T) {
+		r, f, _ := newTestRetrying()
+		// The fake clock stands still, so the 500ms backoff fits before
+		// the caller's deadline; the real wait outlasts it.
+		r.jitter = func(time.Duration) time.Duration { return 0 }
+		r.sleep = func(ctx context.Context, _ time.Duration) error {
+			return sleep(ctx, time.Hour)
+		}
+		ctx, cancel := context.WithTimeout(context.Background(), 600*time.Millisecond)
+		defer cancel()
+		f.QueueGenerate(fail(NewRateLimitedError(0, "")))
+
+		_, err := r.Generate(ctx, retryReq)
+		if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, ErrRateLimited) {
+			t.Fatalf("err = %v, want context.DeadlineExceeded", 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 {