fix(openai): 429 insuffice not retry (#546)

nguyen created 1 month ago

Change summary

internal/llm/provider/openai.go      | 12 +++
internal/llm/provider/openai_test.go | 76 ++++++++++++++++++++++++++++++
2 files changed, 86 insertions(+), 2 deletions(-)

Detailed changes

internal/llm/provider/openai.go

@@ -529,11 +529,19 @@ func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error)
 			return true, 0, nil
 		}
 
-		if apiErr.StatusCode != http.StatusTooManyRequests && apiErr.StatusCode != http.StatusInternalServerError {
+		if apiErr.StatusCode == http.StatusTooManyRequests {
+			// Check if this is an insufficient quota error (permanent)
+			if apiErr.Type == "insufficient_quota" || apiErr.Code == "insufficient_quota" {
+				return false, 0, fmt.Errorf("OpenAI quota exceeded: %s. Please check your plan and billing details", apiErr.Message)
+			}
+			// Other 429 errors (rate limiting) can be retried
+		} else if apiErr.StatusCode != http.StatusInternalServerError {
 			return false, 0, err
 		}
 
-		retryAfterValues = apiErr.Response.Header.Values("Retry-After")
+		if apiErr.Response != nil {
+			retryAfterValues = apiErr.Response.Header.Values("Retry-After")
+		}
 	}
 
 	if apiErr != nil {

internal/llm/provider/openai_test.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -88,3 +89,78 @@ func TestOpenAIClientStreamChoices(t *testing.T) {
 		}
 	}
 }
+
+func TestOpenAIClient429InsufficientQuotaError(t *testing.T) {
+	client := &openaiClient{
+		providerOptions: providerClientOptions{
+			modelType:     config.SelectedModelTypeLarge,
+			apiKey:        "test-key",
+			systemMessage: "test",
+			config: config.ProviderConfig{
+				ID:     "test-openai",
+				APIKey: "test-key",
+			},
+			model: func(config.SelectedModelType) catwalk.Model {
+				return catwalk.Model{
+					ID:   "test-model",
+					Name: "test-model",
+				}
+			},
+		},
+	}
+
+	// Test insufficient_quota error should not retry
+	apiErr := &openai.Error{
+		StatusCode: 429,
+		Message:    "You exceeded your current quota, please check your plan and billing details. For more information on this error, read the docs: https://platform.openai.com/docs/guides/error-codes/api-errors.",
+		Type:       "insufficient_quota",
+		Code:       "insufficient_quota",
+	}
+
+	retry, _, err := client.shouldRetry(1, apiErr)
+	if retry {
+		t.Error("Expected shouldRetry to return false for insufficient_quota error, but got true")
+	}
+	if err == nil {
+		t.Error("Expected shouldRetry to return an error for insufficient_quota, but got nil")
+	}
+	if err != nil && !strings.Contains(err.Error(), "quota") {
+		t.Errorf("Expected error message to mention quota, got: %v", err)
+	}
+}
+
+func TestOpenAIClient429RateLimitError(t *testing.T) {
+	client := &openaiClient{
+		providerOptions: providerClientOptions{
+			modelType:     config.SelectedModelTypeLarge,
+			apiKey:        "test-key",
+			systemMessage: "test",
+			config: config.ProviderConfig{
+				ID:     "test-openai",
+				APIKey: "test-key",
+			},
+			model: func(config.SelectedModelType) catwalk.Model {
+				return catwalk.Model{
+					ID:   "test-model",
+					Name: "test-model",
+				}
+			},
+		},
+	}
+
+	// Test regular rate limit error should retry
+	apiErr := &openai.Error{
+		StatusCode: 429,
+		Message:    "Rate limit reached for requests",
+		Type:       "rate_limit_exceeded",
+		Code:       "rate_limit_exceeded",
+	}
+
+	retry, _, err := client.shouldRetry(1, apiErr)
+	if !retry {
+		t.Error("Expected shouldRetry to return true for rate_limit_exceeded error, but got false")
+	}
+	if err != nil {
+		t.Errorf("Expected shouldRetry to return nil error for rate_limit_exceeded, but got: %v", err)
+	}
+}