seewe/sub2api
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-16 07:33:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(grok): gate OAuth media on paid eligibility

Browse Source
This commit is contained in:
Heatherm Huang
2026-07-17 19:15:16 +08:00
parent 57914967cb
commit e86063155f
16 changed files with 317 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/handler/admin/account_handler.go
+1 -1
View File
@@ -61,7 +61,7 @@ type AccountHandler struct {
sessionLimitCache service.SessionLimitCache
rpmCache service.RPMCache
tokenCacheInvalidator service.TokenCacheInvalidator
grokImportProber grokUsageProber
grokImportProber grokImportProber
upstreamBillingProbe *service.UpstreamBillingProbeService
}
backend/internal/handler/admin/grok_import_probe.go
+6 -6
View File
@@ -15,12 +15,12 @@ const (
grokImportProbeTimeout = 25 * time.Second
)
type grokUsageProber interface {
ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
type grokImportProber interface {
QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
}
type grokImportProbeTask struct {
prober grokUsageProber
prober grokImportProber
accountID int64
}
@@ -51,7 +51,7 @@ func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokIm
}
}
func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) {
func (s *grokImportProbeScheduler) schedule(prober grokImportProber, account *service.Account) {
if s == nil || prober == nil || account == nil || account.ID <= 0 {
return
}
@@ -97,7 +97,7 @@ func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) {
return task, true
}
func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) {
func (s *grokImportProbeScheduler) run(prober grokImportProber, accountID int64) {
defer func() {
if recovered := recover(); recovered != nil {
slog.Error(
@@ -112,7 +112,7 @@ func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64)
// while this timeout only bounds the actual upstream probe execution.
ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
defer cancel()
result, err := prober.ProbeUsage(ctx, accountID)
result, err := prober.QueryQuota(ctx, accountID)
if err != nil {
slog.Warn(
"grok_import_active_probe_failed",
backend/internal/handler/admin/grok_import_probe_test.go
+2 -2
View File
@@ -36,7 +36,7 @@ func newGrokImportProbeStub(buffer int) *grokImportProbeStub {
}
}
func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
func (s *grokImportProbeStub) QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
_, deadlineSeen := ctx.Deadline()
s.mu.Lock()
s.calls[accountID]++
@@ -69,7 +69,7 @@ func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (
return nil, failure
}
return &service.GrokQuotaProbeResult{
Source: "active_probe",
Source: "hybrid_probe",
Model: "grok-4.5",
StatusCode: 200,
ResetSupported: false,
backend/internal/handler/admin/grok_oauth_handler.go
+1 -1
View File
@@ -23,7 +23,7 @@ type GrokOAuthHandler struct {
grokOAuthService *service.GrokOAuthService
adminService service.AdminService
quotaService *service.GrokQuotaService
importProber grokUsageProber
importProber grokImportProber
reconciler service.GrokOAuthReconciler
}
backend/internal/handler/grok_media.go
+36 -1
View File
@@ -167,6 +167,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
sameAccountRetryCount := make(map[int64]int)
var lastFailoverErr *service.UpstreamFailoverError
var oauth429FailoverState service.OpenAIOAuth429FailoverState
mediaEligibilityRejected := false
switchCount := 0
maxAccountSwitches := h.maxAccountSwitches
if maxAccountSwitches <= 0 {
@@ -202,7 +203,8 @@ 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) {
if endpoint.IsGenerationRequest() && errors.Is(err, service.ErrNoAvailableAccounts) &&
(len(failedAccountIDs) == 0 || (mediaEligibilityRejected && lastFailoverErr == nil)) {
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
return
@@ -246,6 +248,25 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
)
account := selection.Account
if endpoint.IsGenerationRequest() {
eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account)
if !eligible {
mediaEligibilityRejected = true
failedAccountIDs[account.ID] = struct{}{}
reqLog.Warn("grok_media.account_eligibility_rejected",
zap.Int64("account_id", account.ID),
zap.String("reason", eligibilityReason),
zap.Bool("probe_failed", eligibilityErr != nil),
)
if switchCount >= maxAccountSwitches {
markOpsRoutingCapacityLimited(c)
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
return
}
switchCount++
continue
}
}
sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account)
setOpsSelectedAccount(c, account.ID, account.Platform)
@@ -365,6 +386,20 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
}
}
func (h *OpenAIGatewayHandler) ensureGrokMediaAccountEligibility(ctx context.Context, account *service.Account) (bool, string, error) {
if account == nil {
return false, "missing_account", errors.New("grok media account is required")
}
eligible, reason := account.GrokMediaGenerationEligibility()
if eligible || reason != "billing_unobserved" {
return eligible, reason, nil
}
if h == nil || h.grokMediaEligibilityProber == nil {
return false, "billing_probe_unavailable", errors.New("grok media eligibility probe is not configured")
}
return h.grokMediaEligibilityProber.ProbeMediaEligibility(ctx, account.ID)
}
func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.OpenAIEndpointCapability {
if endpoint.IsGenerationRequest() {
return service.OpenAIEndpointCapabilityGrokMediaGeneration
backend/internal/handler/grok_media_test.go
+66
View File
@@ -1,12 +1,26 @@
package handler
import (
"context"
"errors"
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
type grokMediaEligibilityProberStub struct {
eligible bool
reason string
err error
calls int
}
func (s *grokMediaEligibilityProberStub) ProbeMediaEligibility(context.Context, int64) (bool, string, error) {
s.calls++
return s.eligible, s.reason, s.err
}
func TestShouldRecordGrokMediaUsage(t *testing.T) {
tests := []struct {
name string
@@ -73,3 +87,55 @@ func TestGrokMediaRequiredCapability(t *testing.T) {
})
}
}
func TestEnsureGrokMediaAccountEligibility(t *testing.T) {
t.Run("non oauth account does not probe", func(t *testing.T) {
prober := &grokMediaEligibilityProberStub{}
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
account := &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey}
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
require.NoError(t, err)
require.True(t, eligible)
require.Equal(t, "non_oauth", reason)
require.Zero(t, prober.calls)
})
t.Run("unobserved oauth is probed before forwarding", func(t *testing.T) {
prober := &grokMediaEligibilityProberStub{eligible: true, reason: "eligible"}
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
account := &service.Account{ID: 7, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
require.NoError(t, err)
require.True(t, eligible)
require.Equal(t, "eligible", reason)
require.Equal(t, 1, prober.calls)
})
t.Run("missing prober fails closed", func(t *testing.T) {
h := &OpenAIGatewayHandler{}
account := &service.Account{ID: 8, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
require.Error(t, err)
require.False(t, eligible)
require.Equal(t, "billing_probe_unavailable", reason)
})
t.Run("probe failure fails closed", func(t *testing.T) {
probeErr := errors.New("probe failed")
prober := &grokMediaEligibilityProberStub{reason: "billing_unobserved", err: probeErr}
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
account := &service.Account{ID: 9, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
require.ErrorIs(t, err, probeErr)
require.False(t, eligible)
require.Equal(t, "billing_unobserved", reason)
})
}
backend/internal/handler/openai_gateway_handler.go
+17 -12
View File
@@ -28,18 +28,23 @@ import (
// OpenAIGatewayHandler handles OpenAI API gateway requests
type OpenAIGatewayHandler struct {
gatewayService *service.OpenAIGatewayService
billingCacheService *service.BillingCacheService
apiKeyService *service.APIKeyService
usageRecordWorkerPool *service.UsageRecordWorkerPool
errorPassthroughService *service.ErrorPassthroughService
contentModerationService *service.ContentModerationService
securityAuditCoordinator *securityaudit.Coordinator
opsService *service.OpsService
concurrencyHelper *ConcurrencyHelper
imageLimiter *imageConcurrencyLimiter
maxAccountSwitches int
cfg *config.Config
gatewayService *service.OpenAIGatewayService
billingCacheService *service.BillingCacheService
apiKeyService *service.APIKeyService
usageRecordWorkerPool *service.UsageRecordWorkerPool
errorPassthroughService *service.ErrorPassthroughService
contentModerationService *service.ContentModerationService
securityAuditCoordinator *securityaudit.Coordinator
grokMediaEligibilityProber grokMediaEligibilityProber
opsService *service.OpsService
concurrencyHelper *ConcurrencyHelper
imageLimiter *imageConcurrencyLimiter
maxAccountSwitches int
cfg *config.Config
}
type grokMediaEligibilityProber interface {
ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error)
}
const maxOpenAIFirstOutputTimeoutSwitches = 1
backend/internal/handler/wire.go
+2
View File
@@ -120,12 +120,14 @@ func ProvideOpenAIGatewayHandler(
errorPassthroughService *service.ErrorPassthroughService,
contentModerationService *service.ContentModerationService,
opsService *service.OpsService,
grokQuotaService *service.GrokQuotaService,
cfg *config.Config,
coordinator *securityaudit.Coordinator,
) *OpenAIGatewayHandler {
h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService,
usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg)
h.securityAuditCoordinator = coordinator
h.grokMediaEligibilityProber = grokQuotaService
return h
}