| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946 |
- package service
- import (
- "context"
- "errors"
- "sync"
- "sync/atomic"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
- "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
- gocache "github.com/patrickmn/go-cache"
- "github.com/stretchr/testify/require"
- )
- type userGroupRateRepoHotpathStub struct {
- UserGroupRateRepository
- rate *float64
- err error
- wait <-chan struct{}
- calls atomic.Int64
- }
- func (s *userGroupRateRepoHotpathStub) GetByUserAndGroup(ctx context.Context, userID, groupID int64) (*float64, error) {
- s.calls.Add(1)
- if s.wait != nil {
- <-s.wait
- }
- if s.err != nil {
- return nil, s.err
- }
- return s.rate, nil
- }
- type usageLogWindowBatchRepoStub struct {
- UsageLogRepository
- batchResult map[int64]*usagestats.AccountStats
- batchErr error
- batchCalls atomic.Int64
- singleResult map[int64]*usagestats.AccountStats
- singleErr error
- singleCalls atomic.Int64
- }
- func (s *usageLogWindowBatchRepoStub) GetAccountWindowStatsBatch(ctx context.Context, accountIDs []int64, startTime time.Time) (map[int64]*usagestats.AccountStats, error) {
- s.batchCalls.Add(1)
- if s.batchErr != nil {
- return nil, s.batchErr
- }
- out := make(map[int64]*usagestats.AccountStats, len(accountIDs))
- for _, id := range accountIDs {
- if stats, ok := s.batchResult[id]; ok {
- out[id] = stats
- }
- }
- return out, nil
- }
- func (s *usageLogWindowBatchRepoStub) GetAccountWindowStats(ctx context.Context, accountID int64, startTime time.Time) (*usagestats.AccountStats, error) {
- s.singleCalls.Add(1)
- if s.singleErr != nil {
- return nil, s.singleErr
- }
- if stats, ok := s.singleResult[accountID]; ok {
- return stats, nil
- }
- return &usagestats.AccountStats{}, nil
- }
- type sessionLimitCacheHotpathStub struct {
- SessionLimitCache
- batchData map[int64]float64
- batchErr error
- setData map[int64]float64
- setErr error
- }
- func (s *sessionLimitCacheHotpathStub) GetWindowCostBatch(ctx context.Context, accountIDs []int64) (map[int64]float64, error) {
- if s.batchErr != nil {
- return nil, s.batchErr
- }
- out := make(map[int64]float64, len(accountIDs))
- for _, id := range accountIDs {
- if v, ok := s.batchData[id]; ok {
- out[id] = v
- }
- }
- return out, nil
- }
- func (s *sessionLimitCacheHotpathStub) SetWindowCost(ctx context.Context, accountID int64, cost float64) error {
- if s.setErr != nil {
- return s.setErr
- }
- if s.setData == nil {
- s.setData = make(map[int64]float64)
- }
- s.setData[accountID] = cost
- return nil
- }
- type modelsListAccountRepoStub struct {
- AccountRepository
- byGroup map[int64][]Account
- all []Account
- err error
- listByGroupCalls atomic.Int64
- listAllCalls atomic.Int64
- }
- type stickyGatewayCacheHotpathStub struct {
- GatewayCache
- stickyID int64
- getCalls atomic.Int64
- }
- func (s *stickyGatewayCacheHotpathStub) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) {
- s.getCalls.Add(1)
- if s.stickyID > 0 {
- return s.stickyID, nil
- }
- return 0, errors.New("not found")
- }
- func (s *stickyGatewayCacheHotpathStub) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error {
- return nil
- }
- func (s *stickyGatewayCacheHotpathStub) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error {
- return nil
- }
- func (s *stickyGatewayCacheHotpathStub) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error {
- return nil
- }
- func (s *stickyGatewayCacheHotpathStub) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
- return nil
- }
- func (s *stickyGatewayCacheHotpathStub) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
- return nil, nil
- }
- func (s *stickyGatewayCacheHotpathStub) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
- return true, nil
- }
- func (s *stickyGatewayCacheHotpathStub) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
- return nil
- }
- func (s *stickyGatewayCacheHotpathStub) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
- return nil
- }
- func (s *stickyGatewayCacheHotpathStub) GetReasoningContent(_ context.Context, _ string) (string, error) {
- return "", ErrReasoningContentNotFound
- }
- func (s *modelsListAccountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
- s.listByGroupCalls.Add(1)
- if s.err != nil {
- return nil, s.err
- }
- accounts, ok := s.byGroup[groupID]
- if !ok {
- return nil, nil
- }
- out := make([]Account, len(accounts))
- copy(out, accounts)
- return out, nil
- }
- func (s *modelsListAccountRepoStub) ListSchedulable(ctx context.Context) ([]Account, error) {
- s.listAllCalls.Add(1)
- if s.err != nil {
- return nil, s.err
- }
- out := make([]Account, len(s.all))
- copy(out, s.all)
- return out, nil
- }
- func resetGatewayHotpathStatsForTest() {
- windowCostPrefetchCacheHitTotal.Store(0)
- windowCostPrefetchCacheMissTotal.Store(0)
- windowCostPrefetchBatchSQLTotal.Store(0)
- windowCostPrefetchFallbackTotal.Store(0)
- windowCostPrefetchErrorTotal.Store(0)
- userGroupRateCacheHitTotal.Store(0)
- userGroupRateCacheMissTotal.Store(0)
- userGroupRateCacheLoadTotal.Store(0)
- userGroupRateCacheSFSharedTotal.Store(0)
- userGroupRateCacheFallbackTotal.Store(0)
- modelsListCacheHitTotal.Store(0)
- modelsListCacheMissTotal.Store(0)
- modelsListCacheStoreTotal.Store(0)
- }
- func TestGetUserGroupRateMultiplier_UsesCacheAndSingleflight(t *testing.T) {
- resetGatewayHotpathStatsForTest()
- rate := 1.7
- unblock := make(chan struct{})
- repo := &userGroupRateRepoHotpathStub{
- rate: &rate,
- wait: unblock,
- }
- svc := &GatewayService{
- userGroupRateRepo: repo,
- userGroupRateCache: gocache.New(time.Minute, time.Minute),
- cfg: &config.Config{
- Gateway: config.GatewayConfig{
- UserGroupRateCacheTTLSeconds: 30,
- },
- },
- }
- const concurrent = 12
- results := make([]float64, concurrent)
- start := make(chan struct{})
- var wg sync.WaitGroup
- wg.Add(concurrent)
- for i := 0; i < concurrent; i++ {
- go func(idx int) {
- defer wg.Done()
- <-start
- results[idx] = svc.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.2)
- }(i)
- }
- close(start)
- // Wait for every caller to have recorded its cache miss before releasing the
- // loader. A fixed sleep raced here: a goroutine that reached the cache after
- // the singleflight load had already finished got a hit instead of a miss, and
- // the miss assertion below saw 11 of 12. The miss counter is the observable
- // that says "all callers are now inside the singleflight group".
- require.Eventually(t, func() bool {
- _, miss, _, _, _ := GatewayUserGroupRateCacheStats()
- return miss == int64(concurrent)
- }, 5*time.Second, time.Millisecond, "all callers must miss the cache before the loader is released")
- close(unblock)
- wg.Wait()
- for _, got := range results {
- require.Equal(t, rate, got)
- }
- require.Equal(t, int64(1), repo.calls.Load())
- // 再次读取应命中缓存,不再回源。
- got := svc.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.2)
- require.Equal(t, rate, got)
- require.Equal(t, int64(1), repo.calls.Load())
- hit, miss, load, sfShared, fallback := GatewayUserGroupRateCacheStats()
- require.GreaterOrEqual(t, hit, int64(1))
- require.Equal(t, int64(12), miss)
- require.Equal(t, int64(1), load)
- require.GreaterOrEqual(t, sfShared, int64(1))
- require.Equal(t, int64(0), fallback)
- }
- func TestGetUserGroupRateMultiplier_FallbackOnRepoError(t *testing.T) {
- resetGatewayHotpathStatsForTest()
- repo := &userGroupRateRepoHotpathStub{
- err: errors.New("db down"),
- }
- svc := &GatewayService{
- userGroupRateRepo: repo,
- userGroupRateCache: gocache.New(time.Minute, time.Minute),
- cfg: &config.Config{
- Gateway: config.GatewayConfig{
- UserGroupRateCacheTTLSeconds: 30,
- },
- },
- }
- got := svc.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.25)
- require.Equal(t, 1.25, got)
- require.Equal(t, int64(1), repo.calls.Load())
- _, _, _, _, fallback := GatewayUserGroupRateCacheStats()
- require.Equal(t, int64(1), fallback)
- }
- func TestGetUserGroupRateMultiplier_CacheHitAndNilRepo(t *testing.T) {
- resetGatewayHotpathStatsForTest()
- repo := &userGroupRateRepoHotpathStub{
- err: errors.New("should not be called"),
- }
- svc := &GatewayService{
- userGroupRateRepo: repo,
- userGroupRateCache: gocache.New(time.Minute, time.Minute),
- }
- key := "101:202"
- svc.userGroupRateCache.Set(key, 2.3, time.Minute)
- got := svc.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.1)
- require.Equal(t, 2.3, got)
- hit, miss, load, _, fallback := GatewayUserGroupRateCacheStats()
- require.Equal(t, int64(1), hit)
- require.Equal(t, int64(0), miss)
- require.Equal(t, int64(0), load)
- require.Equal(t, int64(0), fallback)
- require.Equal(t, int64(0), repo.calls.Load())
- // 无 repo 时直接返回分组默认倍率
- svc2 := &GatewayService{
- userGroupRateCache: gocache.New(time.Minute, time.Minute),
- }
- svc2.userGroupRateCache.Set(key, 1.9, time.Minute)
- require.Equal(t, 1.9, svc2.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.4))
- require.Equal(t, 1.4, svc2.getUserGroupRateMultiplier(context.Background(), 0, 202, 1.4))
- svc2.userGroupRateCache.Delete(key)
- require.Equal(t, 1.4, svc2.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.4))
- }
- func TestWithWindowCostPrefetch_BatchReadAndContextReuse(t *testing.T) {
- resetGatewayHotpathStatsForTest()
- windowStart := time.Now().Add(-30 * time.Minute).Truncate(time.Hour)
- windowEnd := windowStart.Add(5 * time.Hour)
- accounts := []Account{
- {
- ID: 1,
- Platform: PlatformAnthropic,
- Type: AccountTypeOAuth,
- Extra: map[string]any{"window_cost_limit": 100.0},
- SessionWindowStart: &windowStart,
- SessionWindowEnd: &windowEnd,
- },
- {
- ID: 2,
- Platform: PlatformAnthropic,
- Type: AccountTypeSetupToken,
- Extra: map[string]any{"window_cost_limit": 100.0},
- SessionWindowStart: &windowStart,
- SessionWindowEnd: &windowEnd,
- },
- {
- ID: 3,
- Platform: PlatformAnthropic,
- Type: AccountTypeAPIKey,
- Extra: map[string]any{"window_cost_limit": 100.0},
- },
- }
- cache := &sessionLimitCacheHotpathStub{
- batchData: map[int64]float64{
- 1: 11.0,
- },
- }
- repo := &usageLogWindowBatchRepoStub{
- batchResult: map[int64]*usagestats.AccountStats{
- 2: {StandardCost: 22.0},
- },
- }
- svc := &GatewayService{
- sessionLimitCache: cache,
- usageLogRepo: repo,
- }
- outCtx := svc.withWindowCostPrefetch(context.Background(), accounts)
- require.NotNil(t, outCtx)
- cost1, ok1 := windowCostFromPrefetchContext(outCtx, 1)
- require.True(t, ok1)
- require.Equal(t, 11.0, cost1)
- cost2, ok2 := windowCostFromPrefetchContext(outCtx, 2)
- require.True(t, ok2)
- require.Equal(t, 22.0, cost2)
- _, ok3 := windowCostFromPrefetchContext(outCtx, 3)
- require.False(t, ok3)
- require.Equal(t, int64(1), repo.batchCalls.Load())
- require.Equal(t, 22.0, cache.setData[2])
- hit, miss, batchSQL, fallback, errCount := GatewayWindowCostPrefetchStats()
- require.Equal(t, int64(1), hit)
- require.Equal(t, int64(1), miss)
- require.Equal(t, int64(1), batchSQL)
- require.Equal(t, int64(0), fallback)
- require.Equal(t, int64(0), errCount)
- }
- func TestWithWindowCostPrefetch_AllHitNoSQL(t *testing.T) {
- resetGatewayHotpathStatsForTest()
- windowStart := time.Now().Add(-30 * time.Minute).Truncate(time.Hour)
- windowEnd := windowStart.Add(5 * time.Hour)
- accounts := []Account{
- {
- ID: 1,
- Platform: PlatformAnthropic,
- Type: AccountTypeOAuth,
- Extra: map[string]any{"window_cost_limit": 100.0},
- SessionWindowStart: &windowStart,
- SessionWindowEnd: &windowEnd,
- },
- {
- ID: 2,
- Platform: PlatformAnthropic,
- Type: AccountTypeSetupToken,
- Extra: map[string]any{"window_cost_limit": 100.0},
- SessionWindowStart: &windowStart,
- SessionWindowEnd: &windowEnd,
- },
- }
- cache := &sessionLimitCacheHotpathStub{
- batchData: map[int64]float64{
- 1: 11.0,
- 2: 22.0,
- },
- }
- repo := &usageLogWindowBatchRepoStub{}
- svc := &GatewayService{
- sessionLimitCache: cache,
- usageLogRepo: repo,
- }
- outCtx := svc.withWindowCostPrefetch(context.Background(), accounts)
- cost1, ok1 := windowCostFromPrefetchContext(outCtx, 1)
- cost2, ok2 := windowCostFromPrefetchContext(outCtx, 2)
- require.True(t, ok1)
- require.True(t, ok2)
- require.Equal(t, 11.0, cost1)
- require.Equal(t, 22.0, cost2)
- require.Equal(t, int64(0), repo.batchCalls.Load())
- require.Equal(t, int64(0), repo.singleCalls.Load())
- hit, miss, batchSQL, fallback, errCount := GatewayWindowCostPrefetchStats()
- require.Equal(t, int64(2), hit)
- require.Equal(t, int64(0), miss)
- require.Equal(t, int64(0), batchSQL)
- require.Equal(t, int64(0), fallback)
- require.Equal(t, int64(0), errCount)
- }
- func TestWithWindowCostPrefetch_BatchErrorFallbackSingleQuery(t *testing.T) {
- resetGatewayHotpathStatsForTest()
- windowStart := time.Now().Add(-30 * time.Minute).Truncate(time.Hour)
- windowEnd := windowStart.Add(5 * time.Hour)
- accounts := []Account{
- {
- ID: 2,
- Platform: PlatformAnthropic,
- Type: AccountTypeSetupToken,
- Extra: map[string]any{"window_cost_limit": 100.0},
- SessionWindowStart: &windowStart,
- SessionWindowEnd: &windowEnd,
- },
- }
- cache := &sessionLimitCacheHotpathStub{}
- repo := &usageLogWindowBatchRepoStub{
- batchErr: errors.New("batch failed"),
- singleResult: map[int64]*usagestats.AccountStats{
- 2: {StandardCost: 33.0},
- },
- }
- svc := &GatewayService{
- sessionLimitCache: cache,
- usageLogRepo: repo,
- }
- outCtx := svc.withWindowCostPrefetch(context.Background(), accounts)
- cost, ok := windowCostFromPrefetchContext(outCtx, 2)
- require.True(t, ok)
- require.Equal(t, 33.0, cost)
- require.Equal(t, int64(1), repo.batchCalls.Load())
- require.Equal(t, int64(1), repo.singleCalls.Load())
- _, _, _, fallback, errCount := GatewayWindowCostPrefetchStats()
- require.Equal(t, int64(1), fallback)
- require.Equal(t, int64(1), errCount)
- }
- func TestGetAvailableModels_UsesShortCacheAndSupportsInvalidation(t *testing.T) {
- resetGatewayHotpathStatsForTest()
- groupID := int64(9)
- repo := &modelsListAccountRepoStub{
- byGroup: map[int64][]Account{
- groupID: {
- {
- ID: 1,
- Platform: PlatformAnthropic,
- Credentials: map[string]any{
- "model_mapping": map[string]any{
- "claude-3-5-sonnet": "claude-3-5-sonnet",
- "claude-3-5-haiku": "claude-3-5-haiku",
- },
- },
- },
- {
- ID: 2,
- Platform: PlatformGemini,
- Credentials: map[string]any{
- "model_mapping": map[string]any{
- "gemini-2.5-pro": "gemini-2.5-pro",
- },
- },
- },
- },
- },
- }
- svc := &GatewayService{
- accountRepo: repo,
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- models1 := svc.GetAvailableModels(context.Background(), &groupID, PlatformAnthropic)
- require.Equal(t, []string{"claude-3-5-haiku", "claude-3-5-sonnet"}, models1)
- require.Equal(t, int64(1), repo.listByGroupCalls.Load())
- // TTL 内再次请求应命中缓存,不回源。
- models2 := svc.GetAvailableModels(context.Background(), &groupID, PlatformAnthropic)
- require.Equal(t, models1, models2)
- require.Equal(t, int64(1), repo.listByGroupCalls.Load())
- // 更新仓储数据,但缓存未失效前应继续返回旧值。
- repo.byGroup[groupID] = []Account{
- {
- ID: 3,
- Platform: PlatformAnthropic,
- Credentials: map[string]any{
- "model_mapping": map[string]any{
- "claude-3-7-sonnet": "claude-3-7-sonnet",
- },
- },
- },
- }
- models3 := svc.GetAvailableModels(context.Background(), &groupID, PlatformAnthropic)
- require.Equal(t, []string{"claude-3-5-haiku", "claude-3-5-sonnet"}, models3)
- require.Equal(t, int64(1), repo.listByGroupCalls.Load())
- svc.InvalidateAvailableModelsCache(&groupID, PlatformAnthropic)
- models4 := svc.GetAvailableModels(context.Background(), &groupID, PlatformAnthropic)
- require.Equal(t, []string{"claude-3-7-sonnet"}, models4)
- require.Equal(t, int64(2), repo.listByGroupCalls.Load())
- hit, miss, store := GatewayModelsListCacheStats()
- require.Equal(t, int64(2), hit)
- require.Equal(t, int64(2), miss)
- require.Equal(t, int64(2), store)
- }
- // Scenario: 账号模型变更会失效所属平台缓存
- func TestResolveCompositeModelOwnershipUsesModelsCacheInvalidation(t *testing.T) {
- groupID := int64(9)
- repo := &modelsListAccountRepoStub{
- byGroup: map[int64][]Account{
- groupID: {{
- ID: 1,
- Platform: PlatformDeepseek,
- Credentials: map[string]any{"model_mapping": map[string]any{"company-model": "deepseek-v4-pro"}},
- }},
- },
- }
- svc := &GatewayService{
- accountRepo: repo,
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- first, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "company-model")
- require.NoError(t, err)
- require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformDeepseek, Matched: true}, first)
- require.Equal(t, int64(1), repo.listByGroupCalls.Load())
- repo.byGroup[groupID] = []Account{{
- ID: 2,
- Platform: PlatformOpenAI,
- Credentials: map[string]any{"model_mapping": map[string]any{"company-model": "gpt-5"}},
- }}
- cached, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "company-model")
- require.NoError(t, err)
- require.Equal(t, first, cached)
- require.Equal(t, int64(1), repo.listByGroupCalls.Load())
- svc.InvalidateAvailableModelsCache(&groupID, PlatformDeepseek)
- refreshed, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "company-model")
- require.NoError(t, err)
- require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, refreshed)
- require.Equal(t, int64(2), repo.listByGroupCalls.Load())
- }
- func TestGetAvailableModels_ErrorAndGlobalListBranches(t *testing.T) {
- resetGatewayHotpathStatsForTest()
- errRepo := &modelsListAccountRepoStub{
- err: errors.New("db error"),
- }
- svcErr := &GatewayService{
- accountRepo: errRepo,
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- require.Nil(t, svcErr.GetAvailableModels(context.Background(), nil, ""))
- okRepo := &modelsListAccountRepoStub{
- all: []Account{
- {
- ID: 1,
- Platform: PlatformAnthropic,
- Credentials: map[string]any{
- "model_mapping": map[string]any{
- "claude-3-5-sonnet": "claude-3-5-sonnet",
- },
- },
- },
- {
- ID: 2,
- Platform: PlatformGemini,
- Credentials: map[string]any{
- "model_mapping": map[string]any{
- "gemini-2.5-pro": "gemini-2.5-pro",
- },
- },
- },
- },
- }
- svcOK := &GatewayService{
- accountRepo: okRepo,
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- models := svcOK.GetAvailableModels(context.Background(), nil, "")
- require.Equal(t, []string{"claude-3-5-sonnet", "gemini-2.5-pro"}, models)
- require.Equal(t, int64(1), okRepo.listAllCalls.Load())
- }
- func TestGetAvailableModels_OpenAIPassthroughUsesDefaultFallback(t *testing.T) {
- groupID := int64(10)
- tests := []struct {
- name string
- accounts []Account
- want []string
- }{
- {
- name: "passthrough only ignores stale mapping",
- accounts: []Account{
- {
- ID: 1,
- Platform: PlatformOpenAI,
- Credentials: map[string]any{"model_mapping": map[string]any{"stale-model": "upstream-model"}},
- Extra: map[string]any{"openai_passthrough": true},
- },
- },
- want: nil,
- },
- {
- name: "passthrough wins over ordinary account mapping",
- accounts: []Account{
- {
- ID: 2,
- Platform: PlatformOpenAI,
- Credentials: map[string]any{"model_mapping": map[string]any{"configured-model": "configured-upstream"}},
- },
- {
- ID: 3,
- Platform: PlatformOpenAI,
- Credentials: map[string]any{"model_mapping": map[string]any{"stale-model": "upstream-model"}},
- Extra: map[string]any{"openai_passthrough": true},
- },
- },
- want: nil,
- },
- {
- name: "ordinary accounts preserve mapped whitelist",
- accounts: []Account{
- {
- ID: 4,
- Platform: PlatformOpenAI,
- Credentials: map[string]any{"model_mapping": map[string]any{"configured-model": "configured-upstream"}},
- },
- },
- want: []string{"configured-model"},
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- repo := &modelsListAccountRepoStub{byGroup: map[int64][]Account{groupID: tt.accounts}}
- svc := &GatewayService{
- accountRepo: repo,
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- require.Equal(t, tt.want, svc.GetAvailableModels(context.Background(), &groupID, PlatformOpenAI))
- })
- }
- }
- func TestGetAvailableModels_GlobalListPreservesMappedModelsWithOpenAIPassthrough(t *testing.T) {
- groupID := int64(11)
- repo := &modelsListAccountRepoStub{
- byGroup: map[int64][]Account{
- groupID: {
- {
- ID: 1,
- Platform: PlatformOpenAI,
- Extra: map[string]any{"openai_passthrough": true},
- },
- {
- ID: 2,
- Platform: PlatformAnthropic,
- Credentials: map[string]any{"model_mapping": map[string]any{"claude-mapped": "claude-upstream"}},
- },
- },
- },
- }
- svc := &GatewayService{
- accountRepo: repo,
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- require.Equal(t, []string{"claude-mapped"}, svc.GetAvailableModels(context.Background(), &groupID, ""))
- }
- func TestGatewayHotpathHelpers_CacheTTLAndStickyContext(t *testing.T) {
- t.Run("resolve_user_group_rate_cache_ttl", func(t *testing.T) {
- require.Equal(t, defaultUserGroupRateCacheTTL, resolveUserGroupRateCacheTTL(nil))
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- UserGroupRateCacheTTLSeconds: 45,
- },
- }
- require.Equal(t, 45*time.Second, resolveUserGroupRateCacheTTL(cfg))
- })
- t.Run("resolve_models_list_cache_ttl", func(t *testing.T) {
- require.Equal(t, defaultModelsListCacheTTL, resolveModelsListCacheTTL(nil))
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- ModelsListCacheTTLSeconds: 20,
- },
- }
- require.Equal(t, 20*time.Second, resolveModelsListCacheTTL(cfg))
- })
- t.Run("prefetched_sticky_account_id_from_context", func(t *testing.T) {
- require.Equal(t, int64(0), prefetchedStickyAccountIDFromContext(context.TODO(), nil))
- require.Equal(t, int64(0), prefetchedStickyAccountIDFromContext(context.Background(), nil))
- ctx := context.WithValue(context.Background(), ctxkey.PrefetchedStickyAccountID, int64(123))
- ctx = context.WithValue(ctx, ctxkey.PrefetchedStickyGroupID, int64(0))
- require.Equal(t, int64(123), prefetchedStickyAccountIDFromContext(ctx, nil))
- groupID := int64(9)
- ctx2 := context.WithValue(context.Background(), ctxkey.PrefetchedStickyAccountID, 456)
- ctx2 = context.WithValue(ctx2, ctxkey.PrefetchedStickyGroupID, groupID)
- require.Equal(t, int64(456), prefetchedStickyAccountIDFromContext(ctx2, &groupID))
- ctx3 := context.WithValue(context.Background(), ctxkey.PrefetchedStickyAccountID, "invalid")
- ctx3 = context.WithValue(ctx3, ctxkey.PrefetchedStickyGroupID, groupID)
- require.Equal(t, int64(0), prefetchedStickyAccountIDFromContext(ctx3, &groupID))
- ctx4 := context.WithValue(context.Background(), ctxkey.PrefetchedStickyAccountID, int64(789))
- ctx4 = context.WithValue(ctx4, ctxkey.PrefetchedStickyGroupID, int64(10))
- require.Equal(t, int64(0), prefetchedStickyAccountIDFromContext(ctx4, &groupID))
- })
- t.Run("window_cost_from_prefetch_context", func(t *testing.T) {
- require.Equal(t, false, func() bool {
- _, ok := windowCostFromPrefetchContext(context.TODO(), 0)
- return ok
- }())
- require.Equal(t, false, func() bool {
- _, ok := windowCostFromPrefetchContext(context.Background(), 1)
- return ok
- }())
- ctx := context.WithValue(context.Background(), windowCostPrefetchContextKey, map[int64]float64{
- 9: 12.34,
- })
- cost, ok := windowCostFromPrefetchContext(ctx, 9)
- require.True(t, ok)
- require.Equal(t, 12.34, cost)
- })
- }
- func TestInvalidateAvailableModelsCache_ByDimensions(t *testing.T) {
- svc := &GatewayService{
- modelsListCache: gocache.New(time.Minute, time.Minute),
- }
- group9 := int64(9)
- group10 := int64(10)
- svc.modelsListCache.Set(modelsListCacheKey(&group9, PlatformAnthropic), []string{"a"}, time.Minute)
- svc.modelsListCache.Set(modelsListCacheKey(&group9, PlatformGemini), []string{"b"}, time.Minute)
- svc.modelsListCache.Set(modelsListCacheKey(&group10, PlatformAnthropic), []string{"c"}, time.Minute)
- svc.modelsListCache.Set("invalid-key", []string{"d"}, time.Minute)
- t.Run("invalidate_group_and_platform", func(t *testing.T) {
- svc.InvalidateAvailableModelsCache(&group9, PlatformAnthropic)
- _, found := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformAnthropic))
- require.False(t, found)
- _, stillFound := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformGemini))
- require.True(t, stillFound)
- })
- t.Run("invalidate_group_only", func(t *testing.T) {
- svc.InvalidateAvailableModelsCache(&group9, "")
- _, foundA := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformAnthropic))
- _, foundB := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformGemini))
- require.False(t, foundA)
- require.False(t, foundB)
- _, foundOtherGroup := svc.modelsListCache.Get(modelsListCacheKey(&group10, PlatformAnthropic))
- require.True(t, foundOtherGroup)
- })
- t.Run("invalidate_platform_only", func(t *testing.T) {
- // 重建数据后仅按 platform 失效
- svc.modelsListCache.Set(modelsListCacheKey(&group9, PlatformAnthropic), []string{"a"}, time.Minute)
- svc.modelsListCache.Set(modelsListCacheKey(&group9, PlatformGemini), []string{"b"}, time.Minute)
- svc.modelsListCache.Set(modelsListCacheKey(&group10, PlatformAnthropic), []string{"c"}, time.Minute)
- svc.InvalidateAvailableModelsCache(nil, PlatformAnthropic)
- _, found9Anthropic := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformAnthropic))
- _, found10Anthropic := svc.modelsListCache.Get(modelsListCacheKey(&group10, PlatformAnthropic))
- _, found9Gemini := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformGemini))
- require.False(t, found9Anthropic)
- require.False(t, found10Anthropic)
- require.True(t, found9Gemini)
- })
- }
- func TestSelectAccountWithLoadAwareness_StickyReadReuse(t *testing.T) {
- now := time.Now().Add(-time.Minute)
- account := Account{
- ID: 88,
- Platform: PlatformAnthropic,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 4,
- Priority: 1,
- LastUsedAt: &now,
- }
- repo := stubOpenAIAccountRepo{accounts: []Account{account}}
- concurrency := NewConcurrencyService(stubConcurrencyCache{})
- cfg := &config.Config{
- RunMode: config.RunModeStandard,
- Gateway: config.GatewayConfig{
- Scheduling: config.GatewaySchedulingConfig{
- LoadBatchEnabled: true,
- StickySessionMaxWaiting: 3,
- StickySessionWaitTimeout: time.Second,
- FallbackWaitTimeout: time.Second,
- FallbackMaxWaiting: 10,
- },
- },
- }
- baseCtx := context.WithValue(context.Background(), ctxkey.ForcePlatform, PlatformAnthropic)
- t.Run("without_prefetch_reads_cache_once", func(t *testing.T) {
- cache := &stickyGatewayCacheHotpathStub{stickyID: account.ID}
- svc := &GatewayService{
- accountRepo: repo,
- cache: cache,
- cfg: cfg,
- concurrencyService: concurrency,
- userGroupRateCache: gocache.New(time.Minute, time.Minute),
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- result, err := svc.SelectAccountWithLoadAwareness(baseCtx, nil, "sess-hash", "", nil, "", int64(0))
- require.NoError(t, err)
- require.NotNil(t, result)
- require.NotNil(t, result.Account)
- require.Equal(t, account.ID, result.Account.ID)
- require.Equal(t, int64(1), cache.getCalls.Load())
- })
- t.Run("with_prefetch_skips_cache_read", func(t *testing.T) {
- cache := &stickyGatewayCacheHotpathStub{stickyID: account.ID}
- svc := &GatewayService{
- accountRepo: repo,
- cache: cache,
- cfg: cfg,
- concurrencyService: concurrency,
- userGroupRateCache: gocache.New(time.Minute, time.Minute),
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- ctx := context.WithValue(baseCtx, ctxkey.PrefetchedStickyAccountID, account.ID)
- ctx = context.WithValue(ctx, ctxkey.PrefetchedStickyGroupID, int64(0))
- result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sess-hash", "", nil, "", int64(0))
- require.NoError(t, err)
- require.NotNil(t, result)
- require.NotNil(t, result.Account)
- require.Equal(t, account.ID, result.Account.ID)
- require.Equal(t, int64(0), cache.getCalls.Load())
- })
- t.Run("with_prefetch_group_mismatch_reads_cache", func(t *testing.T) {
- cache := &stickyGatewayCacheHotpathStub{stickyID: account.ID}
- svc := &GatewayService{
- accountRepo: repo,
- cache: cache,
- cfg: cfg,
- concurrencyService: concurrency,
- userGroupRateCache: gocache.New(time.Minute, time.Minute),
- modelsListCache: gocache.New(time.Minute, time.Minute),
- modelsListCacheTTL: time.Minute,
- }
- ctx := context.WithValue(baseCtx, ctxkey.PrefetchedStickyAccountID, int64(999))
- ctx = context.WithValue(ctx, ctxkey.PrefetchedStickyGroupID, int64(77))
- result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sess-hash", "", nil, "", int64(0))
- require.NoError(t, err)
- require.NotNil(t, result)
- require.NotNil(t, result.Account)
- require.Equal(t, account.ID, result.Account.ID)
- require.Equal(t, int64(1), cache.getCalls.Load())
- })
- }
|