seewe/sub2api
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-14 12:02:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): 负载感知调度补齐渠道模型限制的逐账号检查

Browse Source 
渠道 billing_model_source=upstream 时，限制模型只能按账号映射后的上游模型
逐账号判定。GatewayService.SelectAccountWithLoadAwareness 的负载感知各层
（模型路由候选、路由内粘性 gate、无路由粘性 gate、负载感知候选及其兜底）
此前都没有这一判定，load_batch_enabled 默认开启，anthropic/gemini/antigravity
分组的渠道白名单实际不生效。

现在入口算一次 needsUpstreamChannelRestrictionCheck，各层共用同一判定；
候选因渠道限制筛空时记 warn 并返回带 channel pricing restriction 的
ErrNoAvailableAccounts，与预检路径一致。
This commit is contained in:
feeeei
2026-09-02 11:46:37 +08:00
parent 5097b31457
commit 9ad3865690
2 changed files with 196 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/gateway_channel_restriction_load_aware_test.go
+165
View File
@@ -0,0 +1,165 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/stretchr/testify/require"
)
// loadAwareRestrictionFixture 构造走负载感知路径（load_batch 开启 + 并发服务存在）的
// GatewayService：分组 10 绑定 upstream 计费基准的渠道，定价列表只允许 claude-sonnet-4-6。
// 账号 1 把 claude-fable-5-1 映射为自身（不在定价列表），账号 2 映射为 claude-sonnet-4-6（在列表）。
type loadAwareRestrictionFixture struct {
svc *GatewayService
ctx context.Context
groupID int64
cache *mockGatewayCacheForPlatform
concurrencyCache *mockConcurrencyCache
}
func newLoadAwareRestrictionFixture(t *testing.T, restrict bool, sessionBindings map[string]int64, routing map[string][]int64) *loadAwareRestrictionFixture {
t.Helper()
groupID := int64(10)
ch := Channel{
ID: 1,
Status: StatusActive,
GroupIDs: []int64{groupID},
RestrictModels: restrict,
BillingModelSource: BillingModelSourceUpstream,
ModelPricing: []ChannelModelPricing{
{Platform: PlatformAnthropic, Models: []string{"claude-sonnet-4-6"}},
},
}
channelSvc := newTestChannelService(makeStandardRepo(ch, map[int64]string{groupID: PlatformAnthropic}))
accountRepo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true, Concurrency: 5,
Credentials: map[string]any{
"model_mapping": map[string]any{"claude-fable-5-1": "claude-fable-5-1"},
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true, Concurrency: 5,
Credentials: map[string]any{
"model_mapping": map[string]any{"claude-fable-5-1": "claude-sonnet-4-6"},
},
},
},
accountsByID: map[int64]*Account{},
}
for i := range accountRepo.accounts {
accountRepo.accountsByID[accountRepo.accounts[i].ID] = &accountRepo.accounts[i]
}
group := &Group{
ID: groupID,
Platform: PlatformAnthropic,
Status: StatusActive,
Hydrated: true,
ModelRoutingEnabled: len(routing) > 0,
ModelRouting: routing,
}
groupRepo := &mockGroupRepoForGateway{groups: map[int64]*Group{groupID: group}}
cfg := testConfig()
cfg.Gateway.Scheduling.LoadBatchEnabled = true
cache := &mockGatewayCacheForPlatform{sessionBindings: sessionBindings}
concurrencyCache := &mockConcurrencyCache{}
svc := &GatewayService{
accountRepo: accountRepo,
groupRepo: groupRepo,
channelService: channelSvc,
cache: cache,
cfg: cfg,
concurrencyService: NewConcurrencyService(concurrencyCache),
}
return &loadAwareRestrictionFixture{
svc: svc,
ctx: context.WithValue(context.Background(), ctxkey.Group, group),
groupID: groupID,
cache: cache,
concurrencyCache: concurrencyCache,
}
}
func TestSelectAccountWithLoadAwareness_UpstreamRestrictionSkipsDisallowedAccount(t *testing.T) {
t.Parallel()
f := newLoadAwareRestrictionFixture(t, true, nil, nil)
result, err := f.svc.SelectAccountWithLoadAwareness(f.ctx, &f.groupID, "", "claude-fable-5-1", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(2), result.Account.ID, "上游模型不在渠道定价列表的高优先级账号必须被跳过")
require.Equal(t, 1, f.concurrencyCache.loadBatchCalls, "应经过 Layer 2 负载感知选择")
}
func TestSelectAccountWithLoadAwareness_UpstreamRestrictionRejectsWhenAllAccountsDisallowed(t *testing.T) {
t.Parallel()
f := newLoadAwareRestrictionFixture(t, true, nil, nil)
// 账号 2 也改为映射到不在定价列表的模型
f.svc.accountRepo.(*mockAccountRepoForPlatform).accounts[1].Credentials = map[string]any{
"model_mapping": map[string]any{"claude-fable-5-1": "claude-fable-5-1"},
}
result, err := f.svc.SelectAccountWithLoadAwareness(f.ctx, &f.groupID, "", "claude-fable-5-1", nil, "", 0)
require.ErrorIs(t, err, ErrNoAvailableAccounts)
require.ErrorContains(t, err, "channel pricing restriction")
require.Nil(t, result)
require.Equal(t, 0, f.concurrencyCache.acquireAccountCalls, "没有合规候选时不应尝试占用任何账号槽位")
}
func TestSelectAccountWithLoadAwareness_UpstreamRestrictionIgnoresStickyAccount(t *testing.T) {
t.Parallel()
f := newLoadAwareRestrictionFixture(t, true, map[string]int64{"sticky": 1}, nil)
result, err := f.svc.SelectAccountWithLoadAwareness(f.ctx, &f.groupID, "sticky", "claude-fable-5-1", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(2), result.Account.ID, "粘性账号的上游模型不在定价列表时不得沿用粘性会话")
require.Equal(t, int64(2), f.cache.sessionBindings["sticky"], "粘性会话应重新绑定到合规账号")
}
func TestSelectAccountWithLoadAwareness_UpstreamRestrictionFiltersRoutedAccounts(t *testing.T) {
t.Parallel()
f := newLoadAwareRestrictionFixture(t, true, nil, map[string][]int64{"claude-fable-5-1": {1}})
result, err := f.svc.SelectAccountWithLoadAwareness(f.ctx, &f.groupID, "", "claude-fable-5-1", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(2), result.Account.ID, "模型路由指向的账号被渠道限制时应回退到合规账号")
}
func TestSelectAccountWithLoadAwareness_UpstreamRestrictionRoutedStickyAccountNotHonored(t *testing.T) {
t.Parallel()
f := newLoadAwareRestrictionFixture(t, true, map[string]int64{"sticky": 1}, map[string][]int64{"claude-fable-5-1": {1, 2}})
result, err := f.svc.SelectAccountWithLoadAwareness(f.ctx, &f.groupID, "sticky", "claude-fable-5-1", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(2), result.Account.ID, "路由列表内的粘性账号被渠道限制时不得沿用")
}
func TestSelectAccountWithLoadAwareness_RestrictModelsDisabledKeepsPriorityOrder(t *testing.T) {
t.Parallel()
f := newLoadAwareRestrictionFixture(t, false, nil, nil)
result, err := f.svc.SelectAccountWithLoadAwareness(f.ctx, &f.groupID, "", "claude-fable-5-1", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(1), result.Account.ID, "未开启限制时不应因定价列表过滤账号")
}
backend/internal/service/gateway_scheduling.go
+31 -4
View File
@@ -243,6 +243,13 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
return excluded
}
// upstream 计费基准的渠道模型限制以账号映射后的上游模型为准，只能逐账号判定；
// 负载感知各层的候选过滤与粘性 gate 共用这一判定。
needsUpstreamCheck := s.needsUpstreamChannelRestrictionCheck(ctx, groupID)
isChannelRestricted := func(account *Account) bool {
return needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel)
}
// 获取模型路由配置（anthropic 目标平台；composite 分组按目标平台判断）
var routingAccountIDs []int64
if group != nil && requestedModel != "" && platform == PlatformAnthropic &&
@@ -270,7 +277,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
if len(routingAccountIDs) > 0 && s.concurrencyService != nil {
// 1. 过滤出路由列表中可调度的账号
var routingCandidates []*Account
var filteredExcluded, filteredMissing, filteredUnsched, filteredPlatform, filteredModelScope, filteredModelMapping, filteredWindowCost int
var filteredExcluded, filteredMissing, filteredUnsched, filteredPlatform, filteredModelScope, filteredModelMapping, filteredChannelRestricted, filteredWindowCost int
var modelScopeSkippedIDs []int64 // 记录因模型限流被跳过的账号 ID
for _, routingAccountID := range routingAccountIDs {
if isExcluded(routingAccountID) {
@@ -297,6 +304,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
filteredModelMapping++
continue
}
if isChannelRestricted(account) {
filteredChannelRestricted++
continue
}
if !s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
filteredModelScope++
modelScopeSkippedIDs = append(modelScopeSkippedIDs, account.ID)
@@ -319,9 +330,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
}
if s.debugModelRoutingEnabled() {
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] routed candidates: group_id=%v model=%s routed=%d candidates=%d filtered(excluded=%d missing=%d unsched=%d platform=%d model_scope=%d model_mapping=%d window_cost=%d)",
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] routed candidates: group_id=%v model=%s routed=%d candidates=%d filtered(excluded=%d missing=%d unsched=%d platform=%d model_scope=%d model_mapping=%d channel_restricted=%d window_cost=%d)",
derefGroupID(groupID), requestedModel, len(routingAccountIDs), len(routingCandidates),
filteredExcluded, filteredMissing, filteredUnsched, filteredPlatform, filteredModelScope, filteredModelMapping, filteredWindowCost)
filteredExcluded, filteredMissing, filteredUnsched, filteredPlatform, filteredModelScope, filteredModelMapping, filteredChannelRestricted, filteredWindowCost)
if len(modelScopeSkippedIDs) > 0 {
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] model_rate_limited accounts skipped: group_id=%v model=%s account_ids=%v",
derefGroupID(groupID), requestedModel, modelScopeSkippedIDs)
@@ -347,6 +358,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
s.isGatewayAccountProfitEligible(ctx, stickyAccount) &&
s.isAccountAllowedForPlatform(stickyAccount, platform, useMixed) &&
(requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, stickyAccount, requestedModel)) &&
!isChannelRestricted(stickyAccount) &&
s.isAccountSchedulableForModelSelection(ctx, stickyAccount, requestedModel) &&
s.isAccountSchedulableForQuota(stickyAccount) &&
s.isAccountSchedulableForWindowCost(ctx, stickyAccount, true)
@@ -531,6 +543,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
platformOK := s.isAccountAllowedForPlatform(account, platform, useMixed)
profitOK := s.isGatewayAccountProfitEligible(ctx, account)
modelSupported := requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)
channelOK := !isChannelRestricted(account)
modelSchedulable := s.isAccountSchedulableForModelSelection(ctx, account, requestedModel)
quotaOK := s.isAccountSchedulableForQuota(account)
windowCostOK := s.isAccountSchedulableForWindowCost(ctx, account, true)
@@ -545,13 +558,14 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
"platform_ok", platformOK,
"profit_ok", profitOK,
"model_supported", modelSupported,
"channel_ok", channelOK,
"model_schedulable", modelSchedulable,
"quota_ok", quotaOK,
"window_cost_ok", windowCostOK,
"rpm_ok", rpmOK,
)
if !clearSticky && platformOK && profitOK && modelSupported && modelSchedulable && quotaOK && windowCostOK && rpmOK && schedulable {
if !clearSticky && platformOK && profitOK && modelSupported && channelOK && modelSchedulable && quotaOK && windowCostOK && rpmOK && schedulable {
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
if err == nil && result.Acquired {
// 会话数量限制检查
@@ -636,6 +650,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
"total_accounts", len(accounts),
)
candidates := make([]*Account, 0, len(accounts))
channelRestrictedCount := 0
for i := range accounts {
acc := &accounts[i]
if isExcluded(acc.ID) {
@@ -656,6 +671,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
if requestedModel != "" && !s.isModelSupportedByAccountWithContext(ctx, acc, requestedModel) {
continue
}
if isChannelRestricted(acc) {
channelRestrictedCount++
continue
}
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
continue
}
@@ -675,6 +694,14 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
}
if len(candidates) == 0 {
if channelRestrictedCount > 0 {
slog.Warn("channel pricing restriction blocked request",
"group_id", derefGroupID(groupID),
"model", requestedModel,
"restricted_accounts", channelRestrictedCount,
"total_accounts", len(accounts))
return nil, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
}
return nil, ErrNoAvailableAccounts
}