mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-16 15:43:13 +08:00
OpenAI 转发 fastpath、HandleUpstreamError 及两处池模式分支对 401 一律跳过
临时不可调度规则,显式 scope=model 的 401 规则在 API Key 账号上仍落到
handleAuthError 整号置 error。
- 新增 tempUnschedRulesApply:401 仅在首条命中规则按模型级生效时交给规则,
其余保持认证错误处理
- 凭据级 401(OAuth 账号与 Spark 影子、OpenAI token_invalidated/token_revoked、
{"detail":"Unauthorized"})不受规则作用范围影响,各入口沿用原有处理
- 模型级命中不参与重复 401 升级;fastpath 不再因模型级命中封整号
616 lines
23 KiB
Go
616 lines
23 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
openAIAccountStateUpdateTimeout = 5 * time.Second
|
|
openAIOAuth429FallbackCooldown = 5 * time.Second
|
|
openAIOAuth429RetryWindow = 2 * time.Minute
|
|
openAIOAuth429RetryDelay = 500 * time.Millisecond
|
|
openAIOAuth429MaxRetryDelay = 8 * time.Second
|
|
openAIOAuth429MaxAccountAttempts = 3
|
|
openAIStopSchedulingBridgeCooldown = 2 * time.Minute
|
|
openAIOAuth429StormWindow = 10 * time.Second
|
|
openAIOAuth429StormMaxAccountSwitches = 1
|
|
)
|
|
|
|
// OpenAIOAuth429FailoverState tracks the request-local follow-up budget after
|
|
// the first Grok OAuth 429. Once that 429 occurs, exactly one different account
|
|
// may be attempted; any failure from that follow-up account ends failover.
|
|
type OpenAIOAuth429FailoverState struct {
|
|
grokOAuth429FollowupPending bool
|
|
}
|
|
|
|
type openAIOAuth429Disposition uint8
|
|
|
|
const (
|
|
openAIOAuth429Transient openAIOAuth429Disposition = iota
|
|
openAIOAuth429Quota5h
|
|
openAIOAuth429Quota7d
|
|
openAIOAuth429QuotaReset
|
|
)
|
|
|
|
// classifyOpenAIOAuth429 区分账号配额耗尽信号与普通瞬时 429。只有窗口达到
|
|
// 100% 或响应体明确给出 reset 时间时,才视为配额限流信号。
|
|
func classifyOpenAIOAuth429(headers http.Header, responseBody []byte) (openAIOAuth429Disposition, *time.Time) {
|
|
if snapshot := ParseCodexRateLimitHeaders(headers); snapshot != nil {
|
|
if normalized := snapshot.Normalize(); normalized != nil {
|
|
if normalized.Used7dPercent != nil && *normalized.Used7dPercent >= 100 {
|
|
if normalized.Reset7dSeconds != nil {
|
|
now := time.Now()
|
|
resetAt := now.Add(time.Duration(*normalized.Reset7dSeconds) * time.Second)
|
|
return openAIOAuth429Quota7d, &resetAt
|
|
}
|
|
return openAIOAuth429Quota7d, nil
|
|
}
|
|
if normalized.Used5hPercent != nil && *normalized.Used5hPercent >= 100 {
|
|
if normalized.Reset5hSeconds != nil {
|
|
now := time.Now()
|
|
resetAt := now.Add(time.Duration(*normalized.Reset5hSeconds) * time.Second)
|
|
return openAIOAuth429Quota5h, &resetAt
|
|
}
|
|
return openAIOAuth429Quota5h, nil
|
|
}
|
|
}
|
|
}
|
|
if resetAt := calculateOpenAI429ResetTime(headers); resetAt != nil {
|
|
return openAIOAuth429QuotaReset, resetAt
|
|
}
|
|
if resetUnix := parseOpenAIRateLimitResetTime(responseBody); resetUnix != nil {
|
|
resetAt := time.Unix(*resetUnix, 0)
|
|
return openAIOAuth429QuotaReset, &resetAt
|
|
}
|
|
return openAIOAuth429Transient, nil
|
|
}
|
|
|
|
func openAIAccountStateContext(ctx context.Context) (context.Context, context.CancelFunc) {
|
|
base := context.Background()
|
|
if ctx != nil {
|
|
base = context.WithoutCancel(ctx)
|
|
}
|
|
return context.WithTimeout(base, openAIAccountStateUpdateTimeout)
|
|
}
|
|
|
|
func isOpenAIOAuthAccount(account *Account) bool {
|
|
return account != nil && account.IsOpenAIOAuthLike()
|
|
}
|
|
|
|
func isGrokOAuthAccount(account *Account) bool {
|
|
return account != nil && account.Platform == PlatformGrok && account.Type == AccountTypeOAuth
|
|
}
|
|
|
|
func isOpenAIAccount(account *Account) bool {
|
|
return account != nil && (account.Platform == PlatformOpenAI || account.Platform == PlatformGrok)
|
|
}
|
|
|
|
// handleOpenAIAccountUpstreamError expects canonicalModel to be the model used
|
|
// for scheduling after applying account mapping exactly once.
|
|
func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, canonicalModel ...string) bool {
|
|
if account != nil && account.Platform == PlatformGrok && isGrokContentPolicyRejection(statusCode, responseBody) {
|
|
return false
|
|
}
|
|
// Any non-2xx upstream HTTP response means the model request was actually sent.
|
|
if s != nil {
|
|
scheduleOllamaCloudUsageActivity(s.deferredService, account)
|
|
scheduleOpenCodeGoUsageActivity(s.deferredService, account)
|
|
}
|
|
// Capacity shedding describes this request, not account health. Keep the
|
|
// account schedulable while the request-local retry budget handles recovery.
|
|
if account != nil && account.Platform == PlatformOpenAI && isOpenAIRequestScopedCapacityShed("", responseBody) {
|
|
return false
|
|
}
|
|
stateCtx, cancel := openAIAccountStateContext(ctx)
|
|
defer cancel()
|
|
if account != nil && account.Platform == PlatformOpenAI && isOpenAIHTTPUpstreamAccessStateError(statusCode, "", responseBody) {
|
|
message := "OpenAI upstream account or workspace is unavailable"
|
|
if upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(responseBody)); upstreamMsg != "" {
|
|
message = upstreamMsg
|
|
}
|
|
if s != nil && s.rateLimitService != nil {
|
|
s.rateLimitService.handleAuthError(stateCtx, account, message)
|
|
}
|
|
if s != nil {
|
|
s.BlockAccountScheduling(account, time.Time{}, "openai_access_state")
|
|
}
|
|
return true
|
|
}
|
|
|
|
if account != nil && account.Platform == PlatformOpenAI && isOpenAIContextWindowError("", responseBody) {
|
|
return false
|
|
}
|
|
|
|
if isOpenAIImageRateLimitError(statusCode, responseBody) {
|
|
if s != nil && s.rateLimitService != nil {
|
|
_ = s.rateLimitService.HandleOpenAIImageRateLimit(stateCtx, account, statusCode, headers, responseBody)
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Self-built images requests always carry a matching image_generation tool, so a
|
|
// "tool choice not found in 'tools'" 400 means upstream revoked this account's
|
|
// image capability. Gated on the self-built marker: passthrough clients control
|
|
// their own tools/tool_choice and could otherwise poison a healthy account.
|
|
if isOpenAIImagesSelfBuiltRequest(ctx) && isOpenAIImageCapabilityLossError(statusCode, responseBody) {
|
|
if s != nil && s.rateLimitService != nil {
|
|
_ = s.rateLimitService.HandleOpenAIImageCapabilityLoss(stateCtx, account, statusCode, responseBody)
|
|
}
|
|
return false
|
|
}
|
|
|
|
if s == nil || account == nil {
|
|
return false
|
|
}
|
|
// Team 联动熔断必须先于 model-not-found 与账户级临时不可调度规则的早退。
|
|
if s.rateLimitService != nil {
|
|
s.rateLimitService.maybeHandleOpenAITeamLinkedError(stateCtx, account, statusCode, responseBody)
|
|
}
|
|
stateCtx = withTempUnschedulableModel(stateCtx, canonicalModel)
|
|
if s.rateLimitService != nil && len(canonicalModel) > 0 && s.rateLimitService.HandleUpstreamModelNotFound(stateCtx, account, canonicalModel[0], statusCode, responseBody) {
|
|
return true
|
|
}
|
|
// Isolate a custom temporary-unschedulable match to the known upstream
|
|
// model before entering the generic account error path. This keeps the
|
|
// account available to other models and avoids the account runtime blocker.
|
|
if s.rateLimitService != nil && len(canonicalModel) > 0 && strings.TrimSpace(canonicalModel[0]) != "" &&
|
|
tempUnschedRulesApply(account, statusCode, responseBody, canonicalModel[0]) &&
|
|
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)
|
|
}
|
|
if s.rateLimitService == nil {
|
|
return false
|
|
}
|
|
shouldDisable := s.rateLimitService.HandleUpstreamError(stateCtx, account, statusCode, headers, responseBody)
|
|
tempModel := tempUnschedulableModel(stateCtx, nil)
|
|
modelTempMatched := tempModel != "" && tempUnschedRulesApply(account, statusCode, responseBody, tempModel) &&
|
|
len(matchTempUnschedulableRules(account, statusCode, responseBody)) > 0
|
|
if shouldDisable && !modelTempMatched {
|
|
s.BlockAccountScheduling(account, time.Time{}, "upstream_disable")
|
|
}
|
|
// Pool-mode retryable upstream errors are already bounded by the request-local
|
|
// same-account retry budget. Recording the generic account+model transient
|
|
// cooldown here would block the next approved retry before that budget is used.
|
|
poolModeRetryable := account.IsPoolMode() && account.IsPoolModeRetryableStatus(statusCode)
|
|
if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey &&
|
|
shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) && !poolModeRetryable {
|
|
model := ""
|
|
if len(canonicalModel) > 0 {
|
|
model = canonicalModel[0]
|
|
}
|
|
decision := s.recordOpenAIAccountModelTransientFailure(account, model, time.Now())
|
|
if decision.FailureStreak > 0 {
|
|
slog.Warn("openai_model_transient_state",
|
|
"account_id", account.ID,
|
|
"model", openAIAccountModelTransientModel(model),
|
|
"failure_streak", decision.FailureStreak,
|
|
"cooldown_ms", decision.Cooldown.Milliseconds(),
|
|
"block_scope", "account_model",
|
|
)
|
|
}
|
|
}
|
|
return shouldDisable
|
|
}
|
|
|
|
func shouldCooldownOpenAITransientUpstreamError(statusCode int, responseBody []byte) bool {
|
|
switch statusCode {
|
|
case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout, 520, 521, 522, 523, 524:
|
|
return true
|
|
case http.StatusBadRequest:
|
|
return isOpenAITransientProcessingError(statusCode, "", responseBody)
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) markOpenAIOAuth429RateLimited(ctx context.Context, account *Account, headers http.Header, responseBody []byte) {
|
|
if s == nil || !isOpenAIOAuthAccount(account) {
|
|
return
|
|
}
|
|
// Spark 影子:不按 /responses 429 的 global x-codex-* 信号做内存运行时熔断(同 handle429,外审第8轮 P1)。
|
|
// 同时避免把 spark 的 429 计入全局 429 storm 计数(recordOpenAIOAuth429),否则会误伤母账号 failover 决策。
|
|
if account.IsShadow() {
|
|
return
|
|
}
|
|
s.recordOpenAIOAuth429()
|
|
disposition, resetAt := classifyOpenAIOAuth429(headers, responseBody)
|
|
if disposition == openAIOAuth429Transient && s.openAIOAuth429RetryWindowActive(account) {
|
|
return
|
|
}
|
|
|
|
now := time.Now()
|
|
cooldownUntil := now.Add(openAIOAuth429FallbackCooldown)
|
|
if resetAt != nil && resetAt.After(now) {
|
|
cooldownUntil = *resetAt
|
|
} else if s.rateLimitService != nil {
|
|
cooldown, ok := s.rateLimitService.get429FallbackCooldown(ctx, account)
|
|
if !ok || cooldown <= 0 {
|
|
s.openaiOAuth429RetryStartedAt.Delete(account.ID)
|
|
return
|
|
}
|
|
cooldownUntil = now.Add(cooldown)
|
|
}
|
|
s.BlockAccountScheduling(account, cooldownUntil, "429")
|
|
s.openaiOAuth429RetryStartedAt.Delete(account.ID)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) shouldRetryOpenAIOAuth429OnSameAccount(account *Account, statusCode int, shouldDisable bool) bool {
|
|
return s.shouldRetryOpenAIOAuth429OnSameAccountWithResponse(account, statusCode, shouldDisable, nil, nil)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) shouldRetryOpenAIOAuth429OnSameAccountWithResponse(account *Account, statusCode int, shouldDisable bool, headers http.Header, responseBody []byte) bool {
|
|
if shouldDisable || statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) || account.IsShadow() {
|
|
return false
|
|
}
|
|
disposition, _ := classifyOpenAIOAuth429(headers, responseBody)
|
|
if disposition != openAIOAuth429Transient {
|
|
return false
|
|
}
|
|
// markOpenAIOAuth429RateLimited parks the account once the window expires.
|
|
// Do not accidentally create a fresh window after that transition.
|
|
if s.isOpenAIAccountRuntimeBlocked(account) {
|
|
return false
|
|
}
|
|
return s.openAIOAuth429RetryWindowActive(account)
|
|
}
|
|
|
|
// ShouldRetryOpenAIOAuth429 lets RateLimitService defer persistent account
|
|
// cooldown until the gateway's same-account retry window is exhausted.
|
|
func (s *OpenAIGatewayService) ShouldRetryOpenAIOAuth429(account *Account, headers http.Header, responseBody []byte) bool {
|
|
if s == nil || !isOpenAIOAuthAccount(account) || account.IsShadow() || s.isOpenAIAccountRuntimeBlocked(account) {
|
|
return false
|
|
}
|
|
disposition, _ := classifyOpenAIOAuth429(headers, responseBody)
|
|
if disposition != openAIOAuth429Transient {
|
|
return false
|
|
}
|
|
return s.openAIOAuth429RetryWindowActive(account)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) openAIOAuth429RetryWindowActive(account *Account) bool {
|
|
if s == nil || !isOpenAIOAuthAccount(account) || account.IsShadow() {
|
|
return false
|
|
}
|
|
now := time.Now()
|
|
value, _ := s.openaiOAuth429RetryStartedAt.LoadOrStore(account.ID, now)
|
|
startedAt, ok := value.(time.Time)
|
|
if !ok {
|
|
s.openaiOAuth429RetryStartedAt.Store(account.ID, now)
|
|
startedAt = now
|
|
}
|
|
return now.Before(startedAt.Add(openAIOAuth429RetryWindow))
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) openAIOAuth429RetryDeadline(account *Account) time.Time {
|
|
if s == nil || !isOpenAIOAuthAccount(account) || account.IsShadow() {
|
|
return time.Time{}
|
|
}
|
|
value, ok := s.openaiOAuth429RetryStartedAt.Load(account.ID)
|
|
if !ok {
|
|
return time.Time{}
|
|
}
|
|
startedAt, ok := value.(time.Time)
|
|
if !ok {
|
|
return time.Time{}
|
|
}
|
|
return startedAt.Add(openAIOAuth429RetryWindow)
|
|
}
|
|
|
|
func openAIOAuth429SameAccountRetryDelay(headers http.Header, deadline time.Time) time.Duration {
|
|
delay := openAIOAuth429RetryDelay
|
|
now := time.Now()
|
|
if resetAt := parseRetryAfterResetTime(headers, now); resetAt != nil && resetAt.After(now) {
|
|
delay = resetAt.Sub(now)
|
|
}
|
|
if delay > openAIOAuth429MaxRetryDelay {
|
|
delay = openAIOAuth429MaxRetryDelay
|
|
}
|
|
if remaining := time.Until(deadline); !deadline.IsZero() && delay > remaining {
|
|
delay = remaining
|
|
}
|
|
if delay < 0 {
|
|
return 0
|
|
}
|
|
return delay
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until time.Time, reason string) {
|
|
if s == nil || !isOpenAIAccount(account) {
|
|
return
|
|
}
|
|
mu := s.openAIAccountRuntimeBlockLock(account.ID)
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
_, _ = s.blockAccountSchedulingLocked(account, until, reason)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) openAIAccountRuntimeBlockLock(accountID int64) *sync.Mutex {
|
|
actual, _ := s.openaiAccountRuntimeBlockLocks.LoadOrStore(accountID, &sync.Mutex{})
|
|
mu, ok := actual.(*sync.Mutex)
|
|
if !ok {
|
|
mu = &sync.Mutex{}
|
|
s.openaiAccountRuntimeBlockLocks.Store(accountID, mu)
|
|
}
|
|
return mu
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) blockAccountSchedulingLocked(account *Account, until time.Time, _ string) (uint64, bool) {
|
|
generation := s.openaiAccountRuntimeBlockSequence.Add(1)
|
|
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, generation)
|
|
now := time.Now()
|
|
blockUntil := until
|
|
if blockUntil.IsZero() || !blockUntil.After(now) {
|
|
blockUntil = now.Add(openAIStopSchedulingBridgeCooldown)
|
|
}
|
|
|
|
for {
|
|
current, loaded := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
|
if !loaded {
|
|
actual, stored := s.openaiAccountRuntimeBlockUntil.LoadOrStore(account.ID, blockUntil)
|
|
if !stored {
|
|
return generation, true
|
|
}
|
|
current = actual
|
|
}
|
|
|
|
currentUntil, ok := current.(time.Time)
|
|
if !ok || currentUntil.IsZero() {
|
|
if s.openaiAccountRuntimeBlockUntil.CompareAndSwap(account.ID, current, blockUntil) {
|
|
return generation, true
|
|
}
|
|
continue
|
|
}
|
|
if !blockUntil.After(currentUntil) {
|
|
return generation, false
|
|
}
|
|
if s.openaiAccountRuntimeBlockUntil.CompareAndSwap(account.ID, current, blockUntil) {
|
|
return generation, true
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) ClearAccountSchedulingBlock(accountID int64) {
|
|
if s == nil || accountID <= 0 {
|
|
return
|
|
}
|
|
mu := s.openAIAccountRuntimeBlockLock(accountID)
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
s.openaiAccountRuntimeBlockUntil.Delete(accountID)
|
|
s.openaiOAuth429RetryStartedAt.Delete(accountID)
|
|
s.openaiAccountRuntimeBlockGeneration.Store(accountID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) isOpenAIAccountRuntimeBlocked(account *Account) bool {
|
|
if s == nil || !isOpenAIAccount(account) {
|
|
return false
|
|
}
|
|
mu := s.openAIAccountRuntimeBlockLock(account.ID)
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
value, ok := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
|
if !ok {
|
|
return false
|
|
}
|
|
cooldownUntil, ok := value.(time.Time)
|
|
if !ok || cooldownUntil.IsZero() {
|
|
s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
|
|
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
|
return false
|
|
}
|
|
if time.Now().Before(cooldownUntil) {
|
|
return true
|
|
}
|
|
s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
|
|
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
|
return false
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) getOpenAIAccountModelTransientState() *openAIAccountModelTransientState {
|
|
if s == nil {
|
|
return nil
|
|
}
|
|
s.openaiModelTransientOnce.Do(func() {
|
|
if s.openaiModelTransient == nil {
|
|
s.openaiModelTransient = newOpenAIAccountModelTransientState(openAIModelTransientDefaultMax)
|
|
}
|
|
})
|
|
return s.openaiModelTransient
|
|
}
|
|
|
|
func canonicalOpenAIAccountSchedulingModel(account *Account, requestedModel string) string {
|
|
model := strings.TrimSpace(requestedModel)
|
|
if account == nil || model == "" {
|
|
return model
|
|
}
|
|
if account.IsOpenAI() {
|
|
return resolveOpenAIAccountUpstreamModelForRequest(account, model, false)
|
|
}
|
|
if mapped := strings.TrimSpace(account.GetMappedModel(model)); mapped != "" {
|
|
return mapped
|
|
}
|
|
return model
|
|
}
|
|
|
|
func openAIAccountModelTransientModel(canonicalModel string) string {
|
|
return normalizeOpenAIAccountModelTransientModel(canonicalModel)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) recordOpenAIAccountModelTransientFailure(account *Account, canonicalModel string, now time.Time) openAIAccountModelTransientDecision {
|
|
if s == nil || account == nil {
|
|
return openAIAccountModelTransientDecision{}
|
|
}
|
|
state := s.getOpenAIAccountModelTransientState()
|
|
if state == nil {
|
|
return openAIAccountModelTransientDecision{}
|
|
}
|
|
return state.recordFailure(account.ID, openAIAccountModelTransientModel(canonicalModel), now)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) clearOpenAIAccountModelTransientState(accountID int64, model string) {
|
|
state := s.getOpenAIAccountModelTransientState()
|
|
if state == nil {
|
|
return
|
|
}
|
|
state.recordSuccess(accountID, model)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) isOpenAIAccountModelRuntimeBlocked(account *Account, requestedModel string) bool {
|
|
if s == nil || account == nil {
|
|
return false
|
|
}
|
|
state := s.getOpenAIAccountModelTransientState()
|
|
if state == nil {
|
|
return false
|
|
}
|
|
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, requestedModel)
|
|
return state.isBlocked(account.ID, openAIAccountModelTransientModel(canonicalModel), time.Now())
|
|
}
|
|
|
|
func accountPersistedSchedulingCooldownActive(account *Account) bool {
|
|
if account == nil {
|
|
return false
|
|
}
|
|
now := time.Now()
|
|
if account.TempUnschedulableUntil != nil && now.Before(*account.TempUnschedulableUntil) {
|
|
return true
|
|
}
|
|
if account.RateLimitResetAt != nil && now.Before(*account.RateLimitResetAt) {
|
|
return true
|
|
}
|
|
if account.OverloadUntil != nil && now.Before(*account.OverloadUntil) {
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
type openAIAccountRuntimeBlockSnapshot struct {
|
|
until time.Time
|
|
generation uint64
|
|
blocked bool
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) peekOpenAIAccountRuntimeBlock(account *Account) openAIAccountRuntimeBlockSnapshot {
|
|
if s == nil || !isOpenAIAccount(account) {
|
|
return openAIAccountRuntimeBlockSnapshot{}
|
|
}
|
|
mu := s.openAIAccountRuntimeBlockLock(account.ID)
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
value, ok := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
|
if !ok {
|
|
return openAIAccountRuntimeBlockSnapshot{}
|
|
}
|
|
until, isTime := value.(time.Time)
|
|
if !isTime || until.IsZero() || !time.Now().Before(until) {
|
|
s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
|
|
s.openaiOAuth429RetryStartedAt.Delete(account.ID)
|
|
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
|
return openAIAccountRuntimeBlockSnapshot{}
|
|
}
|
|
generation, _ := s.openaiAccountRuntimeBlockGeneration.Load(account.ID)
|
|
gen, _ := generation.(uint64)
|
|
return openAIAccountRuntimeBlockSnapshot{until: until, generation: gen, blocked: true}
|
|
}
|
|
|
|
// clearOpenAIAccountRuntimeBlockIfUnchanged deletes the in-process account block
|
|
// only when generation and deadline are unchanged. A newer block installed after
|
|
// peek must be kept even if its deadline happens to match.
|
|
func (s *OpenAIGatewayService) clearOpenAIAccountRuntimeBlockIfUnchanged(accountID int64, snapshot openAIAccountRuntimeBlockSnapshot) {
|
|
if s == nil || accountID <= 0 || !snapshot.blocked {
|
|
return
|
|
}
|
|
mu := s.openAIAccountRuntimeBlockLock(accountID)
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
generation, ok := s.openaiAccountRuntimeBlockGeneration.Load(accountID)
|
|
if !ok || generation != snapshot.generation {
|
|
return
|
|
}
|
|
current, ok := s.openaiAccountRuntimeBlockUntil.Load(accountID)
|
|
currentUntil, isTime := current.(time.Time)
|
|
if !ok || !isTime || !currentUntil.Equal(snapshot.until) {
|
|
return
|
|
}
|
|
s.openaiAccountRuntimeBlockUntil.Delete(accountID)
|
|
s.openaiOAuth429RetryStartedAt.Delete(accountID)
|
|
s.openaiAccountRuntimeBlockGeneration.Store(accountID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
|
}
|
|
|
|
// isOpenAIAccountRequestRuntimeBlocked treats persisted cooldown fields on the
|
|
// scheduling Account as source of truth. When TempUnschedulableUntil,
|
|
// RateLimitResetAt, and OverloadUntil are all inactive, a stale local account
|
|
// block is dropped with generation+deadline CAS. Model-scoped transient blocks
|
|
// are left alone. This is fail-open if a DB write failed or the snapshot has
|
|
// not caught up yet: empty cooldown fields drop the local account-level block.
|
|
func (s *OpenAIGatewayService) isOpenAIAccountRequestRuntimeBlocked(account *Account, requestedModel string) bool {
|
|
if s == nil {
|
|
return false
|
|
}
|
|
snapshot := s.peekOpenAIAccountRuntimeBlock(account)
|
|
if snapshot.blocked {
|
|
if accountPersistedSchedulingCooldownActive(account) {
|
|
return true
|
|
}
|
|
s.clearOpenAIAccountRuntimeBlockIfUnchanged(account.ID, snapshot)
|
|
}
|
|
return s.isOpenAIAccountModelRuntimeBlocked(account, requestedModel)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) recordOpenAIOAuth429() {
|
|
if s == nil {
|
|
return
|
|
}
|
|
now := time.Now()
|
|
windowStart := s.openaiOAuth429WindowStartUnixNano.Load()
|
|
if windowStart == 0 || now.Sub(time.Unix(0, windowStart)) >= openAIOAuth429StormWindow {
|
|
if s.openaiOAuth429WindowStartUnixNano.CompareAndSwap(windowStart, now.UnixNano()) {
|
|
s.openaiOAuth429WindowCount.Store(1)
|
|
return
|
|
}
|
|
}
|
|
s.openaiOAuth429WindowCount.Add(1)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) ShouldStopOpenAIOAuth429Failover(account *Account, statusCode int, failedSwitches int, state *OpenAIOAuth429FailoverState) bool {
|
|
if failedSwitches < openAIOAuth429StormMaxAccountSwitches {
|
|
return false
|
|
}
|
|
if state != nil && state.grokOAuth429FollowupPending {
|
|
// The follow-up budget was armed by a Grok OAuth 429. Consume it on
|
|
// any failing follow-up account, even if a mixed pool selected an API-key
|
|
// account next.
|
|
return true
|
|
}
|
|
if isGrokOAuthAccount(account) {
|
|
if state == nil {
|
|
// Preserve the old threshold for callers that have not adopted the
|
|
// request-local state contract yet.
|
|
return statusCode == http.StatusTooManyRequests && failedSwitches >= 2
|
|
}
|
|
if statusCode == http.StatusTooManyRequests {
|
|
state.grokOAuth429FollowupPending = true
|
|
}
|
|
return false
|
|
}
|
|
if statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) {
|
|
return false
|
|
}
|
|
// Each OpenAI OAuth candidate has already consumed its full same-account
|
|
// retry window before reaching this switch point. A global storm is useful
|
|
// telemetry, but must not prevent trying the bounded next-account budget.
|
|
return failedSwitches >= openAIOAuth429MaxAccountAttempts
|
|
}
|