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:
parent
a2ab67f8c7
commit
5123d92b44
@ -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 {
|
||||
|
||||
@ -41,6 +41,7 @@ func newGatewayRecordUsageServiceForTest(usageRepo UsageLogRepository, userRepo
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@ -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)
|
||||
}
|
||||
|
||||
|
||||
133
backend/internal/service/gateway_tier_fallback.go
Normal file
133
backend/internal/service/gateway_tier_fallback.go
Normal 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")
|
||||
}
|
||||
138
backend/internal/service/gateway_tier_fallback_test.go
Normal file
138
backend/internal/service/gateway_tier_fallback_test.go
Normal 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")
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user