feat(scheduling): add cross-tier fallback chain (subscription → API Key → Bedrock)

Adds an opt-in tier-based fallback scheduling path for Anthropic accounts:
- accountTierLevel(): derives tier from account type without DB migration
  (tier-0=OAuth/SetupToken, tier-1=APIKey, tier-2=Bedrock)
- enableTierFallbackChain(): new config flag
  gateway.scheduling.enable_tier_fallback_chain (default false)
- selectAccountWithTierFallback(): loads all Anthropic accounts, groups by
  tier, honors sticky sessions, applies all existing schedulability guards,
  then tries tiers 0→1→2 in order via tryAcquireByLegacyOrder
- Wired into SelectAccountForModelWithExclusions: Anthropic platform +
  tier fallback enabled → calls new path instead of mixed scheduling
- Fix pre-existing unit-test build break: NewGatewayService now requires
  *RPMTokenBucketService (added in Task #5); add missing nil param
- 7 tests: tier mapping, config toggle, subscription preference,
  APIKey fallback, exclusion handling, empty-pool error, Bedrock last resort
This commit is contained in:
win 2026-04-29 03:23:39 +08:00
parent a2ab67f8c7
commit 5123d92b44
5 changed files with 279 additions and 0 deletions

View File

@ -810,6 +810,10 @@ type GatewaySchedulingConfig struct {
// 全量重建周期配置
// 全量重建周期0 表示禁用
FullRebuildIntervalSeconds int `mapstructure:"full_rebuild_interval_seconds"`
// EnableTierFallbackChain: 启用跨档降级链(订阅 → API Key → Bedrock默认 false
// 仅对 Anthropic 平台生效;启用后账号按类型分层,优先使用订阅账号,依次降级。
EnableTierFallbackChain bool `mapstructure:"enable_tier_fallback_chain"`
}
func (s *ServerConfig) Address() string {

View File

@ -41,6 +41,7 @@ func newGatewayRecordUsageServiceForTest(usageRepo UsageLogRepository, userRepo
nil,
nil,
nil,
nil,
)
}

View File

@ -1192,6 +1192,9 @@ func (s *GatewayService) SelectAccountForModelWithExclusions(ctx context.Context
// anthropic/gemini 分组支持混合调度(包含启用了 mixed_scheduling 的 antigravity 账户)
// 注意:强制平台模式不走混合调度
if (platform == PlatformAnthropic || platform == PlatformGemini) && !hasForcePlatform {
if platform == PlatformAnthropic && s.enableTierFallbackChain() {
return s.selectAccountWithTierFallback(ctx, groupID, sessionHash, requestedModel, excludedIDs)
}
return s.selectAccountWithMixedScheduling(ctx, groupID, sessionHash, requestedModel, excludedIDs, platform)
}

View File

@ -0,0 +1,133 @@
package service
import (
"context"
"errors"
)
// accountTierLevel maps an account type to a scheduling tier:
//
// 0 = subscription (OAuth / SetupToken) — tried first
// 1 = API Key — first fallback
// 2 = Bedrock — last resort
//
// Accounts with an unknown type fall into tier 0 so they participate in the
// primary selection and do not vanish silently.
func accountTierLevel(account *Account) int {
if account == nil {
return 0
}
switch account.Type {
case AccountTypeAPIKey:
return 1
case AccountTypeBedrock:
return 2
default: // OAuth, SetupToken, or unknown
return 0
}
}
// enableTierFallbackChain reports whether the cross-tier fallback chain is
// enabled in config (default false).
func (s *GatewayService) enableTierFallbackChain() bool {
return s != nil && s.cfg != nil && s.cfg.Gateway.Scheduling.EnableTierFallbackChain
}
// selectAccountWithTierFallback tries Anthropic accounts in tier order:
// tier 0 (OAuth/SetupToken subscription) → tier 1 (API Key) → tier 2 (Bedrock).
//
// Sticky sessions are honored within the chain: if the session-bound account is
// in a tier that still has capacity it is returned immediately; otherwise the
// session binding is cleared and the chain proceeds from tier 0.
func (s *GatewayService) selectAccountWithTierFallback(
ctx context.Context,
groupID *int64,
sessionHash string,
requestedModel string,
excludedIDs map[int64]struct{},
) (*Account, error) {
accounts, _, err := s.listSchedulableAccounts(ctx, groupID, PlatformAnthropic, false)
if err != nil {
return nil, err
}
ctx = s.withWindowCostPrefetch(ctx, accounts)
ctx = s.withRPMPrefetch(ctx, accounts)
// Build per-tier candidate lists (pointers into `accounts`).
const numTiers = 3
tierCandidates := [numTiers][]*Account{}
for i := range accounts {
acc := &accounts[i]
if acc.Platform != PlatformAnthropic {
continue
}
if _, excluded := excludedIDs[acc.ID]; excluded {
continue
}
if !s.isAccountSchedulableForSelection(acc) {
continue
}
if requestedModel != "" && !s.isModelSupportedByAccountWithContext(ctx, acc, requestedModel) {
continue
}
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
continue
}
if !s.isAccountSchedulableForQuota(acc) {
continue
}
if !s.isAccountSchedulableForWindowCost(ctx, acc, false) {
continue
}
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
continue
}
tier := accountTierLevel(acc)
if tier < numTiers {
tierCandidates[tier] = append(tierCandidates[tier], acc)
}
}
cfg := s.schedulingConfig()
selectionMode := cfg.FallbackSelectionMode
// Check sticky session: if the bound account is a valid candidate, use it.
if sessionHash != "" && s.cache != nil {
accountID, cacheErr := s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
if cacheErr == nil && accountID > 0 {
if _, excluded := excludedIDs[accountID]; !excluded {
for tier := 0; tier < numTiers; tier++ {
for _, acc := range tierCandidates[tier] {
if acc.ID != accountID {
continue
}
if shouldClearStickySession(acc, requestedModel) {
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
break
}
if s.isAccountSchedulableForWindowCost(ctx, acc, true) &&
s.isAccountSchedulableForRPM(ctx, acc, true) {
return acc, nil
}
}
}
}
}
}
// Try each tier in order.
for tier := 0; tier < numTiers; tier++ {
candidates := tierCandidates[tier]
if len(candidates) == 0 {
continue
}
s.sortCandidatesForFallback(candidates, false, selectionMode)
result, acquired := s.tryAcquireByLegacyOrder(ctx, candidates, groupID, sessionHash, false)
if acquired && result != nil {
return result.Account, nil
}
}
return nil, errors.New("no available accounts in any tier")
}

View File

@ -0,0 +1,138 @@
//go:build unit
package service
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
func TestAccountTierLevel(t *testing.T) {
require.Equal(t, 0, accountTierLevel(nil))
require.Equal(t, 0, accountTierLevel(&Account{Type: AccountTypeOAuth}))
require.Equal(t, 0, accountTierLevel(&Account{Type: AccountTypeSetupToken}))
require.Equal(t, 0, accountTierLevel(&Account{Type: "unknown"}))
require.Equal(t, 1, accountTierLevel(&Account{Type: AccountTypeAPIKey}))
require.Equal(t, 2, accountTierLevel(&Account{Type: AccountTypeBedrock}))
}
func TestGatewayService_EnableTierFallbackChain(t *testing.T) {
require.False(t, (*GatewayService)(nil).enableTierFallbackChain())
require.False(t, (&GatewayService{}).enableTierFallbackChain())
cfgOff := &config.Config{}
cfgOff.Gateway.Scheduling.EnableTierFallbackChain = false
require.False(t, (&GatewayService{cfg: cfgOff}).enableTierFallbackChain())
cfgOn := &config.Config{}
cfgOn.Gateway.Scheduling.EnableTierFallbackChain = true
require.True(t, (&GatewayService{cfg: cfgOn}).enableTierFallbackChain())
}
// TestGatewayService_SelectAccountWithTierFallback_PrefersSubscription verifies
// that when both OAuth (subscription) and APIKey accounts are available, the
// tier-0 OAuth account is always selected first even if APIKey has higher priority.
func TestGatewayService_SelectAccountWithTierFallback_PrefersSubscription(t *testing.T) {
ctx := context.Background()
oauthAcc := Account{ID: 91001, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Priority: 5}
apiKeyAcc := Account{ID: 91002, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Priority: 0}
repo := &mockAccountRepoForPlatform{
accounts: []Account{oauthAcc, apiKeyAcc},
accountsByID: map[int64]*Account{91001: &oauthAcc, 91002: &apiKeyAcc},
}
cache := &mockGatewayCacheForPlatform{}
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", nil)
require.NoError(t, err)
require.NotNil(t, acc)
require.Equal(t, int64(91001), acc.ID, "OAuth (tier-0) account should be preferred over APIKey (tier-1)")
}
// TestGatewayService_SelectAccountWithTierFallback_FallsBackToAPIKey verifies
// that when the subscription tier has no schedulable accounts, the fallback
// selects an API Key account.
func TestGatewayService_SelectAccountWithTierFallback_FallsBackToAPIKey(t *testing.T) {
ctx := context.Background()
rateLimitedUntil := time.Now().Add(30 * time.Minute)
oauthAcc := Account{ID: 92001, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, RateLimitResetAt: &rateLimitedUntil}
apiKeyAcc := Account{ID: 92002, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true}
repo := &mockAccountRepoForPlatform{
accounts: []Account{oauthAcc, apiKeyAcc},
accountsByID: map[int64]*Account{92001: &oauthAcc, 92002: &apiKeyAcc},
}
cache := &mockGatewayCacheForPlatform{}
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", nil)
require.NoError(t, err)
require.NotNil(t, acc)
require.Equal(t, int64(92002), acc.ID, "Should fall back to APIKey when OAuth is rate-limited")
}
// TestGatewayService_SelectAccountWithTierFallback_ExcludesAccounts ensures
// excluded IDs are respected across all tiers.
func TestGatewayService_SelectAccountWithTierFallback_ExcludesAccounts(t *testing.T) {
ctx := context.Background()
oauthAcc := Account{ID: 93001, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
apiKeyAcc := Account{ID: 93002, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true}
repo := &mockAccountRepoForPlatform{
accounts: []Account{oauthAcc, apiKeyAcc},
accountsByID: map[int64]*Account{93001: &oauthAcc, 93002: &apiKeyAcc},
}
cache := &mockGatewayCacheForPlatform{}
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
excluded := map[int64]struct{}{93001: {}}
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", excluded)
require.NoError(t, err)
require.NotNil(t, acc)
require.Equal(t, int64(93002), acc.ID, "Excluded OAuth account should cause APIKey fallback")
}
// TestGatewayService_SelectAccountWithTierFallback_NoAccounts verifies that
// an error is returned when all tiers are empty.
func TestGatewayService_SelectAccountWithTierFallback_NoAccounts(t *testing.T) {
ctx := context.Background()
repo := &mockAccountRepoForPlatform{accounts: nil, accountsByID: map[int64]*Account{}}
cache := &mockGatewayCacheForPlatform{}
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", nil)
require.Error(t, err)
require.Nil(t, acc)
}
// TestGatewayService_SelectAccountWithTierFallback_BedrockLastResort verifies
// that Bedrock accounts are only used when subscription and API Key tiers are exhausted.
func TestGatewayService_SelectAccountWithTierFallback_BedrockLastResort(t *testing.T) {
ctx := context.Background()
rateLimitedUntil := time.Now().Add(30 * time.Minute)
oauthAcc := Account{ID: 94001, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, RateLimitResetAt: &rateLimitedUntil}
apiKeyAcc := Account{ID: 94002, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, RateLimitResetAt: &rateLimitedUntil}
bedrockAcc := Account{ID: 94003, Platform: PlatformAnthropic, Type: AccountTypeBedrock, Status: StatusActive, Schedulable: true}
repo := &mockAccountRepoForPlatform{
accounts: []Account{oauthAcc, apiKeyAcc, bedrockAcc},
accountsByID: map[int64]*Account{94001: &oauthAcc, 94002: &apiKeyAcc, 94003: &bedrockAcc},
}
cache := &mockGatewayCacheForPlatform{}
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", nil)
require.NoError(t, err)
require.NotNil(t, acc)
require.Equal(t, int64(94003), acc.ID, "Bedrock should be selected as last resort")
}