AutoZero/llm/retry.go

package llm

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"time"
)

// RetryConfig 配置重试策略。
type RetryConfig struct {
	MaxRetries     int           // 最大重试次数（默认 3）
	InitialBackoff time.Duration // 初始退避时间（默认 1s）
	MaxBackoff     time.Duration // 最大退避时间（默认 30s）
	RetryOn        []int         // 需要重试的 HTTP 状态码（默认 429, 500-509）
}

// DefaultRetryConfig 默认重试配置。
var DefaultRetryConfig = RetryConfig{
	MaxRetries:     3,
	InitialBackoff: 1 * time.Second,
	MaxBackoff:     30 * time.Second,
	RetryOn:        []int{429, 500, 502, 503, 504},
}

// globalRetryConfig 全局重试配置。
var globalRetryConfig = DefaultRetryConfig

// SetRetryConfig 设置全局重试配置。
func SetRetryConfig(cfg RetryConfig) {
	globalRetryConfig = cfg
}

// GetRetryConfig 获取当前重试配置。
func GetRetryConfig() RetryConfig {
	return globalRetryConfig
}

// doWithRetry 发送请求并对瞬态错误做指数退避重试（尊重 Retry-After）。
// 只覆盖「响应开始前」的阶段：HTTP 状态检查在 SSE 流开始之前，一旦 200 并
// 开始读流，错误不再重试（增量已经发出去了）。
func doWithRetry(ctx context.Context, client *http.Client, req *http.Request, body []byte) (*http.Response, error) {
	return doWithRetryConfig(ctx, client, req, body, globalRetryConfig)
}

// doWithRetryConfig 使用指定配置执行重试。
func doWithRetryConfig(ctx context.Context, client *http.Client, req *http.Request, body []byte, cfg RetryConfig) (*http.Response, error) {
	if cfg.MaxRetries <= 0 {
		cfg.MaxRetries = 1
	}
	if cfg.InitialBackoff <= 0 {
		cfg.InitialBackoff = 1 * time.Second
	}
	if cfg.MaxBackoff <= 0 {
		cfg.MaxBackoff = 30 * time.Second
	}
	if len(cfg.RetryOn) == 0 {
		cfg.RetryOn = DefaultRetryConfig.RetryOn
	}

	// 创建状态码查找表
	retryOnMap := make(map[int]bool, len(cfg.RetryOn))
	for _, code := range cfg.RetryOn {
		retryOnMap[code] = true
	}

	var lastErr error
	var lastRetryAfter time.Duration

	for attempt := 0; attempt < cfg.MaxRetries; attempt++ {
		if attempt > 0 {
			wait := time.Duration(1<<(attempt-1)) * cfg.InitialBackoff
			if wait > cfg.MaxBackoff {
				wait = cfg.MaxBackoff
			}
			if lastRetryAfter > 0 {
				wait = lastRetryAfter
			}
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(wait):
			}
		}

		r := req.Clone(ctx)
		r.Body = io.NopCloser(bytes.NewReader(body))
		r.ContentLength = int64(len(body))
		resp, err := client.Do(r)

		if err != nil {
			lastErr = err
			lastRetryAfter = 0
			continue // 网络错误：可重试
		}

		// 检查是否需要重试
		if resp.StatusCode == http.StatusTooManyRequests || retryOnMap[resp.StatusCode] {
			lastRetryAfter = parseRetryAfter(resp.Header.Get("Retry-After"))
			_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4096))
			resp.Body.Close()
			lastErr = fmt.Errorf("HTTP %d", resp.StatusCode)
			continue
		}

		return resp, nil // 2xx 与 4xx 一律交还调用方（4xx 不可重试）
	}

	return nil, fmt.Errorf("after %d attempts: %w", cfg.MaxRetries, lastErr)
}

// isRetryable 检查错误是否可重试。
func isRetryable(err error) bool {
	if err == nil {
		return false
	}

	// 网络错误可重试
	if _, ok := err.(interface{ Timeout() bool }); ok {
		return true
	}

	// HTTP 错误
	errStr := err.Error()
	return contains(errStr, "timeout") ||
		contains(errStr, "connection refused") ||
		contains(errStr, "connection reset") ||
		contains(errStr, "EOF")
}

// contains 检查字符串是否包含子串。
func contains(s, substr string) bool {
	return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsSubstr(s, substr))
}

func containsSubstr(s, substr string) bool {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}

func parseRetryAfter(v string) time.Duration {
	if v == "" {
		return 0
	}
	if secs, err := strconv.Atoi(v); err == nil && secs > 0 && secs <= 60 {
		return time.Duration(secs) * time.Second
	}
	return 0
}