new-api-custom/service/token_routing_switch_test.go

package service

import (
	"net/http/httptest"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/setting/operation_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"} {
		t.Run(change, func(t *testing.T) {
			db := setupChannelSelectAutoGroupsTest(t)
			const modelName = "manual-switch-model"
			createChannelSelectAutoGroupsChannel(t, db, 6201, "default", modelName)
			createChannelSelectAutoGroupsChannel(t, db, 6202, "vip", modelName)
			model.InitChannelCache()
			ResetGroupAvailability()
			t.Cleanup(ResetGroupAvailability)
			previousRedis := common.RedisEnabled
			common.RedisEnabled = false
			t.Cleanup(func() { common.RedisEnabled = previousRedis })
			require.NoError(t, poolBindings().Purge())
			t.Cleanup(func() { _ = poolBindings().Purge() })
			config := operation_setting.GetChannelAffinitySetting()
			previousConfig := *config
			*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "switch-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
			t.Cleanup(func() { *config = previousConfig })
			previousLogDB := model.LOG_DB
			model.LOG_DB = db
			t.Cleanup(func() { model.LOG_DB = previousLogDB })
			require.NoError(t, db.AutoMigrate(&model.Log{}))
			ctx := func(switched bool) *gin.Context {
				c := setupGroupPoolContext(t)
				c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
				c.Request.Header.Set("Session-Id", "same-session")
				c.Set("id", 100)
				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 switched {
					if change == "priority" {
						policy.Groups[0].Priority = 1
						policy.Groups[1].Priority = 0
					} else {
						policy.ModelRules = []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}
					}
				}
				c.Set(model.TokenRoutingContextKey, policy)
				return c
			}
			old := ctx(false)
			ch, _, selectErr := SelectChannelForRequest(old, modelName, &RetryParam{Ctx: old, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)})
			require.Nil(t, selectErr)
			require.NotNil(t, ch)
			assert.Equal(t, 6201, ch.Id)
			old.Set("channel_id", ch.Id)
			common.SetContextKey(old, constant.ContextKeyAutoGroup, "default")
			RequestPolicy(old).Successful = true
			RecordPoolAffinity(old)
			unchanged, _ := PreferredPoolChannel(ctx(false), modelName)
			require.NotNil(t, unchanged)
			assert.Equal(t, 6201, unchanged.Id)

			details := model.NewLogOther()
			appendRequestDetails(old, nil, details)
			body, err := common.Marshal(details.Snapshot())
			require.NoError(t, err)
			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.
			switched := ctx(true)
			selected, group, selectErr := SelectChannelForRequest(switched, modelName, &RetryParam{Ctx: switched, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)})
			require.Nil(t, selectErr)
			require.NotNil(t, selected)
			assert.Equal(t, 6202, selected.Id)
			assert.Equal(t, "vip", group)
			switched.Set("channel_id", selected.Id)
			common.SetContextKey(switched, constant.ContextKeyAutoGroup, group)
			RequestPolicy(switched).Successful = true
			RecordPoolAffinity(switched)
			unchanged, _ = PreferredPoolChannel(ctx(true), modelName)
			require.NotNil(t, unchanged)
			assert.Equal(t, 6202, unchanged.Id)

			// A request started before the edit may finish after the new one.
			RecordPoolAffinity(old)
			stale, _ := PreferredPoolChannel(ctx(true), modelName)
			assert.Nil(t, stale)

			// Cache expiry/restart must not resurrect the old successful log.
			require.NoError(t, poolBindings().Purge())
			stale, _ = PreferredPoolChannel(ctx(true), modelName)
			assert.Nil(t, stale)

			// A matching historical policy remains recoverable after cache expiry.
			require.NoError(t, poolBindings().Purge())
			recovered, _ := PreferredPoolChannel(ctx(false), modelName)
			require.NotNil(t, recovered)
			assert.Equal(t, 6201, recovered.Id)

			// Legacy bindings carry no policy and must be selected afresh once.
			require.NoError(t, poolBindings().SetWithTTL(old.GetString("pool_affinity_key"), poolBinding{Channel: 6201, Group: "default"}, time.Hour))
			legacy, _ := PreferredPoolChannel(ctx(false), modelName)
			assert.Nil(t, legacy)
		})
	}
}