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/channel_select.go

@@ -2,13 +2,21 @@ package service
 
 import (
 	"errors"
+	"fmt"
+	"net/http"
+	"sort"
+	"sync"
+	"sync/atomic"
 
 	"github.com/QuantumNous/new-api/common"
 	"github.com/QuantumNous/new-api/constant"
 	"github.com/QuantumNous/new-api/dto"
+	"github.com/QuantumNous/new-api/i18n"
 	"github.com/QuantumNous/new-api/logger"
 	"github.com/QuantumNous/new-api/model"
 	"github.com/QuantumNous/new-api/pkg/jsplugin"
+	"github.com/QuantumNous/new-api/relaykit/types"
+	"github.com/QuantumNous/new-api/setting"
 	"github.com/gin-gonic/gin"
 )
 
@@ -28,10 +36,12 @@ func AppendTaskPluginIdentityFilter(c *gin.Context, pluginKey string) {
 	if c == nil {
 		return
 	}
+	channelTypes, pluginKeys := pinnedTaskPluginIdentities(c, pluginKey)
 	GetChannelConstraints(c).AddFilter(dto.ChannelFilter{
 		Kind:                   dto.FilterTaskPluginIdentity,
 		TaskPluginKey:          pluginKey,
-		TaskPluginChannelTypes: pinnedTaskPluginChannelTypes(c, pluginKey),
+		TaskPluginChannelTypes: channelTypes,
+		TaskPluginKeys:         pluginKeys,
 	})
 }
 
@@ -70,6 +80,159 @@ func (p *RetryParam) ResetRetryNextTry() {
 	p.resetNextTry = true
 }
 
+// groupPoolTier is a set of groups sharing one priority. Tiers are tried in
+// ascending priority order, and the groups inside a tier are load-balanced.
+type groupPoolTier struct {
+	priority int
+	groups   []string
+}
+
+var groupPoolRoundRobinCounters sync.Map
+
+// buildGroupPoolTiers resolves the token's routing pool into ordered tiers.
+// Entries the user group cannot select are dropped, and the token's own group is
+// placed ahead of the configured pool when it is missing from it.
+func buildGroupPoolTiers(c *gin.Context, primary, userGroup string) []groupPoolTier {
+	value, exists := common.GetContextKey(c, constant.ContextKeyTokenGroupPool)
+	if !exists {
+		return nil
+	}
+	entries, ok := value.([]model.TokenGroupEntry)
+	if !ok || len(entries) == 0 {
+		return nil
+	}
+	maxCount := setting.GetMaxTokenAutoGroups()
+	ordered := make([]model.TokenGroupEntry, 0, len(entries)+1)
+	seen := make(map[string]struct{}, len(entries)+1)
+	hasPrimary := false
+	for _, entry := range entries {
+		if len(ordered) >= maxCount {
+			break
+		}
+		if !IsUserSelectableGroup(userGroup, entry.Group) {
+			continue
+		}
+		if _, duplicate := seen[entry.Group]; duplicate {
+			continue
+		}
+		seen[entry.Group] = struct{}{}
+		if entry.Group == primary {
+			hasPrimary = true
+		}
+		ordered = append(ordered, entry)
+	}
+	if len(ordered) == 0 {
+		return nil
+	}
+	if !hasPrimary && primary != "" && primary != "auto" {
+		minPriority := ordered[0].Priority
+		for _, entry := range ordered {
+			if entry.Priority < minPriority {
+				minPriority = entry.Priority
+			}
+		}
+		ordered = append([]model.TokenGroupEntry{{Group: primary, Priority: minPriority - 1}}, ordered...)
+	}
+	sort.SliceStable(ordered, func(i, j int) bool { return ordered[i].Priority < ordered[j].Priority })
+	tiers := make([]groupPoolTier, 0, len(ordered))
+	for _, entry := range ordered {
+		if len(tiers) > 0 && tiers[len(tiers)-1].priority == entry.Priority {
+			tiers[len(tiers)-1].groups = append(tiers[len(tiers)-1].groups, entry.Group)
+			continue
+		}
+		tiers = append(tiers, groupPoolTier{priority: entry.Priority, groups: []string{entry.Group}})
+	}
+	return tiers
+}
+
+func failedRequestGroups(c *gin.Context) map[string]bool {
+	failed := make(map[string]bool)
+	groups, ok := common.GetContextKeyType[[]string](c, constant.ContextKeyTokenFailedGroups)
+	if !ok {
+		return failed
+	}
+	for _, group := range groups {
+		failed[group] = true
+	}
+	return failed
+}
+
+// MarkRequestGroupFailed records that a group already failed during the current
+// request so the next selection switches away from it instead of retrying it.
+func MarkRequestGroupFailed(c *gin.Context, group string) {
+	if c == nil || group == "" {
+		return
+	}
+	groups, _ := common.GetContextKeyType[[]string](c, constant.ContextKeyTokenFailedGroups)
+	for _, existing := range groups {
+		if existing == group {
+			return
+		}
+	}
+	common.SetContextKey(c, constant.ContextKeyTokenFailedGroups, append(groups, group))
+}
+
+func nextTierIndex(key string, size int) int {
+	if size <= 1 {
+		return 0
+	}
+	value, _ := groupPoolRoundRobinCounters.LoadOrStore(key, new(uint64))
+	counter, ok := value.(*uint64)
+	if !ok {
+		return 0
+	}
+	return int(atomic.AddUint64(counter, 1)-1) % size
+}
+
+func selectFromGroupPool(param *RetryParam, tiers []groupPoolTier, filters []dto.ChannelFilter) (*model.Channel, string) {
+	failed := failedRequestGroups(param.Ctx)
+	// Each retry advances one priority tier, so a retryable failure reaches the
+	// next available group instead of repeating the group that just failed.
+	start := param.GetRetry()
+	if start < 0 {
+		start = 0
+	}
+	if start > len(tiers)-1 {
+		start = len(tiers) - 1
+	}
+	for _, tier := range tiers[start:] {
+		eligible := make([]string, 0, len(tier.groups))
+		for _, group := range tier.groups {
+			if failed[group] || !IsGroupAvailable(group) {
+				continue
+			}
+			eligible = append(eligible, group)
+		}
+		if len(eligible) == 0 {
+			continue
+		}
+		start := nextTierIndex(tier.groups[0], len(eligible))
+		for offset := 0; offset < len(eligible); offset++ {
+			group := eligible[(start+offset)%len(eligible)]
+			if channel, _ := selectWithTrace(param, group, 0); channel != nil {
+				logger.LogInfo(param.Ctx, fmt.Sprintf("group pool selected: group=%s channel=%d", group, channel.Id))
+				common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group)
+				return channel, group
+			}
+		}
+	}
+	// Availability is a heuristic, so when every tier looks degraded fall back to
+	// the highest-priority group this request has not tried yet rather than
+	// failing without an attempt.
+	for _, tier := range tiers {
+		for _, group := range tier.groups {
+			if failed[group] {
+				continue
+			}
+			if channel, _ := selectWithTrace(param, group, 0); channel != nil {
+				common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group)
+				return channel, group
+			}
+		}
+	}
+	return nil, param.TokenGroup
+}
+
 // CacheGetRandomSatisfiedChannel tries to get a random channel that satisfies the requirements.
 // 尝试获取一个满足要求的随机渠道。
 //
@@ -106,6 +269,7 @@ func (p *RetryParam) ResetRetryNextTry() {
 //	Retry=3: GroupB, priority1 (startRetryIndex=2, priorityRetry=1)
 //	         分组B, 优先级1
 func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, error) {
+	if RequestRoutingPolicy(param.Ctx) != nil { return selectTokenRoutingPool(param) }
 	var channel *model.Channel
 	var err error
 	selectGroup := param.TokenGroup
@@ -113,6 +277,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
 	filters := GetChannelConstraints(param.Ctx).Filters
 	crossGroupRetry := common.GetContextKeyBool(param.Ctx, constant.ContextKeyTokenCrossGroupRetry)
 	orderedGroups := make([]string, 0)
+	tiers := make([]groupPoolTier, 0)
 
 	if param.TokenGroup == "auto" {
 		orderedGroups = GetRequestAutoGroups(param.Ctx, userGroup)
@@ -120,19 +285,31 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
 			return nil, selectGroup, errors.New("auto groups is not enabled")
 		}
 	} else if crossGroupRetry {
-		if value, exists := common.GetContextKey(param.Ctx, constant.ContextKeyTokenAutoGroups); exists {
-			if backupGroups, ok := value.([]string); ok {
-				orderedGroups = append(orderedGroups, param.TokenGroup)
-				for _, group := range FilterUserTokenAutoGroups(userGroup, backupGroups) {
-					if group != param.TokenGroup {
-						orderedGroups = append(orderedGroups, group)
+		if _, exists := common.GetContextKey(param.Ctx, constant.ContextKeyTokenGroupPool); exists {
+			tiers = buildGroupPoolTiers(param.Ctx, param.TokenGroup, userGroup)
+		}
+		if len(tiers) == 0 {
+			if value, exists := common.GetContextKey(param.Ctx, constant.ContextKeyTokenAutoGroups); exists {
+				if backupGroups, ok := value.([]string); ok {
+					orderedGroups = append(orderedGroups, param.TokenGroup)
+					for _, group := range FilterUserTokenAutoGroups(userGroup, backupGroups) {
+						if group != param.TokenGroup {
+							orderedGroups = append(orderedGroups, group)
+						}
 					}
 				}
 			}
 		}
 	}
 
+	if len(tiers) > 0 {
+		logger.LogInfo(param.Ctx, fmt.Sprintf("group pool selection: tokenGroup=%s tiers=%d retry=%d", param.TokenGroup, len(tiers), param.GetRetry()))
+		channel, selectGroup = selectFromGroupPool(param, tiers, filters)
+		return channel, selectGroup, nil
+	}
+
 	if len(orderedGroups) > 0 {
+		logger.LogInfo(param.Ctx, fmt.Sprintf("legacy group selection: groups=%v retry=%d", orderedGroups, param.GetRetry()))
 		// startGroupIndex: the group index to start searching from
 		// startGroupIndex: 开始搜索的分组索引
 		startGroupIndex := 0
@@ -155,12 +332,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
 			}
 			logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry)
 
-			channel, _ = model.GetRandomSatisfiedChannel(
-				autoGroup,
-				param.ModelName,
-				priorityRetry,
-				filters,
-			)
+			channel, _ = selectWithTrace(param, autoGroup, priorityRetry)
 			if channel == nil {
 				// Current group has no available channel for this model, try next group
 				// 当前分组没有该模型的可用渠道，尝试下一个分组
@@ -198,12 +370,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
 			break
 		}
 	} else {
-		channel, err = model.GetRandomSatisfiedChannel(
-			param.TokenGroup,
-			param.ModelName,
-			param.GetRetry(),
-			filters,
-		)
+		channel, err = selectWithTrace(param, param.TokenGroup, param.GetRetry())
 		if err != nil {
 			return nil, param.TokenGroup, err
 		}
@@ -211,15 +378,16 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
 	return channel, selectGroup, nil
 }
 
-func pinnedTaskPluginChannelTypes(c *gin.Context, expected string) []int {
+func pinnedTaskPluginIdentities(c *gin.Context, expected string) ([]int, []string) {
 	if c == nil || expected == "" {
-		return nil
+		return nil, nil
 	}
 	if value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint); exists {
 		pinned, ok := value.(jsplugin.PinnedEndpoint)
 		if ok && pinned.Generation != nil && len(pinned.Candidates) > 1 {
 			expectedFound := false
 			channelTypes := make([]int, 0, len(pinned.Candidates))
+			pluginKeys := make([]string, 0, len(pinned.Candidates))
 			seen := make(map[int]struct{}, len(pinned.Candidates))
 			for _, candidate := range pinned.Candidates {
 				if candidate.Plugin == nil {
@@ -228,6 +396,7 @@ func pinnedTaskPluginChannelTypes(c *gin.Context, expected string) []int {
 				if candidate.Plugin.Meta.Key == expected {
 					expectedFound = true
 				}
+				pluginKeys = append(pluginKeys, candidate.Plugin.Meta.Key)
 				for _, channelType := range candidate.Plugin.Meta.ChannelTypes {
 					if channelType == 0 || channelType == constant.ChannelTypeTaskPlugin {
 						continue
@@ -242,14 +411,14 @@ func pinnedTaskPluginChannelTypes(c *gin.Context, expected string) []int {
 				}
 			}
 			if expectedFound {
-				return channelTypes
+				return channelTypes, pluginKeys
 			}
 		}
 	}
 	value, exists := c.Get(jsplugin.ContextKeyPinnedPlugin)
 	pinned, ok := value.(jsplugin.PinnedPlugin)
 	if !exists || !ok || pinned.Generation == nil || pinned.Plugin == nil || pinned.Plugin.Meta.Key != expected {
-		return nil
+		return nil, nil
 	}
 	channelTypes := make([]int, 0, len(pinned.Plugin.Meta.ChannelTypes))
 	for _, channelType := range pinned.Plugin.Meta.ChannelTypes {
@@ -258,8 +427,157 @@ func pinnedTaskPluginChannelTypes(c *gin.Context, expected string) []int {
 		}
 		channelTypes = append(channelTypes, channelType)
 	}
-	if len(channelTypes) == 0 {
-		return nil
+	return channelTypes, []string{expected}
+}
+
+// ChannelSelectError explains why SelectChannelForRequest found no channel.
+// Callers render it for their transport: the HTTP distributor localizes
+// MessageID with its own helpers and the Responses WebSocket relay wraps it in
+// a NewAPIError. Message is set instead of MessageID when the text is a fixed
+// error code that clients match on.
+type ChannelSelectError struct {
+	StatusCode int
+	Code       types.ErrorCode
+	MessageID  string
+	Params     map[string]any
+	Message    string
+	// FilterKind and Channel identify a candidate rejected by request filters.
+	FilterKind dto.ChannelFilterKind
+	Channel    *model.Channel
+	// NoAvailableChannel marks the "no channel for this group and model"
+	// outcome so the distributor can name the claiming task plugin.
+	NoAvailableChannel bool
+}
+
+// SelectChannelForRequest resolves the channel for one attempt with the rules
+// shared by the HTTP distributor and the Responses WebSocket relay: a pinned
+// channel wins, then session affinity (first attempt only), then a random
+// eligible channel; every candidate must satisfy the request's channel
+// filters. The group the channel was chosen from is returned for auto-group
+// callers. The caller still applies SetupContextForSelectedChannel.
+func SelectChannelForRequest(c *gin.Context, modelName string, retry *RetryParam) (*model.Channel, string, *ChannelSelectError) {
+	c.Set("original_model", modelName)
+	constraints := GetChannelConstraints(c)
+	if pin, found, overridden := constraints.ResolvedPin(); found {
+		for _, lost := range overridden {
+			logger.LogWarn(c, fmt.Sprintf(
+				"channel pin overridden: winning_source=%s winning_channel_id=%d overridden_source=%s overridden_channel_id=%d",
+				pin.Source, pin.ChannelId, lost.Source, lost.ChannelId,
+			))
+		}
+		channel, err := model.CacheGetChannel(pin.ChannelId)
+		if err != nil {
+			recordDirectSelection(c, nil, pin.ChannelId, retry.TokenGroup, modelName, "pinned_channel_missing")
+			return nil, "", pinnedChannelUnavailable(pin, http.StatusBadRequest, i18n.MsgDistributorInvalidChannelId)
+		}
+		if channel.Status != common.ChannelStatusEnabled {
+			recordDirectSelection(c, channel, channel.Id, retry.TokenGroup, modelName, "channel_disabled")
+			return nil, "", pinnedChannelUnavailable(pin, http.StatusForbidden, i18n.MsgDistributorChannelDisabled)
+		}
+		if ok, kind := model.ChannelSatisfiesFilters(channel, modelName, constraints.Filters); !ok {
+			recordDirectSelection(c, channel, channel.Id, retry.TokenGroup, modelName, "excluded_"+string(kind))
+			return nil, "", &ChannelSelectError{
+				StatusCode: http.StatusBadRequest, Code: types.ErrorCode(kind), MessageID: i18n.MsgDistributorNoAvailableChannel,
+				Params:     map[string]any{"Group": common.GetContextKeyString(c, constant.ContextKeyUsingGroup), "Model": modelName},
+				FilterKind: kind, Channel: channel,
+			}
+		}
+		recordDirectSelection(c, channel, channel.Id, retry.TokenGroup, modelName, "pinned_channel")
+		return channel, "", nil
 	}
-	return channelTypes
+
+	usingGroup := retry.TokenGroup
+	if RequestRoutingPolicy(c) != nil {
+		if channel, group := PreferredPoolChannel(c, modelName); channel != nil && retry.GetRetry() == 0 { return channel,group,nil }
+		channel, group, err := selectTokenRoutingPool(retry)
+		if err != nil || channel == nil { return nil,group,&ChannelSelectError{StatusCode:http.StatusServiceUnavailable,Code:types.ErrorCodeModelNotFound,Message:"No eligible channel in the group pool",NoAvailableChannel:true} }
+		return channel,group,nil
+	}
+	var channel *model.Channel
+	var selectGroup string
+	if retry.GetRetry() == 0 {
+		if preferredChannelID, found := GetPreferredChannelByAffinity(c, modelName, usingGroup); found {
+			affinityUsable := false
+			preferred, err := model.CacheGetChannel(preferredChannelID)
+			affinitySatisfied := false
+			if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled {
+				affinitySatisfied, _ = model.ChannelSatisfiesFilters(preferred, modelName, constraints.Filters)
+			}
+			if affinitySatisfied {
+				if usingGroup == "auto" {
+					userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
+					for _, g := range GetRequestAutoGroups(c, userGroup) {
+						if model.IsChannelEnabledForGroupModel(g, modelName, preferred.Id) {
+							selectGroup = g
+							common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
+							channel = preferred
+							affinityUsable = true
+							MarkChannelAffinityUsed(c, g, preferred.Id)
+							break
+						}
+					}
+				} else if model.IsChannelEnabledForGroupModel(usingGroup, modelName, preferred.Id) {
+					channel = preferred
+					selectGroup = usingGroup
+					affinityUsable = true
+					MarkChannelAffinityUsed(c, usingGroup, preferred.Id)
+				}
+			}
+			reason := "affinity_unavailable"
+			if affinityUsable {
+				reason = "affinity_hit"
+			}
+			recordDirectSelection(c, preferred, preferredChannelID, usingGroup, modelName, reason)
+			if !affinityUsable && !ShouldKeepChannelAffinityOnChannelDisabled() {
+				ClearCurrentChannelAffinityCache(c)
+			}
+			if !affinityUsable && RequestPolicy(c).SessionMode == "strict" {
+				return nil, "", &ChannelSelectError{StatusCode: http.StatusServiceUnavailable, Message: "strict_session_binding_unavailable"}
+			}
+		}
+	}
+
+	if channel == nil {
+		var err error
+		channel, selectGroup, err = CacheGetRandomSatisfiedChannel(retry)
+		if err != nil {
+			showGroup := usingGroup
+			if usingGroup == "auto" {
+				showGroup = fmt.Sprintf("auto(%s)", selectGroup)
+			}
+			return nil, selectGroup, &ChannelSelectError{
+				StatusCode: http.StatusServiceUnavailable, Code: types.ErrorCodeModelNotFound, MessageID: i18n.MsgDistributorGetChannelFailed,
+				Params: map[string]any{"Group": showGroup, "Model": modelName, "Error": err.Error()},
+			}
+		}
+		if channel == nil {
+			return nil, selectGroup, &ChannelSelectError{
+				StatusCode: http.StatusServiceUnavailable, Code: types.ErrorCodeModelNotFound, MessageID: i18n.MsgDistributorNoAvailableChannel,
+				Params: map[string]any{"Group": usingGroup, "Model": modelName}, NoAvailableChannel: true,
+			}
+		}
+	}
+	if ok, kind := model.ChannelSatisfiesFilters(channel, modelName, constraints.Filters); !ok {
+		return nil, selectGroup, &ChannelSelectError{
+			StatusCode: http.StatusServiceUnavailable, Code: types.ErrorCodeModelNotFound, MessageID: i18n.MsgDistributorNoAvailableChannel,
+			Params:     map[string]any{"Group": common.GetContextKeyString(c, constant.ContextKeyUsingGroup), "Model": modelName},
+			FilterKind: kind, Channel: channel, NoAvailableChannel: true,
+		}
+	}
+	return channel, selectGroup, nil
+}
+
+// Origin-task pins report a fixed code so task polling can tell a retired
+// channel from a malformed request.
+func pinnedChannelUnavailable(pin dto.ChannelPin, statusCode int, messageID string) *ChannelSelectError {
+	if pin.Source == dto.PinSourceOriginTask {
+		return &ChannelSelectError{StatusCode: http.StatusBadRequest, Code: "origin_task_channel_disabled", Message: "origin_task_channel_disabled"}
+	}
+	return &ChannelSelectError{StatusCode: statusCode, MessageID: messageID}
+}
+
+// AppendUsedChannel records an attempted channel in the request's channel
+// trail, which the retry log and the consume log's admin_info both read.
+func AppendUsedChannel(c *gin.Context, channelID int) {
+	c.Set("use_channel", append(c.GetStringSlice("use_channel"), fmt.Sprintf("%d", channelID)))
 }