feat: add key-level channel affinity

controller/relay.go

@@ -229,6 +229,9 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		relayInfo.LastError = newAPIError
 
 		processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
+		if service.ClearChannelAffinityAfterFailure(c) {
+			logger.LogInfo(c, fmt.Sprintf("渠道亲和命中失败，已清理缓存：channel #%d", channel.Id))
+		}
 
 		if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
 			break
@@ -556,6 +559,9 @@ func RelayTask(c *gin.Context) {
 				*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
 					common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
 				types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
+			if service.ClearChannelAffinityAfterFailure(c) {
+				logger.LogInfo(c, fmt.Sprintf("渠道亲和命中失败，已清理缓存：channel #%d", channel.Id))
+			}
 		}
 
 		if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {

middleware/distributor.go

@@ -101,11 +101,16 @@ func Distribute() func(c *gin.Context) {
 					}
 				}
 
-				if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found {
+				if affinitySelection, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found {
 					affinityUsable := false
-					preferred, err := model.CacheGetChannel(preferredChannelID)
-					if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled {
-						if usingGroup == "auto" {
+					preferred, err := model.CacheGetChannel(affinitySelection.ChannelID)
+					if err == nil && preferred != nil {
+						if preferred.Status != common.ChannelStatusEnabled {
+							if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
+								abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorAffinityChannelDisabled))
+								return
+							}
+						} else if usingGroup == "auto" {
 							userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
 							autoGroups := service.GetUserAutoGroup(userGroup)
 							for _, g := range autoGroups {
@@ -114,15 +119,15 @@ func Distribute() func(c *gin.Context) {
 									common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
 									channel = preferred
 									affinityUsable = true
-									service.MarkChannelAffinityUsed(c, g, preferred.Id)
+									service.MarkChannelAffinityUsed(c, g, affinitySelection)
 									break
 								}
 							}
 						} else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) {
 							channel = preferred
 							selectGroup = usingGroup
 							affinityUsable = true
-							service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id)
+							service.MarkChannelAffinityUsed(c, usingGroup, affinitySelection)
 						}
 					}
 					if !affinityUsable && !service.ShouldKeepChannelAffinityOnChannelDisabled() {
@@ -449,10 +454,19 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
 	common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping())
 	common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
 
-	key, index, newAPIError := channel.GetNextEnabledKey()
-	if newAPIError != nil {
-		return newAPIError
+	var key string
+	var index int
+	var newAPIError *types.NewAPIError
+	if preferredKeyIndex, ok := service.GetChannelAffinityKeyIndex(c, channel.Id); ok {
+		key, index, newAPIError = channel.GetEnabledKeyByIndex(preferredKeyIndex)
+	}
+	if key == "" || newAPIError != nil {
+		key, index, newAPIError = channel.GetNextEnabledKey()
+		if newAPIError != nil {
+			return newAPIError
+		}
 	}
+	service.UpdateChannelAffinitySelectedKeyIndex(c, channel.Id, index)
 	if channel.ChannelInfo.IsMultiKey {
 		common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true)
 		common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, index)

model/channel.go

@@ -196,6 +196,38 @@ func (channel *Channel) GetKeys() []string {
 	return keys
 }
 
+func (channel *Channel) GetEnabledKeyByIndex(index int) (string, int, *types.NewAPIError) {
+	if !channel.ChannelInfo.IsMultiKey {
+		if index != 0 {
+			return "", 0, types.NewError(errors.New("invalid key index"), types.ErrorCodeChannelNoAvailableKey)
+		}
+		return channel.Key, 0, nil
+	}
+
+	keys := channel.GetKeys()
+	if len(keys) == 0 {
+		return "", 0, types.NewError(errors.New("no keys available"), types.ErrorCodeChannelNoAvailableKey)
+	}
+	if index < 0 || index >= len(keys) {
+		return "", 0, types.NewError(errors.New("invalid key index"), types.ErrorCodeChannelNoAvailableKey)
+	}
+
+	lock := GetChannelPollingLock(channel.Id)
+	lock.Lock()
+	defer lock.Unlock()
+
+	status := common.ChannelStatusEnabled
+	if channel.ChannelInfo.MultiKeyStatusList != nil {
+		if s, ok := channel.ChannelInfo.MultiKeyStatusList[index]; ok {
+			status = s
+		}
+	}
+	if status != common.ChannelStatusEnabled {
+		return "", 0, types.NewError(errors.New("key is disabled"), types.ErrorCodeChannelNoAvailableKey)
+	}
+	return keys[index], index, nil
+}
+
 func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) {
 	// If not in multi-key mode, return the original key string directly.
 	if !channel.ChannelInfo.IsMultiKey {

model/channel_key_affinity_test.go

@@ -0,0 +1,48 @@
+package model
+
+import (
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetEnabledKeyByIndex(t *testing.T) {
+	channel := &Channel{
+		Id:  9001,
+		Key: "key-a\nkey-b\nkey-c",
+		ChannelInfo: ChannelInfo{
+			IsMultiKey:           true,
+			MultiKeyMode:         constant.MultiKeyModePolling,
+			MultiKeyPollingIndex: 0,
+			MultiKeyStatusList: map[int]int{
+				1: common.ChannelStatusEnabled,
+			},
+		},
+	}
+
+	key, index, apiErr := channel.GetEnabledKeyByIndex(1)
+	require.Nil(t, apiErr)
+	require.Equal(t, "key-b", key)
+	require.Equal(t, 1, index)
+	require.Equal(t, 0, channel.ChannelInfo.MultiKeyPollingIndex)
+}
+
+func TestGetEnabledKeyByIndexDisabled(t *testing.T) {
+	channel := &Channel{
+		Id:  9002,
+		Key: "key-a\nkey-b",
+		ChannelInfo: ChannelInfo{
+			IsMultiKey: true,
+			MultiKeyStatusList: map[int]int{
+				1: common.ChannelStatusManuallyDisabled,
+			},
+		},
+	}
+
+	key, index, apiErr := channel.GetEnabledKeyByIndex(1)
+	require.NotNil(t, apiErr)
+	require.Empty(t, key)
+	require.Equal(t, 0, index)
+}