diff --git a/internal/provider/errors.go b/internal/provider/errors.go new file mode 100644 index 00000000..083fb39f --- /dev/null +++ b/internal/provider/errors.go @@ -0,0 +1,95 @@ +package provider + +import ( + "fmt" + "net/http" +) + +type ProviderErrorCode string + +const ( + ErrorCodeAuthFailed ProviderErrorCode = "auth_failed" // 认证失败（401） + ErrorCodeForbidden ProviderErrorCode = "forbidden" // 权限不足（403） + ErrorCodeNotFound ProviderErrorCode = "not_found" // 资源不存在（404） + ErrorCodeClient ProviderErrorCode = "client_error" // 客户端请求错误（4xx，排除上述分类） + ErrorCodeRateLimit ProviderErrorCode = "rate_limited" // 限流（429） + ErrorCodeServer ProviderErrorCode = "server_error" // 服务端错误（5xx） + ErrorCodeTimeout ProviderErrorCode = "timeout" // 超时 + ErrorCodeNetwork ProviderErrorCode = "network_error" // 网络错误（连接拒绝、DNS 失败等） + ErrorCodeUnknown ProviderErrorCode = "unknown" // 未知错误 +) + +// ProviderError 是 provider 层的领域错误类型。 +type ProviderError struct { + StatusCode int // HTTP 状态码，0 表示非 HTTP 错误（如网络超时） + Code ProviderErrorCode // 语义化的错误分类 + Message string // 可读错误信息 + Retryable bool // 是否建议重试 +} + +func (e *ProviderError) Error() string { + if e.StatusCode > 0 { + return fmt.Sprintf("provider error (status=%d, code=%s): %s", e.StatusCode, e.Code, e.Message) + } + return fmt.Sprintf("provider error (code=%s): %s", e.Code, e.Message) +} + +// IsRetryableStatus 根据通用 HTTP 语义判断状态码是否可重试。 +// 429（限流）和 5xx（服务端错误）视为可重试。 +func IsRetryableStatus(statusCode int) bool { + return statusCode == http.StatusTooManyRequests || + (statusCode >= http.StatusInternalServerError && statusCode < 600) +} + +// classifyStatus 将 HTTP 状态码映射为 ProviderErrorCode。 +func classifyStatus(statusCode int) ProviderErrorCode { + switch statusCode { + case http.StatusUnauthorized: + return ErrorCodeAuthFailed + case http.StatusForbidden: + return ErrorCodeForbidden + case http.StatusNotFound: + return ErrorCodeNotFound + case http.StatusTooManyRequests: + return ErrorCodeRateLimit + default: + if statusCode >= 400 && statusCode < 500 { + return ErrorCodeClient + } + if statusCode >= 500 && statusCode < 600 { + return ErrorCodeServer + } + } + return ErrorCodeUnknown +} + +// NewProviderErrorFromStatus 根据 HTTP 状态码和消息构造 ProviderError。 +func NewProviderErrorFromStatus(statusCode int, message string) *ProviderError { + code := classifyStatus(statusCode) + return &ProviderError{ + StatusCode: statusCode, + Code: code, + Message: message, + Retryable: IsRetryableStatus(statusCode), + } +} + +// NewNetworkProviderError 构造网络错误类型的 ProviderError（如连接超时、拒绝等）。 +func NewNetworkProviderError(message string) *ProviderError { + return &ProviderError{ + StatusCode: 0, + Code: ErrorCodeNetwork, + Message: message, + Retryable: true, // 网络抖动默认可重试 + } +} + +// NewTimeoutProviderError 构造超时类型的 ProviderError。 +func NewTimeoutProviderError(message string) *ProviderError { + return &ProviderError{ + StatusCode: 0, + Code: ErrorCodeTimeout, + Message: message, + Retryable: true, // 超时默认可重试 + } +} diff --git a/internal/provider/errors_test.go b/internal/provider/errors_test.go new file mode 100644 index 00000000..e6e9ac4d --- /dev/null +++ b/internal/provider/errors_test.go @@ -0,0 +1,172 @@ +package provider + +import ( + "errors" + "fmt" + "net/http" + "strings" + "testing" +) + +func TestProviderError_ErrorFormat(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err *ProviderError + want string + }{ + { + name: "with status code", + err: &ProviderError{StatusCode: 401, Code: ErrorCodeAuthFailed, Message: "invalid key"}, + want: "status=401", + }, + { + name: "without status code", + err: &ProviderError{StatusCode: 0, Code: ErrorCodeNetwork, Message: "connection refused"}, + want: "code=network_error", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := tt.err.Error(); !strings.Contains(got, tt.want) { + t.Fatalf("Error() = %q, want containing %q", got, tt.want) + } + }) + } +} + +func TestIsRetryableStatus(t *testing.T) { + t.Parallel() + + tests := []struct { + status int + expected bool + }{ + {http.StatusBadRequest, false}, + {http.StatusUnauthorized, false}, + {http.StatusForbidden, false}, + {http.StatusNotFound, false}, + {http.StatusUnprocessableEntity, false}, + {http.StatusTooManyRequests, true}, + {http.StatusInternalServerError, true}, + {http.StatusBadGateway, true}, + {http.StatusServiceUnavailable, true}, + {http.StatusGatewayTimeout, true}, + {599, true}, + {200, false}, + {0, false}, + } + + for _, tt := range tests { + tt := tt + t.Run(http.StatusText(tt.status), func(t *testing.T) { + t.Parallel() + if got := IsRetryableStatus(tt.status); got != tt.expected { + t.Fatalf("IsRetryableStatus(%d) = %v, want %v", tt.status, got, tt.expected) + } + }) + } +} + +func TestClassifyStatus(t *testing.T) { + t.Parallel() + + tests := []struct { + status int + want ProviderErrorCode + }{ + {http.StatusUnauthorized, ErrorCodeAuthFailed}, + {http.StatusForbidden, ErrorCodeForbidden}, + {http.StatusNotFound, ErrorCodeNotFound}, + {http.StatusBadRequest, ErrorCodeClient}, + {http.StatusUnprocessableEntity, ErrorCodeClient}, + {http.StatusTooManyRequests, ErrorCodeRateLimit}, + {http.StatusInternalServerError, ErrorCodeServer}, + {http.StatusBadGateway, ErrorCodeServer}, + {http.StatusServiceUnavailable, ErrorCodeServer}, + {200, ErrorCodeUnknown}, + {0, ErrorCodeUnknown}, + } + + for _, tt := range tests { + tt := tt + t.Run(http.StatusText(tt.status), func(t *testing.T) { + t.Parallel() + if got := classifyStatus(tt.status); got != tt.want { + t.Fatalf("classifyStatus(%d) = %q, want %q", tt.status, got, tt.want) + } + }) + } +} + +func TestNewProviderErrorFromStatus(t *testing.T) { + t.Parallel() + + err := NewProviderErrorFromStatus(429, "rate limited") + if err.StatusCode != 429 { + t.Fatalf("expected status 429, got %d", err.StatusCode) + } + if err.Code != ErrorCodeRateLimit { + t.Fatalf("expected code %q, got %q", ErrorCodeRateLimit, err.Code) + } + if !err.Retryable { + t.Fatalf("expected 429 to be retryable") + } + + err = NewProviderErrorFromStatus(401, "unauthorized") + if err.Retryable { + t.Fatalf("expected 401 to not be retryable") + } +} + +func TestNewNetworkProviderError(t *testing.T) { + t.Parallel() + + err := NewNetworkProviderError("connection refused") + if err.StatusCode != 0 { + t.Fatalf("expected status 0, got %d", err.StatusCode) + } + if err.Code != ErrorCodeNetwork { + t.Fatalf("expected code %q, got %q", ErrorCodeNetwork, err.Code) + } + if !err.Retryable { + t.Fatalf("expected network error to be retryable") + } +} + +func TestNewTimeoutProviderError(t *testing.T) { + t.Parallel() + + err := NewTimeoutProviderError("deadline exceeded") + if err.StatusCode != 0 { + t.Fatalf("expected status 0, got %d", err.StatusCode) + } + if err.Code != ErrorCodeTimeout { + t.Fatalf("expected code %q, got %q", ErrorCodeTimeout, err.Code) + } + if !err.Retryable { + t.Fatalf("expected timeout to be retryable") + } +} + +func TestProviderError_As(t *testing.T) { + t.Parallel() + + inner := &ProviderError{StatusCode: 500, Code: ErrorCodeServer, Message: "internal", Retryable: true} + wrapped := fmt.Errorf("wrapped: %w", inner) + + var target *ProviderError + if !errors.As(wrapped, &target) { + t.Fatalf("expected errors.As to match *ProviderError") + } + if target.StatusCode != 500 { + t.Fatalf("expected status 500, got %d", target.StatusCode) + } + if !target.Retryable { + t.Fatalf("expected retryable") + } +} diff --git a/internal/provider/openai/openai.go b/internal/provider/openai/openai.go index a00e8204..700b50a3 100644 --- a/internal/provider/openai/openai.go +++ b/internal/provider/openai/openai.go @@ -16,6 +16,7 @@ import ( "neo-code/internal/config" domain "neo-code/internal/provider" + "neo-code/internal/provider/transport" ) type Provider struct { @@ -23,6 +24,19 @@ type Provider struct { client *http.Client } +type buildOptions struct { + transport http.RoundTripper +} + +type BuildOption func(*buildOptions) + +// WithTransport 注入自定义 HTTP Transport（如 RetryTransport）。 +func WithTransport(rt http.RoundTripper) BuildOption { + return func(o *buildOptions) { + o.transport = rt + } +} + const DriverName = "openai" // Driver 返回 OpenAI 协议驱动的定义。 @@ -30,12 +44,14 @@ func Driver() domain.DriverDefinition { return domain.DriverDefinition{ Name: DriverName, Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (domain.Provider, error) { - return New(cfg) + baseTransport := http.DefaultTransport + retryTransport := transport.NewRetryTransport(baseTransport, transport.DefaultRetryConfig()) + return New(cfg, WithTransport(retryTransport)) }, } } -func New(cfg config.ResolvedProviderConfig) (*Provider, error) { +func New(cfg config.ResolvedProviderConfig, opts ...BuildOption) (*Provider, error) { if err := cfg.Validate(); err != nil { return nil, fmt.Errorf("openai provider: %w", err) } @@ -43,10 +59,18 @@ func New(cfg config.ResolvedProviderConfig) (*Provider, error) { return nil, errors.New("openai provider: api key is empty") } + o := &buildOptions{ + transport: http.DefaultTransport, + } + for _, apply := range opts { + apply(o) + } + return &Provider{ cfg: cfg, client: &http.Client{ - Timeout: 90 * time.Second, + Timeout: 90 * time.Second, + Transport: o.transport, }, }, nil } @@ -233,20 +257,21 @@ func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events cha func (p *Provider) parseError(resp *http.Response) error { data, readErr := io.ReadAll(resp.Body) if readErr != nil { - return fmt.Errorf("openai provider: read error response: %w", readErr) + return domain.NewProviderErrorFromStatus(resp.StatusCode, + fmt.Sprintf("openai provider: read error response: %v", readErr)) } var parsed openAIErrorResponse if err := json.Unmarshal(data, &parsed); err == nil && strings.TrimSpace(parsed.Error.Message) != "" { - return errors.New(parsed.Error.Message) + return domain.NewProviderErrorFromStatus(resp.StatusCode, parsed.Error.Message) } bodyText := strings.TrimSpace(string(data)) if bodyText == "" { - return errors.New(resp.Status) + return domain.NewProviderErrorFromStatus(resp.StatusCode, resp.Status) } - return fmt.Errorf("openai provider: %s", bodyText) + return domain.NewProviderErrorFromStatus(resp.StatusCode, bodyText) } func toOpenAIMessage(message domain.Message) openAIMessage { @@ -436,5 +461,6 @@ type openAIUsage struct { type openAIErrorResponse struct { Error struct { Message string `json:"message"` + Code string `json:"code,omitempty"` } `json:"error"` } diff --git a/internal/provider/transport/retry.go b/internal/provider/transport/retry.go new file mode 100644 index 00000000..dca6a23b --- /dev/null +++ b/internal/provider/transport/retry.go @@ -0,0 +1,183 @@ +package transport + +import ( + "errors" + "io" + "log" + "math/rand/v2" + "net/http" + "strings" + "time" + + "neo-code/internal/provider" +) + +// RetryConfig 控制 RetryTransport 的重试行为。 +type RetryConfig struct { + // MaxRetries 最大重试次数（不含首次请求）。0 表示不重试。 + MaxRetries int + // WaitBase 初始等待时间，默认 500ms。 + WaitBase time.Duration + // MaxWait 最大等待时间，默认 5s。 + MaxWait time.Duration + // RetryableFunc 自定义判断函数。为 nil 时使用 DefaultRetryableFunc。 + RetryableFunc func(error) bool +} + +// DefaultRetryConfig 返回默认重试配置。 +func DefaultRetryConfig() RetryConfig { + return RetryConfig{ + MaxRetries: 2, + WaitBase: 500 * time.Millisecond, + MaxWait: 5 * time.Second, + RetryableFunc: DefaultRetryableFunc, + } +} + +// DefaultRetryableFunc 判断 error 是否可重试。 +// - *ProviderError 且 Retryable == true +// - 网络错误（超时、连接拒绝等） +func DefaultRetryableFunc(err error) bool { + if err == nil { + return false + } + var pErr *provider.ProviderError + if ok := errors.As(err, &pErr); ok { + return pErr.Retryable + } + // 网络层错误（超时、连接拒绝、DNS 失败等）视为可重试。 + if isNetworkError(err) { + return true + } + return false +} + +// isNetworkError 判断是否为网络层错误（不包含 TLS 错误，避免重试证书问题）。 +func isNetworkError(err error) bool { + if err == nil { + return false + } + msg := err.Error() + // 常见网络错误关键字 + networkIndicators := []string{ + "connection refused", + "connection reset", + "connection timed out", + "timeout", + "no such host", + "i/o timeout", + "dial tcp", + } + for _, indicator := range networkIndicators { + if strings.Contains(msg, indicator) { + return true + } + } + return false +} + +// RetryTransport 是一个 http.RoundTripper 中间件，在可重试错误上自动重试。 +// 只在首次请求（尚未读取 resp.Body）阶段生效，SSE 流式传输中途断连不在其重试范围内。 +type RetryTransport struct { + base http.RoundTripper + config RetryConfig +} + +// NewRetryTransport 创建 RetryTransport。 +// base 为底层 transport，通常为 http.DefaultTransport。 +func NewRetryTransport(base http.RoundTripper, cfg RetryConfig) *RetryTransport { + if base == nil { + base = http.DefaultTransport + } + if cfg.RetryableFunc == nil { + cfg.RetryableFunc = DefaultRetryableFunc + } + return &RetryTransport{ + base: base, + config: cfg, + } +} + +// RoundTrip 执行 HTTP 请求，遇到可重试错误时自动重试。 +func (rt *RetryTransport) RoundTrip(req *http.Request) (*http.Response, error) { + var lastErr error + var lastResp *http.Response + + for attempt := 0; attempt <= rt.config.MaxRetries; attempt++ { + // 每次重试前重新读取 request body（如果有 seeker）。 + if attempt > 0 { + if err := rt.drainAndClose(lastResp); err != nil { + log.Printf("retry transport: drain previous response: %v", err) + } + if req.Body != nil { + if seeker, ok := req.Body.(io.Seeker); ok { + if _, err := seeker.Seek(0, io.SeekStart); err != nil { + return nil, err + } + } + } + wait := rt.backoff(attempt) + select { + case <-req.Context().Done(): + return nil, req.Context().Err() + case <-time.After(wait): + } + } + + resp, err := rt.base.RoundTrip(req) + if err == nil { + return resp, nil + } + + lastErr = err + lastResp = nil + + // 检查是否可重试。 + if !rt.config.RetryableFunc(err) { + return nil, err + } + + // 检查 context 是否已取消。 + if req.Context().Err() != nil { + return nil, req.Context().Err() + } + } + + return nil, lastErr +} + +// drainAndClose 读空并关闭前一次失败的 resp.Body，防止连接泄漏。 +func (rt *RetryTransport) drainAndClose(resp *http.Response) error { + if resp == nil || resp.Body == nil { + return nil + } + _, _ = io.Copy(io.Discard, resp.Body) + return resp.Body.Close() +} + +// backoff 计算指数退避 + 随机抖动的等待时间。 +// attempt 从 1 开始（首次重试）。 +func (rt *RetryTransport) backoff(attempt int) time.Duration { + base := rt.config.WaitBase + if base <= 0 { + base = 500 * time.Millisecond + } + + // 指数退避 + wait := base << (attempt - 1) + + // 随机抖动：[0.5, 1.5) * wait + jitter := float64(wait) * (0.5 + rand.Float64()) + wait = time.Duration(jitter) + + // 上限 + maxWait := rt.config.MaxWait + if maxWait <= 0 { + maxWait = 5 * time.Second + } + if wait > maxWait { + wait = maxWait + } + + return wait +} diff --git a/internal/provider/transport/retry_test.go b/internal/provider/transport/retry_test.go new file mode 100644 index 00000000..ac5cc7d7 --- /dev/null +++ b/internal/provider/transport/retry_test.go @@ -0,0 +1,524 @@ +package transport + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net/http" + "strings" + "testing" + "time" + + "neo-code/internal/provider" +) + +// --- isNetworkError --- + +func TestIsNetworkError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"connection refused", errors.New("dial tcp 127.0.0.1:8080: connection refused"), true}, + {"connection reset", errors.New("connection reset by peer"), true}, + {"connection timed out", errors.New("connection timed out"), true}, + {"i/o timeout", errors.New("dial tcp: i/o timeout"), true}, + {"no such host", errors.New("dial tcp: lookup api.openai.com: no such host"), true}, + {"dial tcp", errors.New("dial tcp 127.0.0.1:443: operation was refused"), true}, + {"timeout in message", errors.New("request timeout"), true}, + {"tls cert error", errors.New("x509: certificate signed by unknown authority"), false}, + {"generic error", errors.New("something went wrong"), false}, + {"empty message", errors.New(""), false}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := isNetworkError(tt.err); got != tt.want { + t.Fatalf("isNetworkError() = %v, want %v", got, tt.want) + } + }) + } +} + +// --- DefaultRetryableFunc --- + +func TestDefaultRetryableFunc(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + want bool + }{ + {"nil error", nil, false}, + {"retryable ProviderError", &provider.ProviderError{StatusCode: 429, Code: provider.ErrorCodeRateLimit, Retryable: true}, true}, + {"non-retryable ProviderError", &provider.ProviderError{StatusCode: 401, Code: provider.ErrorCodeAuthFailed, Retryable: false}, false}, + {"wrapped retryable ProviderError", fmt.Errorf("wrapped: %w", &provider.ProviderError{StatusCode: 500, Retryable: true}), true}, + {"wrapped non-retryable ProviderError", fmt.Errorf("wrapped: %w", &provider.ProviderError{StatusCode: 400, Retryable: false}), false}, + {"network error", errors.New("dial tcp 127.0.0.1:8080: connection refused"), true}, + {"non-network non-provider error", errors.New("bad request"), false}, + {"empty error message", errors.New(""), false}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := DefaultRetryableFunc(tt.err); got != tt.want { + t.Fatalf("DefaultRetryableFunc() = %v, want %v", got, tt.want) + } + }) + } +} + +// --- DefaultRetryConfig --- + +func TestDefaultRetryConfig(t *testing.T) { + t.Parallel() + + cfg := DefaultRetryConfig() + if cfg.MaxRetries != 2 { + t.Fatalf("MaxRetries = %d, want 2", cfg.MaxRetries) + } + if cfg.WaitBase != 500*time.Millisecond { + t.Fatalf("WaitBase = %v, want 500ms", cfg.WaitBase) + } + if cfg.MaxWait != 5*time.Second { + t.Fatalf("MaxWait = %v, want 5s", cfg.MaxWait) + } + if cfg.RetryableFunc == nil { + t.Fatal("RetryableFunc should not be nil") + } +} + +// --- NewRetryTransport --- + +func TestNewRetryTransport_NilBase(t *testing.T) { + t.Parallel() + + rt := NewRetryTransport(nil, RetryConfig{}) + if rt.base != http.DefaultTransport { + t.Fatal("nil base should fall back to http.DefaultTransport") + } +} + +func TestNewRetryTransport_NilRetryableFunc(t *testing.T) { + t.Parallel() + + cfg := RetryConfig{MaxRetries: 1} + rt := NewRetryTransport(http.DefaultTransport, cfg) + if rt.config.RetryableFunc == nil { + t.Fatal("nil RetryableFunc should fall back to DefaultRetryableFunc") + } +} + +// --- mockRoundTripper --- + +type mockRoundTripper struct { + //fn returns (response, error). Called sequentially. + fn []func() (*http.Response, error) + // callCount tracks how many times RoundTrip was invoked. + callCount int +} + +func (m *mockRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + idx := m.callCount + m.callCount++ + if idx >= len(m.fn) { + return nil, errors.New("mock: unexpected call") + } + return m.fn[idx]() +} + +func newSuccessResponse(statusCode int, body string) *http.Response { + return &http.Response{ + StatusCode: statusCode, + Body: io.NopCloser(strings.NewReader(body)), + Header: make(http.Header), + } +} + +// --- RoundTrip --- + +func TestRoundTrip_FirstSuccess(t *testing.T) { + t.Parallel() + + wantResp := newSuccessResponse(http.StatusOK, "ok") + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { return wantResp, nil }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{MaxRetries: 2}) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + resp, err := rt.RoundTrip(req) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + if m.callCount != 1 { + t.Fatalf("callCount = %d, want 1", m.callCount) + } +} + +func TestRoundTrip_RetryThenSuccess(t *testing.T) { + t.Parallel() + + wantResp := newSuccessResponse(http.StatusOK, "ok") + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { return nil, errors.New("dial tcp 127.0.0.1:8080: connection refused") }, + func() (*http.Response, error) { return wantResp, nil }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 2, + WaitBase: 1 * time.Millisecond, // 快速完成测试 + MaxWait: 10 * time.Millisecond, + }) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + resp, err := rt.RoundTrip(req) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + if m.callCount != 2 { + t.Fatalf("callCount = %d, want 2", m.callCount) + } +} + +func TestRoundTrip_AllRetriesFail(t *testing.T) { + t.Parallel() + + networkErr := errors.New("connection refused") + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { return nil, networkErr }, + func() (*http.Response, error) { return nil, networkErr }, + func() (*http.Response, error) { return nil, networkErr }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 2, + WaitBase: 1 * time.Millisecond, + MaxWait: 10 * time.Millisecond, + }) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + _, err := rt.RoundTrip(req) + + if err == nil { + t.Fatal("expected error after all retries exhausted") + } + if !errors.Is(err, networkErr) { + t.Fatalf("expected networkErr, got %v", err) + } + if m.callCount != 3 { + t.Fatalf("callCount = %d, want 3 (1 initial + 2 retries)", m.callCount) + } +} + +func TestRoundTrip_NonRetryableError(t *testing.T) { + t.Parallel() + + authErr := provider.NewProviderErrorFromStatus(http.StatusUnauthorized, "invalid api key") + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { return nil, authErr }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{MaxRetries: 2}) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + _, err := rt.RoundTrip(req) + + if err == nil { + t.Fatal("expected error for non-retryable") + } + var pErr *provider.ProviderError + if !errors.As(err, &pErr) { + t.Fatal("expected *ProviderError") + } + if pErr.Code != provider.ErrorCodeAuthFailed { + t.Fatalf("code = %s, want auth_failed", pErr.Code) + } + if m.callCount != 1 { + t.Fatalf("callCount = %d, want 1 (no retry)", m.callCount) + } +} + +func TestRoundTrip_ContextCanceled(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // 立即取消 + + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { return nil, errors.New("connection refused") }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 2, + WaitBase: 1 * time.Millisecond, + MaxWait: 10 * time.Millisecond, + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.com", nil) + _, err := rt.RoundTrip(req) + + if err == nil { + t.Fatal("expected error when context is canceled") + } + // 首次请求会发出，然后 context.Err() 检查应在重试前终止 + if m.callCount < 1 { + t.Fatalf("expected at least 1 call, got %d", m.callCount) + } +} + +func TestRoundTrip_ContextCanceledDuringBackoff(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + + callCount := 0 + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { + callCount++ + return nil, errors.New("connection refused") + }, + func() (*http.Response, error) { + // 在第二次调用之前 context 已取消，这里不应被调用 + callCount++ + return nil, errors.New("should not reach") + }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 3, + WaitBase: 500 * time.Millisecond, // 足够长，便于在 backoff 期间取消 + MaxWait: 1 * time.Second, + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.com", nil) + + // 首次请求失败后进入 backoff 期间取消 + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + _, err := rt.RoundTrip(req) + if err == nil { + t.Fatal("expected error") + } + if callCount != 1 { + t.Fatalf("callCount = %d, want 1 (first call, then canceled during backoff)", callCount) + } +} + +func TestRoundTrip_ZeroMaxRetries(t *testing.T) { + t.Parallel() + + networkErr := errors.New("connection refused") + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { return nil, networkErr }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{MaxRetries: 0}) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + _, err := rt.RoundTrip(req) + + if err == nil { + t.Fatal("expected error") + } + if m.callCount != 1 { + t.Fatalf("callCount = %d, want 1", m.callCount) + } +} + +func TestRoundTrip_SeekerBodyRewound(t *testing.T) { + t.Parallel() + + wantResp := newSuccessResponse(http.StatusOK, "ok") + bodyContent := []byte("request body") + seekedPositions := make([]int64, 0) + + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { + return nil, errors.New("connection refused") + }, + func() (*http.Response, error) { + return wantResp, nil + }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 1, + WaitBase: 1 * time.Millisecond, + MaxWait: 10 * time.Millisecond, + }) + + // seekReadCloser 同时实现 io.ReadCloser + io.Seeker， + // 因为 io.NopCloser 会丢失 Seek 接口。 + body := &seekReadCloser{ + Reader: *bytes.NewReader(bodyContent), + onSeek: func(offset int64, whence int) { + if whence == io.SeekStart && offset == 0 { + seekedPositions = append(seekedPositions, offset) + } + }, + } + req, _ := http.NewRequest(http.MethodPost, "http://example.com", body) + req.Header.Set("Content-Type", "application/json") + + _, err := rt.RoundTrip(req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(seekedPositions) != 1 { + t.Fatalf("expected 1 seek(0) call, got %d: %v", len(seekedPositions), seekedPositions) + } +} + +// seekReadCloser 组合 bytes.Reader（已有 Seek）+ Close，供 http.Request.Body 使用。 +type seekReadCloser struct { + bytes.Reader + onSeek func(offset int64, whence int) +} + +func (s *seekReadCloser) Close() error { return nil } + +func (s *seekReadCloser) Seek(offset int64, whence int) (int64, error) { + if s.onSeek != nil { + s.onSeek(offset, whence) + } + return s.Reader.Seek(offset, whence) +} + +// --- drainAndClose --- + +func TestDrainAndClose_NilResp(t *testing.T) { + t.Parallel() + + rt := &RetryTransport{} + if err := rt.drainAndClose(nil); err != nil { + t.Fatalf("expected nil error, got %v", err) + } +} + +func TestDrainAndClose_NilBody(t *testing.T) { + t.Parallel() + + rt := &RetryTransport{} + if err := rt.drainAndClose(&http.Response{Body: nil}); err != nil { + t.Fatalf("expected nil error, got %v", err) + } +} + +func TestDrainAndClose_ReadsBody(t *testing.T) { + t.Parallel() + + content := "some response body to drain" + rt := &RetryTransport{} + resp := &http.Response{ + Body: io.NopCloser(strings.NewReader(content)), + } + if err := rt.drainAndClose(resp); err != nil { + t.Fatalf("expected nil error, got %v", err) + } +} + +// --- backoff --- + +func TestBackoff_ExponentialGrowth(t *testing.T) { + t.Parallel() + + cfg := RetryConfig{ + WaitBase: 100 * time.Millisecond, + MaxWait: 10 * time.Second, // 足够大，不触发上限截断 + } + rt := NewRetryTransport(http.DefaultTransport, cfg) + + b1 := rt.backoff(1) + _ = rt.backoff(2) + b3 := rt.backoff(3) + + // 指数退避：b1 ≈ 100ms, b2 ≈ 200ms, b3 ≈ 400ms（含抖动） + // 验证大致递增关系（抖动可能让相邻值接近，但 b3 应明显大于 b1） + if b3 <= b1 { + t.Fatalf("backoff(3)=%v should be > backoff(1)=%v", b3, b1) + } +} + +func TestBackoff_RespectsMaxWait(t *testing.T) { + t.Parallel() + + cfg := RetryConfig{ + WaitBase: 500 * time.Millisecond, + MaxWait: 800 * time.Millisecond, // 很小的上限 + } + rt := NewRetryTransport(http.DefaultTransport, cfg) + + // attempt=10: 500ms << 9 = 256000ms，远超 800ms 上限 + b := rt.backoff(10) + if b > 800*time.Millisecond { + t.Fatalf("backoff(10)=%v should be <= 800ms MaxWait", b) + } +} + +func TestBackoff_ZeroWaitBase(t *testing.T) { + t.Parallel() + + cfg := RetryConfig{ + WaitBase: 0, // 应回退到 500ms + MaxWait: 5 * time.Second, + } + rt := NewRetryTransport(http.DefaultTransport, cfg) + + b := rt.backoff(1) + // 500ms * [0.5, 1.5) => [250ms, 750ms) + if b < 200*time.Millisecond || b > 800*time.Millisecond { + t.Fatalf("backoff(1) with zero base = %v, expected ~[250ms, 750ms]", b) + } +} + +func TestBackoff_JitterRange(t *testing.T) { + t.Parallel() + + cfg := RetryConfig{ + WaitBase: 200 * time.Millisecond, + MaxWait: 10 * time.Second, + } + rt := NewRetryTransport(http.DefaultTransport, cfg) + + // 多次采样验证抖动范围 + for i := 0; i < 20; i++ { + b := rt.backoff(1) + // 200ms * [0.5, 1.5) => [100ms, 300ms) + if b < 100*time.Millisecond || b >= 300*time.Millisecond { + t.Fatalf("backoff(1)=%v outside expected jitter range [100ms, 300ms)", b) + } + } +} diff --git a/internal/runtime/events.go b/internal/runtime/events.go index 4cd162c5..0e2ce006 100644 --- a/internal/runtime/events.go +++ b/internal/runtime/events.go @@ -33,4 +33,7 @@ const ( // EventToolCallThinking is emitted when the model decides to call a tool, // before the tool execution begins. TUI can show a transitional indicator. EventToolCallThinking EventType = "tool_call_thinking" + // EventProviderRetry is emitted when runtime retries a provider call due to + // a retryable error (e.g. 429, 5xx). Payload is a human-readable message. + EventProviderRetry EventType = "provider_retry" ) diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 48d685b0..6e0da509 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -3,6 +3,9 @@ package runtime import ( "context" "errors" + "fmt" + "log" + "math/rand/v2" "strings" "sync" "time" @@ -13,6 +16,18 @@ import ( "neo-code/internal/tools" ) +const maxContextTurns = 10 + +const ( + // defaultProviderRetryMax 是 runtime 层对单次 provider.Chat() 的最大重试次数（不含首次调用）。 + // 与 RetryTransport 的 HTTP 层重试互补：Transport 耗尽后 runtime 仍可重试整个 Chat 调用。 + defaultProviderRetryMax = 2 + // providerRetryBaseWait 是 runtime 层重试的初始等待时间。 + providerRetryBaseWait = 1 * time.Second + // providerRetryMaxWait 是 runtime 层重试的最大等待时间。 + providerRetryMaxWait = 5 * time.Second +) + type Runtime interface { Run(ctx context.Context, input UserInput) error CancelActiveRun() bool @@ -113,18 +128,6 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { return err } - resolvedProvider, err := s.configManager.ResolvedSelectedProvider() - if err != nil { - s.emit(ctx, EventError, input.RunID, session.ID, err.Error()) - return err - } - - modelProvider, err := s.providerFactory.Build(ctx, resolvedProvider) - if err != nil { - s.emit(ctx, EventError, input.RunID, session.ID, err.Error()) - return err - } - builtContext, err := s.contextBuilder.Build(ctx, agentcontext.BuildInput{ Messages: session.Messages, Workdir: cfg.Workdir, @@ -133,18 +136,13 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { return s.handleRunError(ctx, input.RunID, session.ID, err) } - streamEvents := make(chan provider.StreamEvent, 32) - streamDone := make(chan struct{}) - go s.forwardProviderEvents(ctx, input.RunID, session.ID, streamEvents, streamDone) + resp, err := s.callProviderWithRetry(ctx, input.RunID, session.ID, provider.ChatRequest{ - resp, err := modelProvider.Chat(ctx, provider.ChatRequest{ Model: cfg.CurrentModel, SystemPrompt: builtContext.SystemPrompt, Messages: builtContext.Messages, Tools: s.toolRegistry.GetSpecs(), - }, streamEvents) - close(streamEvents) - <-streamDone + }) if err != nil { return s.handleRunError(ctx, input.RunID, session.ID, err) } @@ -341,10 +339,103 @@ func (s *Service) handleRunError(ctx context.Context, runID string, sessionID st return context.Canceled } + // 提取 ProviderError 信息用于日志增强。 + var pErr *provider.ProviderError + if errors.As(err, &pErr) { + log.Printf("runtime: provider error (status=%d, code=%s, retryable=%v): %s", + pErr.StatusCode, pErr.Code, pErr.Retryable, pErr.Message) + } + s.emit(ctx, EventError, runID, sessionID, err.Error()) return err } +// isRetryableProviderError 判断 error 是否为可重试的 ProviderError。 +func isRetryableProviderError(err error) bool { + var pErr *provider.ProviderError + if !errors.As(err, &pErr) { + return false + } + return pErr.Retryable +} + +// callProviderWithRetry 在可重试的 ProviderError 上自动重试 provider.Chat() 调用。 +// 每次重试都会重新构建 provider 实例和流式事件转发管道。 +// 非可重试错误、context 取消、重试耗尽时直接返回错误。 +func (s *Service) callProviderWithRetry( + ctx context.Context, + runID string, + sessionID string, + req provider.ChatRequest, +) (provider.ChatResponse, error) { + var lastErr error + + for retryAttempt := 0; retryAttempt <= defaultProviderRetryMax; retryAttempt++ { + if retryAttempt > 0 { + wait := providerRetryBackoff(retryAttempt) + s.emit(ctx, EventProviderRetry, runID, sessionID, + fmt.Sprintf("retrying provider call (attempt %d/%d, wait=%.1fs)...", + retryAttempt, defaultProviderRetryMax, wait.Seconds())) + + select { + case <-ctx.Done(): + return provider.ChatResponse{}, ctx.Err() + case <-time.After(wait): + } + } + + resolvedProvider, err := s.configManager.ResolvedSelectedProvider() + if err != nil { + return provider.ChatResponse{}, err + } + + modelProvider, err := s.providerFactory.Build(ctx, resolvedProvider) + if err != nil { + return provider.ChatResponse{}, err + } + + streamEvents := make(chan provider.StreamEvent, 32) + streamDone := make(chan struct{}) + go s.forwardProviderEvents(ctx, runID, sessionID, streamEvents, streamDone) + + resp, err := modelProvider.Chat(ctx, req, streamEvents) + close(streamEvents) + <-streamDone + + if err == nil { + return resp, nil + } + + lastErr = err + + // 非可重试错误或 context 已取消，立即返回。 + if !isRetryableProviderError(err) { + return provider.ChatResponse{}, err + } + if ctx.Err() != nil { + return provider.ChatResponse{}, ctx.Err() + } + } + + return provider.ChatResponse{}, lastErr +} + +// providerRetryBackoff 计算指数退避 + 随机抖动的等待时间。 +// attempt 从 1 开始（首次重试）。 +func providerRetryBackoff(attempt int) time.Duration { + wait := providerRetryBaseWait << (attempt - 1) + + // 随机抖动：[0.5, 1.5) * wait + jitter := float64(wait) * (0.5 + rand.Float64()) + wait = time.Duration(jitter) + + if wait > providerRetryMaxWait { + wait = providerRetryMaxWait + } + + return wait +} + func (s *Service) isRunCanceled(err error) bool { return errors.Is(err, context.Canceled) } diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index c29b4510..4f3eb869 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -534,6 +534,96 @@ func TestServiceRunErrorPaths(t *testing.T) { } }, }, + { + name: "retryable provider error triggers runtime retry then succeeds", + input: UserInput{RunID: "run-retry-success", Content: "hello"}, + provider: func() *scriptedProvider { + callIdx := 0 + return &scriptedProvider{ + name: "retry-then-success", + chatFn: func(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { + callIdx++ + if callIdx == 1 { + return provider.ChatResponse{}, &provider.ProviderError{ + StatusCode: 500, + Code: provider.ErrorCodeServer, + Message: "internal server error", + Retryable: true, + } + } + return provider.ChatResponse{ + Message: provider.Message{ + Role: "assistant", + Content: "recovered", + }, + FinishReason: "stop", + }, nil + }, + } + }(), + expectEvents: []EventType{EventUserMessage, EventProviderRetry, EventAgentDone}, + assert: func(t *testing.T, store *memoryStore, scripted *scriptedProvider, tool *stubTool) { + t.Helper() + if scripted.callCount < 2 { + t.Fatalf("expected at least 2 provider calls (initial + retry), got %d", scripted.callCount) + } + session := onlySession(t, store) + if len(session.Messages) != 2 { + t.Fatalf("expected user + assistant messages, got %d", len(session.Messages)) + } + if session.Messages[1].Content != "recovered" { + t.Fatalf("expected assistant content %q, got %q", "recovered", session.Messages[1].Content) + } + }, + }, + { + name: "non-retryable provider error does not trigger runtime retry", + input: UserInput{RunID: "run-no-retry", Content: "hello"}, + provider: &scriptedProvider{ + name: "auth-error-no-retry", + chatFn: func(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { + return provider.ChatResponse{}, &provider.ProviderError{ + StatusCode: 401, + Code: provider.ErrorCodeAuthFailed, + Message: "invalid api key", + Retryable: false, + } + }, + }, + expectErr: "invalid api key", + expectEvents: []EventType{EventUserMessage, EventError}, + assert: func(t *testing.T, store *memoryStore, scripted *scriptedProvider, tool *stubTool) { + t.Helper() + if scripted.callCount != 1 { + t.Fatalf("expected exactly 1 provider call (no retry for 401), got %d", scripted.callCount) + } + }, + }, + { + name: "runtime retry exhausted emits error", + input: UserInput{RunID: "run-retry-exhausted", Content: "hello"}, + provider: &scriptedProvider{ + name: "always-500", + chatFn: func(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { + return provider.ChatResponse{}, &provider.ProviderError{ + StatusCode: 500, + Code: provider.ErrorCodeServer, + Message: "internal server error", + Retryable: true, + } + }, + }, + expectErr: "internal server error", + expectEvents: []EventType{EventUserMessage, EventProviderRetry, EventError}, + assert: func(t *testing.T, store *memoryStore, scripted *scriptedProvider, tool *stubTool) { + t.Helper() + // 1 initial + 2 retries = 3 calls + if scripted.callCount != defaultProviderRetryMax+1 { + t.Fatalf("expected %d provider calls (1 initial + %d retries), got %d", + defaultProviderRetryMax+1, defaultProviderRetryMax, scripted.callCount) + } + }, + }, } for _, tt := range tests { @@ -1093,3 +1183,46 @@ func cloneBuildInput(input agentcontext.BuildInput) agentcontext.BuildInput { func containsError(err error, target string) bool { return err != nil && strings.Contains(err.Error(), target) } + +func TestIsRetryableProviderError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + want bool + }{ + {"nil error", nil, false}, + {"retryable provider error", &provider.ProviderError{Retryable: true}, true}, + {"non-retryable provider error", &provider.ProviderError{Retryable: false}, false}, + {"plain error", errors.New("something failed"), false}, + {"wrapped retryable", fmt.Errorf("wrapped: %w", &provider.ProviderError{Retryable: true}), true}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := isRetryableProviderError(tt.err); got != tt.want { + t.Fatalf("isRetryableProviderError() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestProviderRetryBackoff(t *testing.T) { + t.Parallel() + + first := providerRetryBackoff(1) + second := providerRetryBackoff(2) + + if second <= first { + t.Fatalf("expected backoff to increase: first=%v second=%v", first, second) + } + + // 验证不超过 MaxWait + large := providerRetryBackoff(20) + if large > providerRetryMaxWait { + t.Fatalf("expected backoff <= MaxWait, got %v > %v", large, providerRetryMaxWait) + } +}