feat: add channel affinity exclusive lock feature

controller/channel_affinity_cache.go

@@ -2,6 +2,7 @@ package controller
 
 import (
 	"net/http"
+	"strconv"
 	"strings"
 
 	"github.com/QuantumNous/new-api/service"
@@ -59,6 +60,58 @@ func ClearChannelAffinityCache(c *gin.Context) {
 	})
 }
 
+func GetChannelAffinityExclusiveCacheStats(c *gin.Context) {
+	stats := service.GetChannelAffinityExclusiveCacheStats()
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "",
+		"data":    stats,
+	})
+}
+
+func ClearChannelAffinityExclusiveCache(c *gin.Context) {
+	all := strings.TrimSpace(c.Query("all"))
+	channelIDStr := strings.TrimSpace(c.Query("channel_id"))
+
+	if all == "true" {
+		deleted := service.ClearChannelAffinityExclusiveCacheAll()
+		c.JSON(http.StatusOK, gin.H{
+			"success": true,
+			"message": "",
+			"data": gin.H{
+				"deleted": deleted,
+			},
+		})
+		return
+	}
+
+	if channelIDStr == "" {
+		c.JSON(http.StatusBadRequest, gin.H{
+			"success": false,
+			"message": "缺少参数：channel_id，或使用 all=true 清空全部",
+		})
+		return
+	}
+
+	channelID, err := strconv.Atoi(channelIDStr)
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{
+			"success": false,
+			"message": "channel_id 必须是数字",
+		})
+		return
+	}
+
+	ok := service.ClearChannelAffinityExclusiveCacheByChannelID(channelID)
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "",
+		"data": gin.H{
+			"deleted": ok,
+		},
+	})
+}
+
 func GetChannelAffinityUsageCacheStats(c *gin.Context) {
 	ruleName := strings.TrimSpace(c.Query("rule_name"))
 	usingGroup := strings.TrimSpace(c.Query("using_group"))

dto/channel_settings.go

@@ -40,6 +40,7 @@ type ChannelOtherSettings struct {
 	UpstreamModelUpdateLastDetectedModels []string      `json:"upstream_model_update_last_detected_models,omitempty"` // 上次检测到的可加入模型
 	UpstreamModelUpdateLastRemovedModels  []string      `json:"upstream_model_update_last_removed_models,omitempty"`  // 上次检测到的可删除模型
 	UpstreamModelUpdateIgnoredModels      []string      `json:"upstream_model_update_ignored_models,omitempty"`       // 手动忽略的模型
+	AffinityExclusive                     bool          `json:"affinity_exclusive,omitempty"`                         // 亲和性独占：开启后仅允许一个令牌通过亲和性绑定使用此渠道
 }
 
 func (s *ChannelOtherSettings) IsOpenRouterEnterprise() bool {

i18n/keys.go

@@ -290,18 +290,19 @@ const (
 
 // Distributor related messages
 const (
-	MsgDistributorInvalidRequest      = "distributor.invalid_request"
-	MsgDistributorInvalidChannelId    = "distributor.invalid_channel_id"
-	MsgDistributorChannelDisabled     = "distributor.channel_disabled"
-	MsgDistributorTokenNoModelAccess  = "distributor.token_no_model_access"
-	MsgDistributorTokenModelForbidden = "distributor.token_model_forbidden"
-	MsgDistributorModelNameRequired   = "distributor.model_name_required"
-	MsgDistributorInvalidPlayground   = "distributor.invalid_playground_request"
-	MsgDistributorGroupAccessDenied   = "distributor.group_access_denied"
-	MsgDistributorGetChannelFailed    = "distributor.get_channel_failed"
-	MsgDistributorNoAvailableChannel  = "distributor.no_available_channel"
-	MsgDistributorInvalidMidjourney   = "distributor.invalid_midjourney_request"
-	MsgDistributorInvalidParseModel   = "distributor.invalid_request_parse_model"
+	MsgDistributorInvalidRequest         = "distributor.invalid_request"
+	MsgDistributorInvalidChannelId       = "distributor.invalid_channel_id"
+	MsgDistributorChannelDisabled        = "distributor.channel_disabled"
+	MsgDistributorChannelExclusiveLocked = "distributor.channel_exclusive_locked"
+	MsgDistributorTokenNoModelAccess     = "distributor.token_no_model_access"
+	MsgDistributorTokenModelForbidden    = "distributor.token_model_forbidden"
+	MsgDistributorModelNameRequired      = "distributor.model_name_required"
+	MsgDistributorInvalidPlayground      = "distributor.invalid_playground_request"
+	MsgDistributorGroupAccessDenied      = "distributor.group_access_denied"
+	MsgDistributorGetChannelFailed       = "distributor.get_channel_failed"
+	MsgDistributorNoAvailableChannel     = "distributor.no_available_channel"
+	MsgDistributorInvalidMidjourney      = "distributor.invalid_midjourney_request"
+	MsgDistributorInvalidParseModel      = "distributor.invalid_request_parse_model"
 )
 
 // Custom OAuth provider related messages

i18n/locales/en.yaml

@@ -245,6 +245,7 @@ common.invalid_input: "Invalid input"
 distributor.invalid_request: "Invalid request: {{.Error}}"
 distributor.invalid_channel_id: "Invalid channel ID"
 distributor.channel_disabled: "This channel has been disabled"
+distributor.channel_exclusive_locked: "This channel is exclusively occupied by another token"
 distributor.token_no_model_access: "This token has no access to any models"
 distributor.token_model_forbidden: "This token has no access to model {{.Model}}"
 distributor.model_name_required: "Model name not specified, model name cannot be empty"

i18n/locales/zh-CN.yaml

@@ -246,6 +246,7 @@ common.invalid_input: "输入不合法"
 distributor.invalid_request: "无效的请求，{{.Error}}"
 distributor.invalid_channel_id: "无效的渠道 Id"
 distributor.channel_disabled: "该渠道已被禁用"
+distributor.channel_exclusive_locked: "该渠道已被其他令牌独占"
 distributor.token_no_model_access: "该令牌无权访问任何模型"
 distributor.token_model_forbidden: "该令牌无权访问模型 {{.Model}}"
 distributor.model_name_required: "未指定模型名称，模型名称不能为空"

i18n/locales/zh-TW.yaml

@@ -246,6 +246,7 @@ common.invalid_input: "輸入不合法"
 distributor.invalid_request: "無效的請求，{{.Error}}"
 distributor.invalid_channel_id: "無效的管道 Id"
 distributor.channel_disabled: "該管道已被禁用"
+distributor.channel_exclusive_locked: "該管道已被其他令牌獨佔"
 distributor.token_no_model_access: "該令牌無權存取任何模型"
 distributor.token_model_forbidden: "該令牌無權存取模型 {{.Model}}"
 distributor.model_name_required: "未指定模型名稱，模型名稱不能為空"

middleware/distributor.go

@@ -51,6 +51,13 @@ func Distribute() func(c *gin.Context) {
 				abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled))
 				return
 			}
+			if channel.GetOtherSettings().AffinityExclusive {
+				tokenID := c.GetInt(string(constant.ContextKeyTokenId))
+				if !service.IsChannelAffinityExclusiveAvailable(channel.Id, tokenID) {
+					abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelExclusiveLocked))
+					return
+				}
+			}
 		} else {
 			// Select a channel for the user
 			// check token model mapping
@@ -107,22 +114,34 @@ func Distribute() func(c *gin.Context) {
 								abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled))
 								return
 							}
-						} else if usingGroup == "auto" {
-							userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
-							autoGroups := service.GetUserAutoGroup(userGroup)
-							for _, g := range autoGroups {
-								if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) {
-									selectGroup = g
-									common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
+						} else {
+							// 亲和性独占：检查该渠道是否被其他令牌锁定
+							usePreferred := true
+							if preferred.GetOtherSettings().AffinityExclusive {
+								tokenID := c.GetInt(string(constant.ContextKeyTokenId))
+								if !service.IsChannelAffinityExclusiveAvailable(preferred.Id, tokenID) {
+									usePreferred = false // 被其他令牌锁定，回退到随机选择
+								}
+							}
+							if usePreferred {
+								if usingGroup == "auto" {
+									userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
+									autoGroups := service.GetUserAutoGroup(userGroup)
+									for _, g := range autoGroups {
+										if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) {
+											selectGroup = g
+											common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
+											channel = preferred
+											service.MarkChannelAffinityUsed(c, g, preferred.Id)
+											break
+										}
+									}
+								} else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) {
 									channel = preferred
-									service.MarkChannelAffinityUsed(c, g, preferred.Id)
-									break
+									selectGroup = usingGroup
+									service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id)
 								}
 							}
-						} else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) {
-							channel = preferred
-							selectGroup = usingGroup
-							service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id)
 						}
 					}
 				}

model/channel_cache.go

@@ -93,10 +93,113 @@ func SyncChannelCache(frequency int) {
 	}
 }
 
-func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) {
+// ChannelFilter is a predicate that returns true if the channel should be included in selection.
+type ChannelFilter func(*Channel) bool
+
+func passChannelFilters(channel *Channel, filters []ChannelFilter) bool {
+	for _, f := range filters {
+		if f != nil && !f(channel) {
+			return false
+		}
+	}
+	return true
+}
+
+func getRandomSatisfiedChannelFromDB(group string, modelName string, retry int, filters []ChannelFilter) (*Channel, bool, error) {
+	var abilities []Ability
+	if err := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, modelName, true).
+		Order("priority DESC").
+		Order("weight DESC").
+		Find(&abilities).Error; err != nil {
+		return nil, false, err
+	}
+	if len(abilities) == 0 {
+		return nil, false, nil
+	}
+
+	abilitiesByPriority := make(map[int64][]Ability)
+	priorityChannels := make(map[int64][]*Channel)
+	channelIDs := make([]int, 0, len(abilities))
+	for _, ability := range abilities {
+		channelIDs = append(channelIDs, ability.ChannelId)
+	}
+
+	channels, err := GetChannelsByIds(channelIDs)
+	if err != nil {
+		return nil, true, err
+	}
+	channelByID := make(map[int]*Channel, len(channels))
+	for _, channel := range channels {
+		channelByID[channel.Id] = channel
+	}
+
+	availablePriorities := make([]int64, 0)
+	seenPriorities := make(map[int64]bool)
+	for _, ability := range abilities {
+		priority := int64(0)
+		if ability.Priority != nil {
+			priority = *ability.Priority
+		}
+		channel, ok := channelByID[ability.ChannelId]
+		if !ok {
+			return nil, true, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", ability.ChannelId)
+		}
+		if !passChannelFilters(channel, filters) {
+			continue
+		}
+		abilitiesByPriority[priority] = append(abilitiesByPriority[priority], ability)
+		priorityChannels[priority] = append(priorityChannels[priority], channel)
+		if !seenPriorities[priority] {
+			availablePriorities = append(availablePriorities, priority)
+			seenPriorities[priority] = true
+		}
+	}
+
+	if len(availablePriorities) == 0 {
+		return nil, true, nil
+	}
+	sort.Slice(availablePriorities, func(i, j int) bool {
+		return availablePriorities[i] > availablePriorities[j]
+	})
+
+	if retry >= len(availablePriorities) {
+		retry = len(availablePriorities) - 1
+	}
+	targetPriority := availablePriorities[retry]
+	targetAbilities := abilitiesByPriority[targetPriority]
+	targetChannels := priorityChannels[targetPriority]
+	if len(targetAbilities) == 0 || len(targetChannels) == 0 {
+		return nil, true, nil
+	}
+
+	weightSum := uint(0)
+	for _, ability := range targetAbilities {
+		weightSum += ability.Weight + 10
+	}
+	weight := common.GetRandomInt(int(weightSum))
+	for idx, ability := range targetAbilities {
+		weight -= int(ability.Weight) + 10
+		if weight <= 0 {
+			return targetChannels[idx], true, nil
+		}
+	}
+	return targetChannels[len(targetChannels)-1], true, nil
+}
+
+func GetRandomSatisfiedChannel(group string, model string, retry int, filters ...ChannelFilter) (*Channel, error) {
 	// if memory cache is disabled, get channel directly from database
 	if !common.MemoryCacheEnabled {
-		return GetChannel(group, model, retry)
+		channel, modelMatched, err := getRandomSatisfiedChannelFromDB(group, model, retry, filters)
+		if err != nil || channel != nil || modelMatched {
+			return channel, err
+		}
+
+		normalizedModel := ratio_setting.FormatMatchingModelName(model)
+		if normalizedModel == model {
+			return nil, nil
+		}
+		fallbackChannel, _, fallbackErr := getRandomSatisfiedChannelFromDB(group, normalizedModel, retry, filters)
+		return fallbackChannel, fallbackErr
 	}
 
 	channelSyncLock.RLock()
@@ -117,6 +220,9 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
 
 	if len(channels) == 1 {
 		if channel, ok := channelsIDM[channels[0]]; ok {
+			if !passChannelFilters(channel, filters) {
+				return nil, nil
+			}
 			return channel, nil
 		}
 		return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channels[0])
@@ -125,11 +231,17 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
 	uniquePriorities := make(map[int]bool)
 	for _, channelId := range channels {
 		if channel, ok := channelsIDM[channelId]; ok {
+			if !passChannelFilters(channel, filters) {
+				continue
+			}
 			uniquePriorities[int(channel.GetPriority())] = true
 		} else {
 			return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId)
 		}
 	}
+	if len(uniquePriorities) == 0 {
+		return nil, nil
+	}
 	var sortedUniquePriorities []int
 	for priority := range uniquePriorities {
 		sortedUniquePriorities = append(sortedUniquePriorities, priority)
@@ -147,6 +259,9 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
 	for _, channelId := range channels {
 		if channel, ok := channelsIDM[channelId]; ok {
 			if channel.GetPriority() == targetPriority {
+				if !passChannelFilters(channel, filters) {
+					continue
+				}
 				sumWeight += channel.GetWeight()
 				targetChannels = append(targetChannels, channel)
 			}
@@ -156,7 +271,7 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
 	}
 
 	if len(targetChannels) == 0 {
-		return nil, errors.New(fmt.Sprintf("no channel found, group: %s, model: %s, priority: %d", group, model, targetPriority))
+		return nil, nil
 	}
 
 	// smoothing factor and adjustment