mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Composite model catalogs expose exact account model_mapping aliases, but request routing only considered explicit routes or model-name detection. Unprefixed aliases were selectable but unroutable. Resolve unique exact account ownership before the detector, keep explicit routes authoritative, and fail closed when an alias is claimed across platforms.
947 lines
29 KiB
Go
947 lines
29 KiB
Go
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())
|
|
})
|
|
}
|