From 343390057da6fe47b00043cdba6326d2591fa5b1 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Tue, 14 Jul 2026 00:34:03 +0800 Subject: [PATCH 1/6] fix(grok): fail over OAuth credential errors safely --- .../handler/admin/grok_oauth_handler_test.go | 7 +- backend/internal/handler/admin/ops_handler.go | 8 +- backend/internal/handler/failover_loop.go | 6 + .../internal/handler/failover_loop_test.go | 49 +- .../gateway_handler_cancellation_test.go | 98 + .../gateway_handler_chat_completions.go | 11 + .../handler/gateway_handler_responses.go | 11 + backend/internal/handler/gateway_helper.go | 19 +- .../internal/handler/gateway_helper_test.go | 17 + backend/internal/handler/grok_media.go | 14 +- .../handler/openai_chat_completions.go | 14 +- ...i_gateway_credential_failover_loop_test.go | 732 ++++++++ ...openai_gateway_credential_failover_test.go | 211 +++ .../handler/openai_gateway_handler.go | 164 +- backend/internal/handler/ops_error_logger.go | 182 +- backend/internal/repository/account_repo.go | 103 ++ .../account_repo_temp_unsched_test.go | 104 ++ .../internal/repository/grok_oauth_client.go | 23 +- .../repository/grok_oauth_client_test.go | 54 +- .../repository/ops_error_where_test.go | 34 + backend/internal/repository/ops_repo.go | 42 +- .../internal/repository/ops_repo_args_test.go | 37 + ...po_get_error_log_by_id_integration_test.go | 17 + backend/internal/service/account.go | 33 +- .../internal/service/account_base_url_test.go | 34 +- .../internal/service/account_test_service.go | 4 +- .../service/account_test_service_grok_test.go | 15 +- backend/internal/service/gateway_service.go | 62 +- .../service/grok_credential_failure.go | 650 +++++++ .../service/grok_credential_failure_test.go | 1622 +++++++++++++++++ backend/internal/service/grok_media.go | 7 +- .../internal/service/grok_oauth_service.go | 10 +- .../internal/service/grok_quota_service.go | 8 +- .../service/grok_quota_service_test.go | 83 +- .../internal/service/grok_token_provider.go | 210 ++- .../service/grok_token_provider_test.go | 226 ++- .../internal/service/grok_token_refresher.go | 3 +- backend/internal/service/oauth_refresh_api.go | 132 +- .../service/oauth_refresh_api_test.go | 208 ++- .../openai_account_runtime_block_fastpath.go | 34 +- .../service/openai_gateway_cc_pipeline.go | 5 +- .../openai_gateway_chat_completions_raw.go | 5 +- .../internal/service/openai_gateway_grok.go | 8 +- .../openai_gateway_grok_chat_bridge.go | 3 +- .../openai_gateway_grok_chat_bridge_test.go | 49 +- .../service/openai_gateway_grok_test.go | 220 ++- .../service/openai_gateway_messages.go | 2 +- .../service/openai_gateway_service.go | 4 + backend/internal/service/ops_models.go | 15 +- backend/internal/service/ops_service.go | 29 +- .../service/ops_service_batch_test.go | 48 + .../internal/service/ops_upstream_context.go | 8 + backend/internal/service/ops_user_error.go | 4 +- .../internal/service/ops_user_error_test.go | 3 +- backend/internal/service/refresh_policy.go | 8 + .../service/token_refresh_service_test.go | 74 +- backend/internal/service/wire.go | 2 +- .../components/admin/usage/UsageFilters.vue | 1 + frontend/src/i18n/locales/en/admin/ops.ts | 2 + frontend/src/i18n/locales/zh/admin/ops.ts | 2 + frontend/src/utils/errorCategory.ts | 1 + .../ops/components/OpsErrorDetailsModal.vue | 1 + .../admin/ops/components/OpsErrorLogTable.vue | 3 + 63 files changed, 5282 insertions(+), 513 deletions(-) create mode 100644 backend/internal/handler/gateway_handler_cancellation_test.go create mode 100644 backend/internal/handler/openai_gateway_credential_failover_loop_test.go create mode 100644 backend/internal/handler/openai_gateway_credential_failover_test.go create mode 100644 backend/internal/repository/ops_repo_args_test.go create mode 100644 backend/internal/service/grok_credential_failure.go create mode 100644 backend/internal/service/grok_credential_failure_test.go diff --git a/backend/internal/handler/admin/grok_oauth_handler_test.go b/backend/internal/handler/admin/grok_oauth_handler_test.go index 64ea044aa3a..9065e3321b4 100644 --- a/backend/internal/handler/admin/grok_oauth_handler_test.go +++ b/backend/internal/handler/admin/grok_oauth_handler_test.go @@ -90,10 +90,13 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) { ID: 42, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), }, }} upstream := &grokQuotaHandlerUpstream{} diff --git a/backend/internal/handler/admin/ops_handler.go b/backend/internal/handler/admin/ops_handler.go index e820aef0c69..0284dbfd0b6 100644 --- a/backend/internal/handler/admin/ops_handler.go +++ b/backend/internal/handler/admin/ops_handler.go @@ -386,8 +386,8 @@ func (h *OpsHandler) ListRequestErrorUpstreamErrors(c *gin.Context) { filter.EndTime = &endTime } filter.View = "all" - filter.Phase = "upstream" - // 上游错误列表需含 status<400 的 recovered 行,显式豁免客户端可见守卫。 + filter.ErrorPhasesAny = []string{"upstream", "account_auth"} + // Provider-health list includes recovered inference and credential rows. filter.IncludeRecoveredUpstream = true filter.Owner = "provider" filter.Source = strings.TrimSpace(c.Query("error_source")) @@ -470,8 +470,8 @@ func (h *OpsHandler) ListUpstreamErrors(c *gin.Context) { } filter.View = parseOpsViewParam(c) - filter.Phase = "upstream" - // 上游错误列表需含 status<400 的 recovered 行,显式豁免客户端可见守卫。 + filter.ErrorPhasesAny = []string{"upstream", "account_auth"} + // Provider-health list includes recovered inference and credential rows. filter.IncludeRecoveredUpstream = true filter.Owner = "provider" filter.Source = strings.TrimSpace(c.Query("error_source")) diff --git a/backend/internal/handler/failover_loop.go b/backend/internal/handler/failover_loop.go index 5838e58f486..15d29662998 100644 --- a/backend/internal/handler/failover_loop.go +++ b/backend/internal/handler/failover_loop.go @@ -71,7 +71,13 @@ func (s *FailoverState) HandleFailoverError( retryLimit int, failoverErr *service.UpstreamFailoverError, ) FailoverAction { + if ctx != nil && ctx.Err() != nil { + return FailoverCanceled + } s.LastFailoverErr = failoverErr + if failoverErr == nil || !failoverErr.ShouldRetryNextAccount() { + return FailoverExhausted + } // 缓存计费判断 if needForceCacheBilling(s.hasBoundSession, failoverErr) { diff --git a/backend/internal/handler/failover_loop_test.go b/backend/internal/handler/failover_loop_test.go index 9fabe75f25c..bcfc91e07a2 100644 --- a/backend/internal/handler/failover_loop_test.go +++ b/backend/internal/handler/failover_loop_test.go @@ -2,6 +2,7 @@ package handler import ( "context" + "net/http" "testing" "time" @@ -128,6 +129,50 @@ func TestSleepWithContext(t *testing.T) { // --------------------------------------------------------------------------- func TestHandleFailoverError_BasicSwitch(t *testing.T) { + t.Run("显式停止不切换账号且旧错误默认仍切换", func(t *testing.T) { + mock := &mockTempUnscheduler{} + fs := NewFailoverState(3, false) + stopErr := &service.UpstreamFailoverError{ + Stage: service.GatewayFailureStageAccountAuth, + Scope: service.GatewayFailureScopeProvider, + NextAccountAction: service.NextAccountStop, + } + + action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformGrok, maxSameAccountRetries, stopErr) + + require.Equal(t, FailoverExhausted, action) + require.Zero(t, fs.SwitchCount) + require.Empty(t, fs.FailedAccountIDs) + require.Equal(t, stopErr, fs.LastFailoverErr) + + legacyErr := newTestFailoverErr(http.StatusTooManyRequests, false, false) + action = fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformGrok, maxSameAccountRetries, legacyErr) + + require.Equal(t, FailoverContinue, action) + require.Equal(t, 1, fs.SwitchCount) + require.Contains(t, fs.FailedAccountIDs, int64(100)) + }) + + t.Run("已取消的认证失败不改变切换状态", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + mock := &mockTempUnscheduler{} + fs := NewFailoverState(3, false) + err := &service.UpstreamFailoverError{ + Stage: service.GatewayFailureStageAccountAuth, + Scope: service.GatewayFailureScopeAccount, + NextAccountAction: service.NextAccountRetry, + } + + action := fs.HandleFailoverError(ctx, mock, 101, service.PlatformGrok, maxSameAccountRetries, err) + + require.Equal(t, FailoverCanceled, action) + require.Zero(t, fs.SwitchCount) + require.Empty(t, fs.FailedAccountIDs) + require.Nil(t, fs.LastFailoverErr) + require.Empty(t, mock.calls) + }) + t.Run("非重试错误_非Antigravity_直接切换", func(t *testing.T) { mock := &mockTempUnscheduler{} fs := NewFailoverState(3, false) @@ -452,8 +497,8 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) { require.Equal(t, FailoverCanceled, action) require.Less(t, elapsed, 100*time.Millisecond, "应立即返回") - // 重试计数仍应递增 - require.Equal(t, 1, fs.SameAccountRetryCount[100]) + // 入口已取消时不得改变任何重试状态。 + require.Zero(t, fs.SameAccountRetryCount[100]) }) t.Run("Antigravity延迟期间context取消", func(t *testing.T) { diff --git a/backend/internal/handler/gateway_handler_cancellation_test.go b/backend/internal/handler/gateway_handler_cancellation_test.go new file mode 100644 index 00000000000..d53961174ad --- /dev/null +++ b/backend/internal/handler/gateway_handler_cancellation_test.go @@ -0,0 +1,98 @@ +//go:build unit + +package handler + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type countingGatewaySchedulerCache struct { + *fakeSchedulerCache + snapshotCalls atomic.Int64 +} + +func (c *countingGatewaySchedulerCache) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) { + c.snapshotCalls.Add(1) + return c.fakeSchedulerCache.GetSnapshot(ctx, bucket) +} + +func TestGatewayHandlerPreCancelledCompatibleRequestsDoNotSelectAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + groupID := int64(9100) + group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformAnthropic, Status: service.StatusActive} + account := &service.Account{ + ID: 9101, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, + AccountGroups: []service.AccountGroup{{AccountID: 9101, GroupID: groupID}}, + } + schedulerCache := &countingGatewaySchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{account}}} + schedulerSnapshot := service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil) + gatewayService := service.NewGatewayService( + nil, &fakeGroupRepo{group: group}, nil, nil, nil, nil, nil, nil, nil, + schedulerSnapshot, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) + cfg := &config.Config{RunMode: config.RunModeSimple} + billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingCacheService.Stop) + h := &GatewayHandler{ + gatewayService: gatewayService, + billingCacheService: billingCacheService, + concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0), + maxAccountSwitches: 1, + cfg: cfg, + } + apiKey := &service.APIKey{ + ID: 9102, UserID: 9103, GroupID: &groupID, Group: group, Status: service.StatusActive, + User: &service.User{ID: 9103, Concurrency: 10, Balance: 100}, + } + + tests := []struct { + name string + path string + body string + call func(*gin.Context) + }{ + { + name: "responses", path: "/v1/responses", body: `{"model":"claude-test","input":"hello","stream":false}`, + call: h.Responses, + }, + { + name: "chat completions", path: "/v1/chat/completions", body: `{"model":"claude-test","messages":[{"role":"user","content":"hello"}],"stream":false}`, + call: h.ChatCompletions, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + schedulerCache.snapshotCalls.Store(0) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + ctx = context.WithValue(ctx, ctxkey.Group, group) + req := httptest.NewRequest(http.MethodPost, tt.path, bytes.NewBufferString(tt.body)).WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + c.Request = req + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10}) + + tt.call(c) + + require.Zero(t, schedulerCache.snapshotCalls.Load(), "a cancelled request must stop before the account selector") + _, selected := c.Get(opsAccountIDKey) + require.False(t, selected) + }) + } +} diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index af9bcdb344e..d7893b6cd4e 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -159,6 +159,9 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { } for { + if c.Request.Context().Err() != nil { + return + } selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, selectionSessionHash, reqModel, fs.FailedAccountIDs, "", int64(0)) if err != nil { if len(fs.FailedAccountIDs) == 0 { @@ -328,6 +331,14 @@ func (h *GatewayHandler) handleCCFailoverExhausted(c *gin.Context, lastErr *serv if streamStarted { return } + if lastErr != nil { + copyFailoverRetryAfter(c, lastErr.ResponseHeaders) + } + if lastErr != nil && lastErr.IsCredentialFailure() { + status, message := credentialFailoverClientResponse(lastErr) + h.chatCompletionsErrorResponse(c, status, "server_error", message) + return + } statusCode := http.StatusBadGateway if lastErr != nil && lastErr.StatusCode > 0 { statusCode = lastErr.StatusCode diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index 8a88d5fa57a..844de476695 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -157,6 +157,9 @@ func (h *GatewayHandler) Responses(c *gin.Context) { fs := NewFailoverState(h.maxAccountSwitches, false) for { + if requestCtx.Err() != nil { + return + } selection, err := h.gatewayService.SelectAccountWithLoadAwareness(requestCtx, apiKey.GroupID, sessionHash, reqModel, fs.FailedAccountIDs, "", int64(0)) if err != nil { if len(fs.FailedAccountIDs) == 0 { @@ -307,6 +310,14 @@ func (h *GatewayHandler) handleResponsesFailoverExhausted(c *gin.Context, lastEr if streamStarted { return // Can't write error after stream started } + if lastErr != nil { + copyFailoverRetryAfter(c, lastErr.ResponseHeaders) + } + if lastErr != nil && lastErr.IsCredentialFailure() { + status, message := credentialFailoverClientResponse(lastErr) + h.responsesErrorResponse(c, status, "server_error", message) + return + } statusCode := http.StatusBadGateway if lastErr != nil && lastErr.StatusCode > 0 { statusCode = lastErr.StatusCode diff --git a/backend/internal/handler/gateway_helper.go b/backend/internal/handler/gateway_helper.go index 8489b17f6d6..1f6daf52914 100644 --- a/backend/internal/handler/gateway_helper.go +++ b/backend/internal/handler/gateway_helper.go @@ -163,20 +163,15 @@ func wrapReleaseOnDone(ctx context.Context, releaseFunc func()) func() { return nil } var once sync.Once - var stop func() bool - - release := func() { - once.Do(func() { - if stop != nil { - _ = stop() - } - releaseFunc() - }) + releaseOnce := func() { + once.Do(releaseFunc) } + stop := context.AfterFunc(ctx, releaseOnce) - stop = context.AfterFunc(ctx, release) - - return release + return func() { + _ = stop() + releaseOnce() + } } // IncrementWaitCount increments the wait count for a user diff --git a/backend/internal/handler/gateway_helper_test.go b/backend/internal/handler/gateway_helper_test.go index 664258f8c9a..bc0f50f0b6f 100644 --- a/backend/internal/handler/gateway_helper_test.go +++ b/backend/internal/handler/gateway_helper_test.go @@ -6,6 +6,8 @@ import ( "sync/atomic" "testing" "time" + + "github.com/stretchr/testify/require" ) // TestWrapReleaseOnDone_NoGoroutineLeak 验证 wrapReleaseOnDone 修复后不会泄露 goroutine @@ -67,6 +69,21 @@ func TestWrapReleaseOnDone_ContextCancellation(t *testing.T) { } } +func TestWrapReleaseOnDone_AlreadyCancelledReleasesExactlyOnce(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + var releaseCount int32 + release := wrapReleaseOnDone(ctx, func() { + atomic.AddInt32(&releaseCount, 1) + }) + release() + + require.Eventually(t, func() bool { + return atomic.LoadInt32(&releaseCount) == 1 + }, time.Second, time.Millisecond) +} + // TestWrapReleaseOnDone_MultipleCallsOnlyReleaseOnce 验证多次调用 release 只释放一次 func TestWrapReleaseOnDone_MultipleCallsOnlyReleaseOnce(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index b7092fcc42f..408020f0a6f 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -174,6 +174,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. routingStart := time.Now() for { + if requestCtx.Err() != nil { + return + } selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability( requestCtx, apiKey.GroupID, @@ -257,11 +260,20 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. if err != nil { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if requestCtx.Err() != nil { + return + } + if failoverErr.ShouldReportAccountScheduleFailure() { + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + } if c.Writer.Size() != writerSizeBeforeForward { h.handleFailoverExhausted(c, failoverErr, true) return } + if !failoverErr.ShouldRetryNextAccount() { + h.handleFailoverExhausted(c, failoverErr, false) + return + } if failoverErr.RetryableOnSameAccount { retryLimit := account.GetPoolModeRetryCount() if sameAccountRetryCount[account.ID] < retryLimit { diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 636e1437404..fc58b3a25ab 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -135,6 +135,9 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { var lastFailoverErr *service.UpstreamFailoverError for { + if c.Request.Context().Err() != nil { + return + } reqLog.Debug("openai_chat_completions.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs))) selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability( c.Request.Context(), @@ -231,11 +234,20 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { } else { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { + if c.Request.Context().Err() != nil { + return + } if c.Writer.Size() != writerSizeBeforeForward { h.handleFailoverExhausted(c, failoverErr, true) return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if failoverErr.ShouldReportAccountScheduleFailure() { + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + } + if !failoverErr.ShouldRetryNextAccount() { + h.handleFailoverExhausted(c, failoverErr, streamStarted) + return + } // Pool mode: retry on the same account if failoverErr.RetryableOnSameAccount { retryLimit := account.GetPoolModeRetryCount() diff --git a/backend/internal/handler/openai_gateway_credential_failover_loop_test.go b/backend/internal/handler/openai_gateway_credential_failover_loop_test.go new file mode 100644 index 00000000000..e7030d9d2ff --- /dev/null +++ b/backend/internal/handler/openai_gateway_credential_failover_loop_test.go @@ -0,0 +1,732 @@ +//go:build unit + +package handler + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type grokCredentialHandlerRepo struct { + service.AccountRepository + mu sync.Mutex + accounts []service.Account + setErrorIDs []int64 + setTempIDs []int64 + updateExtraIDs []int64 + selectionCalls int + setErrorErr error + setTempErr error + missingOnGet map[int64]bool +} + +func (r *grokCredentialHandlerRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]service.Account, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.selectionCalls++ + out := make([]service.Account, 0, len(r.accounts)) + for _, account := range r.accounts { + if account.Platform == platform && account.IsSchedulable() { + out = append(out, account) + } + } + return out, nil +} + +func (r *grokCredentialHandlerRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, _ int64, platform string) ([]service.Account, error) { + return r.ListSchedulableByPlatform(ctx, platform) +} + +func (r *grokCredentialHandlerRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]service.Account, error) { + return r.ListSchedulableByPlatform(ctx, platform) +} + +func (r *grokCredentialHandlerRepo) GetByID(_ context.Context, id int64) (*service.Account, error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.missingOnGet[id] { + return nil, nil + } + for _, account := range r.accounts { + if account.ID == id { + copy := account + copy.Credentials = cloneCredentialMap(account.Credentials) + return ©, nil + } + } + return nil, nil +} + +func (r *grokCredentialHandlerRepo) SetError(_ context.Context, id int64, message string) error { + r.mu.Lock() + defer r.mu.Unlock() + r.setErrorIDs = append(r.setErrorIDs, id) + if r.setErrorErr != nil { + return r.setErrorErr + } + for i := range r.accounts { + if r.accounts[i].ID == id { + r.accounts[i].Status = service.StatusError + r.accounts[i].Schedulable = false + r.accounts[i].ErrorMessage = message + } + } + return nil +} + +func (r *grokCredentialHandlerRepo) SetTempUnschedulable(_ context.Context, id int64, until time.Time, _ string) error { + r.mu.Lock() + defer r.mu.Unlock() + r.setTempIDs = append(r.setTempIDs, id) + if r.setTempErr != nil { + return r.setTempErr + } + for i := range r.accounts { + if r.accounts[i].ID == id { + value := until + r.accounts[i].TempUnschedulableUntil = &value + } + } + return nil +} + +func (r *grokCredentialHandlerRepo) SetGrokCredentialErrorIfMatch( + _ context.Context, + id int64, + snapshot service.GrokCredentialMutationSnapshot, + message string, +) (bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + for i := range r.accounts { + account := &r.accounts[i] + if account.ID != id || !handlerGrokCredentialSnapshotMatches(account, snapshot) { + continue + } + r.setErrorIDs = append(r.setErrorIDs, id) + if r.setErrorErr != nil { + return false, r.setErrorErr + } + account.Status = service.StatusError + account.Schedulable = false + account.ErrorMessage = message + return true, nil + } + return false, nil +} + +func (r *grokCredentialHandlerRepo) SetGrokCredentialTempUnschedulableIfMatch( + _ context.Context, + id int64, + snapshot service.GrokCredentialMutationSnapshot, + until time.Time, + _ string, +) (bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + for i := range r.accounts { + account := &r.accounts[i] + if account.ID != id || !handlerGrokCredentialSnapshotMatches(account, snapshot) { + continue + } + r.setTempIDs = append(r.setTempIDs, id) + if r.setTempErr != nil { + return false, r.setTempErr + } + value := until + account.TempUnschedulableUntil = &value + return true, nil + } + return false, nil +} + +func handlerGrokCredentialSnapshotMatches(account *service.Account, snapshot service.GrokCredentialMutationSnapshot) bool { + if account == nil { + return false + } + credentialsJSON, err := json.Marshal(account.Credentials) + return err == nil && account.IsGrokOAuth() && account.IsSchedulable() && string(credentialsJSON) == snapshot.CredentialsJSON && + handlerGrokCredentialProxyIDsEqual(account.ProxyID, snapshot.ProxyID) +} + +func handlerGrokCredentialProxyIDsEqual(left, right *int64) bool { + if left == nil || right == nil { + return left == nil && right == nil + } + return *left == *right +} + +func (r *grokCredentialHandlerRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error { + r.mu.Lock() + defer r.mu.Unlock() + r.updateExtraIDs = append(r.updateExtraIDs, id) + for i := range r.accounts { + if r.accounts[i].ID != id { + continue + } + if r.accounts[i].Extra == nil { + r.accounts[i].Extra = map[string]any{} + } + for key, value := range updates { + r.accounts[i].Extra[key] = value + } + } + return nil +} + +func (r *grokCredentialHandlerRepo) errorIDs() []int64 { + r.mu.Lock() + defer r.mu.Unlock() + return append([]int64(nil), r.setErrorIDs...) +} + +func (r *grokCredentialHandlerRepo) selectorCalls() int { + r.mu.Lock() + defer r.mu.Unlock() + return r.selectionCalls +} + +type grokCredentialHandlerTokenCache struct { + service.GrokTokenCache + mu sync.Mutex + deleteErr error +} + +func (c *grokCredentialHandlerTokenCache) GetAccessToken(context.Context, string) (string, error) { + return "", errors.New("not cached") +} + +func (c *grokCredentialHandlerTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error { + return nil +} + +func (c *grokCredentialHandlerTokenCache) DeleteAccessToken(context.Context, string) error { + c.mu.Lock() + defer c.mu.Unlock() + return c.deleteErr +} + +func (c *grokCredentialHandlerTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) { + return true, nil +} + +func (c *grokCredentialHandlerTokenCache) ReleaseRefreshLock(context.Context, string) error { + return nil +} + +func cloneCredentialMap(source map[string]any) map[string]any { + cloned := make(map[string]any, len(source)) + for key, value := range source { + cloned[key] = value + } + return cloned +} + +type grokCredentialHandlerRefresher struct { + mode string + started chan struct{} + once sync.Once +} + +func (r *grokCredentialHandlerRefresher) CacheKey(account *service.Account) string { + return service.GrokTokenCacheKey(account) +} + +func (r *grokCredentialHandlerRefresher) CanRefresh(account *service.Account) bool { + return account != nil && account.IsGrokOAuth() +} + +func (r *grokCredentialHandlerRefresher) NeedsRefresh(account *service.Account, _ time.Duration) bool { + return account != nil && (account.ID == 801 || r.mode == "all_revoked") +} + +func (r *grokCredentialHandlerRefresher) Refresh(ctx context.Context, _ *service.Account) (map[string]any, error) { + switch r.mode { + case "revoked", "all_revoked", "mutation_set_error", "mutation_cache": + return nil, infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant") + case "provider": + return nil, infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_client") + case "cancel": + r.once.Do(func() { close(r.started) }) + <-ctx.Done() + return nil, ctx.Err() + case "transient", "mutation_temp": + return nil, errors.New("temporary refresh transport failure") + default: + return nil, nil + } +} + +type grokCredentialHandlerUpstream struct { + service.HTTPUpstream + mu sync.Mutex + hits []int64 + requestURLs []string + authorization []string + failAccountID int64 + cancelRequest context.CancelFunc +} + +func (u *grokCredentialHandlerUpstream) Do(req *http.Request, _ string, accountID int64, _ int) (*http.Response, error) { + var requestBody []byte + if req.Body != nil { + requestBody, _ = io.ReadAll(req.Body) + } + u.mu.Lock() + u.hits = append(u.hits, accountID) + u.requestURLs = append(u.requestURLs, req.URL.String()) + u.authorization = append(u.authorization, req.Header.Get("Authorization")) + failAccountID := u.failAccountID + cancelRequest := u.cancelRequest + u.mu.Unlock() + if accountID == failAccountID { + if cancelRequest != nil { + cancelRequest() + } + return &http.Response{ + StatusCode: http.StatusPaymentRequired, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewBufferString(`{"error":{"message":"payment required"}}`)), + }, nil + } + if bytes.Contains(requestBody, []byte(`"stream":true`)) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewBufferString( + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_healthy\",\"model\":\"grok-4.5\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n", + )), + }, nil + } + if strings.Contains(req.URL.Path, "/chat/completions") { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewBufferString( + `{"id":"chatcmpl_healthy","object":"chat.completion","model":"grok-4.5","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`, + )), + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewBufferString( + `{"id":"resp_healthy","object":"response","model":"grok-4.5","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1}}`, + )), + }, nil +} + +func (u *grokCredentialHandlerUpstream) accountHits() []int64 { + u.mu.Lock() + defer u.mu.Unlock() + return append([]int64(nil), u.hits...) +} + +func (u *grokCredentialHandlerUpstream) requests() ([]string, []string) { + u.mu.Lock() + defer u.mu.Unlock() + return append([]string(nil), u.requestURLs...), append([]string(nil), u.authorization...) +} + +func TestResponsesCredentialFailoverLoop(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("revoked account selects healthy account", func(t *testing.T) { + h, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "revoked") + defer cleanup() + _ = h + + 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}, repo.errorIDs()) + require.Equal(t, []int64{802}, upstream.accountHits()) + requestURLs, authorization := upstream.requests() + require.Equal(t, []string{xai.DefaultCLIBaseURL + "/responses"}, requestURLs) + require.Equal(t, []string{"Bearer healthy-access"}, authorization) + }) + + t.Run("provider configuration stops before healthy account", func(t *testing.T) { + h, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "provider") + 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.StatusServiceUnavailable, recorder.Code) + require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage) + require.Empty(t, repo.errorIDs()) + require.Empty(t, upstream.accountHits()) + require.Equal(t, 1, repo.selectorCalls()) + require.Zero(t, h.gatewayService.SnapshotOpenAIAccountSchedulerMetrics().RuntimeStatsAccountCount, + "provider-scoped auth failure must not penalize the selected account") + }) + + t.Run("parent cancellation stops before healthy account", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "cancel") + defer cleanup() + + ctx, cancel := context.WithCancel(context.Background()) + req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)).WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + done := make(chan struct{}) + go func() { + defer close(done) + router.ServeHTTP(recorder, req) + }() + + select { + case <-time.After(2 * time.Second): + t.Fatal("credential refresh did not start") + case <-findHandlerRefresherStarted(router): + cancel() + } + select { + case <-time.After(2 * time.Second): + t.Fatal("handler did not stop after cancellation") + case <-done: + } + + require.Empty(t, repo.errorIDs()) + require.Empty(t, upstream.accountHits()) + }) + + t.Run("post-mapping cancellation stops before scheduler mutation or reselection", func(t *testing.T) { + h, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "postmap_cancel") + defer cleanup() + ctx, cancel := context.WithCancel(context.Background()) + upstream.mu.Lock() + upstream.failAccountID = 801 + upstream.cancelRequest = cancel + upstream.mu.Unlock() + + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)).WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, req) + + require.Equal(t, []int64{801}, upstream.accountHits()) + require.Empty(t, repo.errorIDs()) + require.Equal(t, 1, repo.selectorCalls()) + require.Zero(t, h.gatewayService.SnapshotOpenAIAccountSchedulerMetrics().RuntimeStatsAccountCount) + }) + + t.Run("pre-cancelled request never invokes an account selector", func(t *testing.T) { + tests := []struct { + name string + method string + path string + body string + }{ + {name: "responses", method: http.MethodPost, path: "/openai/v1/responses", body: `{"model":"grok","input":"hello","stream":false}`}, + {name: "messages", method: http.MethodPost, path: "/openai/v1/messages", body: `{"model":"grok","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`}, + {name: "chat completions", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stream":false}`}, + {name: "grok media", method: http.MethodGet, path: "/openai/v1/videos/request-1"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "revoked") + defer cleanup() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body)).WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, req) + + require.Zero(t, repo.selectorCalls()) + require.Empty(t, upstream.accountHits()) + }) + } + }) + + t.Run("credential state mutation failures stop before reselection", func(t *testing.T) { + for _, mode := range []string{"mutation_set_error", "mutation_temp", "mutation_cache"} { + t.Run(mode, func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, mode) + 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.StatusServiceUnavailable, recorder.Code, recorder.Body.String()) + require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage) + require.Empty(t, upstream.accountHits()) + require.Equal(t, 1, repo.selectorCalls()) + }) + } + }) + + t.Run("missing credential provider stops before upstream or reselection", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "nil_provider") + 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.StatusServiceUnavailable, recorder.Code, recorder.Body.String()) + require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage) + require.Equal(t, 1, repo.selectorCalls()) + require.Empty(t, upstream.accountHits()) + require.Empty(t, repo.errorIDs()) + }) +} + +func TestGrokOAuthCredentialFailoverAcrossHTTPHandlers(t *testing.T) { + gin.SetMode(gin.TestMode) + endpoints := []struct { + name string + method string + path string + body string + }{ + {name: "messages", method: http.MethodPost, path: "/openai/v1/messages", body: `{"model":"grok","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`}, + {name: "chat completions", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stream":false}`}, + {name: "chat completions raw fallback", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stop":["END"],"stream":false}`}, + {name: "grok media", method: http.MethodGet, path: "/openai/v1/videos/request-1"}, + } + + for _, endpoint := range endpoints { + t.Run(endpoint.name+" revoked selects healthy", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "revoked") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(endpoint.method, endpoint.path, bytes.NewBufferString(endpoint.body)) + req.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801}, repo.errorIDs()) + require.Equal(t, []int64{802}, upstream.accountHits()) + }) + + t.Run(endpoint.name+" all accounts exhausted safely", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_revoked") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(endpoint.method, endpoint.path, bytes.NewBufferString(endpoint.body)) + req.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusServiceUnavailable, recorder.Code, recorder.Body.String()) + require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage) + require.NotContains(t, recorder.Body.String(), "revoked-refresh") + require.NotContains(t, recorder.Body.String(), "healthy-refresh") + require.Equal(t, []int64{801, 802}, repo.errorIDs()) + require.Empty(t, upstream.accountHits()) + }) + } +} + +func TestGrokOAuthMissingSelectedRowRetriesHealthyAccountWithoutMutation(t *testing.T) { + gin.SetMode(gin.TestMode) + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "missing_row") + 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.Equal(t, []int64{802}, upstream.accountHits()) + require.Empty(t, repo.errorIDs()) + require.Empty(t, repo.setTempIDs) +} + +func TestResponsesWebSocketCredentialFailoverLoop(t *testing.T) { + gin.SetMode(gin.TestMode) + dial := func(t *testing.T, router *gin.Engine) (*coderws.Conn, func()) { + t.Helper() + server := httptest.NewServer(router) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + conn, _, err := coderws.Dial(ctx, "ws"+strings.TrimPrefix(server.URL, "http")+"/openai/v1/responses", nil) + cancel() + require.NoError(t, err) + return conn, func() { + _ = conn.CloseNow() + server.Close() + } + } + writeFirst := func(t *testing.T, conn *coderws.Conn) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + require.NoError(t, conn.Write(ctx, coderws.MessageText, []byte(`{"type":"response.create","model":"grok","input":"hello","stream":false}`))) + } + + t.Run("revoked account selects healthy account", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "revoked") + defer cleanup() + conn, closeConn := dial(t, router) + defer closeConn() + writeFirst(t, conn) + + readCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + _, payload, err := conn.Read(readCtx) + cancel() + require.NoError(t, err) + require.Contains(t, string(payload), "resp_healthy") + require.Equal(t, []int64{801}, repo.errorIDs()) + require.Equal(t, 2, repo.selectorCalls()) + require.Equal(t, []int64{802}, upstream.accountHits()) + }) + + t.Run("provider configuration stops", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "provider") + defer cleanup() + conn, closeConn := dial(t, router) + defer closeConn() + writeFirst(t, conn) + + readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + _, _, err := conn.Read(readCtx) + cancel() + var closeErr coderws.CloseError + require.ErrorAs(t, err, &closeErr) + require.Contains(t, closeErr.Reason, service.GrokCredentialUnavailableClientMessage) + require.Equal(t, 1, repo.selectorCalls()) + require.Empty(t, upstream.accountHits()) + }) + + t.Run("parent cancellation prevents reselection", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "cancel") + defer cleanup() + conn, closeConn := dial(t, router) + writeFirst(t, conn) + select { + case <-findHandlerRefresherStarted(router): + case <-time.After(2 * time.Second): + t.Fatal("credential refresh did not start") + } + closeConn() + + require.Eventually(t, func() bool { return repo.selectorCalls() == 1 }, 2*time.Second, 20*time.Millisecond) + require.Empty(t, repo.errorIDs()) + require.Empty(t, upstream.accountHits()) + }) +} + +var handlerRefresherStarted sync.Map + +func findHandlerRefresherStarted(router *gin.Engine) <-chan struct{} { + value, _ := handlerRefresherStarted.Load(router) + return value.(chan struct{}) +} + +func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGatewayHandler, *grokCredentialHandlerRepo, *grokCredentialHandlerUpstream, *gin.Engine, func()) { + t.Helper() + groupID := int64(901) + accounts := []service.Account{ + { + ID: 801, Name: "revoked", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, + Credentials: map[string]any{ + "access_token": "expired", "refresh_token": "revoked-refresh", + "expires_at": time.Now().Add(-time.Minute).UTC().Format(time.RFC3339), + }, + }, + { + ID: 802, Name: "healthy", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, Priority: 2, + Credentials: map[string]any{ + "access_token": "healthy-access", "refresh_token": "healthy-refresh", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), + }, + }, + } + if mode == "postmap_cancel" { + accounts[0].Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + } + if mode == "all_revoked" { + accounts[1].Credentials["expires_at"] = time.Now().Add(-time.Minute).UTC().Format(time.RFC3339) + } + repo := &grokCredentialHandlerRepo{accounts: accounts, missingOnGet: map[int64]bool{}} + if mode == "missing_row" { + repo.missingOnGet[801] = true + } + if mode == "mutation_set_error" { + repo.setErrorErr = errors.New("database write failed") + } + if mode == "mutation_temp" { + repo.setTempErr = errors.New("database write failed") + } + refresher := &grokCredentialHandlerRefresher{mode: mode, started: make(chan struct{})} + tokenCache := &grokCredentialHandlerTokenCache{} + if mode == "mutation_cache" { + tokenCache.deleteErr = errors.New("cache delete failed") + } + var provider *service.GrokTokenProvider + if mode != "nil_provider" { + provider = service.NewGrokTokenProvider(repo, tokenCache) + provider.SetRefreshAPI(service.NewOAuthRefreshAPI(repo, tokenCache), refresher) + } + upstream := &grokCredentialHandlerUpstream{} + cfg := &config.Config{RunMode: config.RunModeSimple} + cfg.Gateway.MaxAccountSwitches = 3 + billingCache := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + gateway := service.NewOpenAIGatewayService( + repo, nil, nil, nil, nil, nil, nil, cfg, nil, nil, + service.NewBillingService(cfg, nil), nil, billingCache, upstream, + &service.DeferredService{}, nil, provider, nil, nil, nil, nil, nil, + ) + cache := &concurrencyCacheMock{ + acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil }, + acquireAccountSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil }, + } + h := NewOpenAIGatewayHandler(gateway, service.NewConcurrencyService(cache), billingCache, &service.APIKeyService{}, nil, nil, nil, nil, cfg) + apiKey := &service.APIKey{ + ID: 902, GroupID: &groupID, + User: &service.User{ID: 903, Status: service.StatusActive}, + Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, Status: service.StatusActive}, + } + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1}) + c.Next() + }) + router.POST("/openai/v1/responses", h.Responses) + router.GET("/openai/v1/responses", h.ResponsesWebSocket) + router.POST("/openai/v1/messages", h.Messages) + router.POST("/openai/v1/chat/completions", h.ChatCompletions) + router.GET("/openai/v1/videos/:request_id", h.GrokVideoStatus) + handlerRefresherStarted.Store(router, refresher.started) + cleanup := func() { + handlerRefresherStarted.Delete(router) + billingCache.Stop() + } + return h, repo, upstream, router, cleanup +} diff --git a/backend/internal/handler/openai_gateway_credential_failover_test.go b/backend/internal/handler/openai_gateway_credential_failover_test.go new file mode 100644 index 00000000000..07a5e4999f6 --- /dev/null +++ b/backend/internal/handler/openai_gateway_credential_failover_test.go @@ -0,0 +1,211 @@ +//go:build unit + +package handler + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestGatewayChatCredentialStopDoesNotSelectAnotherAccountAndReturnsSafe503(t *testing.T) { + gin.SetMode(gin.TestMode) + stopErr := &service.UpstreamFailoverError{ + Stage: service.GatewayFailureStageAccountAuth, + Scope: service.GatewayFailureScopeProvider, + Reason: service.GrokCredentialReasonProviderConfig, + NextAccountAction: service.NextAccountStop, + ClientStatusCode: http.StatusTeapot, + ClientMessage: "invalid_client client_secret=must-not-leak", + } + state := NewFailoverState(3, false) + action := state.HandleFailoverError(context.Background(), &mockTempUnscheduler{}, 71, service.PlatformGrok, 0, stopErr) + + require.Equal(t, FailoverExhausted, action) + require.Zero(t, state.SwitchCount) + require.Empty(t, state.FailedAccountIDs) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + (&GatewayHandler{}).handleCCFailoverExhausted(c, state.LastFailoverErr, false) + + require.Equal(t, http.StatusServiceUnavailable, recorder.Code) + require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage) + require.NotContains(t, recorder.Body.String(), "invalid_client") + require.NotContains(t, recorder.Body.String(), "client_secret") +} + +func TestGatewayChatInferenceExhaustionRestoresRetryAfter(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + + (&GatewayHandler{}).handleCCFailoverExhausted(c, &service.UpstreamFailoverError{ + StatusCode: http.StatusTooManyRequests, + ResponseHeaders: http.Header{"Retry-After": []string{"45"}}, + }, false) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code) + require.Equal(t, "45", recorder.Header().Get("Retry-After")) +} + +func TestCredentialFailoverExhaustionReturnsFixedSafe503(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + h := &OpenAIGatewayHandler{} + + h.handleFailoverExhausted(c, &service.UpstreamFailoverError{ + Stage: service.GatewayFailureStageAccountAuth, + Scope: service.GatewayFailureScopeAccount, + Reason: service.GrokCredentialReasonRevoked, + NextAccountAction: service.NextAccountRetry, + ClientStatusCode: http.StatusTeapot, + ClientMessage: "invalid_grant refresh_token=must-not-leak", + }, false) + + require.Equal(t, http.StatusServiceUnavailable, recorder.Code) + require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage) + require.NotContains(t, strings.ToLower(recorder.Body.String()), "invalid_grant") + require.NotContains(t, strings.ToLower(recorder.Body.String()), "refresh_token") + require.NotContains(t, recorder.Body.String(), "must-not-leak") +} + +func TestInferenceFailoverExhaustionRestoresRetryAfter(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + h := &OpenAIGatewayHandler{} + + h.handleFailoverExhausted(c, &service.UpstreamFailoverError{ + StatusCode: http.StatusTooManyRequests, + ResponseHeaders: http.Header{"Retry-After": []string{"17"}}, + }, false) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code) + require.Equal(t, "17", recorder.Header().Get("Retry-After")) +} + +func TestFailoverExhaustionRejectsSecretBearingRetryAfter(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + h := &OpenAIGatewayHandler{} + + h.handleFailoverExhausted(c, &service.UpstreamFailoverError{ + StatusCode: http.StatusTooManyRequests, + ResponseHeaders: http.Header{"Retry-After": []string{"refresh_token=must-not-leak"}}, + }, false) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code) + require.Empty(t, recorder.Header().Get("Retry-After")) + require.NotContains(t, recorder.Body.String(), "must-not-leak") +} + +func TestFailoverExhaustionRejectsFarFutureRetryAfterDate(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + h := &OpenAIGatewayHandler{} + + h.handleFailoverExhausted(c, &service.UpstreamFailoverError{ + StatusCode: http.StatusTooManyRequests, + ResponseHeaders: http.Header{ + "Retry-After": []string{time.Now().Add(30 * 24 * time.Hour).UTC().Format(http.TimeFormat)}, + }, + }, false) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code) + require.Empty(t, recorder.Header().Get("Retry-After")) +} + +func TestFailoverExhaustionAllowsBoundedRetryAfterDate(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + h := &OpenAIGatewayHandler{} + retryAfter := time.Now().Add(time.Hour).UTC().Format(http.TimeFormat) + + h.handleFailoverExhausted(c, &service.UpstreamFailoverError{ + StatusCode: http.StatusTooManyRequests, + ResponseHeaders: http.Header{"Retry-After": []string{retryAfter}}, + }, false) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code) + require.Equal(t, retryAfter, recorder.Header().Get("Retry-After")) +} + +func TestOpsClassificationTreatsCredentialFailureAsAuthNotInference(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set(service.OpsUpstreamStatusCodeKey, http.StatusForbidden) + c.Set(service.OpsUpstreamErrorMessageKey, "stale inference message") + c.Set(service.OpsUpstreamErrorDetailKey, "stale inference detail") + c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{ + {Stage: string(service.GatewayFailureStageInference), UpstreamStatusCode: http.StatusForbidden, Message: "stale inference message", Detail: "stale inference detail"}, + { + Stage: string(service.GatewayFailureStageAccountAuth), + Scope: string(service.GatewayFailureScopeAccount), + Reason: string(service.GrokCredentialReasonRevoked), + UpstreamStatusCode: 0, + Message: "Grok OAuth credentials require account action", + }, + }) + + phase, _, owner, source := classifyOpsErrorLog(c, "upstream_error", service.GrokCredentialUnavailableClientMessage, "", http.StatusServiceUnavailable) + require.Equal(t, "account_auth", phase) + require.Equal(t, "provider", owner) + require.Equal(t, "gateway", source) + + entry := &service.OpsInsertErrorLogInput{} + applyOpsUpstreamFieldsFromContext(c, entry) + require.NotNil(t, entry.UpstreamStatusCode) + require.Zero(t, *entry.UpstreamStatusCode) + require.NotNil(t, entry.UpstreamErrorMessage) + require.Equal(t, "Grok OAuth credentials require account action", *entry.UpstreamErrorMessage) + require.Nil(t, entry.UpstreamErrorDetail) + require.Len(t, entry.UpstreamErrors, 2) + require.Equal(t, http.StatusForbidden, entry.UpstreamErrors[0].UpstreamStatusCode) +} + +func TestOpsRecoveredCredentialFailoverUsesAccountAuthAttribution(t *testing.T) { + setupOpsErrorLogTestQueue(t, 2) + gin.SetMode(gin.TestMode) + ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + router := gin.New() + router.Use(OpsErrorLoggerMiddleware(ops)) + router.GET("/openai/v1/responses", func(c *gin.Context) { + c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{ + {Stage: string(service.GatewayFailureStageInference), UpstreamStatusCode: http.StatusForbidden, Message: "earlier inference failure"}, + { + Stage: string(service.GatewayFailureStageAccountAuth), Scope: string(service.GatewayFailureScopeAccount), + Reason: string(service.GrokCredentialReasonRevoked), Message: "Grok OAuth credentials require account action", + }, + }) + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/openai/v1/responses", nil)) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Equal(t, int64(1), OpsErrorLogQueueLength()) + job := <-opsErrorLogQueue + require.Equal(t, "account_auth", job.entry.ErrorPhase) + require.Equal(t, "provider", job.entry.ErrorOwner) + require.Equal(t, "gateway", job.entry.ErrorSource) + require.Contains(t, job.entry.ErrorMessage, "Recovered account authentication failure") + require.NotContains(t, job.entry.ErrorMessage, "403") + require.NotContains(t, job.entry.ErrorMessage, "earlier inference failure") + require.NotNil(t, job.entry.UpstreamStatusCode) + require.Zero(t, *job.entry.UpstreamStatusCode) + require.Len(t, job.entry.UpstreamErrors, 2) + require.Equal(t, http.StatusForbidden, job.entry.UpstreamErrors[0].UpstreamStatusCode) +} diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 231363a8a31..6db72c270b8 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -334,6 +334,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { var lastFailoverErr *service.UpstreamFailoverError for { + if c.Request.Context().Err() != nil { + return + } // Select account supporting the requested model reqLog.Debug("openai.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs))) selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability( @@ -443,11 +446,20 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { } else { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { + if c.Request.Context().Err() != nil { + return + } if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward { h.handleFailoverExhausted(c, failoverErr, true) return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if failoverErr.ShouldReportAccountScheduleFailure() { + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + } + if !failoverErr.ShouldRetryNextAccount() { + h.handleFailoverExhausted(c, failoverErr, streamStarted) + return + } // 池模式：同账号重试 if failoverErr.RetryableOnSameAccount { retryLimit := account.GetPoolModeRetryCount() @@ -835,6 +847,9 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { effectiveMappedModel := preferredMappedModel for { + if c.Request.Context().Err() != nil { + return + } currentRoutingModel := routingModel if effectiveMappedModel != "" { currentRoutingModel = effectiveMappedModel @@ -935,11 +950,20 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { } else { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { + if c.Request.Context().Err() != nil { + return + } if c.Writer.Size() != writerSizeBeforeForward { h.handleAnthropicFailoverExhausted(c, failoverErr, true) return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if failoverErr.ShouldReportAccountScheduleFailure() { + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + } + if !failoverErr.ShouldRetryNextAccount() { + h.handleAnthropicFailoverExhausted(c, failoverErr, streamStarted) + return + } // 池模式：同账号重试 if failoverErr.RetryableOnSameAccount { retryLimit := account.GetPoolModeRetryCount() @@ -1093,6 +1117,14 @@ func (h *OpenAIGatewayHandler) anthropicStreamingAwareError(c *gin.Context, stat // handleAnthropicFailoverExhausted maps upstream failover errors to Anthropic format. func (h *OpenAIGatewayHandler) handleAnthropicFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, streamStarted bool) { + if failoverErr != nil { + copyFailoverRetryAfter(c, failoverErr.ResponseHeaders) + } + if failoverErr != nil && failoverErr.IsCredentialFailure() { + status, message := credentialFailoverClientResponse(failoverErr) + h.anthropicStreamingAwareError(c, status, "api_error", message, streamStarted) + return + } status, errType, errMsg := h.mapUpstreamError(failoverErr.StatusCode) h.anthropicStreamingAwareError(c, status, errType, errMsg, streamStarted) } @@ -1464,8 +1496,49 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { switchCount := 0 failedAccountIDs := make(map[int64]struct{}) var lastFailoverErr *service.UpstreamFailoverError + handleWSFailover := func(account *service.Account, failoverErr *service.UpstreamFailoverError) bool { + if ctx.Err() != nil { + return false + } + if failoverErr.ShouldReportAccountScheduleFailure() { + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + } + releaseAccountSlot() + if !failoverErr.ShouldRetryNextAccount() { + closeOpenAIWSFailoverExhausted(wsConn, failoverErr) + return false + } + if ctx.Err() != nil { + return false + } + h.gatewayService.RecordOpenAIAccountSwitch() + failedAccountIDs[account.ID] = struct{}{} + lastFailoverErr = failoverErr + if switchCount >= maxAccountSwitches { + closeOpenAIWSFailoverExhausted(wsConn, failoverErr) + return false + } + switchCount++ + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + closeOpenAIWSFailoverExhausted(wsConn, failoverErr) + return false + } + reqLog.Warn("openai.websocket_upstream_failover_switching", + zap.Int64("account_id", account.ID), + zap.Int("upstream_status", failoverErr.StatusCode), + zap.Int("switch_count", switchCount), + zap.Int("max_switches", maxAccountSwitches), + ) + if ctx.Err() != nil { + return false + } + return ensureUserSlotHeld() + } for { + if ctx.Err() != nil { + return + } reqLog.Debug("openai.websocket_account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs))) selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability( ctx, @@ -1533,9 +1606,19 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { reqLog.Warn("openai.websocket_bind_sticky_session_failed", zap.Int64("account_id", account.ID), zap.Error(err)) } - token, _, err := h.gatewayService.GetAccessToken(ctx, account) + token, _, err := h.gatewayService.GetRequestCredential(ctx, c, account) if err != nil { reqLog.Warn("openai.websocket_get_access_token_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + if ctx.Err() != nil { + return + } + var failoverErr *service.UpstreamFailoverError + if errors.As(err, &failoverErr) { + if handleWSFailover(account, failoverErr) { + continue + } + return + } closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to get access token") return } @@ -1692,30 +1775,10 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { if err := h.gatewayService.ProxyResponsesWebSocketFromClient(ctx, c, wsConn, account, token, wsFirstMessage, hooks); err != nil { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) - releaseAccountSlot() - failedAccountIDs[account.ID] = struct{}{} - lastFailoverErr = failoverErr - if switchCount >= maxAccountSwitches { - closeOpenAIWSFailoverExhausted(wsConn, failoverErr) - return - } - switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { - closeOpenAIWSFailoverExhausted(wsConn, failoverErr) - return - } - h.gatewayService.RecordOpenAIAccountSwitch() - reqLog.Warn("openai.websocket_upstream_failover_switching", - zap.Int64("account_id", account.ID), - zap.Int("upstream_status", failoverErr.StatusCode), - zap.Int("switch_count", switchCount), - zap.Int("max_switches", maxAccountSwitches), - ) - if !ensureUserSlotHeld() { - return + if handleWSFailover(account, failoverErr) { + continue } - continue + return } if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) { @@ -1946,6 +2009,16 @@ func (h *OpenAIGatewayHandler) handleConcurrencyError(c *gin.Context, err error, } func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, streamStarted bool) { + if failoverErr == nil { + h.handleFailoverExhaustedSimple(c, http.StatusBadGateway, streamStarted) + return + } + copyFailoverRetryAfter(c, failoverErr.ResponseHeaders) + if failoverErr.IsCredentialFailure() { + status, message := credentialFailoverClientResponse(failoverErr) + h.handleStreamingAwareError(c, status, "upstream_error", message, streamStarted) + return + } statusCode := failoverErr.StatusCode responseBody := failoverErr.ResponseBody if service.IsOpenAISilentRefusalErrorBody(responseBody) { @@ -1987,6 +2060,41 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted) } +func credentialFailoverClientResponse(failoverErr *service.UpstreamFailoverError) (int, string) { + _ = failoverErr + return http.StatusServiceUnavailable, service.GrokCredentialUnavailableClientMessage +} + +func copyFailoverRetryAfter(c *gin.Context, headers http.Header) { + if c == nil || headers == nil { + return + } + retryAfter := strings.TrimSpace(headers.Get("Retry-After")) + if retryAfter == "" || len(retryAfter) > 128 || strings.ContainsAny(retryAfter, "\r\n") || !isSafeRetryAfter(retryAfter) { + return + } + c.Header("Retry-After", retryAfter) +} + +func isSafeRetryAfter(value string) bool { + digitsOnly := true + for _, char := range value { + if char < '0' || char > '9' { + digitsOnly = false + break + } + } + if digitsOnly { + seconds, err := strconv.ParseUint(value, 10, 32) + return err == nil && seconds <= uint64((7*24*time.Hour)/time.Second) + } + retryAt, err := http.ParseTime(value) + if err != nil { + return false + } + return !retryAt.After(time.Now().Add(7 * 24 * time.Hour)) +} + // handleFailoverExhaustedSimple 简化版本，用于没有响应体的情况 func (h *OpenAIGatewayHandler) handleFailoverExhaustedSimple(c *gin.Context, statusCode int, streamStarted bool) { status, errType, errMsg := h.mapUpstreamError(statusCode) @@ -2194,6 +2302,10 @@ func closeOpenAIWSFailoverExhausted(conn *coderws.Conn, failoverErr *service.Ups closeOpenAIClientWS(conn, coderws.StatusInternalError, "upstream websocket proxy failed") return } + if failoverErr.Stage == service.GatewayFailureStageAccountAuth { + closeOpenAIClientWS(conn, coderws.StatusTryAgainLater, service.GrokCredentialUnavailableClientMessage) + return + } switch failoverErr.StatusCode { case http.StatusTooManyRequests: closeOpenAIClientWS(conn, coderws.StatusTryAgainLater, "upstream rate limit exceeded, please retry later") diff --git a/backend/internal/handler/ops_error_logger.go b/backend/internal/handler/ops_error_logger.go index 64aa3ba4952..f32ca9c0543 100644 --- a/backend/internal/handler/ops_error_logger.go +++ b/backend/internal/handler/ops_error_logger.go @@ -721,10 +721,15 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { var upstreamStatusCode *int var upstreamErrorMessage *string var upstreamErrorDetail *string + finalAccountAuth := false if len(events) > 0 { last := events[len(events)-1] if last != nil { - if last.UpstreamStatusCode > 0 { + finalAccountAuth = last.Stage == string(service.GatewayFailureStageAccountAuth) + if finalAccountAuth { + code := 0 + upstreamStatusCode = &code + } else if last.UpstreamStatusCode > 0 { code := last.UpstreamStatusCode upstreamStatusCode = &code } @@ -737,7 +742,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { } } - if upstreamStatusCode == nil { + if !finalAccountAuth && upstreamStatusCode == nil { if v, ok := c.Get(service.OpsUpstreamStatusCodeKey); ok { switch t := v.(type) { case int: @@ -753,7 +758,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { } } } - if upstreamErrorMessage == nil { + if !finalAccountAuth && upstreamErrorMessage == nil { if v, ok := c.Get(service.OpsUpstreamErrorMessageKey); ok { if s, ok := v.(string); ok && strings.TrimSpace(s) != "" { msg := strings.TrimSpace(s) @@ -761,7 +766,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { } } } - if upstreamErrorDetail == nil { + if !finalAccountAuth && upstreamErrorDetail == nil { if v, ok := c.Get(service.OpsUpstreamErrorDetailKey); ok { if s, ok := v.(string); ok && strings.TrimSpace(s) != "" { detail := strings.TrimSpace(s) @@ -781,13 +786,18 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { } recoveredMsg := "Recovered upstream error" - if effectiveUpstreamStatus > 0 { + if finalAccountAuth { + recoveredMsg = "Recovered account authentication failure" + } else if effectiveUpstreamStatus > 0 { recoveredMsg += " " + strconvItoa(effectiveUpstreamStatus) } if upstreamErrorMessage != nil && strings.TrimSpace(*upstreamErrorMessage) != "" { recoveredMsg += ": " + strings.TrimSpace(*upstreamErrorMessage) } recoveredMsg = truncateString(recoveredMsg, 2048) + recoveredPhase, recoveredBusinessLimited, recoveredOwner, recoveredSource := classifyOpsErrorLog( + c, "upstream_error", recoveredMsg, "", effectiveUpstreamStatus, + ) entry := &service.OpsInsertErrorLogInput{ RequestID: requestID, @@ -828,19 +838,19 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { }(), UserAgent: c.GetHeader("User-Agent"), - ErrorPhase: "upstream", + ErrorPhase: recoveredPhase, ErrorType: "upstream_error", // Severity should reflect the upstream failure, not the final client status (200). Severity: classifyOpsSeverity("upstream_error", effectiveUpstreamStatus), StatusCode: status, - IsBusinessLimited: false, + IsBusinessLimited: recoveredBusinessLimited, IsCountTokens: isCountTokensRequest(c), ErrorMessage: recoveredMsg, ErrorBody: "", - ErrorSource: "upstream_http", - ErrorOwner: "provider", + ErrorSource: recoveredSource, + ErrorOwner: recoveredOwner, UpstreamStatusCode: upstreamStatusCode, UpstreamErrorMessage: upstreamErrorMessage, @@ -850,6 +860,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { CreatedAt: time.Now(), } applyOpsLatencyFieldsFromContext(c, entry) + applyOpsUpstreamFieldsFromContext(c, entry) if apiKey != nil { entry.APIKeyID = &apiKey.ID @@ -987,60 +998,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { CreatedAt: time.Now(), } applyOpsLatencyFieldsFromContext(c, entry) - - // Capture upstream error context set by gateway services (if present). - // This does NOT affect the client response; it enriches Ops troubleshooting data. - { - if v, ok := c.Get(service.OpsUpstreamStatusCodeKey); ok { - switch t := v.(type) { - case int: - if t > 0 { - code := t - entry.UpstreamStatusCode = &code - } - case int64: - if t > 0 { - code := int(t) - entry.UpstreamStatusCode = &code - } - } - } - if v, ok := c.Get(service.OpsUpstreamErrorMessageKey); ok { - if s, ok := v.(string); ok { - if msg := strings.TrimSpace(s); msg != "" { - entry.UpstreamErrorMessage = &msg - } - } - } - if v, ok := c.Get(service.OpsUpstreamErrorDetailKey); ok { - if s, ok := v.(string); ok { - if detail := strings.TrimSpace(s); detail != "" { - entry.UpstreamErrorDetail = &detail - } - } - } - if v, ok := c.Get(service.OpsUpstreamErrorsKey); ok { - if events, ok := v.([]*service.OpsUpstreamErrorEvent); ok && len(events) > 0 { - entry.UpstreamErrors = events - // Best-effort backfill the single upstream fields from the last event when missing. - last := events[len(events)-1] - if last != nil { - if entry.UpstreamStatusCode == nil && last.UpstreamStatusCode > 0 { - code := last.UpstreamStatusCode - entry.UpstreamStatusCode = &code - } - if entry.UpstreamErrorMessage == nil && strings.TrimSpace(last.Message) != "" { - msg := strings.TrimSpace(last.Message) - entry.UpstreamErrorMessage = &msg - } - if entry.UpstreamErrorDetail == nil && strings.TrimSpace(last.Detail) != "" { - detail := strings.TrimSpace(last.Detail) - entry.UpstreamErrorDetail = &detail - } - } - } - } - } + applyOpsUpstreamFieldsFromContext(c, entry) if apiKey != nil { entry.APIKeyID = &apiKey.ID @@ -1235,6 +1193,77 @@ func applyOpsLatencyFieldsFromContext(c *gin.Context, entry *service.OpsInsertEr entry.TimeToFirstTokenMs = getContextLatencyMs(c, service.OpsTimeToFirstTokenMsKey) } +// applyOpsUpstreamFieldsFromContext captures attempt-level upstream context. +// A final account_auth event owns the top-level status and forces it to zero; +// prior inference statuses remain available in UpstreamErrors. +func applyOpsUpstreamFieldsFromContext(c *gin.Context, entry *service.OpsInsertErrorLogInput) { + if c == nil || entry == nil { + return + } + if v, ok := c.Get(service.OpsUpstreamStatusCodeKey); ok { + switch t := v.(type) { + case int: + if t > 0 { + code := t + entry.UpstreamStatusCode = &code + } + case int64: + if t > 0 { + code := int(t) + entry.UpstreamStatusCode = &code + } + } + } + if v, ok := c.Get(service.OpsUpstreamErrorMessageKey); ok { + if value, ok := v.(string); ok { + if message := strings.TrimSpace(value); message != "" { + entry.UpstreamErrorMessage = &message + } + } + } + if v, ok := c.Get(service.OpsUpstreamErrorDetailKey); ok { + if value, ok := v.(string); ok { + if detail := strings.TrimSpace(value); detail != "" { + entry.UpstreamErrorDetail = &detail + } + } + } + if v, ok := c.Get(service.OpsUpstreamErrorsKey); ok { + if events, ok := v.([]*service.OpsUpstreamErrorEvent); ok && len(events) > 0 { + entry.UpstreamErrors = events + last := events[len(events)-1] + if last == nil { + return + } + if last.Stage == string(service.GatewayFailureStageAccountAuth) { + code := 0 + entry.UpstreamStatusCode = &code + entry.UpstreamErrorMessage = nil + if message := strings.TrimSpace(last.Message); message != "" { + entry.UpstreamErrorMessage = &message + } + entry.UpstreamErrorDetail = nil + if detail := strings.TrimSpace(last.Detail); detail != "" { + entry.UpstreamErrorDetail = &detail + } + } else { + if entry.UpstreamStatusCode == nil && last.UpstreamStatusCode > 0 { + code := last.UpstreamStatusCode + entry.UpstreamStatusCode = &code + } + if entry.UpstreamErrorMessage == nil && strings.TrimSpace(last.Message) != "" { + message := strings.TrimSpace(last.Message) + entry.UpstreamErrorMessage = &message + } + if entry.UpstreamErrorDetail == nil && strings.TrimSpace(last.Detail) != "" { + detail := strings.TrimSpace(last.Detail) + entry.UpstreamErrorDetail = &detail + } + } + } + } +} + func getContextLatencyMs(c *gin.Context, key string) *int64 { if c == nil || strings.TrimSpace(key) == "" { return nil @@ -1390,7 +1419,7 @@ func normalizeOpsErrorType(errType string, code string) string { func classifyOpsPhase(errType, message, code string) string { msg := strings.ToLower(message) - // Standardized phases: request|auth|routing|upstream|network|internal + // Standardized phases: request|auth|account_auth|routing|upstream|network|internal // Map billing/concurrency/response => request; scheduling => routing. if isOpsClientAuthError(code, msg) { return "auth" @@ -1445,7 +1474,10 @@ func classifyOpsErrorLog(c *gin.Context, errType, message, code string, status i routingCapacityLimited := isOpsRoutingCapacityLimited(c) clientBusinessLimited := service.HasOpsClientBusinessLimited(c) upstreamError := hasOpsUpstreamErrorContext(c) - if upstreamError && !routingCapacityLimited { + accountAuthFailure := hasOpsAccountAuthFailure(c) + if accountAuthFailure && !routingCapacityLimited { + phase = "account_auth" + } else if upstreamError && !routingCapacityLimited { phase = "upstream" } if clientBusinessLimited && !upstreamError && !routingCapacityLimited { @@ -1570,6 +1602,22 @@ func hasOpsUpstreamErrorContext(c *gin.Context) bool { return false } +func hasOpsAccountAuthFailure(c *gin.Context) bool { + if c == nil { + return false + } + if v, ok := c.Get(service.OpsUpstreamErrorsKey); ok { + if events, ok := v.([]*service.OpsUpstreamErrorEvent); ok { + for i := len(events) - 1; i >= 0; i-- { + if events[i] != nil { + return events[i].Stage == string(service.GatewayFailureStageAccountAuth) + } + } + } + } + return false +} + func isOpsNoAvailableAccountMessage(message string) bool { msg := strings.ToLower(message) return strings.Contains(msg, opsErrNoAvailableAccounts) || @@ -1584,6 +1632,8 @@ func classifyOpsErrorOwner(phase string, message string) string { switch phase { case "upstream", "network": return "provider" + case "account_auth": + return "provider" case "request", "auth": return "client" case "routing", "internal": @@ -1601,6 +1651,8 @@ func classifyOpsErrorSource(phase string, message string) string { switch phase { case "upstream": return "upstream_http" + case "account_auth": + return "gateway" case "network": return "gateway" case "request", "auth": diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 8eb819aeabc..a3eba2eaa94 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -859,6 +859,54 @@ func (r *accountRepository) SetError(ctx context.Context, id int64, errorMsg str return nil } +func (r *accountRepository) SetGrokCredentialErrorIfMatch( + ctx context.Context, + id int64, + snapshot service.GrokCredentialMutationSnapshot, + errorMsg string, +) (bool, error) { + result, err := r.sql.ExecContext(ctx, ` + WITH updated AS ( + UPDATE accounts AS a + SET status = $1, + error_message = $2, + schedulable = false, + updated_at = NOW() + WHERE a.id = $3 + AND a.deleted_at IS NULL + AND a.status = $4 + AND a.platform = $5 + AND a.type = $6 + AND a.schedulable IS TRUE + AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW()) + AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW()) + AND (a.overload_until IS NULL OR a.overload_until <= NOW()) + AND (a.auto_pause_on_expired IS NOT TRUE OR a.expires_at IS NULL OR a.expires_at > NOW()) + AND a.credentials = $7::jsonb + AND a.proxy_id IS NOT DISTINCT FROM $8 + AND ($2 <> $9 OR ( + a.proxy_id IS NOT NULL AND NOT EXISTS ( + SELECT 1 FROM proxies p WHERE p.id = a.proxy_id AND p.deleted_at IS NULL + ) + )) + RETURNING a.id + ) + INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload) + SELECT $10, updated.id, NULL, NULL FROM updated + `, service.StatusError, errorMsg, id, service.StatusActive, service.PlatformGrok, service.AccountTypeOAuth, + snapshot.CredentialsJSON, snapshot.ProxyID, string(service.GrokCredentialReasonProxyInvalid), + service.SchedulerOutboxEventAccountChanged) + if err != nil { + return false, err + } + affected, err := result.RowsAffected() + if err != nil || affected == 0 { + return false, err + } + r.syncSchedulerAccountSnapshotDetached(ctx, id) + return true, nil +} + // syncSchedulerAccountSnapshot 在账号状态变更时主动同步快照到调度器缓存。 // 当账号被设置为错误、禁用、不可调度或临时不可调度时调用， // 确保调度器和粘性会话逻辑能及时感知账号的最新状态，避免继续使用不可用账号。 @@ -881,6 +929,16 @@ func (r *accountRepository) syncSchedulerAccountSnapshot(ctx context.Context, ac } } +func (r *accountRepository) syncSchedulerAccountSnapshotDetached(ctx context.Context, accountID int64) { + base := context.Background() + if ctx != nil { + base = context.WithoutCancel(ctx) + } + propagationCtx, cancel := context.WithTimeout(base, 2*time.Second) + defer cancel() + r.syncSchedulerAccountSnapshot(propagationCtx, accountID) +} + func (r *accountRepository) deleteSchedulerAccountSnapshot(ctx context.Context, accountID int64) { if r == nil || r.schedulerCache == nil || accountID <= 0 { return @@ -1414,6 +1472,51 @@ func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64, return nil } +func (r *accountRepository) SetGrokCredentialTempUnschedulableIfMatch( + ctx context.Context, + id int64, + snapshot service.GrokCredentialMutationSnapshot, + until time.Time, + reason string, +) (bool, error) { + result, err := r.sql.ExecContext(ctx, ` + WITH updated AS ( + UPDATE accounts AS a + SET temp_unschedulable_until = CASE + WHEN a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until < $1 THEN $1 + ELSE a.temp_unschedulable_until + END, + temp_unschedulable_reason = $2, + updated_at = NOW() + WHERE a.id = $3 + AND a.deleted_at IS NULL + AND a.status = $4 + AND a.platform = $5 + AND a.type = $6 + AND a.schedulable IS TRUE + AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW()) + AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW()) + AND (a.overload_until IS NULL OR a.overload_until <= NOW()) + AND (a.auto_pause_on_expired IS NOT TRUE OR a.expires_at IS NULL OR a.expires_at > NOW()) + AND a.credentials = $7::jsonb + AND a.proxy_id IS NOT DISTINCT FROM $8 + RETURNING a.id + ) + INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload) + SELECT $9, updated.id, NULL, NULL FROM updated + `, until, reason, id, service.StatusActive, service.PlatformGrok, service.AccountTypeOAuth, + snapshot.CredentialsJSON, snapshot.ProxyID, service.SchedulerOutboxEventAccountChanged) + if err != nil { + return false, err + } + affected, err := result.RowsAffected() + if err != nil || affected == 0 { + return false, err + } + r.syncSchedulerAccountSnapshotDetached(ctx, id) + return true, nil +} + func (r *accountRepository) ClearTempUnschedulable(ctx context.Context, id int64) error { _, err := r.sql.ExecContext(ctx, ` UPDATE accounts diff --git a/backend/internal/repository/account_repo_temp_unsched_test.go b/backend/internal/repository/account_repo_temp_unsched_test.go index 9c3e4cbdf51..2ea18bd2140 100644 --- a/backend/internal/repository/account_repo_temp_unsched_test.go +++ b/backend/internal/repository/account_repo_temp_unsched_test.go @@ -9,6 +9,7 @@ import ( "time" sqlmock "github.com/DATA-DOG/go-sqlmock" + "github.com/Wei-Shaw/sub2api/internal/service" "github.com/stretchr/testify/require" ) @@ -24,6 +25,103 @@ func TestAccountRepository_SetTempUnschedulable_NoRowsAffectedDoesNotWriteOutbox require.NotContains(t, strings.Join(exec.execQueries, "\n"), "scheduler_outbox") } +func TestAccountRepository_GrokCredentialConditionalMutationsAreEligibleAndAtomicallyPropagated(t *testing.T) { + proxyID := int64(77) + snapshot := service.GrokCredentialMutationSnapshot{ + CredentialsJSON: `{"access_token":"access","refresh_token":"refresh","_token_version":123}`, + ProxyID: &proxyID, + } + + t.Run("permanent", func(t *testing.T) { + exec := &recordingSQLExecutor{result: rowsAffectedResult(0)} + repo := newAccountRepositoryWithSQL(nil, exec, nil) + + updated, err := repo.SetGrokCredentialErrorIfMatch(context.Background(), 42, snapshot, "revoked") + + require.NoError(t, err) + require.False(t, updated) + require.Len(t, exec.execQueries, 1) + normalized := normalizeSQLWhitespace(exec.execQueries[0]) + require.Contains(t, normalized, "WITH updated AS ( UPDATE accounts AS a") + require.Contains(t, normalized, "a.schedulable IS TRUE") + require.Contains(t, normalized, "a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW()") + require.Contains(t, normalized, "a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW()") + require.Contains(t, normalized, "a.overload_until IS NULL OR a.overload_until <= NOW()") + require.Contains(t, normalized, "a.credentials = $7::jsonb") + require.Contains(t, normalized, "a.proxy_id IS NOT DISTINCT FROM $8") + require.Contains(t, normalized, "NOT EXISTS ( SELECT 1 FROM proxies p") + require.Contains(t, normalized, "INSERT INTO scheduler_outbox") + require.Len(t, exec.execArgs[0], 10) + require.Equal(t, snapshot.CredentialsJSON, exec.execArgs[0][6]) + require.Equal(t, &proxyID, exec.execArgs[0][7]) + require.Equal(t, string(service.GrokCredentialReasonProxyInvalid), exec.execArgs[0][8]) + require.Equal(t, service.SchedulerOutboxEventAccountChanged, exec.execArgs[0][9]) + }) + + t.Run("transient", func(t *testing.T) { + exec := &recordingSQLExecutor{result: rowsAffectedResult(0)} + repo := newAccountRepositoryWithSQL(nil, exec, nil) + + updated, err := repo.SetGrokCredentialTempUnschedulableIfMatch( + context.Background(), 42, snapshot, time.Now().Add(time.Minute), "temporary", + ) + + require.NoError(t, err) + require.False(t, updated) + require.Len(t, exec.execQueries, 1) + normalized := normalizeSQLWhitespace(exec.execQueries[0]) + require.Contains(t, normalized, "WITH updated AS ( UPDATE accounts AS a") + require.Contains(t, normalized, "a.schedulable IS TRUE") + require.Contains(t, normalized, "a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW()") + require.Contains(t, normalized, "a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW()") + require.Contains(t, normalized, "a.overload_until IS NULL OR a.overload_until <= NOW()") + require.Contains(t, normalized, "a.credentials = $7::jsonb") + require.Contains(t, normalized, "a.proxy_id IS NOT DISTINCT FROM $8") + require.Contains(t, normalized, "INSERT INTO scheduler_outbox") + require.Len(t, exec.execArgs[0], 9) + require.Equal(t, snapshot.CredentialsJSON, exec.execArgs[0][6]) + require.Equal(t, &proxyID, exec.execArgs[0][7]) + require.Equal(t, service.SchedulerOutboxEventAccountChanged, exec.execArgs[0][8]) + }) +} + +func TestAccountRepository_GrokCredentialCommitCarriesOutboxAcrossCallerCancellation(t *testing.T) { + snapshot := service.GrokCredentialMutationSnapshot{CredentialsJSON: `{"access_token":"access","refresh_token":"refresh"}`} + tests := []struct { + name string + mutate func(context.Context, *accountRepository) (bool, error) + }{ + { + name: "permanent", + mutate: func(ctx context.Context, repo *accountRepository) (bool, error) { + return repo.SetGrokCredentialErrorIfMatch(ctx, 42, snapshot, string(service.GrokCredentialReasonRevoked)) + }, + }, + { + name: "transient", + mutate: func(ctx context.Context, repo *accountRepository) (bool, error) { + return repo.SetGrokCredentialTempUnschedulableIfMatch(ctx, 42, snapshot, time.Now().Add(time.Minute), string(service.GrokCredentialReasonRefreshTransient)) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + exec := &recordingSQLExecutor{result: rowsAffectedResult(1), afterExec: cancel} + repo := newAccountRepositoryWithSQL(nil, exec, nil) + + updated, err := tt.mutate(ctx, repo) + + require.NoError(t, err) + require.True(t, updated) + require.ErrorIs(t, ctx.Err(), context.Canceled) + require.Len(t, exec.execQueries, 1, "state update and scheduler outbox must share one atomic SQL statement") + require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "INSERT INTO scheduler_outbox") + }) + } +} + func TestAccountRepository_ListOAuthRefreshCandidates_SQLFilter(t *testing.T) { db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) require.NoError(t, err) @@ -87,14 +185,20 @@ func (r rowsAffectedResult) RowsAffected() (int64, error) { return int64(r), nil type recordingSQLExecutor struct { result sql.Result err error + afterExec func() execQueries []string + execArgs [][]any } func (e *recordingSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { e.execQueries = append(e.execQueries, query) + e.execArgs = append(e.execArgs, append([]any(nil), args...)) if e.err != nil { return nil, e.err } + if e.afterExec != nil { + e.afterExec() + } return e.result, nil } diff --git a/backend/internal/repository/grok_oauth_client.go b/backend/internal/repository/grok_oauth_client.go index 6c9c2c407f1..38f6cfb96ed 100644 --- a/backend/internal/repository/grok_oauth_client.go +++ b/backend/internal/repository/grok_oauth_client.go @@ -153,14 +153,29 @@ func grokOAuthStatusError(code, message string, resp *req.Response) error { statusCode := http.StatusBadGateway errorCode := code upstreamStatus := 0 - if resp != nil && resp.StatusCode == http.StatusForbidden { - statusCode = http.StatusForbidden - errorCode = "GROK_OAUTH_ENTITLEMENT_DENIED" - } body := "" if resp != nil { upstreamStatus = resp.StatusCode body = logredact.RedactText(resp.String()) + if resp.StatusCode == http.StatusForbidden && grokOAuthHasExplicitEntitlementDenial(body) { + statusCode = http.StatusForbidden + errorCode = "GROK_OAUTH_ENTITLEMENT_DENIED" + } } return infraerrors.Newf(statusCode, errorCode, "%s: status %d, body: %s", message, upstreamStatus, body) } + +func grokOAuthHasExplicitEntitlementDenial(body string) bool { + lower := strings.ToLower(body) + compact := strings.NewReplacer(" ", "", "\n", "", "\r", "", "\t", "").Replace(lower) + for _, field := range []string{"error", "code", "reason"} { + for _, value := range []string{"access_denied", "entitlement_denied", "subscription_required", "no_active_subscription"} { + if strings.Contains(compact, `"`+field+`":"`+value+`"`) { + return true + } + } + } + return strings.Contains(lower, "entitlement denied") || + strings.Contains(lower, "subscription required") || + strings.Contains(lower, "no active grok subscription") +} diff --git a/backend/internal/repository/grok_oauth_client_test.go b/backend/internal/repository/grok_oauth_client_test.go index eabeb641c77..690a8e21390 100644 --- a/backend/internal/repository/grok_oauth_client_test.go +++ b/backend/internal/repository/grok_oauth_client_test.go @@ -50,15 +50,7 @@ func TestGrokOAuthClientExchangeAndRefreshUseFormFields(t *testing.T) { t.Setenv(xai.EnvTokenURL, server.URL) client := NewGrokOAuthClient() - - exchanged, err := client.ExchangeCode( - context.Background(), - "auth-code", - "verifier", - "http://127.0.0.1:56121/callback", - "", - "client-id", - ) + exchanged, err := client.ExchangeCode(context.Background(), "auth-code", "verifier", "http://127.0.0.1:56121/callback", "", "client-id") require.NoError(t, err) require.Equal(t, "exchange-access", exchanged.AccessToken) require.Equal(t, "exchange-refresh", exchanged.RefreshToken) @@ -72,18 +64,30 @@ func TestGrokOAuthClientExchangeAndRefreshUseFormFields(t *testing.T) { require.Equal(t, int64(7200), refreshed.ExpiresIn) } -func TestGrokOAuthClientRefreshForbiddenClassifiesEntitlement(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusForbidden) - _, _ = w.Write([]byte(`{"error":"subscription required"}`)) - })) - defer server.Close() - t.Setenv(xai.EnvTokenURL, server.URL) +func TestGrokOAuthClientRefreshForbiddenClassifiesOnlyExplicitEntitlement(t *testing.T) { + tests := []struct { + name string + body string + wantReason string + }{ + {name: "explicit entitlement", body: `{"error":"access_denied"}`, wantReason: "GROK_OAUTH_ENTITLEMENT_DENIED"}, + {name: "generic forbidden", body: `{"error":"forbidden"}`, wantReason: "GROK_OAUTH_TOKEN_REFRESH_FAILED"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(tt.body)) + })) + defer server.Close() + t.Setenv(xai.EnvTokenURL, server.URL) - client := NewGrokOAuthClient() - _, err := client.RefreshToken(context.Background(), "refresh-token", "", "client-id") - require.Error(t, err) - require.Contains(t, strings.ToUpper(err.Error()), "GROK_OAUTH_ENTITLEMENT_DENIED") + client := NewGrokOAuthClient() + _, err := client.RefreshToken(context.Background(), "refresh-token", "", "client-id") + require.Error(t, err) + require.Contains(t, strings.ToUpper(err.Error()), tt.wantReason) + }) + } } func TestGrokOAuthClientStatusErrorRedactsSensitiveResponseBody(t *testing.T) { @@ -105,3 +109,13 @@ func TestGrokOAuthClientStatusErrorRedactsSensitiveResponseBody(t *testing.T) { require.NotContains(t, errText, "refresh-secret") require.NotContains(t, errText, "verifier-secret") } + +func TestGrokOAuthEntitlementDenialRequiresExplicitEvidence(t *testing.T) { + t.Parallel() + + require.True(t, grokOAuthHasExplicitEntitlementDenial(`{"error":"access_denied"}`)) + require.True(t, grokOAuthHasExplicitEntitlementDenial(`{"code":"entitlement_denied"}`)) + require.True(t, grokOAuthHasExplicitEntitlementDenial(`{"message":"no active Grok subscription"}`)) + require.False(t, grokOAuthHasExplicitEntitlementDenial(`{"error":"forbidden","message":"request forbidden"}`)) + require.False(t, grokOAuthHasExplicitEntitlementDenial(`403 Forbidden`)) +} diff --git a/backend/internal/repository/ops_error_where_test.go b/backend/internal/repository/ops_error_where_test.go index c997865ba4a..c52461e8ff0 100644 --- a/backend/internal/repository/ops_error_where_test.go +++ b/backend/internal/repository/ops_error_where_test.go @@ -101,6 +101,40 @@ func TestBuildOpsErrorLogsWhere_CyberPolicyStatusExemption(t *testing.T) { if strings.Contains(whereRecovered, "status_code") { t.Fatalf("upstream phase with IncludeRecoveredUpstream must not add any status_code clause\nfull: %s", whereRecovered) } + + // account_auth uses the same explicit provider-health opt-in but remains a + // distinct phase from inference upstream errors. + whereAccountAuth, _ := buildOpsErrorLogsWhere(&service.OpsErrorLogFilter{Phase: "account_auth", IncludeRecoveredUpstream: true}) + if strings.Contains(whereAccountAuth, "status_code") { + t.Fatalf("account_auth phase with IncludeRecoveredUpstream must expose recovered rows\nfull: %s", whereAccountAuth) + } + if !strings.Contains(whereAccountAuth, "e.error_phase = $") { + t.Fatalf("account_auth recovered filter must retain its explicit phase\nfull: %s", whereAccountAuth) + } + + whereProviderHealth, _ := buildOpsErrorLogsWhere(&service.OpsErrorLogFilter{ + ErrorPhasesAny: []string{"upstream", "account_auth"}, + IncludeRecoveredUpstream: true, + }) + if strings.Contains(whereProviderHealth, "status_code") { + t.Fatalf("provider-health ANY filter must expose recovered inference and credential rows\nfull: %s", whereProviderHealth) + } + if !strings.Contains(whereProviderHealth, "e.error_phase = ANY($") { + t.Fatalf("provider-health filter must preserve distinct phase values\nfull: %s", whereProviderHealth) + } + + whereUserAccountAuth, _ := buildOpsErrorLogsWhere(&service.OpsErrorLogFilter{ErrorPhasesAny: []string{"account_auth"}}) + if !strings.Contains(whereUserAccountAuth, "COALESCE(e.status_code, 0) >= 400") { + t.Fatalf("request-error account_auth filters must exclude recovered successes\nfull: %s", whereUserAccountAuth) + } + + whereMixed, _ := buildOpsErrorLogsWhere(&service.OpsErrorLogFilter{ + ErrorPhasesAny: []string{"account_auth", "request"}, + IncludeRecoveredUpstream: true, + }) + if !strings.Contains(whereMixed, "COALESCE(e.status_code, 0) >= 400") { + t.Fatalf("recovered opt-in must not bypass the guard for non-provider phases\nfull: %s", whereMixed) + } } func TestBuildOpsErrorLogsWhere_MatchDeletedKeyOwner(t *testing.T) { diff --git a/backend/internal/repository/ops_repo.go b/backend/internal/repository/ops_repo.go index 900abcf212f..635bcd1a1f3 100644 --- a/backend/internal/repository/ops_repo.go +++ b/backend/internal/repository/ops_repo.go @@ -160,7 +160,7 @@ func opsInsertErrorLogArgs(input *service.OpsInsertErrorLogInput) []any { opsNullString(input.ErrorBody), opsNullString(input.ErrorSource), opsNullString(input.ErrorOwner), - opsNullInt(input.UpstreamStatusCode), + opsNullableIntPointer(input.UpstreamStatusCode), opsNullString(input.UpstreamErrorMessage), opsNullString(input.UpstreamErrorDetail), opsNullString(input.UpstreamErrorsJSON), @@ -582,7 +582,7 @@ LIMIT 1` s := clientIP.String out.ClientIP = &s } - if upstreamStatusCode.Valid && upstreamStatusCode.Int64 > 0 { + if upstreamStatusCode.Valid { v := int(upstreamStatusCode.Int64) out.UpstreamStatusCode = &v } @@ -978,13 +978,13 @@ func buildOpsErrorLogsWhere(filter *service.OpsErrorLogFilter) (string, []any) { resolvedFilter = filter.Resolved } // Keep list endpoints scoped to client errors unless the caller explicitly opts - // into recovered upstream rows (Phase=="upstream" + IncludeRecoveredUpstream, - // ops 专用上游列表)。请求错误语义的端点即便过滤 phase=upstream 也保留该守卫。 + // into recovered provider-health rows (upstream/account_auth). Request-error + // endpoints never set the opt-in and retain this guard. // cyber_policy is exempt from the status >= 400 guard: streaming cyber hits arrive with // status 200 (the SSE stream opened successfully before upstream returned response.failed), // but they are always client-visible blocked requests that belong in admin + user error // lists. Without the exemption the entire streaming-path cyber sink would be invisible. - if phaseFilter != "upstream" || filter == nil || !filter.IncludeRecoveredUpstream { + if !opsFilterIncludesRecoveredProviderRows(filter, phaseFilter) { clauses = append(clauses, "(COALESCE(e.status_code, 0) >= 400 OR e.error_type = 'cyber_policy')") } @@ -1117,6 +1117,28 @@ func buildOpsErrorLogsWhere(filter *service.OpsErrorLogFilter) (string, []any) { return "WHERE " + strings.Join(clauses, " AND "), args } +func opsFilterIncludesRecoveredProviderRows(filter *service.OpsErrorLogFilter, phaseFilter string) bool { + if filter == nil || !filter.IncludeRecoveredUpstream { + return false + } + if phaseFilter != "" { + return phaseFilter == "upstream" || phaseFilter == "account_auth" + } + if len(filter.ErrorPhasesAny) == 0 { + return false + } + sawProviderPhase := false + for _, rawPhase := range filter.ErrorPhasesAny { + switch strings.TrimSpace(strings.ToLower(rawPhase)) { + case "upstream", "account_auth": + sawProviderPhase = true + default: + return false + } + } + return sawProviderPhase +} + func buildOpsSystemLogsWhere(filter *service.OpsSystemLogFilter) (string, []any, bool) { clauses := make([]string, 0, 10) args := make([]any, 0, 10) @@ -1269,6 +1291,16 @@ func opsNullInt(v any) any { } } +// opsNullableIntPointer distinguishes an absent value from an explicitly +// observed zero. Credential-stage failures intentionally persist upstream +// status 0 because no inference request was sent. +func opsNullableIntPointer(v *int) any { + if v == nil { + return sql.NullInt64{} + } + return sql.NullInt64{Int64: int64(*v), Valid: true} +} + func opsNullInt16(v *int16) any { if v == nil { return sql.NullInt64{} diff --git a/backend/internal/repository/ops_repo_args_test.go b/backend/internal/repository/ops_repo_args_test.go new file mode 100644 index 00000000000..5a88758f0ad --- /dev/null +++ b/backend/internal/repository/ops_repo_args_test.go @@ -0,0 +1,37 @@ +//go:build unit + +package repository + +import ( + "database/sql" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestOpsInsertErrorLogArgsPreservesExplicitZeroUpstreamStatus(t *testing.T) { + zero := 0 + args := opsInsertErrorLogArgs(&service.OpsInsertErrorLogInput{UpstreamStatusCode: &zero}) + + require.Len(t, args, 41) + encoded, ok := args[27].(sql.NullInt64) + require.True(t, ok) + require.True(t, encoded.Valid) + require.Zero(t, encoded.Int64) +} + +func TestOpsNullableIntPointerDistinguishesNilZeroAndStatus(t *testing.T) { + missing := opsNullableIntPointer(nil).(sql.NullInt64) + require.False(t, missing.Valid) + + zeroValue := 0 + zero := opsNullableIntPointer(&zeroValue).(sql.NullInt64) + require.True(t, zero.Valid) + require.Zero(t, zero.Int64) + + statusValue := 503 + status := opsNullableIntPointer(&statusValue).(sql.NullInt64) + require.True(t, status.Valid) + require.EqualValues(t, 503, status.Int64) +} diff --git a/backend/internal/repository/ops_repo_get_error_log_by_id_integration_test.go b/backend/internal/repository/ops_repo_get_error_log_by_id_integration_test.go index 470b1c0d0bf..1411002be0c 100644 --- a/backend/internal/repository/ops_repo_get_error_log_by_id_integration_test.go +++ b/backend/internal/repository/ops_repo_get_error_log_by_id_integration_test.go @@ -91,4 +91,21 @@ func TestGetErrorLogByID_DeletedKeyOwner(t *testing.T) { require.Equal(t, "sk-valid", valid.APIKeyPrefix) require.Empty(t, valid.AttemptedKeyPrefix, "attempted prefix and api key prefix are mutually exclusive") require.Nil(t, valid.DeletedKeyOwnerUserID, "valid key error has no deleted owner") + + // ── Case 4: account_auth with no inference attempt preserves explicit 0 ── + zero := 0 + credentialFailureID, err := repo.InsertErrorLog(ctx, &service.OpsInsertErrorLogInput{ + ErrorPhase: "account_auth", + ErrorType: "upstream_error", + Severity: "error", + StatusCode: 503, + UpstreamStatusCode: &zero, + CreatedAt: time.Now(), + }) + require.NoError(t, err) + + credentialFailure, err := repo.GetErrorLogByID(ctx, credentialFailureID) + require.NoError(t, err) + require.NotNil(t, credentialFailure.UpstreamStatusCode) + require.Zero(t, *credentialFailure.UpstreamStatusCode) } diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 3e67fa59828..539c92741dc 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -1268,16 +1268,13 @@ func (a *Account) GetGrokBaseURL() string { if !a.IsGrok() { return "" } - baseURL := a.GetCredential("base_url") if a.IsGrokOAuth() { - if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) { - return xai.DefaultCLIBaseURL - } - if _, err := xai.ValidateTrustedBaseURL(baseURL); err == nil { - return baseURL - } + // OAuth bearer credentials are subscription credentials and may only be + // sent to the supported CLI gateway. Stored base_url values and unsafe + // development overrides apply exclusively to API-key accounts. return xai.DefaultCLIBaseURL } + baseURL := a.GetCredential("base_url") if baseURL != "" { return baseURL } @@ -1286,27 +1283,17 @@ func (a *Account) GetGrokBaseURL() string { // GetGrokMediaBaseURL selects the upstream used by Grok Imagine APIs. // -// OAuth text requests need the CLI subscription proxy, but that proxy has a -// smaller request-body limit than the official Imagine API. Media requests can -// contain large base64 inputs, so default OAuth accounts must use api.x.ai. -// API-key accounts and explicit unsafe development overrides retain their -// configured base URL. +// OAuth media credentials have the same trust boundary as OAuth text traffic: +// they are pinned to the supported CLI gateway even for large request bodies. +// API-key accounts retain their configured public/custom upstream behavior. func (a *Account) GetGrokMediaBaseURL() string { if !a.IsGrok() { return "" } - if !a.IsGrokOAuth() { - return a.GetGrokBaseURL() - } - - baseURL := a.GetCredential("base_url") - if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) || isOfficialGrokCLIBaseURL(baseURL) { - return xai.DefaultBaseURL - } - if _, err := xai.ValidateTrustedBaseURL(baseURL); err == nil { - return baseURL + if a.IsGrokOAuth() { + return xai.DefaultCLIBaseURL } - return xai.DefaultBaseURL + return a.GetGrokBaseURL() } func isOfficialGrokAPIBaseURL(raw string) bool { diff --git a/backend/internal/service/account_base_url_test.go b/backend/internal/service/account_base_url_test.go index 59f53db3db0..701a37e1313 100644 --- a/backend/internal/service/account_base_url_test.go +++ b/backend/internal/service/account_base_url_test.go @@ -255,7 +255,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) { expected: xai.DefaultCLIBaseURL, }, { - name: "oauth non-default API port remains an explicit override", + name: "oauth non-default API port remains pinned to CLI proxy", account: Account{ Type: AccountTypeOAuth, Platform: PlatformGrok, @@ -263,7 +263,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) { "base_url": "https://api.x.ai:8443/v1", }, }, - expected: "https://api.x.ai:8443/v1", + expected: xai.DefaultCLIBaseURL, }, { name: "oauth explicit custom base_url stays pinned to CLI proxy by default", @@ -294,7 +294,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) { } } -func TestGetGrokBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *testing.T) { +func TestGetGrokBaseURLPinsOAuthWhenUnsafeOverridesEnabled(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") account := Account{ Type: AccountTypeOAuth, @@ -304,26 +304,26 @@ func TestGetGrokBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t * }, } - require.Equal(t, "https://custom.example.com/v1", account.GetGrokBaseURL()) + require.Equal(t, xai.DefaultCLIBaseURL, account.GetGrokBaseURL()) } -func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) { +func TestGetGrokMediaBaseURLPinsOAuthMediaToCLIProxy(t *testing.T) { tests := []struct { name string account Account expected string }{ { - name: "oauth without base_url uses official media API", + name: "oauth without base_url uses CLI subscription proxy", account: Account{ Type: AccountTypeOAuth, Platform: PlatformGrok, Credentials: map[string]any{}, }, - expected: xai.DefaultBaseURL, + expected: xai.DefaultCLIBaseURL, }, { - name: "oauth stored CLI proxy uses official media API", + name: "oauth stored CLI proxy stays on CLI subscription proxy", account: Account{ Type: AccountTypeOAuth, Platform: PlatformGrok, @@ -331,10 +331,10 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) { "base_url": xai.DefaultCLIBaseURL, }, }, - expected: xai.DefaultBaseURL, + expected: xai.DefaultCLIBaseURL, }, { - name: "oauth stored CLI proxy variant uses official media API", + name: "oauth stored CLI proxy variant is canonicalized to CLI proxy", account: Account{ Type: AccountTypeOAuth, Platform: PlatformGrok, @@ -342,10 +342,10 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) { "base_url": "HTTPS://CLI-CHAT-PROXY.GROK.COM:443/%76%31/", }, }, - expected: xai.DefaultBaseURL, + expected: xai.DefaultCLIBaseURL, }, { - name: "oauth legacy official API remains on official media API", + name: "oauth legacy official API is pinned to CLI proxy", account: Account{ Type: AccountTypeOAuth, Platform: PlatformGrok, @@ -353,10 +353,10 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) { "base_url": xai.DefaultBaseURL, }, }, - expected: xai.DefaultBaseURL, + expected: xai.DefaultCLIBaseURL, }, { - name: "oauth untrusted custom base_url is pinned to official media API", + name: "oauth untrusted custom base_url is pinned to CLI proxy", account: Account{ Type: AccountTypeOAuth, Platform: PlatformGrok, @@ -364,7 +364,7 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) { "base_url": "https://custom.example.com/v1", }, }, - expected: xai.DefaultBaseURL, + expected: xai.DefaultCLIBaseURL, }, { name: "API key retains its configured media API", @@ -395,7 +395,7 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) { } } -func TestGetGrokMediaBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *testing.T) { +func TestGetGrokMediaBaseURLPinsOAuthWhenUnsafeOverridesEnabled(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") account := Account{ Type: AccountTypeOAuth, @@ -405,5 +405,5 @@ func TestGetGrokMediaBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnable }, } - require.Equal(t, "https://custom.example.com/v1", account.GetGrokMediaBaseURL()) + require.Equal(t, xai.DefaultCLIBaseURL, account.GetGrokMediaBaseURL()) } diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 222b3f8a4d9..a9eef998e49 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -724,7 +724,9 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/event-stream") req.Header.Set("Authorization", "Bearer "+authToken) - applyGrokCLIHeaders(req.Header) + if account.IsGrokOAuth() { + applyGrokCLIHeaders(req.Header) + } proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { diff --git a/backend/internal/service/account_test_service_grok_test.go b/backend/internal/service/account_test_service_grok_test.go index 4b0890ff441..f311402874e 100644 --- a/backend/internal/service/account_test_service_grok_test.go +++ b/backend/internal/service/account_test_service_grok_test.go @@ -40,8 +40,9 @@ func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testin Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "grok-access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "access_token": "grok-access-token", + "refresh_token": "grok-refresh-token", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), "model_mapping": map[string]any{ "grok": "grok-4.3", }, @@ -133,8 +134,9 @@ func TestAccountTestService_Grok429PersistsRateLimitReset(t *testing.T) { Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "grok-access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "access_token": "grok-access-token", + "refresh_token": "grok-refresh-token", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), }, } baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} @@ -166,8 +168,9 @@ func TestAccountTestService_Grok429WithoutQuotaHeadersUsesFallback(t *testing.T) ID: 15, Name: "grok-oauth-limited-no-headers", Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "grok-access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "access_token": "grok-access-token", + "refresh_token": "grok-refresh-token", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), }, } baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index c1513756237..49ef63c0c1f 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -558,19 +558,79 @@ type ForwardResult struct { ImageSizeBreakdown map[string]int } -// UpstreamFailoverError indicates an upstream error that should trigger account failover. +// GatewayFailureStage identifies which request stage failed. The zero value is +// intentionally treated as inference so existing UpstreamFailoverError callers +// retain their current behavior. +type GatewayFailureStage string + +const ( + GatewayFailureStageInference GatewayFailureStage = "inference" + GatewayFailureStageAccountAuth GatewayFailureStage = "account_auth" +) + +// GatewayFailureScope identifies whether selecting another account can help. +type GatewayFailureScope string + +const ( + GatewayFailureScopeAccount GatewayFailureScope = "account" + GatewayFailureScopeProvider GatewayFailureScope = "provider" + GatewayFailureScopeRequest GatewayFailureScope = "request" +) + +// NextAccountAction is tri-state for backwards compatibility. The zero value +// means legacy retry behavior; only NextAccountStop explicitly short-circuits. +type NextAccountAction uint8 + +const ( + NextAccountLegacyRetry NextAccountAction = iota + NextAccountRetry + NextAccountStop +) + +type GatewayFailureReason string + +// UpstreamFailoverError indicates an upstream or credential error that may +// trigger account failover. Additive metadata keeps existing composite literals +// source-compatible and preserves their legacy retry-next-account behavior. type UpstreamFailoverError struct { StatusCode int ResponseBody []byte // 上游响应体，用于错误透传规则匹配 ResponseHeaders http.Header // 上游响应头，用于透传 cf-ray/cf-mitigated/content-type 等诊断信息 ForceCacheBilling bool // Antigravity 粘性会话切换时设为 true RetryableOnSameAccount bool // 临时性错误（如 Google 间歇性 400、空响应），应在同一账号上重试 N 次再切换 + Stage GatewayFailureStage + Scope GatewayFailureScope + Reason GatewayFailureReason + NextAccountAction NextAccountAction + ClientStatusCode int + ClientMessage string } func (e *UpstreamFailoverError) Error() string { + if e != nil && e.Stage == GatewayFailureStageAccountAuth { + return fmt.Sprintf("credential failure: %s (failover)", e.Reason) + } return fmt.Sprintf("upstream error: %d (failover)", e.StatusCode) } +func (e *UpstreamFailoverError) ShouldRetryNextAccount() bool { + return e != nil && e.NextAccountAction != NextAccountStop +} + +func (e *UpstreamFailoverError) IsCredentialFailure() bool { + return e != nil && e.Stage == GatewayFailureStageAccountAuth +} + +// ShouldReportAccountScheduleFailure prevents provider- and request-scoped +// credential failures from being misattributed to the selected account. Legacy +// and inference failures retain their existing scheduler-health behavior. +func (e *UpstreamFailoverError) ShouldReportAccountScheduleFailure() bool { + if e == nil { + return false + } + return !e.IsCredentialFailure() || e.Scope == GatewayFailureScopeAccount +} + // sseStreamErrorEventError 表示上游 SSE 流体内出现 event:error 帧。 // RawData 是该事件 data: 行的原始 JSON 字符串 // （Anthropic 标准结构 {"type":"error","error":{"type":"...","message":"..."}}）。 diff --git a/backend/internal/service/grok_credential_failure.go b/backend/internal/service/grok_credential_failure.go new file mode 100644 index 00000000000..8886855f12f --- /dev/null +++ b/backend/internal/service/grok_credential_failure.go @@ -0,0 +1,650 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + "strings" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/gin-gonic/gin" +) + +const ( + grokCredentialFailoverDeadlineKey = "grok_credential_failover_deadline" + grokCredentialFailoverBudget = 15 * time.Second + grokCredentialMutationTimeout = 5 * time.Second + grokCredentialMutationConfirmWait = 250 * time.Millisecond + grokCredentialCacheCleanupTimeout = 500 * time.Millisecond + + GrokCredentialUnavailableClientMessage = "No healthy Grok OAuth account is currently available" + + GrokCredentialReasonRevoked GatewayFailureReason = "grok_oauth_credential_revoked" + GrokCredentialReasonMissing GatewayFailureReason = "grok_oauth_credentials_missing" + GrokCredentialReasonEntitlement GatewayFailureReason = "grok_oauth_entitlement_action_required" + GrokCredentialReasonProxyInvalid GatewayFailureReason = "grok_oauth_proxy_invalid" + GrokCredentialReasonRefreshTransient GatewayFailureReason = "grok_oauth_refresh_transient" + GrokCredentialReasonProviderConfig GatewayFailureReason = "grok_oauth_provider_config" + GrokCredentialReasonProviderDown GatewayFailureReason = "grok_oauth_provider_unavailable" + GrokCredentialReasonAccountChanged GatewayFailureReason = "grok_oauth_account_state_changed" + GrokCredentialReasonStateUpdate GatewayFailureReason = "grok_oauth_account_state_update_failed" + GrokCredentialReasonFailoverTimeout GatewayFailureReason = "grok_oauth_failover_timeout" +) + +var errGrokCredentialStateUpdateFailed = errors.New("grok oauth account state update failed") + +type grokCredentialFailureClass struct { + scope GatewayFailureScope + reason GatewayFailureReason + action NextAccountAction + permanent bool + transient bool + message string + snapshot *GrokCredentialMutationSnapshot +} + +// GrokCredentialMutationSnapshot is the credential identity observed when the +// request selected an account. Repository mutations compare all fields before +// quarantining that account so a concurrent refresh cannot be overwritten. +type GrokCredentialMutationSnapshot struct { + CredentialsJSON string + AccessToken string + RefreshToken string + TokenVersion int64 + ProxyID *int64 +} + +type grokCredentialFailureSnapshotError struct { + cause error + snapshot GrokCredentialMutationSnapshot +} + +func (e *grokCredentialFailureSnapshotError) Error() string { return e.cause.Error() } +func (e *grokCredentialFailureSnapshotError) Unwrap() error { return e.cause } + +func withGrokCredentialFailureSnapshot(err error, account *Account) error { + if err == nil || account == nil || !account.IsGrokOAuth() { + return err + } + var existing *grokCredentialFailureSnapshotError + if errors.As(err, &existing) { + return err + } + return &grokCredentialFailureSnapshotError{cause: err, snapshot: grokCredentialMutationSnapshot(account)} +} + +func grokCredentialFailureSnapshot(err error) (GrokCredentialMutationSnapshot, bool) { + var snapshotErr *grokCredentialFailureSnapshotError + if !errors.As(err, &snapshotErr) || snapshotErr == nil { + return GrokCredentialMutationSnapshot{}, false + } + return snapshotErr.snapshot, true +} + +type grokCredentialConditionalStateRepository interface { + SetGrokCredentialErrorIfMatch(context.Context, int64, GrokCredentialMutationSnapshot, string) (bool, error) + SetGrokCredentialTempUnschedulableIfMatch(context.Context, int64, GrokCredentialMutationSnapshot, time.Time, string) (bool, error) +} + +// GetRequestCredential applies the request-path credential and failover contract +// before any upstream transport is opened. +func (s *OpenAIGatewayService) GetRequestCredential(ctx context.Context, c *gin.Context, account *Account) (string, string, error) { + return s.getRequestCredential(ctx, c, account) +} + +func (s *OpenAIGatewayService) getRequestCredential(ctx context.Context, c *gin.Context, account *Account) (string, string, error) { + if ctx == nil { + ctx = context.Background() + } + if account == nil { + return "", "", errors.New("account is nil") + } + if !account.IsGrokOAuth() { + return s.GetAccessToken(ctx, account) + } + if err := ctx.Err(); err != nil { + return "", "", err + } + if s == nil || s.grokTokenProvider == nil { + return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{ + scope: GatewayFailureScopeProvider, + reason: GrokCredentialReasonProviderConfig, + action: NextAccountStop, + message: "Grok OAuth credential provider is unavailable", + }) + } + if s.isOpenAIAccountRuntimeBlocked(account) { + return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{ + scope: GatewayFailureScopeAccount, + reason: GrokCredentialReasonAccountChanged, + action: NextAccountRetry, + message: "Grok OAuth account is not currently schedulable", + }) + } + + credentialCtx, cancel, budgetExpired := grokCredentialAcquisitionContext(ctx, c) + if cancel != nil { + defer cancel() + } + if budgetExpired { + return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{ + scope: GatewayFailureScopeRequest, + reason: GrokCredentialReasonFailoverTimeout, + action: NextAccountStop, + message: "Grok OAuth credential failover budget exhausted", + }) + } + + token, kind, err := s.GetAccessToken(credentialCtx, account) + if err == nil { + if s.isOpenAIAccountRuntimeBlocked(account) { + return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{ + scope: GatewayFailureScopeAccount, + reason: GrokCredentialReasonAccountChanged, + action: NextAccountRetry, + message: "Grok OAuth account is not currently schedulable", + }) + } + return token, kind, nil + } + if parentErr := ctx.Err(); parentErr != nil { + return "", "", parentErr + } + if credentialCtx.Err() != nil { + return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{ + scope: GatewayFailureScopeRequest, + reason: GrokCredentialReasonFailoverTimeout, + action: NextAccountStop, + message: "Grok OAuth credential failover budget exhausted", + }) + } + + class := classifyGrokCredentialFailure(account, err) + if snapshot, ok := grokCredentialFailureSnapshot(err); ok { + class.snapshot = &snapshot + } + if ctx.Err() != nil { + return "", "", ctx.Err() + } + if class.permanent || class.transient { + freshToken, mutationErr := s.applyGrokCredentialAccountFailure(credentialCtx, account, class) + if freshToken != "" { + return freshToken, "oauth", nil + } + if mutationErr != nil { + if ctx.Err() != nil { + return "", "", ctx.Err() + } + if credentialCtx.Err() != nil { + return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{ + scope: GatewayFailureScopeRequest, + reason: GrokCredentialReasonFailoverTimeout, + action: NextAccountStop, + message: "Grok OAuth credential failover budget exhausted", + }) + } + if errors.Is(mutationErr, errOAuthRefreshAccountStateChanged) { + class = grokCredentialFailureClass{ + scope: GatewayFailureScopeAccount, + reason: GrokCredentialReasonAccountChanged, + action: NextAccountRetry, + message: "Grok OAuth account eligibility changed", + } + } else if errors.Is(mutationErr, errOAuthRefreshAccountRereadFailed) { + class = grokCredentialFailureClass{ + scope: GatewayFailureScopeProvider, + reason: GrokCredentialReasonProviderDown, + action: NextAccountStop, + message: "Grok OAuth account state is temporarily unavailable", + } + } else { + class = grokCredentialFailureClass{ + scope: GatewayFailureScopeProvider, + reason: GrokCredentialReasonStateUpdate, + action: NextAccountStop, + message: "Grok OAuth account state could not be updated safely", + } + } + } + } + return "", "", s.newGrokCredentialFailover(c, account, class) +} + +func grokCredentialAcquisitionContext(ctx context.Context, c *gin.Context) (context.Context, context.CancelFunc, bool) { + if c == nil { + return ctx, nil, false + } + deadline := time.Time{} + if raw, ok := c.Get(grokCredentialFailoverDeadlineKey); ok { + deadline, _ = raw.(time.Time) + } + if deadline.IsZero() { + deadline = time.Now().Add(grokCredentialFailoverBudget) + c.Set(grokCredentialFailoverDeadlineKey, deadline) + } + if !time.Now().Before(deadline) { + return ctx, nil, true + } + acquireCtx, cancel := context.WithDeadline(ctx, deadline) + return acquireCtx, cancel, false +} + +func classifyGrokCredentialFailure(account *Account, err error) grokCredentialFailureClass { + stableReason := strings.ToLower(strings.TrimSpace(infraerrors.Reason(err))) + message := "" + if err != nil { + message = strings.ToLower(err.Error()) + } + contains := func(values ...string) bool { + for _, value := range values { + if strings.Contains(stableReason, value) || strings.Contains(message, value) { + return true + } + } + return false + } + + switch { + case errors.Is(err, errGrokOAuthRefreshTokenMissing), errors.Is(err, errGrokOAuthAccessTokenMissing), errors.Is(err, errGrokOAuthAccessTokenExpired): + return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonMissing, action: NextAccountRetry, permanent: true, message: "Grok OAuth credentials are missing or expired"} + case contains("invalid_grant", "invalid_refresh_token", "token_expired", "refresh_token_reused", "refresh_token_invalidated", "app_session_terminated"): + return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRevoked, action: NextAccountRetry, permanent: true, message: "Grok OAuth credentials require account action"} + case contains("grok_oauth_entitlement_denied", "entitlement_denied", "access_denied", "subscription required", "no active grok subscription"): + return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonEntitlement, action: NextAccountRetry, permanent: true, message: "Grok OAuth entitlement requires account action"} + case errors.Is(err, errGrokOAuthConfiguredProxyMiss), contains("grok_oauth_proxy_not_found"): + return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonProxyInvalid, action: NextAccountRetry, permanent: true, message: "Grok OAuth account proxy configuration is invalid"} + case errors.Is(err, errOAuthRefreshAccountRereadFailed): + return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth account state is temporarily unavailable"} + case errors.Is(err, errOAuthRefreshCredentialPersist): + return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth shared credential state is temporarily unavailable"} + case errors.Is(err, errOAuthRefreshAccountStateChanged): + return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonAccountChanged, action: NextAccountRetry, message: "Grok OAuth account eligibility changed"} + case errors.Is(err, errGrokOAuthRefreshNotConfigured), contains("invalid_client", "unauthorized_client", "invalid_scope", "unknown scope", "grok oauth service is not configured", "grok_oauth_proxy_not_available"): + return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"} + case contains("grok_oauth_proxy_lookup_failed"), + contains("grok_oauth_token_refresh_failed") && contains("status 403") && (account == nil || account.ProxyID == nil): + return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"} + case contains("grok_oauth_client_init_failed") && (account == nil || account.ProxyID == nil): + return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"} + case contains("grok_oauth_request_failed") && (account == nil || account.ProxyID == nil): + return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"} + case contains("status 429", "status 500", "status 502", "status 503", "status 504") && (account == nil || account.ProxyID == nil): + return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"} + default: + return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRefreshTransient, action: NextAccountRetry, transient: true, message: "Grok OAuth credential refresh is temporarily unavailable"} + } +} + +func (s *OpenAIGatewayService) applyGrokCredentialAccountFailure(ctx context.Context, account *Account, class grokCredentialFailureClass) (string, error) { + if s == nil || account == nil || ctx == nil || ctx.Err() != nil { + if ctx != nil { + return "", ctx.Err() + } + return "", context.Canceled + } + mutationMu := s.grokCredentialMutationLock(account.ID) + if err := mutationMu.Lock(ctx); err != nil { + return "", err + } + defer mutationMu.Unlock() + stateRepo, hasConditionalStateRepo := s.accountRepo.(grokCredentialConditionalStateRepository) + snapshot := grokCredentialMutationSnapshot(account) + if class.snapshot != nil { + snapshot = *class.snapshot + } + if token, err := s.validateCurrentGrokCredentialFailure(ctx, account.ID, snapshot, class); err != nil || token != "" { + return token, err + } + + if class.permanent { + if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, account.ID, snapshot); ok { + return token, nil + } + if ctx.Err() != nil { + return "", ctx.Err() + } + rollbackRuntime := s.blockGrokCredentialRuntime(account, time.Time{}, string(class.reason)) + keepRuntimeBlock := false + runtimeRollbackDone := false + defer func() { + if !keepRuntimeBlock && !runtimeRollbackDone { + rollbackRuntime() + } + }() + if s.accountRepo == nil { + keepRuntimeBlock = true + return "", fmt.Errorf("%w: account repository is not configured", errGrokCredentialStateUpdateFailed) + } + if !hasConditionalStateRepo { + keepRuntimeBlock = true + return "", fmt.Errorf("%w: conditional account repository is not configured", errGrokCredentialStateUpdateFailed) + } + stateCtx, cancel := context.WithTimeout(ctx, grokCredentialMutationTimeout) + if err := ctx.Err(); err != nil { + cancel() + return "", err + } + updated, err := stateRepo.SetGrokCredentialErrorIfMatch(stateCtx, account.ID, snapshot, string(class.reason)) + requestErr := ctx.Err() + cancel() + if err != nil { + slog.Warn("grok_credential_failure.set_error_failed", "account_id", account.ID, "reason", class.reason, "error", err) + if requestErr != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + if s.grokCredentialMutationCommitted(account.ID, class, time.Time{}) { + updated = true + } else if requestErr != nil { + return "", requestErr + } else { + keepRuntimeBlock = true + return "", fmt.Errorf("%w: permanent state commit could not be confirmed: %v", errGrokCredentialStateUpdateFailed, err) + } + } else { + keepRuntimeBlock = true + return "", fmt.Errorf("%w: persist permanent state: %v", errGrokCredentialStateUpdateFailed, err) + } + } + if !updated { + rollbackRuntime() + runtimeRollbackDone = true + return s.resolveGrokCredentialCASMiss(ctx, account.ID, snapshot) + } + // SetError is the linearization point: the durable quarantine is now + // committed and this node's runtime block must not be rolled back. + keepRuntimeBlock = true + if s.grokTokenProvider == nil { + if ctx.Err() != nil { + return "", ctx.Err() + } + return "", fmt.Errorf("%w: token provider is not configured", errGrokCredentialStateUpdateFailed) + } + invalidateCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), grokCredentialCacheCleanupTimeout) + err = s.grokTokenProvider.InvalidateToken(invalidateCtx, account) + cancel() + if err != nil { + slog.Warn("grok_credential_failure.invalidate_token_failed", "account_id", account.ID, "reason", class.reason, "error", err) + if ctx.Err() != nil { + return "", ctx.Err() + } + return "", fmt.Errorf("%w: invalidate cached credential: %v", errGrokCredentialStateUpdateFailed, err) + } + if ctx.Err() != nil { + return "", ctx.Err() + } + return "", nil + } + + if class.transient { + until := time.Now().Add(tokenRefreshTempUnschedDuration) + if ctx.Err() != nil { + return "", ctx.Err() + } + rollbackRuntime := s.blockGrokCredentialRuntime(account, until, string(class.reason)) + keepRuntimeBlock := false + runtimeRollbackDone := false + defer func() { + if !keepRuntimeBlock && !runtimeRollbackDone { + rollbackRuntime() + } + }() + stateCtx, cancel := context.WithTimeout(ctx, grokCredentialMutationTimeout) + if s.accountRepo == nil { + cancel() + keepRuntimeBlock = true + return "", fmt.Errorf("%w: account repository is not configured", errGrokCredentialStateUpdateFailed) + } + if !hasConditionalStateRepo { + cancel() + keepRuntimeBlock = true + return "", fmt.Errorf("%w: conditional account repository is not configured", errGrokCredentialStateUpdateFailed) + } + if err := ctx.Err(); err != nil { + cancel() + return "", err + } + updated, err := stateRepo.SetGrokCredentialTempUnschedulableIfMatch(stateCtx, account.ID, snapshot, until, string(class.reason)) + requestErr := ctx.Err() + cancel() + if err != nil { + slog.Warn("grok_credential_failure.set_temp_unschedulable_failed", "account_id", account.ID, "reason", class.reason, "error", err) + if requestErr != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + if s.grokCredentialMutationCommitted(account.ID, class, until) { + updated = true + } else if requestErr != nil { + return "", requestErr + } else { + keepRuntimeBlock = true + return "", fmt.Errorf("%w: transient state commit could not be confirmed: %v", errGrokCredentialStateUpdateFailed, err) + } + } else { + keepRuntimeBlock = true + return "", fmt.Errorf("%w: persist transient state: %v", errGrokCredentialStateUpdateFailed, err) + } + } + if !updated { + rollbackRuntime() + runtimeRollbackDone = true + return s.resolveGrokCredentialCASMiss(ctx, account.ID, snapshot) + } + // The temporary quarantine is durable after SetTempUnschedulable succeeds. + keepRuntimeBlock = true + if ctx.Err() != nil { + return "", ctx.Err() + } + return "", nil + } + + return "", nil +} + +func (s *OpenAIGatewayService) validateCurrentGrokCredentialFailure( + ctx context.Context, + accountID int64, + snapshot GrokCredentialMutationSnapshot, + class grokCredentialFailureClass, +) (string, error) { + if s == nil || s.accountRepo == nil || accountID <= 0 || ctx == nil || ctx.Err() != nil { + if ctx != nil { + return "", ctx.Err() + } + return "", context.Canceled + } + checkCtx, cancel := context.WithTimeout(ctx, time.Second) + defer cancel() + latest, err := s.accountRepo.GetByID(checkCtx, accountID) + if err != nil { + if errors.Is(err, ErrAccountNotFound) { + return "", errOAuthRefreshAccountStateChanged + } + return "", fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, err) + } + if latest == nil || !latest.IsGrokOAuth() || !latest.IsSchedulable() || s.isOpenAIAccountRuntimeBlocked(latest) { + return "", errOAuthRefreshAccountStateChanged + } + latestSnapshot := grokCredentialMutationSnapshot(latest) + if latestSnapshot.CredentialsJSON != snapshot.CredentialsJSON || + !grokCredentialProxyIDsEqual(latestSnapshot.ProxyID, snapshot.ProxyID) { + if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, accountID, snapshot); ok { + return token, nil + } + return "", errOAuthRefreshAccountStateChanged + } + + // A configured proxy is external to the account-row CAS identity. Recheck + // the hydrated proxy object so restoring a deleted row under the same ID + // wins over a stale proxy-invalid failure. + if class.reason == GrokCredentialReasonProxyInvalid { + if latest.ProxyID == nil || latest.Proxy != nil { + return "", errOAuthRefreshAccountStateChanged + } + } else if latest.ProxyID != nil && latest.Proxy == nil { + return "", errOAuthRefreshAccountStateChanged + } + + if class.reason == GrokCredentialReasonMissing { + expiresAt := latest.GetCredentialAsTime("expires_at") + credentialsStillMissing := strings.TrimSpace(latest.GetGrokAccessToken()) == "" || + strings.TrimSpace(latest.GetGrokRefreshToken()) == "" || expiresAt == nil || !time.Now().Before(*expiresAt) + if !credentialsStillMissing { + return "", errOAuthRefreshAccountStateChanged + } + } + return "", nil +} + +func (s *OpenAIGatewayService) grokCredentialMutationLock(accountID int64) *oauthRefreshLocalLock { + actual, _ := s.grokCredentialMutationLocks.LoadOrStore(accountID, newOAuthRefreshLocalLock()) + return actual.(*oauthRefreshLocalLock) +} + +func (s *OpenAIGatewayService) grokCredentialMutationCommitted(accountID int64, class grokCredentialFailureClass, until time.Time) bool { + if s == nil || s.accountRepo == nil || accountID <= 0 { + return false + } + confirmCtx, cancel := context.WithTimeout(context.Background(), grokCredentialMutationConfirmWait) + defer cancel() + latest, err := s.accountRepo.GetByID(confirmCtx, accountID) + if err != nil || latest == nil { + return false + } + if class.permanent { + return latest.Status == StatusError && !latest.Schedulable && latest.ErrorMessage == string(class.reason) + } + if class.transient { + return latest.TempUnschedulableUntil != nil && !latest.TempUnschedulableUntil.Before(until) && + latest.TempUnschedulableReason == string(class.reason) + } + return false +} + +func grokCredentialMutationSnapshot(account *Account) GrokCredentialMutationSnapshot { + if account == nil { + return GrokCredentialMutationSnapshot{} + } + credentialsJSON := "null" + if encoded, err := json.Marshal(account.Credentials); err == nil { + credentialsJSON = string(encoded) + } + snapshot := GrokCredentialMutationSnapshot{ + CredentialsJSON: credentialsJSON, + AccessToken: strings.TrimSpace(account.GetGrokAccessToken()), + RefreshToken: strings.TrimSpace(account.GetGrokRefreshToken()), + TokenVersion: account.GetCredentialAsInt64("_token_version"), + } + if account.ProxyID != nil { + proxyID := *account.ProxyID + snapshot.ProxyID = &proxyID + } + return snapshot +} + +func (s *OpenAIGatewayService) resolveGrokCredentialCASMiss(ctx context.Context, accountID int64, snapshot GrokCredentialMutationSnapshot) (string, error) { + if ctx.Err() != nil { + return "", ctx.Err() + } + if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, accountID, snapshot); ok { + return token, nil + } + return "", errOAuthRefreshAccountStateChanged +} + +func (s *OpenAIGatewayService) blockGrokCredentialRuntime(account *Account, until time.Time, reason string) func() { + if s == nil || account == nil { + return func() {} + } + mu := s.openAIAccountRuntimeBlockLock(account.ID) + mu.Lock() + before, hadBefore := s.openaiAccountRuntimeBlockUntil.Load(account.ID) + installedGeneration, changed := s.blockAccountSchedulingLocked(account, until, reason) + installed, installedOK := s.openaiAccountRuntimeBlockUntil.Load(account.ID) + installedUntil, isTime := installed.(time.Time) + mu.Unlock() + if !changed || !installedOK || !isTime { + return func() {} + } + if hadBefore { + if beforeUntil, ok := before.(time.Time); ok && beforeUntil.Equal(installedUntil) { + return func() {} + } + } + return func() { + mu.Lock() + defer mu.Unlock() + generation, ok := s.openaiAccountRuntimeBlockGeneration.Load(account.ID) + if !ok || generation != installedGeneration { + return + } + current, ok := s.openaiAccountRuntimeBlockUntil.Load(account.ID) + currentUntil, isTime := current.(time.Time) + if !ok || !isTime || !currentUntil.Equal(installedUntil) { + return + } + if hadBefore { + s.openaiAccountRuntimeBlockUntil.Store(account.ID, before) + s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1)) + return + } + s.openaiAccountRuntimeBlockUntil.Delete(account.ID) + s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1)) + } +} + +func (s *OpenAIGatewayService) grokCredentialConcurrentlyRefreshedToken(ctx context.Context, accountID int64, baseline GrokCredentialMutationSnapshot) (string, bool) { + if s == nil || s.accountRepo == nil || accountID <= 0 || ctx == nil || ctx.Err() != nil { + return "", false + } + checkCtx, cancel := context.WithTimeout(ctx, time.Second) + defer cancel() + latest, err := s.accountRepo.GetByID(checkCtx, accountID) + if err != nil || latest == nil { + return "", false + } + latestSnapshot := grokCredentialMutationSnapshot(latest) + if !grokCredentialProxyIDsEqual(latestSnapshot.ProxyID, baseline.ProxyID) || + latestSnapshot.CredentialsJSON == baseline.CredentialsJSON || !latest.IsSchedulable() || + (latest.ProxyID != nil && latest.Proxy == nil) || s.isOpenAIAccountRuntimeBlocked(latest) { + return "", false + } + latestToken := strings.TrimSpace(latest.GetGrokAccessToken()) + if latestToken == "" || strings.TrimSpace(latest.GetGrokRefreshToken()) == "" { + return "", false + } + expiresAt := latest.GetCredentialAsTime("expires_at") + if expiresAt == nil || !time.Now().Before(*expiresAt) { + return "", false + } + return latestToken, true +} + +func grokCredentialProxyIDsEqual(left, right *int64) bool { + if left == nil || right == nil { + return left == nil && right == nil + } + return *left == *right +} + +func (s *OpenAIGatewayService) newGrokCredentialFailover(c *gin.Context, account *Account, class grokCredentialFailureClass) error { + if strings.TrimSpace(class.message) == "" { + class.message = "Grok OAuth credentials are unavailable" + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: PlatformGrok, + AccountID: account.ID, + Stage: string(GatewayFailureStageAccountAuth), + Scope: string(class.scope), + Reason: string(class.reason), + Kind: "credential_failover", + Message: class.message, + }) + return &UpstreamFailoverError{ + Stage: GatewayFailureStageAccountAuth, + Scope: class.scope, + Reason: class.reason, + NextAccountAction: class.action, + ClientStatusCode: http.StatusServiceUnavailable, + ClientMessage: GrokCredentialUnavailableClientMessage, + } +} diff --git a/backend/internal/service/grok_credential_failure_test.go b/backend/internal/service/grok_credential_failure_test.go new file mode 100644 index 00000000000..c32559ea12c --- /dev/null +++ b/backend/internal/service/grok_credential_failure_test.go @@ -0,0 +1,1622 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type grokCredentialPersistingRepo struct { + *tokenRefreshAccountRepo +} + +func (r *grokCredentialPersistingRepo) SetError(ctx context.Context, id int64, message string) error { + if err := ctx.Err(); err != nil { + return err + } + if err := r.tokenRefreshAccountRepo.SetError(ctx, id, message); err != nil { + return err + } + if account := r.accountsByID[id]; account != nil { + account.Status = StatusError + account.Schedulable = false + account.ErrorMessage = message + } + return nil +} + +type grokCredentialProxyRepoStub struct { + ProxyRepository + proxy *Proxy + err error +} + +func (r *grokCredentialProxyRepoStub) GetByID(context.Context, int64) (*Proxy, error) { + return r.proxy, r.err +} + +type grokCredentialBlockingRepo struct { + *tokenRefreshAccountRepo + setErrorStarted chan struct{} + setTempStarted chan struct{} + onceError sync.Once + onceTemp sync.Once +} + +type grokCredentialCommitThenCancelRepo struct { + *tokenRefreshAccountRepo + returnErr error +} + +type grokCredentialUncommittedDeadlineRepo struct { + *tokenRefreshAccountRepo +} + +func (r *grokCredentialUncommittedDeadlineRepo) SetGrokCredentialErrorIfMatch( + context.Context, + int64, + GrokCredentialMutationSnapshot, + string, +) (bool, error) { + return false, context.DeadlineExceeded +} + +func (r *grokCredentialUncommittedDeadlineRepo) SetGrokCredentialTempUnschedulableIfMatch( + context.Context, + int64, + GrokCredentialMutationSnapshot, + time.Time, + string, +) (bool, error) { + return false, context.DeadlineExceeded +} + +func (r *grokCredentialCommitThenCancelRepo) SetGrokCredentialErrorIfMatch( + ctx context.Context, + id int64, + _ GrokCredentialMutationSnapshot, + reason string, +) (bool, error) { + account := r.accountsByID[id] + account.Status = StatusError + account.Schedulable = false + account.ErrorMessage = reason + if r.returnErr != nil { + return false, r.returnErr + } + <-ctx.Done() + return false, ctx.Err() +} + +func (r *grokCredentialCommitThenCancelRepo) SetGrokCredentialTempUnschedulableIfMatch( + ctx context.Context, + id int64, + _ GrokCredentialMutationSnapshot, + until time.Time, + reason string, +) (bool, error) { + account := r.accountsByID[id] + account.TempUnschedulableUntil = &until + account.TempUnschedulableReason = reason + if r.returnErr != nil { + return false, r.returnErr + } + <-ctx.Done() + return false, ctx.Err() +} + +func (r *grokCredentialBlockingRepo) SetError(ctx context.Context, _ int64, _ string) error { + r.onceError.Do(func() { close(r.setErrorStarted) }) + <-ctx.Done() + return ctx.Err() +} + +func (r *grokCredentialBlockingRepo) SetTempUnschedulable(ctx context.Context, _ int64, _ time.Time, _ string) error { + r.onceTemp.Do(func() { close(r.setTempStarted) }) + <-ctx.Done() + return ctx.Err() +} + +func (r *grokCredentialBlockingRepo) SetGrokCredentialErrorIfMatch( + ctx context.Context, + _ int64, + _ GrokCredentialMutationSnapshot, + _ string, +) (bool, error) { + r.onceError.Do(func() { close(r.setErrorStarted) }) + <-ctx.Done() + return false, ctx.Err() +} + +func (r *grokCredentialBlockingRepo) SetGrokCredentialTempUnschedulableIfMatch( + ctx context.Context, + _ int64, + _ GrokCredentialMutationSnapshot, + _ time.Time, + _ string, +) (bool, error) { + r.onceTemp.Do(func() { close(r.setTempStarted) }) + <-ctx.Done() + return false, ctx.Err() +} + +type grokCredentialBlockingCache struct { + GrokTokenCache + deleteStarted chan struct{} + releaseDelete chan struct{} + once sync.Once + mu sync.Mutex + deleted bool +} + +type grokCredentialSequencedRepo struct { + *tokenRefreshAccountRepo + mu sync.Mutex + latest *Account + getCall int +} + +type grokCredentialRereadFailureRepo struct { + *tokenRefreshAccountRepo + account *Account + err error +} + +type grokCredentialCountingRefresher struct { + refreshCalls int +} + +func (r *grokCredentialCountingRefresher) CacheKey(account *Account) string { + return GrokTokenCacheKey(account) +} + +func (r *grokCredentialCountingRefresher) CanRefresh(*Account) bool { return true } + +func (r *grokCredentialCountingRefresher) NeedsRefresh(*Account, time.Duration) bool { return true } + +func (r *grokCredentialCountingRefresher) Refresh(context.Context, *Account) (map[string]any, error) { + r.refreshCalls++ + return map[string]any{"access_token": "must-not-be-used"}, nil +} + +func (r *grokCredentialRereadFailureRepo) GetByID(context.Context, int64) (*Account, error) { + return r.account, r.err +} + +func (r *grokCredentialSequencedRepo) GetByID(ctx context.Context, id int64) (*Account, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.getCall++ + if r.getCall > 1 && r.latest != nil { + return r.latest, nil + } + return r.tokenRefreshAccountRepo.GetByID(ctx, id) +} + +func (c *grokCredentialBlockingCache) DeleteAccessToken(ctx context.Context, _ string) error { + c.once.Do(func() { close(c.deleteStarted) }) + select { + case <-c.releaseDelete: + c.mu.Lock() + c.deleted = true + c.mu.Unlock() + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (c *grokCredentialBlockingCache) wasDeleted() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.deleted +} + +func TestUpstreamFailoverErrorNextAccountActionPreservesLegacyRetry(t *testing.T) { + t.Parallel() + + require.True(t, (&UpstreamFailoverError{}).ShouldRetryNextAccount()) + require.True(t, (&UpstreamFailoverError{NextAccountAction: NextAccountRetry}).ShouldRetryNextAccount()) + require.False(t, (&UpstreamFailoverError{NextAccountAction: NextAccountStop}).ShouldRetryNextAccount()) +} + +func TestGetRequestCredentialMapsPermanentGrokOAuthFailureAndRedactsSecrets(t *testing.T) { + gin.SetMode(gin.TestMode) + account := expiredGrokOAuthAccountForCredentialTest(701) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{ + err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant access_token=leaked-access refresh_token=leaked-refresh"), + }) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + token, kind, err := svc.getRequestCredential(context.Background(), c, account) + require.Error(t, err) + require.Empty(t, token) + require.Empty(t, kind) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureStageAccountAuth, failoverErr.Stage) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonRevoked, failoverErr.Reason) + require.True(t, failoverErr.ShouldRetryNextAccount()) + require.Equal(t, 0, failoverErr.StatusCode) + require.Equal(t, http.StatusServiceUnavailable, failoverErr.ClientStatusCode) + require.NotContains(t, err.Error(), "leaked-access") + require.NotContains(t, err.Error(), "leaked-refresh") + + require.Equal(t, 1, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys) + require.NotContains(t, repo.lastErrorMessage, "leaked-access") + require.NotContains(t, repo.lastErrorMessage, "leaked-refresh") + + rawEvents, ok := c.Get(OpsUpstreamErrorsKey) + require.True(t, ok) + events, ok := rawEvents.([]*OpsUpstreamErrorEvent) + require.True(t, ok) + require.Len(t, events, 1) + require.Equal(t, string(GatewayFailureStageAccountAuth), events[0].Stage) + require.Equal(t, string(GatewayFailureScopeAccount), events[0].Scope) + require.Equal(t, string(GrokCredentialReasonRevoked), events[0].Reason) + require.Zero(t, events[0].UpstreamStatusCode) + require.NotContains(t, events[0].Message, "leaked-access") + require.NotContains(t, events[0].Message, "leaked-refresh") +} + +func TestGetRequestCredentialPermanentMappingsPersistAndInvalidate(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + prepare func(*Account) + refreshErr error + wantReason GatewayFailureReason + cachedToken string + }{ + { + name: "missing refresh credential", + prepare: func(account *Account) { + delete(account.Credentials, "refresh_token") + }, + wantReason: GrokCredentialReasonMissing, + }, + { + name: "missing access credential", + prepare: func(account *Account) { + delete(account.Credentials, "access_token") + account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + }, + wantReason: GrokCredentialReasonMissing, + cachedToken: "stale-cached-access", + }, + { + name: "explicit entitlement action required", + prepare: func(*Account) {}, + refreshErr: infraerrors.New(http.StatusForbidden, "GROK_OAUTH_ENTITLEMENT_DENIED", "access_denied"), + wantReason: GrokCredentialReasonEntitlement, + }, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(720 + index)) + account.Status = StatusActive + account.Schedulable = true + tt.prepare(account) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: true, token: tt.cachedToken} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr}) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, tt.wantReason, failoverErr.Reason) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, 1, baseRepo.setErrorCalls) + require.Equal(t, StatusError, account.Status) + require.False(t, account.Schedulable) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys) + }) + } +} + +func TestGetRequestCredentialMissingAccessNeverRefreshesAndPermanentlyFailsOver(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + expiresAt *time.Time + }{ + {name: "expiry missing"}, + {name: "expired", expiresAt: func() *time.Time { value := time.Now().Add(-time.Minute); return &value }()}, + {name: "near expiry", expiresAt: func() *time.Time { value := time.Now().Add(30 * time.Minute); return &value }()}, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(760 + index)) + account.Schedulable = true + delete(account.Credentials, "access_token") + if tt.expiresAt == nil { + delete(account.Credentials, "expires_at") + } else { + account.Credentials["expires_at"] = tt.expiresAt.UTC().Format(time.RFC3339) + } + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: true, token: "stale-cache-must-not-win"} + refresher := &grokCredentialCountingRefresher{} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + token, kind, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Empty(t, token) + require.Empty(t, kind) + require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Zero(t, refresher.refreshCalls, "structurally missing access credentials must not reach the token endpoint") + require.Equal(t, 1, baseRepo.setErrorCalls) + require.Equal(t, StatusError, account.Status) + require.False(t, account.Schedulable) + require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + } +} + +func TestGetRequestCredentialWarmCachedAccessWithMissingRefreshPermanentlyFailsOver(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(764) + account.Credentials["access_token"] = "valid-access" + account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + delete(account.Credentials, "refresh_token") + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: true, token: "valid-access"} + refresher := &grokCredentialCountingRefresher{} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Zero(t, refresher.refreshCalls) + require.Equal(t, 1, baseRepo.setErrorCalls) + require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestGetRequestCredentialMapsTransientAndProviderFailuresSeparately(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("account transient temporarily unschedules", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(702) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: errors.New("temporary refresh transport failure")}) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonRefreshTransient, failoverErr.Reason) + require.True(t, failoverErr.ShouldRetryNextAccount()) + require.Zero(t, repo.setErrorCalls) + require.Equal(t, 1, repo.setTempUnschedCalls) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("shared provider configuration stops without mutation", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(703) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + provider := NewGrokTokenProvider(repo, nil) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope) + require.Equal(t, NextAccountStop, failoverErr.NextAccountAction) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("account reread failures preserve shared versus missing-row scope", func(t *testing.T) { + for _, tt := range []struct { + name string + account *Account + err error + }{ + {name: "repository error", err: errors.New("database temporarily unavailable")}, + {name: "missing row"}, + } { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(712) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: baseRepo, account: tt.account, err: tt.err} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil)) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + if tt.err != nil { + require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason) + require.Equal(t, NextAccountStop, failoverErr.NextAccountAction) + } else { + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + } + require.Zero(t, baseRepo.setErrorCalls) + require.Zero(t, baseRepo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + } + }) + + t.Run("fresh account eligibility changes retry without mutating stale state", func(t *testing.T) { + for _, tt := range []struct { + name string + mutate func(*Account) + }{ + { + name: "account disabled", + mutate: func(account *Account) { + account.Status = StatusDisabled + }, + }, + { + name: "account converted", + mutate: func(account *Account) { + account.Type = AccountTypeUpstream + }, + }, + { + name: "account manually unschedulable", + mutate: func(account *Account) { + account.Schedulable = false + }, + }, + { + name: "account temporarily unschedulable", + mutate: func(account *Account) { + until := time.Now().Add(time.Minute) + account.TempUnschedulableUntil = &until + }, + }, + } { + t.Run(tt.name, func(t *testing.T) { + staleAccount := expiredGrokOAuthAccountForCredentialTest(713) + freshAccount := *staleAccount + freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials) + tt.mutate(&freshAccount) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil)) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, staleAccount) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount)) + }) + } + }) + + t.Run("fresh missing refresh credential permanently blocks the account", func(t *testing.T) { + staleAccount := expiredGrokOAuthAccountForCredentialTest(714) + staleAccount.Schedulable = true + freshAccount := *staleAccount + freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials) + delete(freshAccount.Credentials, "refresh_token") + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount} + repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil)) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, staleAccount) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Equal(t, 1, baseRepo.setErrorCalls) + require.Zero(t, baseRepo.setTempUnschedCalls) + require.Equal(t, StatusError, freshAccount.Status) + require.False(t, freshAccount.Schedulable) + require.Equal(t, []string{GrokTokenCacheKey(staleAccount)}, cache.deletedKeys) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount)) + }) + + t.Run("refresh added after locked structural failure wins conditional mutation", func(t *testing.T) { + staleAccount := expiredGrokOAuthAccountForCredentialTest(717) + freshAccount := *staleAccount + freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials) + delete(freshAccount.Credentials, "refresh_token") + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount} + repo.beforeConditionalState = func() { + repaired := freshAccount + repaired.Credentials = shallowCopyMap(freshAccount.Credentials) + repaired.Credentials["refresh_token"] = "repaired-refresh-token" + repaired.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339) + repo.accountsByID[staleAccount.ID] = &repaired + } + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil)) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + token, kind, err := svc.getRequestCredential(context.Background(), c, staleAccount) + + require.NoError(t, err) + require.Equal(t, "expired-access-token", token) + require.Equal(t, "oauth", kind) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount)) + }) + + t.Run("expiry-only repair wins full credential fingerprint CAS", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(718) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + repo.beforeConditionalState = func() { + repaired := *account + repaired.Credentials = shallowCopyMap(account.Credentials) + repaired.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339) + repo.accountsByID[account.ID] = &repaired + } + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{ + err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"), + }) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + token, kind, err := svc.getRequestCredential(context.Background(), c, account) + + require.NoError(t, err) + require.Equal(t, "expired-access-token", token) + require.Equal(t, "oauth", kind) + require.Zero(t, repo.setErrorCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("generic token endpoint 403 stops as shared provider failure", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(708) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{ + err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "token refresh failed: status 403, body: forbidden"), + }) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason) + require.Equal(t, NextAccountStop, failoverErr.NextAccountAction) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("account proxy generic 403 remains bounded account transient", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(711) + proxyID := int64(43) + account.ProxyID = &proxyID + account.Proxy = &Proxy{} + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{ + err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "token refresh failed: status 403, body: forbidden"), + }) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonRefreshTransient, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Zero(t, repo.setErrorCalls) + require.Equal(t, 1, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("proxy repository read failure stops without account mutation", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(709) + proxyID := int64(41) + account.ProxyID = &proxyID + account.Proxy = &Proxy{} + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{lockResult: true} + oauthSvc := NewGrokOAuthService(&grokCredentialProxyRepoStub{err: errors.New("database temporarily unavailable")}, &grokOAuthClientStub{}) + defer oauthSvc.Stop() + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc)) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason) + require.Equal(t, NextAccountStop, failoverErr.NextAccountAction) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("structurally missing configured proxy permanently blocks only that account", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(710) + account.Status = StatusActive + account.Schedulable = true + proxyID := int64(42) + account.ProxyID = &proxyID + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: true} + oauthSvc := NewGrokOAuthService(&grokCredentialProxyRepoStub{err: ErrProxyNotFound}, &grokOAuthClientStub{}) + defer oauthSvc.Stop() + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc)) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonProxyInvalid, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Equal(t, 1, baseRepo.setErrorCalls) + require.Equal(t, StatusError, account.Status) + require.False(t, account.Schedulable) + require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) +} + +func TestGetRequestCredentialRuntimeBlockWinsBeforeWarmTokenCache(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(716) + account.Credentials["access_token"] = "valid-access" + account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{token: "valid-access"} + provider := NewGrokTokenProvider(repo, cache) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + svc.BlockAccountScheduling(account, time.Now().Add(time.Minute), "independent") + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Zero(t, cache.getCalls) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) +} + +func TestGetRequestCredentialWarmCachedAccessWithMissingConfiguredProxyPermanentlyFailsOver(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(715) + account.Credentials["access_token"] = "valid-access" + account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + proxyID := int64(44) + account.ProxyID = &proxyID + account.Proxy = nil + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: true, token: "valid-access"} + refresher := &grokCredentialCountingRefresher{} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GrokCredentialReasonProxyInvalid, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Zero(t, refresher.refreshCalls) + require.Equal(t, 1, baseRepo.setErrorCalls) + require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestGetRequestCredentialCancellationAndBudgetDoNotMutateAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + account := expiredGrokOAuthAccountForCredentialTest(704) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + provider := NewGrokTokenProvider(repo, nil) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + + t.Run("parent cancellation is returned directly", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(ctx, c, account) + require.ErrorIs(t, err, context.Canceled) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr)) + }) + + t.Run("request credential budget stops safely", func(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set(grokCredentialFailoverDeadlineKey, time.Now().Add(-time.Second)) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeRequest, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonFailoverTimeout, failoverErr.Reason) + require.False(t, failoverErr.ShouldRetryNextAccount()) + }) + + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestGetRequestCredentialStateMutationFailureStopsAndKeepsRuntimeBlock(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + refreshErr error + configure func(*tokenRefreshAccountRepo, *grokTokenCacheForProviderTest) + wantSetError int + wantSetTemp int + wantCacheDelete int + }{ + { + name: "permanent state persistence", + refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"), + configure: func(repo *tokenRefreshAccountRepo, _ *grokTokenCacheForProviderTest) { + repo.setErrorErr = errors.New("database write failed") + }, + wantSetError: 1, + }, + { + name: "transient state persistence", + refreshErr: errors.New("temporary refresh transport failure"), + configure: func(repo *tokenRefreshAccountRepo, _ *grokTokenCacheForProviderTest) { + repo.setTempUnschedErr = errors.New("database write failed") + }, + wantSetTemp: 1, + }, + { + name: "permanent token cache invalidation", + refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"), + configure: func(_ *tokenRefreshAccountRepo, cache *grokTokenCacheForProviderTest) { + cache.deleteErr = errors.New("cache delete failed") + }, + wantSetError: 1, + wantCacheDelete: 1, + }, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(740 + index)) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{lockResult: true} + tt.configure(repo, cache) + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr}) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonStateUpdate, failoverErr.Reason) + require.Equal(t, NextAccountStop, failoverErr.NextAccountAction) + require.Equal(t, tt.wantSetError, repo.setErrorCalls) + require.Equal(t, tt.wantSetTemp, repo.setTempUnschedCalls) + require.Len(t, cache.deletedKeys, tt.wantCacheDelete) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "failed mutation must retain the immediate local block") + }) + } +} + +func TestGrokCredentialMutationBoundariesHonorParentCancellation(t *testing.T) { + t.Run("blocked SetError cancellation prevents cache and runtime mutation", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(730) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialBlockingRepo{ + tokenRefreshAccountRepo: baseRepo, + setErrorStarted: make(chan struct{}), + setTempStarted: make(chan struct{}), + } + cache := &grokTokenCacheForProviderTest{} + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)} + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { + _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{ + reason: GrokCredentialReasonRevoked, permanent: true, + }) + result <- err + }() + <-repo.setErrorStarted + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede persistent SetError") + cancel() + + require.ErrorIs(t, <-result, context.Canceled) + require.Zero(t, baseRepo.setErrorCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("blocked temporary unschedule cancellation prevents runtime mutation", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(731) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialBlockingRepo{ + tokenRefreshAccountRepo: baseRepo, + setErrorStarted: make(chan struct{}), + setTempStarted: make(chan struct{}), + } + svc := &OpenAIGatewayService{accountRepo: repo} + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { + _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{ + reason: GrokCredentialReasonRefreshTransient, transient: true, + }) + result <- err + }() + <-repo.setTempStarted + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede temporary unscheduling") + cancel() + + require.ErrorIs(t, <-result, context.Canceled) + require.Zero(t, baseRepo.setTempUnschedCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("post-commit cancellation finishes cache cleanup and retains quarantine", func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(732) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokCredentialBlockingCache{deleteStarted: make(chan struct{}), releaseDelete: make(chan struct{})} + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)} + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { + _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{ + reason: GrokCredentialReasonRevoked, permanent: true, + }) + result <- err + }() + <-cache.deleteStarted + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede cache invalidation") + cancel() + close(cache.releaseDelete) + + require.ErrorIs(t, <-result, context.Canceled) + require.Equal(t, 1, repo.setErrorCalls) + require.True(t, cache.wasDeleted()) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) +} + +func TestGrokCredentialMutationLockWaitHonorsCredentialBudget(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(735) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})} + mutationLock := svc.grokCredentialMutationLock(account.ID) + require.NoError(t, mutationLock.Lock(context.Background())) + defer mutationLock.Unlock() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + + startedAt := time.Now() + token, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{ + reason: GrokCredentialReasonRevoked, permanent: true, + }) + + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Empty(t, token) + require.Less(t, time.Since(startedAt), 500*time.Millisecond) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestGetRequestCredentialBudgetBoundsBlockedConditionalMutation(t *testing.T) { + gin.SetMode(gin.TestMode) + account := expiredGrokOAuthAccountForCredentialTest(736) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialBlockingRepo{ + tokenRefreshAccountRepo: baseRepo, + setErrorStarted: make(chan struct{}), + setTempStarted: make(chan struct{}), + } + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{ + err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"), + }) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set(grokCredentialFailoverDeadlineKey, time.Now().Add(40*time.Millisecond)) + + startedAt := time.Now() + token, kind, err := svc.getRequestCredential(context.Background(), c, account) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeRequest, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonFailoverTimeout, failoverErr.Reason) + require.Equal(t, NextAccountStop, failoverErr.NextAccountAction) + require.Empty(t, token) + require.Empty(t, kind) + require.Less(t, time.Since(startedAt), 500*time.Millisecond) + require.Zero(t, baseRepo.setErrorCalls) + require.Zero(t, baseRepo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestGetRequestCredentialLockHeldTimeoutDoesNotQuarantineAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + buildRepo func(*Account) AccountRepository + wantScope GatewayFailureScope + wantReason GatewayFailureReason + wantAction NextAccountAction + }{ + { + name: "authoritative row unchanged", + buildRepo: func(account *Account) AccountRepository { + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + return repo + }, + wantScope: GatewayFailureScopeAccount, + wantReason: GrokCredentialReasonAccountChanged, + wantAction: NextAccountRetry, + }, + { + name: "selected account was deleted", + buildRepo: func(account *Account) AccountRepository { + base := &tokenRefreshAccountRepo{} + base.accountsByID = map[int64]*Account{account.ID: account} + return &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: base} + }, + wantScope: GatewayFailureScopeAccount, + wantReason: GrokCredentialReasonAccountChanged, + wantAction: NextAccountRetry, + }, + { + name: "shared account store unavailable", + buildRepo: func(account *Account) AccountRepository { + base := &tokenRefreshAccountRepo{} + base.accountsByID = map[int64]*Account{account.ID: account} + return &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: base, err: errors.New("database unavailable")} + }, + wantScope: GatewayFailureScopeProvider, + wantReason: GrokCredentialReasonProviderDown, + wantAction: NextAccountStop, + }, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(7400 + index)) + repo := tt.buildRepo(account) + cache := &grokTokenCacheForProviderTest{lockResult: false} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{}) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + startedAt := time.Now() + _, _, err := svc.getRequestCredential(context.Background(), c, account) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, tt.wantScope, failoverErr.Scope) + require.Equal(t, tt.wantReason, failoverErr.Reason) + require.Equal(t, tt.wantAction, failoverErr.NextAccountAction) + require.Less(t, time.Since(startedAt), 3*time.Second) + switch countingRepo := repo.(type) { + case *tokenRefreshAccountRepo: + require.Zero(t, countingRepo.setErrorCalls) + require.Zero(t, countingRepo.setTempUnschedCalls) + case *grokCredentialRereadFailureRepo: + require.Zero(t, countingRepo.tokenRefreshAccountRepo.setErrorCalls) + require.Zero(t, countingRepo.tokenRefreshAccountRepo.setTempUnschedCalls) + } + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + } +} + +func TestGrokCredentialMutationCancellationAmbiguityConfirmsDurableCommit(t *testing.T) { + tests := []struct { + name string + class grokCredentialFailureClass + committed func(*Account) bool + }{ + { + name: "permanent quarantine", + class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}, + committed: func(account *Account) bool { + return account.Status == StatusError && !account.Schedulable && account.ErrorMessage == string(GrokCredentialReasonRevoked) + }, + }, + { + name: "temporary quarantine", + class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}, + committed: func(account *Account) bool { + return account.TempUnschedulableUntil != nil && account.TempUnschedulableReason == string(GrokCredentialReasonRefreshTransient) + }, + }, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(737 + index)) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialCommitThenCancelRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{} + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)} + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + + token, err := svc.applyGrokCredentialAccountFailure(ctx, account, tt.class) + + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Empty(t, token) + require.True(t, tt.committed(account), "the detached confirmation must recognize the durable mutation") + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "a confirmed durable quarantine must retain its runtime block") + if tt.class.permanent { + require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys) + } else { + require.Empty(t, cache.deletedKeys) + } + }) + } +} + +func TestGrokCredentialInnerStateDeadlineAmbiguityConfirmsDurableCommit(t *testing.T) { + tests := []struct { + name string + class grokCredentialFailureClass + }{ + {name: "permanent quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}}, + {name: "temporary quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}}, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(7500 + index)) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialCommitThenCancelRepo{ + tokenRefreshAccountRepo: baseRepo, + returnErr: context.DeadlineExceeded, + } + cache := &grokTokenCacheForProviderTest{} + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)} + + token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, tt.class) + + require.NoError(t, err, "the detached readback must resolve the inner timeout's commit ambiguity") + require.Empty(t, token) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + if tt.class.permanent { + require.Equal(t, StatusError, account.Status) + require.False(t, account.Schedulable) + require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys) + } else { + require.NotNil(t, account.TempUnschedulableUntil) + require.Equal(t, string(GrokCredentialReasonRefreshTransient), account.TempUnschedulableReason) + require.Empty(t, cache.deletedKeys) + } + }) + } +} + +func TestGrokCredentialUnconfirmedInnerStateDeadlineStopsAndRetainsSafetyBlock(t *testing.T) { + tests := []struct { + name string + class grokCredentialFailureClass + }{ + {name: "permanent quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}}, + {name: "temporary quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}}, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(7600 + index)) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialUncommittedDeadlineRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{} + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)} + + token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, tt.class) + + require.ErrorIs(t, err, errGrokCredentialStateUpdateFailed) + require.Empty(t, token) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "an unknown commit outcome must retain the local safety block") + require.Equal(t, StatusActive, account.Status) + require.True(t, account.Schedulable) + require.Nil(t, account.TempUnschedulableUntil) + require.Empty(t, cache.deletedKeys) + }) + } +} + +func TestGrokCredentialRuntimeRollbackOwnership(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(734) + t.Run("later extending block survives", func(t *testing.T) { + svc := &OpenAIGatewayService{} + until := time.Now().Add(time.Minute) + rollbackFirst := svc.blockGrokCredentialRuntime(account, until, "first") + + secondInstalled := make(chan struct{}) + go func() { + svc.BlockAccountScheduling(account, until.Add(time.Minute), "independent") + close(secondInstalled) + }() + <-secondInstalled + rollbackFirst() + + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), + "rollback owned by the first invocation must not remove a later extending block") + }) + + t.Run("independent shorter block steals rollback ownership", func(t *testing.T) { + svc := &OpenAIGatewayService{} + until := time.Now().Add(2 * time.Minute) + rollbackFirst := svc.blockGrokCredentialRuntime(account, until, "first") + svc.BlockAccountScheduling(account, until.Add(-time.Minute), "shorter-no-op") + + rollbackFirst() + + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + + t.Run("serialized tentative rollbacks leave no block", func(t *testing.T) { + svc := &OpenAIGatewayService{} + for i := 0; i < 2; i++ { + mu := svc.grokCredentialMutationLock(account.ID) + require.NoError(t, mu.Lock(context.Background())) + rollback := svc.blockGrokCredentialRuntime(account, time.Now().Add(time.Minute), "tentative") + rollback() + mu.Unlock() + } + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) +} + +func TestGetRequestCredentialAPIKeyBypassesOAuthFailureMapping(t *testing.T) { + gin.SetMode(gin.TestMode) + account := &Account{ + ID: 705, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "third-party-key", + "base_url": "https://grok.example.test/v1", + }, + } + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + svc := &OpenAIGatewayService{} + + token, kind, err := svc.getRequestCredential(context.Background(), c, account) + require.NoError(t, err) + require.Equal(t, "third-party-key", token) + require.Equal(t, "apikey", kind) + _, hasEvents := c.Get(OpsUpstreamErrorsKey) + require.False(t, hasEvents) +} + +func TestPermanentCredentialFailureDoesNotDisableConcurrentlyRefreshedAccount(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(707) + latest := *account + latest.Credentials = shallowCopyMap(account.Credentials) + latest.Credentials["access_token"] = "fresh-access-token" + latest.Credentials["refresh_token"] = "rotated-refresh-token" + latest.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339) + latest.Credentials["_token_version"] = time.Now().UnixMilli() + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: &latest} + cache := &grokTokenCacheForProviderTest{} + svc := &OpenAIGatewayService{ + accountRepo: repo, + grokTokenProvider: NewGrokTokenProvider(repo, cache), + } + + _, mutationErr := svc.applyGrokCredentialAccountFailure(context.Background(), account, grokCredentialFailureClass{ + scope: GatewayFailureScopeAccount, + reason: GrokCredentialReasonRevoked, + action: NextAccountRetry, + permanent: true, + }) + + require.NoError(t, mutationErr) + require.Zero(t, repo.setErrorCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestCredentialFailureConditionalMutationLosesToConcurrentRefresh(t *testing.T) { + for index, tt := range []struct { + name string + refreshErr error + }{ + {name: "permanent", refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant")}, + {name: "transient", refreshErr: errors.New("temporary refresh transport failure")}, + } { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(770 + index)) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + repo.beforeConditionalState = func() { + fresh := *account + fresh.Credentials = shallowCopyMap(account.Credentials) + fresh.Credentials["access_token"] = "refresh-won-token" + fresh.Credentials["refresh_token"] = "refresh-won-refresh" + fresh.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + fresh.Credentials["_token_version"] = time.Now().UnixMilli() + repo.accountsByID[account.ID] = &fresh + } + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr}) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + token, kind, err := svc.getRequestCredential(context.Background(), c, account) + + require.NoError(t, err) + require.Equal(t, "refresh-won-token", token) + require.Equal(t, "oauth", kind) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + } +} + +func TestCredentialFailureConditionalMutationLosesToConcurrentProxyRepair(t *testing.T) { + for index, tt := range []struct { + name string + refreshErr error + }{ + {name: "permanent", refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant")}, + {name: "transient", refreshErr: errors.New("temporary refresh transport failure")}, + } { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(780 + index)) + oldProxyID := int64(10) + account.ProxyID = &oldProxyID + account.Proxy = &Proxy{} + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + repo.beforeConditionalState = func() { + fresh := *account + fresh.Credentials = shallowCopyMap(account.Credentials) + repairedProxyID := int64(11) + fresh.ProxyID = &repairedProxyID + repo.accountsByID[account.ID] = &fresh + } + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr}) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + } +} + +func TestCredentialFailureConditionalMutationLosesToSameIDProxyRestoration(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(790) + proxyID := int64(10) + account.ProxyID = &proxyID + account.Proxy = nil + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + repo.beforeConditionalState = func() { + account.Proxy = &Proxy{ID: proxyID} + } + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestCredentialFailureConditionalMutationLosesToConcurrentUnschedulableState(t *testing.T) { + future := time.Now().Add(time.Hour) + states := []struct { + name string + mutate func(*Account) + }{ + {name: "admin schedulable false", mutate: func(account *Account) { account.Schedulable = false }}, + {name: "temporary cooldown", mutate: func(account *Account) { account.TempUnschedulableUntil = &future }}, + {name: "rate limit cooldown", mutate: func(account *Account) { account.RateLimitResetAt = &future }}, + {name: "overload cooldown", mutate: func(account *Account) { account.OverloadUntil = &future }}, + } + classes := []struct { + name string + class grokCredentialFailureClass + }{ + {name: "permanent", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}}, + {name: "transient", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}}, + } + + for classIndex, classCase := range classes { + for stateIndex, stateCase := range states { + t.Run(classCase.name+"/"+stateCase.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(791 + classIndex*10 + stateIndex)) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + repo.beforeConditionalState = func() { stateCase.mutate(account) } + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})} + + token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, classCase.class) + + require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged) + require.Empty(t, token) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + } + } +} + +func TestCredentialFailureCASMissDoesNotRecoverIneligibleLatestCredential(t *testing.T) { + future := time.Now().Add(time.Hour) + states := []struct { + name string + mutate func(*OpenAIGatewayService, *Account) + wantRuntimeBlocked bool + }{ + {name: "disabled", mutate: func(_ *OpenAIGatewayService, account *Account) { account.Status = StatusDisabled }}, + {name: "not schedulable", mutate: func(_ *OpenAIGatewayService, account *Account) { account.Schedulable = false }}, + {name: "temporarily unschedulable", mutate: func(_ *OpenAIGatewayService, account *Account) { account.TempUnschedulableUntil = &future }}, + {name: "rate limited", mutate: func(_ *OpenAIGatewayService, account *Account) { account.RateLimitResetAt = &future }}, + {name: "overloaded", mutate: func(_ *OpenAIGatewayService, account *Account) { account.OverloadUntil = &future }}, + { + name: "independently runtime blocked", + mutate: func(svc *OpenAIGatewayService, account *Account) { + svc.BlockAccountScheduling(account, time.Now().Add(24*time.Hour), "independent") + }, + wantRuntimeBlocked: true, + }, + } + classes := []struct { + name string + class grokCredentialFailureClass + }{ + {name: "permanent", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}}, + {name: "transient", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}}, + } + + for classIndex, classCase := range classes { + for stateIndex, stateCase := range states { + t.Run(classCase.name+"/"+stateCase.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(800 + classIndex*20 + stateIndex)) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})} + repo.beforeConditionalState = func() { + latest := *account + latest.Credentials = shallowCopyMap(account.Credentials) + latest.Credentials["access_token"] = "fresh-but-ineligible-token" + latest.Credentials["refresh_token"] = "fresh-but-ineligible-refresh" + latest.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + latest.Credentials["_token_version"] = time.Now().UnixMilli() + stateCase.mutate(svc, &latest) + repo.accountsByID[account.ID] = &latest + } + + token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, classCase.class) + + require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged) + require.Empty(t, token) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Equal(t, stateCase.wantRuntimeBlocked, svc.isOpenAIAccountRuntimeBlocked(account)) + }) + } + } +} + +func TestGetRequestCredentialSharedCredentialPersistenceFailureStopsWithoutAccountMutation(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(782) + repo := &tokenRefreshAccountRepo{updateErr: errors.New("database unavailable")} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{credentials: map[string]any{ + "access_token": "new-access-token", + "refresh_token": "new-refresh-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }}) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + _, _, err := svc.getRequestCredential(context.Background(), c, account) + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope) + require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason) + require.Equal(t, NextAccountStop, failoverErr.NextAccountAction) + require.Equal(t, 1, repo.updateCredentialsCalls) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestGetRequestCredentialRecoversConcurrentRefreshWithoutFailover(t *testing.T) { + gin.SetMode(gin.TestMode) + account := expiredGrokOAuthAccountForCredentialTest(733) + latest := *account + latest.Credentials = shallowCopyMap(account.Credentials) + latest.Credentials["access_token"] = "fresh-concurrent-access" + latest.Credentials["refresh_token"] = "fresh-concurrent-refresh" + latest.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + latest.Credentials["_token_version"] = time.Now().UnixMilli() + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialSequencedRepo{tokenRefreshAccountRepo: baseRepo, latest: &latest} + cache := &grokTokenCacheForProviderTest{lockResult: true} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{ + err: infraerrors.New(http.StatusForbidden, "GROK_OAUTH_ENTITLEMENT_DENIED", "access_denied"), + }) + svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + token, kind, err := svc.getRequestCredential(context.Background(), c, account) + + require.NoError(t, err) + require.Equal(t, "fresh-concurrent-access", token) + require.Equal(t, "oauth", kind) + require.Zero(t, baseRepo.setErrorCalls) + require.Zero(t, baseRepo.setTempUnschedCalls) + require.Empty(t, cache.deletedKeys) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + _, hasEvents := c.Get(OpsUpstreamErrorsKey) + require.False(t, hasEvents) +} + +func expiredGrokOAuthAccountForCredentialTest(id int64) *Account { + return &Account{ + ID: id, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Credentials: map[string]any{ + "access_token": "expired-access-token", + "refresh_token": "refresh-token", + "expires_at": time.Now().Add(-time.Minute).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, + }, + } +} diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 100d720659c..5cc8e2d05e9 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -304,7 +304,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( return nil, fmt.Errorf("account platform %s is not supported for grok media", account.Platform) } - token, _, err := s.GetAccessToken(ctx, account) + token, _, err := s.getRequestCredential(ctx, c, account) if err != nil { return nil, err } @@ -339,7 +339,9 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( } upstreamReq.Header.Set("Authorization", "Bearer "+token) upstreamReq.Header.Set("Accept", "application/json") - applyGrokCLIHeaders(upstreamReq.Header) + if account.IsGrokOAuth() { + applyGrokCLIHeaders(upstreamReq.Header) + } if endpoint.RequiresRequestBody() { contentType = strings.TrimSpace(contentType) if contentType == "" { @@ -635,6 +637,7 @@ func (s *OpenAIGatewayService) handleGrokMediaErrorResponse( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: body, + ResponseHeaders: resp.Header.Clone(), RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } diff --git a/backend/internal/service/grok_oauth_service.go b/backend/internal/service/grok_oauth_service.go index 136e5e3b21d..49e77bcfd77 100644 --- a/backend/internal/service/grok_oauth_service.go +++ b/backend/internal/service/grok_oauth_service.go @@ -3,6 +3,9 @@ package service import ( "context" "crypto/subtle" + "encoding/base64" + "encoding/json" + "errors" "net/http" "strings" "time" @@ -314,10 +317,13 @@ func (s *GrokOAuthService) proxyURL(ctx context.Context, proxyID *int64) (string } proxy, err := s.proxyRepo.GetByID(ctx, *proxyID) if err != nil { - return "", infraerrors.Newf(http.StatusBadRequest, "GROK_OAUTH_PROXY_NOT_FOUND", "proxy not found: %v", err) + if errors.Is(err, ErrProxyNotFound) { + return "", infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_PROXY_NOT_FOUND", "configured proxy was not found") + } + return "", infraerrors.New(http.StatusServiceUnavailable, "GROK_OAUTH_PROXY_LOOKUP_FAILED", "proxy lookup is temporarily unavailable") } if proxy == nil { - return "", nil + return "", infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_PROXY_NOT_FOUND", "configured proxy was not found") } return proxy.URL(), nil } diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 2219a75c155..ef8e3e4c035 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -149,7 +149,9 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json") - applyGrokCLIHeaders(req.Header) + if account.IsGrokOAuth() { + applyGrokCLIHeaders(req.Header) + } resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 1)) if err != nil { @@ -387,6 +389,7 @@ func (s *GrokQuotaService) prepareProbe(ctx context.Context, accountID int64) (* if err != nil { return nil, "", "", err } + proxyURL := s.resolveProxyURL(ctx, account) token, err := s.tokenProvider.GetAccessToken(ctx, account) if err != nil { @@ -396,7 +399,7 @@ func (s *GrokQuotaService) prepareProbe(ctx context.Context, accountID int64) (* return nil, "", "", infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_TOKEN_UNAVAILABLE", "access token is empty") } - return account, token, s.resolveProxyURL(ctx, account), nil + return account, token, proxyURL, nil } func (s *GrokQuotaService) resolveProxyURL(ctx context.Context, account *Account) string { @@ -408,6 +411,7 @@ func (s *GrokQuotaService) resolveProxyURL(ctx context.Context, account *Account return account.Proxy.URL() case s != nil && s.proxyRepo != nil: if proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && proxy != nil { + account.Proxy = proxy return proxy.URL() } } diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index ba9b6cbcb82..4e06fe156ce 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -180,19 +180,26 @@ func (r *grokQuotaProxyRepo) GetByID(_ context.Context, id int64) (*Proxy, error return r.proxies[id], nil } -func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { - t.Parallel() - - account := &Account{ - ID: 42, +func healthyGrokQuotaOAuthAccount(id int64) *Account { + return &Account{ + ID: id, Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_at": time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339), }, } +} + +func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { + t.Parallel() + + account := healthyGrokQuotaOAuthAccount(42) repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{42: account}, @@ -236,19 +243,10 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { func TestGrokQuotaServiceProbeUsageIgnoresAccountGrokMapping(t *testing.T) { t.Parallel() - account := &Account{ - ID: 47, - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "model_mapping": map[string]any{ - "grok": "grok-composer", - "grok-composer": "grok-composer-2.5-fast", - }, - }, + account := healthyGrokQuotaOAuthAccount(47) + account.Credentials["model_mapping"] = map[string]any{ + "grok": "grok-composer", + "grok-composer": "grok-composer-2.5-fast", } repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ @@ -272,16 +270,7 @@ func TestGrokQuotaServiceProbeUsageIgnoresAccountGrokMapping(t *testing.T) { func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T) { t.Parallel() - account := &Account{ - ID: 48, - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(48) repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{48: account}, @@ -352,17 +341,8 @@ func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T t.Parallel() proxyID := int64(7) - account := &Account{ - ID: 46, - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - ProxyID: &proxyID, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(46) + account.ProxyID = &proxyID repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{46: account}, @@ -394,16 +374,7 @@ func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) { t.Parallel() - account := &Account{ - ID: 45, - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(45) repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{45: account}, @@ -435,15 +406,7 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) { func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { t.Parallel() - account := &Account{ - ID: 43, - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(43) repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{43: account}, diff --git a/backend/internal/service/grok_token_provider.go b/backend/internal/service/grok_token_provider.go index f19ee42b889..c4dc4d4b12c 100644 --- a/backend/internal/service/grok_token_provider.go +++ b/backend/internal/service/grok_token_provider.go @@ -3,19 +3,26 @@ package service import ( "context" "errors" - "log/slog" + "fmt" "strconv" "strings" "time" - - "github.com/Wei-Shaw/sub2api/internal/util/logredact" ) const ( - grokTokenCacheSkew = 5 * time.Minute - grokRequestRefreshTimeout = 8 * time.Second - grokTokenProviderLogComponent = "grok_token_provider" - grokTempUnschedulableErrorCode = "token_refresh_failed" + grokTokenCacheSkew = 5 * time.Minute + grokRequestRefreshTimeout = 8 * time.Second + grokRefreshLockWaitTimeout = 2 * time.Second + grokRefreshLockPollInterval = 25 * time.Millisecond +) + +var ( + errGrokOAuthRefreshNotConfigured = errors.New("grok oauth refresh is not configured") + errGrokOAuthRefreshTokenMissing = errors.New("grok oauth refresh token is missing") + errGrokOAuthAccessTokenMissing = errors.New("grok oauth access token is missing") + errGrokOAuthAccessTokenExpired = errors.New("grok oauth access token is expired") + errGrokOAuthConfiguredProxyMiss = errors.New("grok oauth configured proxy is missing") + errGrokOAuthRefreshLockTimeout = errors.New("grok oauth refresh lock wait timed out") ) type GrokTokenCache = GeminiTokenCache @@ -36,7 +43,7 @@ func NewGrokTokenProvider( return &GrokTokenProvider{ accountRepo: accountRepo, tokenCache: tokenCache, - refreshPolicy: AntigravityProviderRefreshPolicy(), + refreshPolicy: GrokProviderRefreshPolicy(), } } @@ -60,32 +67,57 @@ func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account if account.Platform != PlatformGrok || account.Type != AccountTypeOAuth { return "", errors.New("not a grok oauth account") } + selectedProxyID := cloneGrokProxyID(account.ProxyID) + if eligibilityErr := grokOAuthRequestAccountEligibilityError(account); eligibilityErr != nil { + return "", withGrokCredentialFailureSnapshot(eligibilityErr, account) + } + expiresAt := account.GetCredentialAsTime("expires_at") + accountAccessToken := strings.TrimSpace(account.GetGrokAccessToken()) + if accountAccessToken == "" { + return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenMissing, account) + } + if strings.TrimSpace(account.GetGrokRefreshToken()) == "" { + return "", withGrokCredentialFailureSnapshot(errGrokOAuthRefreshTokenMissing, account) + } cacheKey := GrokTokenCacheKey(account) if p.tokenCache != nil { - if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && strings.TrimSpace(token) != "" { - return token, nil + if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil { + cachedToken := strings.TrimSpace(token) + if cachedToken != "" && accountAccessToken != "" && cachedToken == accountAccessToken && + expiresAt != nil && time.Until(*expiresAt) > grokTokenRefreshSkew { + return cachedToken, nil + } } } - expiresAt := account.GetCredentialAsTime("expires_at") needsRefresh := expiresAt == nil || time.Until(*expiresAt) <= grokTokenRefreshSkew - if needsRefresh && strings.TrimSpace(account.GetGrokRefreshToken()) == "" { - if expiresAt == nil || !time.Now().Before(*expiresAt) { - return "", errors.New("grok access_token expired and refresh_token is missing") + if needsRefresh { + if p.refreshAPI == nil || p.executor == nil { + return "", errGrokOAuthRefreshNotConfigured } - needsRefresh = false - } - if needsRefresh && p.refreshAPI != nil && p.executor != nil { refreshCtx, cancel := context.WithTimeout(ctx, grokRequestRefreshTimeout) defer cancel() - result, err := p.refreshAPI.RefreshIfNeeded(refreshCtx, account, p.executor, grokTokenRefreshSkew) + result, err := p.refreshAPI.RefreshIfNeeded(withOAuthRefreshRequestPath(refreshCtx), account, p.executor, grokTokenRefreshSkew) if err != nil { - p.markTempUnschedulable(account, err) if p.refreshPolicy.OnRefreshError == ProviderRefreshErrorReturn { return "", err } - } else if !result.LockHeld && result.Account != nil { + } else if result != nil && result.LockHeld { + if p.refreshPolicy.OnLockHeld == ProviderLockHeldWaitForCache { + token, waitErr := p.waitForRefreshedToken(refreshCtx, account, cacheKey) + return token, withGrokCredentialFailureSnapshot(waitErr, account) + } + if expiresAt == nil || !time.Now().Before(*expiresAt) { + return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenExpired, account) + } + } else if result != nil && result.Account != nil { + if eligibilityErr := grokOAuthRequestAccountEligibilityError(result.Account); eligibilityErr != nil { + return "", withGrokCredentialFailureSnapshot(eligibilityErr, result.Account) + } + if !grokCredentialProxyIDsEqual(result.Account.ProxyID, selectedProxyID) { + return "", withGrokCredentialFailureSnapshot(errOAuthRefreshAccountStateChanged, result.Account) + } account = result.Account expiresAt = account.GetCredentialAsTime("expires_at") } @@ -93,15 +125,28 @@ func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account accessToken := account.GetGrokAccessToken() if strings.TrimSpace(accessToken) == "" { - return "", errors.New("access_token not found in credentials") + return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenMissing, account) + } + if expiresAt != nil && !time.Now().Before(*expiresAt) { + return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenExpired, account) } if p.tokenCache != nil { latestAccount, isStale := CheckTokenVersion(ctx, account, p.accountRepo) if isStale && latestAccount != nil { + if eligibilityErr := grokOAuthRequestAccountEligibilityError(latestAccount); eligibilityErr != nil { + return "", withGrokCredentialFailureSnapshot(eligibilityErr, latestAccount) + } + if !grokCredentialProxyIDsEqual(latestAccount.ProxyID, selectedProxyID) { + return "", withGrokCredentialFailureSnapshot(errOAuthRefreshAccountStateChanged, latestAccount) + } accessToken = latestAccount.GetGrokAccessToken() if strings.TrimSpace(accessToken) == "" { - return "", errors.New("access_token not found after version check") + return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenMissing, latestAccount) + } + latestExpiry := latestAccount.GetCredentialAsTime("expires_at") + if latestExpiry == nil || !time.Now().Before(*latestExpiry) { + return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenExpired, latestAccount) } } else { ttl := 30 * time.Minute @@ -123,40 +168,105 @@ func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account return accessToken, nil } -func (p *GrokTokenProvider) markTempUnschedulable(account *Account, refreshErr error) { - if p == nil || p.accountRepo == nil || account == nil { - return - } - now := time.Now() - until := now.Add(tokenRefreshTempUnschedDuration) - redactedErr := "unknown error" - if refreshErr != nil { - redactedErr = logredact.RedactText(refreshErr.Error()) - } - if isNonRetryableRefreshError(refreshErr) { - if err := p.accountRepo.SetError(context.Background(), account.ID, "grok token refresh failed (non-retryable): "+redactedErr); err != nil { - slog.Warn(grokTokenProviderLogComponent+".set_error_status_failed", "account_id", account.ID, "error", err) +func (p *GrokTokenProvider) waitForRefreshedToken(ctx context.Context, account *Account, cacheKey string) (string, error) { + waitCtx, cancel := context.WithTimeout(ctx, grokRefreshLockWaitTimeout) + defer cancel() + + initialToken := strings.TrimSpace(account.GetGrokAccessToken()) + initialVersion := account.GetCredentialAsInt64("_token_version") + selectedProxyID := cloneGrokProxyID(account.ProxyID) + sawAuthoritativeState := false + var lastAccountReadErr error + ticker := time.NewTicker(grokRefreshLockPollInterval) + defer ticker.Stop() + + for { + cachedToken := "" + if p.tokenCache != nil { + if token, err := p.tokenCache.GetAccessToken(waitCtx, cacheKey); err == nil { + cachedToken = strings.TrimSpace(token) + } } - return - } - reason := "grok token refresh failed on request path: " + redactedErr - bgCtx := context.Background() - if err := p.accountRepo.SetTempUnschedulable(bgCtx, account.ID, until, reason); err != nil { - slog.Warn(grokTokenProviderLogComponent+".set_temp_unschedulable_failed", "account_id", account.ID, "error", err) - return - } - if p.tempUnschedCache != nil { - state := &TempUnschedState{ - UntilUnix: until.Unix(), - TriggeredAtUnix: now.Unix(), - ErrorMessage: grokTempUnschedulableErrorCode + ": " + reason, + + if p.accountRepo != nil { + latest, err := p.accountRepo.GetByID(waitCtx, account.ID) + if err != nil { + lastAccountReadErr = err + } else if latest == nil { + return "", errOAuthRefreshAccountStateChanged + } else { + sawAuthoritativeState = true + if eligibilityErr := grokOAuthRequestAccountEligibilityError(latest); eligibilityErr != nil { + return "", withGrokCredentialFailureSnapshot(eligibilityErr, latest) + } + if !grokCredentialProxyIDsEqual(latest.ProxyID, selectedProxyID) { + return "", withGrokCredentialFailureSnapshot(errOAuthRefreshAccountStateChanged, latest) + } + token := strings.TrimSpace(latest.GetGrokAccessToken()) + version := latest.GetCredentialAsInt64("_token_version") + expiresAt := latest.GetCredentialAsTime("expires_at") + changed := token != initialToken || (version > 0 && version > initialVersion) + valid := expiresAt != nil && time.Now().Before(*expiresAt) + if token != "" && changed && valid { + // The versioned DB credential is authoritative. A stale cache must + // not hold the request on the old expired token; repair it best-effort. + if cachedToken != "" && cachedToken != token { + ttl := time.Until(*expiresAt) + if ttl > grokTokenCacheSkew { + ttl -= grokTokenCacheSkew + } + _ = p.tokenCache.SetAccessToken(waitCtx, cacheKey, token, ttl) + } + return token, nil + } + } } - if err := p.tempUnschedCache.SetTempUnsched(bgCtx, account.ID, state); err != nil { - slog.Warn(grokTokenProviderLogComponent+".temp_unsched_cache_set_failed", "account_id", account.ID, "error", err) + + select { + case <-waitCtx.Done(): + if ctx.Err() != nil { + return "", ctx.Err() + } + if !sawAuthoritativeState { + if lastAccountReadErr == nil { + lastAccountReadErr = waitCtx.Err() + } + return "", fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, lastAccountReadErr) + } + // Another worker still owns the refresh and the authoritative row is + // unchanged. Do not quarantine the old credential: its refresh may + // commit immediately after this bounded wait. + return "", errOAuthRefreshAccountStateChanged + case <-ticker.C: } } } +func grokOAuthRequestAccountEligibilityError(account *Account) error { + if account == nil || !account.IsGrokOAuth() || !account.IsSchedulable() { + return errOAuthRefreshAccountStateChanged + } + if account.ProxyID != nil && account.Proxy == nil { + return errGrokOAuthConfiguredProxyMiss + } + return nil +} + +func cloneGrokProxyID(proxyID *int64) *int64 { + if proxyID == nil { + return nil + } + value := *proxyID + return &value +} + +func (p *GrokTokenProvider) InvalidateToken(ctx context.Context, account *Account) error { + if p == nil || p.tokenCache == nil || account == nil { + return nil + } + return p.tokenCache.DeleteAccessToken(ctx, GrokTokenCacheKey(account)) +} + func GrokTokenCacheKey(account *Account) string { if account == nil { return "grok:account:0" diff --git a/backend/internal/service/grok_token_provider_test.go b/backend/internal/service/grok_token_provider_test.go index 1f7e2d8e0e0..7c68e55455d 100644 --- a/backend/internal/service/grok_token_provider_test.go +++ b/backend/internal/service/grok_token_provider_test.go @@ -5,6 +5,7 @@ package service import ( "context" "errors" + "sync" "testing" "time" @@ -19,9 +20,33 @@ type grokTokenCacheForProviderTest struct { setTTL time.Duration lockResult bool releaseCalls int + deletedKeys []string + deleteErr error + getCalls int + mu sync.Mutex +} + +type grokCredentialRaceRepo struct { + *tokenRefreshAccountRepo + mu sync.RWMutex +} + +func (r *grokCredentialRaceRepo) GetByID(ctx context.Context, id int64) (*Account, error) { + r.mu.RLock() + defer r.mu.RUnlock() + return r.tokenRefreshAccountRepo.GetByID(ctx, id) +} + +func (r *grokCredentialRaceRepo) setAccount(account *Account) { + r.mu.Lock() + defer r.mu.Unlock() + r.accountsByID[account.ID] = account } func (c *grokTokenCacheForProviderTest) GetAccessToken(context.Context, string) (string, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.getCalls++ if c.token == "" { return "", errors.New("not cached") } @@ -29,14 +54,19 @@ func (c *grokTokenCacheForProviderTest) GetAccessToken(context.Context, string) } func (c *grokTokenCacheForProviderTest) SetAccessToken(_ context.Context, key string, token string, ttl time.Duration) error { + c.mu.Lock() + defer c.mu.Unlock() c.setKey = key c.setToken = token c.setTTL = ttl return nil } -func (c *grokTokenCacheForProviderTest) DeleteAccessToken(context.Context, string) error { - return nil +func (c *grokTokenCacheForProviderTest) DeleteAccessToken(_ context.Context, key string) error { + c.mu.Lock() + defer c.mu.Unlock() + c.deletedKeys = append(c.deletedKeys, key) + return c.deleteErr } func (c *grokTokenCacheForProviderTest) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) { @@ -53,9 +83,11 @@ func TestGrokTokenProviderRefreshesExpiredTokenOnRequestPath(t *testing.T) { expiredAt := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339) account := &Account{ - ID: 54, - Platform: PlatformGrok, - Type: AccountTypeOAuth, + ID: 54, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, Credentials: map[string]any{ "access_token": "expired-access-token", "refresh_token": "refresh-token", @@ -95,9 +127,11 @@ func TestGrokTokenProviderRefreshesExpiredTokenOnRequestPath(t *testing.T) { func TestGrokTokenProviderRefreshFailureUnschedulesWithRedactedReason(t *testing.T) { expiredAt := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339) account := &Account{ - ID: 55, - Platform: PlatformGrok, - Type: AccountTypeOAuth, + ID: 55, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, Credentials: map[string]any{ "access_token": "expired-access-token", "refresh_token": "refresh-token", @@ -108,24 +142,178 @@ func TestGrokTokenProviderRefreshFailureUnschedulesWithRedactedReason(t *testing repo := &tokenRefreshAccountRepo{} repo.accountsByID = map[int64]*Account{55: account} cache := &grokTokenCacheForProviderTest{lockResult: true} - tempCache := &tempUnschedCacheStub{} provider := NewGrokTokenProvider(repo, cache) provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{ err: errors.New("temporary refresh failure access_token=leaked-access refresh_token=leaked-refresh"), }) - provider.SetTempUnschedCache(tempCache) token, err := provider.GetAccessToken(context.Background(), account) require.Error(t, err) require.Empty(t, token) - require.Equal(t, 1, repo.setTempUnschedCalls) + require.Equal(t, 0, repo.setTempUnschedCalls) require.Equal(t, 0, repo.setErrorCalls) - require.Contains(t, repo.lastTempUnschedReason, "access_token=***") - require.Contains(t, repo.lastTempUnschedReason, "refresh_token=***") - require.NotContains(t, repo.lastTempUnschedReason, "leaked-access") - require.NotContains(t, repo.lastTempUnschedReason, "leaked-refresh") - require.Equal(t, 1, tempCache.setCalls) - require.NotNil(t, tempCache.lastState) - require.NotContains(t, tempCache.lastState.ErrorMessage, "leaked-access") - require.NotContains(t, tempCache.lastState.ErrorMessage, "leaked-refresh") +} + +func TestGrokTokenProviderLockHeldWaitsForRefreshedCacheAndNeverUsesExpiredToken(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(56) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialRaceRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: false, token: "expired-access-token"} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{}) + + go func() { + time.Sleep(40 * time.Millisecond) + refreshed := *account + refreshed.Credentials = shallowCopyMap(account.Credentials) + refreshed.Credentials["access_token"] = "refreshed-after-lock" + refreshed.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339) + refreshed.Credentials["_token_version"] = time.Now().UnixMilli() + repo.setAccount(&refreshed) + cache.mu.Lock() + cache.token = "refreshed-after-lock" + cache.mu.Unlock() + }() + + startedAt := time.Now() + token, err := provider.GetAccessToken(context.Background(), account) + require.NoError(t, err) + require.Equal(t, "refreshed-after-lock", token) + require.NotEqual(t, "expired-access-token", token) + require.GreaterOrEqual(t, time.Since(startedAt), 25*time.Millisecond, + "expired account metadata must prevent returning the old cached token") +} + +func TestGrokTokenProviderLockHeldTimeoutDoesNotReturnExpiredToken(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(57) + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: account} + cache := &grokTokenCacheForProviderTest{lockResult: false} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{}) + ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond) + defer cancel() + + token, err := provider.GetAccessToken(ctx, account) + require.Error(t, err) + require.Empty(t, token) +} + +func TestGrokTokenProviderLockHeldRejectsChangedTokenWithoutExpiry(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(58) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialRaceRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: false, token: "expired-access-token"} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{}) + + go func() { + time.Sleep(30 * time.Millisecond) + refreshed := *account + refreshed.Credentials = shallowCopyMap(account.Credentials) + refreshed.Credentials["access_token"] = "changed-without-expiry" + delete(refreshed.Credentials, "expires_at") + refreshed.Credentials["_token_version"] = time.Now().UnixMilli() + repo.setAccount(&refreshed) + cache.mu.Lock() + cache.token = "changed-without-expiry" + cache.mu.Unlock() + }() + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + token, err := provider.GetAccessToken(ctx, account) + + require.Error(t, err) + require.Empty(t, token, "an unbounded credential must not win the lock-held race") +} + +func TestGrokTokenProviderLockHeldUsesVersionedDBTokenAndRepairsStaleCache(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(60) + baseRepo := &tokenRefreshAccountRepo{} + baseRepo.accountsByID = map[int64]*Account{account.ID: account} + repo := &grokCredentialRaceRepo{tokenRefreshAccountRepo: baseRepo} + cache := &grokTokenCacheForProviderTest{lockResult: false, token: "expired-access-token"} + provider := NewGrokTokenProvider(repo, cache) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{}) + + go func() { + time.Sleep(30 * time.Millisecond) + refreshed := *account + refreshed.Credentials = shallowCopyMap(account.Credentials) + refreshed.Credentials["access_token"] = "db-authoritative-token" + refreshed.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + refreshed.Credentials["_token_version"] = time.Now().UnixMilli() + repo.setAccount(&refreshed) + }() + + token, err := provider.GetAccessToken(context.Background(), account) + + require.NoError(t, err) + require.Equal(t, "db-authoritative-token", token) + require.Equal(t, "db-authoritative-token", cache.setToken) + require.Greater(t, cache.setTTL, time.Duration(0)) +} + +func TestGrokTokenProviderRejectsStaleDBTokenWithoutExpiry(t *testing.T) { + expiresAt := time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339) + account := &Account{ + ID: 59, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Credentials: map[string]any{ + "access_token": "old-access-token", + "refresh_token": "refresh-token", + "expires_at": expiresAt, + }, + } + latest := *account + latest.Credentials = shallowCopyMap(account.Credentials) + latest.Credentials["access_token"] = "new-access-token-without-expiry" + latest.Credentials["_token_version"] = time.Now().UnixMilli() + delete(latest.Credentials, "expires_at") + repo := &tokenRefreshAccountRepo{} + repo.accountsByID = map[int64]*Account{account.ID: &latest} + cache := &grokTokenCacheForProviderTest{} + provider := NewGrokTokenProvider(repo, cache) + + token, err := provider.GetAccessToken(context.Background(), account) + + require.ErrorIs(t, err, errGrokOAuthAccessTokenExpired) + require.Empty(t, token) +} + +func TestGrokTokenProviderRejectsIneligibleSelectedAccountBeforeWarmCache(t *testing.T) { + future := time.Now().Add(time.Hour) + tests := []struct { + name string + mutate func(*Account) + }{ + {name: "disabled", mutate: func(account *Account) { account.Status = StatusDisabled }}, + {name: "not schedulable", mutate: func(account *Account) { account.Schedulable = false }}, + {name: "temporarily unschedulable", mutate: func(account *Account) { account.TempUnschedulableUntil = &future }}, + {name: "rate limited", mutate: func(account *Account) { account.RateLimitResetAt = &future }}, + {name: "overloaded", mutate: func(account *Account) { account.OverloadUntil = &future }}, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := expiredGrokOAuthAccountForCredentialTest(int64(90 + index)) + account.Credentials["access_token"] = "warm-cache-token" + account.Credentials["expires_at"] = time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339) + tt.mutate(account) + cache := &grokTokenCacheForProviderTest{token: "warm-cache-token"} + provider := NewGrokTokenProvider(&tokenRefreshAccountRepo{}, cache) + + token, err := provider.GetAccessToken(context.Background(), account) + + require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged) + require.Empty(t, token) + require.Zero(t, cache.getCalls, "an ineligible selected account must be rejected before cache lookup") + }) + } } diff --git a/backend/internal/service/grok_token_refresher.go b/backend/internal/service/grok_token_refresher.go index 92d88cc058c..0b87fcac781 100644 --- a/backend/internal/service/grok_token_refresher.go +++ b/backend/internal/service/grok_token_refresher.go @@ -22,7 +22,8 @@ func (r *GrokTokenRefresher) CacheKey(account *Account) string { } func (r *GrokTokenRefresher) CanRefresh(account *Account) bool { - return account != nil && account.Platform == PlatformGrok && account.Type == AccountTypeOAuth + return account != nil && account.Platform == PlatformGrok && account.Type == AccountTypeOAuth && + strings.TrimSpace(account.GetGrokRefreshToken()) != "" } func (r *GrokTokenRefresher) NeedsRefresh(account *Account, refreshWindow time.Duration) bool { diff --git a/backend/internal/service/oauth_refresh_api.go b/backend/internal/service/oauth_refresh_api.go index 571e9ecdb7c..b80cf84b645 100644 --- a/backend/internal/service/oauth_refresh_api.go +++ b/backend/internal/service/oauth_refresh_api.go @@ -2,6 +2,7 @@ package service import ( "context" + "errors" "fmt" "log/slog" "strconv" @@ -20,6 +21,45 @@ type OAuthRefreshExecutor interface { } const defaultRefreshLockTTL = 60 * time.Second +const oauthRefreshLockCleanupTimeout = 2 * time.Second + +var ( + errOAuthRefreshAccountRereadFailed = errors.New("oauth refresh account reread failed") + errOAuthRefreshAccountStateChanged = errors.New("oauth refresh account state changed") + errOAuthRefreshCredentialPersist = errors.New("oauth refresh credential persistence failed") +) + +type oauthRefreshRequestPathKey struct{} + +func withOAuthRefreshRequestPath(ctx context.Context) context.Context { + return context.WithValue(ctx, oauthRefreshRequestPathKey{}, true) +} + +func isOAuthRefreshRequestPath(ctx context.Context) bool { + requestPath, _ := ctx.Value(oauthRefreshRequestPathKey{}).(bool) + return requestPath +} + +type oauthRefreshLocalLock struct { + semaphore chan struct{} +} + +func newOAuthRefreshLocalLock() *oauthRefreshLocalLock { + return &oauthRefreshLocalLock{semaphore: make(chan struct{}, 1)} +} + +func (l *oauthRefreshLocalLock) Lock(ctx context.Context) error { + select { + case l.semaphore <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (l *oauthRefreshLocalLock) Unlock() { + <-l.semaphore +} // OAuthRefreshResult 统一刷新结果 type OAuthRefreshResult struct { @@ -35,7 +75,7 @@ type OAuthRefreshAPI struct { accountRepo AccountRepository tokenCache GeminiTokenCache // 可选，nil = 无分布式锁 lockTTL time.Duration - localLocks sync.Map // key: cacheKey string -> value: *sync.Mutex + localLocks sync.Map // key: cacheKey string -> value: *oauthRefreshLocalLock } // NewOAuthRefreshAPI 创建统一刷新 API @@ -53,11 +93,11 @@ func NewOAuthRefreshAPI(accountRepo AccountRepository, tokenCache GeminiTokenCac } // getLocalLock 返回指定 cacheKey 的进程内互斥锁 -func (api *OAuthRefreshAPI) getLocalLock(cacheKey string) *sync.Mutex { - actual, _ := api.localLocks.LoadOrStore(cacheKey, &sync.Mutex{}) - mu, ok := actual.(*sync.Mutex) +func (api *OAuthRefreshAPI) getLocalLock(cacheKey string) *oauthRefreshLocalLock { + actual, _ := api.localLocks.LoadOrStore(cacheKey, newOAuthRefreshLocalLock()) + mu, ok := actual.(*oauthRefreshLocalLock) if !ok { - mu = &sync.Mutex{} + mu = newOAuthRefreshLocalLock() api.localLocks.Store(cacheKey, mu) } return mu @@ -78,15 +118,25 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded( executor OAuthRefreshExecutor, refreshWindow time.Duration, ) (*OAuthRefreshResult, error) { + if api == nil || api.accountRepo == nil { + return nil, errors.New("oauth refresh account repository is not configured") + } + if account == nil { + return nil, errors.New("oauth refresh account is nil") + } + if executor == nil { + return nil, errors.New("oauth refresh executor is nil") + } cacheKey := executor.CacheKey(account) // 0. 获取进程内互斥锁（防止同一进程内的并发刷新竞争） localMu := api.getLocalLock(cacheKey) - localMu.Lock() + if err := localMu.Lock(ctx); err != nil { + return nil, fmt.Errorf("oauth refresh local lock: %w", err) + } defer localMu.Unlock() // 1. 获取分布式锁 - lockAcquired := false if api.tokenCache != nil { acquired, lockErr := api.tokenCache.AcquireRefreshLock(ctx, cacheKey, api.lockTTL) if lockErr != nil { @@ -100,22 +150,38 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded( // 锁被其他 worker 持有 return &OAuthRefreshResult{LockHeld: true}, nil } else { - lockAcquired = true - defer func() { _ = api.tokenCache.ReleaseRefreshLock(ctx, cacheKey) }() + defer func() { + cleanupCtx, cancel := context.WithTimeout(context.Background(), oauthRefreshLockCleanupTimeout) + defer cancel() + _ = api.tokenCache.ReleaseRefreshLock(cleanupCtx, cacheKey) + }() } } // 2. 从 DB 重读最新 account（锁保护下，确保使用最新的 refresh_token） freshAccount, err := api.accountRepo.GetByID(ctx, account.ID) if err != nil { - slog.Warn("oauth_refresh_db_reread_failed", - "account_id", account.ID, - "error", err, - ) - // 降级使用传入的 account - freshAccount = account - } else if freshAccount == nil { - freshAccount = account + return nil, fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, err) + } + if freshAccount == nil { + return nil, fmt.Errorf("%w: account not found", errOAuthRefreshAccountStateChanged) + } + if freshAccount.ID != account.ID { + return nil, fmt.Errorf("%w: account identity mismatch", errOAuthRefreshAccountRereadFailed) + } + if !freshAccount.IsActive() { + return nil, fmt.Errorf("%w: account is not active", errOAuthRefreshAccountStateChanged) + } + if isOAuthRefreshRequestPath(ctx) && freshAccount.Platform == PlatformGrok { + if eligibilityErr := grokOAuthRequestAccountEligibilityError(freshAccount); eligibilityErr != nil { + return nil, withGrokCredentialFailureSnapshot(eligibilityErr, freshAccount) + } + } + if !executor.CanRefresh(freshAccount) { + if freshAccount.IsGrokOAuth() && strings.TrimSpace(freshAccount.GetGrokRefreshToken()) == "" { + return nil, withGrokCredentialFailureSnapshot(errGrokOAuthRefreshTokenMissing, freshAccount) + } + return nil, fmt.Errorf("%w: account is no longer refreshable", errOAuthRefreshAccountStateChanged) } // 3. 二次检查是否仍需刷新（另一条路径可能已刷新） @@ -127,11 +193,19 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded( // 4. 执行平台特定刷新逻辑 newCredentials, refreshErr := executor.Refresh(ctx, freshAccount) + if err := ctx.Err(); err != nil { + return nil, err + } if refreshErr != nil { // 竞争恢复：invalid_grant 可能是另一个 worker 已消费了旧 refresh_token // 重新读取 DB，如果 refresh_token 已更新则说明是竞争，返回成功 if isInvalidGrantError(refreshErr) { if recoveredAccount, recovered := api.tryRecoverFromRefreshRace(ctx, freshAccount); recovered { + if isOAuthRefreshRequestPath(ctx) && recoveredAccount.Platform == PlatformGrok { + if eligibilityErr := grokOAuthRequestAccountEligibilityError(recoveredAccount); eligibilityErr != nil { + return nil, withGrokCredentialFailureSnapshot(eligibilityErr, recoveredAccount) + } + } slog.Info("oauth_refresh_race_recovered", "account_id", freshAccount.ID, "platform", freshAccount.Platform, @@ -141,7 +215,7 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded( }, nil } } - return nil, refreshErr + return nil, withGrokCredentialFailureSnapshot(refreshErr, freshAccount) } // 5. 设置版本号 + 更新 DB @@ -152,16 +226,30 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded( "account_id", freshAccount.ID, "error", updateErr, ) - return nil, fmt.Errorf("oauth refresh succeeded but DB update failed: %w", updateErr) + return nil, withGrokCredentialFailureSnapshot( + fmt.Errorf("%w: %v", errOAuthRefreshCredentialPersist, updateErr), freshAccount, + ) } } - - _ = lockAcquired // suppress unused warning when tokenCache is nil + resultAccount := freshAccount + if isOAuthRefreshRequestPath(ctx) && freshAccount.Platform == PlatformGrok { + latestAccount, rereadErr := api.accountRepo.GetByID(ctx, freshAccount.ID) + if rereadErr != nil { + return nil, fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, rereadErr) + } + if latestAccount == nil { + return nil, fmt.Errorf("%w: account not found after refresh", errOAuthRefreshAccountStateChanged) + } + if eligibilityErr := grokOAuthRequestAccountEligibilityError(latestAccount); eligibilityErr != nil { + return nil, withGrokCredentialFailureSnapshot(eligibilityErr, latestAccount) + } + resultAccount = latestAccount + } return &OAuthRefreshResult{ Refreshed: true, NewCredentials: newCredentials, - Account: freshAccount, + Account: resultAccount, }, nil } diff --git a/backend/internal/service/oauth_refresh_api_test.go b/backend/internal/service/oauth_refresh_api_test.go index cacffddc595..adf089f0bd9 100644 --- a/backend/internal/service/oauth_refresh_api_test.go +++ b/backend/internal/service/oauth_refresh_api_test.go @@ -51,13 +51,14 @@ func (r *refreshAPIAccountRepo) UpdateCredentials(_ context.Context, id int64, c // refreshAPIExecutorStub implements OAuthRefreshExecutor for tests. type refreshAPIExecutorStub struct { - needsRefresh bool - credentials map[string]any - err error - refreshCalls int + needsRefresh bool + cannotRefresh bool + credentials map[string]any + err error + refreshCalls int } -func (e *refreshAPIExecutorStub) CanRefresh(_ *Account) bool { return true } +func (e *refreshAPIExecutorStub) CanRefresh(_ *Account) bool { return !e.cannotRefresh } func (e *refreshAPIExecutorStub) NeedsRefresh(_ *Account, _ time.Duration) bool { return e.needsRefresh @@ -77,9 +78,10 @@ func (e *refreshAPIExecutorStub) CacheKey(account *Account) string { // refreshAPICacheStub implements GeminiTokenCache for OAuthRefreshAPI tests. type refreshAPICacheStub struct { - lockResult bool - lockErr error - releaseCalls int + lockResult bool + lockErr error + releaseCalls int + releaseCtxErr error } func (c *refreshAPICacheStub) GetAccessToken(context.Context, string) (string, error) { @@ -96,15 +98,16 @@ func (c *refreshAPICacheStub) AcquireRefreshLock(context.Context, string, time.D return c.lockResult, c.lockErr } -func (c *refreshAPICacheStub) ReleaseRefreshLock(context.Context, string) error { +func (c *refreshAPICacheStub) ReleaseRefreshLock(ctx context.Context, _ string) error { c.releaseCalls++ + c.releaseCtxErr = ctx.Err() return nil } // ========== RefreshIfNeeded tests ========== func TestRefreshIfNeeded_Success(t *testing.T) { - account := &Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeOAuth} + account := &Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive} repo := &refreshAPIAccountRepo{account: account} cache := &refreshAPICacheStub{lockResult: true} executor := &refreshAPIExecutorStub{ @@ -132,6 +135,7 @@ func TestRefreshIfNeeded_UpdateCredentialsPreservesRateLimitState(t *testing.T) ID: 11, Platform: PlatformGemini, Type: AccountTypeOAuth, + Status: StatusActive, RateLimitResetAt: &resetAt, } repo := &refreshAPIAccountRepo{account: account} @@ -152,7 +156,7 @@ func TestRefreshIfNeeded_UpdateCredentialsPreservesRateLimitState(t *testing.T) } func TestRefreshIfNeeded_LockHeld(t *testing.T) { - account := &Account{ID: 2, Platform: PlatformAnthropic} + account := &Account{ID: 2, Platform: PlatformAnthropic, Status: StatusActive} repo := &refreshAPIAccountRepo{account: account} cache := &refreshAPICacheStub{lockResult: false} // lock not acquired executor := &refreshAPIExecutorStub{needsRefresh: true} @@ -168,7 +172,7 @@ func TestRefreshIfNeeded_LockHeld(t *testing.T) { } func TestRefreshIfNeeded_LockErrorDegrades(t *testing.T) { - account := &Account{ID: 3, Platform: PlatformGemini, Type: AccountTypeOAuth} + account := &Account{ID: 3, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive} repo := &refreshAPIAccountRepo{account: account} cache := &refreshAPICacheStub{lockErr: errors.New("redis down")} // lock error executor := &refreshAPIExecutorStub{ @@ -187,7 +191,7 @@ func TestRefreshIfNeeded_LockErrorDegrades(t *testing.T) { } func TestRefreshIfNeeded_NoCacheNoLock(t *testing.T) { - account := &Account{ID: 4, Platform: PlatformGemini, Type: AccountTypeOAuth} + account := &Account{ID: 4, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive} repo := &refreshAPIAccountRepo{account: account} executor := &refreshAPIExecutorStub{ needsRefresh: true, @@ -203,7 +207,7 @@ func TestRefreshIfNeeded_NoCacheNoLock(t *testing.T) { } func TestRefreshIfNeeded_AlreadyRefreshed(t *testing.T) { - account := &Account{ID: 5, Platform: PlatformAnthropic} + account := &Account{ID: 5, Platform: PlatformAnthropic, Status: StatusActive} repo := &refreshAPIAccountRepo{account: account} cache := &refreshAPICacheStub{lockResult: true} executor := &refreshAPIExecutorStub{needsRefresh: false} // already refreshed @@ -220,7 +224,7 @@ func TestRefreshIfNeeded_AlreadyRefreshed(t *testing.T) { } func TestRefreshIfNeeded_RefreshError(t *testing.T) { - account := &Account{ID: 6, Platform: PlatformAnthropic} + account := &Account{ID: 6, Platform: PlatformAnthropic, Status: StatusActive} repo := &refreshAPIAccountRepo{account: account} cache := &refreshAPICacheStub{lockResult: true} executor := &refreshAPIExecutorStub{ @@ -239,7 +243,7 @@ func TestRefreshIfNeeded_RefreshError(t *testing.T) { } func TestRefreshIfNeeded_DBUpdateError(t *testing.T) { - account := &Account{ID: 7, Platform: PlatformGemini, Type: AccountTypeOAuth} + account := &Account{ID: 7, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive} repo := &refreshAPIAccountRepo{ account: account, updateErr: errors.New("db connection lost"), @@ -255,12 +259,12 @@ func TestRefreshIfNeeded_DBUpdateError(t *testing.T) { require.Error(t, err) require.Nil(t, result) - require.Contains(t, err.Error(), "DB update failed") + require.ErrorIs(t, err, errOAuthRefreshCredentialPersist) require.Equal(t, 1, repo.updateCalls) // attempted } func TestRefreshIfNeeded_DBRereadFails(t *testing.T) { - account := &Account{ID: 8, Platform: PlatformAnthropic, Type: AccountTypeOAuth} + account := &Account{ID: 8, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive} repo := &refreshAPIAccountRepo{ account: nil, // GetByID returns nil getByIDErr: errors.New("db timeout"), @@ -274,13 +278,95 @@ func TestRefreshIfNeeded_DBRereadFails(t *testing.T) { api := NewOAuthRefreshAPI(repo, cache) result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute) - require.NoError(t, err) - require.True(t, result.Refreshed) - require.Equal(t, 1, executor.refreshCalls) // still refreshes using passed-in account + require.ErrorContains(t, err, "oauth refresh account reread") + require.Nil(t, result) + require.Zero(t, executor.refreshCalls, "must not refresh with the stale caller snapshot") + require.Zero(t, repo.updateCalls) + require.Equal(t, 1, cache.releaseCalls) +} + +func TestRefreshIfNeeded_DBRereadNilFailsClosed(t *testing.T) { + account := &Account{ID: 81, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive} + repo := &refreshAPIAccountRepo{} + cache := &refreshAPICacheStub{lockResult: true} + executor := &refreshAPIExecutorStub{needsRefresh: true} + + api := NewOAuthRefreshAPI(repo, cache) + result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute) + + require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged) + require.Nil(t, result) + require.Zero(t, executor.refreshCalls) + require.Zero(t, repo.updateCalls) + require.Equal(t, 1, cache.releaseCalls) +} + +func TestRefreshIfNeeded_DBRereadInactiveFailsClosed(t *testing.T) { + account := &Account{ID: 82, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive} + freshAccount := &Account{ID: account.ID, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusDisabled} + repo := &refreshAPIAccountRepo{account: freshAccount} + executor := &refreshAPIExecutorStub{needsRefresh: true} + + api := NewOAuthRefreshAPI(repo, nil) + result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute) + + require.ErrorContains(t, err, "account is not active") + require.Nil(t, result) + require.Zero(t, executor.refreshCalls) + require.Zero(t, repo.updateCalls) +} + +func TestRefreshIfNeeded_DBRereadRevalidatesExecutorContract(t *testing.T) { + tests := []struct { + name string + freshPlatform string + freshType string + }{ + {name: "platform changed", freshPlatform: PlatformAnthropic, freshType: AccountTypeOAuth}, + {name: "type changed", freshPlatform: PlatformGrok, freshType: AccountTypeUpstream}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := &Account{ID: 83, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive} + freshAccount := &Account{ID: account.ID, Platform: tt.freshPlatform, Type: tt.freshType, Status: StatusActive} + repo := &refreshAPIAccountRepo{account: freshAccount} + executor := NewGrokTokenRefresher(nil) + + api := NewOAuthRefreshAPI(repo, nil) + result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute) + + require.ErrorContains(t, err, "no longer refreshable") + require.Nil(t, result) + require.Zero(t, repo.updateCalls) + }) + } +} + +func TestRefreshIfNeeded_DBRereadMissingGrokRefreshCredentialReturnsPermanentSignal(t *testing.T) { + account := &Account{ + ID: 84, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{ + "refresh_token": "caller-snapshot-refresh-token", + }, + } + freshAccount := &Account{ID: account.ID, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive} + repo := &refreshAPIAccountRepo{account: freshAccount} + executor := NewGrokTokenRefresher(nil) + + api := NewOAuthRefreshAPI(repo, nil) + result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute) + + require.ErrorIs(t, err, errGrokOAuthRefreshTokenMissing) + require.Nil(t, result) + require.Zero(t, repo.updateCalls) } func TestRefreshIfNeeded_NilCredentials(t *testing.T) { - account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeOAuth} + account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive} repo := &refreshAPIAccountRepo{account: account} cache := &refreshAPICacheStub{lockResult: true} executor := &refreshAPIExecutorStub{ @@ -413,6 +499,7 @@ func TestRefreshIfNeeded_InvalidGrantRaceRecovered(t *testing.T) { ID: 10, Platform: PlatformAnthropic, Type: AccountTypeOAuth, + Status: StatusActive, Credentials: map[string]any{"refresh_token": "old-rt", "access_token": "old-at"}, } // After race, DB has new refresh token from another worker @@ -420,6 +507,7 @@ func TestRefreshIfNeeded_InvalidGrantRaceRecovered(t *testing.T) { ID: 10, Platform: PlatformAnthropic, Type: AccountTypeOAuth, + Status: StatusActive, Credentials: map[string]any{"refresh_token": "new-rt", "access_token": "new-at"}, } repo := &refreshAPIAccountRepoWithRace{ @@ -449,6 +537,7 @@ func TestRefreshIfNeeded_InvalidGrantGenuine(t *testing.T) { ID: 11, Platform: PlatformAnthropic, Type: AccountTypeOAuth, + Status: StatusActive, Credentials: map[string]any{"refresh_token": "revoked-rt", "access_token": "old-at"}, } repo := &refreshAPIAccountRepoWithRace{ @@ -474,6 +563,7 @@ func TestRefreshIfNeeded_InvalidGrantDBRereadFailsOnRecovery(t *testing.T) { ID: 12, Platform: PlatformAnthropic, Type: AccountTypeOAuth, + Status: StatusActive, Credentials: map[string]any{"refresh_token": "old-rt"}, } repo := &refreshAPIAccountRepoWithRace{ @@ -500,6 +590,7 @@ func TestRefreshIfNeeded_LocalMutexSerializesConcurrent(t *testing.T) { ID: 20, Platform: PlatformAnthropic, Type: AccountTypeOAuth, + Status: StatusActive, Credentials: map[string]any{"refresh_token": "new-rt", "access_token": "new-at"}, } callCount := 0 @@ -507,6 +598,7 @@ func TestRefreshIfNeeded_LocalMutexSerializesConcurrent(t *testing.T) { ID: 20, Platform: PlatformAnthropic, Type: AccountTypeOAuth, + Status: StatusActive, Credentials: map[string]any{"refresh_token": "old-rt"}, }} @@ -556,6 +648,78 @@ func TestRefreshIfNeeded_LocalMutexSerializesConcurrent(t *testing.T) { mu.Unlock() } +func TestRefreshIfNeeded_LocalLockWaitHonorsContextCancellation(t *testing.T) { + account := &Account{ID: 21, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive} + repo := &refreshAPIAccountRepo{account: account} + refreshStarted := make(chan struct{}) + releaseRefresh := make(chan struct{}) + var once sync.Once + executor := &dynamicRefreshExecutor{ + canRefresh: true, + cacheKey: "test:context-lock:grok", + needsRefreshFunc: func() bool { return true }, + refreshFunc: func(context.Context, *Account) (map[string]any, error) { + once.Do(func() { close(refreshStarted) }) + <-releaseRefresh + return map[string]any{"access_token": "new-at"}, nil + }, + } + api := NewOAuthRefreshAPI(repo, nil) + firstDone := make(chan error, 1) + go func() { + _, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute) + firstDone <- err + }() + <-refreshStarted + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + startedAt := time.Now() + result, err := api.RefreshIfNeeded(ctx, account, executor, 3*time.Minute) + + require.Nil(t, result) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Less(t, time.Since(startedAt), 500*time.Millisecond) + close(releaseRefresh) + require.NoError(t, <-firstDone) +} + +func TestRefreshIfNeeded_ReleasesDistributedLockWithCleanupContext(t *testing.T) { + account := &Account{ + ID: 22, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{ + "access_token": "old-access", + "refresh_token": "old-refresh", + }, + } + repo := &refreshAPIAccountRepo{account: account} + cache := &refreshAPICacheStub{lockResult: true} + ctx, cancel := context.WithCancel(context.Background()) + executor := &dynamicRefreshExecutor{ + canRefresh: true, + cacheKey: "test:cleanup:grok", + needsRefreshFunc: func() bool { return true }, + refreshFunc: func(context.Context, *Account) (map[string]any, error) { + cancel() + return map[string]any{"access_token": "new-at"}, nil + }, + } + api := NewOAuthRefreshAPI(repo, cache) + + result, err := api.RefreshIfNeeded(ctx, account, executor, 3*time.Minute) + + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, result) + require.Zero(t, repo.updateCalls) + require.Equal(t, "old-access", account.GetGrokAccessToken()) + require.Zero(t, account.GetCredentialAsInt64("_token_version")) + require.Equal(t, 1, cache.releaseCalls) + require.NoError(t, cache.releaseCtxErr) +} + // dynamicRefreshExecutor is a test helper with function-based NeedsRefresh and Refresh. type dynamicRefreshExecutor struct { canRefresh bool diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index a3f60fee2fa..0b39ba8a2a5 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -3,6 +3,7 @@ package service import ( "context" "net/http" + "sync" "time" ) @@ -96,6 +97,20 @@ func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until ti 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{}) + return actual.(*sync.Mutex) +} + +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) { @@ -107,7 +122,7 @@ func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until ti if !loaded { actual, stored := s.openaiAccountRuntimeBlockUntil.LoadOrStore(account.ID, blockUntil) if !stored { - return + return generation, true } current = actual } @@ -115,15 +130,15 @@ func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until ti currentUntil, ok := current.(time.Time) if !ok || currentUntil.IsZero() { if s.openaiAccountRuntimeBlockUntil.CompareAndSwap(account.ID, current, blockUntil) { - return + return generation, true } continue } - if currentUntil.After(blockUntil) { - return + if !blockUntil.After(currentUntil) { + return generation, false } if s.openaiAccountRuntimeBlockUntil.CompareAndSwap(account.ID, current, blockUntil) { - return + return generation, true } } } @@ -132,13 +147,20 @@ 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.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 @@ -146,12 +168,14 @@ func (s *OpenAIGatewayService) isOpenAIAccountRuntimeBlocked(account *Account) b 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 } diff --git a/backend/internal/service/openai_gateway_cc_pipeline.go b/backend/internal/service/openai_gateway_cc_pipeline.go index 690b9991bbd..2391e25ba58 100644 --- a/backend/internal/service/openai_gateway_cc_pipeline.go +++ b/backend/internal/service/openai_gateway_cc_pipeline.go @@ -118,6 +118,7 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( return &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, + ResponseHeaders: resp.Header.Clone(), RetryableOnSameAccount: account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), } } @@ -197,7 +198,9 @@ func (s *OpenAIGatewayService) sendCCUpstreamRequest( // 账号级请求头覆写（仅 openai api_key 账号启用时生效） account.ApplyHeaderOverrides(upstreamReq.Header) if account.Platform == PlatformGrok { - applyGrokCLIHeaders(upstreamReq.Header) + if account.IsGrokOAuth() { + applyGrokCLIHeaders(upstreamReq.Header) + } applyGrokCacheHeaders(upstreamReq.Header, grokCacheIdentity) } diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 7693986b0d1..095d544f89b 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -109,7 +109,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( // Grok Composer does not accept image_url parts directly, but Grok Build // can describe the images first. Bridge only this exact failure mode. - token, tokenKind, err := s.GetAccessToken(ctx, account) + token, tokenKind, err := s.getRequestCredential(ctx, c, account) if err != nil { return nil, err } @@ -162,7 +162,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( } SetActualOpenAIUpstreamEndpoint(c, grokChatRawEndpoint) customUA := account.GetOpenAIUserAgent() - if customUA == "" && account.Platform == PlatformGrok { + if customUA == "" && account.IsGrokOAuth() { customUA = "sub2api-grok/1.0" } resp, err := s.sendCCUpstreamRequest(ctx, c, account, targetURL, upstreamBody, clientStream, token, customUA, grokCacheIdentity) @@ -189,6 +189,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, + ResponseHeaders: resp.Header.Clone(), RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 8bd1d456dcc..675b4bc9f95 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -54,7 +54,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( return nil, fmt.Errorf("apply grok prompt cache identity: %w", err) } - token, _, err := s.GetAccessToken(ctx, account) + token, _, err := s.getRequestCredential(ctx, c, account) if err != nil { return nil, err } @@ -100,6 +100,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, + ResponseHeaders: resp.Header.Clone(), RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } @@ -610,6 +611,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( return "", OpenAIUsage{}, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, + ResponseHeaders: resp.Header.Clone(), RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } @@ -750,7 +752,9 @@ func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Acc req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/event-stream") - applyGrokCLIHeaders(req.Header) + if account.IsGrokOAuth() { + applyGrokCLIHeaders(req.Header) + } applyGrokCacheHeaders(req.Header, cacheIdentity) if c != nil { if v := c.GetHeader("OpenAI-Beta"); strings.TrimSpace(v) != "" { diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge.go b/backend/internal/service/openai_gateway_grok_chat_bridge.go index fd6b4265360..fe9cfb09fd6 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge.go @@ -253,7 +253,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( } responsesBody = updatedBody - token, _, err := s.GetAccessToken(ctx, account) + token, _, err := s.getRequestCredential(ctx, c, account) if err != nil { return nil, fmt.Errorf("get grok access token: %w", err) } @@ -294,6 +294,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, + ResponseHeaders: resp.Header.Clone(), RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go index 0843f1ea427..4fc31461b4d 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go @@ -284,6 +284,7 @@ func TestForwardGrokChatViaResponses429UsesGrokRateLimitPolicy(t *testing.T) { var failoverErr *UpstreamFailoverError require.True(t, errors.As(err, &failoverErr)) require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Equal(t, "45", failoverErr.ResponseHeaders.Get("Retry-After")) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) require.Equal(t, grokChatResponsesEndpoint, GetActualOpenAIUpstreamEndpoint(c)) require.Equal(t, 1, repo.rateLimitedCalls) @@ -292,6 +293,45 @@ func TestForwardGrokChatViaResponses429UsesGrokRateLimitPolicy(t *testing.T) { require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) } +func TestForwardGrokRawChat429PreservesRetryAfter(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done"}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 7551}) + + account := grokChatBridgeTestAccount(755) + account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Retry-After": []string{"45"}, + }, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + + require.Error(t, err) + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Equal(t, "45", failoverErr.ResponseHeaders.Get("Retry-After")) + require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String()) +} + func TestForwardGrokRawChatErrorRecordsActualEndpoint(t *testing.T) { gin.SetMode(gin.TestMode) @@ -329,11 +369,14 @@ func grokChatBridgeTestAccount(id int64) *Account { Name: "grok-cache-bridge", Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, }, } } diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 3b64cbc16b6..4f7ed05edfa 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -259,7 +259,7 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id") require.NoError(t, err) require.Equal(t, http.MethodPost, req.Method) - require.Equal(t, "https://xai.test/v1/responses", req.URL.String()) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String()) require.Equal(t, "Bearer access-token", req.Header.Get("Authorization")) require.Equal(t, "application/json", req.Header.Get("Content-Type")) require.Contains(t, req.Header.Get("Accept"), "text/event-stream") @@ -284,6 +284,8 @@ func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing. require.NoError(t, err) require.Equal(t, "https://grok.example.test/v1/responses", req.URL.String()) require.Equal(t, "Bearer api-key", req.Header.Get("Authorization")) + require.Empty(t, req.Header.Get("X-Grok-Client-Version")) + require.NotEqual(t, grokUpstreamUserAgent, req.Header.Get("User-Agent")) } func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T) { @@ -425,7 +427,8 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) { require.Equal(t, http.MethodPost, upstream.lastReq.Method) require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type")) - require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent")) require.JSONEq(t, `{"model":"grok-imagine-image-quality","prompt":"draw a cat"}`, string(upstream.lastBody)) require.Equal(t, http.StatusOK, recorder.Code) require.JSONEq(t, `{"data":[]}`, recorder.Body.String()) @@ -611,7 +614,7 @@ func TestForwardGrokMediaVideoGenerationPreservesImageToVideoModel(t *testing.T) require.Equal(t, VideoBillingDefaultDurationSeconds, result.VideoDurationSeconds) } -func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing.T) { +func TestForwardGrokMediaOAuthImageToVideoKeepsCLIGatewayForLargeBody(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() @@ -626,10 +629,14 @@ func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing Name: "grok-oauth", Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "oauth-access-token", - "base_url": xai.DefaultCLIBaseURL, + "access_token": "oauth-access-token", + "refresh_token": "oauth-refresh-token", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, }, } upstream := &httpUpstreamRecorder{resp: &http.Response{ @@ -639,11 +646,11 @@ func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing }, Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-oauth"}`)), }} - svc := &OpenAIGatewayService{httpUpstream: upstream} + svc := &OpenAIGatewayService{httpUpstream: upstream, grokTokenProvider: NewGrokTokenProvider(nil, nil)} _, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json") require.NoError(t, err) - require.Equal(t, xai.DefaultBaseURL+"/videos/generations", upstream.lastReq.URL.String()) + require.Equal(t, xai.DefaultCLIBaseURL+"/videos/generations", upstream.lastReq.URL.String()) require.Equal(t, "data:image/png;base64,"+imageData, gjson.GetBytes(upstream.lastBody, "image.image_url").String()) } @@ -681,6 +688,8 @@ func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) { require.Equal(t, "https://xai.test/v1/videos/request-123", upstream.lastReq.URL.String()) require.Equal(t, http.MethodGet, upstream.lastReq.Method) require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization")) + require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent")) require.Empty(t, upstream.lastReq.Header.Get("Content-Type")) require.Empty(t, upstream.lastBody) require.Equal(t, http.StatusOK, recorder.Code) @@ -792,27 +801,64 @@ func TestForwardGrokMedia429ReconcilesRateLimitBeforeCustomErrorBypass(t *testin require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) } -func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *testing.T) { +func TestGrokMedia429FailoverPreservesRetryAfter(t *testing.T) { gin.SetMode(gin.TestMode) - recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) - body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done","prompt_cache_key":"raw-client-cache-key"}`) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) - c.Set("api_key", &APIKey{ID: 5101}) - + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) account := &Account{ - ID: 51, + ID: 641, Name: "grok-oauth", Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, Schedulable: true, + Credentials: map[string]any{ + "custom_error_codes_enabled": true, + "custom_error_codes": []any{float64(http.StatusTooManyRequests)}, + }, + } + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + resp := &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"45"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + } + + result, err := svc.handleGrokMediaErrorResponse(context.Background(), resp, c, account, "request-id", "grok-imagine") + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Equal(t, "45", failoverErr.ResponseHeaders.Get("Retry-After")) +} + +func healthyGrokOAuthGatewayTestAccount(id int64, token string) *Account { + return &Account{ + ID: id, Name: "grok", Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, + "access_token": token, + "refresh_token": "refresh-token", + "expires_at": time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, }, } +} + +func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done","prompt_cache_key":"raw-client-cache-key"}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 5101}) + + account := healthyGrokOAuthGatewayTestAccount(51, "access-token") repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{51: account}, @@ -864,18 +910,7 @@ func TestForwardGrokResponsesStreamingDefaultsEmptyModelTo45AndSnapshots(t *test c.Request.Header.Set("OpenAI-Beta", "responses=experimental") c.Set("api_key", &APIKey{ID: 5201}) - account := &Account{ - ID: 52, - Name: "grok", - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, - }, - } + account := healthyGrokOAuthGatewayTestAccount(52, "access-token") repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{52: account}, @@ -968,12 +1003,46 @@ func TestForwardGrokResponsesAPIKeyUsesXAIResponses(t *testing.T) { require.NoError(t, err) require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer xai-test-key", upstream.lastReq.Header.Get("Authorization")) + require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent")) require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) require.Equal(t, "resp_grok_api_key", result.ResponseID) require.Equal(t, 2, result.Usage.InputTokens) require.Equal(t, 1, result.Usage.OutputTokens) } +func TestForwardAsChatCompletionsForGrokAPIKeyUsesConfiguredRawEndpointWithoutOAuthIdentity(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + account := &Account{ + ID: 706, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "third-party-key", + "base_url": "https://grok.example.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.5","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)), + }} + svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream} + + _, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.NoError(t, err) + require.Equal(t, "https://grok.example.test/v1/chat/completions", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization")) + require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent")) +} + func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) { gin.SetMode(gin.TestMode) @@ -1017,18 +1086,7 @@ func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *te c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) c.Request.Header.Set("Content-Type", "application/json") - account := &Account{ - ID: 53, - Name: "grok", - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, - }, - } + account := healthyGrokOAuthGatewayTestAccount(53, "access-token") repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{53: account}, @@ -1085,18 +1143,7 @@ func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *test c.Request.Header.Set("Content-Type", "application/json") c.Set("api_key", &APIKey{ID: 5202}) - account := &Account{ - ID: 56, - Name: "grok", - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, - }, - } + account := healthyGrokOAuthGatewayTestAccount(56, "access-token") repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{56: account}, @@ -1145,18 +1192,9 @@ func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing c.Set("api_key", &APIKey{ID: 5203}) newAccount := func(id int64, token string) *Account { - return &Account{ - ID: id, - Name: fmt.Sprintf("grok-%d", id), - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": token, - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, - }, - } + account := healthyGrokOAuthGatewayTestAccount(id, token) + account.Name = fmt.Sprintf("grok-%d", id) + return account } firstAccount := newAccount(58, "access-token-a") secondAccount := newAccount(59, "access-token-b") @@ -1213,18 +1251,7 @@ func TestForwardAsChatCompletionsForGrokStreamingStopFallsBackToRawXAIChatComple c.Request.Header.Set(grokConversationIDHeader, "native-client-conversation") c.Set("api_key", &APIKey{ID: 5301}) - account := &Account{ - ID: 53, - Name: "grok", - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, - }, - } + account := healthyGrokOAuthGatewayTestAccount(53, "access-token") repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{53: account}, @@ -1284,18 +1311,7 @@ func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T) c.Request.Header.Set("Content-Type", "application/json") c.Set("api_key", &APIKey{ID: 5501}) - account := &Account{ - ID: 55, - Name: "grok", - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, - }, - } + account := healthyGrokOAuthGatewayTestAccount(55, "access-token") repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{55: account}, @@ -1358,18 +1374,7 @@ func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) { c.Request.Header.Set("OpenAI-Beta", "grok-experimental") c.Request.Header.Set("originator", "opencode") - account := &Account{ - ID: 54, - Name: "grok", - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, - }, - } + account := healthyGrokOAuthGatewayTestAccount(54, "access-token") repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{54: account}, @@ -1419,18 +1424,7 @@ func TestForwardAsAnthropicForGrokStreamingPreservesCacheUsage(t *testing.T) { c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) c.Set("api_key", &APIKey{ID: 5402}) - account := &Account{ - ID: 57, - Name: "grok", - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - "base_url": xai.DefaultCLIBaseURL, - }, - } + account := healthyGrokOAuthGatewayTestAccount(57, "access-token") repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{57: account}, diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 8cf62bdc652..e2add90ff2c 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -258,7 +258,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( } // 5. Get access token - token, _, err := s.GetAccessToken(ctx, account) + token, _, err := s.getRequestCredential(ctx, c, account) if err != nil { return nil, fmt.Errorf("get access token: %w", err) } diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 11a8255bfad..12c9a5bb971 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -401,6 +401,10 @@ type OpenAIGatewayService struct { openaiWSFallbackUntil sync.Map // key: int64(accountID), value: time.Time openaiAccountRuntimeBlockUntil sync.Map // key: int64(accountID), value: time.Time + openaiAccountRuntimeBlockLocks sync.Map // key: int64(accountID), value: *sync.Mutex + openaiAccountRuntimeBlockGeneration sync.Map // key: int64(accountID), value: uint64 + openaiAccountRuntimeBlockSequence atomic.Uint64 + grokCredentialMutationLocks sync.Map // key: int64(accountID), value: *sync.Mutex openaiOAuth429WindowStartUnixNano atomic.Int64 openaiOAuth429WindowCount atomic.Int64 openaiWSRetryMetrics openAIWSRetryMetrics diff --git a/backend/internal/service/ops_models.go b/backend/internal/service/ops_models.go index d95bbadbb8f..411ca7595b0 100644 --- a/backend/internal/service/ops_models.go +++ b/backend/internal/service/ops_models.go @@ -27,7 +27,7 @@ type OpsErrorLog struct { CreatedAt time.Time `json:"created_at"` // Standardized classification - // - phase: request|auth|routing|upstream|network|internal + // - phase: request|auth|account_auth|routing|upstream|network|internal // - owner: client|provider|platform // - source: client_request|upstream_http|gateway Phase string `json:"phase"` @@ -121,7 +121,7 @@ type OpsErrorLogFilter struct { StatusCodes []int StatusCodesOther bool - Phase string // Special: Phase=="upstream" bypasses status>=400 clause; do not set together with ErrorPhasesAny. + Phase string // Recovered provider rows bypass status>=400 only with the explicit opt-in below. Owner string Source string Resolved *bool @@ -150,14 +150,15 @@ type OpsErrorLogFilter struct { // ExcludeCountTokens drops count_tokens probe errors (is_count_tokens=true). ExcludeCountTokens bool - // IncludeRecoveredUpstream 显式豁免 status>=400 守卫（仅在 Phase=="upstream" 时生效）： - // ops 专用上游错误列表需要看到 status<400 的 recovered upstream 行。 - // 请求错误语义的端点不设此开关，phase=upstream 过滤照常生效且守卫保留。 + // IncludeRecoveredUpstream explicitly exempts provider-health phases + // (upstream and account_auth) from the status>=400 guard. Ops provider + // health lists need status<400 recovered rows; request-error endpoints do + // not set this flag and retain client-error semantics. IncludeRecoveredUpstream bool // ErrorPhasesAny / ErrorTypesAny add plain ANY() filters WITHOUT touching the - // special-cased single `Phase` field (only Phase=="upstream" with - // IncludeRecoveredUpstream bypasses the status>=400 clause). + // special-cased single `Phase` field. With IncludeRecoveredUpstream, an ANY + // list containing only upstream/account_auth also bypasses status>=400. // NOTE: these ANY filters do NOT bypass status>=400; records with error_phase='upstream' // but status_code<400 (recovered upstream errors) remain excluded. // Used to map user-facing coarse categories to backend conditions. diff --git a/backend/internal/service/ops_service.go b/backend/internal/service/ops_service.go index 61f85ef9048..3b272526655 100644 --- a/backend/internal/service/ops_service.go +++ b/backend/internal/service/ops_service.go @@ -222,6 +222,33 @@ func (s *OpsService) prepareErrorLogInput(ctx context.Context, entry *OpsInsertE entry.ErrorType = "api_error" } + // Credential acquisition is a gateway/account-auth stage, not an inference + // HTTP attempt. Enforce that ownership at the persistence boundary so an + // earlier inference attempt cannot leak its status or text into top-level + // auth fields even if a caller supplied stale single-value context. + for i := len(entry.UpstreamErrors) - 1; i >= 0; i-- { + last := entry.UpstreamErrors[i] + if last == nil { + continue + } + if last.Stage == string(GatewayFailureStageAccountAuth) { + entry.ErrorPhase = string(GatewayFailureStageAccountAuth) + entry.ErrorOwner = "provider" + entry.ErrorSource = "gateway" + code := 0 + entry.UpstreamStatusCode = &code + entry.UpstreamErrorMessage = nil + if message := strings.TrimSpace(last.Message); message != "" { + entry.UpstreamErrorMessage = &message + } + entry.UpstreamErrorDetail = nil + if detail := strings.TrimSpace(last.Detail); detail != "" { + entry.UpstreamErrorDetail = &detail + } + } + break + } + // Sanitize + truncate error_body to avoid storing sensitive data. if strings.TrimSpace(entry.ErrorBody) != "" { sanitized, _ := sanitizeErrorBodyForStorage(entry.ErrorBody, opsMaxStoredErrorBodyBytes) @@ -229,7 +256,7 @@ func (s *OpsService) prepareErrorLogInput(ctx context.Context, entry *OpsInsertE } // Sanitize upstream error context if provided by gateway services. - if entry.UpstreamStatusCode != nil && *entry.UpstreamStatusCode <= 0 { + if entry.UpstreamStatusCode != nil && *entry.UpstreamStatusCode <= 0 && entry.ErrorPhase != string(GatewayFailureStageAccountAuth) { entry.UpstreamStatusCode = nil } if entry.UpstreamErrorMessage != nil { diff --git a/backend/internal/service/ops_service_batch_test.go b/backend/internal/service/ops_service_batch_test.go index a9419ad7bfa..a027f487372 100644 --- a/backend/internal/service/ops_service_batch_test.go +++ b/backend/internal/service/ops_service_batch_test.go @@ -97,6 +97,54 @@ func TestOpsServiceRecordErrorBatch_FallsBackToSingleInsert(t *testing.T) { require.Equal(t, 2, singleCalls) } +func TestOpsServiceRecordErrorPersistsExplicitAccountAuthStatusZero(t *testing.T) { + t.Parallel() + + var captured *OpsInsertErrorLogInput + repo := &opsRepoMock{ + InsertErrorLogFn: func(_ context.Context, input *OpsInsertErrorLogInput) (int64, error) { + captured = input + return 1, nil + }, + } + svc := NewOpsService(repo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + staleStatus := 403 + staleMessage := "stale inference message" + staleDetail := "stale inference detail" + + err := svc.RecordError(context.Background(), &OpsInsertErrorLogInput{ + ErrorPhase: "upstream", + ErrorType: "upstream_error", + ErrorOwner: "provider", + ErrorSource: "upstream_http", + UpstreamStatusCode: &staleStatus, + UpstreamErrorMessage: &staleMessage, + UpstreamErrorDetail: &staleDetail, + UpstreamErrors: []*OpsUpstreamErrorEvent{ + {Stage: string(GatewayFailureStageInference), UpstreamStatusCode: 403, Message: staleMessage, Detail: staleDetail}, + { + Stage: string(GatewayFailureStageAccountAuth), Scope: string(GatewayFailureScopeAccount), + Reason: string(GrokCredentialReasonRevoked), Message: "Grok OAuth credentials require account action", + }, + }, + }) + + require.NoError(t, err) + require.NotNil(t, captured) + require.Equal(t, "account_auth", captured.ErrorPhase) + require.Equal(t, "provider", captured.ErrorOwner) + require.Equal(t, "gateway", captured.ErrorSource) + require.NotNil(t, captured.UpstreamStatusCode) + require.Zero(t, *captured.UpstreamStatusCode) + require.NotNil(t, captured.UpstreamErrorMessage) + require.Equal(t, "Grok OAuth credentials require account action", *captured.UpstreamErrorMessage) + require.Nil(t, captured.UpstreamErrorDetail) + require.Nil(t, captured.UpstreamErrors) + require.NotNil(t, captured.UpstreamErrorsJSON) + require.Contains(t, *captured.UpstreamErrorsJSON, `"upstream_status_code":403`) + require.Contains(t, *captured.UpstreamErrorsJSON, `"stage":"account_auth"`) +} + func strPtr(v string) *string { return &v } diff --git a/backend/internal/service/ops_upstream_context.go b/backend/internal/service/ops_upstream_context.go index 6389f1ecb56..2d28f1a168c 100644 --- a/backend/internal/service/ops_upstream_context.go +++ b/backend/internal/service/ops_upstream_context.go @@ -188,6 +188,11 @@ type OpsUpstreamErrorEvent struct { // Kind: http_error | request_error | retry_exhausted | failover Kind string `json:"kind,omitempty"` + // Stage/Scope/Reason distinguish credential acquisition from inference + // without overloading upstream_status_code with a synthetic HTTP status. + Stage string `json:"stage,omitempty"` + Scope string `json:"scope,omitempty"` + Reason string `json:"reason,omitempty"` Message string `json:"message,omitempty"` Detail string `json:"detail,omitempty"` @@ -204,6 +209,9 @@ func appendOpsUpstreamError(c *gin.Context, ev OpsUpstreamErrorEvent) { ev.UpstreamRequestID = strings.TrimSpace(ev.UpstreamRequestID) ev.UpstreamResponseBody = strings.TrimSpace(ev.UpstreamResponseBody) ev.Kind = strings.TrimSpace(ev.Kind) + ev.Stage = strings.TrimSpace(ev.Stage) + ev.Scope = strings.TrimSpace(ev.Scope) + ev.Reason = strings.TrimSpace(ev.Reason) ev.UpstreamURL = strings.TrimSpace(ev.UpstreamURL) ev.Message = strings.TrimSpace(ev.Message) ev.Detail = strings.TrimSpace(ev.Detail) diff --git a/backend/internal/service/ops_user_error.go b/backend/internal/service/ops_user_error.go index e3055c23927..22d160bf347 100644 --- a/backend/internal/service/ops_user_error.go +++ b/backend/internal/service/ops_user_error.go @@ -44,7 +44,7 @@ func MapUserErrorCategory(phase, errType string) string { return "auth" case "routing": return "service_unavailable" - case "upstream", "network": + case "account_auth", "upstream", "network": return "upstream" case "internal": return "internal" @@ -73,7 +73,7 @@ func CategoryToFilter(category string) (phases []string, errorTypes []string) { case "service_unavailable": return []string{"routing"}, nil case "upstream": - return []string{"upstream", "network"}, nil + return []string{"account_auth", "upstream", "network"}, nil case "internal": return []string{"internal"}, nil case "rate_limit": diff --git a/backend/internal/service/ops_user_error_test.go b/backend/internal/service/ops_user_error_test.go index 9e0bc164b42..c72cf5dbfb6 100644 --- a/backend/internal/service/ops_user_error_test.go +++ b/backend/internal/service/ops_user_error_test.go @@ -17,6 +17,7 @@ func TestMapUserErrorCategory(t *testing.T) { {"request", "subscription_error", "quota"}, {"request", "invalid_request_error", "invalid_request"}, {"routing", "api_error", "service_unavailable"}, + {"account_auth", "upstream_error", "upstream"}, {"upstream", "upstream_error", "upstream"}, {"network", "api_error", "upstream"}, {"internal", "api_error", "internal"}, @@ -43,7 +44,7 @@ func TestCategoryToFilter(t *testing.T) { t.Fatalf("service_unavailable => phases=%v types=%v", phases, types) } phases, types = CategoryToFilter("upstream") - if len(phases) != 2 || phases[0] != "upstream" || phases[1] != "network" || len(types) != 0 { + if len(phases) != 3 || phases[0] != "account_auth" || phases[1] != "upstream" || phases[2] != "network" || len(types) != 0 { t.Fatalf("upstream => phases=%v types=%v", phases, types) } phases, types = CategoryToFilter("internal") diff --git a/backend/internal/service/refresh_policy.go b/backend/internal/service/refresh_policy.go index 7f299be0193..6b90077e3b8 100644 --- a/backend/internal/service/refresh_policy.go +++ b/backend/internal/service/refresh_policy.go @@ -61,6 +61,14 @@ func AntigravityProviderRefreshPolicy() ProviderRefreshPolicy { } } +func GrokProviderRefreshPolicy() ProviderRefreshPolicy { + return ProviderRefreshPolicy{ + OnRefreshError: ProviderRefreshErrorReturn, + OnLockHeld: ProviderLockHeldWaitForCache, + FailureTTL: 0, + } +} + // BackgroundSkipAction 定义后台刷新服务在“未实际刷新”场景的计数方式。 type BackgroundSkipAction int diff --git a/backend/internal/service/token_refresh_service_test.go b/backend/internal/service/token_refresh_service_test.go index bd6a521576b..c4c0b35cebc 100644 --- a/backend/internal/service/token_refresh_service_test.go +++ b/backend/internal/service/token_refresh_service_test.go @@ -26,6 +26,9 @@ type tokenRefreshAccountRepo struct { lastExtraUpdates map[string]any lastAccount *Account updateErr error + setErrorErr error + setTempUnschedErr error + beforeConditionalState func() } func (r *tokenRefreshAccountRepo) Update(ctx context.Context, account *Account) error { @@ -56,7 +59,7 @@ func (r *tokenRefreshAccountRepo) UpdateCredentials(ctx context.Context, id int6 func (r *tokenRefreshAccountRepo) SetError(ctx context.Context, id int64, errorMsg string) error { r.setErrorCalls++ r.lastErrorMessage = errorMsg - return nil + return r.setErrorErr } func (r *tokenRefreshAccountRepo) ClearTempUnschedulable(ctx context.Context, id int64) error { @@ -67,7 +70,66 @@ func (r *tokenRefreshAccountRepo) ClearTempUnschedulable(ctx context.Context, id func (r *tokenRefreshAccountRepo) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error { r.setTempUnschedCalls++ r.lastTempUnschedReason = reason - return nil + return r.setTempUnschedErr +} + +func (r *tokenRefreshAccountRepo) SetGrokCredentialErrorIfMatch( + _ context.Context, + id int64, + snapshot GrokCredentialMutationSnapshot, + errorMsg string, +) (bool, error) { + if r.beforeConditionalState != nil { + hook := r.beforeConditionalState + r.beforeConditionalState = nil + hook() + } + account := r.accountsByID[id] + if !grokCredentialSnapshotMatchesAccount(account, snapshot) || + (errorMsg == string(GrokCredentialReasonProxyInvalid) && account.Proxy != nil) { + return false, nil + } + r.setErrorCalls++ + r.lastErrorMessage = errorMsg + if r.setErrorErr != nil { + return false, r.setErrorErr + } + account.Status = StatusError + account.Schedulable = false + account.ErrorMessage = errorMsg + return true, nil +} + +func (r *tokenRefreshAccountRepo) SetGrokCredentialTempUnschedulableIfMatch( + _ context.Context, + id int64, + snapshot GrokCredentialMutationSnapshot, + until time.Time, + reason string, +) (bool, error) { + if r.beforeConditionalState != nil { + hook := r.beforeConditionalState + r.beforeConditionalState = nil + hook() + } + account := r.accountsByID[id] + if !grokCredentialSnapshotMatchesAccount(account, snapshot) { + return false, nil + } + r.setTempUnschedCalls++ + r.lastTempUnschedReason = reason + if r.setTempUnschedErr != nil { + return false, r.setTempUnschedErr + } + value := until + account.TempUnschedulableUntil = &value + return true, nil +} + +func grokCredentialSnapshotMatchesAccount(account *Account, snapshot GrokCredentialMutationSnapshot) bool { + return account != nil && account.IsGrokOAuth() && account.IsSchedulable() && + grokCredentialMutationSnapshot(account).CredentialsJSON == snapshot.CredentialsJSON && + grokCredentialProxyIDsEqual(account.ProxyID, snapshot.ProxyID) } func (r *tokenRefreshAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error { @@ -750,6 +812,7 @@ func TestPathA_Success(t *testing.T) { ID: 100, Platform: PlatformGemini, Type: AccountTypeOAuth, + Status: StatusActive, } repo := &tokenRefreshAccountRepo{} repo.accountsByID = map[int64]*Account{account.ID: account} @@ -771,6 +834,7 @@ func TestPathA_LockHeld(t *testing.T) { ID: 101, Platform: PlatformGemini, Type: AccountTypeOAuth, + Status: StatusActive, } repo := &tokenRefreshAccountRepo{} invalidator := &tokenCacheInvalidatorStub{} @@ -791,6 +855,7 @@ func TestPathA_AlreadyRefreshed(t *testing.T) { ID: 102, Platform: PlatformGemini, Type: AccountTypeOAuth, + Status: StatusActive, } repo := &tokenRefreshAccountRepo{} repo.accountsByID = map[int64]*Account{account.ID: account} @@ -830,6 +895,7 @@ func TestPathA_NonRetryableError(t *testing.T) { ID: 103, Platform: PlatformGemini, Type: AccountTypeOAuth, + Status: StatusActive, } repo := &tokenRefreshAccountRepo{} repo.accountsByID = map[int64]*Account{account.ID: account} @@ -855,6 +921,7 @@ func TestPathA_RetryableErrorExhausted(t *testing.T) { ID: 104, Platform: PlatformGemini, Type: AccountTypeOAuth, + Status: StatusActive, } repo := &tokenRefreshAccountRepo{} repo.accountsByID = map[int64]*Account{account.ID: account} @@ -888,6 +955,7 @@ func TestPathA_DBUpdateFailed(t *testing.T) { ID: 105, Platform: PlatformGemini, Type: AccountTypeOAuth, + Status: StatusActive, } repo := &tokenRefreshAccountRepo{updateErr: errors.New("db connection lost")} repo.accountsByID = map[int64]*Account{account.ID: account} @@ -898,7 +966,7 @@ func TestPathA_DBUpdateFailed(t *testing.T) { err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour) require.Error(t, err) - require.Contains(t, err.Error(), "DB update failed") + require.ErrorIs(t, err, errOAuthRefreshCredentialPersist) require.Equal(t, 1, repo.updateCalls) // DB 更新被尝试 require.Equal(t, 0, invalidator.calls) // DB 失败时不应触发缓存失效 } diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index d5d9124ec2a..96956b2f5d4 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -186,7 +186,7 @@ func ProvideGrokTokenProvider( p := NewGrokTokenProvider(accountRepo, tokenCache) executor := NewGrokTokenRefresher(grokOAuthService) p.SetRefreshAPI(refreshAPI, executor) - p.SetRefreshPolicy(AntigravityProviderRefreshPolicy()) + p.SetRefreshPolicy(GrokProviderRefreshPolicy()) p.SetTempUnschedCache(tempUnschedCache) return p } diff --git a/frontend/src/components/admin/usage/UsageFilters.vue b/frontend/src/components/admin/usage/UsageFilters.vue index aa981ac1875..d1e15d2cf41 100644 --- a/frontend/src/components/admin/usage/UsageFilters.vue +++ b/frontend/src/components/admin/usage/UsageFilters.vue @@ -277,6 +277,7 @@ const billingTypeOptions = ref([ const errorPhaseOptions = computed(() => [ { value: null, label: t('admin.usage.allTypes') }, { value: 'upstream', label: t('admin.ops.errorLog.typeUpstream') }, + { value: 'account_auth', label: t('admin.ops.errorLog.typeAccountAuth') }, { value: 'request', label: t('admin.ops.errorLog.typeRequest') }, { value: 'auth', label: t('admin.ops.errorLog.typeAuth') }, { value: 'routing', label: t('admin.ops.errorLog.typeRouting') }, diff --git a/frontend/src/i18n/locales/en/admin/ops.ts b/frontend/src/i18n/locales/en/admin/ops.ts index 588f6735180..6a6d02a81c3 100644 --- a/frontend/src/i18n/locales/en/admin/ops.ts +++ b/frontend/src/i18n/locales/en/admin/ops.ts @@ -266,6 +266,7 @@ export default { typeUpstream: 'Upstream', typeRequest: 'Request', typeAuth: 'Auth', + typeAccountAuth: 'Account Auth', typeRouting: 'Routing', typeInternal: 'Internal', endpoint: 'Endpoint', @@ -291,6 +292,7 @@ export default { phase: { request: 'Request', auth: 'Auth', + account_auth: 'Account Auth', routing: 'Routing', upstream: 'Upstream', network: 'Network', diff --git a/frontend/src/i18n/locales/zh/admin/ops.ts b/frontend/src/i18n/locales/zh/admin/ops.ts index 830d797a318..2816ec32403 100644 --- a/frontend/src/i18n/locales/zh/admin/ops.ts +++ b/frontend/src/i18n/locales/zh/admin/ops.ts @@ -266,6 +266,7 @@ export default { typeUpstream: '上游', typeRequest: '请求', typeAuth: '认证', + typeAccountAuth: '账号认证', typeRouting: '路由', typeInternal: '内部', endpoint: '端点', @@ -291,6 +292,7 @@ export default { phase: { request: '请求', auth: '认证', + account_auth: '账号认证', routing: '路由', upstream: '上游', network: '网络', diff --git a/frontend/src/utils/errorCategory.ts b/frontend/src/utils/errorCategory.ts index 58e97e14c43..7150c18265a 100644 --- a/frontend/src/utils/errorCategory.ts +++ b/frontend/src/utils/errorCategory.ts @@ -9,6 +9,7 @@ export function mapErrorCategory(phase?: string | null, errType?: string | null) return 'auth' case 'routing': return 'service_unavailable' + case 'account_auth': case 'upstream': case 'network': return 'upstream' diff --git a/frontend/src/views/admin/ops/components/OpsErrorDetailsModal.vue b/frontend/src/views/admin/ops/components/OpsErrorDetailsModal.vue index 0400033e898..f04876ffaed 100644 --- a/frontend/src/views/admin/ops/components/OpsErrorDetailsModal.vue +++ b/frontend/src/views/admin/ops/components/OpsErrorDetailsModal.vue @@ -72,6 +72,7 @@ const phaseSelectOptions = computed(() => { { value: '', label: t('common.all') }, { value: 'request', label: t('admin.ops.errorDetails.phase.request') || 'request' }, { value: 'auth', label: t('admin.ops.errorDetails.phase.auth') || 'auth' }, + { value: 'account_auth', label: t('admin.ops.errorDetails.phase.account_auth') || 'account_auth' }, { value: 'routing', label: t('admin.ops.errorDetails.phase.routing') || 'routing' }, { value: 'upstream', label: t('admin.ops.errorDetails.phase.upstream') || 'upstream' }, { value: 'network', label: t('admin.ops.errorDetails.phase.network') || 'network' }, diff --git a/frontend/src/views/admin/ops/components/OpsErrorLogTable.vue b/frontend/src/views/admin/ops/components/OpsErrorLogTable.vue index 856b8810d7c..6cb0bb1b4cf 100644 --- a/frontend/src/views/admin/ops/components/OpsErrorLogTable.vue +++ b/frontend/src/views/admin/ops/components/OpsErrorLogTable.vue @@ -277,6 +277,9 @@ function getTypeBadge(log: OpsErrorLog): { label: string; className: string } { if (phase === 'auth' && owner === 'client') { return { label: t('admin.ops.errorLog.typeAuth'), className: 'bg-blue-100 text-blue-800 dark:bg-blue-900 dark:text-blue-200' } } + if (phase === 'account_auth') { + return { label: t('admin.ops.errorLog.typeAccountAuth'), className: 'bg-orange-100 text-orange-800 dark:bg-orange-900 dark:text-orange-200' } + } if (phase === 'routing' && owner === 'platform') { return { label: t('admin.ops.errorLog.typeRouting'), className: 'bg-purple-100 text-purple-800 dark:bg-purple-900 dark:text-purple-200' } } From 91ae671a583abd1d331a020ba7421c8e726015c5 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Tue, 14 Jul 2026 01:02:16 +0800 Subject: [PATCH 2/6] fix(grok): satisfy lint checks --- backend/internal/service/account.go | 31 ------------------- .../service/grok_credential_failure.go | 7 ++++- .../internal/service/grok_token_provider.go | 1 - .../openai_account_runtime_block_fastpath.go | 7 ++++- 4 files changed, 12 insertions(+), 34 deletions(-) diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 539c92741dc..5214518fb52 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -6,7 +6,6 @@ import ( "errors" "hash/fnv" "log/slog" - "net/url" "reflect" "sort" "strconv" @@ -1296,36 +1295,6 @@ func (a *Account) GetGrokMediaBaseURL() string { return a.GetGrokBaseURL() } -func isOfficialGrokAPIBaseURL(raw string) bool { - return isOfficialGrokBaseURL(raw, xai.DefaultBaseURL) -} - -func isOfficialGrokCLIBaseURL(raw string) bool { - return isOfficialGrokBaseURL(raw, xai.DefaultCLIBaseURL) -} - -func isOfficialGrokBaseURL(raw, expected string) bool { - parsed, err := url.Parse(strings.TrimSpace(raw)) - if err != nil || parsed == nil || parsed.Opaque != "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { - return false - } - defaultURL, err := url.Parse(expected) - if err != nil { - return false - } - if !strings.EqualFold(parsed.Scheme, defaultURL.Scheme) || !strings.EqualFold(parsed.Hostname(), defaultURL.Hostname()) { - return false - } - if port := parsed.Port(); port != "" { - portNumber, err := strconv.Atoi(port) - if err != nil || portNumber != 443 { - return false - } - } - path := strings.TrimRight(parsed.Path, "/") - return path == "" || path == strings.TrimRight(defaultURL.Path, "/") -} - func (a *Account) GetGrokAccessToken() string { if !a.IsGrok() { return "" diff --git a/backend/internal/service/grok_credential_failure.go b/backend/internal/service/grok_credential_failure.go index 8886855f12f..65dd8440721 100644 --- a/backend/internal/service/grok_credential_failure.go +++ b/backend/internal/service/grok_credential_failure.go @@ -497,7 +497,12 @@ func (s *OpenAIGatewayService) validateCurrentGrokCredentialFailure( func (s *OpenAIGatewayService) grokCredentialMutationLock(accountID int64) *oauthRefreshLocalLock { actual, _ := s.grokCredentialMutationLocks.LoadOrStore(accountID, newOAuthRefreshLocalLock()) - return actual.(*oauthRefreshLocalLock) + mu, ok := actual.(*oauthRefreshLocalLock) + if !ok { + mu = newOAuthRefreshLocalLock() + s.grokCredentialMutationLocks.Store(accountID, mu) + } + return mu } func (s *OpenAIGatewayService) grokCredentialMutationCommitted(accountID int64, class grokCredentialFailureClass, until time.Time) bool { diff --git a/backend/internal/service/grok_token_provider.go b/backend/internal/service/grok_token_provider.go index c4dc4d4b12c..ae57cdc59e4 100644 --- a/backend/internal/service/grok_token_provider.go +++ b/backend/internal/service/grok_token_provider.go @@ -22,7 +22,6 @@ var ( errGrokOAuthAccessTokenMissing = errors.New("grok oauth access token is missing") errGrokOAuthAccessTokenExpired = errors.New("grok oauth access token is expired") errGrokOAuthConfiguredProxyMiss = errors.New("grok oauth configured proxy is missing") - errGrokOAuthRefreshLockTimeout = errors.New("grok oauth refresh lock wait timed out") ) type GrokTokenCache = GeminiTokenCache diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 0b39ba8a2a5..978b740a562 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -105,7 +105,12 @@ func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until ti func (s *OpenAIGatewayService) openAIAccountRuntimeBlockLock(accountID int64) *sync.Mutex { actual, _ := s.openaiAccountRuntimeBlockLocks.LoadOrStore(accountID, &sync.Mutex{}) - return actual.(*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) { From b32b815e46fe97901179d776387f23d6ba7df0c7 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Tue, 14 Jul 2026 09:43:56 +0800 Subject: [PATCH 3/6] fix(grok): harden OAuth pool recovery --- backend/internal/handler/grok_media.go | 5 + .../internal/handler/openai_alpha_search.go | 3 +- .../handler/openai_chat_completions.go | 3 +- ...i_gateway_credential_failover_loop_test.go | 196 +++++++++++++++- .../handler/openai_gateway_handler.go | 9 +- backend/internal/handler/openai_images.go | 3 +- backend/internal/pkg/xai/oauth.go | 78 ++++++- backend/internal/pkg/xai/oauth_test.go | 39 ++++ backend/internal/repository/account_repo.go | 29 +++ .../account_repo_integration_test.go | 50 +++++ .../internal/service/account_test_service.go | 8 +- backend/internal/service/grok_media.go | 24 +- .../internal/service/grok_quota_service.go | 6 +- .../service/grok_quota_service_test.go | 19 ++ backend/internal/service/grok_upstream_url.go | 90 ++++++++ .../service/grok_upstream_url_test.go | 122 ++++++++++ .../openai_account_runtime_block_fastpath.go | 29 ++- ...nai_account_runtime_block_fastpath_test.go | 39 +++- .../openai_gateway_chat_completions_raw.go | 5 +- .../internal/service/openai_gateway_grok.go | 97 +++++++- .../openai_gateway_grok_chat_bridge.go | 5 +- .../service/openai_gateway_grok_test.go | 211 +++++++++++++++++- .../service/openai_gateway_messages.go | 5 +- .../internal/service/openai_ws_http_bridge.go | 5 +- 24 files changed, 1001 insertions(+), 79 deletions(-) create mode 100644 backend/internal/service/grok_upstream_url.go create mode 100644 backend/internal/service/grok_upstream_url_test.go diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 408020f0a6f..6b10f7bcbe5 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -166,6 +166,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState switchCount := 0 maxAccountSwitches := h.maxAccountSwitches if maxAccountSwitches <= 0 { @@ -300,6 +301,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } switchCount++ + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { + h.handleFailoverExhausted(c, failoverErr, false) + return + } reqLog.Warn("grok_media.upstream_failover_switching", zap.Int64("account_id", account.ID), zap.Int("upstream_status", failoverErr.StatusCode), diff --git a/backend/internal/handler/openai_alpha_search.go b/backend/internal/handler/openai_alpha_search.go index a808145235d..4f5c07d8c5d 100644 --- a/backend/internal/handler/openai_alpha_search.go +++ b/backend/internal/handler/openai_alpha_search.go @@ -106,6 +106,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) var lastFailoverErr *service.UpstreamFailoverError switchCount := 0 + var oauth429FailoverState service.OpenAIOAuth429FailoverState routingStart := time.Now() for { @@ -188,7 +189,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleFailoverExhausted(c, failoverErr, false) return } diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index fc58b3a25ab..004aff2bb72 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -133,6 +133,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState for { if c.Request.Context().Err() != nil { @@ -275,7 +276,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleFailoverExhausted(c, failoverErr, streamStarted) return } 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 e7030d9d2ff..42a6bf6cd36 100644 --- a/backend/internal/handler/openai_gateway_credential_failover_loop_test.go +++ b/backend/internal/handler/openai_gateway_credential_failover_loop_test.go @@ -31,6 +31,7 @@ type grokCredentialHandlerRepo struct { accounts []service.Account setErrorIDs []int64 setTempIDs []int64 + rateLimitIDs []int64 updateExtraIDs []int64 selectionCalls int setErrorErr error @@ -108,6 +109,34 @@ func (r *grokCredentialHandlerRepo) SetTempUnschedulable(_ context.Context, id i return nil } +func (r *grokCredentialHandlerRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error { + r.mu.Lock() + defer r.mu.Unlock() + r.rateLimitIDs = append(r.rateLimitIDs, id) + for i := range r.accounts { + if r.accounts[i].ID != id { + continue + } + now := time.Now() + r.accounts[i].RateLimitedAt = &now + value := resetAt + r.accounts[i].RateLimitResetAt = &value + } + return nil +} + +func (r *grokCredentialHandlerRepo) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error { + r.mu.Lock() + for i := range r.accounts { + if r.accounts[i].ID == id && r.accounts[i].RateLimitResetAt != nil && !resetAt.After(*r.accounts[i].RateLimitResetAt) { + r.mu.Unlock() + return nil + } + } + r.mu.Unlock() + return r.SetRateLimited(ctx, id, resetAt) +} + func (r *grokCredentialHandlerRepo) SetGrokCredentialErrorIfMatch( _ context.Context, id int64, @@ -204,6 +233,12 @@ func (r *grokCredentialHandlerRepo) selectorCalls() int { return r.selectionCalls } +func (r *grokCredentialHandlerRepo) rateLimitedAccountIDs() []int64 { + r.mu.Lock() + defer r.mu.Unlock() + return append([]int64(nil), r.rateLimitIDs...) +} + type grokCredentialHandlerTokenCache struct { service.GrokTokenCache mu sync.Mutex @@ -282,6 +317,8 @@ type grokCredentialHandlerUpstream struct { requestURLs []string authorization []string failAccountID int64 + rateLimitIDs map[int64]bool + failureStatus map[int64]int cancelRequest context.CancelFunc } @@ -295,8 +332,27 @@ func (u *grokCredentialHandlerUpstream) Do(req *http.Request, _ string, accountI u.requestURLs = append(u.requestURLs, req.URL.String()) u.authorization = append(u.authorization, req.Header.Get("Authorization")) failAccountID := u.failAccountID + rateLimited := u.rateLimitIDs[accountID] + failureStatus := u.failureStatus[accountID] cancelRequest := u.cancelRequest u.mu.Unlock() + if rateLimited { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Retry-After": []string{"60"}, + }, + Body: io.NopCloser(bytes.NewBufferString(`{"error":{"message":"rate limited"}}`)), + }, nil + } + if failureStatus > 0 { + return &http.Response{ + StatusCode: failureStatus, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewBufferString(`{"error":{"message":"upstream unavailable"}}`)), + }, nil + } if accountID == failAccountID { if cancelRequest != nil { cancelRequest() @@ -502,6 +558,115 @@ func TestResponsesCredentialFailoverLoop(t *testing.T) { }) } +func TestResponsesGrok429FailoverIsBounded(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("first rate limited account selects healthy account", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429") + 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.rateLimitedAccountIDs()) + }) + + t.Run("two rate limited accounts stop without sweeping the pool", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429") + 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.StatusTooManyRequests, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + require.Equal(t, []int64{801, 802}, repo.rateLimitedAccountIDs()) + require.NotContains(t, recorder.Body.String(), "expired") + require.NotContains(t, recorder.Body.String(), "healthy-access") + require.NotContains(t, recorder.Body.String(), "rate limited") + }) +} + +func TestResponsesGrok429FailoverHandlesMixedStatuses(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("429 then 500 stops after the bounded followup", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "mixed_429_500") + 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.StatusBadGateway, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + require.NotContains(t, recorder.Body.String(), "upstream unavailable") + }) + + t.Run("500 then 429 permits one healthy followup", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "mixed_500_429") + 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.Equal(t, []int64{801, 802, 803}, upstream.accountHits()) + }) + + t.Run("OAuth 429 then API-key failure cannot bypass the bound", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "oauth_429_apikey_500") + 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.StatusBadGateway, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + }) +} + +func TestGrokMedia429FailoverIsBounded(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("first 429 selects one healthy followup", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil) + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + }) + + t.Run("second 429 stops without sweeping a third account", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil) + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + require.NotContains(t, recorder.Body.String(), "rate limited") + }) +} + func TestGrokOAuthCredentialFailoverAcrossHTTPHandlers(t *testing.T) { gin.SetMode(gin.TestMode) endpoints := []struct { @@ -667,9 +832,23 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway }, }, } - if mode == "postmap_cancel" { + if mode == "postmap_cancel" || 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" { + accounts = append(accounts, service.Account{ + ID: 803, Name: "untried-healthy", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, Priority: 3, + Credentials: map[string]any{ + "access_token": "untried-healthy-access", "refresh_token": "untried-healthy-refresh", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), + }, + }) + } + if mode == "oauth_429_apikey_500" { + accounts[1].Type = service.AccountTypeAPIKey + accounts[1].Credentials = map[string]any{"api_key": "third-party-key"} + } if mode == "all_revoked" { accounts[1].Credentials["expires_at"] = time.Now().Add(-time.Minute).UTC().Format(time.RFC3339) } @@ -694,6 +873,21 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway provider.SetRefreshAPI(service.NewOAuthRefreshAPI(repo, tokenCache), refresher) } upstream := &grokCredentialHandlerUpstream{} + switch mode { + case "first_429": + upstream.rateLimitIDs = map[int64]bool{801: true} + case "all_429": + upstream.rateLimitIDs = map[int64]bool{801: true, 802: true} + case "mixed_429_500": + upstream.rateLimitIDs = map[int64]bool{801: true} + upstream.failureStatus = map[int64]int{802: http.StatusInternalServerError} + case "mixed_500_429": + upstream.failureStatus = map[int64]int{801: http.StatusInternalServerError} + upstream.rateLimitIDs = map[int64]bool{802: true} + case "oauth_429_apikey_500": + upstream.rateLimitIDs = map[int64]bool{801: true} + upstream.failureStatus = map[int64]int{802: http.StatusInternalServerError} + } cfg := &config.Config{RunMode: config.RunModeSimple} cfg.Gateway.MaxAccountSwitches = 3 billingCache := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 6db72c270b8..33bef42cef4 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -332,6 +332,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState for { if c.Request.Context().Err() != nil { @@ -487,7 +488,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleFailoverExhausted(c, failoverErr, streamStarted) return } @@ -844,6 +845,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState effectiveMappedModel := preferredMappedModel for { @@ -991,7 +993,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleAnthropicFailoverExhausted(c, failoverErr, streamStarted) return } @@ -1496,6 +1498,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { switchCount := 0 failedAccountIDs := make(map[int64]struct{}) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState handleWSFailover := func(account *service.Account, failoverErr *service.UpstreamFailoverError) bool { if ctx.Err() != nil { return false @@ -1519,7 +1522,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return false } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { closeOpenAIWSFailoverExhausted(wsConn, failoverErr) return false } diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index c5982fb7d1d..d18af28966e 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -145,6 +145,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { stopJSONKeepalive := func() {} jsonKeepaliveStarted := false defer func() { stopJSONKeepalive() }() + var oauth429FailoverState service.OpenAIOAuth429FailoverState for { reqLog.Debug("openai.images.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs))) @@ -299,7 +300,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleFailoverExhausted(c, failoverErr, streamStarted) return } diff --git a/backend/internal/pkg/xai/oauth.go b/backend/internal/pkg/xai/oauth.go index a8a549c9dcf..ea837fd78a5 100644 --- a/backend/internal/pkg/xai/oauth.go +++ b/backend/internal/pkg/xai/oauth.go @@ -5,6 +5,7 @@ import ( "crypto/sha256" "encoding/base64" "encoding/hex" + "errors" "fmt" "net/url" "os" @@ -164,6 +165,24 @@ func ValidatedBaseURL(override string) (string, error) { return ValidateBaseURL(EffectiveBaseURL(override)) } +// BaseURLValidator applies the caller's outbound URL trust policy before xAI +// endpoint paths are appended. The service layer uses this for API-key accounts +// so the global security.url_allowlist policy remains the single source of +// truth; OAuth callers keep using the strict trusted-host validator. +type BaseURLValidator func(string) (string, error) + +func validatedBaseURLWithValidator(override string, validator BaseURLValidator) (string, error) { + if validator == nil { + return ValidatedBaseURL(override) + } + raw := EffectiveBaseURL(override) + validated, err := validator(raw) + if err != nil { + return "", err + } + return normalizeKnownBaseURLPath(validated) +} + type RuntimeSanityCheck struct { Value string `json:"value"` Valid bool `json:"valid"` @@ -282,7 +301,16 @@ func ValidateTrustedBaseURL(raw string) (string, error) { func normalizeKnownBaseURLPath(raw string) (string, error) { parsed, err := url.Parse(raw) if err != nil || parsed.Scheme == "" || parsed.Host == "" { - return "", fmt.Errorf("invalid url: %s", raw) + return "", errors.New("invalid base URL") + } + if parsed.User != nil { + return "", errors.New("base URL must not include userinfo") + } + if parsed.RawQuery != "" { + return "", errors.New("base URL must not include a query") + } + if parsed.Fragment != "" { + return "", errors.New("base URL must not include a fragment") } path := strings.TrimRight(parsed.Path, "/") if path == "" { @@ -435,7 +463,11 @@ func ParseAuthorizationInput(raw string) AuthorizationInput { } func BuildResponsesURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildResponsesURLWithValidator(baseURL, nil) +} + +func BuildResponsesURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -443,7 +475,11 @@ func BuildResponsesURL(baseURL string) (string, error) { } func BuildChatCompletionsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildChatCompletionsURLWithValidator(baseURL, nil) +} + +func BuildChatCompletionsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -451,7 +487,11 @@ func BuildChatCompletionsURL(baseURL string) (string, error) { } func BuildImagesGenerationsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildImagesGenerationsURLWithValidator(baseURL, nil) +} + +func BuildImagesGenerationsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -459,7 +499,11 @@ func BuildImagesGenerationsURL(baseURL string) (string, error) { } func BuildImagesEditsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildImagesEditsURLWithValidator(baseURL, nil) +} + +func BuildImagesEditsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -467,7 +511,11 @@ func BuildImagesEditsURL(baseURL string) (string, error) { } func BuildVideosGenerationsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildVideosGenerationsURLWithValidator(baseURL, nil) +} + +func BuildVideosGenerationsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -475,7 +523,11 @@ func BuildVideosGenerationsURL(baseURL string) (string, error) { } func BuildVideosEditsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildVideosEditsURLWithValidator(baseURL, nil) +} + +func BuildVideosEditsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -483,7 +535,11 @@ func BuildVideosEditsURL(baseURL string) (string, error) { } func BuildVideosExtensionsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildVideosExtensionsURLWithValidator(baseURL, nil) +} + +func BuildVideosExtensionsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -491,7 +547,11 @@ func BuildVideosExtensionsURL(baseURL string) (string, error) { } func BuildVideoURL(baseURL, requestID string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildVideoURLWithValidator(baseURL, requestID, nil) +} + +func BuildVideoURLWithValidator(baseURL, requestID string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } diff --git a/backend/internal/pkg/xai/oauth_test.go b/backend/internal/pkg/xai/oauth_test.go index 39200ffc398..a82a26e2995 100644 --- a/backend/internal/pkg/xai/oauth_test.go +++ b/backend/internal/pkg/xai/oauth_test.go @@ -6,6 +6,7 @@ import ( "net/url" "testing" + "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" "github.com/stretchr/testify/require" ) @@ -165,6 +166,44 @@ func TestValidateBaseURLAllowsPublicThirdPartyGrokAPI(t *testing.T) { require.Error(t, err) } +func TestBuildResponsesURLWithValidatorUsesCallerPolicy(t *testing.T) { + validator := func(raw string) (string, error) { + return urlvalidator.ValidateURLFormat(raw, true) + } + + target, err := BuildResponsesURLWithValidator("http://grok.example.test/v1/", validator) + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/responses", target) +} + +func TestBuildResponsesURLPreservesUnsafeOverrideCustomPath(t *testing.T) { + t.Setenv(EnvAllowUnsafeURLOverrides, "true") + + target, err := BuildResponsesURL("http://localhost:8080/custom") + require.NoError(t, err) + require.Equal(t, "http://localhost:8080/custom/responses", target) +} + +func TestBuildResponsesURLWithValidatorRejectsBaseURLComponents(t *testing.T) { + permissive := func(raw string) (string, error) { return raw, nil } + tests := []struct { + name string + raw string + }{ + {name: "userinfo", raw: "https://user:secret@grok.example.test/v1"}, + {name: "query", raw: "https://grok.example.test/v1?token=secret"}, + {name: "fragment", raw: "https://grok.example.test/v1#secret"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := BuildResponsesURLWithValidator(tt.raw, permissive) + require.Error(t, err) + require.NotContains(t, err.Error(), "secret") + }) + } +} + func TestValidateXAIURLsAllowUnsafeDevOverride(t *testing.T) { t.Setenv(EnvAllowUnsafeURLOverrides, "true") diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index a3eba2eaa94..8527b7010d6 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -1377,6 +1377,35 @@ func (r *accountRepository) SetRateLimitedIfLater(ctx context.Context, id int64, return nil } +// ClearRateLimitIfObserved clears exactly the Grok rate-limit generation seen +// by a successful request. Matching both timestamps prevents a stale success +// from erasing a later clear/re-arm generation with an equal or shorter reset. +func (r *accountRepository) ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error) { + updated, err := r.client.Account.Update(). + Where( + dbaccount.IDEQ(id), + dbaccount.PlatformEQ(service.PlatformGrok), + dbaccount.TypeEQ(service.AccountTypeOAuth), + dbaccount.RateLimitedAtEQ(observedLimitedAt), + dbaccount.RateLimitResetAtEQ(observedResetAt), + ). + ClearRateLimitedAt(). + ClearRateLimitResetAt(). + Save(ctx) + if err != nil { + return false, err + } + if updated == 0 { + r.syncSchedulerAccountSnapshot(ctx, id) + return false, nil + } + if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil { + logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue observed rate-limit clear failed: account=%d err=%v", id, err) + } + r.syncSchedulerAccountSnapshot(ctx, id) + return true, nil +} + func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error { if scope == "" { return nil diff --git a/backend/internal/repository/account_repo_integration_test.go b/backend/internal/repository/account_repo_integration_test.go index f82ad4ca280..ae519f3df44 100644 --- a/backend/internal/repository/account_repo_integration_test.go +++ b/backend/internal/repository/account_repo_integration_test.go @@ -722,6 +722,56 @@ func (s *AccountRepoSuite) TestSetRateLimitedIfLaterDoesNotShortenReset() { s.Require().WithinDuration(later, *cacheRecorder.setAccounts[1].RateLimitResetAt, time.Second) } +func (s *AccountRepoSuite) TestClearRateLimitIfObservedProtectsRearmed429Generation() { + account := mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "acc-rl-conditional-clear", + Platform: service.PlatformGrok, + Type: service.AccountTypeOAuth, + }) + firstReset := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second) + rearmedReset := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second) + + s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, firstReset)) + staleGeneration, err := s.repo.GetByID(s.ctx, account.ID) + s.Require().NoError(err) + s.Require().NotNil(staleGeneration.RateLimitedAt) + s.Require().NotNil(staleGeneration.RateLimitResetAt) + cleared, err := s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt) + s.Require().NoError(err) + s.Require().True(cleared) + + // A newer generation may legitimately re-arm a shorter boundary after the + // first generation was cleared. The stale success must not erase it. + s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, rearmedReset)) + cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt) + s.Require().NoError(err) + s.Require().False(cleared) + + got, err := s.repo.GetByID(s.ctx, account.ID) + s.Require().NoError(err) + s.Require().NotNil(got.RateLimitedAt) + s.Require().NotNil(got.RateLimitResetAt) + s.Require().WithinDuration(rearmedReset, *got.RateLimitResetAt, time.Second) + + // An admin can retype the row while the successful OAuth request is still + // in flight. The stale OAuth recovery must not cross into API-key state even + // when both observed timestamps still match. + _, err = s.client.Account.UpdateOneID(account.ID). + SetType(service.AccountTypeAPIKey). + Save(s.ctx) + s.Require().NoError(err) + cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *got.RateLimitedAt, *got.RateLimitResetAt) + s.Require().NoError(err) + s.Require().False(cleared) + + retyped, err := s.repo.GetByID(s.ctx, account.ID) + s.Require().NoError(err) + s.Require().Equal(service.AccountTypeAPIKey, retyped.Type) + s.Require().NotNil(retyped.RateLimitedAt) + s.Require().NotNil(retyped.RateLimitResetAt) + s.Require().WithinDuration(rearmedReset, *retyped.RateLimitResetAt, time.Second) +} + func (s *AccountRepoSuite) TestClearRateLimit() { account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-clear"}) until := time.Now().Add(1 * time.Hour) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index a9eef998e49..a739bcc120c 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -695,7 +695,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Grok account type: %s", account.Type)) } - apiURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) + apiURL, err := buildGrokResponsesURL(account, s.cfg) if err != nil { return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok base URL: %s", err.Error())) } @@ -742,7 +742,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * now := time.Now() snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now) if snapshot != nil && s.accountRepo != nil { - resetAt, limited := grokRateLimitResetAt(snapshot, now) + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) if limited { normalizeGrokExhaustedWindowResets(snapshot, resetAt, now) } @@ -751,7 +751,11 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * }) if limited { persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) + } else if isSuccessfulGrokRateLimitRecovery(account, snapshot) { + clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account) } + } else if s.accountRepo != nil && isSuccessfulGrokRateLimitRecovery(account, &xai.QuotaSnapshot{StatusCode: resp.StatusCode}) { + clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account) } if resp.StatusCode != http.StatusOK { diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 5cc8e2d05e9..f7c6bd81530 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -13,7 +13,6 @@ import ( "strings" "time" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" @@ -268,25 +267,6 @@ func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount(ctx context.Cont return s.BindStickySession(ctx, groupID, GrokMediaVideoRequestSessionHash(requestID), accountID) } -func (e GrokMediaEndpoint) upstreamURL(baseURL, requestID string) (string, error) { - switch e { - case GrokMediaEndpointImagesGenerations: - return xai.BuildImagesGenerationsURL(baseURL) - case GrokMediaEndpointImagesEdits: - return xai.BuildImagesEditsURL(baseURL) - case GrokMediaEndpointVideosGenerations: - return xai.BuildVideosGenerationsURL(baseURL) - case GrokMediaEndpointVideosEdits: - return xai.BuildVideosEditsURL(baseURL) - case GrokMediaEndpointVideosExtensions: - return xai.BuildVideosExtensionsURL(baseURL) - case GrokMediaEndpointVideoStatus: - return xai.BuildVideoURL(baseURL, requestID) - default: - return "", fmt.Errorf("unsupported grok media endpoint: %s", e) - } -} - func (s *OpenAIGatewayService) ForwardGrokMedia( ctx context.Context, c *gin.Context, @@ -308,7 +288,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( if err != nil { return nil, err } - targetURL, err := endpoint.upstreamURL(account.GetGrokMediaBaseURL(), requestID) + targetURL, err := buildGrokMediaURL(account, s.cfg, endpoint, requestID) if err != nil { return nil, err } @@ -368,7 +348,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel) } - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) if err != nil { return nil, err diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index ef8e3e4c035..a95b7787078 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -135,7 +135,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr if err != nil { return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_PROBE_BODY_ERROR", "failed to build probe body: %v", err) } - targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) + targetURL, err := buildGrokResponsesURL(account, nil) if err != nil { return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err) } @@ -160,7 +160,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr defer func() { _ = resp.Body.Close() }() snapshot := xai.ObserveQuotaHeaders(resp.Header, resp.StatusCode, "active_probe") - resetAt, limited := grokRateLimitResetAt(snapshot, time.Now()) + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, time.Now()) if limited { normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now()) } @@ -169,6 +169,8 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr }) if limited { persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) + } else if isSuccessfulGrokRateLimitRecovery(account, snapshot) { + clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account) } result := &GrokQuotaProbeResult{ diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 4e06fe156ce..ad3c85ba191 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -32,6 +32,10 @@ type grokQuotaAccountRepo struct { lastTempUnschedID int64 lastTempUnschedUntil time.Time lastTempUnschedReason string + recoveryClearCalls int + recoveryObservedAt time.Time + recoveryObservedReset time.Time + recoveryClearResult bool } func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error { @@ -54,6 +58,13 @@ func (r *grokQuotaAccountRepo) SetRateLimitedIfLater(ctx context.Context, id int return r.SetRateLimited(ctx, id, resetAt) } +func (r *grokQuotaAccountRepo) ClearRateLimitIfObserved(_ context.Context, _ int64, observedLimitedAt, observedResetAt time.Time) (bool, error) { + r.recoveryClearCalls++ + r.recoveryObservedAt = observedLimitedAt + r.recoveryObservedReset = observedResetAt + return r.recoveryClearResult, nil +} + func (r *grokQuotaAccountRepo) SetTempUnschedulable(_ context.Context, id int64, until time.Time, reason string) error { r.tempUnschedCalls++ r.lastTempUnschedID = id @@ -375,10 +386,15 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) { t.Parallel() account := healthyGrokQuotaOAuthAccount(45) + observedResetAt := time.Now().Add(-time.Second).UTC().Truncate(time.Second) + observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown) + account.RateLimitedAt = &observedLimitedAt + account.RateLimitResetAt = &observedResetAt repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{45: account}, }, + recoveryClearResult: true, } upstream := &httpUpstreamRecorder{resp: &http.Response{ StatusCode: http.StatusOK, @@ -401,6 +417,9 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) { require.True(t, ok) require.False(t, stored.HeadersObserved) require.Equal(t, http.StatusOK, stored.StatusCode) + require.Equal(t, 1, repo.recoveryClearCalls) + require.Equal(t, observedLimitedAt, repo.recoveryObservedAt) + require.Equal(t, observedResetAt, repo.recoveryObservedReset) } func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { diff --git a/backend/internal/service/grok_upstream_url.go b/backend/internal/service/grok_upstream_url.go new file mode 100644 index 00000000000..ca2ba6b64dd --- /dev/null +++ b/backend/internal/service/grok_upstream_url.go @@ -0,0 +1,90 @@ +package service + +import ( + "errors" + "fmt" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" +) + +func grokBaseURLValidator(account *Account, cfg *config.Config) (xai.BaseURLValidator, error) { + if account == nil || !account.IsGrok() { + return nil, fmt.Errorf("grok account is required") + } + switch account.Type { + case AccountTypeOAuth: + // Subscription credentials are never governed by the operator's API-key + // URL policy. They stay pinned to the supported CLI gateway. + return redactedGrokBaseURLValidator(xai.ValidateTrustedBaseURL), nil + case AccountTypeAPIKey: + if cfg == nil { + return redactedGrokBaseURLValidator(xai.ValidateBaseURL), nil + } + if !cfg.Security.URLAllowlist.Enabled { + return redactedGrokBaseURLValidator(func(raw string) (string, error) { + return urlvalidator.ValidateURLFormat(raw, cfg.Security.URLAllowlist.AllowInsecureHTTP) + }), nil + } + return redactedGrokBaseURLValidator(func(raw string) (string, error) { + return urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{ + AllowedHosts: cfg.Security.URLAllowlist.UpstreamHosts, + RequireAllowlist: true, + AllowPrivate: cfg.Security.URLAllowlist.AllowPrivateHosts, + }) + }), nil + default: + return nil, fmt.Errorf("unsupported grok account type: %s", account.Type) + } +} + +func redactedGrokBaseURLValidator(validator xai.BaseURLValidator) xai.BaseURLValidator { + return func(raw string) (string, error) { + validated, err := validator(raw) + if err != nil { + return "", errors.New("base URL rejected by URL security policy") + } + return validated, nil + } +} + +func buildGrokResponsesURL(account *Account, cfg *config.Config) (string, error) { + validator, err := grokBaseURLValidator(account, cfg) + if err != nil { + return "", err + } + return xai.BuildResponsesURLWithValidator(account.GetGrokBaseURL(), validator) +} + +func buildGrokChatCompletionsURL(account *Account, cfg *config.Config) (string, error) { + validator, err := grokBaseURLValidator(account, cfg) + if err != nil { + return "", err + } + return xai.BuildChatCompletionsURLWithValidator(account.GetGrokBaseURL(), validator) +} + +func buildGrokMediaURL(account *Account, cfg *config.Config, endpoint GrokMediaEndpoint, requestID string) (string, error) { + validator, err := grokBaseURLValidator(account, cfg) + if err != nil { + return "", err + } + baseURL := account.GetGrokMediaBaseURL() + switch endpoint { + case GrokMediaEndpointImagesGenerations: + return xai.BuildImagesGenerationsURLWithValidator(baseURL, validator) + case GrokMediaEndpointImagesEdits: + return xai.BuildImagesEditsURLWithValidator(baseURL, validator) + case GrokMediaEndpointVideosGenerations: + return xai.BuildVideosGenerationsURLWithValidator(baseURL, validator) + case GrokMediaEndpointVideosEdits: + return xai.BuildVideosEditsURLWithValidator(baseURL, validator) + case GrokMediaEndpointVideosExtensions: + return xai.BuildVideosExtensionsURLWithValidator(baseURL, validator) + case GrokMediaEndpointVideoStatus: + return xai.BuildVideoURLWithValidator(baseURL, requestID, validator) + default: + return "", fmt.Errorf("unsupported grok media endpoint: %s", endpoint) + } +} diff --git a/backend/internal/service/grok_upstream_url_test.go b/backend/internal/service/grok_upstream_url_test.go new file mode 100644 index 00000000000..7919e472020 --- /dev/null +++ b/backend/internal/service/grok_upstream_url_test.go @@ -0,0 +1,122 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/stretchr/testify/require" +) + +func TestGrokAPIKeyURLPolicyFollowsGlobalSecurityConfig(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "base_url": "http://grok.example.test/v1", + }, + } + + t.Run("insecure HTTP enabled with allowlist disabled", func(t *testing.T) { + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + + responsesURL, err := buildGrokResponsesURL(account, cfg) + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/responses", responsesURL) + + chatURL, err := buildGrokChatCompletionsURL(account, cfg) + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/chat/completions", chatURL) + + mediaURL, err := buildGrokMediaURL(account, cfg, GrokMediaEndpointImagesGenerations, "") + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/images/generations", mediaURL) + }) + + t.Run("insecure HTTP disabled", func(t *testing.T) { + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = false + + _, err := buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + }) + + t.Run("enabled allowlist remains HTTPS only", func(t *testing.T) { + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = true + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Security.URLAllowlist.UpstreamHosts = []string{"grok.example.test"} + + _, err := buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + }) +} + +func TestGrokAPIKeyURLPolicyAppliesAllowlistAndPrivateHostControls(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "base_url": "https://grok.example.test/v1", + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = true + cfg.Security.URLAllowlist.UpstreamHosts = []string{"grok.example.test"} + + target, err := buildGrokResponsesURL(account, cfg) + require.NoError(t, err) + require.Equal(t, "https://grok.example.test/v1/responses", target) + + cfg.Security.URLAllowlist.UpstreamHosts = []string{"other.example.test"} + _, err = buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + + account.Credentials["base_url"] = "https://127.0.0.1/v1" + cfg.Security.URLAllowlist.UpstreamHosts = []string{"127.0.0.1"} + _, err = buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + + cfg.Security.URLAllowlist.AllowPrivateHosts = true + target, err = buildGrokResponsesURL(account, cfg) + require.NoError(t, err) + require.Equal(t, "https://127.0.0.1/v1/responses", target) +} + +func TestGrokAPIKeyURLPolicyRedactsMalformedConfiguredURL(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "base_url": "https://%zz:secret@grok.example.test/v1", + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + + _, err := buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + require.NotContains(t, err.Error(), "secret") +} + +func TestGrokOAuthURLPolicyIgnoresAPIKeyOverrides(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "base_url": "http://attacker.example.test/v1", + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + + target, err := buildGrokResponsesURL(account, cfg) + require.NoError(t, err) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", target) +} diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 978b740a562..00d971ba224 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -16,6 +16,13 @@ const ( 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 +} + func openAIAccountStateContext(ctx context.Context) (context.Context, context.CancelFunc) { base := context.Background() if ctx != nil { @@ -210,14 +217,28 @@ func (s *OpenAIGatewayService) isOpenAIOAuth429Storm() bool { return s.openaiOAuth429WindowCount.Load() >= openAIOAuth429StormThreshold } -func (s *OpenAIGatewayService) ShouldStopOpenAIOAuth429Failover(account *Account, statusCode int, failedSwitches int) bool { - if statusCode != http.StatusTooManyRequests || failedSwitches < openAIOAuth429StormMaxAccountSwitches { +func (s *OpenAIGatewayService) ShouldStopOpenAIOAuth429Failover(account *Account, statusCode int, failedSwitches int, state *OpenAIOAuth429FailoverState) bool { + if failedSwitches < openAIOAuth429StormMaxAccountSwitches { return false } - if isGrokOAuthAccount(account) { + 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 !isOpenAIOAuthAccount(account) { + 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 } return s.isOpenAIOAuth429Storm() 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 a51c13c5314..7d7cfbda500 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath_test.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath_test.go @@ -135,26 +135,45 @@ func TestShouldStopOpenAIOAuth429Failover_OnlyDuringStorm(t *testing.T) { svc := &OpenAIGatewayService{} account := &Account{ID: 42, Platform: PlatformOpenAI, Type: AccountTypeOAuth} apiKeyAccount := &Account{ID: 43, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + var state OpenAIOAuth429FailoverState - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state)) for i := 0; i < openAIOAuth429StormThreshold; i++ { svc.recordOpenAIOAuth429() } - require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0)) + require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0, &state)) } -func TestShouldStopOpenAIOAuth429Failover_StopsGrokAfterFirst429Switch(t *testing.T) { +func TestShouldStopOpenAIOAuth429Failover_TracksOneGrokFollowupAttempt(t *testing.T) { svc := &OpenAIGatewayService{} account := &Account{ID: 44, Platform: PlatformGrok, Type: AccountTypeOAuth} apiKeyAccount := &Account{ID: 45, Platform: PlatformGrok, Type: AccountTypeAPIKey} - require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1)) + t.Run("429 then 500 stops after one followup", func(t *testing.T) { + var state OpenAIOAuth429FailoverState + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state)) + require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 2, &state)) + }) + + t.Run("500 then 429 still allows one followup", func(t *testing.T) { + var state OpenAIOAuth429FailoverState + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 2, &state)) + require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusBadGateway, 3, &state)) + }) + + t.Run("OAuth 429 then API-key failure consumes the same followup", func(t *testing.T) { + var state OpenAIOAuth429FailoverState + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state)) + require.True(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusInternalServerError, 2, &state)) + }) + + var state OpenAIOAuth429FailoverState + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 2, &state)) } diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 095d544f89b..800d53be324 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -9,7 +9,6 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" @@ -202,7 +201,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( } if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) } // 8. Forward response @@ -222,7 +221,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( func (s *OpenAIGatewayService) rawChatCompletionsURL(account *Account) (string, error) { if account.Platform == PlatformGrok { - targetURL, err := xai.BuildChatCompletionsURL(account.GetGrokBaseURL()) + targetURL, err := buildGrokChatCompletionsURL(account, s.cfg) if err != nil { return "", fmt.Errorf("invalid grok base_url: %w", err) } diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 675b4bc9f95..92cf2910c60 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -11,6 +11,7 @@ import ( "strings" "time" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" @@ -25,6 +26,10 @@ const ( grokCLIVersion = "0.2.93" grokDefaultResponsesModel = "grok-4.5" grokRateLimitFallbackCooldown = 2 * time.Minute + grokRateLimitRepeatCooldown = 10 * time.Minute + grokRateLimitSustainedCooldown = 30 * time.Minute + grokRateLimitMaxAdaptiveCooldown = time.Hour + grokRateLimitBackoffQuietPeriod = time.Hour ) func (s *OpenAIGatewayService) forwardGrokResponses( @@ -61,7 +66,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) defer releaseUpstreamCtx() - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity) + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg) if err != nil { return nil, err } @@ -107,7 +112,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( return s.handleErrorResponse(ctx, resp, c, account, patchedBody, upstreamModel) } - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) var usage *OpenAIUsage var firstTokenMs *int @@ -574,7 +579,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) // Image-description probes are auxiliary requests, not conversation turns. // Do not bind them to the caller's Grok prompt-cache identity. - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "") + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "", s.cfg) releaseUpstreamCtx() if err != nil { return "", OpenAIUsage{}, fmt.Errorf("build grok composer image bridge request: %w", err) @@ -618,7 +623,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg) } - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil) if err != nil { return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err) @@ -740,8 +745,8 @@ func addOpenAIUsage(dst *OpenAIUsage, usage OpenAIUsage) { dst.ImageOutputTokens += usage.ImageOutputTokens } -func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string) (*http.Request, error) { - targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) +func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string, cfg *config.Config) (*http.Request, error) { + targetURL, err := buildGrokResponsesURL(account, cfg) if err != nil { return nil, err } @@ -780,11 +785,12 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco } accountID := account.ID now := time.Now() - resetAt, hasActiveLimit := grokRateLimitResetAt(snapshot, now) + resetAt, hasActiveLimit := grokRateLimitResetAtForAccount(account, snapshot, now) if hasActiveLimit { normalizeGrokExhaustedWindowResets(snapshot, resetAt, now) } - critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit + recovery := isSuccessfulGrokRateLimitRecovery(account, snapshot) + critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit || recovery if s.codexSnapshotThrottle != nil { allowed := s.codexSnapshotThrottle.Allow(accountID, now) if !critical && !allowed { @@ -810,6 +816,23 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco // on the passive snapshot scheduler check. if hasActiveLimit { s.rateLimitGrok(stateCtx, account, resetAt) + } else if recovery { + clearGrokRateLimitAfterRecovery(stateCtx, s.accountRepo, account) + } +} + +func (s *OpenAIGatewayService) updateGrokUsageFromResponse(ctx context.Context, account *Account, headers http.Header, statusCode int) { + snapshot := parseGrokQuotaSnapshot(headers, statusCode, time.Now()) + if snapshot != nil { + s.updateGrokUsageSnapshot(ctx, account, snapshot) + return + } + // Successful responses are recovery evidence even when the upstream omits + // optional quota headers. Do not replace an informative stored snapshot with + // an empty one; only clear the exact observed cooldown generation. + recoverySnapshot := &xai.QuotaSnapshot{StatusCode: statusCode} + if isSuccessfulGrokRateLimitRecovery(account, recoverySnapshot) { + clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account) } } @@ -900,6 +923,37 @@ func grokRateLimitResetAt(snapshot *xai.QuotaSnapshot, now time.Time) (time.Time return time.Time{}, false } +func grokRateLimitResetAtForAccount(account *Account, snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) { + resetAt, limited := grokRateLimitResetAt(snapshot, now) + if !limited || !isGrokOAuthAccount(account) || snapshot == nil || snapshot.StatusCode != http.StatusTooManyRequests { + return resetAt, limited + } + if account.RateLimitedAt == nil || account.RateLimitResetAt == nil { + return resetAt, true + } + previousResetAt := *account.RateLimitResetAt + if previousResetAt.After(now) || now.Sub(previousResetAt) > grokRateLimitBackoffQuietPeriod { + return resetAt, true + } + previousCooldown := previousResetAt.Sub(*account.RateLimitedAt) + if previousCooldown <= 0 { + return resetAt, true + } + + adaptiveCooldown := grokRateLimitRepeatCooldown + switch { + case previousCooldown >= grokRateLimitSustainedCooldown: + adaptiveCooldown = grokRateLimitMaxAdaptiveCooldown + case previousCooldown >= grokRateLimitRepeatCooldown: + adaptiveCooldown = grokRateLimitSustainedCooldown + } + adaptiveResetAt := now.Add(adaptiveCooldown) + if adaptiveResetAt.After(resetAt) { + resetAt = adaptiveResetAt + } + return resetAt, true +} + func normalizeGrokRateLimitResetAt(account *Account, resetAt, now time.Time) time.Time { if !resetAt.After(now) { resetAt = now.Add(grokRateLimitFallbackCooldown) @@ -914,6 +968,33 @@ type grokRateLimitExtendingRepository interface { SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error } +type grokRateLimitRecoveryRepository interface { + ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error) +} + +func isSuccessfulGrokRateLimitRecovery(account *Account, snapshot *xai.QuotaSnapshot) bool { + return isGrokOAuthAccount(account) && + account.RateLimitedAt != nil && + account.RateLimitResetAt != nil && + snapshot != nil && + snapshot.StatusCode >= http.StatusOK && + snapshot.StatusCode < http.StatusMultipleChoices +} + +func clearGrokRateLimitAfterRecovery(ctx context.Context, repo AccountRepository, account *Account) { + if repo == nil || account == nil || account.RateLimitedAt == nil || account.RateLimitResetAt == nil || ctx.Err() != nil { + return + } + recoveryRepo, ok := repo.(grokRateLimitRecoveryRepository) + if !ok { + return + } + _, err := recoveryRepo.ClearRateLimitIfObserved(ctx, account.ID, *account.RateLimitedAt, *account.RateLimitResetAt) + if err != nil { + slog.Warn("grok_rate_limit_recovery_clear_failed", "account_id", account.ID, "error", err) + } +} + func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account *Account, resetAt time.Time) { if repo == nil || account == nil || account.ID <= 0 { return diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge.go b/backend/internal/service/openai_gateway_grok_chat_bridge.go index fe9cfb09fd6..6fd58b2db18 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge.go @@ -10,7 +10,6 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" ) @@ -258,7 +257,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( return nil, fmt.Errorf("get grok access token: %w", err) } upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, cacheIdentity) + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, cacheIdentity, s.cfg) releaseUpstreamCtx() if err != nil { return nil, fmt.Errorf("build grok responses bridge request: %w", err) @@ -301,7 +300,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( return s.handleChatCompletionsErrorResponse(resp, c, account, billingModel) } - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) var result *OpenAIForwardResult if clientStream { diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 4f7ed05edfa..d93415b2e03 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -256,7 +256,7 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) }, } - req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id") + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id", nil) require.NoError(t, err) require.Equal(t, http.MethodPost, req.Method) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String()) @@ -280,7 +280,7 @@ func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing. }, } - req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "") + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "", nil) require.NoError(t, err) require.Equal(t, "https://grok.example.test/v1/responses", req.URL.String()) require.Equal(t, "Bearer api-key", req.Header.Get("Authorization")) @@ -299,7 +299,7 @@ func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T) }, } - req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "") + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "", nil) require.NoError(t, err) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String()) } @@ -1077,6 +1077,41 @@ func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) { require.Contains(t, recorder.Body.String(), `"type":"test_complete"`) } +func TestAccountTestServiceGrokAPIKeyAllowsConfiguredHTTPWhenGlobalPolicyDoes(t *testing.T) { + gin.SetMode(gin.TestMode) + + account := &Account{ + ID: 55, + Name: "grok-api-key-http", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "third-party-key", + "base_url": "http://grok.example.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" + + "data: {\"type\":\"response.completed\"}\n\n", + )), + }} + svc := &AccountTestService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/55/test", nil) + + err := svc.testGrokAccountConnection(c, account, "grok") + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization")) + require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.Contains(t, recorder.Body.String(), `"type":"test_complete"`) +} + func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) { gin.SetMode(gin.TestMode) @@ -1144,10 +1179,15 @@ func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *test c.Set("api_key", &APIKey{ID: 5202}) account := healthyGrokOAuthGatewayTestAccount(56, "access-token") + observedResetAt := time.Now().Add(-time.Second).UTC().Truncate(time.Second) + observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown) + account.RateLimitedAt = &observedLimitedAt + account.RateLimitResetAt = &observedResetAt repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{56: account}, }, + recoveryClearResult: true, } upstream := &httpUpstreamRecorder{resp: &http.Response{ StatusCode: http.StatusOK, @@ -1180,6 +1220,9 @@ func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *test require.False(t, gjson.GetBytes(upstream.lastBody, "tools").Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists()) require.Equal(t, "resp_grok_non_stream", gjson.Get(recorder.Body.String(), "id").String()) + require.Equal(t, 1, repo.recoveryClearCalls) + require.Equal(t, observedLimitedAt, repo.recoveryObservedAt) + require.Equal(t, observedResetAt, repo.recoveryObservedReset) } func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing.T) { @@ -1565,6 +1608,102 @@ func TestHandleGrokAccountUpstreamError429UsesFallbackReset(t *testing.T) { require.Zero(t, repo.tempUnschedCalls) } +func TestGrokRateLimitResetAtForAccountEscalatesRepeated429s(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + retryAfter := 45 + snapshot := &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + RetryAfterSeconds: &retryAfter, + UpdatedAt: now.Format(time.RFC3339), + } + tests := []struct { + name string + previousCooldown time.Duration + wantCooldown time.Duration + }{ + {name: "repeat after short boundary", previousCooldown: 45 * time.Second, wantCooldown: grokRateLimitRepeatCooldown}, + {name: "sustained repeat", previousCooldown: grokRateLimitRepeatCooldown, wantCooldown: grokRateLimitSustainedCooldown}, + {name: "capped repeat", previousCooldown: grokRateLimitSustainedCooldown, wantCooldown: grokRateLimitMaxAdaptiveCooldown}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + previousReset := now.Add(-time.Second) + previousLimited := previousReset.Add(-tt.previousCooldown) + account := &Account{ + ID: 630, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + RateLimitedAt: &previousLimited, + RateLimitResetAt: &previousReset, + } + + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) + + require.True(t, limited) + require.WithinDuration(t, now.Add(tt.wantCooldown), resetAt, time.Second) + }) + } +} + +func TestGrokRateLimitResetAtForAccountPreservesAuthoritativeAndQuietRecovery(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + retryAfter := 45 + previousReset := now.Add(-grokRateLimitBackoffQuietPeriod - time.Second) + previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown) + account := &Account{ + ID: 631, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + RateLimitedAt: &previousLimited, + RateLimitResetAt: &previousReset, + } + snapshot := &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + RetryAfterSeconds: &retryAfter, + UpdatedAt: now.Format(time.RFC3339), + } + + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) + require.True(t, limited) + require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second) + + authoritativeReset := now.Add(2 * time.Hour) + remaining := int64(0) + snapshot.Requests = &xai.QuotaWindow{Remaining: &remaining, ResetUnix: grokInt64PtrForTest(authoritativeReset.Unix())} + recentReset := now.Add(-time.Second) + recentLimited := recentReset.Add(-grokRateLimitSustainedCooldown) + account.RateLimitResetAt = &recentReset + account.RateLimitedAt = &recentLimited + + resetAt, limited = grokRateLimitResetAtForAccount(account, snapshot, now) + require.True(t, limited) + require.WithinDuration(t, authoritativeReset, resetAt, time.Second) +} + +func TestGrokRateLimitResetAtForAccountLeavesAPIKey429PolicyUnchanged(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + retryAfter := 45 + previousReset := now.Add(-time.Second) + previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown) + account := &Account{ + ID: 632, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + RateLimitedAt: &previousLimited, + RateLimitResetAt: &previousReset, + } + snapshot := &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + RetryAfterSeconds: &retryAfter, + UpdatedAt: now.Format(time.RFC3339), + } + + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) + require.True(t, limited) + require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second) +} + func TestGrokRateLimitResetAtUsesFutureWindowAfterRetryAfterExpires(t *testing.T) { now := time.Now().UTC().Truncate(time.Second) observedAt := now.Add(-2 * time.Minute) @@ -1666,6 +1805,72 @@ func TestUpdateGrokUsageSnapshotAvailableSuccessDoesNotSetRateLimited(t *testing require.Zero(t, repo.rateLimitedCalls) } +func TestUpdateGrokUsageFromResponseHeaderlessSuccessClearsObservedCooldown(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + limitedAt := now.Add(-grokRateLimitRepeatCooldown) + observedResetAt := now.Add(-time.Second) + account := &Account{ + ID: 660, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + RateLimitedAt: &limitedAt, + RateLimitResetAt: &observedResetAt, + } + repo := &grokQuotaAccountRepo{recoveryClearResult: true} + svc := &OpenAIGatewayService{ + accountRepo: repo, + codexSnapshotThrottle: newAccountWriteThrottle(time.Hour), + } + + svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK) + + require.Zero(t, repo.updateCalls, "headerless success must not overwrite an informative quota snapshot") + require.Equal(t, 1, repo.recoveryClearCalls) + require.Equal(t, limitedAt, repo.recoveryObservedAt) + require.Equal(t, observedResetAt, repo.recoveryObservedReset) + require.Same(t, &observedResetAt, account.RateLimitResetAt, "shared account snapshots must not be mutated in place") +} + +func TestUpdateGrokUsageFromResponseRecoveryRespectsCancellationAndAPIKeyBoundary(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + observedResetAt := now.Add(-time.Second) + observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown) + + t.Run("parent cancellation does not mutate account state", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + account := &Account{ + ID: 661, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + RateLimitedAt: &observedLimitedAt, + RateLimitResetAt: &observedResetAt, + } + repo := &grokQuotaAccountRepo{recoveryClearResult: true} + svc := &OpenAIGatewayService{accountRepo: repo} + + svc.updateGrokUsageFromResponse(ctx, account, nil, http.StatusOK) + + require.Zero(t, repo.recoveryClearCalls) + }) + + t.Run("API key success does not alter OAuth cooldown state", func(t *testing.T) { + account := &Account{ + ID: 662, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + RateLimitedAt: &observedLimitedAt, + RateLimitResetAt: &observedResetAt, + } + repo := &grokQuotaAccountRepo{recoveryClearResult: true} + svc := &OpenAIGatewayService{accountRepo: repo} + + svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK) + + require.Zero(t, repo.recoveryClearCalls) + }) +} + func TestUpdateGrokUsageSnapshotExhaustedSuccessWithoutResetUsesFallback(t *testing.T) { repo := &grokQuotaAccountRepo{} svc := &OpenAIGatewayService{accountRepo: repo} diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index e2add90ff2c..067d458b890 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -14,7 +14,6 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/claude" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -273,7 +272,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) var upstreamReq *http.Request if account.Platform == PlatformGrok { - upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, grokCacheIdentity) + upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, grokCacheIdentity, s.cfg) } else { upstreamReq, err = s.buildUpstreamRequest(upstreamCtx, c, account, responsesBody, token, isStream, promptCacheKey, false) } @@ -344,7 +343,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( return s.handleAnthropicErrorResponse(resp, c, account, billingModel) } if account.Platform == PlatformGrok && account.Type == AccountTypeOAuth && !account.IsShadow() { - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) } if account.Type == AccountTypeOAuth && promptCacheKey != "" { diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index dbc24c850af..77a5e593299 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -12,7 +12,6 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" ) @@ -193,7 +192,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( releaseUpstreamCtx() return nil, fmt.Errorf("apply grok prompt cache identity: %w", err) } - upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, body, token, grokCacheIdentity) + upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, body, token, grokCacheIdentity, s.cfg) } else { upstreamReq, err = s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token) } @@ -236,7 +235,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg) } if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) } responseID := "" From 81b8b78378f026b88181d30bda4bea468add6109 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Tue, 14 Jul 2026 14:00:00 +0800 Subject: [PATCH 4/6] fix(grok): resolve upstream OAuth imports --- backend/internal/service/grok_oauth_service.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/backend/internal/service/grok_oauth_service.go b/backend/internal/service/grok_oauth_service.go index 49e77bcfd77..92f39379672 100644 --- a/backend/internal/service/grok_oauth_service.go +++ b/backend/internal/service/grok_oauth_service.go @@ -3,8 +3,6 @@ package service import ( "context" "crypto/subtle" - "encoding/base64" - "encoding/json" "errors" "net/http" "strings" From cb7dfea5a9e54f9ddfb9dac55bdd7cc56485b9f5 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Tue, 14 Jul 2026 14:11:05 +0800 Subject: [PATCH 5/6] test(grok): update default model fixture --- backend/internal/service/account_test_service_grok_test.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/backend/internal/service/account_test_service_grok_test.go b/backend/internal/service/account_test_service_grok_test.go index f311402874e..72f995a0c7e 100644 --- a/backend/internal/service/account_test_service_grok_test.go +++ b/backend/internal/service/account_test_service_grok_test.go @@ -93,8 +93,9 @@ func TestAccountTestService_TestAccountConnection_GrokDefaultsEmptyModelTo45(t * Schedulable: true, Concurrency: 1, Credentials: map[string]any{ - "access_token": "grok-access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "access_token": "grok-access-token", + "refresh_token": "grok-refresh-token", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), }, } repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} From 4c764d3fc97e895355799ebf2ffb8a166b68d5d1 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Tue, 14 Jul 2026 14:56:49 +0800 Subject: [PATCH 6/6] test(grok): use healthy quota account fixture --- .../service/grok_quota_service_test.go | 67 +++---------------- 1 file changed, 8 insertions(+), 59 deletions(-) diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index ad3c85ba191..52def9f101b 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -302,16 +302,7 @@ func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T func TestGrokQuotaServiceProbeUsageRedactsUpstreamErrorBodyFromErrorAndLogs(t *testing.T) { const upstreamSecret = "upstream-secret-refresh-token" - account := &Account{ - ID: 49, - Platform: PlatformGrok, - Type: AccountTypeOAuth, - Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(49) repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{49: account}, @@ -453,13 +444,7 @@ func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { func TestGrokQuotaServiceQueryQuotaFreeFallsBackToGrok45(t *testing.T) { t.Parallel() - account := &Account{ - ID: 51, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(51) repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{account.ID: account}, }} @@ -500,13 +485,7 @@ func TestGrokQuotaServiceQueryQuotaFreeFallsBackToGrok45(t *testing.T) { func TestGrokQuotaServiceQueryQuotaPaidBillingSkipsActiveProbe(t *testing.T) { t.Parallel() - account := &Account{ - ID: 52, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(52) repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{account.ID: account}, }} @@ -534,13 +513,7 @@ func TestGrokQuotaServiceQueryQuotaPaidBillingSkipsActiveProbe(t *testing.T) { func TestGrokQuotaServiceQueryQuotaCustomPaidMonthlyLimitSkipsActiveProbe(t *testing.T) { t.Parallel() - account := &Account{ - ID: 57, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(57) repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{account.ID: account}, }} @@ -673,13 +646,7 @@ func TestGrokLocalUsageForBillingOnlyReturnsAvailableWindows(t *testing.T) { func TestAccountUsageServiceGrokRefreshUsesBillingOnly(t *testing.T) { t.Parallel() - account := &Account{ - ID: 54, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(54) repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{account.ID: account}, }} @@ -714,13 +681,7 @@ func TestAccountUsageServiceGrokRefreshUsesBillingOnly(t *testing.T) { func TestGrokQuotaServiceProbeFlightsDeduplicateBillingAndSeparateActive(t *testing.T) { t.Parallel() - account := &Account{ - ID: 55, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(55) repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{account.ID: account}, }} @@ -776,13 +737,7 @@ func TestGrokQuotaServiceProbeFlightsDeduplicateBillingAndSeparateActive(t *test func TestGrokQuotaServiceBilling429DoesNotPauseModelScheduling(t *testing.T) { t.Parallel() - account := &Account{ - ID: 56, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(56) repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{account.ID: account}, }} @@ -802,13 +757,7 @@ func TestGrokQuotaServiceBilling429DoesNotPauseModelScheduling(t *testing.T) { func TestGrokQuotaServiceQueryQuotaFree429PersistsLimitAndKeepsBilling(t *testing.T) { t.Parallel() - account := &Account{ - ID: 53, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, - Credentials: map[string]any{ - "access_token": "access-token", - "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), - }, - } + account := healthyGrokQuotaOAuthAccount(53) repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{account.ID: account}, }}