fix: make checkAndWaitForRateLimit context-cancellable

pkg/cli/logs_orchestrator.go

@@ -367,7 +367,13 @@ outerLoop:
 		// always sleeps at least APICallCooldown but will also block until the
 		// reset window when the remaining budget is nearly exhausted.
 		if iteration > 0 {
-			if rlErr := checkAndWaitForRateLimit(verbose); rlErr != nil {
+			if rlErr := checkAndWaitForRateLimit(activeCtx, verbose); rlErr != nil {
+				if errors.Is(rlErr, context.Canceled) || errors.Is(rlErr, context.DeadlineExceeded) {
+					// Context was cancelled or timed out during the rate-limit wait.
+					// Use continue (not break) so the top-of-loop activeCtx.Done() path
+					// preserves Canceled vs DeadlineExceeded behavior.
+					continue
+				}
 				logsOrchestratorLog.Printf("Rate limit check failed (using static cooldown): %v", rlErr)
 			}
 		}

pkg/cli/logs_rate_limit.go

@@ -11,7 +11,9 @@
 package cli
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"os"
 	"time"
@@ -22,6 +24,7 @@ import (
 )
 
 var logsRateLimitLog = logger.New("cli:logs_rate_limit")
+var fetchRateLimitFunc = fetchRateLimit
 
 // rateLimitResponse models the JSON returned by `gh api rate_limit`.
 // Only the "core" resource bucket is used because log downloads and
@@ -70,21 +73,49 @@ func fetchRateLimit() (rateLimitResource, error) {
 	return resp.Resources.Core, nil
 }
 
+// sleepWithContext pauses for duration d and returns nil when the timer fires.
+// If ctx is cancelled before the timer expires, it stops the timer and returns
+// ctx.Err() so callers can propagate cancellation immediately.
+func sleepWithContext(ctx context.Context, d time.Duration) error {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
+	timer := time.NewTimer(d)
+	defer timer.Stop()
+	select {
+	case <-timer.C:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 // checkAndWaitForRateLimit queries the GitHub API rate limit and sleeps until
 // the reset window when the remaining core request budget falls at or below
 // RateLimitThreshold.  It always waits at least APICallCooldown between
 // successive calls so that even when requests are plentiful the orchestrator
 // does not hammer the API.
 //
+// ctx is checked on every sleep so that a user cancellation (Ctrl-C) or
+// deadline expiry wakes the function early and propagates the context error.
+//
 // If the rate limit cannot be fetched (e.g. network error) the function falls
 // back to the static APICallCooldown sleep and returns the error so callers
 // can decide whether to surface it.
-func checkAndWaitForRateLimit(verbose bool) error {
-	rl, err := fetchRateLimit()
+func checkAndWaitForRateLimit(ctx context.Context, verbose bool) error {
+	rl, err := fetchRateLimitFunc()
 	if err != nil {
 		// Best-effort: fall back to static cooldown so the caller can continue.
 		logsRateLimitLog.Printf("Could not fetch rate limit, using static cooldown: %v", err)
-		time.Sleep(APICallCooldown)
+		if sleepErr := sleepWithContext(ctx, APICallCooldown); sleepErr != nil {
+			return fmt.Errorf("rate-limit fetch failed and context was canceled or timed out during fallback cooldown: %w", errors.Join(err, sleepErr))
+		}
 		return err
 	}
 
@@ -94,8 +125,7 @@ func checkAndWaitForRateLimit(verbose bool) error {
 		if waitDur <= 0 {
 			// Reset has already passed; apply minimal cooldown and carry on.
 			logsRateLimitLog.Print("Rate limit reset has already passed, applying minimal cooldown")
-			time.Sleep(APICallCooldown)
-			return nil
+			return sleepWithContext(ctx, APICallCooldown)
 		}
 
 		// Add a small buffer so we don't resume right on the boundary.
@@ -107,8 +137,7 @@ func checkAndWaitForRateLimit(verbose bool) error {
 		)
 		fmt.Fprintln(os.Stderr, console.FormatWarningMessage(msg))
 		logsRateLimitLog.Printf("Sleeping for rate limit reset: duration=%s", waitDur)
-		time.Sleep(waitDur)
-		return nil
+		return sleepWithContext(ctx, waitDur)
 	}
 
 	if verbose {
@@ -118,6 +147,5 @@ func checkAndWaitForRateLimit(verbose bool) error {
 	}
 
 	// Even when budget is healthy, apply the minimum inter-call cooldown.
-	time.Sleep(APICallCooldown)
-	return nil
+	return sleepWithContext(ctx, APICallCooldown)
 }

pkg/cli/logs_rate_limit_test.go

@@ -3,7 +3,9 @@
 package cli
 
 import (
+	"context"
 	"encoding/json"
+	stderrors "errors"
 	"strconv"
 	"testing"
 	"time"
@@ -83,3 +85,100 @@ func TestRateLimitResourceIsBelowThreshold(t *testing.T) {
 func jsonInt(n int64) string {
 	return strconv.FormatInt(n, 10)
 }
+
+// TestSleepWithContextCancellation verifies that sleepWithContext returns ctx.Err()
+// immediately when the context is cancelled before the timer fires.
+func TestSleepWithContextCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel immediately
+
+	start := time.Now()
+	err := sleepWithContext(ctx, 10*time.Second)
+	elapsed := time.Since(start)
+
+	require.ErrorIs(t, err, context.Canceled)
+	assert.Less(t, elapsed, time.Second, "sleepWithContext should return quickly when context is already cancelled")
+}
+
+// TestSleepWithContextDeadlineExceeded verifies that sleepWithContext respects a
+// deadline that expires before the sleep duration.
+func TestSleepWithContextDeadlineExceeded(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+
+	start := time.Now()
+	err := sleepWithContext(ctx, 10*time.Second)
+	elapsed := time.Since(start)
+
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	assert.Less(t, elapsed, time.Second, "sleepWithContext should return as soon as the deadline expires")
+}
+
+// TestSleepWithContextTimerFires verifies that sleepWithContext returns nil when the
+// timer fires before context cancellation.
+func TestSleepWithContextTimerFires(t *testing.T) {
+	ctx := context.Background()
+	start := time.Now()
+	err := sleepWithContext(ctx, 5*time.Millisecond)
+	elapsed := time.Since(start)
+	require.NoError(t, err, "sleepWithContext should return nil when timer fires normally")
+	assert.Less(t, elapsed, time.Second, "timer should have fired and returned promptly")
+}
+
+func TestSleepWithContextNilContext(t *testing.T) {
+	var nilCtx context.Context
+	start := time.Now()
+	err := sleepWithContext(nilCtx, 2*time.Millisecond)
+	elapsed := time.Since(start)
+	require.NoError(t, err, "nil context should fall back to background context")
+	assert.Less(t, elapsed, time.Second, "nil context should not block longer than timer duration")
+}
+
+func TestSleepWithContextAlreadyCanceled(t *testing.T) {
+	canceledCtx, cancel := context.WithCancel(context.Background())
+	cancel()
+	start := time.Now()
+	err := sleepWithContext(canceledCtx, 2*time.Millisecond)
+	elapsed := time.Since(start)
+	require.ErrorIs(t, err, context.Canceled)
+	assert.Less(t, elapsed, time.Second, "already-canceled context should return promptly")
+}
+
+func TestCheckAndWaitForRateLimitContextCancelled(t *testing.T) {
+	oldFetchRateLimitFunc := fetchRateLimitFunc
+	fetchRateLimitFunc = func() (rateLimitResource, error) {
+		return rateLimitResource{
+			Limit:     5000,
+			Remaining: 0,
+			Reset:     time.Now().Add(10 * time.Minute).Unix(),
+		}, nil
+	}
+	t.Cleanup(func() { fetchRateLimitFunc = oldFetchRateLimitFunc })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	start := time.Now()
+	err := checkAndWaitForRateLimit(ctx, false)
+	elapsed := time.Since(start)
+
+	require.ErrorIs(t, err, context.Canceled)
+	assert.Less(t, elapsed, 100*time.Millisecond, "cancelled context should interrupt rate-limit wait promptly")
+}
+
+func TestCheckAndWaitForRateLimitFetchErrorAndContextDone(t *testing.T) {
+	oldFetchRateLimitFunc := fetchRateLimitFunc
+	expectedFetchErr := stderrors.New("fetch failure")
+	fetchRateLimitFunc = func() (rateLimitResource, error) {
+		return rateLimitResource{}, expectedFetchErr
+	}
+	t.Cleanup(func() { fetchRateLimitFunc = oldFetchRateLimitFunc })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err := checkAndWaitForRateLimit(ctx, false)
+	require.Error(t, err)
+	require.ErrorIs(t, err, expectedFetchErr)
+	require.ErrorIs(t, err, context.Canceled)
+}