fix: address review feedback — double-sleep bug, dead test field, ctx-cancel test, nits

Assisted-By: docker-agent

Author: simonferquel-clanker

pkg/modelerrors/modelerrors_test.go

@@ -330,24 +330,25 @@ func TestClassifyModelError(t *testing.T) {
 		retryAfter      time.Duration // value to pass as the retryAfter param
 		wantRetryable   bool
 		wantRateLimited bool
+		wantRetryAfter  time.Duration // expected retryAfterOut
 	}{
-		{name: "nil", err: nil, wantRetryable: false, wantRateLimited: false},
-		{name: "context canceled", err: context.Canceled, wantRetryable: false, wantRateLimited: false},
-		{name: "context deadline exceeded", err: context.DeadlineExceeded, wantRetryable: false, wantRateLimited: false},
-		{name: "context overflow", err: errors.New("prompt is too long: 200000 tokens > 100000 maximum"), wantRetryable: false, wantRateLimited: false},
+		{name: "nil", err: nil, wantRetryable: false, wantRateLimited: false, wantRetryAfter: 0},
+		{name: "context canceled", err: context.Canceled, wantRetryable: false, wantRateLimited: false, wantRetryAfter: 0},
+		{name: "context deadline exceeded", err: context.DeadlineExceeded, wantRetryable: false, wantRateLimited: false, wantRetryAfter: 0},
+		{name: "context overflow", err: errors.New("prompt is too long: 200000 tokens > 100000 maximum"), wantRetryable: false, wantRateLimited: false, wantRetryAfter: 0},
 		// 429 rate limit cases (retryAfter passed in by caller from provider.ExtractRetryAfter)
-		{name: "429 message only, no header", err: errors.New("POST /v1/chat: 429 Too Many Requests"), retryAfter: 0, wantRetryable: false, wantRateLimited: true},
-		{name: "429 message only, with header", err: errors.New("POST /v1/chat: 429 Too Many Requests"), retryAfter: 30 * time.Second, wantRetryable: false, wantRateLimited: true},
+		{name: "429 message only, no header", err: errors.New("POST /v1/chat: 429 Too Many Requests"), retryAfter: 0, wantRetryable: false, wantRateLimited: true, wantRetryAfter: 0},
+		{name: "429 message only, with header", err: errors.New("POST /v1/chat: 429 Too Many Requests"), retryAfter: 30 * time.Second, wantRetryable: false, wantRateLimited: true, wantRetryAfter: 30 * time.Second},
 		// Retryable server errors
-		{name: "500 internal server error", err: errors.New("500 internal server error"), wantRetryable: true, wantRateLimited: false},
-		{name: "502 bad gateway", err: errors.New("502 bad gateway"), wantRetryable: true, wantRateLimited: false},
-		{name: "503 service unavailable", err: errors.New("503 service unavailable"), wantRetryable: true, wantRateLimited: false},
-		{name: "504 gateway timeout", err: errors.New("504 gateway timeout"), wantRetryable: true, wantRateLimited: false},
+		{name: "500 internal server error", err: errors.New("500 internal server error"), wantRetryable: true, wantRateLimited: false, wantRetryAfter: 0},
+		{name: "502 bad gateway", err: errors.New("502 bad gateway"), wantRetryable: true, wantRateLimited: false, wantRetryAfter: 0},
+		{name: "503 service unavailable", err: errors.New("503 service unavailable"), wantRetryable: true, wantRateLimited: false, wantRetryAfter: 0},
+		{name: "504 gateway timeout", err: errors.New("504 gateway timeout"), wantRetryable: true, wantRateLimited: false, wantRetryAfter: 0},
 		// Non-retryable errors (message-based, no SDK types needed here)
-		{name: "401 unauthorized", err: errors.New("401 unauthorized"), wantRetryable: false, wantRateLimited: false},
-		{name: "403 forbidden", err: errors.New("403 forbidden"), wantRetryable: false, wantRateLimited: false},
+		{name: "401 unauthorized", err: errors.New("401 unauthorized"), wantRetryable: false, wantRateLimited: false, wantRetryAfter: 0},
+		{name: "403 forbidden", err: errors.New("403 forbidden"), wantRetryable: false, wantRateLimited: false, wantRetryAfter: 0},
 		// Network errors
-		{name: "network timeout", err: &mockTimeoutError{}, wantRetryable: true, wantRateLimited: false},
+		{name: "network timeout", err: &mockTimeoutError{}, wantRetryable: true, wantRateLimited: false, wantRetryAfter: 0},
 	}
 
 	for _, tt := range tests {
@@ -356,7 +357,7 @@ func TestClassifyModelError(t *testing.T) {
 			retryable, rateLimited, retryAfterOut := ClassifyModelError(tt.err, tt.retryAfter)
 			assert.Equal(t, tt.wantRetryable, retryable, "retryable mismatch")
 			assert.Equal(t, tt.wantRateLimited, rateLimited, "rateLimited mismatch")
-			assert.GreaterOrEqual(t, retryAfterOut, time.Duration(0), "retryAfterOut should never be negative")
+			assert.Equal(t, tt.wantRetryAfter, retryAfterOut, "retryAfter mismatch")
 		})
 	}
 

pkg/runtime/fallback.go

@@ -69,15 +69,6 @@ func logFallbackAttempt(agentName string, model modelWithFallback, attempt, maxR
 	}
 }
 
-// logRetryBackoff logs when we're backing off before a retry
-func logRetryBackoff(agentName, modelID string, attempt int, backoff time.Duration) {
-	slog.Debug("Backing off before retry",
-		"agent", agentName,
-		"model", modelID,
-		"attempt", attempt+1,
-		"backoff", backoff)
-}
-
 // getCooldownState returns the current cooldown state for an agent (thread-safe).
 // Returns nil if no cooldown is active or if cooldown has expired.
 // Expired entries are evicted to prevent stale state accumulation.
@@ -228,15 +219,6 @@ func (r *LocalRuntime) tryModelWithFallback(
 				return streamResult{}, nil, ctx.Err()
 			}
 
-			// Apply backoff before retry (not on first attempt of each model)
-			if attempt > 0 {
-				backoff := modelerrors.CalculateBackoff(attempt - 1)
-				logRetryBackoff(a.Name(), modelEntry.provider.ID(), attempt, backoff)
-				if !modelerrors.SleepWithContext(ctx, backoff) {
-					return streamResult{}, nil, ctx.Err()
-				}
-			}
-
 			// Emit fallback event when transitioning to a new model (but not when starting in cooldown)
 			if chainIdx > startIndex && attempt == 0 {
 				logFallbackAttempt(a.Name(), modelEntry, attempt, fallbackRetries, lastErr)
@@ -272,7 +254,7 @@ func (r *LocalRuntime) tryModelWithFallback(
 					return streamResult{}, nil, err
 				}
 
-				decision := r.handleModelError(ctx, err, a, modelEntry, attempt, hasFallbacks, &primaryFailedWithNonRetryable)
+				decision := handleModelError(ctx, err, a, modelEntry, attempt, hasFallbacks, &primaryFailedWithNonRetryable)
 				if decision == retryDecisionReturn {
 					return streamResult{}, nil, ctx.Err()
 				} else if decision == retryDecisionBreak {
@@ -292,7 +274,7 @@ func (r *LocalRuntime) tryModelWithFallback(
 					return streamResult{}, nil, err
 				}
 
-				decision := r.handleModelError(ctx, err, a, modelEntry, attempt, hasFallbacks, &primaryFailedWithNonRetryable)
+				decision := handleModelError(ctx, err, a, modelEntry, attempt, hasFallbacks, &primaryFailedWithNonRetryable)
 				if decision == retryDecisionReturn {
 					return streamResult{}, nil, ctx.Err()
 				} else if decision == retryDecisionBreak {
@@ -335,10 +317,12 @@ func (r *LocalRuntime) tryModelWithFallback(
 type retryDecision int
 
 const (
-	// retryDecisionContinue means retry the same model (backoff already applied).
-	retryDecisionContinue retryDecision = iota
 	// retryDecisionBreak means skip to the next model in the fallback chain.
-	retryDecisionBreak
+	// This is the zero value — safe default: skip to next model rather than
+	// accidentally retrying or returning early.
+	retryDecisionBreak retryDecision = iota
+	// retryDecisionContinue means retry the same model (sleep already applied).
+	retryDecisionContinue
 	// retryDecisionReturn means context was cancelled; return immediately.
 	retryDecisionReturn
 )
@@ -348,9 +332,12 @@ const (
 //   - retryDecisionBreak    — non-retryable error or 429 with fallbacks; skip to next model
 //   - retryDecisionContinue — retryable error or 429 without fallbacks; retry same model
 //
+// All sleeping (both 5xx backoff and 429 Retry-After) is performed here so the
+// outer loop never needs its own sleep path.
+//
 // Side-effect: sets *primaryFailedWithNonRetryable when the primary model fails with a
 // non-retryable (or rate-limited-with-fallbacks) error.
-func (r *LocalRuntime) handleModelError(
+func handleModelError(
 	ctx context.Context,
 	err error,
 	a *agent.Agent,
@@ -363,11 +350,12 @@ func (r *LocalRuntime) handleModelError(
 
 	if rateLimited {
 		if hasFallbacks {
-			// Fallbacks available → skip to next model immediately (existing behaviour).
-			slog.Warn("Rate limited with fallbacks available, skipping to next model",
+			// Fallbacks available → skip to next model immediately.
+			slog.Warn("Rate limited, skipping model",
 				"agent", a.Name(),
 				"model", modelEntry.provider.ID(),
-				"retry_after", retryAfter)
+				"retry_after", retryAfter,
+				"error", err)
 			if !modelEntry.isFallback {
 				*primaryFailedWithNonRetryable = true
 			}
@@ -391,7 +379,8 @@ func (r *LocalRuntime) handleModelError(
 			"model", modelEntry.provider.ID(),
 			"attempt", attempt+1,
 			"wait", waitDuration,
-			"retry_after_from_header", retryAfter > 0)
+			"retry_after_from_header", retryAfter > 0,
+			"error", err)
 		if !modelerrors.SleepWithContext(ctx, waitDuration) {
 			return retryDecisionReturn
 		}
@@ -409,10 +398,16 @@ func (r *LocalRuntime) handleModelError(
 		return retryDecisionBreak
 	}
 
+	// Retryable (5xx, timeouts): sleep with backoff then retry same model.
+	waitDuration := modelerrors.CalculateBackoff(attempt)
 	slog.Warn("Retryable error from model",
 		"agent", a.Name(),
 		"model", modelEntry.provider.ID(),
 		"attempt", attempt+1,
+		"wait", waitDuration,
 		"error", err)
+	if !modelerrors.SleepWithContext(ctx, waitDuration) {
+		return retryDecisionReturn
+	}
 	return retryDecisionContinue
 }

pkg/runtime/fallback_test.go

@@ -613,3 +613,54 @@ func TestFallback500RetryableWithBackoff(t *testing.T) {
 		assert.Equal(t, 2, primary.callCount, "primary should be called twice: 1 failure + 1 success")
 	})
 }
+
+func TestFallback429WithoutFallbacksContextCancelled(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		// Model always returns 429 with no fallbacks; handleModelError will sleep before
+		// retrying. We cancel the context while it is sleeping and verify that RunStream
+		// returns promptly (stream channel closed) rather than hanging until the backoff
+		// expires.
+		primary := &failingProvider{
+			id:  "primary/always-429",
+			err: errors.New("POST /v1/chat/completions: 429 Too Many Requests"),
+		}
+
+		root := agent.New("root", "test",
+			agent.WithModel(primary),
+			// No fallback models; 429 will be retried with backoff.
+			// Use many retries to ensure the runtime would block for a long time
+			// without context cancellation.
+			agent.WithFallbackRetries(5),
+		)
+
+		tm := team.New(team.WithAgents(root))
+		rt, err := NewLocalRuntime(tm, WithSessionCompaction(false), WithModelStore(mockModelStore{}))
+		require.NoError(t, err)
+
+		ctx, cancel := context.WithCancel(t.Context())
+		defer cancel()
+
+		sess := session.New(session.WithUserMessage("test"))
+		sess.Title = "429 Context Cancel Test"
+
+		// Cancel the context from a goroutine once all goroutines in the bubble are
+		// durably blocked (i.e., the retry sleep has started). synctest.Wait() returns
+		// only when every goroutine is blocked, so at that point the runtime is mid-sleep.
+		go func() {
+			synctest.Wait()
+			cancel()
+		}()
+
+		// Drain the stream. If context cancellation is properly handled, RunStream
+		// must close the channel promptly; if not, the bubble's fake time would never
+		// advance and the test would deadlock.
+		var eventCount int
+		for range rt.RunStream(ctx, sess) {
+			eventCount++
+		}
+		// The primary was called at least once before the sleep started.
+		// We can't assert on eventCount because no content is produced — just verify
+		// the channel closed (loop above completed without deadlock).
+		_ = eventCount
+	})
+}