From 8bc0a277f205cbd5a3a7b1199c1bcc3f7d6a39dc Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Fri, 17 Jul 2026 11:51:16 +0800 Subject: [PATCH] fix(grok): quarantine ineligible media accounts --- backend/internal/handler/grok_media.go | 20 ++- backend/internal/handler/grok_media_test.go | 21 +++ backend/internal/pkg/xai/billing.go | 2 + backend/internal/service/account.go | 60 ++++++++- .../account_grok_media_eligibility_test.go | 121 ++++++++++++++++++ backend/internal/service/admin_account.go | 66 ++++++++++ .../internal/service/grok_quota_service.go | 27 +++- .../service/grok_quota_service_test.go | 29 +++++ .../service/openai_gateway_scheduling.go | 4 + 9 files changed, 346 insertions(+), 4 deletions(-) create mode 100644 backend/internal/service/account_grok_media_eligibility_test.go diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index eea6149e8c..32cc855333 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -173,6 +173,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. maxAccountSwitches = 3 } routingStart := time.Now() + requiredCapability := grokMediaRequiredCapability(endpoint) for { if failoverClientGone(c) { @@ -186,7 +187,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. requestModel, failedAccountIDs, service.OpenAIUpstreamTransportHTTPSSE, - "", + requiredCapability, false, false, false, @@ -201,6 +202,11 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. zap.Error(err), zap.Int("excluded_account_count", len(failedAccountIDs)), ) + if endpoint.IsGenerationRequest() && len(failedAccountIDs) == 0 && errors.Is(err, service.ErrNoAvailableAccounts) { + markOpsRoutingCapacityLimitedIfNoAvailable(c, err) + h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") + return + } if len(failedAccountIDs) == 0 { cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok) if !cls.ModelNotFound { @@ -217,6 +223,11 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } if selection == nil || selection.Account == nil { + if endpoint.IsGenerationRequest() { + markOpsRoutingCapacityLimited(c) + h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") + return + } cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok) if !cls.ModelNotFound { markOpsRoutingCapacityLimited(c) @@ -354,6 +365,13 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. } } +func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.OpenAIEndpointCapability { + if endpoint.IsGenerationRequest() { + return service.OpenAIEndpointCapabilityGrokMediaGeneration + } + return "" +} + func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string) bool { return endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) != "" } diff --git a/backend/internal/handler/grok_media_test.go b/backend/internal/handler/grok_media_test.go index 1b82f8fa6f..3c58742c78 100644 --- a/backend/internal/handler/grok_media_test.go +++ b/backend/internal/handler/grok_media_test.go @@ -52,3 +52,24 @@ func TestShouldRecordGrokMediaUsage(t *testing.T) { }) } } + +func TestGrokMediaRequiredCapability(t *testing.T) { + tests := []struct { + name string + endpoint service.GrokMediaEndpoint + want service.OpenAIEndpointCapability + }{ + {name: "image generation", endpoint: service.GrokMediaEndpointImagesGenerations, want: service.OpenAIEndpointCapabilityGrokMediaGeneration}, + {name: "image edit", endpoint: service.GrokMediaEndpointImagesEdits, want: service.OpenAIEndpointCapabilityGrokMediaGeneration}, + {name: "video generation", endpoint: service.GrokMediaEndpointVideosGenerations, want: service.OpenAIEndpointCapabilityGrokMediaGeneration}, + {name: "video edit", endpoint: service.GrokMediaEndpointVideosEdits, want: service.OpenAIEndpointCapabilityGrokMediaGeneration}, + {name: "video extension", endpoint: service.GrokMediaEndpointVideosExtensions, want: service.OpenAIEndpointCapabilityGrokMediaGeneration}, + {name: "video status preserves lookup", endpoint: service.GrokMediaEndpointVideoStatus, want: ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, grokMediaRequiredCapability(tt.endpoint)) + }) + } +} diff --git a/backend/internal/pkg/xai/billing.go b/backend/internal/pkg/xai/billing.go index f6f0c0948a..d8580ebab0 100644 --- a/backend/internal/pkg/xai/billing.go +++ b/backend/internal/pkg/xai/billing.go @@ -76,6 +76,8 @@ type BillingSummary struct { UsedPercent *float64 `json:"used_percent,omitempty"` Plan string `json:"plan,omitempty"` // SuperGrok | SuperGrok Heavy | "" StatusCode int `json:"status_code,omitempty"` + WeeklyStatusCode int `json:"weekly_status_code,omitempty"` + MonthlyStatusCode int `json:"monthly_status_code,omitempty"` Source string `json:"source,omitempty"` FetchedAt string `json:"fetched_at,omitempty"` UpdatedAt string `json:"updated_at,omitempty"` diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 701e14d742..f5e0c07d04 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -89,6 +89,11 @@ const ( OpenAIEndpointCapabilityChatCompletions OpenAIEndpointCapability = "chat_completions" OpenAIEndpointCapabilityEmbeddings OpenAIEndpointCapability = "embeddings" OpenAIEndpointCapabilityAlphaSearch OpenAIEndpointCapability = "alpha_search" + // OpenAIEndpointCapabilityGrokMediaGeneration keeps image/video generation + // away from Grok accounts that are explicitly disabled or whose billing + // entitlement probe was forbidden. Video status lookups intentionally do not + // require this capability so already-submitted requests remain queryable. + OpenAIEndpointCapabilityGrokMediaGeneration OpenAIEndpointCapability = "grok_media_generation" // OpenAIEndpointCapabilityResponses 表示上游确实提供 /v1/responses 端点。 // 与其他能力不同：支持状态来自 accounts.extra 的自动探测标记 // （openai_responses_supported / openai_responses_mode），而非 @@ -99,6 +104,11 @@ const ( const openAIEndpointCapabilitiesCredentialKey = "openai_capabilities" +// GrokMediaEligibleExtraKey is an optional per-account override stored in +// accounts.extra. true forces media routing on, false disables it, and an +// absent/null value uses provider observations. +const GrokMediaEligibleExtraKey = "grok_media_eligible" + const ( OpenAIAuthModePersonalAccessToken = "personalAccessToken" openAIAuthModeCredentialKey = "auth_mode" @@ -1409,7 +1419,15 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa return false } if a.IsGrok() { - return capability == OpenAIEndpointCapabilityChatCompletions + switch capability { + case OpenAIEndpointCapabilityChatCompletions: + return true + case OpenAIEndpointCapabilityGrokMediaGeneration: + eligible, _ := a.GrokMediaGenerationEligibility() + return eligible + default: + return false + } } switch capability { case OpenAIEndpointCapabilityChatCompletions: @@ -1449,6 +1467,46 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa return configured[string(capability)] } +// GrokMediaGenerationEligibility reports whether a Grok account may receive +// new image/video generation requests. Missing observations preserve legacy +// routing; operators can fail closed for known-bad accounts with the explicit +// override. A successful override takes precedence over stale probe data. +func (a *Account) GrokMediaGenerationEligibility() (bool, string) { + if a == nil || !a.IsGrok() { + return false, "not_grok" + } + if override, ok := grokMediaEligibilityOverride(a.Extra); ok { + if override { + return true, "override_enabled" + } + return false, "override_disabled" + } + if a.Type != AccountTypeOAuth { + return true, "non_oauth" + } + + billing, err := grokBillingSnapshotFromExtra(a.Extra) + if err != nil || billing == nil { + return true, "billing_unobserved" + } + if billing.StatusCode == 403 || billing.WeeklyStatusCode == 403 || billing.MonthlyStatusCode == 403 { + return false, "billing_forbidden" + } + return true, "eligible" +} + +func grokMediaEligibilityOverride(extra map[string]any) (bool, bool) { + if extra == nil { + return false, false + } + raw, exists := extra[GrokMediaEligibleExtraKey] + if !exists || raw == nil { + return false, false + } + value, ok := raw.(bool) + return value, ok +} + func (a *Account) openAIEndpointCapabilitySet() (map[string]bool, bool) { if a == nil || a.Credentials == nil { return nil, false diff --git a/backend/internal/service/account_grok_media_eligibility_test.go b/backend/internal/service/account_grok_media_eligibility_test.go new file mode 100644 index 0000000000..86bdb67939 --- /dev/null +++ b/backend/internal/service/account_grok_media_eligibility_test.go @@ -0,0 +1,121 @@ +//go:build unit + +package service + +import ( + "context" + "net/http" + "testing" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/stretchr/testify/require" +) + +func TestGrokMediaGenerationEligibility(t *testing.T) { + forbiddenBilling := &xai.BillingSummary{ + StatusCode: http.StatusForbidden, + WeeklyStatusCode: http.StatusForbidden, + MonthlyStatusCode: http.StatusForbidden, + } + weeklyAllowance := &xai.BillingSummary{ + PeriodType: "weekly", + StatusCode: http.StatusOK, + WeeklyStatusCode: http.StatusOK, + } + + tests := []struct { + name string + account *Account + want bool + wantReason string + }{ + {name: "nil account", account: nil, want: false, wantReason: "not_grok"}, + {name: "non grok account", account: &Account{Platform: PlatformOpenAI}, want: false, wantReason: "not_grok"}, + {name: "non oauth grok account stays eligible", account: &Account{Platform: PlatformGrok, Type: AccountTypeAPIKey}, want: true, wantReason: "non_oauth"}, + {name: "unobserved oauth preserves legacy routing", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: true, wantReason: "billing_unobserved"}, + {name: "weekly allowance is not treated as weekly subscription", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"}, + {name: "billing forbidden is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: forbiddenBilling}}, want: false, wantReason: "billing_forbidden"}, + {name: "explicit disable wins", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}, want: false, wantReason: "override_disabled"}, + {name: "explicit enable wins over forbidden probe", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: true, grokBillingExtraKey: forbiddenBilling}}, want: true, wantReason: "override_enabled"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, reason := tt.account.GrokMediaGenerationEligibility() + require.Equal(t, tt.want, got) + require.Equal(t, tt.wantReason, reason) + }) + } +} + +func TestGrokMediaCapabilityFiltersOnlyGeneration(t *testing.T) { + account := &Account{ + ID: 1, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Extra: map[string]any{GrokMediaEligibleExtraKey: false}, + } + + require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions)) + require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration)) + require.False(t, isOpenAICompatibleAccountEligibleForRequest( + context.Background(), account, PlatformGrok, "grok-imagine-video", false, + OpenAIEndpointCapabilityGrokMediaGeneration, + )) +} + +func TestNormalizeGrokMediaEligibilityExtra(t *testing.T) { + t.Run("boolean override is accepted", func(t *testing.T) { + extra, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: false}) + + require.NoError(t, err) + require.Equal(t, false, extra[GrokMediaEligibleExtraKey]) + }) + + t.Run("null clears override", func(t *testing.T) { + extra, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: nil}) + + require.NoError(t, err) + require.NotContains(t, extra, GrokMediaEligibleExtraKey) + }) + + t.Run("malformed override is rejected", func(t *testing.T) { + _, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: "false"}) + + require.Error(t, err) + require.Equal(t, http.StatusBadRequest, infraerrors.Code(err)) + }) + + t.Run("other platforms ignore provider owned value", func(t *testing.T) { + extra := map[string]any{GrokMediaEligibleExtraKey: "provider-owned"} + normalized, err := normalizeGrokMediaEligibilityExtra(PlatformOpenAI, extra) + + require.NoError(t, err) + require.Equal(t, extra, normalized) + }) +} + +func TestNormalizeGrokMediaEligibilityUpdateExtra(t *testing.T) { + account := &Account{Platform: PlatformGrok, Extra: map[string]any{GrokMediaEligibleExtraKey: false}} + + t.Run("omitted override preserves current value", func(t *testing.T) { + input := &UpdateAccountInput{Extra: map[string]any{"quota_used": float64(1)}} + normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{"quota_used": float64(1)}) + + require.NoError(t, err) + require.Equal(t, false, normalized[GrokMediaEligibleExtraKey]) + }) + + t.Run("null removes current override", func(t *testing.T) { + input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: nil}} + normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{GrokMediaEligibleExtraKey: nil}) + + require.NoError(t, err) + require.NotContains(t, normalized, GrokMediaEligibleExtraKey) + require.Contains(t, input.Extra, GrokMediaEligibleExtraKey) + }) +} diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go index 904153f5aa..5935a8cf6f 100644 --- a/backend/internal/service/admin_account.go +++ b/backend/internal/service/admin_account.go @@ -394,6 +394,64 @@ func normalizeOpenAILongContextBillingUpdateExtra(account *Account, input *Updat return normalized, nil } +// ValidateGrokMediaEligibilityExtra validates the optional media-routing +// override. null removes the override and returns the account to automatic +// provider-observation based routing. +func ValidateGrokMediaEligibilityExtra(platform string, extra map[string]any) error { + if platform != PlatformGrok || extra == nil { + return nil + } + raw, exists := extra[GrokMediaEligibleExtraKey] + if !exists || raw == nil { + return nil + } + if _, ok := raw.(bool); !ok { + return infraerrors.BadRequest( + "GROK_MEDIA_ELIGIBILITY_INVALID", + "grok_media_eligible must be a boolean or null", + ) + } + return nil +} + +func normalizeGrokMediaEligibilityExtra(platform string, extra map[string]any) (map[string]any, error) { + if platform != PlatformGrok { + return extra, nil + } + if err := ValidateGrokMediaEligibilityExtra(platform, extra); err != nil { + return nil, err + } + normalized := maps.Clone(extra) + if normalized != nil && normalized[GrokMediaEligibleExtraKey] == nil { + delete(normalized, GrokMediaEligibleExtraKey) + } + return normalized, nil +} + +func normalizeGrokMediaEligibilityUpdateExtra(account *Account, input *UpdateAccountInput, normalized map[string]any) (map[string]any, error) { + if account == nil || account.Platform != PlatformGrok { + return normalized, nil + } + if err := ValidateGrokMediaEligibilityExtra(account.Platform, input.Extra); err != nil { + return nil, err + } + normalized = maps.Clone(normalized) + if normalized == nil { + normalized = make(map[string]any) + } + raw, provided := input.Extra[GrokMediaEligibleExtraKey] + if provided { + if raw == nil { + delete(normalized, GrokMediaEligibleExtraKey) + } + return normalized, nil + } + if current, ok := account.Extra[GrokMediaEligibleExtraKey].(bool); ok { + normalized[GrokMediaEligibleExtraKey] = current + } + return normalized, nil +} + func buildAccountForCreate(input *CreateAccountInput, accountExtra map[string]any) (*Account, error) { // Probe state is system-managed. New accounts always start with auto probe disabled. delete(accountExtra, UpstreamBillingProbeEnabledExtraKey) @@ -448,6 +506,10 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou if err != nil { return nil, err } + accountExtra, err = normalizeGrokMediaEligibilityExtra(input.Platform, accountExtra) + if err != nil { + return nil, err + } // 绑定分组 groupIDs := input.GroupIDs @@ -535,6 +597,10 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U if err != nil { return nil, err } + normalizedExtra, err = normalizeGrokMediaEligibilityUpdateExtra(account, input, normalizedExtra) + if err != nil { + return nil, err + } } previousProbeIdentity := upstreamBillingProbeIdentity(account) // 安全/身份不变量(影子账号):通用更新路径被 edit/re-auth/refresh/batch 共用, diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 4ce9fafcfb..45e7be83ac 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -249,12 +249,25 @@ func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (* weeklyOK := weekly.summary != nil monthlyOK := monthly.summary != nil + previous, _ := grokBillingSnapshotFromExtra(account.Extra) if !weeklyOK && !monthlyOK { - return nil, mergeGrokBillingProbeErrors(weekly.status, monthly.status, weekly.err, monthly.err) + probeErr := mergeGrokBillingProbeErrors(weekly.status, monthly.status, weekly.err, monthly.err) + billing := xai.MergeBillingProbeResult(previous, nil, nil, false, false) + if billing == nil { + billing = &xai.BillingSummary{Partial: true, FailedWindows: []string{"weekly", "monthly"}} + } + billing.WeeklyStatusCode = weekly.status + billing.MonthlyStatusCode = monthly.status + billing = xai.StampBillingSummary(billing, preferBillingObservationStatus(weekly.status, monthly.status), "billing_probe") + if persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{grokBillingExtraKey: billing}); persistErr != nil { + slog.Warn("grok_billing_failure_persist_failed", "account_id", account.ID, "error", persistErr) + } + return nil, probeErr } statusCode := preferSuccessfulBillingStatus(weekly.status, monthly.status, weeklyOK, monthlyOK) - previous, _ := grokBillingSnapshotFromExtra(account.Extra) billing := xai.MergeBillingProbeResult(previous, weekly.summary, monthly.summary, weeklyOK, monthlyOK) + billing.WeeklyStatusCode = weekly.status + billing.MonthlyStatusCode = monthly.status billing = xai.StampBillingSummary(billing, statusCode, "billing_probe") persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ grokBillingExtraKey: billing, @@ -276,6 +289,16 @@ func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (* }, nil } +func preferBillingObservationStatus(weeklyStatus, monthlyStatus int) int { + if weeklyStatus == http.StatusForbidden || monthlyStatus == http.StatusForbidden { + return http.StatusForbidden + } + if weeklyStatus != 0 { + return weeklyStatus + } + return monthlyStatus +} + func (s *GrokQuotaService) runProbeFlight( ctx context.Context, key string, diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index f1d7001a31..f4b84a8b3e 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -755,6 +755,35 @@ func TestGrokQuotaServiceBilling429DoesNotPauseModelScheduling(t *testing.T) { require.Zero(t, repo.rateLimitedCalls) } +func TestGrokQuotaServiceBilling403PersistsMediaEligibilitySignal(t *testing.T) { + t.Parallel() + + account := healthyGrokQuotaOAuthAccount(58) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{billingStatus: http.StatusForbidden} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil) + + result, err := svc.ProbeBilling(context.Background(), account.ID) + + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, 1, repo.updateCalls) + raw := repo.updates[account.ID][grokBillingExtraKey] + billing, ok := raw.(*xai.BillingSummary) + require.True(t, ok) + require.Equal(t, http.StatusForbidden, billing.StatusCode) + require.Equal(t, http.StatusForbidden, billing.WeeklyStatusCode) + require.Equal(t, http.StatusForbidden, billing.MonthlyStatusCode) + require.True(t, billing.Partial) + + account.Extra = map[string]any{grokBillingExtraKey: billing} + eligible, reason := account.GrokMediaGenerationEligibility() + require.False(t, eligible) + require.Equal(t, "billing_forbidden", reason) +} + func TestGrokQuotaServiceQueryQuotaFree429PersistsLimitAndKeepsBilling(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index a8ed5fd2b5..1ce86249b6 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -235,6 +235,10 @@ func isOpenAICompatibleAccountEligibleForRequest(ctx context.Context, account *A return false } if !account.SupportsOpenAIEndpointCapability(requiredCapability) { + if account.IsGrok() && requiredCapability == OpenAIEndpointCapabilityGrokMediaGeneration { + _, reason := account.GrokMediaGenerationEligibility() + slog.Debug("grok_media_account_ineligible", "account_id", account.ID, "reason", reason) + } return false } if requireCompact && (!account.IsOpenAI() || openAICompactSupportTier(account) == 0) {