diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index d3fbe9c847..e5ab1dd04c 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -198,30 +198,57 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. } routingStart := time.Now() requiredCapability := grokMediaRequiredCapability(endpoint) + var accountReleaseFunc func() + releaseAccount := func() { + if accountReleaseFunc != nil { + accountReleaseFunc() + accountReleaseFunc = nil + } + } + defer releaseAccount() for { + releaseAccount() if failoverClientGone(c) { return } - selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability( - requestCtx, - apiKey.GroupID, - "", - sessionHash, - routingModel, - failedAccountIDs, - service.OpenAIUpstreamTransportHTTPSSE, - requiredCapability, - false, - false, - false, - service.PlatformGrok, - ) + var selection *service.AccountSelectionResult + var scheduleDecision service.OpenAIAccountScheduleDecision + if boundLookupAccountID > 0 { + selection, scheduleDecision, err = h.gatewayService.SelectGrokMediaVideoRequestAccount( + requestCtx, apiKey.GroupID, sessionHash, boundLookupAccountID, routingModel, + ) + } else { + selection, scheduleDecision, err = h.gatewayService.SelectAccountWithSchedulerForCapability( + requestCtx, + apiKey.GroupID, + "", + sessionHash, + routingModel, + failedAccountIDs, + service.OpenAIUpstreamTransportHTTPSSE, + requiredCapability, + false, + false, + false, + service.PlatformGrok, + ) + } + // Own an eagerly acquired slot before any rejection or eligibility probe. + // Forwarding takes over the same once-only release after admission. + if selection != nil && selection.Acquired { + selection.ReleaseFunc = wrapReleaseOnDone(requestCtx, selection.ReleaseFunc) + accountReleaseFunc = selection.ReleaseFunc + } if err != nil { if failoverClientGone(c) { reqLog.Info("grok_media.account_select_aborted_client_disconnected", zap.Error(err)) return } + if boundLookupAccountID > 0 && errors.Is(err, service.ErrNoAvailableAccounts) { + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found") + return + } reqLog.Warn("grok_media.account_select_failed", zap.Error(err), zap.Int("excluded_account_count", len(failedAccountIDs)), @@ -260,6 +287,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. h.errorResponse(c, cls.Status, cls.ErrType, cls.Message) return } + if failoverClientGone(c) { + return + } if boundLookupAccountID > 0 && selection.Account.ID != boundLookupAccountID { reqLog.Warn("grok_media.video_lookup_bound_account_unavailable", zap.Int64("bound_account_id", boundLookupAccountID), @@ -282,6 +312,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. if endpoint.IsGenerationRequest() { eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account) if !eligible { + releaseAccount() mediaEligibilityRejected = true failedAccountIDs[account.ID] = struct{}{} reqLog.Warn("grok_media.account_eligibility_rejected", @@ -298,10 +329,19 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. continue } } + if failoverClientGone(c) { + return + } sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account) setOpsSelectedAccount(c, account.ID, account.Platform) - accountReleaseFunc, slotResult := h.acquireResponsesAccountSlot(c, apiKey.GroupID, sessionHash, selection, false, &streamStarted, reqLog) + admissionSessionHash := sessionHash + if boundLookupAccountID > 0 { + // Waiting must not replace the video ownership TTL with text sticky TTL. + admissionSessionHash = "" + } + var slotResult openAISlotAcquireResult + accountReleaseFunc, slotResult = h.acquireResponsesAccountSlot(c, apiKey.GroupID, admissionSessionHash, selection, false, &streamStarted, reqLog) if slotResult == openAISlotAcquireProfitVetoed { // 媒体路径已显式豁免利润门（suppress 标记），此分支仅防御性兜底， // 同样受否决上限约束。 @@ -319,11 +359,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. forwardStart := time.Now() writerSizeBeforeForward := c.Writer.Size() result, err := func() (*service.OpenAIForwardResult, error) { - defer func() { - if accountReleaseFunc != nil { - accountReleaseFunc() - } - }() + defer releaseAccount() return h.gatewayService.ForwardGrokMedia(requestCtx, c, account, endpoint, requestID, body, contentType) }() diff --git a/backend/internal/handler/grok_media_slots_test.go b/backend/internal/handler/grok_media_slots_test.go new file mode 100644 index 0000000000..3393e18344 --- /dev/null +++ b/backend/internal/handler/grok_media_slots_test.go @@ -0,0 +1,405 @@ +//go:build unit + +package handler + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/Wei-Shaw/sub2api/internal/testutil" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +type grokMediaSlotsCache struct { + testutil.StubConcurrencyCache + mu sync.Mutex + accounts map[string]int64 + users map[string]int64 + acquired, released, userAcquired, userReleased int + waiting, denied int + full bool + queueFull bool + onWait func() + duplicateRelease bool + peak int +} + +func (s *grokMediaSlotsCache) AcquireAccountSlot(_ context.Context, id int64, _ int, request string) (bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.full || s.denied > 0 { + if s.denied > 0 { + s.denied-- + } + return false, nil + } + s.accounts[request] = id + s.acquired++ + if len(s.accounts) > s.peak { + s.peak = len(s.accounts) + } + return true, nil +} + +func (s *grokMediaSlotsCache) ReleaseAccountSlot(_ context.Context, _ int64, request string) error { + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.accounts[request]; !ok { + s.duplicateRelease = true + } + delete(s.accounts, request) + s.released++ + return nil +} + +func (s *grokMediaSlotsCache) AcquireUserSlot(_ context.Context, id int64, _ int, request string) (bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.users[request] = id + s.userAcquired++ + return true, nil +} + +func (s *grokMediaSlotsCache) ReleaseUserSlot(_ context.Context, _ int64, request string) error { + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.users[request]; !ok { + s.duplicateRelease = true + } + delete(s.users, request) + s.userReleased++ + return nil +} + +func (s *grokMediaSlotsCache) IncrementAccountWaitCount(context.Context, int64, int) (bool, error) { + s.mu.Lock() + if s.queueFull { + s.mu.Unlock() + return false, nil + } + s.waiting++ + onWait := s.onWait + s.mu.Unlock() + if onWait != nil { + onWait() + } + return true, nil +} + +func (s *grokMediaSlotsCache) DecrementAccountWaitCount(context.Context, int64) error { + s.mu.Lock() + defer s.mu.Unlock() + s.waiting-- + return nil +} + +func (s *grokMediaSlotsCache) assertReleased(t *testing.T) { + t.Helper() + s.mu.Lock() + defer s.mu.Unlock() + require.Empty(t, s.accounts) + require.Empty(t, s.users) + require.Equal(t, s.acquired, s.released) + require.Equal(t, s.userAcquired, s.userReleased) + require.False(t, s.duplicateRelease) + require.Zero(t, s.waiting) +} + +type grokMediaSlotBindings struct { + testutil.StubGatewayCache + owner int64 + writes int + key string + billed map[string]bool +} + +func (s *grokMediaSlotBindings) GetSessionAccountID(_ context.Context, groupID int64, key string) (int64, error) { + if groupID != 24 || key != s.key { + return 0, service.ErrStickySessionNotFound + } + return s.owner, nil +} +func (s *grokMediaSlotBindings) SetSessionAccountID(_ context.Context, _ int64, key string, owner int64, _ time.Duration) error { + s.key, s.owner = key, owner + s.writes++ + return nil +} +func (s *grokMediaSlotBindings) RefreshSessionTTL(context.Context, int64, string, time.Duration) error { + s.writes++ + return nil +} +func (s *grokMediaSlotBindings) DeleteSessionAccountID(context.Context, int64, string) error { + s.writes++ + return nil +} + +func (s *grokMediaSlotBindings) ClaimGrokVideoBilled(_ context.Context, key string, _ time.Duration) (bool, error) { + if s.billed == nil { + s.billed = make(map[string]bool) + } + if s.billed[key] { + return false, nil + } + s.billed[key] = true + return true, nil +} + +type grokMediaSlotRepo struct { + openAIImagesFailoverAccountRepo + mismatch bool +} + +func (r grokMediaSlotRepo) GetByID(ctx context.Context, id int64) (*service.Account, error) { + if r.mismatch { + return r.openAIImagesFailoverAccountRepo.GetByID(ctx, 2) + } + return r.openAIImagesFailoverAccountRepo.GetByID(ctx, id) +} + +type grokMediaSlotUpstream struct { + service.HTTPUpstream + call func(*http.Request, int64) (*http.Response, error) + calls int +} + +func (u *grokMediaSlotUpstream) Do(req *http.Request, _ string, id int64, _ int) (*http.Response, error) { + u.calls++ + return u.call(req, id) +} + +type grokMediaSlotProber func(context.Context, int64) (bool, string, error) + +func (p grokMediaSlotProber) ProbeMediaEligibility(ctx context.Context, id int64) (bool, string, error) { + return p(ctx, id) +} + +func newGrokMediaSlotHandler(t *testing.T, oauth, mismatch bool) (*OpenAIGatewayHandler, *grokMediaSlotsCache, *grokMediaSlotBindings, *grokMediaSlotUpstream) { + t.Helper() + gin.SetMode(gin.TestMode) + accounts := make([]service.Account, 3) + for i := range accounts { + accounts[i] = service.Account{ID: int64(i + 1), Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey, + Status: service.StatusActive, Schedulable: true, Concurrency: 50, Priority: i, + GroupIDs: []int64{24}, Credentials: map[string]any{"api_key": "test-key", "access_token": "test-token"}} + if oauth { + accounts[i].Type = service.AccountTypeOAuth + accounts[i].Credentials["refresh_token"] = "test-refresh" + accounts[i].Credentials["expires_at"] = time.Now().Add(7 * 24 * time.Hour).UTC().Format(time.RFC3339) + } + } + slots := &grokMediaSlotsCache{accounts: map[string]int64{}, users: map[string]int64{}} + concurrency := service.NewConcurrencyService(slots) + bindings := &grokMediaSlotBindings{owner: 1} + upstream := &grokMediaSlotUpstream{call: func(*http.Request, int64) (*http.Response, error) { + return &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"request_id":"task","status":"pending"}`))}, nil + }} + cfg := &config.Config{RunMode: config.RunModeSimple} + cfg.Gateway.Scheduling.StickySessionWaitTimeout = 20 * time.Millisecond + cfg.Gateway.Scheduling.StickySessionMaxWaiting = 3 + cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true + repo := grokMediaSlotRepo{openAIImagesFailoverAccountRepo: openAIImagesFailoverAccountRepo{accounts: accounts}, mismatch: mismatch} + provider := service.NewGrokTokenProvider(repo, nil) + if oauth { + _, err := provider.GetAccessToken(context.Background(), &accounts[1]) + require.NoError(t, err) + } + gateway := service.NewOpenAIGatewayService(repo, nil, nil, nil, nil, nil, bindings, cfg, nil, concurrency, nil, nil, nil, upstream, nil, nil, provider, nil, nil, nil, nil, nil) + groupID := int64(24) + require.NoError(t, gateway.BindGrokMediaVideoRequestAccount(context.Background(), &groupID, "task", 10, 20, 1)) + bindings.writes = 0 + billing := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billing.Stop) + handler := NewOpenAIGatewayHandler(gateway, concurrency, billing, service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg), nil, nil, nil, nil, cfg) + return handler, slots, bindings, upstream +} + +func grokMediaSlotContext(ctx context.Context, generation bool) (*gin.Context, *httptest.ResponseRecorder) { + groupID := int64(24) + method, path, body := http.MethodGet, "/v1/videos/task", "" + if generation { + method, path, body = http.MethodPost, "/v1/videos/generations", `{"model":"grok-imagine-video","prompt":"test","duration":6}` + } + req := httptest.NewRequest(method, path, strings.NewReader(body)).WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + if !generation { + req.Header.Set("session-id", "unrelated-client-session") + } + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = req + c.Params = gin.Params{{Key: "request_id", Value: "task"}} + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ID: 20, UserID: 10, GroupID: &groupID, + Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, AllowImageGeneration: true}, User: &service.User{ID: 10}}) + c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 10, Concurrency: 5}) + return c, w +} + +func TestGrokMediaLookupSlotLifecycle(t *testing.T) { + for _, scenario := range []string{"normal", "mismatch acquired", "mismatch wait", "full", "queue full", "cancel while waiting", "wait then acquired", "upstream error", "cancel", "panic"} { + t.Run(scenario, func(t *testing.T) { + h, slots, bindings, upstream := newGrokMediaSlotHandler(t, false, strings.HasPrefix(scenario, "mismatch")) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if scenario == "full" || scenario == "mismatch wait" || scenario == "queue full" || scenario == "cancel while waiting" { + slots.full = true + } + if scenario == "queue full" { + slots.queueFull = true + } + if scenario == "cancel while waiting" { + slots.onWait = cancel + } + if scenario == "wait then acquired" { + slots.denied = 1 + } + original := upstream.call + upstream.call = func(req *http.Request, id int64) (*http.Response, error) { + require.Equal(t, int64(1), id) + switch scenario { + case "upstream error": + return nil, errors.New("upstream unavailable") + case "cancel": + cancel() + return nil, context.Canceled + case "panic": + panic("test upstream panic") + } + return original(req, id) + } + for range 20 { + c, w := grokMediaSlotContext(ctx, false) + h.GrokVideoStatus(c) + slots.assertReleased(t) + if strings.HasPrefix(scenario, "mismatch") { + require.Equal(t, 404, w.Code) + } + if scenario == "normal" || scenario == "wait then acquired" { + require.Equal(t, 200, w.Code) + } + if scenario == "full" || scenario == "queue full" { + require.Equal(t, http.StatusTooManyRequests, w.Code) + } + } + require.Zero(t, bindings.writes, "lookups must preserve owner and TTL") + if scenario == "mismatch acquired" { + require.Equal(t, 20, slots.released) + } + if slots.full { + require.Zero(t, slots.released) + require.Zero(t, upstream.calls) + } + if scenario == "full" || strings.HasPrefix(scenario, "mismatch") { + require.Zero(t, upstream.calls) + } + switch scenario { + case "normal", "wait then acquired", "upstream error", "panic": + require.Equal(t, 20, upstream.calls) + require.Equal(t, 20, slots.released) + case "cancel": + require.Equal(t, 1, upstream.calls) + require.Equal(t, 1, slots.released) + } + }) + } +} + +func TestGrokMediaEligibilityReleasesBeforeSwitch(t *testing.T) { + for _, scenario := range []string{"switch", "exhausted", "cancel during probe"} { + t.Run(scenario, func(t *testing.T) { + h, slots, _, upstream := newGrokMediaSlotHandler(t, true, false) + h.maxAccountSwitches = 1 + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + probes := 0 + h.grokMediaEligibilityProber = grokMediaSlotProber(func(context.Context, int64) (bool, string, error) { + probes++ + if scenario == "cancel during probe" { + cancel() + require.Eventually(t, func() bool { slots.mu.Lock(); defer slots.mu.Unlock(); return len(slots.accounts) == 0 }, time.Second, time.Millisecond) + } + if scenario == "switch" && probes == 2 { + return true, "eligible", nil + } + return false, "billing_unobserved", errors.New("probe failed") + }) + c, w := grokMediaSlotContext(ctx, true) + h.GrokVideoGeneration(c) + slots.assertReleased(t) + require.Equal(t, 1, slots.peak, "rejected attempts must release before the next selection") + if scenario == "switch" { + events, _ := c.Get(service.OpsUpstreamErrorsKey) + detail, err := json.Marshal(events) + require.NoError(t, err) + require.Equal(t, 200, w.Code, "body=%s probes=%d upstream=%d events=%s", w.Body.String(), probes, upstream.calls, detail) + require.Equal(t, 1, upstream.calls) + } + if scenario == "exhausted" { + require.Equal(t, 503, w.Code) + require.Equal(t, 2, slots.released) + require.Zero(t, upstream.calls) + } + }) + } +} + +func TestGrokMediaVideoLookupOwnerIsolation(t *testing.T) { + for _, other := range []string{"user", "api key", "group", "task"} { + t.Run(other, func(t *testing.T) { + h, slots, bindings, upstream := newGrokMediaSlotHandler(t, false, false) + c, w := grokMediaSlotContext(context.Background(), false) + key, ok := middleware2.GetAPIKeyFromContext(c) + require.True(t, ok) + switch other { + case "user": + c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 11, Concurrency: 5}) + case "api key": + key.ID = 21 + case "group": + groupID := int64(25) + key.GroupID = &groupID + case "task": + c.Params = gin.Params{{Key: "request_id", Value: "other-task"}} + } + h.GrokVideoStatus(c) + require.Equal(t, http.StatusNotFound, w.Code) + slots.assertReleased(t) + require.Zero(t, slots.acquired) + require.Zero(t, upstream.calls) + require.Zero(t, bindings.writes) + }) + } +} + +func TestGrokMediaVideoCompletionStillClaimsBillingOnce(t *testing.T) { + h, _, bindings, _ := newGrokMediaSlotHandler(t, false, false) + c, _ := grokMediaSlotContext(context.Background(), false) + key, ok := middleware2.GetAPIKeyFromContext(c) + require.True(t, ok) + subject := middleware2.AuthSubject{UserID: 10, Concurrency: 5} + result := &service.OpenAIForwardResult{ResponseID: "task", Model: "grok-imagine-video", + VideoCount: 1, VideoDurationSeconds: 6} + for i := range 20 { + bill := prepareGrokVideoCompletionBilling(c.Request.Context(), h, zap.NewNop(), key, subject, "task", result) + if i == 0 { + require.NotNil(t, bill) + require.Equal(t, service.StableGrokVideoBillingRequestID("task"), bill.RequestID) + } else { + require.Nil(t, bill) + } + } + require.Len(t, bindings.billed, 1) +} diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index ff4187d987..24b4dcb1b3 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -332,6 +332,36 @@ func (s *OpenAIGatewayService) ResolveGrokMediaVideoRequestAccount( return s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), cacheKey) } +// SelectGrokMediaVideoRequestAccount only admits the already authenticated +// task owner. Generic sticky fallback can query another account and overwrite +// the ownership key; video lookups must neither escape nor refresh that key. +func (s *OpenAIGatewayService) SelectGrokMediaVideoRequestAccount( + ctx context.Context, groupID *int64, sessionHash string, accountID int64, requestedModel string, +) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { + decision := OpenAIAccountScheduleDecision{Layer: openAIAccountScheduleLayerSessionSticky} + if accountID <= 0 || strings.TrimSpace(sessionHash) == "" { + return nil, decision, ErrNoAvailableAccounts + } + ctx = s.withOpenAIGroupPrivacyRequirement(WithOpenAIProfitControlSuppressed(ctx), groupID) + scheduler := &defaultOpenAIAccountScheduler{service: s} + selection, _, err := scheduler.selectBySessionHash(ctx, OpenAIAccountScheduleRequest{ + GroupID: groupID, Platform: PlatformGrok, SessionHash: sessionHash, + StickyAccountID: accountID, PreserveStickyBinding: true, DisableStickyEscape: true, + RequestedModel: requestedModel, RequiredTransport: OpenAIUpstreamTransportHTTPSSE, + RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID), + }) + if err != nil { + return nil, decision, err + } + if selection == nil || selection.Account == nil { + return nil, decision, ErrNoAvailableAccounts + } + decision.StickySessionHit = true + decision.SelectedAccountID = selection.Account.ID + decision.SelectedAccountType = selection.Account.Type + return selection, decision, nil +} + // GrokVideoPendingBilling is the create-time snapshot used when status polling // first observes a completed video URL. Status may omit model/duration; we fall // back to this snapshot, then defaults. diff --git a/backend/internal/service/grok_media_selection_test.go b/backend/internal/service/grok_media_selection_test.go new file mode 100644 index 0000000000..537c7a62c7 --- /dev/null +++ b/backend/internal/service/grok_media_selection_test.go @@ -0,0 +1,110 @@ +package service + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestSelectGrokMediaVideoRequestAccountPreservesOwner(t *testing.T) { + for _, state := range []string{"available", "full", "unavailable", "wrong group", "missing", "invalid id"} { + t.Run(state, func(t *testing.T) { + groupID := int64(24) + ownerID := int64(1) + owner := Account{ID: ownerID, Platform: PlatformGrok, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, Concurrency: 50, GroupIDs: []int64{groupID}} + other := owner + other.ID = 2 + if state == "unavailable" { + until := time.Now().Add(time.Minute) + owner.TempUnschedulableUntil = &until + } + if state == "wrong group" { + owner.GroupIDs = []int64{25} + } + accounts := []Account{owner, other} + if state == "missing" { + accounts = []Account{other} + } + if state == "invalid id" { + ownerID = 0 + } + var acquired, released []int64 + cache := &schedulerTestGatewayCache{} + cfg := &config.Config{} + cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true + cfg.Gateway.Scheduling.StickySessionWaitTimeout = time.Second + cfg.Gateway.Scheduling.StickySessionMaxWaiting = 3 + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, cache: cache, cfg: cfg, + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{ + acquireResults: map[int64]bool{1: state != "full", 2: true}, + acquiredIDs: &acquired, releasedIDs: &released, + }), + } + ctx := context.Background() + require.NoError(t, svc.BindGrokMediaVideoRequestAccount(ctx, &groupID, "task", 10, 20, 1)) + sessionHash := GrokMediaVideoRequestSessionHash("task", 10, 20) + for range 20 { + selection, decision, err := svc.SelectGrokMediaVideoRequestAccount(ctx, &groupID, sessionHash, ownerID, "") + switch state { + case "available": + require.NoError(t, err) + require.Equal(t, int64(1), selection.Account.ID) + require.True(t, selection.Acquired) + require.True(t, decision.StickySessionHit) + selection.ReleaseFunc() + case "full": + require.NoError(t, err) + require.Equal(t, int64(1), selection.Account.ID) + require.False(t, selection.Acquired) + require.Nil(t, selection.ReleaseFunc) + require.Equal(t, int64(1), selection.WaitPlan.AccountID) + default: + require.ErrorIs(t, err, ErrNoAvailableAccounts) + require.Nil(t, selection) + } + bound, err := svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "task", 10, 20) + require.NoError(t, err) + require.Equal(t, int64(1), bound) + } + require.NotContains(t, acquired, int64(2)) + require.Empty(t, cache.deletedSessions) + if state == "available" { + require.Len(t, released, 20) + } else { + require.Empty(t, released) + } + }) + } +} + +func TestGrokVideoStickySelectionIgnoresHealthEscape(t *testing.T) { + groupID := int64(24) + account := Account{ID: 1, Platform: PlatformGrok, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, Concurrency: 50, GroupIDs: []int64{groupID}} + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}}, + cache: &schedulerTestGatewayCache{}, + } + stats := newOpenAIAccountRuntimeStats() + for range 20 { + stats.report(1, false, nil) + } + scheduler := &defaultOpenAIAccountScheduler{service: svc, stats: stats} + req := OpenAIAccountScheduleRequest{GroupID: &groupID, Platform: PlatformGrok, + SessionHash: "task", StickyAccountID: 1, PreserveStickyBinding: true} + selection, escaped, err := scheduler.selectBySessionHash(context.Background(), req) + require.NoError(t, err) + require.Nil(t, selection) + require.True(t, escaped) + req.DisableStickyEscape = true + selection, escaped, err = scheduler.selectBySessionHash(context.Background(), req) + require.NoError(t, err) + require.False(t, escaped) + require.True(t, selection.Acquired) + selection.ReleaseFunc() +} diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 2986f6abeb..5ad63d30f8 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -77,6 +77,9 @@ type OpenAIAccountScheduleRequest struct { StickyWeighted bool SubscriptionPriority bool PreserveStickyBinding bool + // DisableStickyEscape keeps task-owner lookups on their account even when + // generic sticky health or concurrency heuristics would prefer another. + DisableStickyEscape bool RequirePrivacySet bool PreviousResponseID string PreviousResponseCanMove bool @@ -555,7 +558,7 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( return nil, false, nil } escapeCfg := s.service.openAIStickyEscapeConfig() - if reason, errorRate, ttft, shouldEscape := s.shouldEscapeStickyAccount(accountID, escapeCfg); shouldEscape { + if reason, errorRate, ttft, shouldEscape := s.shouldEscapeStickyAccount(accountID, escapeCfg); shouldEscape && !req.DisableStickyEscape { slog.Info("sticky_escape_triggered", "account_id", accountID, "reason", reason, @@ -565,6 +568,9 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( return nil, true, nil } result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency) + if acquireErr != nil && req.DisableStickyEscape { + return nil, false, acquireErr + } if acquireErr == nil && result != nil && result.Acquired { if !req.PreserveStickyBinding { _ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL()) @@ -579,7 +585,7 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( cfg := s.service.schedulingConfig() // WaitPlan.MaxConcurrency 使用 Concurrency（非 EffectiveLoadFactor），因为 WaitPlan 控制的是 Redis 实际并发槽位等待。 if s.service.concurrencyService != nil { - if escapeCfg.enabled && acquireErr == nil && result != nil && !result.Acquired { + if escapeCfg.enabled && !req.DisableStickyEscape && acquireErr == nil && result != nil && !result.Acquired { errorRate, ttft, _ := s.stats.snapshot(accountID) slog.Info("sticky_escape_triggered", "account_id", accountID,