Merge pull request #6661 from whhphd/fix/grok-media-slot-release

fix(grok): release abandoned media slots and preserve video ownership

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

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

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.

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

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,