feat(retry): Add configurable status codes and callback mechanism

claude[bot] · edenreich · web-flow · commit 1d3cefa4a826 · 2025-08-22T14:11:47.000Z
- Add RetryableStatusCodes field to RetryConfig for custom status codes
- Add OnRetry callback for retry event logging and monitoring
- Maintain backward compatibility with default behavior
- Update isRetryableStatusCode to use configurable codes
- Add comprehensive test coverage for new functionality

Co-authored-by: Eden Reich &lt;edenreich@users.noreply.github.com&gt;
diff --git a/sdk.go b/sdk.go
@@ -71,7 +71,18 @@ func isRetryableError(err error) bool {
 }
 
 // isRetryableStatusCode determines if an HTTP status code should trigger a retry
-func isRetryableStatusCode(statusCode int) bool {
+func isRetryableStatusCode(statusCode int, config *RetryConfig) bool {
+	// Use custom status codes if provided
+	if len(config.RetryableStatusCodes) > 0 {
+		for _, code := range config.RetryableStatusCodes {
+			if statusCode == code {
+				return true
+			}
+		}
+		return false
+	}
+
+	// Use default status codes
 	switch statusCode {
 	case
 		http.StatusRequestTimeout,      // 408
@@ -179,6 +190,11 @@ func (c *clientImpl) executeWithRetry(ctx context.Context, request func() (*rest
 		if attempt > 0 {
 			delay := calculateBackoff(attempt, c.retryConfig)
 
+			// Call OnRetry callback if provided
+			if c.retryConfig.OnRetry != nil {
+				c.retryConfig.OnRetry(attempt, lastErr, delay)
+			}
+
 			select {
 			case <-ctx.Done():
 				return nil, ctx.Err()
@@ -189,13 +205,13 @@ func (c *clientImpl) executeWithRetry(ctx context.Context, request func() (*rest
 		resp, lastErr = request()
 
 		if lastErr == nil {
-			if !resp.IsError() || !isRetryableStatusCode(resp.StatusCode()) {
+			if !resp.IsError() || !isRetryableStatusCode(resp.StatusCode(), c.retryConfig) {
 				return resp, nil
 			}
 			lastErr = fmt.Errorf("HTTP %d", resp.StatusCode())
 		}
 
-		if !isRetryableError(lastErr) && (resp == nil || !isRetryableStatusCode(resp.StatusCode())) {
+		if !isRetryableError(lastErr) && (resp == nil || !isRetryableStatusCode(resp.StatusCode(), c.retryConfig)) {
 			break
 		}
 
diff --git a/sdk_test.go b/sdk_test.go
@@ -1454,6 +1454,8 @@ func TestCalculateBackoff(t *testing.T) {
 }
 
 func TestIsRetryableStatusCode(t *testing.T) {
+	defaultConfig := getDefaultRetryConfig()
+	
 	tests := []struct {
 		statusCode int
 		expected   bool
@@ -1473,7 +1475,7 @@ func TestIsRetryableStatusCode(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(fmt.Sprintf("status_%d", tt.statusCode), func(t *testing.T) {
-			result := isRetryableStatusCode(tt.statusCode)
+			result := isRetryableStatusCode(tt.statusCode, defaultConfig)
 			assert.Equal(t, tt.expected, result)
 		})
 	}
@@ -1510,3 +1512,262 @@ func TestRetryWithContext(t *testing.T) {
 	assert.GreaterOrEqual(t, callCount, 1)
 	assert.LessOrEqual(t, callCount, 5)
 }
+
+func TestCustomRetryableStatusCodes(t *testing.T) {
+	tests := []struct {
+		name              string
+		customStatusCodes []int
+		statusCode        int
+		expected          bool
+	}{
+		{
+			name:              "custom codes include 418",
+			customStatusCodes: []int{418, 422},
+			statusCode:        418,
+			expected:          true,
+		},
+		{
+			name:              "custom codes exclude 500",
+			customStatusCodes: []int{418, 422},
+			statusCode:        500,
+			expected:          false,
+		},
+		{
+			name:              "custom codes include 422",
+			customStatusCodes: []int{418, 422},
+			statusCode:        422,
+			expected:          true,
+		},
+		{
+			name:              "custom codes exclude 200",
+			customStatusCodes: []int{418, 422},
+			statusCode:        200,
+			expected:          false,
+		},
+		{
+			name:              "empty custom codes use defaults",
+			customStatusCodes: []int{},
+			statusCode:        500,
+			expected:          true,
+		},
+		{
+			name:              "nil custom codes use defaults",
+			customStatusCodes: nil,
+			statusCode:        500,
+			expected:          true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := &RetryConfig{
+				RetryableStatusCodes: tt.customStatusCodes,
+			}
+			result := isRetryableStatusCode(tt.statusCode, config)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func TestRetryCallback(t *testing.T) {
+	var callbackCalls []struct {
+		attempt int
+		err     error
+		delay   time.Duration
+	}
+
+	retryConfig := &RetryConfig{
+		Enabled:           true,
+		MaxAttempts:       3,
+		InitialBackoffSec: 1,
+		MaxBackoffSec:     10,
+		BackoffMultiplier: 2,
+		OnRetry: func(attempt int, err error, delay time.Duration) {
+			callbackCalls = append(callbackCalls, struct {
+				attempt int
+				err     error
+				delay   time.Duration
+			}{attempt, err, delay})
+		},
+	}
+
+	callCount := 0
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if callCount < 2 {
+			w.WriteHeader(http.StatusInternalServerError)
+			err := json.NewEncoder(w).Encode(Error{Error: stringPtr("Server error")})
+			assert.NoError(t, err)
+		} else {
+			w.WriteHeader(http.StatusOK)
+			response := ListModelsResponse{Object: "list", Data: []Model{}}
+			err := json.NewEncoder(w).Encode(response)
+			assert.NoError(t, err)
+		}
+		callCount++
+	}))
+	defer server.Close()
+
+	baseURL := server.URL + "/v1"
+	client := NewClient(&ClientOptions{
+		BaseURL:     baseURL,
+		RetryConfig: retryConfig,
+	})
+
+	ctx := context.Background()
+	_, err := client.ListModels(ctx)
+
+	assert.NoError(t, err)
+	assert.Equal(t, 3, callCount)
+	assert.Len(t, callbackCalls, 2) // Two retries after initial failure
+
+	// Verify callback calls
+	assert.Equal(t, 1, callbackCalls[0].attempt)
+	assert.Contains(t, callbackCalls[0].err.Error(), "HTTP 500")
+	assert.Equal(t, 1*time.Second, callbackCalls[0].delay)
+
+	assert.Equal(t, 2, callbackCalls[1].attempt)
+	assert.Contains(t, callbackCalls[1].err.Error(), "HTTP 500")
+	assert.Equal(t, 2*time.Second, callbackCalls[1].delay)
+}
+
+func TestRetryWithCustomStatusCodesAndCallback(t *testing.T) {
+	var callbackCalls []struct {
+		attempt int
+		err     error
+		delay   time.Duration
+	}
+
+	retryConfig := &RetryConfig{
+		Enabled:              true,
+		MaxAttempts:          3,
+		InitialBackoffSec:    1,
+		MaxBackoffSec:        10,
+		BackoffMultiplier:    2,
+		RetryableStatusCodes: []int{418, 503}, // Custom codes: I'm a teapot and Service Unavailable
+		OnRetry: func(attempt int, err error, delay time.Duration) {
+			callbackCalls = append(callbackCalls, struct {
+				attempt int
+				err     error
+				delay   time.Duration
+			}{attempt, err, delay})
+		},
+	}
+
+	tests := []struct {
+		name           string
+		statusCodes    []int
+		expectRetries  bool
+		expectedCalls  int
+		callbackCounts int
+	}{
+		{
+			name:           "retry on custom 418 status code",
+			statusCodes:    []int{418, 418, 200},
+			expectRetries:  true,
+			expectedCalls:  3,
+			callbackCounts: 2,
+		},
+		{
+			name:           "retry on custom 503 status code",
+			statusCodes:    []int{503, 200},
+			expectRetries:  true,
+			expectedCalls:  2,
+			callbackCounts: 1,
+		},
+		{
+			name:           "no retry on non-custom 500 status code",
+			statusCodes:    []int{500},
+			expectRetries:  false,
+			expectedCalls:  1,
+			callbackCounts: 0,
+		},
+		{
+			name:           "no retry on 400 status code",
+			statusCodes:    []int{400},
+			expectRetries:  false,
+			expectedCalls:  1,
+			callbackCounts: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			callbackCalls = nil // Reset callback calls
+			callCount := 0
+
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				if callCount < len(tt.statusCodes) {
+					w.WriteHeader(tt.statusCodes[callCount])
+					if tt.statusCodes[callCount] == 200 {
+						response := ListModelsResponse{Object: "list", Data: []Model{}}
+						err := json.NewEncoder(w).Encode(response)
+						assert.NoError(t, err)
+					} else {
+						err := json.NewEncoder(w).Encode(Error{Error: stringPtr("Server error")})
+						assert.NoError(t, err)
+					}
+				}
+				callCount++
+			}))
+			defer server.Close()
+
+			baseURL := server.URL + "/v1"
+			client := NewClient(&ClientOptions{
+				BaseURL:     baseURL,
+				RetryConfig: retryConfig,
+			})
+
+			ctx := context.Background()
+			_, err := client.ListModels(ctx)
+
+			assert.Equal(t, tt.expectedCalls, callCount)
+			assert.Len(t, callbackCalls, tt.callbackCounts)
+
+			if tt.expectRetries && tt.statusCodes[len(tt.statusCodes)-1] == 200 {
+				assert.NoError(t, err)
+			} else if !tt.expectRetries || tt.statusCodes[len(tt.statusCodes)-1] != 200 {
+				assert.Error(t, err)
+			}
+		})
+	}
+}
+
+func TestRetryConfigWithNilCallback(t *testing.T) {
+	retryConfig := &RetryConfig{
+		Enabled:           true,
+		MaxAttempts:       2,
+		InitialBackoffSec: 1,
+		MaxBackoffSec:     10,
+		BackoffMultiplier: 2,
+		OnRetry:           nil, // Explicitly nil callback
+	}
+
+	callCount := 0
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if callCount == 0 {
+			w.WriteHeader(http.StatusInternalServerError)
+			err := json.NewEncoder(w).Encode(Error{Error: stringPtr("Server error")})
+			assert.NoError(t, err)
+		} else {
+			w.WriteHeader(http.StatusOK)
+			response := ListModelsResponse{Object: "list", Data: []Model{}}
+			err := json.NewEncoder(w).Encode(response)
+			assert.NoError(t, err)
+		}
+		callCount++
+	}))
+	defer server.Close()
+
+	baseURL := server.URL + "/v1"
+	client := NewClient(&ClientOptions{
+		BaseURL:     baseURL,
+		RetryConfig: retryConfig,
+	})
+
+	ctx := context.Background()
+	_, err := client.ListModels(ctx)
+
+	// Should work without callback
+	assert.NoError(t, err)
+	assert.Equal(t, 2, callCount)
+}
diff --git a/types.go b/types.go
@@ -76,6 +76,12 @@ type RetryConfig struct {
 	MaxBackoffSec int
 	// BackoffMultiplier is the multiplier for exponential backoff
 	BackoffMultiplier int
+	// RetryableStatusCodes is a custom list of HTTP status codes that should trigger a retry.
+	// If nil or empty, uses default status codes (408, 429, 500, 502, 503, 504)
+	RetryableStatusCodes []int
+	// OnRetry is called before each retry attempt with attempt number, error, and delay.
+	// The attempt number starts from 1 for the first retry (after initial request fails)
+	OnRetry func(attempt int, err error, delay time.Duration)
 }
 
 // MiddlewareOptions represents options for controlling middleware behavior
