cneicy/new-api-custom
1
0 Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

feat: publish sanitized custom.97 source

Browse Source
This commit is contained in:
cneicy committed 2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
1022 files changed
+69903 -19207
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

service/token_routing_switch_test.go
+116 -2
View File
@@ -8,14 +8,17 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
for _, change := range []string{"priority", "model_rule"} {
for _, change := range []string{"priority", "model_rule", "blacklist"} {
t.Run(change, func(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
const modelName = "manual-switch-model"
@@ -45,10 +48,17 @@ func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
c.Set("token_id", 100)
c.Set("original_model", modelName)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}}
if change == "blacklist" {
policy.Mode = "smart"
policy.Groups = nil
automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
}
if switched {
if change == "priority" {
policy.Groups[0].Priority = 1
policy.Groups[1].Priority = 0
} else if change == "blacklist" {
policy.ExcludedGroups = []string{"default"}
} else {
policy.ModelRules = []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}
}
@@ -76,7 +86,7 @@ func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
log := model.Log{UserId: 100, TokenId: 100, ModelName: modelName, CorrelationKey: old.GetString("log_correlation_key"), Type: model.LogTypeConsume, CreatedAt: time.Now().Unix(), ChannelId: 6201, Group: "default", Other: string(body)}
require.NoError(t, db.Create(&log).Error)
// The old group is still eligible: only its priority/preference changed.
// The edited policy must take precedence over the existing binding.
switched := ctx(true)
selected, group, selectErr := SelectChannelForRequest(switched, modelName, &RetryParam{Ctx: switched, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)})
require.Nil(t, selectErr)
@@ -114,3 +124,107 @@ func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
})
}
}
func TestAutomaticPoolKeepsHealthySession(t *testing.T) {
for _, scenario := range []string{"cheaper_recovered", "ranking_changed", "candidate_added", "degraded", "cooldown", "permission_removed", "channel_disabled", "policy_changed", "new_session", "session_off", "custom_recovery"} {
t.Run(scenario, func(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
const modelName = "smart-sticky-model"
createChannelSelectAutoGroupsChannel(t, db, 6301, "default", modelName)
createChannelSelectAutoGroupsChannel(t, db, 6302, "vip", modelName)
model.InitChannelCache()
oldRedis, oldLogDB := common.RedisEnabled, model.LOG_DB
common.RedisEnabled, model.LOG_DB = false, nil
config := operation_setting.GetChannelAffinitySetting()
oldConfig := *config
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "prefer", Rules: []operation_setting.ChannelAffinityRule{{Name: "smart-sticky", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "prefer"}}}
require.NoError(t, poolBindings().Purge())
require.NoError(t, poolCooldowns().Purge())
automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
t.Cleanup(func() {
*config = oldConfig
common.RedisEnabled, model.LOG_DB = oldRedis, oldLogDB
_ = poolBindings().Purge()
_ = poolCooldowns().Purge()
automaticMetricCache.Purge()
})
ctx := func() *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Request.Header.Set("Session-Id", "same-session")
c.Set("id", 101)
c.Set("token_id", 102)
c.Set("original_model", modelName)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "cost", SkipDegradedNonAstra: true}
if scenario == "custom_recovery" {
policy.Mode = "custom"
policy.Groups = []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}
}
c.Set(model.TokenRoutingContextKey, policy)
c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{})
c.Set("pool_group_visions", map[string]DetectionHistory{})
return c
}
// A temporary failure sends the session to the more expensive group.
first := ctx()
cheapKey := poolCooldownKey(first, modelName, "default")
require.NoError(t, poolCooldowns().SetWithTTL(cheapKey, poolCooldown{Until: time.Now().Add(time.Minute).UnixMilli()}, 0))
if scenario == "candidate_added" {
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"vip":"VIP"}`))
}
ch, group, selectErr := SelectChannelForRequest(first, modelName, &RetryParam{Ctx: first, ModelName: modelName, Retry: common.GetPointer(0)})
require.Nil(t, selectErr)
require.NotNil(t, ch)
require.Equal(t, "vip", group)
first.Set("channel_id", ch.Id)
RequestPolicy(first).Successful = true
RecordPoolAffinity(first)
require.NoError(t, poolCooldowns().SetWithTTL(cheapKey, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
next := ctx()
want := "vip"
switch scenario {
case "ranking_changed":
// The initial pool already placed default first. Reverse it to prove ranking changes are irrelevant.
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":4,"vip":3}`))
case "candidate_added":
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
case "degraded":
next.Set("pool_group_visions", map[string]DetectionHistory{"vip": {Quality: "degraded"}})
want = "default"
case "cooldown":
require.NoError(t, poolCooldowns().SetWithTTL(poolCooldownKey(next, modelName, "vip"), poolCooldown{Until: time.Now().Add(time.Minute).UnixMilli()}, 0))
want = "default"
case "permission_removed":
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
want = "default"
case "channel_disabled":
require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", 6302).Update("status", common.ChannelStatusAutoDisabled).Error)
model.InitChannelCache()
want = "default"
case "policy_changed":
RequestRoutingPolicy(next).AutoPreference = "stable"
want = "default"
case "new_session":
next.Request.Header.Set("Session-Id", "new-session")
want = "default"
case "session_off":
config.Rules[0].SessionMode = "off"
config.SessionMode = "off"
want = "default"
case "custom_recovery":
want = "default"
}
selected, selectedGroup, selectErr := SelectChannelForRequest(next, modelName, &RetryParam{Ctx: next, ModelName: modelName, Retry: common.GetPointer(0)})
require.Nil(t, selectErr)
require.NotNil(t, selected)
assert.Equal(t, want, selectedGroup)
if want == "vip" {
assert.Equal(t, 6302, selected.Id)
state, found, err := poolCooldowns().Get(cheapKey)
require.NoError(t, err)
require.True(t, found)
assert.Empty(t, state.ProbeID, "keeping a healthy route must not claim a cheaper group's recovery probe")
}
})
}
}