diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 7b59c28fd47c..482a7cad880f 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -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() { diff --git a/backend/internal/service/openai_account_runtime_block_fastpath_test.go b/backend/internal/service/openai_account_runtime_block_fastpath_test.go index 89d8c5b5ec44..7446be26fd20 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath_test.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath_test.go @@ -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} diff --git a/backend/internal/service/openai_account_runtime_transient_test.go b/backend/internal/service/openai_account_runtime_transient_test.go index 4ba1f16a5500..b25794fe77a0 100644 --- a/backend/internal/service/openai_account_runtime_transient_test.go +++ b/backend/internal/service/openai_account_runtime_transient_test.go @@ -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)) } diff --git a/backend/internal/service/openai_upstream_transport_error.go b/backend/internal/service/openai_upstream_transport_error.go index ed5c5f9c2330..22a1f6e6a782 100644 --- a/backend/internal/service/openai_upstream_transport_error.go +++ b/backend/internal/service/openai_upstream_transport_error.go @@ -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 {