mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-22 06:40:21 +08:00
fix(openai): fail over API-key passthrough 5xx
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user