fix(grok): avoid pool cooldown on 5xx

This commit is contained in:
wucm667
2026-07-23 12:24:30 +00:00
parent cb24522dd5
commit 51f354f5c2
2 changed files with 38 additions and 1 deletions
@@ -1363,7 +1363,7 @@ func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Contex
case http.StatusTooManyRequests:
// updateGrokUsageSnapshot installs both runtime and durable rate-limit state.
default:
if statusCode >= 500 {
if statusCode >= 500 && !account.IsPoolMode() {
s.tempUnscheduleGrok(ctx, account, 2*time.Minute, "grok upstream temporary error")
}
}
@@ -2475,6 +2475,43 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *test
}
}
func TestHandleGrokAccountUpstreamError5xxRespectsPoolMode(t *testing.T) {
t.Run("pool mode keeps scheduling state", func(t *testing.T) {
account := &Account{
ID: 611,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"pool_mode": true,
},
}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadGateway, nil, nil)
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Zero(t, repo.tempUnschedCalls)
require.Nil(t, account.TempUnschedulableUntil)
require.Empty(t, account.TempUnschedulableReason)
})
t.Run("non-pool mode keeps two minute cooldown", func(t *testing.T) {
account := &Account{ID: 612, Platform: PlatformGrok, Type: AccountTypeAPIKey}
repo := &grokQuotaAccountRepo{}
svc := &OpenAIGatewayService{accountRepo: repo}
before := time.Now()
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadGateway, nil, nil)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
require.Equal(t, 1, repo.tempUnschedCalls)
require.Equal(t, account.ID, repo.lastTempUnschedID)
require.Equal(t, "grok upstream temporary error", repo.lastTempUnschedReason)
require.WithinDuration(t, before.Add(2*time.Minute), repo.lastTempUnschedUntil, time.Second)
})
}
func TestHandleGrokAccountUpstreamError429SetsRateLimitedFromRetryAfter(t *testing.T) {
account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth}
repo := &grokQuotaAccountRepo{}