mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 14:38:27 +08:00
fix(openai): scope Spark quota 429 to model
This commit is contained in:
@@ -12,6 +12,7 @@ const (
|
||||
modelRateLimitsKey = "model_rate_limits"
|
||||
antigravityGeminiModelRateLimitKey = "antigravity:gemini"
|
||||
openAIImageGenerationRateLimitKey = "openai:image_generation"
|
||||
openAICodexSparkRateLimitReason = "openai_codex_spark_rate_limit"
|
||||
// anthropicFableRateLimitKey 是 Anthropic 7d_oi(Fable 专属 7d 窗口)限流的
|
||||
// 家族级 scope:命中后所有 Fable 变体(含 [1m] 等后缀)都不再调度到该账号。
|
||||
anthropicFableRateLimitKey = "claude-fable-5"
|
||||
|
||||
@@ -162,6 +162,10 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont
|
||||
s.rateLimitService.HandleTempUnschedulable(stateCtx, account, statusCode, responseBody, canonicalModel[0]) {
|
||||
return true
|
||||
}
|
||||
if statusCode == http.StatusTooManyRequests && s.rateLimitService != nil && len(canonicalModel) > 0 &&
|
||||
s.rateLimitService.HandleOpenAICodexSparkRateLimit(stateCtx, account, canonicalModel[0], statusCode, headers, responseBody) {
|
||||
return false
|
||||
}
|
||||
if statusCode == http.StatusTooManyRequests {
|
||||
s.markOpenAIOAuth429RateLimited(stateCtx, account, headers, responseBody)
|
||||
}
|
||||
|
||||
@@ -15,8 +15,11 @@ import (
|
||||
|
||||
type oauth429RateLimitRepo struct {
|
||||
mockAccountRepoForGemini
|
||||
setRateLimitedCalls int
|
||||
lastRateLimitedUntil time.Time
|
||||
setRateLimitedCalls int
|
||||
lastRateLimitedUntil time.Time
|
||||
setModelRateLimitCalls int
|
||||
lastModelRateLimitKey string
|
||||
lastModelRateLimitedUntil time.Time
|
||||
}
|
||||
|
||||
func (r *oauth429RateLimitRepo) SetRateLimited(_ context.Context, _ int64, until time.Time) error {
|
||||
@@ -25,6 +28,13 @@ func (r *oauth429RateLimitRepo) SetRateLimited(_ context.Context, _ int64, until
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *oauth429RateLimitRepo) SetModelRateLimit(_ context.Context, _ int64, scope string, until time.Time, _ ...string) error {
|
||||
r.setModelRateLimitCalls++
|
||||
r.lastModelRateLimitKey = scope
|
||||
r.lastModelRateLimitedUntil = until
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestOpenAI429FastPath_KeepsOAuthAccountSchedulableDuringRetryWindow(t *testing.T) {
|
||||
repo := &oauth429RateLimitRepo{}
|
||||
rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
@@ -85,6 +95,59 @@ func TestOpenAI429FastPath_BlocksOAuthImmediatelyWhenSevenDayQuotaIsExhausted(t
|
||||
require.False(t, svc.ShouldRetryOpenAIOAuth429(account, headers, nil))
|
||||
}
|
||||
|
||||
func TestOpenAI429FastPath_SparkQuotaOnlyBlocksSparkModel(t *testing.T) {
|
||||
repo := &oauth429RateLimitRepo{}
|
||||
rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
svc := &OpenAIGatewayService{rateLimitService: rateLimits}
|
||||
rateLimits.SetAccountRuntimeBlocker(svc)
|
||||
account := &Account{ID: 425, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
headers := http.Header{}
|
||||
headers.Set("x-codex-primary-used-percent", "100")
|
||||
headers.Set("x-codex-primary-reset-after-seconds", "604800")
|
||||
headers.Set("x-codex-primary-window-minutes", "10080")
|
||||
headers.Set("x-codex-secondary-used-percent", "20")
|
||||
headers.Set("x-codex-secondary-reset-after-seconds", "3600")
|
||||
headers.Set("x-codex-secondary-window-minutes", "300")
|
||||
|
||||
shouldDisable := svc.handleOpenAIAccountUpstreamError(
|
||||
context.Background(), account, http.StatusTooManyRequests, headers,
|
||||
[]byte(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`),
|
||||
"gpt-5.3-codex-spark",
|
||||
)
|
||||
|
||||
require.False(t, shouldDisable)
|
||||
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account), "Spark quota must not create an account runtime block")
|
||||
require.Equal(t, 0, repo.setRateLimitedCalls, "Spark quota must not persist account-level rate limit")
|
||||
require.Equal(t, 1, repo.setModelRateLimitCalls)
|
||||
require.Equal(t, "gpt-5.3-codex-spark", repo.lastModelRateLimitKey)
|
||||
require.Greater(t, time.Until(repo.lastModelRateLimitedUntil), 6*24*time.Hour)
|
||||
}
|
||||
|
||||
func TestOpenAI429FastPath_SparkShadowQuotaStaysModelScoped(t *testing.T) {
|
||||
repo := &oauth429RateLimitRepo{}
|
||||
rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
svc := &OpenAIGatewayService{rateLimitService: rateLimits}
|
||||
rateLimits.SetAccountRuntimeBlocker(svc)
|
||||
parentID := int64(426)
|
||||
shadow := &Account{ID: 427, Platform: PlatformOpenAI, Type: AccountTypeOAuth, ParentAccountID: &parentID, QuotaDimension: QuotaDimensionSpark}
|
||||
headers := http.Header{}
|
||||
headers.Set("x-codex-primary-used-percent", "100")
|
||||
headers.Set("x-codex-primary-reset-after-seconds", "604800")
|
||||
headers.Set("x-codex-primary-window-minutes", "10080")
|
||||
|
||||
shouldDisable := svc.handleOpenAIAccountUpstreamError(
|
||||
context.Background(), shadow, http.StatusTooManyRequests, headers,
|
||||
[]byte(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`),
|
||||
"gpt-5.3-codex-spark",
|
||||
)
|
||||
|
||||
require.False(t, shouldDisable)
|
||||
require.False(t, svc.isOpenAIAccountRuntimeBlocked(shadow))
|
||||
require.Equal(t, 0, repo.setRateLimitedCalls)
|
||||
require.Equal(t, 1, repo.setModelRateLimitCalls)
|
||||
require.Equal(t, "gpt-5.3-codex-spark", repo.lastModelRateLimitKey)
|
||||
}
|
||||
|
||||
func TestOpenAI429FastPath_RetriesOAuthWhenNoQuotaSignalExists(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{}
|
||||
account := &Account{ID: 424, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
|
||||
@@ -1558,6 +1558,7 @@ func (s *OpenAIGatewayService) handleOpenAIStreamTerminalAccountSideEffects(
|
||||
payload []byte,
|
||||
message string,
|
||||
headers http.Header,
|
||||
canonicalModel ...string,
|
||||
) (int, bool) {
|
||||
statusCode := openAIStreamFailureStatus(payload, message)
|
||||
switch statusCode {
|
||||
@@ -1578,7 +1579,11 @@ func (s *OpenAIGatewayService) handleOpenAIStreamTerminalAccountSideEffects(
|
||||
// carried by a stream terminal event.
|
||||
accountHeaders = nil
|
||||
}
|
||||
return statusCode, s.handleOpenAIAccountUpstreamError(ctx, account, statusCode, accountHeaders, payload)
|
||||
model := firstNonEmpty(canonicalModel...)
|
||||
if model == "" {
|
||||
model = firstNonEmpty(gjson.GetBytes(payload, "model").String(), gjson.GetBytes(payload, "response.model").String())
|
||||
}
|
||||
return statusCode, s.handleOpenAIAccountUpstreamError(ctx, account, statusCode, accountHeaders, payload, model)
|
||||
default:
|
||||
return statusCode, false
|
||||
}
|
||||
@@ -1873,7 +1878,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough(
|
||||
return
|
||||
}
|
||||
if bareErrorAccountSideEffectsPending {
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(c, account, bareErrorPayload, failedMessage, resp.Header)
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(c, account, bareErrorPayload, failedMessage, resp.Header, mappedModel)
|
||||
bareErrorAccountSideEffectsPending = false
|
||||
}
|
||||
if clientDisconnected || !writePendingLines() {
|
||||
@@ -2005,7 +2010,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough(
|
||||
// account health; EOF synthesis applies the pending effect.
|
||||
bareErrorAccountSideEffectsPending = true
|
||||
} else {
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(c, account, dataBytes, failedMessage, resp.Header)
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(c, account, dataBytes, failedMessage, resp.Header, mappedModel)
|
||||
bareErrorAccountSideEffectsPending = false
|
||||
}
|
||||
}
|
||||
|
||||
@@ -372,7 +372,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.
|
||||
completeGuardedEvent(true)
|
||||
}
|
||||
if codexFailureTerminal && sawBareError && !sawResponseFailed && bareErrorAccountSideEffectsPending {
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(c, account, bareErrorPayload, failedMessage, resp.Header)
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(c, account, bareErrorPayload, failedMessage, resp.Header, mappedModel)
|
||||
bareErrorAccountSideEffectsPending = false
|
||||
}
|
||||
if codexFailureTerminal && sawBareError && !sawResponseFailed && !clientDisconnected {
|
||||
@@ -557,7 +557,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.
|
||||
// Defer account health updates so the pair is applied once.
|
||||
bareErrorAccountSideEffectsPending = true
|
||||
} else {
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(c, account, dataBytes, failedMessage, resp.Header)
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(c, account, dataBytes, failedMessage, resp.Header, mappedModel)
|
||||
bareErrorAccountSideEffectsPending = false
|
||||
}
|
||||
}
|
||||
|
||||
@@ -866,7 +866,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
)
|
||||
var dialErr *openAIWSDialError
|
||||
if errors.As(acquireErr, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests {
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(acquireErr.Error()))
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(acquireErr.Error()), canonicalModel)
|
||||
return nil, s.newOpenAIWSRateLimitFailoverError(account, dialErr.ResponseHeaders, nil, acquireErr.Error())
|
||||
}
|
||||
if errors.Is(acquireErr, errOpenAIWSPreferredConnUnavailable) {
|
||||
@@ -1019,7 +1019,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
}
|
||||
}
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw)
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw, mappedModel)
|
||||
fallbackReason, _ := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw)
|
||||
errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw)
|
||||
recoverablePrevNotFound := fallbackReason == openAIWSIngressStagePreviousResponseNotFound &&
|
||||
|
||||
@@ -131,7 +131,8 @@ func (s *OpenAIGatewayService) performOpenAIWSGeneratePrewarm(
|
||||
|
||||
if eventType == "error" {
|
||||
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message)
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw)
|
||||
prewarmModel, _ := reqBody["model"].(string)
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw, prewarmModel)
|
||||
errMsg := strings.TrimSpace(errMsgRaw)
|
||||
if errMsg == "" {
|
||||
errMsg = "OpenAI websocket prewarm error"
|
||||
@@ -336,13 +337,13 @@ func (s *OpenAIGatewayService) handleOpenAIWSFailureAccountSideEffects(ctx conte
|
||||
status := openAIStreamFailureStatus(payload, message)
|
||||
switch status {
|
||||
case http.StatusUnauthorized, http.StatusTooManyRequests, 529:
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, message, headers)
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, message, headers, canonicalModel)
|
||||
return true
|
||||
case http.StatusForbidden:
|
||||
if !openAIStream403AccountFailure(payload, message) {
|
||||
return false
|
||||
}
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, message, headers)
|
||||
s.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, message, headers, canonicalModel)
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -683,14 +684,18 @@ func isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) persistOpenAIWSRateLimitSignal(ctx context.Context, account *Account, headers http.Header, responseBody []byte, codeRaw, errTypeRaw, msgRaw string) {
|
||||
func (s *OpenAIGatewayService) persistOpenAIWSRateLimitSignal(ctx context.Context, account *Account, headers http.Header, responseBody []byte, codeRaw, errTypeRaw, msgRaw string, canonicalModel ...string) {
|
||||
if s == nil || s.rateLimitService == nil || account == nil || account.Platform != PlatformOpenAI {
|
||||
return
|
||||
}
|
||||
if !isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) {
|
||||
return
|
||||
}
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, http.StatusTooManyRequests, headers, responseBody)
|
||||
model := firstNonEmpty(canonicalModel...)
|
||||
if model == "" {
|
||||
model = firstNonEmpty(gjson.GetBytes(responseBody, "model").String(), gjson.GetBytes(responseBody, "response.model").String())
|
||||
}
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, http.StatusTooManyRequests, headers, responseBody, model)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) newOpenAIWSRateLimitFailoverError(account *Account, headers http.Header, responseBody []byte, message string) *UpstreamFailoverError {
|
||||
|
||||
@@ -245,7 +245,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
)
|
||||
var dialErr *openAIWSDialError
|
||||
if errors.As(err, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests {
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error()))
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error()), mappedModel)
|
||||
}
|
||||
return nil, wrapOpenAIWSFallback(classifyOpenAIWSAcquireError(err), err)
|
||||
}
|
||||
@@ -601,7 +601,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
if eventType == "error" {
|
||||
s.handleOpenAIWSErrorEventTransientFailure(ctx, account, mappedModel, lease.HandshakeHeaders(), message)
|
||||
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message)
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw)
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw, mappedModel)
|
||||
errMsg := strings.TrimSpace(errMsgRaw)
|
||||
if errMsg == "" {
|
||||
errMsg = "Upstream websocket error"
|
||||
|
||||
@@ -902,7 +902,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
|
||||
)
|
||||
s.handleOpenAIWSDialTransientFailure(ctx, account, capturedSessionModel, dialErr)
|
||||
if statusCode == http.StatusTooManyRequests {
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, handshakeHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error()))
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, handshakeHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error()), capturedSessionModel)
|
||||
return s.newOpenAIWSRateLimitFailoverError(account, handshakeHeaders, nil, err.Error())
|
||||
}
|
||||
return s.mapOpenAIWSPassthroughDialError(err, statusCode, handshakeHeaders)
|
||||
@@ -1291,7 +1291,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
|
||||
if wroteDownstream || !isOpenAIWSRateLimitError(errCodeRaw, errTypeRaw, errMsgRaw) {
|
||||
return nil
|
||||
}
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, handshakeHeaders, payload, errCodeRaw, errTypeRaw, errMsgRaw)
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, handshakeHeaders, payload, errCodeRaw, errTypeRaw, errMsgRaw, capturedSessionModel)
|
||||
logOpenAIWSV2Passthrough(
|
||||
"relay_rate_limit_failover account_id=%d err_code=%s err_type=%s err_message=%s",
|
||||
account.ID,
|
||||
|
||||
@@ -2192,6 +2192,38 @@ func (s *RateLimitService) HandleOpenAIImageRateLimit(ctx context.Context, accou
|
||||
return true
|
||||
}
|
||||
|
||||
// HandleOpenAICodexSparkRateLimit 将 Spark 独立配额窗口记录为模型级限流。
|
||||
// Spark 的 x-codex-* 使用率和 reset 时间只代表 Spark 模型维度,不能写入账号级
|
||||
// RateLimitResetAt,否则同一 OAuth 账号上的其他模型也会被错误停调。
|
||||
func (s *RateLimitService) HandleOpenAICodexSparkRateLimit(ctx context.Context, account *Account, requestedModel string, statusCode int, headers http.Header, responseBody []byte) bool {
|
||||
if s == nil || account == nil || s.accountRepo == nil || statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) {
|
||||
return false
|
||||
}
|
||||
if !isCodexSparkModel(requestedModel) || !account.ShouldHandleErrorCode(statusCode) {
|
||||
return false
|
||||
}
|
||||
|
||||
modelKey := normalizeCodexModel(modelRateLimitKeyForUpstreamModelNotFound(ctx, account, requestedModel))
|
||||
if modelKey == "" {
|
||||
return false
|
||||
}
|
||||
now := time.Now()
|
||||
_, resetAt := classifyOpenAIOAuth429(headers, responseBody)
|
||||
if resetAt == nil || !resetAt.After(now) {
|
||||
cooldown, ok := s.get429FallbackCooldown(ctx, account)
|
||||
if !ok || cooldown <= 0 {
|
||||
cooldown = time.Duration(defaultRateLimit429CooldownSeconds) * time.Second
|
||||
}
|
||||
reset := now.Add(cooldown)
|
||||
resetAt = &reset
|
||||
}
|
||||
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, *resetAt, openAICodexSparkRateLimitReason); err != nil {
|
||||
slog.Warn("openai_codex_spark_model_rate_limit_set_failed", "account_id", account.ID, "model", modelKey, "error", err)
|
||||
}
|
||||
slog.Info("openai_codex_spark_model_rate_limited", "account_id", account.ID, "model", modelKey, "reset_at", *resetAt)
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *RateLimitService) HandleOpenAIImageCapabilityLoss(ctx context.Context, account *Account, statusCode int, responseBody []byte) bool {
|
||||
if s == nil || account == nil || s.accountRepo == nil {
|
||||
return false
|
||||
|
||||
Reference in New Issue
Block a user