fix(openai): fail over API-key passthrough 5xx

This commit is contained in:
王鹏
2026-07-15 04:21:38 +08:00
parent da85cc7e47
commit 91dc31ef39
3 changed files with 289 additions and 11 deletions
@@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/http/httptest"
"strings"
@@ -1272,6 +1273,29 @@ type openAIWSFailoverHandlerAccountRepoStub struct {
rateLimitedIDs []int64
}
type openAIHTTPPassthroughFailoverUpstream struct {
service.HTTPUpstream
mu sync.Mutex
accountIDs []int64
}
func (u *openAIHTTPPassthroughFailoverUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
u.mu.Lock()
u.accountIDs = append(u.accountIDs, accountID)
u.mu.Unlock()
return &http.Response{
StatusCode: http.StatusBadGateway,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"temporary upstream failure"}}`)),
}, nil
}
func (u *openAIHTTPPassthroughFailoverUpstream) calls() []int64 {
u.mu.Lock()
defer u.mu.Unlock()
return append([]int64(nil), u.accountIDs...)
}
func (s *openAIWSFailoverHandlerAccountRepoStub) ListSchedulableByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
out := make([]service.Account, 0, len(s.accounts))
for _, account := range s.accounts {
@@ -1344,6 +1368,96 @@ func (s *openAIWSUsageHandlerChannelRepoStub) GetGroupPlatforms(ctx context.Cont
return out, nil
}
func TestOpenAIResponses_APIKeyPassthroughPool5xxRetriesThenExhaustsMaxSwitches(t *testing.T) {
gin.SetMode(gin.TestMode)
groupID := int64(4203)
accounts := []service.Account{
{
ID: 9910, Name: "pool-api-key", Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 1,
Credentials: map[string]any{
"api_key": "sk-pool",
"base_url": "https://api.example.test",
"pool_mode": true,
"pool_mode_retry_count": float64(1),
"pool_mode_retry_status_codes": []any{float64(http.StatusBadGateway)},
},
Extra: map[string]any{"openai_passthrough": true},
},
{
ID: 9911, Name: "fallback-api-key", Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 2,
Credentials: map[string]any{
"api_key": "sk-fallback",
"base_url": "https://api.example.test",
},
Extra: map[string]any{"openai_passthrough": true},
},
}
cfg := &config.Config{RunMode: config.RunModeSimple}
cfg.Default.RateMultiplier = 1
cfg.Security.URLAllowlist.Enabled = false
cfg.Gateway.MaxAccountSwitches = 1
accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
upstream := &openAIHTTPPassthroughFailoverUpstream{}
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
t.Cleanup(billingCacheSvc.Stop)
gatewaySvc := service.NewOpenAIGatewayService(
accountRepo,
nil,
nil,
nil,
nil,
nil,
nil,
cfg,
nil,
nil,
service.NewBillingService(cfg, nil),
nil,
billingCacheSvc,
upstream,
&service.DeferredService{},
nil,
nil,
nil,
nil,
nil,
nil,
nil,
)
h := NewOpenAIGatewayHandler(
gatewaySvc,
service.NewConcurrencyService(nil),
billingCacheSvc,
service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
nil,
nil,
nil,
nil,
cfg,
)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5.2","input":"hello","stream":false}`))
c.Request.Header.Set("Content-Type", "application/json")
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
ID: 1803, GroupID: &groupID,
User: &service.User{ID: 1703, Status: service.StatusActive},
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
})
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1703, Concurrency: 0})
h.Responses(c)
require.Equal(t, []int64{9910, 9910, 9911}, upstream.calls())
require.Equal(t, http.StatusBadGateway, rec.Code)
require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
require.Equal(t, "Upstream service temporarily unavailable", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
}
func TestOpenAIResponsesWebSocket_FailoverOnUpstreamUsageLimitEvent(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -189,12 +189,13 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= 400 {
// 透传模式默认保持原样代理;但 429/529 属于网关必须兜底的
// 上游容量错误,应先触发多账号 failover 以维持基础 SLA。
if shouldFailoverOpenAIPassthroughResponse(resp.StatusCode) {
return nil, s.handleFailoverErrorResponsePassthrough(ctx, resp, c, account, body)
responseBody := s.readUpstreamErrorBody(resp)
// 透传模式默认保持原样代理;容量错误以及 API-key 上游的瞬时
// 5xx 应先触发多账号 failover,且此时尚未写入下游响应。
if shouldFailoverOpenAIPassthroughResponse(account, resp.StatusCode, responseBody) {
return nil, s.handleFailoverErrorResponsePassthrough(ctx, resp, c, account, body, responseBody)
}
return nil, s.handleErrorResponsePassthrough(ctx, resp, c, account, body)
return nil, s.handleErrorResponsePassthrough(ctx, resp, c, account, body, responseBody)
}
serviceTier := extractOpenAIServiceTierFromBody(body)
@@ -417,10 +418,24 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough(
return req, nil
}
func shouldFailoverOpenAIPassthroughResponse(statusCode int) bool {
func shouldFailoverOpenAIPassthroughResponse(account *Account, statusCode int, responseBody []byte) bool {
if isOpenAIContextWindowError("", responseBody) {
return false
}
switch statusCode {
case http.StatusTooManyRequests, 529:
return true
}
if account == nil || account.Type != AccountTypeAPIKey {
return false
}
switch statusCode {
case http.StatusInternalServerError,
http.StatusBadGateway,
http.StatusServiceUnavailable,
http.StatusGatewayTimeout,
520, 521, 522, 523, 524:
return true
default:
return false
}
@@ -432,8 +447,9 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough(
c *gin.Context,
account *Account,
requestBody []byte,
responseBody []byte,
) error {
body := s.readUpstreamErrorBody(resp)
body := responseBody
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
@@ -462,9 +478,10 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough(
UpstreamResponseBody: upstreamDetail,
})
return &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: body,
ResponseHeaders: resp.Header.Clone(),
StatusCode: resp.StatusCode,
ResponseBody: body,
ResponseHeaders: resp.Header.Clone(),
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
@@ -474,9 +491,10 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough(
c *gin.Context,
account *Account,
requestBody []byte,
responseBody []byte,
) error {
MarkResponseCommitted(c)
body := s.readUpstreamErrorBody(resp)
body := responseBody
// cyber_policy:透传账号本就把原始 body 回给客户端(下方 c.Data),此处仅打标记,
// 供 handler 事后写风控/邮件。cyber 是上游网络安全策略拦截,不冷却账号,
@@ -41,6 +41,16 @@ type passthroughErrReadCloser struct {
err error
}
type passthroughCloseTrackingReadCloser struct {
io.Reader
closed bool
}
func (r *passthroughCloseTrackingReadCloser) Close() error {
r.closed = true
return nil
}
func (r passthroughErrReadCloser) Read(_ []byte) (int, error) {
if r.err != nil {
return 0, r.err
@@ -1194,6 +1204,142 @@ func TestOpenAIGatewayService_OpenAIPassthrough_RetryableStatusesTriggerFailover
}
}
func TestOpenAIGatewayService_APIKeyPassthrough_Transient5xxTriggersFailover(t *testing.T) {
gin.SetMode(gin.TestMode)
requestBody := []byte(`{"model":"gpt-5.2","stream":false,"input":"hello"}`)
for _, statusCode := range []int{
http.StatusInternalServerError,
http.StatusBadGateway,
http.StatusServiceUnavailable,
http.StatusGatewayTimeout,
520, 521, 522, 523, 524,
} {
t.Run(fmt.Sprintf("status_%d", statusCode), func(t *testing.T) {
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil))
c.Request.Header.Set("User-Agent", "codex_cli_rs/0.1.0")
upstreamBody := fmt.Sprintf(`{"error":{"message":"temporary upstream failure","status":%d}}`, statusCode)
body := &passthroughCloseTrackingReadCloser{Reader: strings.NewReader(upstreamBody)}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: statusCode,
Header: http.Header{
"Content-Type": []string{"application/json"},
"X-Request-Id": []string{"rid-api-key-5xx"},
},
Body: body,
}}
svc := &OpenAIGatewayService{
cfg: &config.Config{Gateway: config.GatewayConfig{ForceCodexCLI: false}},
httpUpstream: upstream,
}
account := &Account{
ID: 124,
Name: "api-key-transient-5xx",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://api.example.test",
},
Extra: map[string]any{"openai_passthrough": true},
Status: StatusActive,
Schedulable: true,
}
result, err := svc.Forward(context.Background(), c, account, requestBody)
require.Nil(t, result, "failed attempts must not report usage or success metadata")
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, statusCode, failoverErr.StatusCode)
require.JSONEq(t, upstreamBody, string(failoverErr.ResponseBody))
require.Equal(t, "rid-api-key-5xx", failoverErr.ResponseHeaders.Get("x-request-id"))
require.False(t, c.Writer.Written(), "failover must happen before downstream output is committed")
require.True(t, body.closed, "the failed upstream response body must be closed")
require.Equal(t, requestBody, upstream.lastBody, "the request body remains available for the outer account retry")
value, ok := c.Get(OpsUpstreamErrorsKey)
require.True(t, ok)
events, ok := value.([]*OpsUpstreamErrorEvent)
require.True(t, ok)
require.NotEmpty(t, events)
require.Equal(t, "failover", events[len(events)-1].Kind)
require.Equal(t, account.ID, events[len(events)-1].AccountID)
})
}
}
func TestOpenAIGatewayService_APIKeyPassthrough_ContextWindow502DoesNotFailover(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil))
const upstreamBody = `{"error":{"message":"Your input exceeds the context window of this model. Please adjust your input and try again.","type":"upstream_error"}}`
body := &passthroughCloseTrackingReadCloser{Reader: strings.NewReader(upstreamBody)}
svc := &OpenAIGatewayService{
cfg: &config.Config{Gateway: config.GatewayConfig{ForceCodexCLI: false}},
httpUpstream: &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusBadGateway,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: body,
}},
}
account := &Account{
ID: 127, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1,
Credentials: map[string]any{"api_key": "sk-test", "base_url": "https://api.example.test"},
Extra: map[string]any{"openai_passthrough": true}, Status: StatusActive, Schedulable: true,
}
result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.2","input":"hello"}`))
require.Nil(t, result)
require.Error(t, err)
var failoverErr *UpstreamFailoverError
require.False(t, errors.As(err, &failoverErr), "context-window errors are deterministic request failures")
require.True(t, c.Writer.Written())
require.Equal(t, http.StatusBadGateway, rec.Code)
require.Contains(t, rec.Body.String(), "exceeds the context window")
require.True(t, body.closed)
}
func TestOpenAIGatewayService_APIKeyPassthrough_PoolModeConfigured5xxRetriesSameAccount(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil))
svc := &OpenAIGatewayService{
cfg: &config.Config{Gateway: config.GatewayConfig{ForceCodexCLI: false}},
httpUpstream: &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusBadGateway,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"temporary upstream failure"}}`)),
}},
}
account := &Account{
ID: 128, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://api.example.test",
"pool_mode": true,
"pool_mode_retry_status_codes": []any{float64(http.StatusBadGateway)},
},
Extra: map[string]any{"openai_passthrough": true}, Status: StatusActive, Schedulable: true,
}
_, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.2","input":"hello"}`))
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.True(t, failoverErr.RetryableOnSameAccount)
require.False(t, c.Writer.Written())
}
func TestOpenAIGatewayService_OpenAIPassthrough_CompactNetworkErrorsTriggerFailover(t *testing.T) {
gin.SetMode(gin.TestMode)