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" ) func GetChannelConstraints(c *gin.Context) *dto.ChannelConstraints { if c == nil { return &dto.ChannelConstraints{} } if existing, ok := common.GetContextKeyType[*dto.ChannelConstraints](c, constant.ContextKeyChannelConstraints); ok && existing != nil { return existing } constraints := &dto.ChannelConstraints{} common.SetContextKey(c, constant.ContextKeyChannelConstraints, constraints) return constraints } 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: channelTypes, TaskPluginKeys: pluginKeys, }) } type RetryParam struct { Ctx *gin.Context TokenGroup string ModelName string RequestPath string Retry *int resetNextTry bool } func (p *RetryParam) GetRetry() int { if p.Retry == nil { return 0 } return *p.Retry } func (p *RetryParam) SetRetry(retry int) { p.Retry = &retry } func (p *RetryParam) IncreaseRetry() { if p.resetNextTry { p.resetNextTry = false return } if p.Retry == nil { p.Retry = new(int) } *p.Retry++ } 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. // 尝试获取一个满足要求的随机渠道。 // // For ordered Auto groups or a fixed token group with a backup group: // 对于有序 Auto 分组,或配置了备用分组的固定令牌分组: // // - Each group will exhaust all its priorities before moving to the next group. // 每个分组会用完所有优先级后才会切换到下一个分组。 // // - Uses ContextKeyAutoGroupIndex to track current group index. // 使用 ContextKeyAutoGroupIndex 跟踪当前分组索引。 // // - Uses ContextKeyAutoGroupRetryIndex to track the global Retry count when current group started. // 使用 ContextKeyAutoGroupRetryIndex 跟踪当前分组开始时的全局重试次数。 // // - priorityRetry = Retry - startRetryIndex, represents the priority level within current group. // priorityRetry = Retry - startRetryIndex,表示当前分组内的优先级级别。 // // - When GetRandomSatisfiedChannel returns nil (priorities exhausted), moves to next group. // 当 GetRandomSatisfiedChannel 返回 nil(优先级用完)时,切换到下一个分组。 // // Example flow (2 groups, each with 2 priorities, RetryTimes=3): // 示例流程(2个分组,每个有2个优先级,RetryTimes=3): // // Retry=0: GroupA, priority0 (startRetryIndex=0, priorityRetry=0) // 分组A, 优先级0 // // Retry=1: GroupA, priority1 (startRetryIndex=0, priorityRetry=1) // 分组A, 优先级1 // // Retry=2: GroupA exhausted → GroupB, priority0 (startRetryIndex=2, priorityRetry=0) // 分组A用完 → 分组B, 优先级0 // // 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) } if group := imageGenerationRequestGroup(param.Ctx, param.ModelName); group != "" { channel, err := selectWithTrace(param, group, param.GetRetry()) common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group) return channel, group, err } var channel *model.Channel var err error selectGroup := param.TokenGroup userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup) 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) if len(orderedGroups) == 0 { return nil, selectGroup, errors.New("auto groups is not enabled") } } else if crossGroupRetry { 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 if lastGroupIndex, exists := common.GetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex); exists { if idx, ok := lastGroupIndex.(int); ok { startGroupIndex = idx } } for i := startGroupIndex; i < len(orderedGroups); i++ { autoGroup := orderedGroups[i] // Calculate priorityRetry for current group // 计算当前分组的 priorityRetry priorityRetry := param.GetRetry() // If moved to a new group, reset priorityRetry and update startRetryIndex // 如果切换到新分组,重置 priorityRetry 并更新 startRetryIndex if i > startGroupIndex { priorityRetry = 0 } logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) channel, _ = selectWithTrace(param, autoGroup, priorityRetry) if channel == nil { // Current group has no available channel for this model, try next group // 当前分组没有该模型的可用渠道,尝试下一个分组 logger.LogDebug(param.Ctx, "No available channel in group %s for model %s at priorityRetry %d, trying next group", autoGroup, param.ModelName, priorityRetry) // 重置状态以尝试下一个分组 common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i+1) common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupRetryIndex, 0) // Reset retry counter so outer loop can continue for next group // 重置重试计数器,以便外层循环可以为下一个分组继续 param.SetRetry(0) continue } common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, autoGroup) selectGroup = autoGroup logger.LogDebug(param.Ctx, "Auto selected group: %s", autoGroup) // Prepare state for next retry // 为下一次重试准备状态 if crossGroupRetry && priorityRetry >= common.RetryTimes { // Current group has exhausted all retries, prepare to switch to next group // This request still uses current group, but next retry will use next group // 当前分组已用完所有重试次数,准备切换到下一个分组 // 本次请求仍使用当前分组,但下次重试将使用下一个分组 logger.LogDebug(param.Ctx, "Current group %s retries exhausted (priorityRetry=%d >= RetryTimes=%d), preparing switch to next group for next retry", autoGroup, priorityRetry, common.RetryTimes) common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i+1) // Reset retry counter so outer loop can continue for next group // 重置重试计数器,以便外层循环可以为下一个分组继续 param.SetRetry(0) param.ResetRetryNextTry() } else { // Stay in current group, save current state // 保持在当前分组,保存当前状态 common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i) } break } } else { channel, err = selectWithTrace(param, param.TokenGroup, param.GetRetry()) if err != nil { return nil, param.TokenGroup, err } } return channel, selectGroup, nil } func pinnedTaskPluginIdentities(c *gin.Context, expected string) ([]int, []string) { if c == nil || expected == "" { 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 { continue } 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 } if _, duplicate := seen[channelType]; duplicate { continue } if plugin, indexed := pinned.Generation.GetByChannelType(channelType); indexed && plugin == candidate.Plugin { seen[channelType] = struct{}{} channelTypes = append(channelTypes, channelType) } } } if expectedFound { 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, nil } channelTypes := make([]int, 0, len(pinned.Plugin.Meta.ChannelTypes)) for _, channelType := range pinned.Plugin.Meta.ChannelTypes { if channelType == 0 || channelType == constant.ChannelTypeTaskPlugin { continue } channelTypes = append(channelTypes, channelType) } 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, common.LogText( "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 } usingGroup := retry.TokenGroup if RequestRoutingPolicy(c) != nil { if retry.GetRetry() == 0 { if channel, group := PreferredPoolChannel(c, modelName); channel != nil { 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))) }