fix: improve retry

Carlos Alexandro Becker created 1 month ago

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

Change summary

internal/llm/provider/anthropic.go | 13 ++++++++++++-
internal/llm/provider/gemini.go    | 16 +++++++++++++++-
internal/llm/provider/openai.go    | 13 ++++++++++++-
3 files changed, 39 insertions(+), 3 deletions(-)

Detailed changes

internal/llm/provider/anthropic.go

@@ -509,7 +509,18 @@ func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, err
 	isOverloaded := strings.Contains(apiErr.Error(), "overloaded") || strings.Contains(apiErr.Error(), "rate limit exceeded")
 	// 529 (unofficial): The service is overloaded
 	if apiErr.StatusCode != http.StatusTooManyRequests && apiErr.StatusCode != 529 && !isOverloaded {
-		return false, 0, err
+		prev := a.providerOptions.apiKey
+		// in case the key comes from a script, we try to re-evaluate it.
+		a.providerOptions.apiKey, err = config.Get().Resolve(a.providerOptions.config.APIKey)
+		if err != nil {
+			return false, 0, fmt.Errorf("failed to resolve API key: %w", err)
+		}
+		// if it didn't change, do not retry.
+		if prev == a.providerOptions.apiKey {
+			return false, 0, err
+		}
+		a.client = createAnthropicClient(a.providerOptions, a.tp)
+		return true, 0, nil
 	}
 
 	retryMs := 0

internal/llm/provider/gemini.go

@@ -436,7 +436,21 @@ func (g *geminiClient) shouldRetry(attempts int, err error) (bool, int64, error)
 
 	// Check for token expiration (401 Unauthorized)
 	if contains(errMsg, "unauthorized", "invalid api key", "api key expired") {
-		return false, 0, err
+		prev := g.providerOptions.apiKey
+		// in case the key comes from a script, we try to re-evaluate it.
+		g.providerOptions.apiKey, err = config.Get().Resolve(g.providerOptions.config.APIKey)
+		if err != nil {
+			return false, 0, fmt.Errorf("failed to resolve API key: %w", err)
+		}
+		// if it didn't change, do not retry.
+		if prev == g.providerOptions.apiKey {
+			return false, 0, err
+		}
+		g.client, err = createGeminiClient(g.providerOptions)
+		if err != nil {
+			return false, 0, fmt.Errorf("failed to create Gemini client after API key refresh: %w", err)
+		}
+		return true, 0, nil
 	}
 
 	// Check for common rate limit error messages

internal/llm/provider/openai.go

@@ -515,7 +515,18 @@ func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error)
 	if errors.As(err, &apiErr) {
 		// Check for token expiration (401 Unauthorized)
 		if apiErr.StatusCode == http.StatusUnauthorized {
-			return false, 0, err
+			prev := o.providerOptions.apiKey
+			// in case the key comes from a script, we try to re-evaluate it.
+			o.providerOptions.apiKey, err = config.Get().Resolve(o.providerOptions.config.APIKey)
+			if err != nil {
+				return false, 0, fmt.Errorf("failed to resolve API key: %w", err)
+			}
+			// if it didn't change, do not retry.
+			if prev == o.providerOptions.apiKey {
+				return false, 0, err
+			}
+			o.client = createOpenAIClient(o.providerOptions)
+			return true, 0, nil
 		}
 
 		if apiErr.StatusCode != http.StatusTooManyRequests && apiErr.StatusCode != http.StatusInternalServerError {