diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index c6341d5294..ae4e2c2321 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -1776,6 +1776,11 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx con if account == nil { return false, "account_nil" } + if source, ok := CompositeRouteSourceFromContext(ctx); ok && source == CompositeRouteSourceAccount { + if publicModel, modelOK := RequestedPublicModelFromContext(ctx); modelOK && !explicitModelMappingClaims(*account, publicModel) { + return false, "account_model_not_owned" + } + } if req.RequirePrivacySet && !account.IsPrivacySet() { return false, "privacy_not_set" } diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index 081cf30f9a..7a047d6069 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -44,6 +44,10 @@ func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx return result, nil } +func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) { + return append([]Account(nil), r.accounts...), nil +} + func (r schedulerTestOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) { var result []Account for _, acc := range r.accounts { @@ -1257,6 +1261,116 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedSessionIn } } +func TestOpenAIGatewayService_SelectAccountWithScheduler_AccountModelRouteOnlyUsesClaimingAccount(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + + ctx := context.Background() + groupID := int64(1010711) + owner := Account{ + ID: 371011, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10, + GroupIDs: []int64{groupID}, + Credentials: map[string]any{"model_mapping": map[string]any{"team-alias": "gpt-5.1"}}, + } + nonOwner := Account{ + ID: 371012, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, + GroupIDs: []int64{groupID}, + } + repo := schedulerTestOpenAIAccountRepo{accounts: []Account{owner, nonOwner}} + ownership, err := (&GatewayService{accountRepo: repo}).resolveCompositeModelOwnership(ctx, groupID, "team-alias") + require.NoError(t, err) + require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, ownership) + ctx = WithCompositeRouteDecision(ctx, CompositeRouteDecision{ + Matched: true, Source: CompositeRouteSourceAccount, GroupID: groupID, + PublicModel: "team-alias", TargetPlatform: ownership.TargetPlatform, UpstreamModel: "team-alias", + }) + + for _, tt := range []struct { + name string + sessionHash string + sessionBindings map[string]int64 + }{ + {name: "initial filter excludes non owner"}, + {name: "non owner sticky hit is rejected", sessionHash: "non-owner", sessionBindings: map[string]int64{"openai:non-owner": nonOwner.ID}}, + } { + t.Run(tt.name, func(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.Scheduling.LoadBatchEnabled = false + svc := &OpenAIGatewayService{ + accountRepo: repo, + cache: &schedulerTestGatewayCache{sessionBindings: tt.sessionBindings}, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"), + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + + selection, _, selectErr := svc.SelectAccountWithScheduler( + ctx, &groupID, "", tt.sessionHash, "team-alias", nil, OpenAIUpstreamTransportAny, false, + ) + require.NoError(t, selectErr) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, owner.ID, selection.Account.ID) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } + }) + } +} + +func TestOpenAIGatewayService_SelectAccountWithScheduler_AccountModelRouteDBRecheckFailsOverToOwner(t *testing.T) { + ctx := context.Background() + groupID := int64(1010712) + mapping := map[string]any{"model_mapping": map[string]any{"team-alias": "gpt-5.1"}} + stalePrimary := &Account{ + ID: 371021, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, + GroupIDs: []int64{groupID}, Credentials: mapping, + } + staleBackup := &Account{ + ID: 371022, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10, + GroupIDs: []int64{groupID}, Credentials: mapping, + } + dbPrimary := *stalePrimary + dbPrimary.Credentials = nil + dbBackup := *staleBackup + repo := schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbBackup}} + ctx = WithCompositeRouteDecision(ctx, CompositeRouteDecision{ + Matched: true, Source: CompositeRouteSourceAccount, GroupID: groupID, + PublicModel: "team-alias", TargetPlatform: PlatformOpenAI, UpstreamModel: "team-alias", + }) + acquiredIDs, releasedIDs := []int64{}, []int64{} + svc := &OpenAIGatewayService{ + accountRepo: repo, + cfg: &config.Config{RunMode: config.RunModeStandard}, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"), + schedulerSnapshot: &SchedulerSnapshotService{cache: &openAISnapshotCacheStub{ + snapshotAccounts: []*Account{stalePrimary, staleBackup}, + accountsByID: map[int64]*Account{stalePrimary.ID: stalePrimary, staleBackup.ID: staleBackup}, + }}, + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{ + acquiredIDs: &acquiredIDs, releasedIDs: &releasedIDs, + }), + } + + selection, _, err := svc.SelectAccountWithScheduler( + ctx, &groupID, "", "", "team-alias", nil, OpenAIUpstreamTransportAny, false, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, dbBackup.ID, selection.Account.ID) + if containsInt64(acquiredIDs, dbPrimary.ID) { + require.Contains(t, releasedIDs, dbPrimary.ID) + } + require.Contains(t, acquiredIDs, dbBackup.ID) + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedPreviousRequiresMovableContext(t *testing.T) { resetOpenAIAdvancedSchedulerSettingCacheForTest()