fix(grok): harden media quarantine edge cases

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)
+	})
 }

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,
 		} {

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{

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 {

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()
 

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()
 

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) {

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.