From d006b24f2f9cdd4ea15b460d3030ccd96be5612f Mon Sep 17 00:00:00 2001 From: fengsilin Date: Tue, 24 Mar 2026 11:22:57 +0800 Subject: [PATCH] feat: use BoundChannelId for channel selection in distributor When a token has BoundChannelId set, the distributor now uses that channel instead of the normal channel selection process. The bound channel is still validated to ensure it is enabled and supports the requested model. Changes: - Add ContextKeyTokenBoundChannelId constant - Set bound channel in context during token authentication - Check and use bound channel in distributor before normal selection Co-Authored-By: Claude Opus 4.6 --- constant/context_key.go | 1 + middleware/auth.go | 3 + middleware/distributor.go | 176 ++++++++++++++++++++++---------------- 3 files changed, 106 insertions(+), 74 deletions(-) diff --git a/constant/context_key.go b/constant/context_key.go index 2ba2fe2..2cdce7e 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -19,6 +19,7 @@ const ( ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled" ContextKeyTokenModelLimit ContextKey = "token_model_limit" ContextKeyTokenCrossGroupRetry ContextKey = "token_cross_group_retry" + ContextKeyTokenBoundChannelId ContextKey = "token_bound_channel_id" /* channel related keys */ ContextKeyChannelId ContextKey = "channel_id" diff --git a/middleware/auth.go b/middleware/auth.go index 342e7f4..6cb20ab 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -389,6 +389,9 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e } common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group) common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry) + if token.BoundChannelId != nil { + common.SetContextKey(c, constant.ContextKeyTokenBoundChannelId, *token.BoundChannelId) + } if len(parts) > 1 { if model.IsAdmin(token.UserId) { c.Set("specific_channel_id", parts[1]) diff --git a/middleware/distributor.go b/middleware/distributor.go index 9e66cb8..fe08db6 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -53,99 +53,127 @@ func Distribute() func(c *gin.Context) { } } else { // Select a channel for the user - // check token model mapping - modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) - if modelLimitEnable { - s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) + // Check if token has a bound channel + boundChannelId, hasBoundChannel := common.GetContextKey(c, constant.ContextKeyTokenBoundChannelId) + if hasBoundChannel { + id, ok := boundChannelId.(int) if !ok { - // token model limit is empty, all models are not allowed - abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess)) + abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId)) return } - var tokenModelLimit map[string]bool - tokenModelLimit, ok = s.(map[string]bool) - if !ok { - tokenModelLimit = map[string]bool{} - } - matchName := ratio_setting.FormatMatchingModelName(modelRequest.Model) // match gpts & thinking-* - if _, ok := tokenModelLimit[matchName]; !ok { - abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) + channel, err = model.GetChannelById(id, true) + if err != nil { + abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId)) return } - } - - if shouldSelectChannel { - if modelRequest.Model == "" { - abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorModelNameRequired)) + if channel.Status != common.ChannelStatusEnabled { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled)) return } - var selectGroup string - usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) - // check path is /pg/chat/completions - if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { - playgroundRequest := &dto.PlayGroundRequest{} - err = common.UnmarshalBodyReusable(c, playgroundRequest) - if err != nil { - abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidPlayground, map[string]any{"Error": err.Error()})) + // Verify the bound channel supports the requested model + if shouldSelectChannel && modelRequest.Model != "" { + usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) + if !model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, channel.Id) { + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound) return } - if playgroundRequest.Group != "" { - if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.Group != usingGroup { - abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorGroupAccessDenied)) - return - } - usingGroup = playgroundRequest.Group - common.SetContextKey(c, constant.ContextKeyUsingGroup, usingGroup) + } + } else { + // Normal channel selection logic + // check token model mapping + modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) + if modelLimitEnable { + s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) + if !ok { + // token model limit is empty, all models are not allowed + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess)) + return + } + var tokenModelLimit map[string]bool + tokenModelLimit, ok = s.(map[string]bool) + if !ok { + tokenModelLimit = map[string]bool{} + } + matchName := ratio_setting.FormatMatchingModelName(modelRequest.Model) // match gpts & thinking-* + if _, ok := tokenModelLimit[matchName]; !ok { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) + return } } - if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { - preferred, err := model.CacheGetChannel(preferredChannelID) - if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled { - 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 - } + if shouldSelectChannel { + if modelRequest.Model == "" { + abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorModelNameRequired)) + return + } + var selectGroup string + usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) + // check path is /pg/chat/completions + if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { + playgroundRequest := &dto.PlayGroundRequest{} + err = common.UnmarshalBodyReusable(c, playgroundRequest) + if err != nil { + abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidPlayground, map[string]any{"Error": err.Error()})) + return + } + if playgroundRequest.Group != "" { + if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.Group != usingGroup { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorGroupAccessDenied)) + return } - } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) { - channel = preferred - selectGroup = usingGroup - service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) + usingGroup = playgroundRequest.Group + common.SetContextKey(c, constant.ContextKeyUsingGroup, usingGroup) } } - } - if channel == nil { - channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ - Ctx: c, - ModelName: modelRequest.Model, - TokenGroup: usingGroup, - Retry: common.GetPointer(0), - }) - if err != nil { - showGroup := usingGroup - if usingGroup == "auto" { - showGroup = fmt.Sprintf("auto(%s)", selectGroup) + if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { + preferred, err := model.CacheGetChannel(preferredChannelID) + if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled { + 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 + selectGroup = usingGroup + service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) + } } - message := i18n.T(c, i18n.MsgDistributorGetChannelFailed, map[string]any{"Group": showGroup, "Model": modelRequest.Model, "Error": err.Error()}) - // 如果错误,但是渠道不为空,说明是数据库一致性问题 - //if channel != nil { - // common.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id)) - // message = "数据库一致性已被破坏,请联系管理员" - //} - abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, types.ErrorCodeModelNotFound) - return } + if channel == nil { - abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound) - return + channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ + Ctx: c, + ModelName: modelRequest.Model, + TokenGroup: usingGroup, + Retry: common.GetPointer(0), + }) + if err != nil { + showGroup := usingGroup + if usingGroup == "auto" { + showGroup = fmt.Sprintf("auto(%s)", selectGroup) + } + message := i18n.T(c, i18n.MsgDistributorGetChannelFailed, map[string]any{"Group": showGroup, "Model": modelRequest.Model, "Error": err.Error()}) + // 如果错误,但是渠道不为空,说明是数据库一致性问题 + //if channel != nil { + // common.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id)) + // message = "数据库一致性已被破坏,请联系管理员" + //} + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, types.ErrorCodeModelNotFound) + return + } + if channel == nil { + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound) + return + } } } }