Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -469,8 +469,94 @@ func (s *OpenAIGatewayService) isOpenAIAccountModelRuntimeBlocked(account *Accou
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 {
return s != nil && (s.isOpenAIAccountRuntimeBlocked(account) || s.isOpenAIAccountModelRuntimeBlocked(account, requestedModel))
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() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -545,6 +545,45 @@ func TestOpenAIRuntimeBlock_ClearAccountSchedulingBlock(t *testing.T) {
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
}

func TestRuntimeBlockHonorsClearedPersistedCooldown(t *testing.T) {
svc := &OpenAIGatewayService{}
account := &Account{ID: 92, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
svc.BlockAccountScheduling(account, time.Now().Add(30*time.Minute), "grok payment required")
require.False(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "grok-3"))
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
}

func TestRuntimeBlockConditionalClearSkipsNewerGeneration(t *testing.T) {
svc := &OpenAIGatewayService{}
account := &Account{ID: 94, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
firstUntil := time.Now().Add(10 * time.Minute)
svc.BlockAccountScheduling(account, firstUntil, "stale")
snapshot := svc.peekOpenAIAccountRuntimeBlock(account)
require.True(t, snapshot.blocked)
newerUntil := time.Now().Add(30 * time.Minute)
svc.BlockAccountScheduling(account, newerUntil, "fresh")
svc.clearOpenAIAccountRuntimeBlockIfUnchanged(account.ID, snapshot)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.False(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "grok-3"))
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
}

func TestRuntimeBlockKeepsActivePersistedCooldown(t *testing.T) {
svc := &OpenAIGatewayService{}
until := time.Now().Add(30 * time.Minute)
account := &Account{
ID: 93,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
TempUnschedulableUntil: &until,
}
svc.BlockAccountScheduling(account, until, "grok payment required")
require.True(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "grok-3"))
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
}

func TestShouldStopOpenAIOAuth429Failover_AfterBoundedFullWindows(t *testing.T) {
svc := &OpenAIGatewayService{}
account := &Account{ID: 42, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,8 @@ func TestHandleOpenAITransientError_HardDisableStillBlocksWholeAccount(t *testin

svc.BlockAccountScheduling(account, time.Now().Add(time.Minute), "upstream_disable")

require.True(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5"))
require.True(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.6-sol"))
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.False(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5"))
require.False(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.6-sol"))
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
}
5 changes: 3 additions & 2 deletions backend/internal/service/openai_upstream_transport_error.go
Original file line number Diff line number Diff line change
Expand Up @@ -163,8 +163,9 @@ func (s *OpenAIGatewayService) tempUnscheduleOpenAITransportError(ctx context.Co
until := time.Now().Add(openAITransportErrorTempUnschedDuration)
reason := "upstream transport error (proxy/network): " + safeErr

// Immediate in-memory block (honoured by the scheduler at selection time),
// effective even if the DB write below fails or the account cache lags.
// Immediate in-memory block so this process skips the account until the
// persisted cooldown is visible on the scheduling Account. Selection is
// fail-open: empty snapshot/DB cooldown fields drop a stale local block.
s.BlockAccountScheduling(account, until, "transport_error")

if s.accountRepo == nil {
Expand Down
Loading