From b97a798ebcb2571be2f9e08917f4cf10fe322a27 Mon Sep 17 00:00:00 2001 From: whhphd Date: Sat, 5 Sep 2026 19:45:39 +0800 Subject: [PATCH] fix(grok): release abandoned media slots and preserve video ownership Claim scheduler-acquired slots before media eligibility checks and task-owner validation. Release abandoned attempts immediately and retain once-only cleanup on cancellation and forwarding completion. Select only the resolved video owner, reuse bounded account waiting without sticky escape, and preserve the task binding and TTL. Keep generation failover and owner/billing isolation intact. Regression tests reproduce unreleased account slots on upstream ab99d56e9 and pass with this change. Validation: go test ./..., unit and integration suites, focused race tests, backend build, and golangci-lint (0 issues). --- backend/internal/handler/grok_media.go | 76 +++- .../internal/handler/grok_media_slots_test.go | 405 ++++++++++++++++++ backend/internal/service/grok_media.go | 30 ++ .../service/grok_media_selection_test.go | 110 +++++ .../service/openai_account_scheduler.go | 10 +- 5 files changed, 609 insertions(+), 22 deletions(-) create mode 100644 backend/internal/handler/grok_media_slots_test.go create mode 100644 backend/internal/service/grok_media_selection_test.go 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,