fix: address grok review feedback

backend/cmd/server/wire_gen.go

@@ -199,7 +199,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
 	openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService, openAIQuotaService)
 	geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService)
 	antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService)
-	grokQuotaService := service.ProvideGrokQuotaService(accountRepository, grokTokenProvider, httpUpstream)
+	grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream)
 	grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService)
 	proxyHandler := admin.NewProxyHandler(adminService)
 	adminRedeemHandler := admin.NewRedeemHandler(adminService, redeemService)

backend/internal/handler/admin/grok_oauth_handler_test.go

@@ -85,7 +85,7 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
 		},
 		Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
 	}}
-	quotaService := service.NewGrokQuotaService(repo, service.NewGrokTokenProvider(repo, nil, nil), upstream)
+	quotaService := service.NewGrokQuotaService(repo, nil, service.NewGrokTokenProvider(repo, nil), upstream)
 	handler := NewGrokOAuthHandler(nil, nil, quotaService)
 
 	router := gin.New()
@@ -112,7 +112,7 @@ func TestGrokOAuthHandlerResetQuotaReturnsUnsupported(t *testing.T) {
 		Platform: service.PlatformGrok,
 		Type:     service.AccountTypeOAuth,
 	}}
-	quotaService := service.NewGrokQuotaService(repo, nil, nil)
+	quotaService := service.NewGrokQuotaService(repo, nil, nil, nil)
 	handler := NewGrokOAuthHandler(nil, nil, quotaService)
 
 	router := gin.New()

backend/internal/service/account_usage_service.go

@@ -864,7 +864,8 @@ func (s *AccountUsageService) getAntigravityUsage(ctx context.Context, account *
 
 func (s *AccountUsageService) getGrokUsage(ctx context.Context, account *Account) (*UsageInfo, error) {
 	if s.grokQuotaFetcher == nil {
-		s.grokQuotaFetcher = NewGrokQuotaFetcher()
+		now := time.Now()
+		return &UsageInfo{UpdatedAt: &now}, nil
 	}
 	usage := s.grokQuotaFetcher.BuildUsageInfo(account)
 	if usage.GrokQuotaSnapshotState == "" {

backend/internal/service/admin_account_concurrency_test.go

@@ -0,0 +1,27 @@
+//go:build unit
+
+package service
+
+import (
+	"testing"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNormalizeAccountConcurrencyCapsGrokOAuthUnlessUnsafe(t *testing.T) {
+	t.Setenv(xai.EnvUnsafeAllowHighConcurrency, "")
+
+	require.Equal(t, 1, normalizeAccountConcurrency(PlatformGrok, AccountTypeOAuth, 0))
+	require.Equal(t, 1, normalizeAccountConcurrency(PlatformGrok, AccountTypeOAuth, -5))
+	require.Equal(t, 1, normalizeAccountConcurrency(PlatformGrok, AccountTypeOAuth, 50))
+	require.Equal(t, 2, normalizeAccountConcurrency(PlatformOpenAI, AccountTypeOAuth, 2))
+	require.Equal(t, 2, normalizeAccountConcurrency(PlatformGrok, AccountTypeAPIKey, 2))
+}
+
+func TestNormalizeAccountConcurrencyAllowsGrokOAuthUnsafeOverride(t *testing.T) {
+	t.Setenv(xai.EnvUnsafeAllowHighConcurrency, "true")
+
+	require.Equal(t, 50, normalizeAccountConcurrency(PlatformGrok, AccountTypeOAuth, 50))
+	require.Equal(t, 1, normalizeAccountConcurrency(PlatformGrok, AccountTypeOAuth, 0))
+}

backend/internal/service/admin_service.go

@@ -2573,8 +2573,13 @@ func (s *adminServiceImpl) GetAccountsByIDs(ctx context.Context, ids []int64) ([
 }
 
 func normalizeAccountConcurrency(platform, accountType string, concurrency int) int {
-	if platform == PlatformGrok && accountType == AccountTypeOAuth && concurrency <= 0 {
-		return 1
+	if platform == PlatformGrok && accountType == AccountTypeOAuth {
+		if concurrency <= 0 {
+			return 1
+		}
+		if concurrency > 1 && !xai.AllowUnsafeHighConcurrency() {
+			return 1
+		}
 	}
 	return concurrency
 }

backend/internal/service/grok_quota_service.go

@@ -37,17 +37,20 @@ type GrokQuotaResetResult struct {
 
 type GrokQuotaService struct {
 	accountRepo   AccountRepository
+	proxyRepo     ProxyRepository
 	tokenProvider *GrokTokenProvider
 	httpUpstream  HTTPUpstream
 }
 
 func NewGrokQuotaService(
 	accountRepo AccountRepository,
+	proxyRepo ProxyRepository,
 	tokenProvider *GrokTokenProvider,
 	httpUpstream HTTPUpstream,
 ) *GrokQuotaService {
 	return &GrokQuotaService{
 		accountRepo:   accountRepo,
+		proxyRepo:     proxyRepo,
 		tokenProvider: tokenProvider,
 		httpUpstream:  httpUpstream,
 	}
@@ -134,11 +137,22 @@ func (s *GrokQuotaService) prepareProbe(ctx context.Context, accountID int64) (*
 		return nil, "", "", infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_TOKEN_UNAVAILABLE", "access token is empty")
 	}
 
-	proxyURL := ""
-	if account.ProxyID != nil && account.Proxy != nil {
-		proxyURL = account.Proxy.URL()
+	return account, token, s.resolveProxyURL(ctx, account), nil
+}
+
+func (s *GrokQuotaService) resolveProxyURL(ctx context.Context, account *Account) string {
+	if account == nil || account.ProxyID == nil {
+		return ""
 	}
-	return account, token, proxyURL, nil
+	switch {
+	case account.Proxy != nil:
+		return account.Proxy.URL()
+	case s != nil && s.proxyRepo != nil:
+		if proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && proxy != nil {
+			return proxy.URL()
+		}
+	}
+	return ""
 }
 
 func (s *GrokQuotaService) loadGrokOAuthAccount(ctx context.Context, accountID int64) (*Account, error) {

backend/internal/service/grok_quota_service_test.go

@@ -40,6 +40,17 @@ func (r *grokQuotaAccountRepo) SetTempUnschedulable(_ context.Context, id int64,
 	return nil
 }
 
+type grokQuotaProxyRepo struct {
+	proxyRepoStub
+	proxies map[int64]*Proxy
+	calls   int
+}
+
+func (r *grokQuotaProxyRepo) GetByID(_ context.Context, id int64) (*Proxy, error) {
+	r.calls++
+	return r.proxies[id], nil
+}
+
 func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
 	t.Parallel()
 
@@ -69,7 +80,7 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
 		},
 		Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
 	}}
-	svc := NewGrokQuotaService(repo, NewGrokTokenProvider(repo, nil, nil), upstream)
+	svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
 
 	result, err := svc.ProbeUsage(context.Background(), 42)
 	require.NoError(t, err)
@@ -90,6 +101,49 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
 	require.NotNil(t, repo.updates[42][grokQuotaSnapshotExtraKey])
 }
 
+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),
+		},
+	}
+	repo := &grokQuotaAccountRepo{
+		mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
+			accountsByID: map[int64]*Account{46: account},
+		},
+	}
+	proxyRepo := &grokQuotaProxyRepo{
+		proxies: map[int64]*Proxy{
+			proxyID: {
+				ID:       proxyID,
+				Protocol: "http",
+				Host:     "proxy.test",
+				Port:     3128,
+			},
+		},
+	}
+	upstream := &httpUpstreamRecorder{resp: &http.Response{
+		StatusCode: http.StatusOK,
+		Header:     http.Header{},
+		Body:       io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
+	}}
+	svc := NewGrokQuotaService(repo, proxyRepo, NewGrokTokenProvider(repo, nil), upstream)
+
+	_, err := svc.ProbeUsage(context.Background(), 46)
+	require.NoError(t, err)
+	require.Equal(t, 1, proxyRepo.calls)
+	require.Equal(t, "http://proxy.test:3128", upstream.lastProxyURL)
+}
+
 func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) {
 	t.Parallel()
 
@@ -113,7 +167,7 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) {
 		Header:     http.Header{},
 		Body:       io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
 	}}
-	svc := NewGrokQuotaService(repo, NewGrokTokenProvider(repo, nil, nil), upstream)
+	svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
 
 	result, err := svc.ProbeUsage(context.Background(), 45)
 	require.NoError(t, err)
@@ -153,7 +207,7 @@ func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) {
 		Header:     http.Header{"Retry-After": []string{"45"}},
 		Body:       io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
 	}}
-	svc := NewGrokQuotaService(repo, NewGrokTokenProvider(repo, nil, nil), upstream)
+	svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
 
 	result, err := svc.ProbeUsage(context.Background(), 43)
 	require.NoError(t, err)
@@ -176,7 +230,7 @@ func TestGrokQuotaServiceResetQuotaUnsupported(t *testing.T) {
 			accountsByID: map[int64]*Account{44: account},
 		},
 	}
-	svc := NewGrokQuotaService(repo, nil, nil)
+	svc := NewGrokQuotaService(repo, nil, nil, nil)
 
 	_, err := svc.ResetQuota(context.Background(), 44)
 	require.Error(t, err)

backend/internal/service/grok_token_provider.go

@@ -23,7 +23,6 @@ type GrokTokenCache = GeminiTokenCache
 type GrokTokenProvider struct {
 	accountRepo      AccountRepository
 	tokenCache       GrokTokenCache
-	grokOAuthService *GrokOAuthService
 	refreshAPI       *OAuthRefreshAPI
 	executor         OAuthRefreshExecutor
 	refreshPolicy    ProviderRefreshPolicy
@@ -33,13 +32,11 @@ type GrokTokenProvider struct {
 func NewGrokTokenProvider(
 	accountRepo AccountRepository,
 	tokenCache GrokTokenCache,
-	grokOAuthService *GrokOAuthService,
 ) *GrokTokenProvider {
 	return &GrokTokenProvider{
-		accountRepo:      accountRepo,
-		tokenCache:       tokenCache,
-		grokOAuthService: grokOAuthService,
-		refreshPolicy:    AntigravityProviderRefreshPolicy(),
+		accountRepo:   accountRepo,
+		tokenCache:    tokenCache,
+		refreshPolicy: AntigravityProviderRefreshPolicy(),
 	}
 }
 

backend/internal/service/grok_token_provider_test.go

@@ -76,7 +76,7 @@ func TestGrokTokenProviderRefreshesExpiredTokenOnRequestPath(t *testing.T) {
 	})
 	defer oauthSvc.Stop()
 
-	provider := NewGrokTokenProvider(repo, cache, oauthSvc)
+	provider := NewGrokTokenProvider(repo, cache)
 	provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc))
 
 	token, err := provider.GetAccessToken(context.Background(), account)
@@ -109,7 +109,7 @@ func TestGrokTokenProviderRefreshFailureUnschedulesWithRedactedReason(t *testing
 	repo.accountsByID = map[int64]*Account{55: account}
 	cache := &grokTokenCacheForProviderTest{lockResult: true}
 	tempCache := &tempUnschedCacheStub{}
-	provider := NewGrokTokenProvider(repo, cache, nil)
+	provider := NewGrokTokenProvider(repo, cache)
 	provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
 		err: errors.New("temporary refresh failure access_token=leaked-access refresh_token=leaked-refresh"),
 	})

backend/internal/service/grok_token_refresher.go

@@ -10,10 +10,10 @@ import (
 const grokTokenRefreshSkew = time.Hour
 
 type GrokTokenRefresher struct {
-	grokOAuthService *GrokOAuthService
+	grokOAuthService GrokOAuthTokenService
 }
 
-func NewGrokTokenRefresher(grokOAuthService *GrokOAuthService) *GrokTokenRefresher {
+func NewGrokTokenRefresher(grokOAuthService GrokOAuthTokenService) *GrokTokenRefresher {
 	return &GrokTokenRefresher{grokOAuthService: grokOAuthService}
 }
 

backend/internal/service/oauth_service.go

@@ -24,6 +24,12 @@ type GrokOAuthClient interface {
 	RefreshToken(ctx context.Context, refreshToken, proxyURL, clientID string) (*xai.TokenResponse, error)
 }
 
+// GrokOAuthTokenService is the narrow refresh port used by Grok token providers.
+type GrokOAuthTokenService interface {
+	RefreshAccountToken(ctx context.Context, account *Account) (*GrokTokenInfo, error)
+	BuildAccountCredentials(tokenInfo *GrokTokenInfo) map[string]any
+}
+
 // ClaudeOAuthClient handles HTTP requests for Claude OAuth flows
 type ClaudeOAuthClient interface {
 	GetOrganizationUUID(ctx context.Context, sessionKey, proxyURL string) (string, error)

backend/internal/service/openai_gateway_grok_test.go

@@ -118,7 +118,7 @@ func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *te
 	}}
 	svc := &OpenAIGatewayService{
 		httpUpstream:      upstream,
-		grokTokenProvider: NewGrokTokenProvider(repo, nil, nil),
+		grokTokenProvider: NewGrokTokenProvider(repo, nil),
 		accountRepo:       repo,
 	}
 
@@ -182,7 +182,7 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T)
 	}}
 	svc := &OpenAIGatewayService{
 		httpUpstream:      upstream,
-		grokTokenProvider: NewGrokTokenProvider(repo, nil, nil),
+		grokTokenProvider: NewGrokTokenProvider(repo, nil),
 		accountRepo:       repo,
 	}
 
@@ -251,7 +251,7 @@ func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *te
 	svc := &OpenAIGatewayService{
 		cfg:               rawChatCompletionsTestConfig(),
 		httpUpstream:      upstream,
-		grokTokenProvider: NewGrokTokenProvider(repo, nil, nil),
+		grokTokenProvider: NewGrokTokenProvider(repo, nil),
 		accountRepo:       repo,
 	}
 

backend/internal/service/openai_oauth_passthrough_test.go

@@ -24,10 +24,11 @@ import (
 func f64p(v float64) *float64 { return &v }
 
 type httpUpstreamRecorder struct {
-	lastReq  *http.Request
-	lastBody []byte
-	requests []*http.Request
-	bodies   [][]byte
+	lastReq      *http.Request
+	lastBody     []byte
+	lastProxyURL string
+	requests     []*http.Request
+	bodies       [][]byte
 
 	resp      *http.Response
 	responses []*http.Response
@@ -36,6 +37,7 @@ type httpUpstreamRecorder struct {
 
 func (u *httpUpstreamRecorder) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
 	u.lastReq = req
+	u.lastProxyURL = proxyURL
 	if req != nil && req.Body != nil {
 		b, _ := io.ReadAll(req.Body)
 		u.lastBody = b

backend/internal/service/wire.go

@@ -127,10 +127,11 @@ func ProvideOpenAIQuotaService(
 
 func ProvideGrokQuotaService(
 	accountRepo AccountRepository,
+	proxyRepo ProxyRepository,
 	tokenProvider *GrokTokenProvider,
 	httpUpstream HTTPUpstream,
 ) *GrokQuotaService {
-	return NewGrokQuotaService(accountRepo, tokenProvider, httpUpstream)
+	return NewGrokQuotaService(accountRepo, proxyRepo, tokenProvider, httpUpstream)
 }
 
 // ProvideGeminiTokenProvider creates GeminiTokenProvider with OAuthRefreshAPI injection
@@ -171,7 +172,7 @@ func ProvideGrokTokenProvider(
 	refreshAPI *OAuthRefreshAPI,
 	tempUnschedCache TempUnschedCache,
 ) *GrokTokenProvider {
-	p := NewGrokTokenProvider(accountRepo, tokenCache, grokOAuthService)
+	p := NewGrokTokenProvider(accountRepo, tokenCache)
 	executor := NewGrokTokenRefresher(grokOAuthService)
 	p.SetRefreshAPI(refreshAPI, executor)
 	p.SetRefreshPolicy(AntigravityProviderRefreshPolicy())