mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-15 23:22:59 +08:00
fix(openai): retry SSE rate limits as HTTP 429
OpenAI Responses may return code=rate_limit_exceeded in a response.failed SSE event while the HTTP status remains 200. Classify these failures as 429 so configured pool-mode retries and account failover are applied. No account-level rate-limit state is written on this path: the 200-stream response headers carry normal quota snapshots, and retry semantics stay owned by the failover engine. Ignore compact keepalive bytes when determining whether semantic output has started, preserving safe retries before real output.
This commit is contained in:
@@ -1812,6 +1812,40 @@ func (u *openAIHTTPPassthroughFailoverUpstream) calls() []int64 {
|
||||
return append([]int64(nil), u.accountIDs...)
|
||||
}
|
||||
|
||||
type openAIHTTPPassthroughSSERateLimitUpstream struct {
|
||||
service.HTTPUpstream
|
||||
mu sync.Mutex
|
||||
accountIDs []int64
|
||||
}
|
||||
|
||||
func (u *openAIHTTPPassthroughSSERateLimitUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
|
||||
u.mu.Lock()
|
||||
u.accountIDs = append(u.accountIDs, accountID)
|
||||
u.mu.Unlock()
|
||||
body := strings.Join([]string{
|
||||
"event: response.created",
|
||||
`data: {"type":"response.created","response":{"id":"resp_rate_limited"}}`,
|
||||
"",
|
||||
"event: response.failed",
|
||||
`data: {"type":"response.failed","response":{"id":"resp_rate_limited","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
|
||||
"",
|
||||
}, "\n")
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"text/event-stream"},
|
||||
"Retry-After": []string{"1"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (u *openAIHTTPPassthroughSSERateLimitUpstream) 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 {
|
||||
@@ -2012,6 +2046,88 @@ func TestOpenAIResponses_APIKeyPassthroughPool5xxRetriesThenExhaustsMaxSwitches(
|
||||
require.Equal(t, "Upstream service temporarily unavailable", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
|
||||
}
|
||||
|
||||
func TestOpenAIResponses_APIKeyPassthroughSSERateLimitUsesConfiguredPoolRetry(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
groupID := int64(4204)
|
||||
accounts := []service.Account{
|
||||
{
|
||||
ID: 9912, Name: "pool-sse-rate-limit", 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.StatusTooManyRequests)},
|
||||
},
|
||||
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 := &openAIHTTPPassthroughSSERateLimitUpstream{}
|
||||
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.6-sol","input":"hello","stream":true}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 1804, GroupID: &groupID,
|
||||
User: &service.User{ID: 1704, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
|
||||
})
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1704, Concurrency: 0})
|
||||
|
||||
h.Responses(c)
|
||||
|
||||
require.Equal(t, []int64{9912, 9912}, upstream.calls())
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
require.Equal(t, "1", rec.Header().Get("Retry-After"))
|
||||
require.Equal(t, "rate_limit_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
|
||||
require.Equal(t, "Upstream rate limit exceeded, please retry later", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_FailoverOnUpstreamUsageLimitEvent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user