From d2ac6b2c979db65d6dfcee0dd93be52e16834620 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Fri, 17 Jul 2026 12:08:50 +0800 Subject: [PATCH] fix(grok): harden media quarantine edge cases --- .../account_grok_media_eligibility_test.go | 39 ++++++++ backend/internal/service/admin_account.go | 1 + ...min_account_upstream_billing_probe_test.go | 30 +++++++ backend/internal/service/grok_media.go | 2 +- .../service/grok_quota_service_test.go | 90 ++++++++++++++++--- .../service/openai_account_scheduler_test.go | 65 ++++++++++++++ .../service/openai_gateway_grok_test.go | 20 ++++- .../service/openai_gateway_scheduling.go | 17 +++- 8 files changed, 244 insertions(+), 20 deletions(-) diff --git a/backend/internal/service/account_grok_media_eligibility_test.go b/backend/internal/service/account_grok_media_eligibility_test.go index 86bdb6793..be7b4759a 100644 --- a/backend/internal/service/account_grok_media_eligibility_test.go +++ b/backend/internal/service/account_grok_media_eligibility_test.go @@ -23,6 +23,16 @@ func TestGrokMediaGenerationEligibility(t *testing.T) { StatusCode: http.StatusOK, WeeklyStatusCode: http.StatusOK, } + weeklyForbidden := &xai.BillingSummary{ + StatusCode: http.StatusOK, + WeeklyStatusCode: http.StatusForbidden, + MonthlyStatusCode: http.StatusOK, + } + monthlyForbidden := &xai.BillingSummary{ + StatusCode: http.StatusOK, + WeeklyStatusCode: http.StatusOK, + MonthlyStatusCode: http.StatusForbidden, + } tests := []struct { name string @@ -36,6 +46,10 @@ func TestGrokMediaGenerationEligibility(t *testing.T) { {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: "weekly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyForbidden}}, want: false, wantReason: "billing_forbidden"}, + {name: "monthly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: monthlyForbidden}}, want: false, wantReason: "billing_forbidden"}, + {name: "malformed billing observation preserves legacy routing", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: make(chan int)}}, want: true, wantReason: "billing_unobserved"}, + {name: "malformed override falls back to observations", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: "false", grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"}, {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"}, } @@ -118,4 +132,29 @@ func TestNormalizeGrokMediaEligibilityUpdateExtra(t *testing.T) { require.NotContains(t, normalized, GrokMediaEligibleExtraKey) require.Contains(t, input.Extra, GrokMediaEligibleExtraKey) }) + + t.Run("provided boolean replaces current override", func(t *testing.T) { + input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: true}} + normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{GrokMediaEligibleExtraKey: true}) + + require.NoError(t, err) + require.Equal(t, true, normalized[GrokMediaEligibleExtraKey]) + }) + + t.Run("malformed override is rejected on update", func(t *testing.T) { + input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: "false"}} + _, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, nil) + + require.Error(t, err) + require.Equal(t, http.StatusBadRequest, infraerrors.Code(err)) + }) + + t.Run("non grok update is unchanged", func(t *testing.T) { + input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}} + normalized := map[string]any{GrokMediaEligibleExtraKey: "provider-owned"} + got, err := normalizeGrokMediaEligibilityUpdateExtra(&Account{Platform: PlatformOpenAI}, input, normalized) + + require.NoError(t, err) + require.Equal(t, normalized, got) + }) } diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go index 5935a8cf6..899679221 100644 --- a/backend/internal/service/admin_account.go +++ b/backend/internal/service/admin_account.go @@ -673,6 +673,7 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U "quota_daily_start", "quota_weekly_used", "quota_weekly_start", + grokBillingExtraKey, UpstreamBillingProbeEnabledExtraKey, UpstreamBillingProbeExtraKey, } { diff --git a/backend/internal/service/admin_account_upstream_billing_probe_test.go b/backend/internal/service/admin_account_upstream_billing_probe_test.go index 317423aae..727a66cf1 100644 --- a/backend/internal/service/admin_account_upstream_billing_probe_test.go +++ b/backend/internal/service/admin_account_upstream_billing_probe_test.go @@ -2,8 +2,10 @@ package service import ( "context" + "net/http" "testing" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/stretchr/testify/require" ) @@ -62,6 +64,34 @@ func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit( require.Equal(t, "value", updated.Extra["custom"]) } +func TestUpdateAccountPreservesGrokBillingSnapshotForUnrelatedEdit(t *testing.T) { + accountID := int64(112) + billing := &xai.BillingSummary{ + StatusCode: http.StatusForbidden, + WeeklyStatusCode: http.StatusForbidden, + } + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + accountID: { + ID: accountID, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Extra: map[string]any{grokBillingExtraKey: billing}, + }, + }} + + updated, err := (&adminServiceImpl{accountRepo: repo}).UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + Extra: map[string]any{"custom": "value"}, + }) + + require.NoError(t, err) + require.Equal(t, billing, updated.Extra[grokBillingExtraKey]) + require.Equal(t, "value", updated.Extra["custom"]) + eligible, reason := updated.GrokMediaGenerationEligibility() + require.False(t, eligible) + require.Equal(t, "billing_forbidden", reason) +} + func TestUpdateAccountPreservesProbeSnapshotWhenIdentityValuesAreUnchanged(t *testing.T) { accountID := int64(119) repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 2d0630230..c75bd271f 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -515,7 +515,7 @@ func canonicalizeGrokMediaImageURLObject(body []byte, path string) ([]byte, erro } out := body - if !gjson.GetBytes(out, path+".url").Exists() { + if strings.TrimSpace(gjson.GetBytes(out, path+".url").String()) == "" { var err error out, err = sjson.SetBytes(out, path+".url", legacy.Value()) if err != nil { diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index f4b84a8b3..59b19ef27 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -99,18 +99,20 @@ func (r *grokQuotaUsageLogRepo) GetAccountTodayStats(context.Context, int64) (*u type grokHybridUpstream struct { httpUpstreamRecorder - mu sync.Mutex - requests []*http.Request - bodies [][]byte - weeklyUsagePercent *float64 - monthlyLimitCents *float64 - activeStatus int - activeHeaders http.Header - billingStarted chan struct{} - billingRelease <-chan struct{} - billingStartOnce sync.Once - billingStatus int - billingHeaders http.Header + mu sync.Mutex + requests []*http.Request + bodies [][]byte + weeklyUsagePercent *float64 + monthlyLimitCents *float64 + activeStatus int + activeHeaders http.Header + billingStarted chan struct{} + billingRelease <-chan struct{} + billingStartOnce sync.Once + billingStatus int + weeklyBillingStatus int + monthlyBillingStatus int + billingHeaders http.Header } func (u *grokHybridUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { @@ -147,9 +149,16 @@ func (u *grokHybridUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*h return nil, req.Context().Err() } } - if u.billingStatus != 0 && u.billingStatus != http.StatusOK { + billingStatus := u.billingStatus + if req.URL.RawQuery == "format=credits" && u.weeklyBillingStatus != 0 { + billingStatus = u.weeklyBillingStatus + } + if req.URL.RawQuery != "format=credits" && u.monthlyBillingStatus != 0 { + billingStatus = u.monthlyBillingStatus + } + if billingStatus != 0 && billingStatus != http.StatusOK { return &http.Response{ - StatusCode: u.billingStatus, + StatusCode: billingStatus, Header: u.billingHeaders, Body: io.NopCloser(strings.NewReader(`{"error":{"message":"billing limited"}}`)), }, nil @@ -784,6 +793,59 @@ func TestGrokQuotaServiceBilling403PersistsMediaEligibilitySignal(t *testing.T) require.Equal(t, "billing_forbidden", reason) } +func TestGrokQuotaServicePartialBilling403PersistsMediaEligibilitySignal(t *testing.T) { + t.Parallel() + + account := healthyGrokQuotaOAuthAccount(59) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{ + weeklyBillingStatus: http.StatusForbidden, + monthlyBillingStatus: http.StatusOK, + } + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil) + + result, err := svc.ProbeBilling(context.Background(), account.ID) + + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.Billing) + require.Equal(t, http.StatusOK, result.StatusCode) + require.Equal(t, http.StatusForbidden, result.Billing.WeeklyStatusCode) + require.Equal(t, http.StatusOK, result.Billing.MonthlyStatusCode) + require.True(t, result.Billing.Partial) + require.Contains(t, result.Billing.FailedWindows, "weekly") + require.Equal(t, 1, repo.updateCalls) + + account.Extra = map[string]any{grokBillingExtraKey: result.Billing} + eligible, reason := account.GrokMediaGenerationEligibility() + require.False(t, eligible) + require.Equal(t, "billing_forbidden", reason) +} + +func TestPreferBillingObservationStatus(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + weeklyStatus int + monthlyStatus int + want int + }{ + {name: "weekly forbidden wins", weeklyStatus: http.StatusForbidden, monthlyStatus: http.StatusBadGateway, want: http.StatusForbidden}, + {name: "monthly forbidden wins", weeklyStatus: http.StatusBadGateway, monthlyStatus: http.StatusForbidden, want: http.StatusForbidden}, + {name: "weekly observation otherwise wins", weeklyStatus: http.StatusTooManyRequests, monthlyStatus: http.StatusBadGateway, want: http.StatusTooManyRequests}, + {name: "monthly observation is fallback", monthlyStatus: http.StatusBadGateway, want: http.StatusBadGateway}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, preferBillingObservationStatus(tt.weeklyStatus, tt.monthlyStatus)) + }) + } +} + func TestGrokQuotaServiceQueryQuotaFree429PersistsLimitAndKeepsBilling(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index 25b563281..dfc182703 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -761,6 +761,71 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_AllowsG require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) } +func TestOpenAIGatewayService_SelectAccountWithScheduler_GrokMediaCapabilityFiltersIneligibleAccounts(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + + ctx := context.Background() + groupID := int64(10114) + ineligible := Account{ + ID: 36051, Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, + Extra: map[string]any{GrokMediaEligibleExtraKey: false}, + } + eligible := Account{ + ID: 36052, Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, + Extra: map[string]any{GrokMediaEligibleExtraKey: true}, + } + newService := func(accounts []Account) *OpenAIGatewayService { + cfg := &config.Config{} + cfg.Gateway.Scheduling.LoadBatchEnabled = false + return &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, + cache: &schedulerTestGatewayCache{}, + cfg: cfg, + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + } + + t.Run("media generation skips higher priority ineligible account", func(t *testing.T) { + selection, _, err := newService([]Account{ineligible, eligible}).SelectAccountWithSchedulerForCapability( + ctx, &groupID, "", "", "grok-imagine-video", nil, + OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityGrokMediaGeneration, + false, false, false, PlatformGrok, + ) + + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, eligible.ID, selection.Account.ID) + }) + + t.Run("media generation fails closed when all accounts are ineligible", func(t *testing.T) { + selection, _, err := newService([]Account{ineligible}).SelectAccountWithSchedulerForCapability( + ctx, &groupID, "", "", "grok-imagine-video", nil, + OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityGrokMediaGeneration, + false, false, false, PlatformGrok, + ) + + require.Error(t, err) + require.ErrorIs(t, err, ErrNoAvailableAccounts) + require.Nil(t, selection) + }) + + t.Run("chat remains routable on media-ineligible account", func(t *testing.T) { + selection, _, err := newService([]Account{ineligible}).SelectAccountWithSchedulerForCapability( + ctx, &groupID, "", "", "grok-4.3", nil, + OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityChatCompletions, + false, false, false, PlatformGrok, + ) + + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, ineligible.ID, selection.Account.ID) + }) +} + func TestOpenAIGatewayService_SelectAccountWithScheduler_EnabledUsesAdvancedPreviousResponseRouting(t *testing.T) { resetOpenAIAdvancedSchedulerSettingCacheForTest() diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index df574d46a..67ad71214 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -471,10 +471,12 @@ func TestCanonicalizeGrokMediaImageURLFieldsPreservesOfficialURL(t *testing.T) { "images":[ {"image_url":"https://example.com/first.png"}, {"url":"https://example.com/second.png"} - ] + ], + "reference_images":[{"image_url":"https://example.com/reference.png"}], + "mask":{"image_url":"https://example.com/mask.png"} }`) - out, err := canonicalizeGrokMediaImageURLFields(body, "image", "images") + out, err := canonicalizeGrokMediaImageURLFields(body, "image", "images", "reference_images", "mask") require.NoError(t, err) require.Equal(t, "https://example.com/official.png", gjson.GetBytes(out, "image.url").String()) @@ -482,6 +484,20 @@ func TestCanonicalizeGrokMediaImageURLFieldsPreservesOfficialURL(t *testing.T) { require.Equal(t, "https://example.com/first.png", gjson.GetBytes(out, "images.0.url").String()) require.False(t, gjson.GetBytes(out, "images.0.image_url").Exists()) require.Equal(t, "https://example.com/second.png", gjson.GetBytes(out, "images.1.url").String()) + require.Equal(t, "https://example.com/reference.png", gjson.GetBytes(out, "reference_images.0.url").String()) + require.False(t, gjson.GetBytes(out, "reference_images.0.image_url").Exists()) + require.Equal(t, "https://example.com/mask.png", gjson.GetBytes(out, "mask.url").String()) + require.False(t, gjson.GetBytes(out, "mask.image_url").Exists()) +} + +func TestCanonicalizeGrokMediaImageURLFieldsReplacesEmptyOfficialURL(t *testing.T) { + body := []byte(`{"image":{"url":" ","image_url":"https://example.com/legacy.png"}}`) + + out, err := canonicalizeGrokMediaImageURLFields(body, "image") + + require.NoError(t, err) + require.Equal(t, "https://example.com/legacy.png", gjson.GetBytes(out, "image.url").String()) + require.False(t, gjson.GetBytes(out, "image.image_url").Exists()) } func TestNormalizeGrokMediaModelForEndpoint(t *testing.T) { diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index 1ce86249b..63a40255b 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -6,7 +6,6 @@ package service import ( "context" "encoding/json" - "errors" "fmt" "log/slog" "sort" @@ -176,9 +175,21 @@ func noAvailableOpenAISelectionError(requestedModel string, compactBlocked bool) return ErrNoAvailableCompactAccounts } if requestedModel != "" { - return fmt.Errorf("no available OpenAI accounts supporting model: %s", requestedModel) + return openAINoAvailableSelectionError{message: fmt.Sprintf("no available OpenAI accounts supporting model: %s", requestedModel)} } - return errors.New("no available OpenAI accounts") + return openAINoAvailableSelectionError{message: "no available OpenAI accounts"} +} + +type openAINoAvailableSelectionError struct { + message string +} + +func (e openAINoAvailableSelectionError) Error() string { + return e.message +} + +func (e openAINoAvailableSelectionError) Unwrap() error { + return ErrNoAvailableAccounts } // openAICompactSupportTier classifies an OpenAI account by compact capability.