From 382975a7af82199760ec722a880f86f7df6cabd2 Mon Sep 17 00:00:00 2001
From: XYSK-lilong007 <267018309+XYSK-lilong007@users.noreply.github.com>
Date: Fri, 13 Mar 2026 08:31:04 +0800
Subject: [PATCH] fix(agent): retry transient provider failures

---
 pkg/agent/loop.go      | 14 +++++++++----
 pkg/agent/loop_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 33da33e92..16f8b702d 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -1128,8 +1128,9 @@ func (al *AgentLoop) runLLMIteration(
 			return agent.Provider.Chat(ctx, messages, providerToolDefs, activeModel, llmOpts)
 		}
 
-		// Retry loop for context/token errors
+		// Retry loop for context/token errors and transient single-provider failures.
 		maxRetries := 2
+		singleCandidate := len(activeCandidates) == 1
 		for retry := 0; retry <= maxRetries; retry++ {
 			response, err = callLLM()
 			if err == nil {
@@ -1137,6 +1138,10 @@ func (al *AgentLoop) runLLMIteration(
 			}
 
 			errMsg := strings.ToLower(err.Error())
+			var classifiedErr *providers.FailoverError
+			if singleCandidate {
+				classifiedErr = providers.ClassifyError(err, activeCandidates[0].Provider, activeModel)
+			}
 
 			// Check if this is a network/HTTP timeout — not a context window error.
 			isTimeoutError := errors.Is(err, context.DeadlineExceeded) ||
@@ -1144,9 +1149,10 @@ func (al *AgentLoop) runLLMIteration(
 				strings.Contains(errMsg, "client.timeout") ||
 				strings.Contains(errMsg, "timed out") ||
 				strings.Contains(errMsg, "timeout exceeded")
+			isTransientProviderError := classifiedErr != nil && classifiedErr.Reason == providers.FailoverTimeout
 
 			// Detect real context window / token limit errors, excluding network timeouts.
-			isContextError := !isTimeoutError && (strings.Contains(errMsg, "context_length_exceeded") ||
+			isContextError := !isTimeoutError && !isTransientProviderError && (strings.Contains(errMsg, "context_length_exceeded") ||
 				strings.Contains(errMsg, "context window") ||
 				strings.Contains(errMsg, "maximum context length") ||
 				strings.Contains(errMsg, "token limit") ||
@@ -1156,9 +1162,9 @@ func (al *AgentLoop) runLLMIteration(
 				strings.Contains(errMsg, "prompt is too long") ||
 				strings.Contains(errMsg, "request too large"))
 
-			if isTimeoutError && retry < maxRetries {
+			if (isTimeoutError || isTransientProviderError) && retry < maxRetries {
 				backoff := time.Duration(retry+1) * 5 * time.Second
-				logger.WarnCF("agent", "Timeout error, retrying after backoff", map[string]any{
+				logger.WarnCF("agent", "Transient provider error, retrying after backoff", map[string]any{
 					"error":   err.Error(),
 					"retry":   retry,
 					"backoff": backoff.String(),
diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go
index 8432ccac4..f23f4b22f 100644
--- a/pkg/agent/loop_test.go
+++ b/pkg/agent/loop_test.go
@@ -845,6 +845,52 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) {
 	}
 }
 
+func TestAgentLoop_TransientProviderRetry(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "agent-test-*")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	cfg := &config.Config{
+		Agents: config.AgentsConfig{
+			Defaults: config.AgentDefaults{
+				Workspace:         tmpDir,
+				Model:             "test-model",
+				MaxTokens:         4096,
+				MaxToolIterations: 10,
+			},
+		},
+	}
+
+	msgBus := bus.NewMessageBus()
+	provider := &failFirstMockProvider{
+		failures:    1,
+		failError:   fmt.Errorf("API error: status: 500 temporary upstream failure"),
+		successResp: "Recovered after transient provider failure",
+	}
+
+	al := NewAgentLoop(cfg, msgBus, provider)
+
+	response, err := al.ProcessDirectWithChannel(
+		context.Background(),
+		"retry this request",
+		"test-session-transient",
+		"cli",
+		"direct",
+	)
+	if err != nil {
+		t.Fatalf("Expected success after transient retry, got error: %v", err)
+	}
+
+	if response != "Recovered after transient provider failure" {
+		t.Fatalf("response = %q, want %q", response, "Recovered after transient provider failure")
+	}
+	if provider.currentCall != 2 {
+		t.Fatalf("expected 2 calls (1 fail + 1 success), got %d", provider.currentCall)
+	}
+}
+
 // TestProcessDirectWithChannel_TriggersMCPInitialization verifies that
 // ProcessDirectWithChannel triggers MCP initialization when MCP is enabled.
 // Note: Manager is only initialized when at least one MCP server is configured