feat: sync deployed v1.0.0-rc.38-custom.42

Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.

service/token_routing_test.go

@@ -0,0 +1,69 @@
+package service
+
+import (
+	"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/require"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestTokenRoutingPoolPreferenceExclusionAndAffinityIsolation(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	const modelName = "pool-model"
+	createChannelSelectAutoGroupsChannel(t, db, 6101, "default", modelName)
+	createChannelSelectAutoGroupsChannel(t, db, 6102, "vip", modelName)
+	model.InitChannelCache()
+	ResetGroupAvailability()
+	t.Cleanup(ResetGroupAvailability)
+	config := operation_setting.GetChannelAffinitySetting()
+	previous := *config
+	*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "prefer", Rules: []operation_setting.ChannelAffinityRule{{Name: "pool-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "prefer"}}}
+	t.Cleanup(func() { *config = previous; _ = poolBindings().Purge() })
+	ctx := func(token int, session string) *gin.Context {
+		c := setupGroupPoolContext(t)
+		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
+		c.Request.Header.Set("Session-Id", session)
+		c.Set("id", 100)
+		c.Set("token_id", token)
+		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}, ModelRules: []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}})
+		return c
+	}
+	c := ctx(100, "session-a")
+	param := &RetryParam{Ctx: c, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)}
+	first, group, err := SelectChannelForRequest(c, modelName, param)
+	require.Nil(t, err)
+	require.NotNil(t, first)
+	require.Equal(t, 6102, first.Id)
+	require.Equal(t, "vip", group)
+	AppendUsedChannel(c, first.Id)
+	second, group, err2 := CacheGetRandomSatisfiedChannel(param)
+	require.NoError(t, err2)
+	require.NotNil(t, second)
+	require.Equal(t, 6101, second.Id)
+	c.Set("channel_id", second.Id)
+	common.SetContextKey(c, constant.ContextKeyAutoGroup, group)
+	RequestPolicy(c).Successful = true
+	RecordPoolAffinity(c)
+	next := ctx(100, "session-a")
+	bound, boundGroup := PreferredPoolChannel(next, modelName)
+	require.NotNil(t, bound)
+	require.Equal(t, 6101, bound.Id)
+	require.Equal(t, "default", boundGroup)
+	isolated, _ := PreferredPoolChannel(ctx(101, "session-a"), modelName)
+	require.Nil(t, isolated)
+	isolated, _ = PreferredPoolChannel(ctx(100, "session-b"), modelName)
+	require.Nil(t, isolated)
+	isolated, _ = PreferredPoolChannel(ctx(100, "session-a"), "another-model")
+	require.Nil(t, isolated)
+	ClearPoolAffinity(next)
+	bound, _ = PreferredPoolChannel(ctx(100, "session-a"), modelName)
+	require.Nil(t, bound)
+	AppendUsedChannel(c, second.Id)
+	last, _, err2 := CacheGetRandomSatisfiedChannel(param)
+	require.NoError(t, err2)
+	require.Nil(t, last)
+}