Codex
aa6c3f50f4
Fixed-group tokens with a backup group now take precedence over the session affinity skip-retry-on-failure rule, so retryable upstream errors (503) fall through to the backup group instead of stopping on the primary channel. Non-retryable errors, explicit channel pins, and auto groups are unchanged. Adds controller regression tests for the affinity override and a build-time Go test gate.
228 lines
8.2 KiB
Go
228 lines
8.2 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/setting"
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func setupChannelSelectAutoGroupsTest(t *testing.T) *gorm.DB {
|
|
t.Helper()
|
|
|
|
originalDB := model.DB
|
|
originalMemoryCacheEnabled := common.MemoryCacheEnabled
|
|
originalRetryTimes := common.RetryTimes
|
|
originalAutoGroups := setting.AutoGroups2JsonString()
|
|
originalUsableGroups := setting.UserUsableGroups2JSONString()
|
|
originalGroupRatios := ratio_setting.GroupRatio2JSONString()
|
|
originalMaxTokenAutoGroups := setting.GetMaxTokenAutoGroups()
|
|
|
|
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
|
|
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}))
|
|
model.DB = db
|
|
common.MemoryCacheEnabled = true
|
|
common.RetryTimes = 0
|
|
|
|
require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`))
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2}`))
|
|
require.NoError(t, setting.UpdateMaxTokenAutoGroups("2"))
|
|
|
|
t.Cleanup(func() {
|
|
model.DB = originalDB
|
|
common.MemoryCacheEnabled = originalMemoryCacheEnabled
|
|
common.RetryTimes = originalRetryTimes
|
|
require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups))
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatios))
|
|
require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMaxTokenAutoGroups)))
|
|
|
|
if originalMemoryCacheEnabled && originalDB != nil &&
|
|
originalDB.Migrator().HasTable(&model.Channel{}) && originalDB.Migrator().HasTable(&model.Ability{}) {
|
|
model.InitChannelCache()
|
|
}
|
|
sqlDB, err := db.DB()
|
|
if err == nil {
|
|
require.NoError(t, sqlDB.Close())
|
|
}
|
|
})
|
|
|
|
return db
|
|
}
|
|
|
|
func createChannelSelectAutoGroupsChannel(t *testing.T, db *gorm.DB, id int, group, modelName string) {
|
|
t.Helper()
|
|
priority := int64(0)
|
|
weight := uint(100)
|
|
require.NoError(t, db.Create(&model.Channel{
|
|
Id: id,
|
|
Type: constant.ChannelTypeOpenAI,
|
|
Key: fmt.Sprintf("key-%d", id),
|
|
Status: common.ChannelStatusEnabled,
|
|
Name: fmt.Sprintf("channel-%d", id),
|
|
Weight: &weight,
|
|
Models: modelName,
|
|
Group: group,
|
|
Priority: &priority,
|
|
}).Error)
|
|
require.NoError(t, db.Create(&model.Ability{
|
|
Group: group,
|
|
Model: modelName,
|
|
ChannelId: id,
|
|
Enabled: true,
|
|
Priority: &priority,
|
|
Weight: weight,
|
|
}).Error)
|
|
}
|
|
|
|
func TestCacheGetRandomSatisfiedChannelUsesTokenAutoGroupsWhenGlobalAutoIsEmpty(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
const modelName = "auto-groups-runtime-model"
|
|
createChannelSelectAutoGroupsChannel(t, db, 2101, "vip", modelName)
|
|
createChannelSelectAutoGroupsChannel(t, db, 2102, "default", modelName)
|
|
model.InitChannelCache()
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip", "default"})
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenCrossGroupRetry, true)
|
|
|
|
retry := 0
|
|
param := &RetryParam{
|
|
Ctx: ctx,
|
|
TokenGroup: "auto",
|
|
ModelName: modelName,
|
|
RequestPath: "/v1/chat/completions",
|
|
Retry: &retry,
|
|
}
|
|
|
|
first, selectedGroup, err := CacheGetRandomSatisfiedChannel(param)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, first)
|
|
assert.Equal(t, 2101, first.Id)
|
|
assert.Equal(t, "vip", selectedGroup)
|
|
assert.Equal(t, "vip", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup))
|
|
assert.Empty(t, setting.GetAutoGroups(), "the selection must not depend on the global Auto list")
|
|
|
|
param.IncreaseRetry()
|
|
second, selectedGroup, err := CacheGetRandomSatisfiedChannel(param)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, second)
|
|
assert.Equal(t, 2102, second.Id)
|
|
assert.Equal(t, "default", selectedGroup)
|
|
assert.Equal(t, "default", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup))
|
|
}
|
|
|
|
func TestCacheGetRandomSatisfiedChannelUsesFixedTokenBackupGroup(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
const modelName = "fixed-backup-runtime-model"
|
|
createChannelSelectAutoGroupsChannel(t, db, 2201, "default", modelName)
|
|
createChannelSelectAutoGroupsChannel(t, db, 2202, "vip", modelName)
|
|
model.InitChannelCache()
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"})
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenCrossGroupRetry, true)
|
|
|
|
retry := 0
|
|
param := &RetryParam{
|
|
Ctx: ctx,
|
|
TokenGroup: "default",
|
|
ModelName: modelName,
|
|
RequestPath: "/v1/chat/completions",
|
|
Retry: &retry,
|
|
}
|
|
|
|
primary, selectedGroup, err := CacheGetRandomSatisfiedChannel(param)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, primary)
|
|
assert.Equal(t, 2201, primary.Id)
|
|
assert.Equal(t, "default", selectedGroup)
|
|
assert.Equal(t, "default", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup))
|
|
|
|
param.IncreaseRetry()
|
|
backup, selectedGroup, err := CacheGetRandomSatisfiedChannel(param)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, backup)
|
|
assert.Equal(t, 2202, backup.Id)
|
|
assert.Equal(t, "vip", selectedGroup)
|
|
assert.Equal(t, "vip", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup))
|
|
}
|
|
|
|
func TestCacheGetRandomSatisfiedChannelSkipsUnavailableFixedPrimary(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
const modelName = "fixed-backup-no-primary-model"
|
|
createChannelSelectAutoGroupsChannel(t, db, 2302, "vip", modelName)
|
|
model.InitChannelCache()
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"})
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenCrossGroupRetry, true)
|
|
|
|
retry := 0
|
|
channel, selectedGroup, err := CacheGetRandomSatisfiedChannel(&RetryParam{
|
|
Ctx: ctx,
|
|
TokenGroup: "default",
|
|
ModelName: modelName,
|
|
RequestPath: "/v1/responses",
|
|
Retry: &retry,
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, channel)
|
|
assert.Equal(t, 2302, channel.Id)
|
|
assert.Equal(t, "vip", selectedGroup)
|
|
assert.Equal(t, "vip", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup))
|
|
}
|
|
|
|
func TestFixedBackupAfterAffinityFailureWithOneRetry(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
common.RetryTimes = 1
|
|
const modelName = "fixed-backup-affinity"
|
|
createChannelSelectAutoGroupsChannel(t, db, 13, "default", modelName)
|
|
createChannelSelectAutoGroupsChannel(t, db, 14, "vip", modelName)
|
|
model.InitChannelCache()
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "default")
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"})
|
|
common.SetContextKey(ctx, constant.ContextKeyTokenCrossGroupRetry, true)
|
|
ctx.Set(ginKeyChannelAffinitySkipRetry, true)
|
|
require.False(t, ShouldSkipRetryAfterChannelAffinityFailure(ctx))
|
|
|
|
// The distributor's affinity hit bypasses the ordered group selector.
|
|
param := &RetryParam{Ctx: ctx, TokenGroup: "default", ModelName: modelName}
|
|
param.SetRetry(1)
|
|
first, _, err := CacheGetRandomSatisfiedChannel(param)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, first)
|
|
require.Equal(t, 13, first.Id)
|
|
require.Equal(t, 0, param.GetRetry())
|
|
param.IncreaseRetry()
|
|
backup, group, err := CacheGetRandomSatisfiedChannel(param)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, backup)
|
|
require.Equal(t, 14, backup.Id)
|
|
require.Equal(t, "vip", group)
|
|
require.Equal(t, "vip", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup))
|
|
}
|