fix(openai): scope Spark quota 429 to model

This commit is contained in:
shaw
2026-08-29 09:57:08 +08:00
parent ac18c588c8
commit 5d9c7abed5
10 changed files with 128 additions and 18 deletions
@@ -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