From ca0d3314cf67aa9342ae9aa2fd75a443859cf638 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Thu, 23 Jul 2026 00:06:04 +0800 Subject: [PATCH] fix(grok): cool down accounts after 402 --- ...i_gateway_credential_failover_loop_test.go | 30 ++++++++++++++++++- .../internal/service/openai_gateway_grok.go | 2 ++ .../service/openai_gateway_grok_test.go | 27 +++++++++++++++++ 3 files changed, 58 insertions(+), 1 deletion(-) diff --git a/backend/internal/handler/openai_gateway_credential_failover_loop_test.go b/backend/internal/handler/openai_gateway_credential_failover_loop_test.go index b0aeb19f9d..093a00e095 100644 --- a/backend/internal/handler/openai_gateway_credential_failover_loop_test.go +++ b/backend/internal/handler/openai_gateway_credential_failover_loop_test.go @@ -594,6 +594,32 @@ func TestResponsesGrok429FailoverIsBounded(t *testing.T) { }) } +func TestResponsesGrok402FailoverCooldown(t *testing.T) { + gin.SetMode(gin.TestMode) + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_402") + defer cleanup() + + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Contains(t, recorder.Body.String(), "resp_healthy") + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + require.Equal(t, []int64{801}, repo.setTempIDs) + before := repo.selectorCalls() + + second := httptest.NewRecorder() + secondReq := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"again","stream":false}`)) + secondReq.Header.Set("Content-Type", "application/json") + router.ServeHTTP(second, secondReq) + + require.Equal(t, http.StatusOK, second.Code, second.Body.String()) + require.Equal(t, before+1, repo.selectorCalls()) + require.Equal(t, []int64{801, 802, 802}, upstream.accountHits(), "cooldown must exclude the 402 account from later requests") +} + func TestResponsesGrok429FailoverHandlesMixedStatuses(t *testing.T) { gin.SetMode(gin.TestMode) @@ -836,7 +862,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway Extra: map[string]any{service.GrokMediaEligibleExtraKey: true}, }, } - if mode == "postmap_cancel" || mode == "first_429" || mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" { + if mode == "postmap_cancel" || mode == "first_402" || mode == "first_429" || mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" { accounts[0].Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) } if mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" { @@ -879,6 +905,8 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway } upstream := &grokCredentialHandlerUpstream{} switch mode { + case "first_402": + upstream.failAccountID = 801 case "first_429": upstream.rateLimitIDs = map[int64]bool{801: true} case "all_429": diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 08109d5dc8..d887ff73fc 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -1353,6 +1353,8 @@ func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Contex switch statusCode { case http.StatusUnauthorized: s.tempUnscheduleGrok(ctx, account, 10*time.Minute, "grok credentials unauthorized") + case http.StatusPaymentRequired: + s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok payment required") case http.StatusForbidden: if s.applyGrokForbiddenPolicy(ctx, account, responseBody) { return diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 922a57be09..3a8e347332 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -2439,6 +2439,13 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *test wantMinCooldown: 30*time.Minute - time.Second, wantMaxCooldown: 30*time.Minute + time.Second, }, + { + name: "payment required", + status: http.StatusPaymentRequired, + wantReason: "grok payment required", + wantMinCooldown: 30*time.Minute - time.Second, + wantMaxCooldown: 30*time.Minute + time.Second, + }, { name: "upstream temporary error", status: http.StatusInternalServerError, @@ -2483,6 +2490,26 @@ func TestHandleGrokAccountUpstreamError429SetsRateLimitedFromRetryAfter(t *testi require.Zero(t, repo.tempUnschedCalls) } +func TestHandleGrokAccountUpstreamError402RecoversAfterCooldownExpiry(t *testing.T) { + account := &Account{ + ID: 610, Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, Schedulable: true, + } + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusPaymentRequired, nil, nil) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Equal(t, 1, repo.tempUnschedCalls) + + expired := time.Now().Add(-time.Second) + account.TempUnschedulableUntil = &expired + svc.openaiAccountRuntimeBlockUntil.Store(account.ID, expired) + + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.True(t, account.IsSchedulable()) +} + func TestHandleGrokAccountUpstreamError429UsesLatestExhaustedWindowReset(t *testing.T) { now := time.Now() requestReset := now.Add(10 * time.Minute).Truncate(time.Second)