From 21077fd3d50d5abea7efaef6dce7faae83170175 Mon Sep 17 00:00:00 2001 From: faithleysath <120073078+faithleysath@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:34:09 +0800 Subject: [PATCH] feat(channel): add RPM and concurrency admission limits --- controller/channel_admission_test.go | 144 ++++ controller/relay.go | 142 +++- docs/channel/other_setting.md | 30 +- i18n/keys.go | 1 + i18n/locales/en.yaml | 1 + i18n/locales/zh-CN.yaml | 1 + i18n/locales/zh-TW.yaml | 1 + middleware/channel_admission_test.go | 349 ++++++++++ middleware/distributor.go | 144 ++-- model/ability.go | 170 ++--- model/channel.go | 3 + model/channel_cache.go | 158 +++-- model/channel_cache_test.go | 17 + model/channel_settings_test.go | 9 + relaykit/dto/channel_settings.go | 27 +- relaykit/dto/channel_settings_test.go | 27 + relaykit/types/error.go | 15 +- service/channel_admission.go | 549 +++++++++++++++ service/channel_admission_test.go | 656 ++++++++++++++++++ service/channel_select.go | 281 +++++--- service/lua/channel_admission_acquire.lua | 51 ++ service/lua/channel_admission_release.lua | 5 + service/lua/channel_admission_renew.lua | 12 + service/lua/channel_admission_snapshot.lua | 20 + .../drawers/channel-mutate-drawer.tsx | 83 ++- .../channel-form-admission-limits.test.ts | 113 +++ .../channels/lib/channel-form-errors.ts | 2 + web/src/features/channels/lib/channel-form.ts | 33 +- web/src/features/channels/types.ts | 2 + web/src/i18n/locales/en.json | 5 + web/src/i18n/locales/fr.json | 5 + web/src/i18n/locales/ja.json | 5 + web/src/i18n/locales/ru.json | 5 + web/src/i18n/locales/vi.json | 5 + web/src/i18n/locales/zh-TW.json | 5 + web/src/i18n/locales/zh.json | 5 + web/src/i18n/static-keys.ts | 3 + 37 files changed, 2711 insertions(+), 373 deletions(-) create mode 100644 controller/channel_admission_test.go create mode 100644 middleware/channel_admission_test.go create mode 100644 model/channel_cache_test.go create mode 100644 service/channel_admission.go create mode 100644 service/channel_admission_test.go create mode 100644 service/lua/channel_admission_acquire.lua create mode 100644 service/lua/channel_admission_release.lua create mode 100644 service/lua/channel_admission_renew.lua create mode 100644 service/lua/channel_admission_snapshot.lua create mode 100644 web/src/features/channels/lib/__tests__/channel-form-admission-limits.test.ts diff --git a/controller/channel_admission_test.go b/controller/channel_admission_test.go new file mode 100644 index 000000000000..d52aa373b53e --- /dev/null +++ b/controller/channel_admission_test.go @@ -0,0 +1,144 @@ +package controller + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + taskdto "github.com/QuantumNous/new-api/dto" + appI18n "github.com/QuantumNous/new-api/i18n" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/service" + + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestGetInitialUnlimitedChannelUsesContextSnapshot(t *testing.T) { + previousMemoryCache := common.MemoryCacheEnabled + previousDB := model.DB + common.MemoryCacheEnabled = false + model.DB = nil + t.Cleanup(func() { + common.MemoryCacheEnabled = previousMemoryCache + model.DB = previousDB + }) + + gin.SetMode(gin.TestMode) + response := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(response) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + common.SetContextKey(ctx, constant.ContextKeyChannelId, 601) + common.SetContextKey(ctx, constant.ContextKeyChannelType, 1) + common.SetContextKey(ctx, constant.ContextKeyChannelName, "snapshot-channel") + common.SetContextKey(ctx, constant.ContextKeyChannelAutoBan, true) + common.SetContextKey(ctx, constant.ContextKeyChannelIsMultiKey, true) + common.SetContextKey(ctx, constant.ContextKeyChannelSetting, dto.ChannelSettings{}) + + channel, lease, newAPIError := getChannel(ctx, &relaycommon.RelayInfo{}, &service.RetryParam{}) + + require.Nil(t, newAPIError) + require.NotNil(t, channel) + assert.Equal(t, 601, channel.Id) + assert.Equal(t, "snapshot-channel", channel.Name) + assert.True(t, channel.GetAutoBan()) + assert.True(t, channel.ChannelInfo.IsMultiKey) + assert.Nil(t, lease) +} + +func TestGetInitialLimitedChannelReturnsCapacityError(t *testing.T) { + require.NoError(t, appI18n.Init()) + dsn := fmt.Sprintf("file:controller-limited-channel-%d?mode=memory&cache=shared", time.Now().UnixNano()) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{})) + + previousDB := model.DB + previousMemoryCache := common.MemoryCacheEnabled + previousRedisEnabled := common.RedisEnabled + model.DB = db + common.MemoryCacheEnabled = false + common.RedisEnabled = false + t.Cleanup(func() { + model.DB = previousDB + common.MemoryCacheEnabled = previousMemoryCache + common.RedisEnabled = previousRedisEnabled + }) + + limitedChannel := &model.Channel{ + Id: 602, + Name: "limited-channel", + Type: constant.ChannelTypeOpenAI, + Status: common.ChannelStatusEnabled, + } + limitedChannel.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + require.NoError(t, db.Create(limitedChannel).Error) + + heldLease, decision, err := service.AcquireChannelAdmission(context.Background(), limitedChannel) + require.NoError(t, err) + require.True(t, decision.Allowed) + defer func() { require.NoError(t, heldLease.Release()) }() + + response := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(response) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + common.SetContextKey(ctx, constant.ContextKeyChannelId, limitedChannel.Id) + common.SetContextKey(ctx, constant.ContextKeyChannelSetting, limitedChannel.GetSetting()) + + channel, lease, newAPIError := getChannel(ctx, &relaycommon.RelayInfo{}, &service.RetryParam{Ctx: ctx}) + + assert.Nil(t, channel) + assert.Nil(t, lease) + require.NotNil(t, newAPIError) + assert.Equal(t, http.StatusTooManyRequests, newAPIError.StatusCode) + assert.Equal(t, types.ErrorCodeChannelCapacityExhausted, newAPIError.GetErrorCode()) + assert.Equal(t, "1", response.Header().Get("Retry-After")) +} + +func TestChannelCapacityAPIErrorIsLocalAndNonRetryable(t *testing.T) { + require.NoError(t, appI18n.Init()) + gin.SetMode(gin.TestMode) + response := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(response) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + err := channelCapacityAPIError(ctx, 1500*time.Millisecond) + + assert.Equal(t, http.StatusTooManyRequests, err.StatusCode) + assert.Equal(t, types.ErrorCodeChannelCapacityExhausted, err.GetErrorCode()) + assert.Equal(t, "2", response.Header().Get("Retry-After")) + assert.True(t, types.IsSkipRetryError(err)) + assert.False(t, types.IsChannelError(err)) + assert.False(t, types.IsRecordErrorLog(err)) + assert.False(t, service.ShouldDisableChannel(err)) +} + +func TestShouldRetryTaskRelaySkipsChannelCapacity(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil) + + capacityErr := &taskdto.TaskError{ + Code: string(types.ErrorCodeChannelCapacityExhausted), + StatusCode: http.StatusTooManyRequests, + LocalError: true, + } + upstreamRateLimit := &taskdto.TaskError{ + Code: "upstream_rate_limit", + StatusCode: http.StatusTooManyRequests, + } + + assert.False(t, shouldRetryTaskRelay(ctx, 1, capacityErr, 1)) + assert.True(t, shouldRetryTaskRelay(ctx, 1, upstreamRateLimit, 1)) +} diff --git a/controller/relay.go b/controller/relay.go index 8dccfe76dddd..628d8dda8f32 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -12,6 +12,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" taskdto "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" @@ -190,15 +191,18 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } relayInfo.RetryIndex = 0 relayInfo.LastError = nil + var activeAdmissionLease *service.ChannelAdmissionLease + defer func() { releaseChannelAdmissionLease(c, activeAdmissionLease) }() for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { relayInfo.RetryIndex = retryParam.GetRetry() - channel, channelErr := getChannel(c, relayInfo, retryParam) + channel, admissionLease, channelErr := getChannel(c, relayInfo, retryParam) if channelErr != nil { logger.LogError(c, channelErr.Error()) newAPIError = channelErr break } + activeAdmissionLease = admissionLease addUsedChannel(c, channel.Id) if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil { newAPIError = billingErr @@ -216,6 +220,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { break } c.Request.Body = io.NopCloser(bodyStorage) + activeAdmissionLease.Commit() switch relayFormat { case types.RelayFormatOpenAIRealtime: @@ -227,6 +232,8 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { default: newAPIError = relayHandler(c, relayInfo) } + releaseChannelAdmissionLease(c, activeAdmissionLease) + activeAdmissionLease = nil if newAPIError == nil { relayInfo.LastError = nil @@ -297,35 +304,84 @@ func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta { return meta } -func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) { +func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *service.ChannelAdmissionLease, *types.NewAPIError) { if info.ChannelMeta == nil { - autoBan := c.GetBool("auto_ban") - autoBanInt := 1 - if !autoBan { - autoBanInt = 0 - } - return &model.Channel{ - Id: c.GetInt("channel_id"), - Type: c.GetInt("channel_type"), - Name: c.GetString("channel_name"), - AutoBan: &autoBanInt, - }, nil - } - channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam) + lease := service.GetChannelAdmissionLease(c) + channelSetting, hasChannelSetting := common.GetContextKeyType[dto.ChannelSettings](c, constant.ContextKeyChannelSetting) + if hasChannelSetting && (lease != nil || (channelSetting.MaxConcurrency <= 0 && channelSetting.RPMLimit <= 0)) { + autoBanInt := 1 + if !common.GetContextKeyBool(c, constant.ContextKeyChannelAutoBan) { + autoBanInt = 0 + } + return &model.Channel{ + Id: common.GetContextKeyInt(c, constant.ContextKeyChannelId), + Type: common.GetContextKeyInt(c, constant.ContextKeyChannelType), + Name: common.GetContextKeyString(c, constant.ContextKeyChannelName), + AutoBan: &autoBanInt, + ChannelInfo: model.ChannelInfo{ + IsMultiKey: common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey), + }, + }, lease, nil + } + + channelID := common.GetContextKeyInt(c, constant.ContextKeyChannelId) + channel, err := model.CacheGetChannel(channelID) + if err != nil { + return nil, nil, types.NewError(fmt.Errorf("get initial channel #%d: %w", channelID, err), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + } + var decision service.ChannelAdmissionDecision + lease, decision, err = service.AcquireChannelAdmission(c.Request.Context(), channel) + if err != nil { + return nil, nil, types.NewError(fmt.Errorf("acquire initial channel #%d admission: %w", channelID, err), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + } + if !decision.Allowed { + return nil, nil, channelCapacityAPIError(c, decision.RetryAfter) + } + return channel, lease, nil + } + + selection, err := service.SelectChannelWithAdmission(retryParam) if err != nil { - return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + var capacityErr *service.ChannelCapacityError + if errors.As(err, &capacityErr) { + return nil, nil, channelCapacityAPIError(c, time.Duration(capacityErr.RetryAfterSeconds())*time.Second) + } + return nil, nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", retryParam.TokenGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) } - if channel == nil { - return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在（retry）", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + if selection == nil || selection.Channel == nil { + return nil, nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在（retry）", retryParam.TokenGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) } info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info) - - newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName) + newAPIError := middleware.SetupContextForSelectedChannel(c, selection.Channel, info.OriginModelName) if newAPIError != nil { - return nil, newAPIError + if releaseErr := selection.Lease.Release(); releaseErr != nil { + logger.LogWarn(c, fmt.Sprintf("release channel admission after setup failure: %v", releaseErr)) + } + return nil, nil, newAPIError + } + return selection.Channel, selection.Lease, nil +} + +func channelCapacityAPIError(c *gin.Context, retryAfter time.Duration) *types.NewAPIError { + retryAfterSeconds := int64((retryAfter + time.Second - 1) / time.Second) + if retryAfterSeconds < 1 { + retryAfterSeconds = 1 + } + c.Header("Retry-After", fmt.Sprintf("%d", retryAfterSeconds)) + return types.NewErrorWithStatusCode( + errors.New(i18n.T(c, i18n.MsgDistributorChannelCapacityExceeded)), + types.ErrorCodeChannelCapacityExhausted, + http.StatusTooManyRequests, + types.ErrOptionWithSkipRetry(), + types.ErrOptionWithNoRecordErrorLog(), + ) +} + +func releaseChannelAdmissionLease(c *gin.Context, lease *service.ChannelAdmissionLease) { + if err := lease.Release(); err != nil { + logger.LogWarn(c, fmt.Sprintf("release channel admission lease failed: %v", err)) } - return channel, nil } func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { @@ -520,13 +576,36 @@ func RelayTask(c *gin.Context) { RequestPath: c.Request.URL.Path, Retry: common.GetPointer(0), } + var activeAdmissionLease *service.ChannelAdmissionLease + defer func() { releaseChannelAdmissionLease(c, activeAdmissionLease) }() for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { var channel *model.Channel + var admissionLease *service.ChannelAdmissionLease if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil { channel = lockedCh - if retryParam.GetRetry() > 0 { + initialLease := service.GetChannelAdmissionLease(c) + if retryParam.GetRetry() == 0 && initialLease != nil && initialLease.ChannelID() == channel.Id { + admissionLease = initialLease + } else { + if retryParam.GetRetry() == 0 && initialLease != nil { + releaseChannelAdmissionLease(c, initialLease) + } + lease, decision, admissionErr := service.AcquireChannelAdmission(c.Request.Context(), channel) + if admissionErr != nil { + taskErr = service.TaskErrorWrapperLocal(admissionErr, "channel_admission_failed", http.StatusServiceUnavailable) + break + } + if !decision.Allowed { + capacityErr := channelCapacityAPIError(c, decision.RetryAfter) + taskErr = service.TaskErrorWrapperLocal(capacityErr.Err, string(capacityErr.GetErrorCode()), capacityErr.StatusCode) + break + } + admissionLease = lease + } + activeAdmissionLease = admissionLease + if retryParam.GetRetry() > 0 || common.GetContextKeyInt(c, constant.ContextKeyChannelId) != channel.Id { if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil { taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError) break @@ -534,12 +613,17 @@ func RelayTask(c *gin.Context) { } } else { var channelErr *types.NewAPIError - channel, channelErr = getChannel(c, relayInfo, retryParam) + channel, admissionLease, channelErr = getChannel(c, relayInfo, retryParam) if channelErr != nil { logger.LogError(c, channelErr.Error()) - taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError) + statusCode := channelErr.StatusCode + if statusCode <= 0 { + statusCode = http.StatusInternalServerError + } + taskErr = service.TaskErrorWrapperLocal(channelErr.Err, string(channelErr.GetErrorCode()), statusCode) break } + activeAdmissionLease = admissionLease } addUsedChannel(c, channel.Id) @@ -553,8 +637,11 @@ func RelayTask(c *gin.Context) { break } c.Request.Body = io.NopCloser(bodyStorage) + activeAdmissionLease.Commit() result, taskErr = relay.RelayTaskSubmit(c, relayInfo) + releaseChannelAdmissionLease(c, activeAdmissionLease) + activeAdmissionLease = nil if taskErr == nil { break } @@ -613,7 +700,7 @@ func RelayTask(c *gin.Context) { // respondTaskError 统一输出 Task 错误响应（含 429 限流提示改写） func respondTaskError(c *gin.Context, taskErr *taskdto.TaskError) { - if taskErr.StatusCode == http.StatusTooManyRequests { + if taskErr.StatusCode == http.StatusTooManyRequests && taskErr.Code != string(types.ErrorCodeChannelCapacityExhausted) { taskErr.Message = "当前分组上游负载已饱和，请稍后再试" } c.JSON(taskErr.StatusCode, taskErr) @@ -632,6 +719,9 @@ func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *taskdto.TaskEr if _, ok := c.Get("specific_channel_id"); ok { return false } + if taskErr.Code == string(types.ErrorCodeChannelCapacityExhausted) { + return false + } if taskErr.StatusCode == http.StatusTooManyRequests { return true } diff --git a/docs/channel/other_setting.md b/docs/channel/other_setting.md index 4eb8e7755cd8..eb43edaa3c25 100644 --- a/docs/channel/other_setting.md +++ b/docs/channel/other_setting.md @@ -1,6 +1,6 @@ # 渠道额外设置说明 -该配置用于设置一些额外的渠道参数，可以通过 JSON 对象进行配置。主要包含以下三个设置项： +该配置用于设置一些额外的渠道参数，可以通过 JSON 对象进行配置。常用设置项包括： 1. force_format - 用于标识是否对数据进行强制格式化为 OpenAI 格式 @@ -16,6 +16,16 @@ - 用于标识是否将思考内容`reasoning_content`转换为``标签拼接到内容中返回 - 类型为布尔值，设置为 true 时启用思考内容转换 +4. max_concurrency + - 限制该渠道同时在途的中转请求数 + - 类型为非负整数，`0` 或不填写表示不限 + - 多 Key 渠道的所有 Key 共享同一个渠道并发上限 + +5. rpm_limit + - 限制该渠道在滚动 60 秒窗口内开始的中转请求数 + - 类型为非负整数，`0` 或不填写表示不限 + - 多 Key 渠道的所有 Key 共享同一个渠道 RPM 上限 + -------------------------------------------------------------- ## JSON 格式示例 @@ -26,16 +36,30 @@ { "force_format": true, "thinking_to_content": true, - "proxy": "socks5://proxy.example:1080" + "proxy": "socks5://proxy.example:1080", + "max_concurrency": 20, + "rpm_limit": 120 } ``` -------------------------------------------------------------- -通过调整上述 JSON 配置中的值，可以灵活控制渠道的额外行为，比如是否进行格式化以及使用特定的网络代理。 +通过调整上述 JSON 配置中的值，可以灵活控制渠道的额外行为，比如是否进行格式化、使用特定网络代理，以及限制渠道容量。 + +## 渠道容量限制语义 + +- 容量检查和预留发生在发送上游请求之前。候选渠道达到并发或 RPM 上限时，路由器会先尝试同优先级的其他渠道，再尝试较低优先级或下一个自动分组。 +- Token 固定渠道和渠道亲和绑定不会因容量不足而静默改道；绑定渠道满载时直接返回本地容量错误，原亲和绑定保持不变。 +- 本地容量跳过不占用上游重试次数，不触发渠道自动禁用，也不会作为上游错误记录。所有候选均满载时返回 `429 channel_capacity_exhausted` 和 `Retry-After`；本功能不提供请求排队。 +- 并发租约覆盖流式、非流式、WebSocket 和任务提交的完整上游调用。正常结束、错误、取消和 panic 都会释放；Redis 租约带过期保护，长请求会自动续租。 +- RPM 在候选预留后、真正开始调用上游前仍可回退，例如本地请求校验或渠道上下文初始化失败。上游调用一旦开始，RPM 计数不会因成功、失败、取消或重试而回退。 +- 启用 Redis 时，并发与 RPM 在所有实例间全局共享，并通过原子脚本同时检查和预留。未配置 Redis 时使用进程内原子限制，每个实例独立计算；Redis 运行时故障会降级到进程内模式并记录告警，此时不再保证跨实例全局上限。 +- 修改限制后，新请求立即按新值判断。降低并发上限不会中断现有请求，而是暂停该渠道的新准入，直到在途请求数回落。 ## 升级兼容性 +`max_concurrency` 和 `rpm_limit` 直接存放在现有渠道设置 JSON 中，不需要数据库迁移。已有渠道缺省这两个字段时保持不限流，因此升级本身不会改变现有路由行为。 + 旧版本会忽略代理地址中的 path、query 和 fragment。为避免升级后中断已有渠道流量，运行时会继续剥离这些遗留后缀，并对同一代理地址每个进程记录一次不含凭证和后缀的警告。该兼容逻辑不会改写数据库；再次保存渠道时必须按上述严格规则修正代理地址。 代理连接使用 30 秒 TCP 拨号超时和 30 秒 KeepAlive；TLS 握手超时为 10 秒。这些超时同样适用于未配置渠道代理的中转请求。 diff --git a/i18n/keys.go b/i18n/keys.go index 64a835e1a942..c0f4f42f5995 100644 --- a/i18n/keys.go +++ b/i18n/keys.go @@ -322,6 +322,7 @@ const ( MsgDistributorGroupAccessDenied = "distributor.group_access_denied" MsgDistributorGetChannelFailed = "distributor.get_channel_failed" MsgDistributorNoAvailableChannel = "distributor.no_available_channel" + MsgDistributorChannelCapacityExceeded = "distributor.channel_capacity_exhausted" MsgDistributorInvalidMidjourney = "distributor.invalid_midjourney_request" MsgDistributorInvalidParseModel = "distributor.invalid_request_parse_model" ) diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index c533daecc32d..36ce7c3e1cc1 100644 --- a/i18n/locales/en.yaml +++ b/i18n/locales/en.yaml @@ -272,6 +272,7 @@ distributor.invalid_playground_request: "Invalid playground request: {{.Error}}" distributor.group_access_denied: "No permission to access this group" distributor.get_channel_failed: "Failed to get available channel for model {{.Model}} under group {{.Group}} (distributor): {{.Error}}" distributor.no_available_channel: "No available channel for model {{.Model}} under group {{.Group}} (distributor)" +distributor.channel_capacity_exhausted: "All matching channels are at their configured capacity. Please retry later." distributor.invalid_midjourney_request: "Invalid Midjourney request: {{.Error}}" distributor.invalid_request_parse_model: "Invalid request, unable to parse model" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index a2f5275be9a8..dd53ee27b997 100644 --- a/i18n/locales/zh-CN.yaml +++ b/i18n/locales/zh-CN.yaml @@ -273,6 +273,7 @@ distributor.invalid_playground_request: "无效的playground请求，{{.Error}}" distributor.group_access_denied: "无权访问该分组" distributor.get_channel_failed: "获取分组 {{.Group}} 下模型 {{.Model}} 的可用渠道失败（distributor）：{{.Error}}" distributor.no_available_channel: "分组 {{.Group}} 下模型 {{.Model}} 无可用渠道（distributor）" +distributor.channel_capacity_exhausted: "匹配的渠道均已达到配置的容量上限，请稍后重试" distributor.invalid_midjourney_request: "无效的midjourney请求，{{.Error}}" distributor.invalid_request_parse_model: "无效的请求，无法解析模型" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index 84ebd57ed587..9d626b8918c0 100644 --- a/i18n/locales/zh-TW.yaml +++ b/i18n/locales/zh-TW.yaml @@ -273,6 +273,7 @@ distributor.invalid_playground_request: "無效的playground請求，{{.Error}}" distributor.group_access_denied: "無權存取該分組" distributor.get_channel_failed: "獲取分組 {{.Group}} 下模型 {{.Model}} 的可用管道失敗（distributor）：{{.Error}}" distributor.no_available_channel: "分組 {{.Group}} 下模型 {{.Model}} 無可用管道（distributor）" +distributor.channel_capacity_exhausted: "符合條件的渠道均已達到設定的容量上限，請稍後重試" distributor.invalid_midjourney_request: "無效的midjourney請求，{{.Error}}" distributor.invalid_request_parse_model: "無效的請求，無法解析模型" diff --git a/middleware/channel_admission_test.go b/middleware/channel_admission_test.go new file mode 100644 index 000000000000..48f95e5449ef --- /dev/null +++ b/middleware/channel_admission_test.go @@ -0,0 +1,349 @@ +package middleware + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + appI18n "github.com/QuantumNous/new-api/i18n" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/service" + + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" + "gorm.io/gorm" +) + +func TestDistributorRejectsSaturatedSpecificChannelBeforeDownstream(t *testing.T) { + require.NoError(t, appI18n.Init()) + dsn := fmt.Sprintf("file:distributor-specific-admission-%d?mode=memory&cache=shared", time.Now().UnixNano()) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{})) + + previousDB := model.DB + previousMemoryCache := common.MemoryCacheEnabled + previousRedisEnabled := common.RedisEnabled + model.DB = db + common.MemoryCacheEnabled = false + common.RedisEnabled = false + t.Cleanup(func() { + model.DB = previousDB + common.MemoryCacheEnabled = previousMemoryCache + common.RedisEnabled = previousRedisEnabled + if previousMemoryCache && previousDB != nil { + model.InitChannelCache() + } + }) + + priority := int64(0) + weight := uint(1) + channel := &model.Channel{ + Id: 501, + Name: "specific", + Type: constant.ChannelTypeOpenAI, + Key: "test-key", + Status: common.ChannelStatusEnabled, + Models: "gpt-test", + Group: "default", + Priority: &priority, + Weight: &weight, + } + channel.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + require.NoError(t, db.Create(channel).Error) + + activeLease, decision, err := service.AcquireChannelAdmission(context.Background(), channel) + require.NoError(t, err) + require.True(t, decision.Allowed) + activeLease.Commit() + defer func() { require.NoError(t, activeLease.Release()) }() + + gin.SetMode(gin.TestMode) + downstreamCalled := false + router := gin.New() + router.Use(func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyTokenSpecificChannelId, "501") + c.Next() + }) + router.Use(Distribute()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + downstreamCalled = true + c.Status(http.StatusNoContent) + }) + + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-test"}`)) + request.Header.Set("Content-Type", "application/json") + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + assert.Equal(t, http.StatusTooManyRequests, response.Code) + assert.Equal(t, "1", response.Header().Get("Retry-After")) + assert.Equal(t, "channel_capacity_exhausted", gjson.Get(response.Body.String(), "error.code").String()) + assert.False(t, downstreamCalled) +} + +func TestDistributorRejectsNonStringSpecificChannelID(t *testing.T) { + require.NoError(t, appI18n.Init()) + gin.SetMode(gin.TestMode) + downstreamCalled := false + router := gin.New() + router.Use(func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyTokenSpecificChannelId, 501) + c.Next() + }) + router.Use(Distribute()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + downstreamCalled = true + c.Status(http.StatusNoContent) + }) + + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-test"}`)) + request.Header.Set("Content-Type", "application/json") + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + assert.Equal(t, http.StatusBadRequest, response.Code) + assert.False(t, downstreamCalled) +} + +func TestDistributorReleasesAdmissionAcrossDownstreamExitPaths(t *testing.T) { + require.NoError(t, appI18n.Init()) + dsn := fmt.Sprintf("file:distributor-lifecycle-admission-%d?mode=memory&cache=shared", time.Now().UnixNano()) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{})) + + previousDB := model.DB + previousMemoryCache := common.MemoryCacheEnabled + previousRedisEnabled := common.RedisEnabled + model.DB = db + common.MemoryCacheEnabled = false + common.RedisEnabled = false + t.Cleanup(func() { + model.DB = previousDB + common.MemoryCacheEnabled = previousMemoryCache + common.RedisEnabled = previousRedisEnabled + if previousMemoryCache && previousDB != nil { + model.InitChannelCache() + } + }) + + priority := int64(0) + weight := uint(1) + channel := &model.Channel{ + Id: 504, + Name: "lifecycle", + Type: constant.ChannelTypeOpenAI, + Key: "test-key", + Status: common.ChannelStatusEnabled, + Models: "gpt-test", + Group: "default", + Priority: &priority, + Weight: &weight, + } + channel.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + require.NoError(t, db.Create(channel).Error) + + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(gin.Recovery()) + router.Use(func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyTokenSpecificChannelId, "504") + c.Next() + }) + router.Use(Distribute()) + var cancelDownstream context.CancelFunc + router.POST("/v1/chat/completions", func(c *gin.Context) { + lease := service.GetChannelAdmissionLease(c) + require.NotNil(t, lease) + lease.Commit() + switch c.GetHeader("X-Test-Exit") { + case "error": + c.Status(http.StatusBadGateway) + case "cancel": + cancelDownstream() + <-c.Request.Context().Done() + c.Status(499) + case "timeout": + <-c.Request.Context().Done() + c.Status(499) + case "stream": + c.Header("Content-Type", "text/event-stream") + _, _ = c.Writer.WriteString("data: done\n\n") + case "panic": + panic("test downstream panic") + default: + c.Status(http.StatusNoContent) + } + }) + + tests := []struct { + name string + exit string + wantStatus int + context func() (context.Context, context.CancelFunc) + }{ + {name: "success", wantStatus: http.StatusNoContent}, + {name: "upstream error", exit: "error", wantStatus: http.StatusBadGateway}, + {name: "cancellation", exit: "cancel", wantStatus: 499, context: func() (context.Context, context.CancelFunc) { + ctx, cancel := context.WithCancel(context.Background()) + cancelDownstream = cancel + return ctx, cancel + }}, + {name: "timeout", exit: "timeout", wantStatus: 499, context: func() (context.Context, context.CancelFunc) { + return context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) + }}, + {name: "stream end", exit: "stream", wantStatus: http.StatusOK}, + {name: "panic", exit: "panic", wantStatus: http.StatusInternalServerError}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + requestContext := context.Background() + cancel := func() {} + if test.context != nil { + requestContext, cancel = test.context() + } + defer cancel() + + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-test"}`)).WithContext(requestContext) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("X-Test-Exit", test.exit) + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + assert.Equal(t, test.wantStatus, response.Code) + snapshot, snapshotErr := service.GetChannelAdmissionSnapshot(context.Background(), channel) + require.NoError(t, snapshotErr) + assert.Equal(t, 0, snapshot.CurrentConcurrency) + }) + } +} + +func TestDistributorAllowsRoutesThatResolveTheirChannelDownstream(t *testing.T) { + gin.SetMode(gin.TestMode) + downstreamCalled := false + originalModelSet := false + router := gin.New() + router.Use(Distribute()) + router.GET("/suno/fetch/:id", func(c *gin.Context) { + downstreamCalled = true + _, originalModelSet = c.Get(string(constant.ContextKeyOriginalModel)) + c.Status(http.StatusNoContent) + }) + + request := httptest.NewRequest(http.MethodGet, "/suno/fetch/task-1", nil) + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + assert.Equal(t, http.StatusNoContent, response.Code) + assert.True(t, downstreamCalled) + assert.True(t, originalModelSet) +} + +func TestDistributorRejectsSaturatedAffinityAndPreservesBinding(t *testing.T) { + require.NoError(t, appI18n.Init()) + dsn := fmt.Sprintf("file:distributor-reroute-admission-%d?mode=memory&cache=shared", time.Now().UnixNano()) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{})) + + previousDB := model.DB + previousMemoryCache := common.MemoryCacheEnabled + previousRedisEnabled := common.RedisEnabled + model.DB = db + common.MemoryCacheEnabled = true + common.RedisEnabled = false + t.Cleanup(func() { + model.DB = previousDB + common.MemoryCacheEnabled = previousMemoryCache + common.RedisEnabled = previousRedisEnabled + if previousMemoryCache && previousDB != nil { + model.InitChannelCache() + } + }) + + priority := int64(10) + weight := uint(1) + channels := []model.Channel{ + {Id: 502, Name: "full", Type: constant.ChannelTypeOpenAI, Key: "full-key", Status: common.ChannelStatusEnabled, Models: "gpt-test", Group: "default", Priority: &priority, Weight: &weight}, + {Id: 503, Name: "available", Type: constant.ChannelTypeOpenAI, Key: "available-key", Status: common.ChannelStatusEnabled, Models: "gpt-test", Group: "default", Priority: &priority, Weight: &weight}, + } + for index := range channels { + channels[index].SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + require.NoError(t, db.Create(&channels[index]).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: "default", + Model: "gpt-test", + ChannelId: channels[index].Id, + Enabled: true, + Priority: &priority, + Weight: weight, + }).Error) + } + model.InitChannelCache() + service.ClearChannelAffinityCacheAll() + t.Cleanup(func() { service.ClearChannelAffinityCacheAll() }) + + requestBody := `{"model":"gpt-test","prompt_cache_key":"sticky-capacity-test"}` + seedResponse := httptest.NewRecorder() + seedContext, _ := gin.CreateTestContext(seedResponse) + seedContext.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(requestBody)) + seedContext.Request.Header.Set("Content-Type", "application/json") + _, found := service.GetPreferredChannelByAffinity(seedContext, "gpt-test", "default") + require.False(t, found) + service.RecordChannelAffinity(seedContext, channels[0].Id) + + activeLease, decision, err := service.AcquireChannelAdmission(context.Background(), &channels[0]) + require.NoError(t, err) + require.True(t, decision.Allowed) + activeLease.Commit() + defer func() { require.NoError(t, activeLease.Release()) }() + + gin.SetMode(gin.TestMode) + downstreamCalled := false + router := gin.New() + router.Use(func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyUsingGroup, "default") + common.SetContextKey(c, constant.ContextKeyUserGroup, "default") + c.Next() + }) + router.Use(Distribute()) + router.POST("/v1/responses", func(c *gin.Context) { + downstreamCalled = true + c.Status(http.StatusNoContent) + }) + + request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(requestBody)) + request.Header.Set("Content-Type", "application/json") + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + assert.Equal(t, http.StatusTooManyRequests, response.Code) + assert.Equal(t, "channel_capacity_exhausted", gjson.Get(response.Body.String(), "error.code").String()) + assert.False(t, downstreamCalled) + checkResponse := httptest.NewRecorder() + checkContext, _ := gin.CreateTestContext(checkResponse) + checkContext.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(requestBody)) + checkContext.Request.Header.Set("Content-Type", "application/json") + preferredChannelID, found := service.GetPreferredChannelByAffinity(checkContext, "gpt-test", "default") + require.True(t, found) + assert.Equal(t, 502, preferredChannelID) + + fullSnapshot, err := service.GetChannelAdmissionSnapshot(context.Background(), &channels[0]) + require.NoError(t, err) + availableSnapshot, err := service.GetChannelAdmissionSnapshot(context.Background(), &channels[1]) + require.NoError(t, err) + assert.Equal(t, 1, fullSnapshot.CurrentConcurrency) + assert.Equal(t, 0, availableSnapshot.CurrentConcurrency) +} diff --git a/middleware/distributor.go b/middleware/distributor.go index 7decf0e28728..865d09733c42 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -14,6 +14,7 @@ import ( "github.com/QuantumNous/new-api/constant" taskdto "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/i18n" + "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/model" relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relaykit/dto" @@ -33,15 +34,27 @@ type ModelRequest struct { func Distribute() func(c *gin.Context) { return func(c *gin.Context) { var channel *model.Channel - channelId, ok := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId) + var admissionLease *service.ChannelAdmissionLease + defer func() { + if err := admissionLease.Release(); err != nil { + logger.LogWarn(c.Request.Context(), fmt.Sprintf("release initial channel admission lease failed: %v", err)) + } + }() + + channelID, specificChannel := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId) modelRequest, shouldSelectChannel, err := getModelRequest(c) if err != nil { abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()})) return } - if ok { - id, err := strconv.Atoi(channelId.(string)) - if err != nil { + if specificChannel { + rawChannelID, isString := channelID.(string) + if !isString { + abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId)) + return + } + id, parseErr := strconv.Atoi(rawChannelID) + if parseErr != nil { abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId)) return } @@ -54,24 +67,30 @@ func Distribute() func(c *gin.Context) { abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled)) return } + var decision service.ChannelAdmissionDecision + admissionLease, decision, err = service.AcquireChannelAdmission(c.Request.Context(), channel) + if err != nil { + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, err.Error(), types.ErrorCodeGetChannelFailed) + return + } + if !decision.Allowed { + abortWithChannelCapacity(c, decision.RetryAfter) + return + } } 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) - if !ok { - // token model limit is empty, all models are not allowed + s, exists := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) + if !exists { abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess)) return } - var tokenModelLimit map[string]bool - tokenModelLimit, ok = s.(map[string]bool) - if !ok { + tokenModelLimit, valid := s.(map[string]bool) + if !valid { tokenModelLimit = map[string]bool{} } - matchName := ratio_setting.FormatMatchingModelName(modelRequest.Model) // match gpts & thinking-* - if _, ok := tokenModelLimit[matchName]; !ok { + matchName := ratio_setting.FormatMatchingModelName(modelRequest.Model) + if _, allowed := tokenModelLimit[matchName]; !allowed { abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) return } @@ -82,9 +101,7 @@ func Distribute() func(c *gin.Context) { 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) @@ -104,28 +121,37 @@ func Distribute() func(c *gin.Context) { if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { affinityUsable := false - preferred, err := model.CacheGetChannel(preferredChannelID) - if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled && + preferred, preferredErr := model.CacheGetChannel(preferredChannelID) + preferredGroup := "" + if preferredErr == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled && channelSupportsRequestPath(preferred, c.Request.URL.Path, modelRequest.Model) { if usingGroup == "auto" { userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) - autoGroups := service.GetRequestAutoGroups(c, userGroup) - for _, g := range autoGroups { - if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { - selectGroup = g - common.SetContextKey(c, constant.ContextKeyAutoGroup, g) - channel = preferred - affinityUsable = true - service.MarkChannelAffinityUsed(c, g, preferred.Id) + for _, group := range service.GetRequestAutoGroups(c, userGroup) { + if model.IsChannelEnabledForGroupModel(group, modelRequest.Model, preferred.Id) { + preferredGroup = group break } } } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) { - channel = preferred - selectGroup = usingGroup - affinityUsable = true - service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) + preferredGroup = usingGroup + } + } + if preferredGroup != "" { + lease, decision, admissionErr := service.AcquireChannelAdmission(c.Request.Context(), preferred) + if admissionErr != nil { + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, admissionErr.Error(), types.ErrorCodeGetChannelFailed) + return } + if !decision.Allowed { + abortWithChannelCapacity(c, decision.RetryAfter) + return + } + channel = preferred + admissionLease = lease + affinityUsable = true + common.SetContextKey(c, constant.ContextKeyAutoGroup, preferredGroup) + service.MarkChannelAffinityUsed(c, preferredGroup, preferred.Id) } if !affinityUsable && !service.ShouldKeepChannelAffinityOnChannelDisabled() { service.ClearCurrentChannelAffinityCache(c) @@ -133,43 +159,71 @@ func Distribute() func(c *gin.Context) { } if channel == nil { - channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ + selection, selectionErr := service.SelectChannelWithAdmission(&service.RetryParam{ Ctx: c, ModelName: modelRequest.Model, TokenGroup: usingGroup, RequestPath: c.Request.URL.Path, Retry: common.GetPointer(0), }) - if err != nil { - showGroup := usingGroup - if usingGroup == "auto" { - showGroup = fmt.Sprintf("auto(%s)", selectGroup) + if selectionErr != nil { + var capacityErr *service.ChannelCapacityError + if errors.As(selectionErr, &capacityErr) { + abortWithChannelCapacity(c, capacityErr.RetryAfter) + return } - 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) + message := i18n.T(c, i18n.MsgDistributorGetChannelFailed, map[string]any{"Group": usingGroup, "Model": modelRequest.Model, "Error": selectionErr.Error()}) + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, types.ErrorCodeGetChannelFailed) return } - if channel == nil { + if selection == nil || selection.Channel == nil { abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound) return } + channel = selection.Channel + admissionLease = selection.Lease } } } + common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now()) - SetupContextForSelectedChannel(c, channel, modelRequest.Model) + if channel == nil { + common.SetContextKey(c, constant.ContextKeyOriginalModel, modelRequest.Model) + } else { + if setupErr := SetupContextForSelectedChannel(c, channel, modelRequest.Model); setupErr != nil { + statusCode := setupErr.StatusCode + if statusCode <= 0 { + statusCode = http.StatusInternalServerError + } + abortWithOpenAiMessage(c, statusCode, setupErr.Error(), setupErr.GetErrorCode()) + return + } + service.SetChannelAdmissionLease(c, admissionLease) + } c.Next() if channel != nil && c.Writer != nil && c.Writer.Status() < http.StatusBadRequest { - service.RecordChannelAffinity(c, channel.Id) + finalChannelID := common.GetContextKeyInt(c, constant.ContextKeyChannelId) + if finalChannelID > 0 { + service.RecordChannelAffinity(c, finalChannelID) + } } } } +func abortWithChannelCapacity(c *gin.Context, retryAfter time.Duration) { + retryAfterSeconds := int64((retryAfter + time.Second - 1) / time.Second) + if retryAfterSeconds < 1 { + retryAfterSeconds = 1 + } + c.Header("Retry-After", strconv.FormatInt(retryAfterSeconds, 10)) + abortWithOpenAiMessage( + c, + http.StatusTooManyRequests, + i18n.T(c, i18n.MsgDistributorChannelCapacityExceeded), + types.ErrorCodeChannelCapacityExhausted, + ) +} + // channelSupportsRequestPath reports whether a channel can serve the request path. // Only Advanced Custom (type 58) channels are path-checked; all other channel types // always pass. A type-58 channel is usable only when one of its routes matches. diff --git a/model/ability.go b/model/ability.go index d950a6adbfc4..18fe8fdf822a 100644 --- a/model/ability.go +++ b/model/ability.go @@ -3,12 +3,12 @@ package model import ( "errors" "fmt" + "sort" "strings" "sync" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" - "github.com/QuantumNous/new-api/relaykit/dto" "github.com/samber/lo" "gorm.io/gorm" @@ -60,137 +60,81 @@ func GetAllEnableAbilities() []Ability { return abilities } -func getPriority(group string, model string, retry int) (int, error) { - - var priorities []int - err := DB.Model(&Ability{}). - Select("DISTINCT(priority)"). - Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true). - Order("priority DESC"). // 按优先级降序排序 - Pluck("priority", &priorities).Error // Pluck用于将查询的结果直接扫描到一个切片中 - - if err != nil { - // 处理错误 - return 0, err - } - - if len(priorities) == 0 { - // 如果没有查询到优先级，则返回错误 - return 0, errors.New("数据库一致性被破坏") +func getDatabaseSatisfiedChannelTiers(group string, modelName string, requestPath string) ([]ChannelCandidateTier, error) { + var abilities []Ability + if err := DB.Where(&Ability{Group: group, Model: modelName, Enabled: true}). + Order("priority DESC"). + Find(&abilities).Error; err != nil { + return nil, err } - - // 确定要使用的优先级 - var priorityToUse int - if retry >= len(priorities) { - // 如果重试次数大于优先级数，则使用最小的优先级 - priorityToUse = priorities[len(priorities)-1] - } else { - priorityToUse = priorities[retry] + if len(abilities) == 0 { + return nil, nil } - return priorityToUse, nil -} -func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { - maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true) - channelQuery := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = (?)", group, model, true, maxPrioritySubQuery) - if retry != 0 { - priority, err := getPriority(group, model, retry) - if err != nil { - return nil, err - } else { - channelQuery = DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = ?", group, model, true, priority) - } + channelIDs := make([]int, 0, len(abilities)) + for _, ability := range abilities { + channelIDs = append(channelIDs, ability.ChannelId) } - - return channelQuery, nil -} - -func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) { - var abilities []Ability - - var err error = nil - channelQuery, err := getChannelQuery(group, model, retry) - if err != nil { + var channels []*Channel + if err := DB.Where("id IN ?", channelIDs).Find(&channels).Error; err != nil { return nil, err } - if common.UsingMainDatabase(common.DatabaseTypeSQLite) || common.UsingMainDatabase(common.DatabaseTypePostgreSQL) { - err = channelQuery.Order("weight DESC").Find(&abilities).Error - } else { - err = channelQuery.Order("weight DESC").Find(&abilities).Error - } - if err != nil { - return nil, err + channelsByID := make(map[int]*Channel, len(channels)) + for _, channel := range channels { + channelsByID[channel.Id] = channel } - abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, model) - channel := Channel{} - if len(abilities) > 0 { - // Randomly choose one - weightSum := uint(0) - for _, ability_ := range abilities { - weightSum += ability_.Weight + 10 + + tiersByPriority := make(map[int64][]ChannelCandidate) + priorities := make([]int64, 0) + for _, ability := range abilities { + channel, ok := channelsByID[ability.ChannelId] + if !ok { + return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", ability.ChannelId) } - // Randomly choose one - weight := common.GetRandomInt(int(weightSum)) - for _, ability_ := range abilities { - weight -= int(ability_.Weight) + 10 - //log.Printf("weight: %d, ability weight: %d", weight, *ability_.Weight) - if weight <= 0 { - channel.Id = ability_.ChannelId - break + if requestPath != "" && channel.Type == constant.ChannelTypeAdvancedCustom { + config := channel.GetOtherSettings().AdvancedCustom + if config == nil || !config.SupportsPathForModel(requestPath, modelName) { + continue } } - } else { - return nil, nil - } - err = DB.First(&channel, "id = ?", channel.Id).Error - return &channel, err -} -// filterAbilitiesByRequestPathAndModel restricts candidates by request path and -// model for the DB (non-memory-cache) selection path. Only Advanced Custom -// (type 58) channels are path-checked: kept only when one of their routes matches -// requestPath and model; all other channel types always pass. When requestPath is -// empty, filtering is skipped. -func filterAbilitiesByRequestPathAndModel(abilities []Ability, requestPath string, model string) []Ability { - if requestPath == "" || len(abilities) == 0 { - return abilities + priority := int64(0) + if ability.Priority != nil { + priority = *ability.Priority + } + if _, exists := tiersByPriority[priority]; !exists { + priorities = append(priorities, priority) + } + tiersByPriority[priority] = append(tiersByPriority[priority], ChannelCandidate{ + Channel: channel, + Weight: int(ability.Weight) + 10, + }) } + sort.Slice(priorities, func(i, j int) bool { return priorities[i] > priorities[j] }) - channelIds := make([]int, 0, len(abilities)) - seen := make(map[int]struct{}, len(abilities)) - for _, ability := range abilities { - if _, ok := seen[ability.ChannelId]; ok { - continue + tiers := make([]ChannelCandidateTier, 0, len(priorities)) + for _, priority := range priorities { + candidates := tiersByPriority[priority] + if len(candidates) > 0 { + tiers = append(tiers, ChannelCandidateTier{Priority: priority, Candidates: candidates}) } - seen[ability.ChannelId] = struct{}{} - channelIds = append(channelIds, ability.ChannelId) } + return tiers, nil +} - var channels []*Channel - if err := DB.Where("id IN ?", channelIds).Find(&channels).Error; err != nil { - // On error, fall back to unfiltered candidates to avoid blocking selection - return abilities +func GetChannel(group string, modelName string, retry int, requestPath string) (*Channel, error) { + tiers, err := getDatabaseSatisfiedChannelTiers(group, modelName, requestPath) + if err != nil || len(tiers) == 0 { + return nil, err } - - advancedConfigs := make(map[int]*dto.AdvancedCustomConfig) - for _, channel := range channels { - if channel.Type == constant.ChannelTypeAdvancedCustom { - advancedConfigs[channel.Id] = channel.GetOtherSettings().AdvancedCustom - } + if retry < 0 { + retry = 0 } - - filtered := make([]Ability, 0, len(abilities)) - for _, ability := range abilities { - config, isAdvancedCustom := advancedConfigs[ability.ChannelId] - if !isAdvancedCustom { - filtered = append(filtered, ability) - continue - } - if config != nil && config.SupportsPathForModel(requestPath, model) { - filtered = append(filtered, ability) - } + if retry >= len(tiers) { + retry = len(tiers) - 1 } - return filtered + candidate, _ := PickWeightedChannelCandidate(tiers[retry].Candidates) + return candidate.Channel, nil } func (channel *Channel) AddAbilities(tx *gorm.DB) error { diff --git a/model/channel.go b/model/channel.go index 2cd7c3115ff6..d95cf99c9646 100644 --- a/model/channel.go +++ b/model/channel.go @@ -957,6 +957,9 @@ func (channel *Channel) ValidateSettings() error { if err := channelParams.ValidateHTTPTransport(); err != nil { return err } + if err := channelParams.ValidateAdmissionLimits(); err != nil { + return err + } channelOtherSettings := &dto.ChannelOtherSettings{} if channel.OtherSettings != "" { err := common.UnmarshalJsonStr(channel.OtherSettings, channelOtherSettings) diff --git a/model/channel_cache.go b/model/channel_cache.go index 86c594384d50..9dc204ac33dc 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -1,7 +1,6 @@ package model import ( - "errors" "fmt" "math/rand" "sort" @@ -111,101 +110,120 @@ func SyncChannelCache(frequency int) { } } -func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) { - // if memory cache is disabled, get channel directly from database +type ChannelCandidate struct { + Channel *Channel + Weight int +} + +type ChannelCandidateTier struct { + Priority int64 + Candidates []ChannelCandidate +} + +// GetSatisfiedChannelTiers loads every eligible channel once and groups them by +// descending priority. Candidate weights preserve the existing cache and direct +// database selection semantics. +func GetSatisfiedChannelTiers(group string, modelName string, requestPath string) ([]ChannelCandidateTier, error) { if !common.MemoryCacheEnabled { - return GetChannel(group, model, retry, requestPath) + return getDatabaseSatisfiedChannelTiers(group, modelName, requestPath) } channelSyncLock.RLock() defer channelSyncLock.RUnlock() - // First, try to find channels with the exact model name. - channels := filterChannelsByRequestPathAndModel(group2model2channels[group][model], requestPath, model) - - // If no channels found, try to find channels with the normalized model name. - if len(channels) == 0 { - normalizedModel := ratio_setting.FormatMatchingModelName(model) - channels = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, model) + channelIDs := filterChannelsByRequestPathAndModel(group2model2channels[group][modelName], requestPath, modelName) + if len(channelIDs) == 0 { + normalizedModel := ratio_setting.FormatMatchingModelName(modelName) + channelIDs = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, modelName) } - - if len(channels) == 0 { + if len(channelIDs) == 0 { return nil, nil } - if len(channels) == 1 { - if channel, ok := channelsIDM[channels[0]]; ok { - return channel, nil + channelsByPriority := make(map[int64][]*Channel) + priorities := make([]int64, 0) + for _, channelID := range channelIDs { + channel, ok := channelsIDM[channelID] + if !ok { + return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelID) + } + priority := channel.GetPriority() + if _, exists := channelsByPriority[priority]; !exists { + priorities = append(priorities, priority) } - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channels[0]) + channelsByPriority[priority] = append(channelsByPriority[priority], channel) } + sort.Slice(priorities, func(i, j int) bool { return priorities[i] > priorities[j] }) - uniquePriorities := make(map[int]bool) - for _, channelId := range channels { - if channel, ok := channelsIDM[channelId]; ok { - uniquePriorities[int(channel.GetPriority())] = true - } else { - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) + tiers := make([]ChannelCandidateTier, 0, len(priorities)) + for _, priority := range priorities { + channels := channelsByPriority[priority] + sumWeight := 0 + for _, channel := range channels { + sumWeight += channel.GetWeight() } - } - var sortedUniquePriorities []int - for priority := range uniquePriorities { - sortedUniquePriorities = append(sortedUniquePriorities, priority) - } - sort.Sort(sort.Reverse(sort.IntSlice(sortedUniquePriorities))) - if retry >= len(uniquePriorities) { - retry = len(uniquePriorities) - 1 - } - targetPriority := int64(sortedUniquePriorities[retry]) + smoothingFactor := 1 + smoothingAdjustment := 0 + if sumWeight == 0 { + smoothingAdjustment = 100 + } else if sumWeight/len(channels) < 10 { + smoothingFactor = 100 + } - // get the priority for the given retry number - var sumWeight = 0 - var targetChannels []*Channel - for _, channelId := range channels { - if channel, ok := channelsIDM[channelId]; ok { - if channel.GetPriority() == targetPriority { - sumWeight += channel.GetWeight() - targetChannels = append(targetChannels, channel) - } - } else { - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) + candidates := make([]ChannelCandidate, 0, len(channels)) + for _, channel := range channels { + candidates = append(candidates, ChannelCandidate{ + Channel: channel, + Weight: channel.GetWeight()*smoothingFactor + smoothingAdjustment, + }) } + tiers = append(tiers, ChannelCandidateTier{Priority: priority, Candidates: candidates}) } + return tiers, nil +} - if len(targetChannels) == 0 { - return nil, errors.New(fmt.Sprintf("no channel found, group: %s, model: %s, priority: %d", group, model, targetPriority)) +func PickWeightedChannelCandidate(candidates []ChannelCandidate) (ChannelCandidate, int) { + if len(candidates) == 0 { + return ChannelCandidate{}, -1 + } + totalWeight := 0 + for _, candidate := range candidates { + if candidate.Weight > 0 { + totalWeight += candidate.Weight + } } - - // smoothing factor and adjustment - smoothingFactor := 1 - smoothingAdjustment := 0 - - if sumWeight == 0 { - // when all channels have weight 0, set sumWeight to the number of channels and set smoothing adjustment to 100 - // each channel's effective weight = 100 - sumWeight = len(targetChannels) * 100 - smoothingAdjustment = 100 - } else if sumWeight/len(targetChannels) < 10 { - // when the average weight is less than 10, set smoothing factor to 100 - smoothingFactor = 100 + if totalWeight <= 0 { + index := rand.Intn(len(candidates)) + return candidates[index], index } - // Calculate the total weight of all channels up to endIdx - totalWeight := sumWeight * smoothingFactor - - // Generate a random value in the range [0, totalWeight) randomWeight := rand.Intn(totalWeight) - - // Find a channel based on its weight - for _, channel := range targetChannels { - randomWeight -= channel.GetWeight()*smoothingFactor + smoothingAdjustment + for index, candidate := range candidates { + if candidate.Weight <= 0 { + continue + } + randomWeight -= candidate.Weight if randomWeight < 0 { - return channel, nil + return candidate, index } } - // return null if no channel is not found - return nil, errors.New("channel not found") + return candidates[len(candidates)-1], len(candidates) - 1 +} + +func GetRandomSatisfiedChannel(group string, modelName string, retry int, requestPath string) (*Channel, error) { + tiers, err := GetSatisfiedChannelTiers(group, modelName, requestPath) + if err != nil || len(tiers) == 0 { + return nil, err + } + if retry < 0 { + retry = 0 + } + if retry >= len(tiers) { + retry = len(tiers) - 1 + } + candidate, _ := PickWeightedChannelCandidate(tiers[retry].Candidates) + return candidate.Channel, nil } // filterChannelsByRequestPathAndModel restricts candidates by request path and diff --git a/model/channel_cache_test.go b/model/channel_cache_test.go new file mode 100644 index 000000000000..13ed7d53eb8e --- /dev/null +++ b/model/channel_cache_test.go @@ -0,0 +1,17 @@ +package model + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestPickWeightedChannelCandidatePreservesZeroWeightSemantics(t *testing.T) { + zeroWeight := ChannelCandidate{Channel: &Channel{Id: 207}, Weight: 0} + positiveWeight := ChannelCandidate{Channel: &Channel{Id: 208}, Weight: 1} + + candidate, index := PickWeightedChannelCandidate([]ChannelCandidate{zeroWeight, positiveWeight}) + + assert.Equal(t, 1, index) + assert.Equal(t, positiveWeight.Channel.Id, candidate.Channel.Id) +} diff --git a/model/channel_settings_test.go b/model/channel_settings_test.go index 7612e697080f..d1ec44accf85 100644 --- a/model/channel_settings_test.go +++ b/model/channel_settings_test.go @@ -24,6 +24,15 @@ func TestChannelValidateSettingsRejectsInvalidHTTPTransport(t *testing.T) { setting: dto.ChannelSettings{HTTPProtocol: "http1", HTTP2ConnectionShards: 2}, wantErr: "http2_connection_shards", }, + { + name: "channel admission limits are valid", + setting: dto.ChannelSettings{MaxConcurrency: 20, RPMLimit: 120}, + }, + { + name: "negative channel admission limit rejected", + setting: dto.ChannelSettings{RPMLimit: -1}, + wantErr: "rpm_limit", + }, } for _, tt := range tests { diff --git a/relaykit/dto/channel_settings.go b/relaykit/dto/channel_settings.go index d3ede20d69c5..92c9c085e5ae 100644 --- a/relaykit/dto/channel_settings.go +++ b/relaykit/dto/channel_settings.go @@ -23,12 +23,20 @@ type ChannelSettings struct { // HTTP2ConnectionShards spreads HTTP/2 traffic across N independent transports // (1-8). Zero/unset means 1. Ignored when HTTPProtocol is "http1". HTTP2ConnectionShards int `json:"http2_connection_shards,omitempty"` + // MaxConcurrency limits simultaneous in-flight relay attempts for the channel. + // Zero means unlimited. + MaxConcurrency int `json:"max_concurrency,omitempty"` + // RPMLimit limits relay attempts admitted in a rolling 60-second window. + // Zero means unlimited. + RPMLimit int `json:"rpm_limit,omitempty"` } const ( - HTTPProtocolAuto = "auto" - HTTPProtocolHTTP1 = "http1" - MaxHTTP2ConnectionShards = 8 + HTTPProtocolAuto = "auto" + HTTPProtocolHTTP1 = "http1" + MaxHTTP2ConnectionShards = 8 + MaxChannelConcurrencyLimit = 1_000_000 + MaxChannelRequestsPerMinute = 1_000_000 ) // ValidateHTTPTransport validates save-time HTTP transport channel settings. @@ -51,6 +59,19 @@ func (s *ChannelSettings) ValidateHTTPTransport() error { return nil } +func (s *ChannelSettings) ValidateAdmissionLimits() error { + if s == nil { + return nil + } + if s.MaxConcurrency < 0 || s.MaxConcurrency > MaxChannelConcurrencyLimit { + return fmt.Errorf("invalid max_concurrency: %d", s.MaxConcurrency) + } + if s.RPMLimit < 0 || s.RPMLimit > MaxChannelRequestsPerMinute { + return fmt.Errorf("invalid rpm_limit: %d", s.RPMLimit) + } + return nil +} + type VertexKeyType string const ( diff --git a/relaykit/dto/channel_settings_test.go b/relaykit/dto/channel_settings_test.go index d482679a1a4c..9c0eef2cfa3c 100644 --- a/relaykit/dto/channel_settings_test.go +++ b/relaykit/dto/channel_settings_test.go @@ -10,6 +10,33 @@ import ( "github.com/stretchr/testify/require" ) +func TestChannelSettingsValidateAdmissionLimits(t *testing.T) { + tests := []struct { + name string + setting ChannelSettings + wantErr string + }{ + {name: "unset is unlimited"}, + {name: "positive limits", setting: ChannelSettings{MaxConcurrency: 20, RPMLimit: 120}}, + {name: "negative concurrency", setting: ChannelSettings{MaxConcurrency: -1}, wantErr: "max_concurrency"}, + {name: "excessive concurrency", setting: ChannelSettings{MaxConcurrency: MaxChannelConcurrencyLimit + 1}, wantErr: "max_concurrency"}, + {name: "negative rpm", setting: ChannelSettings{RPMLimit: -1}, wantErr: "rpm_limit"}, + {name: "excessive rpm", setting: ChannelSettings{RPMLimit: MaxChannelRequestsPerMinute + 1}, wantErr: "rpm_limit"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.setting.ValidateAdmissionLimits() + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + func TestAdvancedCustomValidateResponsesToChatConverterPath(t *testing.T) { valid := &AdvancedCustomConfig{ Routes: []AdvancedCustomRoute{ diff --git a/relaykit/types/error.go b/relaykit/types/error.go index 387fdad76948..ef5d1dbc6042 100644 --- a/relaykit/types/error.go +++ b/relaykit/types/error.go @@ -43,13 +43,14 @@ const ( ErrorCodeViolationFeeGrokCSAM ErrorCode = "violation_fee.grok.csam" // new api error - ErrorCodeCountTokenFailed ErrorCode = "count_token_failed" - ErrorCodeModelPriceError ErrorCode = "model_price_error" - ErrorCodeInvalidApiType ErrorCode = "invalid_api_type" - ErrorCodeJsonMarshalFailed ErrorCode = "json_marshal_failed" - ErrorCodeDoRequestFailed ErrorCode = "do_request_failed" - ErrorCodeGetChannelFailed ErrorCode = "get_channel_failed" - ErrorCodeGenRelayInfoFailed ErrorCode = "gen_relay_info_failed" + ErrorCodeCountTokenFailed ErrorCode = "count_token_failed" + ErrorCodeModelPriceError ErrorCode = "model_price_error" + ErrorCodeInvalidApiType ErrorCode = "invalid_api_type" + ErrorCodeJsonMarshalFailed ErrorCode = "json_marshal_failed" + ErrorCodeDoRequestFailed ErrorCode = "do_request_failed" + ErrorCodeGetChannelFailed ErrorCode = "get_channel_failed" + ErrorCodeGenRelayInfoFailed ErrorCode = "gen_relay_info_failed" + ErrorCodeChannelCapacityExhausted ErrorCode = "channel_capacity_exhausted" // channel error ErrorCodeChannelNoAvailableKey ErrorCode = "channel:no_available_key" diff --git a/service/channel_admission.go b/service/channel_admission.go new file mode 100644 index 000000000000..22111571c8c3 --- /dev/null +++ b/service/channel_admission.go @@ -0,0 +1,549 @@ +package service + +import ( + "context" + _ "embed" + "errors" + "fmt" + "strconv" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + + "github.com/gin-gonic/gin" + "github.com/go-redis/redis/v8" +) + +const ( + channelAdmissionNamespace = "new-api:channel_admission:v1" + channelAdmissionRPMWindow = time.Minute + channelAdmissionLeaseTTL = 2 * time.Minute + channelAdmissionBackendTimeout = 3 * time.Second + channelAdmissionFallbackLogRate = time.Minute + channelAdmissionLeaseContextKey = "channel_admission_lease" +) + +//go:embed lua/channel_admission_acquire.lua +var channelAdmissionAcquireLua string + +//go:embed lua/channel_admission_renew.lua +var channelAdmissionRenewLua string + +//go:embed lua/channel_admission_release.lua +var channelAdmissionReleaseLua string + +//go:embed lua/channel_admission_snapshot.lua +var channelAdmissionSnapshotLua string + +var ( + channelAdmissionAcquireScript = redis.NewScript(channelAdmissionAcquireLua) + channelAdmissionRenewScript = redis.NewScript(channelAdmissionRenewLua) + channelAdmissionReleaseScript = redis.NewScript(channelAdmissionReleaseLua) + channelAdmissionSnapshotScript = redis.NewScript(channelAdmissionSnapshotLua) +) + +type ChannelAdmissionMode string + +const ( + ChannelAdmissionModeDisabled ChannelAdmissionMode = "disabled" + ChannelAdmissionModeRedis ChannelAdmissionMode = "redis" + ChannelAdmissionModeMemory ChannelAdmissionMode = "memory" + ChannelAdmissionModeMemoryFallback ChannelAdmissionMode = "memory_fallback" +) + +type ChannelAdmissionReason string + +const ( + ChannelAdmissionReasonConcurrency ChannelAdmissionReason = "concurrency" + ChannelAdmissionReasonRPM ChannelAdmissionReason = "rpm" +) + +type ChannelAdmissionDecision struct { + Allowed bool + Reason ChannelAdmissionReason + RetryAfter time.Duration + Mode ChannelAdmissionMode + CurrentConcurrency int + CurrentRPM int +} + +type ChannelAdmissionSnapshot struct { + Mode ChannelAdmissionMode `json:"mode"` + CurrentConcurrency int `json:"current_concurrency"` + MaxConcurrency int `json:"max_concurrency"` + CurrentRPM int `json:"current_rpm"` + RPMLimit int `json:"rpm_limit"` +} + +type memoryRPMAdmission struct { + leaseID string + startedAt time.Time +} + +type memoryChannelAdmissionState struct { + mu sync.Mutex + concurrencyIDs map[string]struct{} + rpmAdmissions []memoryRPMAdmission +} + +// pruneRPM removes expired rolling-window admissions. The caller must hold s.mu. +func (s *memoryChannelAdmissionState) pruneRPM(cutoff time.Time) { + firstActive := 0 + for firstActive < len(s.rpmAdmissions) && !s.rpmAdmissions[firstActive].startedAt.After(cutoff) { + firstActive++ + } + if firstActive > 0 { + s.rpmAdmissions = append([]memoryRPMAdmission(nil), s.rpmAdmissions[firstActive:]...) + } +} + +type channelAdmissionManager struct { + redisClient func() *redis.Client + redisEnabled func() bool + now func() time.Time + leaseTTL time.Duration + rpmWindow time.Duration + renewLeases bool + memoryStates sync.Map + lastFallbackAt atomic.Int64 +} + +var defaultChannelAdmissionManager = &channelAdmissionManager{ + redisClient: func() *redis.Client { return common.RDB }, + redisEnabled: func() bool { return common.RedisEnabled && common.RDB != nil }, + now: time.Now, + leaseTTL: channelAdmissionLeaseTTL, + rpmWindow: channelAdmissionRPMWindow, + renewLeases: true, +} + +type ChannelAdmissionLease struct { + manager *channelAdmissionManager + channelID int + leaseID string + mode ChannelAdmissionMode + redisClient *redis.Client + tracksConcurrency bool + tracksRPM bool + renewalDone chan struct{} + mu sync.Mutex + renewalStopped bool + committed bool + released bool +} + +func (l *ChannelAdmissionLease) ChannelID() int { + if l == nil { + return 0 + } + return l.channelID +} + +func (l *ChannelAdmissionLease) Mode() ChannelAdmissionMode { + if l == nil { + return ChannelAdmissionModeDisabled + } + return l.mode +} + +// Commit marks the admission as an upstream attempt. Once committed, releasing +// the lease never refunds its rolling-window RPM entry. +func (l *ChannelAdmissionLease) Commit() { + if l == nil { + return + } + l.mu.Lock() + l.committed = true + l.mu.Unlock() +} + +func (l *ChannelAdmissionLease) Release() error { + if l == nil { + return nil + } + + l.mu.Lock() + defer l.mu.Unlock() + if l.released { + return nil + } + if !l.renewalStopped { + close(l.renewalDone) + l.renewalStopped = true + } + rollbackRPM := l.tracksRPM && !l.committed + if !l.tracksConcurrency && !rollbackRPM { + l.released = true + return nil + } + + switch l.mode { + case ChannelAdmissionModeRedis: + ctx, cancel := context.WithTimeout(context.Background(), channelAdmissionBackendTimeout) + defer cancel() + rollbackValue := 0 + if rollbackRPM { + rollbackValue = 1 + } + if err := channelAdmissionReleaseScript.Run( + ctx, + l.redisClient, + []string{channelAdmissionConcurrencyKey(l.channelID), channelAdmissionRPMKey(l.channelID)}, + l.leaseID, + rollbackValue, + ).Err(); err != nil { + return fmt.Errorf("release channel admission lease: %w", err) + } + case ChannelAdmissionModeMemory, ChannelAdmissionModeMemoryFallback: + state := l.manager.memoryState(l.channelID) + state.mu.Lock() + if l.tracksConcurrency { + delete(state.concurrencyIDs, l.leaseID) + } + if rollbackRPM { + for index, admission := range state.rpmAdmissions { + if admission.leaseID == l.leaseID { + state.rpmAdmissions = append(state.rpmAdmissions[:index], state.rpmAdmissions[index+1:]...) + break + } + } + } + state.mu.Unlock() + } + l.released = true + return nil +} + +func AcquireChannelAdmission(ctx context.Context, channel *model.Channel) (*ChannelAdmissionLease, ChannelAdmissionDecision, error) { + if channel == nil || channel.Id <= 0 { + return nil, ChannelAdmissionDecision{}, errors.New("channel admission requires a persisted channel") + } + settings := channel.GetSetting() + if err := settings.ValidateAdmissionLimits(); err != nil { + return nil, ChannelAdmissionDecision{}, err + } + return defaultChannelAdmissionManager.acquire(ctx, channel.Id, settings.MaxConcurrency, settings.RPMLimit) +} + +func GetChannelAdmissionSnapshot(ctx context.Context, channel *model.Channel) (ChannelAdmissionSnapshot, error) { + if channel == nil || channel.Id <= 0 { + return ChannelAdmissionSnapshot{}, errors.New("channel admission snapshot requires a persisted channel") + } + settings := channel.GetSetting() + if err := settings.ValidateAdmissionLimits(); err != nil { + return ChannelAdmissionSnapshot{}, err + } + return defaultChannelAdmissionManager.snapshot(ctx, channel.Id, settings.MaxConcurrency, settings.RPMLimit) +} + +func SetChannelAdmissionLease(c *gin.Context, lease *ChannelAdmissionLease) { + if c == nil || lease == nil { + return + } + c.Set(channelAdmissionLeaseContextKey, lease) +} + +func GetChannelAdmissionLease(c *gin.Context) *ChannelAdmissionLease { + if c == nil { + return nil + } + value, exists := c.Get(channelAdmissionLeaseContextKey) + if !exists { + return nil + } + lease, _ := value.(*ChannelAdmissionLease) + return lease +} + +func (m *channelAdmissionManager) acquire(ctx context.Context, channelID int, maxConcurrency int, rpmLimit int) (*ChannelAdmissionLease, ChannelAdmissionDecision, error) { + if maxConcurrency <= 0 && rpmLimit <= 0 { + return nil, ChannelAdmissionDecision{Allowed: true, Mode: ChannelAdmissionModeDisabled}, nil + } + if ctx == nil { + ctx = context.Background() + } + + if m.redisEnabled != nil && m.redisEnabled() { + client := m.redisClient() + lease, decision, err := m.acquireRedis(ctx, client, channelID, maxConcurrency, rpmLimit) + if err == nil { + return lease, decision, nil + } + if ctx.Err() != nil { + return nil, ChannelAdmissionDecision{}, ctx.Err() + } + m.logMemoryFallback(ctx, err) + return m.acquireMemory(channelID, maxConcurrency, rpmLimit, ChannelAdmissionModeMemoryFallback) + } + return m.acquireMemory(channelID, maxConcurrency, rpmLimit, ChannelAdmissionModeMemory) +} + +func (m *channelAdmissionManager) acquireRedis(ctx context.Context, client *redis.Client, channelID int, maxConcurrency int, rpmLimit int) (*ChannelAdmissionLease, ChannelAdmissionDecision, error) { + if client == nil { + return nil, ChannelAdmissionDecision{}, errors.New("Redis client is not initialized") + } + leaseID := common.GetUUID() + values, err := channelAdmissionAcquireScript.Run( + ctx, + client, + []string{channelAdmissionConcurrencyKey(channelID), channelAdmissionRPMKey(channelID)}, + maxConcurrency, + rpmLimit, + leaseID, + m.leaseTTL.Milliseconds(), + m.rpmWindow.Milliseconds(), + ).Slice() + if err != nil { + return nil, ChannelAdmissionDecision{}, err + } + if len(values) != 5 { + return nil, ChannelAdmissionDecision{}, fmt.Errorf("unexpected channel admission reply length %d", len(values)) + } + + allowed, err := redisAdmissionInteger(values[0]) + if err != nil { + return nil, ChannelAdmissionDecision{}, err + } + reasonValue, err := redisAdmissionInteger(values[1]) + if err != nil { + return nil, ChannelAdmissionDecision{}, err + } + concurrencyUsed, err := redisAdmissionInteger(values[2]) + if err != nil { + return nil, ChannelAdmissionDecision{}, err + } + rpmUsed, err := redisAdmissionInteger(values[3]) + if err != nil { + return nil, ChannelAdmissionDecision{}, err + } + retryAfterSeconds, err := redisAdmissionInteger(values[4]) + if err != nil { + return nil, ChannelAdmissionDecision{}, err + } + + decision := ChannelAdmissionDecision{ + Allowed: allowed == 1, + Mode: ChannelAdmissionModeRedis, + CurrentConcurrency: int(concurrencyUsed), + CurrentRPM: int(rpmUsed), + RetryAfter: time.Duration(retryAfterSeconds) * time.Second, + } + if !decision.Allowed { + if reasonValue == 1 { + decision.Reason = ChannelAdmissionReasonConcurrency + } else { + decision.Reason = ChannelAdmissionReasonRPM + } + return nil, decision, nil + } + + lease := &ChannelAdmissionLease{ + manager: m, + channelID: channelID, + leaseID: leaseID, + mode: ChannelAdmissionModeRedis, + redisClient: client, + tracksConcurrency: maxConcurrency > 0, + tracksRPM: rpmLimit > 0, + renewalDone: make(chan struct{}), + } + if lease.tracksConcurrency && m.renewLeases { + go lease.renewRedisLoop() + } + return lease, decision, nil +} + +func (m *channelAdmissionManager) acquireMemory(channelID int, maxConcurrency int, rpmLimit int, mode ChannelAdmissionMode) (*ChannelAdmissionLease, ChannelAdmissionDecision, error) { + state := m.memoryState(channelID) + now := m.now() + state.mu.Lock() + defer state.mu.Unlock() + + state.pruneRPM(now.Add(-m.rpmWindow)) + + decision := ChannelAdmissionDecision{ + Allowed: false, + Mode: mode, + CurrentConcurrency: len(state.concurrencyIDs), + CurrentRPM: len(state.rpmAdmissions), + } + if maxConcurrency > 0 && len(state.concurrencyIDs) >= maxConcurrency { + decision.Reason = ChannelAdmissionReasonConcurrency + decision.RetryAfter = time.Second + return nil, decision, nil + } + if rpmLimit > 0 && len(state.rpmAdmissions) >= rpmLimit { + decision.Reason = ChannelAdmissionReasonRPM + decision.RetryAfter = state.rpmAdmissions[0].startedAt.Add(m.rpmWindow).Sub(now) + if decision.RetryAfter < time.Second { + decision.RetryAfter = time.Second + } + return nil, decision, nil + } + + leaseID := common.GetUUID() + if maxConcurrency > 0 { + state.concurrencyIDs[leaseID] = struct{}{} + decision.CurrentConcurrency++ + } + if rpmLimit > 0 { + state.rpmAdmissions = append(state.rpmAdmissions, memoryRPMAdmission{leaseID: leaseID, startedAt: now}) + decision.CurrentRPM++ + } + decision.Allowed = true + lease := &ChannelAdmissionLease{ + manager: m, + channelID: channelID, + leaseID: leaseID, + mode: mode, + tracksConcurrency: maxConcurrency > 0, + tracksRPM: rpmLimit > 0, + renewalDone: make(chan struct{}), + } + return lease, decision, nil +} + +func (m *channelAdmissionManager) snapshot(ctx context.Context, channelID int, maxConcurrency int, rpmLimit int) (ChannelAdmissionSnapshot, error) { + snapshot := ChannelAdmissionSnapshot{MaxConcurrency: maxConcurrency, RPMLimit: rpmLimit} + if maxConcurrency <= 0 && rpmLimit <= 0 { + snapshot.Mode = ChannelAdmissionModeDisabled + return snapshot, nil + } + if ctx == nil { + ctx = context.Background() + } + if m.redisEnabled != nil && m.redisEnabled() { + client := m.redisClient() + var values []interface{} + var err error + if client == nil { + err = errors.New("Redis client is not initialized") + } else { + values, err = channelAdmissionSnapshotScript.Run( + ctx, + client, + []string{channelAdmissionConcurrencyKey(channelID), channelAdmissionRPMKey(channelID)}, + maxConcurrency, + rpmLimit, + m.rpmWindow.Milliseconds(), + ).Slice() + } + if err == nil && len(values) != 2 { + err = fmt.Errorf("unexpected channel admission snapshot reply length %d", len(values)) + } + if err == nil { + concurrencyUsed, concurrencyErr := redisAdmissionInteger(values[0]) + rpmUsed, rpmErr := redisAdmissionInteger(values[1]) + if concurrencyErr == nil && rpmErr == nil { + snapshot.Mode = ChannelAdmissionModeRedis + snapshot.CurrentConcurrency = int(concurrencyUsed) + snapshot.CurrentRPM = int(rpmUsed) + return snapshot, nil + } + err = errors.Join(concurrencyErr, rpmErr) + } + if ctx.Err() != nil { + return ChannelAdmissionSnapshot{}, ctx.Err() + } + m.logMemoryFallback(ctx, err) + snapshot.Mode = ChannelAdmissionModeMemoryFallback + } else { + snapshot.Mode = ChannelAdmissionModeMemory + } + + state := m.memoryState(channelID) + now := m.now() + state.mu.Lock() + defer state.mu.Unlock() + state.pruneRPM(now.Add(-m.rpmWindow)) + snapshot.CurrentConcurrency = len(state.concurrencyIDs) + snapshot.CurrentRPM = len(state.rpmAdmissions) + return snapshot, nil +} + +func (m *channelAdmissionManager) memoryState(channelID int) *memoryChannelAdmissionState { + state := &memoryChannelAdmissionState{concurrencyIDs: make(map[string]struct{})} + actual, _ := m.memoryStates.LoadOrStore(channelID, state) + return actual.(*memoryChannelAdmissionState) +} + +func (m *channelAdmissionManager) logMemoryFallback(ctx context.Context, err error) { + now := m.now().Unix() + last := m.lastFallbackAt.Load() + if last != 0 && now-last < int64(channelAdmissionFallbackLogRate/time.Second) { + return + } + if m.lastFallbackAt.CompareAndSwap(last, now) { + logger.LogWarn(ctx, fmt.Sprintf("channel admission Redis unavailable; using per-process fallback: %v", err)) + } +} + +func (l *ChannelAdmissionLease) renewRedisLoop() { + interval := l.manager.leaseTTL / 3 + if interval <= 0 { + interval = time.Second + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-l.renewalDone: + return + case <-ticker.C: + renewed, err := l.renewRedis() + if err != nil { + logger.LogWarn(context.Background(), fmt.Sprintf("renew channel admission lease failed: channel_id=%d error=%v", l.channelID, err)) + continue + } + if !renewed { + logger.LogWarn(context.Background(), fmt.Sprintf("channel admission lease expired before renewal: channel_id=%d", l.channelID)) + return + } + } + } +} + +func (l *ChannelAdmissionLease) renewRedis() (bool, error) { + if l == nil || !l.tracksConcurrency || l.mode != ChannelAdmissionModeRedis { + return false, nil + } + ctx, cancel := context.WithTimeout(context.Background(), channelAdmissionBackendTimeout) + defer cancel() + result, err := channelAdmissionRenewScript.Run( + ctx, + l.redisClient, + []string{channelAdmissionConcurrencyKey(l.channelID)}, + l.leaseID, + l.manager.leaseTTL.Milliseconds(), + ).Int64() + if err != nil { + return false, err + } + return result == 1, nil +} + +func channelAdmissionConcurrencyKey(channelID int) string { + return fmt.Sprintf("%s:{channel:%d}:concurrency", channelAdmissionNamespace, channelID) +} + +func channelAdmissionRPMKey(channelID int) string { + return fmt.Sprintf("%s:{channel:%d}:rpm", channelAdmissionNamespace, channelID) +} + +func redisAdmissionInteger(value interface{}) (int64, error) { + switch typed := value.(type) { + case int64: + return typed, nil + case string: + return strconv.ParseInt(typed, 10, 64) + case []byte: + return strconv.ParseInt(string(typed), 10, 64) + default: + return 0, fmt.Errorf("unexpected Redis integer reply type %T", value) + } +} diff --git a/service/channel_admission_test.go b/service/channel_admission_test.go new file mode 100644 index 000000000000..369f72c9bca7 --- /dev/null +++ b/service/channel_admission_test.go @@ -0,0 +1,656 @@ +package service + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/setting" + + "github.com/alicebob/miniredis/v2" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/go-redis/redis/v8" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func newMemoryChannelAdmissionManager(now func() time.Time) *channelAdmissionManager { + return &channelAdmissionManager{ + redisClient: func() *redis.Client { return nil }, + redisEnabled: func() bool { return false }, + now: now, + leaseTTL: channelAdmissionLeaseTTL, + rpmWindow: channelAdmissionRPMWindow, + renewLeases: false, + } +} + +func newRedisChannelAdmissionManager(client *redis.Client, leaseTTL time.Duration) *channelAdmissionManager { + return &channelAdmissionManager{ + redisClient: func() *redis.Client { return client }, + redisEnabled: func() bool { return true }, + now: time.Now, + leaseTTL: leaseTTL, + rpmWindow: channelAdmissionRPMWindow, + renewLeases: false, + } +} + +func TestChannelAdmissionMemoryConcurrencyUsesCurrentLimit(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + + unlimited, decision, err := manager.acquire(context.Background(), 101, 0, 0) + require.NoError(t, err) + assert.True(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionModeDisabled, decision.Mode) + assert.Nil(t, unlimited) + + first, decision, err := manager.acquire(context.Background(), 101, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionModeMemory, decision.Mode) + + _, decision, err = manager.acquire(context.Background(), 101, 1, 0) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonConcurrency, decision.Reason) + + second, decision, err := manager.acquire(context.Background(), 101, 2, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + assert.Equal(t, 2, decision.CurrentConcurrency) + + _, decision, err = manager.acquire(context.Background(), 101, 1, 0) + require.NoError(t, err) + assert.False(t, decision.Allowed) + + require.NoError(t, first.Release()) + require.NoError(t, first.Release()) + _, decision, err = manager.acquire(context.Background(), 101, 1, 0) + require.NoError(t, err) + assert.False(t, decision.Allowed) + + require.NoError(t, second.Release()) + third, decision, err := manager.acquire(context.Background(), 101, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + require.NoError(t, third.Release()) +} + +func TestChannelAdmissionMemoryRPMUsesCurrentLimit(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + + for range 2 { + lease, decision, err := manager.acquire(context.Background(), 106, 0, 2) + require.NoError(t, err) + require.True(t, decision.Allowed) + lease.Commit() + require.NoError(t, lease.Release()) + } + + _, decision, err := manager.acquire(context.Background(), 106, 0, 1) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonRPM, decision.Reason) + + third, decision, err := manager.acquire(context.Background(), 106, 0, 3) + require.NoError(t, err) + require.True(t, decision.Allowed) + third.Commit() + require.NoError(t, third.Release()) + + unlimited, decision, err := manager.acquire(context.Background(), 106, 0, 0) + require.NoError(t, err) + assert.True(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionModeDisabled, decision.Mode) + assert.Nil(t, unlimited) +} + +func TestChannelAdmissionMemoryConcurrentAcquireIsAtomic(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + start := make(chan struct{}) + leases := make(chan *ChannelAdmissionLease, 2) + decisions := make(chan ChannelAdmissionDecision, 2) + errorsCh := make(chan error, 2) + var waitGroup sync.WaitGroup + + for range 2 { + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + <-start + lease, decision, err := manager.acquire(context.Background(), 105, 1, 0) + errorsCh <- err + leases <- lease + decisions <- decision + }() + } + close(start) + waitGroup.Wait() + close(leases) + close(decisions) + close(errorsCh) + + for err := range errorsCh { + require.NoError(t, err) + } + allowed := 0 + for decision := range decisions { + if decision.Allowed { + allowed++ + } + } + assert.Equal(t, 1, allowed) + for lease := range leases { + if lease != nil { + require.NoError(t, lease.Release()) + } + } +} + +func TestChannelAdmissionMemoryRPMRollsBackOnlyUnstartedAttempts(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + + unstarted, decision, err := manager.acquire(context.Background(), 102, 0, 1) + require.NoError(t, err) + require.True(t, decision.Allowed) + require.NoError(t, unstarted.Release()) + + started, decision, err := manager.acquire(context.Background(), 102, 0, 1) + require.NoError(t, err) + require.True(t, decision.Allowed) + started.Commit() + require.NoError(t, started.Release()) + + _, decision, err = manager.acquire(context.Background(), 102, 0, 1) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonRPM, decision.Reason) + assert.Equal(t, time.Minute, decision.RetryAfter) + + now = now.Add(time.Minute) + afterWindow, decision, err := manager.acquire(context.Background(), 102, 0, 1) + require.NoError(t, err) + require.True(t, decision.Allowed) + require.NoError(t, afterWindow.Release()) +} + +func TestChannelAdmissionCombinedRejectionDoesNotSpendRPM(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + + first, decision, err := manager.acquire(context.Background(), 103, 1, 2) + require.NoError(t, err) + require.True(t, decision.Allowed) + first.Commit() + + _, decision, err = manager.acquire(context.Background(), 103, 1, 2) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonConcurrency, decision.Reason) + assert.Equal(t, 1, decision.CurrentRPM) + + require.NoError(t, first.Release()) + second, decision, err := manager.acquire(context.Background(), 103, 1, 2) + require.NoError(t, err) + require.True(t, decision.Allowed) + assert.Equal(t, 2, decision.CurrentRPM) + second.Commit() + require.NoError(t, second.Release()) + + _, decision, err = manager.acquire(context.Background(), 103, 1, 2) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonRPM, decision.Reason) +} + +func TestChannelAdmissionRedisFailureUsesMemoryFallback(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := &channelAdmissionManager{ + redisClient: func() *redis.Client { return nil }, + redisEnabled: func() bool { return true }, + now: func() time.Time { return now }, + leaseTTL: channelAdmissionLeaseTTL, + rpmWindow: channelAdmissionRPMWindow, + renewLeases: false, + } + + lease, decision, err := manager.acquire(context.Background(), 104, 1, 1) + require.NoError(t, err) + require.True(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionModeMemoryFallback, decision.Mode) + require.NoError(t, lease.Release()) +} + +func TestChannelAdmissionRedisIsGlobalAcrossClients(t *testing.T) { + server := miniredis.RunT(t) + clientA := redis.NewClient(&redis.Options{Addr: server.Addr()}) + clientB := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { + _ = clientA.Close() + _ = clientB.Close() + }) + managerA := newRedisChannelAdmissionManager(clientA, 10*time.Second) + managerB := newRedisChannelAdmissionManager(clientB, 10*time.Second) + + first, decision, err := managerA.acquire(context.Background(), 201, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionModeRedis, decision.Mode) + + _, decision, err = managerB.acquire(context.Background(), 201, 1, 0) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonConcurrency, decision.Reason) + + snapshot, err := managerB.snapshot(context.Background(), 201, 1, 0) + require.NoError(t, err) + assert.Equal(t, 1, snapshot.CurrentConcurrency) + + first.Commit() + require.NoError(t, first.Release()) + second, decision, err := managerB.acquire(context.Background(), 201, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + require.NoError(t, second.Release()) +} + +func TestChannelAdmissionRedisRPMIsGlobalAcrossClients(t *testing.T) { + server := miniredis.RunT(t) + baseTime := time.Unix(1_700_000_000, 0) + server.SetTime(baseTime) + clientA := redis.NewClient(&redis.Options{Addr: server.Addr()}) + clientB := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { + _ = clientA.Close() + _ = clientB.Close() + }) + managerA := newRedisChannelAdmissionManager(clientA, 10*time.Second) + managerB := newRedisChannelAdmissionManager(clientB, 10*time.Second) + + first, decision, err := managerA.acquire(context.Background(), 205, 0, 1) + require.NoError(t, err) + require.True(t, decision.Allowed) + first.Commit() + require.NoError(t, first.Release()) + + _, decision, err = managerB.acquire(context.Background(), 205, 0, 1) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonRPM, decision.Reason) + assert.Equal(t, time.Minute, decision.RetryAfter) + + server.SetTime(baseTime.Add(time.Minute)) + afterWindow, decision, err := managerB.acquire(context.Background(), 205, 0, 1) + require.NoError(t, err) + require.True(t, decision.Allowed) + require.NoError(t, afterWindow.Release()) +} + +func TestChannelAdmissionRedisRPMRollbackAndCommit(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + manager := newRedisChannelAdmissionManager(client, 10*time.Second) + + unstarted, decision, err := manager.acquire(context.Background(), 202, 0, 1) + require.NoError(t, err) + require.True(t, decision.Allowed) + require.NoError(t, unstarted.Release()) + + started, decision, err := manager.acquire(context.Background(), 202, 0, 1) + require.NoError(t, err) + require.True(t, decision.Allowed) + started.Commit() + require.NoError(t, started.Release()) + + _, decision, err = manager.acquire(context.Background(), 202, 0, 1) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonRPM, decision.Reason) +} + +func TestChannelAdmissionRedisLeaseExpiresAndRenews(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + manager := newRedisChannelAdmissionManager(client, 2*time.Second) + baseTime := time.Unix(1_700_000_000, 0) + server.SetTime(baseTime) + + expiring, decision, err := manager.acquire(context.Background(), 203, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + server.SetTime(baseTime.Add(3 * time.Second)) + replacement, decision, err := manager.acquire(context.Background(), 203, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + require.NoError(t, replacement.Release()) + + renewed, decision, err := manager.acquire(context.Background(), 204, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + server.SetTime(baseTime.Add(4500 * time.Millisecond)) + ok, err := renewed.renewRedis() + require.NoError(t, err) + require.True(t, ok) + server.SetTime(baseTime.Add(6 * time.Second)) + _, decision, err = manager.acquire(context.Background(), 204, 1, 0) + require.NoError(t, err) + assert.False(t, decision.Allowed) + require.NoError(t, renewed.Release()) + require.NoError(t, expiring.Release()) +} + +func TestAcquireChannelAdmissionReadsUpdatedChannelSettings(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + previousManager := defaultChannelAdmissionManager + defaultChannelAdmissionManager = manager + t.Cleanup(func() { defaultChannelAdmissionManager = previousManager }) + + channel := &model.Channel{Id: 209} + unlimited, decision, err := AcquireChannelAdmission(context.Background(), channel) + require.NoError(t, err) + assert.True(t, decision.Allowed) + assert.Nil(t, unlimited) + + channel.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + first, decision, err := AcquireChannelAdmission(context.Background(), channel) + require.NoError(t, err) + require.True(t, decision.Allowed) + + channel.SetSetting(dto.ChannelSettings{MaxConcurrency: 2}) + second, decision, err := AcquireChannelAdmission(context.Background(), channel) + require.NoError(t, err) + require.True(t, decision.Allowed) + + channel.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + _, decision, err = AcquireChannelAdmission(context.Background(), channel) + require.NoError(t, err) + assert.False(t, decision.Allowed) + + channel.SetSetting(dto.ChannelSettings{}) + unlimited, decision, err = AcquireChannelAdmission(context.Background(), channel) + require.NoError(t, err) + assert.True(t, decision.Allowed) + assert.Nil(t, unlimited) + + require.NoError(t, first.Release()) + require.NoError(t, second.Release()) +} + +func TestChannelAdmissionMultiKeyChannelSharesOneLimit(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + previousManager := defaultChannelAdmissionManager + defaultChannelAdmissionManager = manager + t.Cleanup(func() { defaultChannelAdmissionManager = previousManager }) + + channel := &model.Channel{ + Id: 206, + Keys: []string{"first-key", "second-key"}, + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + }, + } + channel.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + + first, decision, err := AcquireChannelAdmission(context.Background(), channel) + require.NoError(t, err) + require.True(t, decision.Allowed) + + _, decision, err = AcquireChannelAdmission(context.Background(), channel) + require.NoError(t, err) + assert.False(t, decision.Allowed) + assert.Equal(t, ChannelAdmissionReasonConcurrency, decision.Reason) + require.NoError(t, first.Release()) +} + +func TestSelectAdmittedChannelTriesSameTierThenLowerPriority(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + previousManager := defaultChannelAdmissionManager + defaultChannelAdmissionManager = manager + t.Cleanup(func() { defaultChannelAdmissionManager = previousManager }) + + highA := &model.Channel{Id: 301} + highA.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + highB := &model.Channel{Id: 302} + highB.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + low := &model.Channel{Id: 303} + low.SetSetting(dto.ChannelSettings{}) + + highALease, decision, err := manager.acquire(context.Background(), highA.Id, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + + tiers := []model.ChannelCandidateTier{ + {Priority: 10, Candidates: []model.ChannelCandidate{{Channel: highA, Weight: 1}, {Channel: highB, Weight: 1}}}, + {Priority: 0, Candidates: []model.ChannelCandidate{{Channel: low, Weight: 1}}}, + } + selection, err := selectAdmittedChannel(context.Background(), "default", tiers, 0) + require.NoError(t, err) + require.NotNil(t, selection) + assert.Equal(t, highB.Id, selection.Channel.Id) + require.NoError(t, selection.Lease.Release()) + + highBLease, decision, err := manager.acquire(context.Background(), highB.Id, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + selection, err = selectAdmittedChannel(context.Background(), "default", tiers, 0) + require.NoError(t, err) + require.NotNil(t, selection) + assert.Equal(t, low.Id, selection.Channel.Id) + require.NoError(t, selection.Lease.Release()) + + require.NoError(t, highALease.Release()) + require.NoError(t, highBLease.Release()) +} + +func TestSelectAdmittedChannelReturnsCapacityErrorWhenAllCandidatesAreFull(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + previousManager := defaultChannelAdmissionManager + defaultChannelAdmissionManager = manager + t.Cleanup(func() { defaultChannelAdmissionManager = previousManager }) + + channels := []*model.Channel{{Id: 304}, {Id: 305}} + leases := make([]*ChannelAdmissionLease, 0, len(channels)) + for _, channel := range channels { + channel.SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + lease, decision, err := manager.acquire(context.Background(), channel.Id, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + leases = append(leases, lease) + } + defer func() { + for _, lease := range leases { + require.NoError(t, lease.Release()) + } + }() + + selection, err := selectAdmittedChannel(context.Background(), "default", []model.ChannelCandidateTier{{ + Priority: 10, + Candidates: []model.ChannelCandidate{ + {Channel: channels[0], Weight: 1}, + {Channel: channels[1], Weight: 1}, + }, + }}, 0) + assert.Nil(t, selection) + var capacityErr *ChannelCapacityError + require.True(t, errors.As(err, &capacityErr)) + assert.Equal(t, 2, capacityErr.ConcurrencyRejects) + assert.Equal(t, int64(1), capacityErr.RetryAfterSeconds()) +} + +func TestSelectChannelWithAdmissionFallsThroughAutoGroupsOnCapacity(t *testing.T) { + dsn := fmt.Sprintf("file:channel-admission-auto-%d?mode=memory&cache=shared", time.Now().UnixNano()) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{})) + + previousDB := model.DB + previousMemoryCache := common.MemoryCacheEnabled + previousRedisEnabled := common.RedisEnabled + previousManager := defaultChannelAdmissionManager + previousAutoGroups := setting.AutoGroups2JsonString() + previousUsableGroups := setting.UserUsableGroups2JSONString() + model.DB = db + common.MemoryCacheEnabled = true + common.RedisEnabled = false + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + defaultChannelAdmissionManager = manager + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["default","vip"]`)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + t.Cleanup(func() { + model.DB = previousDB + common.MemoryCacheEnabled = previousMemoryCache + common.RedisEnabled = previousRedisEnabled + defaultChannelAdmissionManager = previousManager + require.NoError(t, setting.UpdateAutoGroupsByJsonString(previousAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(previousUsableGroups)) + if previousMemoryCache && previousDB != nil { + model.InitChannelCache() + } + }) + + priority := int64(10) + weight := uint(1) + channels := []model.Channel{ + {Id: 404, Name: "default-full", Key: "key-default", Status: common.ChannelStatusEnabled, Models: "gpt-test", Group: "default", Priority: &priority, Weight: &weight}, + {Id: 405, Name: "vip-available", Key: "key-vip", Status: common.ChannelStatusEnabled, Models: "gpt-test", Group: "vip", Priority: &priority, Weight: &weight}, + } + for index := range channels { + channels[index].SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + require.NoError(t, db.Create(&channels[index]).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: channels[index].Group, + Model: "gpt-test", + ChannelId: channels[index].Id, + Enabled: true, + Priority: &priority, + Weight: weight, + }).Error) + } + model.InitChannelCache() + + fullLease, decision, err := manager.acquire(context.Background(), channels[0].Id, 1, 0) + require.NoError(t, err) + require.True(t, decision.Allowed) + defer func() { require.NoError(t, fullLease.Release()) }() + + gin.SetMode(gin.TestMode) + ctx := &gin.Context{} + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + retry := 0 + selection, err := SelectChannelWithAdmission(&RetryParam{ + Ctx: ctx, + TokenGroup: "auto", + ModelName: "gpt-test", + RequestPath: "/v1/chat/completions", + Retry: &retry, + }) + require.NoError(t, err) + require.NotNil(t, selection) + assert.Equal(t, channels[1].Id, selection.Channel.Id) + assert.Equal(t, "vip", selection.Group) + assert.Equal(t, 0, retry) + require.NoError(t, selection.Lease.Release()) +} + +func TestSelectChannelWithAdmissionDoesNotAdvanceRetryOnCapacitySkip(t *testing.T) { + dsn := fmt.Sprintf("file:channel-admission-%d?mode=memory&cache=shared", time.Now().UnixNano()) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{})) + + previousDB := model.DB + previousMemoryCache := common.MemoryCacheEnabled + previousRedisEnabled := common.RedisEnabled + previousManager := defaultChannelAdmissionManager + model.DB = db + common.MemoryCacheEnabled = false + common.RedisEnabled = false + now := time.Unix(1_700_000_000, 0) + manager := newMemoryChannelAdmissionManager(func() time.Time { return now }) + defaultChannelAdmissionManager = manager + t.Cleanup(func() { + model.DB = previousDB + common.MemoryCacheEnabled = previousMemoryCache + common.RedisEnabled = previousRedisEnabled + defaultChannelAdmissionManager = previousManager + if previousMemoryCache && previousDB != nil { + model.InitChannelCache() + } + }) + + highPriority := int64(10) + lowPriority := int64(0) + weight := uint(1) + channels := []model.Channel{ + {Id: 401, Name: "high-a", Key: "key-a", Status: common.ChannelStatusEnabled, Models: "gpt-test", Group: "default", Priority: &highPriority, Weight: &weight}, + {Id: 402, Name: "high-b", Key: "key-b", Status: common.ChannelStatusEnabled, Models: "gpt-test", Group: "default", Priority: &highPriority, Weight: &weight}, + {Id: 403, Name: "low", Key: "key-c", Status: common.ChannelStatusEnabled, Models: "gpt-test", Group: "default", Priority: &lowPriority, Weight: &weight}, + } + for index := range channels { + if channels[index].Id != 403 { + channels[index].SetSetting(dto.ChannelSettings{MaxConcurrency: 1}) + } + require.NoError(t, db.Create(&channels[index]).Error) + priority := channels[index].GetPriority() + require.NoError(t, db.Create(&model.Ability{ + Group: "default", + Model: "gpt-test", + ChannelId: channels[index].Id, + Enabled: true, + Priority: &priority, + Weight: weight, + }).Error) + } + model.InitChannelCache() + + leases := make([]*ChannelAdmissionLease, 0, 2) + for _, channelID := range []int{401, 402} { + lease, decision, acquireErr := manager.acquire(context.Background(), channelID, 1, 0) + require.NoError(t, acquireErr) + require.True(t, decision.Allowed) + leases = append(leases, lease) + } + defer func() { + for _, lease := range leases { + require.NoError(t, lease.Release()) + } + }() + + gin.SetMode(gin.TestMode) + ctx := &gin.Context{} + retry := 0 + selection, err := SelectChannelWithAdmission(&RetryParam{ + Ctx: ctx, + TokenGroup: "default", + ModelName: "gpt-test", + RequestPath: "/v1/chat/completions", + Retry: &retry, + }) + require.NoError(t, err) + require.NotNil(t, selection) + assert.Equal(t, 403, selection.Channel.Id) + assert.Equal(t, 0, retry) + require.NoError(t, selection.Lease.Release()) +} diff --git a/service/channel_select.go b/service/channel_select.go index 0ab88dc84ff2..e949723f741d 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -1,7 +1,10 @@ package service import ( + "context" "errors" + "fmt" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -45,118 +48,194 @@ func (p *RetryParam) ResetRetryNextTry() { p.resetNextTry = true } -// CacheGetRandomSatisfiedChannel tries to get a random channel that satisfies the requirements. -// 尝试获取一个满足要求的随机渠道。 -// -// For "auto" tokenGroup with cross-group Retry enabled: -// 对于启用了跨分组重试的 "auto" tokenGroup： -// -// - 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) { - var channel *model.Channel - var err error - selectGroup := param.TokenGroup +type ChannelSelection struct { + Channel *model.Channel + Group string + Lease *ChannelAdmissionLease +} + +type ChannelCapacityError struct { + RetryAfter time.Duration + ConcurrencyRejects int + RPMRejects int +} + +func (e *ChannelCapacityError) Error() string { + return "all matching channels are at their configured capacity" +} + +func (e *ChannelCapacityError) RetryAfterSeconds() int64 { + if e == nil || e.RetryAfter <= 0 { + return 1 + } + seconds := int64((e.RetryAfter + time.Second - 1) / time.Second) + if seconds < 1 { + return 1 + } + return seconds +} + +func (e *ChannelCapacityError) addDecision(decision ChannelAdmissionDecision) { + if decision.Reason == ChannelAdmissionReasonConcurrency { + e.ConcurrencyRejects++ + } else if decision.Reason == ChannelAdmissionReasonRPM { + e.RPMRejects++ + } + if decision.RetryAfter > 0 && (e.RetryAfter <= 0 || decision.RetryAfter < e.RetryAfter) { + e.RetryAfter = decision.RetryAfter + } +} + +func (e *ChannelCapacityError) merge(other *ChannelCapacityError) { + if other == nil { + return + } + e.ConcurrencyRejects += other.ConcurrencyRejects + e.RPMRejects += other.RPMRejects + if other.RetryAfter > 0 && (e.RetryAfter <= 0 || other.RetryAfter < e.RetryAfter) { + e.RetryAfter = other.RetryAfter + } +} + +// SelectChannelWithAdmission selects and reserves a channel before any upstream +// request is sent. Capacity rejections are handled inside this call, so they do +// not advance RetryParam or consume an upstream retry. +func SelectChannelWithAdmission(param *RetryParam) (*ChannelSelection, error) { + if param == nil || param.Ctx == nil { + return nil, errors.New("channel selection requires a request context") + } + if param.TokenGroup != "auto" { + tiers, err := model.GetSatisfiedChannelTiers(param.TokenGroup, param.ModelName, param.RequestPath) + if err != nil { + return nil, err + } + selection, err := selectAdmittedChannel(param.Ctx, param.TokenGroup, tiers, param.GetRetry()) + return selection, err + } + userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup) + autoGroups := GetRequestAutoGroups(param.Ctx, userGroup) + if len(autoGroups) == 0 { + return nil, errors.New("auto groups is not enabled") + } + startGroupIndex := 0 + if lastGroupIndex, exists := common.GetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex); exists { + if index, ok := lastGroupIndex.(int); ok && index >= 0 { + startGroupIndex = index + } + } + if startGroupIndex >= len(autoGroups) { + return nil, nil + } - if param.TokenGroup == "auto" { - autoGroups := GetRequestAutoGroups(param.Ctx, userGroup) - if len(autoGroups) == 0 { - return nil, selectGroup, errors.New("auto groups is not enabled") + crossGroupRetry := common.GetContextKeyBool(param.Ctx, constant.ContextKeyTokenCrossGroupRetry) + capacityErr := &ChannelCapacityError{} + for groupIndex := startGroupIndex; groupIndex < len(autoGroups); groupIndex++ { + selectGroup := autoGroups[groupIndex] + priorityRetry := param.GetRetry() + if groupIndex > startGroupIndex { + priorityRetry = 0 } + logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", selectGroup, priorityRetry) - // startGroupIndex: the group index to start searching from - // startGroupIndex: 开始搜索的分组索引 - startGroupIndex := 0 - crossGroupRetry := common.GetContextKeyBool(param.Ctx, constant.ContextKeyTokenCrossGroupRetry) + tiers, err := model.GetSatisfiedChannelTiers(selectGroup, param.ModelName, param.RequestPath) + if err != nil { + return nil, err + } + if len(tiers) == 0 { + common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, groupIndex+1) + common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupRetryIndex, 0) + param.SetRetry(0) + continue + } - if lastGroupIndex, exists := common.GetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex); exists { - if idx, ok := lastGroupIndex.(int); ok { - startGroupIndex = idx + selection, err := selectAdmittedChannel(param.Ctx, selectGroup, tiers, priorityRetry) + if err != nil { + var groupCapacityErr *ChannelCapacityError + if !errors.As(err, &groupCapacityErr) { + return nil, err } + capacityErr.merge(groupCapacityErr) + common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, groupIndex+1) + common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupRetryIndex, 0) + param.SetRetry(0) + continue + } + if selection == nil { + continue + } + + common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, selectGroup) + logger.LogDebug(param.Ctx, "Auto selected group: %s", selectGroup) + if crossGroupRetry && priorityRetry >= common.RetryTimes { + common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, groupIndex+1) + param.SetRetry(0) + param.ResetRetryNextTry() + } else { + common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, groupIndex) } + return selection, nil + } - for i := startGroupIndex; i < len(autoGroups); i++ { - autoGroup := autoGroups[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 + if capacityErr.ConcurrencyRejects > 0 || capacityErr.RPMRejects > 0 { + return nil, capacityErr + } + return nil, nil +} + +// CacheGetRandomSatisfiedChannel preserves the legacy selection contract for +// callers that cannot take ownership of an admission lease. +func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, error) { + group := "" + if param != nil { + group = param.TokenGroup + } + selection, err := SelectChannelWithAdmission(param) + if err != nil { + return nil, group, err + } + if selection == nil || selection.Channel == nil { + return nil, group, nil + } + if err := selection.Lease.Release(); err != nil { + return nil, selection.Group, err + } + return selection.Channel, selection.Group, nil +} + +func selectAdmittedChannel(ctx context.Context, group string, tiers []model.ChannelCandidateTier, startTier int) (*ChannelSelection, error) { + if len(tiers) == 0 { + return nil, nil + } + if startTier < 0 { + startTier = 0 + } + if startTier >= len(tiers) { + startTier = len(tiers) - 1 + } + + capacityErr := &ChannelCapacityError{} + for tierIndex := startTier; tierIndex < len(tiers); tierIndex++ { + candidates := append([]model.ChannelCandidate(nil), tiers[tierIndex].Candidates...) + for len(candidates) > 0 { + candidate, candidateIndex := model.PickWeightedChannelCandidate(candidates) + if candidateIndex < 0 || candidate.Channel == nil { + break } - logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) - - channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath) - 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 + candidates = append(candidates[:candidateIndex], candidates[candidateIndex+1:]...) + + lease, decision, err := AcquireChannelAdmission(ctx, candidate.Channel) + if err != nil { + return nil, fmt.Errorf("acquire channel #%d admission: %w", candidate.Channel.Id, err) } - 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) + if decision.Allowed { + return &ChannelSelection{Channel: candidate.Channel, Group: group, Lease: lease}, nil } - break - } - } else { - channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath) - if err != nil { - return nil, param.TokenGroup, err + capacityErr.addDecision(decision) } } - return channel, selectGroup, nil + if capacityErr.ConcurrencyRejects > 0 || capacityErr.RPMRejects > 0 { + return nil, capacityErr + } + return nil, nil } diff --git a/service/lua/channel_admission_acquire.lua b/service/lua/channel_admission_acquire.lua new file mode 100644 index 000000000000..2e33681ba002 --- /dev/null +++ b/service/lua/channel_admission_acquire.lua @@ -0,0 +1,51 @@ +local max_concurrency = tonumber(ARGV[1]) or 0 +local rpm_limit = tonumber(ARGV[2]) or 0 +local lease_id = ARGV[3] +local lease_ttl_ms = tonumber(ARGV[4]) +local rpm_window_ms = tonumber(ARGV[5]) + +local redis_time = redis.call('TIME') +local now_ms = tonumber(redis_time[1]) * 1000 + math.floor(tonumber(redis_time[2]) / 1000) + +local concurrency_used = 0 +if max_concurrency > 0 then + redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now_ms) + concurrency_used = redis.call('ZCARD', KEYS[1]) +end + +local rpm_used = 0 +if rpm_limit > 0 then + redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', now_ms - rpm_window_ms) + rpm_used = redis.call('ZCARD', KEYS[2]) +end + +if max_concurrency > 0 and concurrency_used >= max_concurrency then + local oldest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES') + local retry_after = 1 + if #oldest >= 2 then + retry_after = math.max(1, math.floor((tonumber(oldest[2]) - now_ms + 999) / 1000)) + end + return {0, 1, concurrency_used, rpm_used, retry_after} +end + +if rpm_limit > 0 and rpm_used >= rpm_limit then + local oldest = redis.call('ZRANGE', KEYS[2], 0, 0, 'WITHSCORES') + local retry_after = 1 + if #oldest >= 2 then + retry_after = math.max(1, math.floor((tonumber(oldest[2]) + rpm_window_ms - now_ms + 999) / 1000)) + end + return {0, 2, concurrency_used, rpm_used, retry_after} +end + +if max_concurrency > 0 then + redis.call('ZADD', KEYS[1], now_ms + lease_ttl_ms, lease_id) + redis.call('PEXPIRE', KEYS[1], lease_ttl_ms * 2) + concurrency_used = concurrency_used + 1 +end +if rpm_limit > 0 then + redis.call('ZADD', KEYS[2], now_ms, lease_id) + redis.call('PEXPIRE', KEYS[2], rpm_window_ms + 5000) + rpm_used = rpm_used + 1 +end + +return {1, 0, concurrency_used, rpm_used, 0} diff --git a/service/lua/channel_admission_release.lua b/service/lua/channel_admission_release.lua new file mode 100644 index 000000000000..760e466a160f --- /dev/null +++ b/service/lua/channel_admission_release.lua @@ -0,0 +1,5 @@ +local removed = redis.call('ZREM', KEYS[1], ARGV[1]) +if tonumber(ARGV[2]) == 1 then + removed = removed + redis.call('ZREM', KEYS[2], ARGV[1]) +end +return removed diff --git a/service/lua/channel_admission_renew.lua b/service/lua/channel_admission_renew.lua new file mode 100644 index 000000000000..b6cc5c6dfda1 --- /dev/null +++ b/service/lua/channel_admission_renew.lua @@ -0,0 +1,12 @@ +local lease_id = ARGV[1] +local lease_ttl_ms = tonumber(ARGV[2]) + +if not redis.call('ZSCORE', KEYS[1], lease_id) then + return 0 +end + +local redis_time = redis.call('TIME') +local now_ms = tonumber(redis_time[1]) * 1000 + math.floor(tonumber(redis_time[2]) / 1000) +redis.call('ZADD', KEYS[1], 'XX', now_ms + lease_ttl_ms, lease_id) +redis.call('PEXPIRE', KEYS[1], lease_ttl_ms * 2) +return 1 diff --git a/service/lua/channel_admission_snapshot.lua b/service/lua/channel_admission_snapshot.lua new file mode 100644 index 000000000000..04257abdfb59 --- /dev/null +++ b/service/lua/channel_admission_snapshot.lua @@ -0,0 +1,20 @@ +local max_concurrency = tonumber(ARGV[1]) or 0 +local rpm_limit = tonumber(ARGV[2]) or 0 +local rpm_window_ms = tonumber(ARGV[3]) + +local redis_time = redis.call('TIME') +local now_ms = tonumber(redis_time[1]) * 1000 + math.floor(tonumber(redis_time[2]) / 1000) + +local concurrency_used = 0 +if max_concurrency > 0 then + redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now_ms) + concurrency_used = redis.call('ZCARD', KEYS[1]) +end + +local rpm_used = 0 +if rpm_limit > 0 then + redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', now_ms - rpm_window_ms) + rpm_used = redis.call('ZCARD', KEYS[2]) +end + +return {concurrency_used, rpm_used} diff --git a/web/src/features/channels/components/drawers/channel-mutate-drawer.tsx b/web/src/features/channels/components/drawers/channel-mutate-drawer.tsx index 3380d9e52c24..200a6e565632 100644 --- a/web/src/features/channels/components/drawers/channel-mutate-drawer.tsx +++ b/web/src/features/channels/components/drawers/channel-mutate-drawer.tsx @@ -150,6 +150,7 @@ import { useChannelMutateForm } from '../../hooks/use-channel-mutate-form' import { CHANNEL_FORM_DEFAULT_VALUES, CHANNEL_TYPE_ADVANCED_CUSTOM, + MAX_CHANNEL_ADMISSION_LIMIT, channelFormSchema, channelsQueryKeys, getAdvancedCustomStats, @@ -286,6 +287,8 @@ const SENSITIVE_FORM_FIELDS = [ 'proxy', 'http_protocol', 'http2_connection_shards', + 'max_concurrency', + 'rpm_limit', 'pass_through_body_enabled', 'system_prompt', 'system_prompt_override', @@ -344,6 +347,8 @@ function hasAdvancedSettingsValues(values: ChannelFormValues): boolean { (values.http_protocol && values.http_protocol !== 'auto') || (values.http2_connection_shards != null && values.http2_connection_shards > 1) || + (values.max_concurrency ?? 0) > 0 || + (values.rpm_limit ?? 0) > 0 || values.claude_beta_query || values.upstream_model_update_check_enabled || values.upstream_model_update_auto_sync_enabled || @@ -736,6 +741,8 @@ export function ChannelMutateDrawer({ const currentAdvancedCustom = form.watch('advanced_custom') const currentPriority = form.watch('priority') const currentWeight = form.watch('weight') + const currentMaxConcurrency = form.watch('max_concurrency') + const currentRPMLimit = form.watch('rpm_limit') const currentTestModel = form.watch('test_model') const currentAutoBan = form.watch('auto_ban') const currentTag = form.watch('tag') @@ -1003,6 +1010,8 @@ export function ChannelMutateDrawer({ const routingStrategyConfigured = Boolean( currentPriority || currentWeight || + currentMaxConcurrency || + currentRPMLimit || currentTestModel?.trim() || (currentAutoBan ?? 1) !== 1 ) @@ -3688,6 +3697,76 @@ export function ChannelMutateDrawer({ /> +
+ ( + + + {t('Maximum concurrency')} + + + + field.onChange( + event.target.value === '' + ? 0 + : Number(event.target.value) + ) + } + /> + + + {t( + 'Maximum in-flight requests for this channel. 0 means unlimited.' + )} + + + + )} + /> + + ( + + {t('RPM limit')} + + + field.onChange( + event.target.value === '' + ? 0 + : Number(event.target.value) + ) + } + /> + + + {t( + 'Maximum request starts in a rolling 60-second window. 0 means unlimited.' + )} + + + + )} + /> +
+ - + {t('Auto')} diff --git a/web/src/features/channels/lib/__tests__/channel-form-admission-limits.test.ts b/web/src/features/channels/lib/__tests__/channel-form-admission-limits.test.ts new file mode 100644 index 000000000000..965f0eaccbde --- /dev/null +++ b/web/src/features/channels/lib/__tests__/channel-form-admission-limits.test.ts @@ -0,0 +1,113 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { describe, test } from 'node:test' + +import type { Channel } from '../../types' +import { + CHANNEL_FORM_DEFAULT_VALUES, + MAX_CHANNEL_ADMISSION_LIMIT, + buildSettingJSON, + channelFormSchema, + transformChannelToFormDefaults, +} from '../channel-form' + +function channelWithSetting(setting: string): Channel { + return { + id: 1, + type: 1, + key: '', + status: 1, + name: 'limited upstream', + created_time: 1, + test_time: 0, + response_time: 0, + other: '', + balance: 0, + balance_updated_time: 0, + models: 'gpt-test', + group: 'default', + used_quota: 0, + other_info: '', + setting, + remark: '', + max_input_tokens: 0, + channel_info: { + is_multi_key: false, + multi_key_size: 0, + multi_key_polling_index: 0, + multi_key_mode: 'random', + }, + settings: '{}', + } +} + +function validForm() { + return { + ...CHANNEL_FORM_DEFAULT_VALUES, + name: 'limited upstream', + models: 'gpt-test', + } +} + +describe('channel admission limit form', () => { + test('loads and serializes positive channel limits', () => { + const values = transformChannelToFormDefaults( + channelWithSetting('{"max_concurrency":20,"rpm_limit":120}') + ) + + assert.equal(values.max_concurrency, 20) + assert.equal(values.rpm_limit, 120) + assert.deepEqual(JSON.parse(buildSettingJSON(values)), { + force_format: false, + thinking_to_content: false, + proxy: '', + pass_through_body_enabled: false, + system_prompt: '', + system_prompt_override: false, + max_concurrency: 20, + rpm_limit: 120, + }) + }) + + test('omits zero limits so existing channels remain unlimited', () => { + const setting = JSON.parse(buildSettingJSON(validForm())) + + assert.equal('max_concurrency' in setting, false) + assert.equal('rpm_limit' in setting, false) + }) + + test('rejects negative, fractional, and excessive limits', () => { + const invalidValues = [-1, 1.5, MAX_CHANNEL_ADMISSION_LIMIT + 1] + + for (const value of invalidValues) { + const concurrencyResult = channelFormSchema.safeParse({ + ...validForm(), + max_concurrency: value, + }) + const rpmResult = channelFormSchema.safeParse({ + ...validForm(), + rpm_limit: value, + }) + + assert.equal(concurrencyResult.success, false) + assert.equal(rpmResult.success, false) + } + }) +}) diff --git a/web/src/features/channels/lib/channel-form-errors.ts b/web/src/features/channels/lib/channel-form-errors.ts index 92716038462e..4d885d666108 100644 --- a/web/src/features/channels/lib/channel-form-errors.ts +++ b/web/src/features/channels/lib/channel-form-errors.ts @@ -41,6 +41,8 @@ const ADVANCED_SETTINGS_FIELDS = new Set>([ 'proxy', 'http_protocol', 'http2_connection_shards', + 'max_concurrency', + 'rpm_limit', 'system_prompt', 'system_prompt_override', 'allow_service_tier', diff --git a/web/src/features/channels/lib/channel-form.ts b/web/src/features/channels/lib/channel-form.ts index 22f07931e4e2..cba22cee0885 100644 --- a/web/src/features/channels/lib/channel-form.ts +++ b/web/src/features/channels/lib/channel-form.ts @@ -73,6 +73,9 @@ function isOptionalProxyURL(value: string | undefined): boolean { export const HTTP_PROTOCOL_AUTO = 'auto' export const HTTP_PROTOCOL_HTTP1 = 'http1' export const MAX_HTTP2_CONNECTION_SHARDS = 8 +export const MAX_CHANNEL_ADMISSION_LIMIT = 1_000_000 +const CHANNEL_ADMISSION_LIMIT_ERROR = + 'Channel limits must be whole numbers from 0 to 1000000' export function normalizeHttpProtocol( value: string | undefined | null @@ -258,6 +261,18 @@ export const channelFormSchema = z .refine(isOptionalProxyURL, ERROR_MESSAGES.INVALID_PROXY), http_protocol: z.enum(['auto', 'http1']).optional(), http2_connection_shards: z.number().int().optional(), + max_concurrency: z + .number() + .int(CHANNEL_ADMISSION_LIMIT_ERROR) + .min(0, CHANNEL_ADMISSION_LIMIT_ERROR) + .max(MAX_CHANNEL_ADMISSION_LIMIT, CHANNEL_ADMISSION_LIMIT_ERROR) + .optional(), + rpm_limit: z + .number() + .int(CHANNEL_ADMISSION_LIMIT_ERROR) + .min(0, CHANNEL_ADMISSION_LIMIT_ERROR) + .max(MAX_CHANNEL_ADMISSION_LIMIT, CHANNEL_ADMISSION_LIMIT_ERROR) + .optional(), pass_through_body_enabled: z.boolean().optional(), system_prompt: z.string().optional(), system_prompt_override: z.boolean().optional(), @@ -430,6 +445,8 @@ export const CHANNEL_FORM_DEFAULT_VALUES: ChannelFormValues = { proxy: '', http_protocol: HTTP_PROTOCOL_AUTO, http2_connection_shards: 1, + max_concurrency: 0, + rpm_limit: 0, pass_through_body_enabled: false, system_prompt: '', system_prompt_override: false, @@ -470,6 +487,8 @@ export function transformChannelToFormDefaults( proxy: '', http_protocol: HTTP_PROTOCOL_AUTO as 'auto' | 'http1', http2_connection_shards: 1, + max_concurrency: 0, + rpm_limit: 0, pass_through_body_enabled: false, system_prompt: '', system_prompt_override: false, @@ -487,8 +506,12 @@ export function transformChannelToFormDefaults( thinking_to_content: parsed.thinking_to_content || false, proxy: parsed.proxy || '', http_protocol: protocol, - http2_connection_shards: - protocol === HTTP_PROTOCOL_HTTP1 ? 1 : shards, + http2_connection_shards: protocol === HTTP_PROTOCOL_HTTP1 ? 1 : shards, + max_concurrency: + typeof parsed.max_concurrency === 'number' + ? parsed.max_concurrency + : 0, + rpm_limit: typeof parsed.rpm_limit === 'number' ? parsed.rpm_limit : 0, pass_through_body_enabled: parsed.pass_through_body_enabled || false, system_prompt: parsed.system_prompt || '', system_prompt_override: parsed.system_prompt_override || false, @@ -623,6 +646,12 @@ export function buildSettingJSON(formData: ChannelFormValues): string { } else if (shards > 1) { settingObj.http2_connection_shards = shards } + if ((formData.max_concurrency ?? 0) > 0) { + settingObj.max_concurrency = formData.max_concurrency + } + if ((formData.rpm_limit ?? 0) > 0) { + settingObj.rpm_limit = formData.rpm_limit + } return JSON.stringify(settingObj) } diff --git a/web/src/features/channels/types.ts b/web/src/features/channels/types.ts index f7747fa21210..2ed9620dec18 100644 --- a/web/src/features/channels/types.ts +++ b/web/src/features/channels/types.ts @@ -88,6 +88,8 @@ export interface ChannelSettings { system_prompt_override?: boolean http_protocol?: 'auto' | 'http1' | string http2_connection_shards?: number + max_concurrency?: number + rpm_limit?: number } export interface ChannelOtherSettings { diff --git a/web/src/i18n/locales/en.json b/web/src/i18n/locales/en.json index b48096a7fcfa..a53b2fba2187 100644 --- a/web/src/i18n/locales/en.json +++ b/web/src/i18n/locales/en.json @@ -756,6 +756,7 @@ "Channel ID is required": "Channel ID is required", "Channel key": "Channel key", "Channel key unlocked": "Channel key unlocked", + "Channel limits must be whole numbers from 0 to 1000000": "Channel limits must be whole numbers from 0 to 1000000", "Channel Management": "Channel Management", "Channel models": "Channel models", "Channel name is required": "Channel name is required", @@ -2616,11 +2617,14 @@ "Maximum 200 characters": "Maximum 200 characters", "Maximum 500 characters. Supports Markdown and HTML.": "Maximum 500 characters. Supports Markdown and HTML.", "Maximum check-in quota": "Maximum check-in quota", + "Maximum concurrency": "Maximum concurrency", "Maximum custom groups per token": "Maximum custom groups per token", + "Maximum in-flight requests for this channel. 0 means unlimited.": "Maximum in-flight requests for this channel. 0 means unlimited.", "Maximum input window": "Maximum input window", "Maximum number of tokens each user can create. Default 1000. Setting too large may affect performance.": "Maximum number of tokens each user can create. Default 1000. Setting too large may affect performance.", "Maximum number of tokens in the response": "Maximum number of tokens in the response", "Maximum quota amount awarded for check-in": "Maximum quota amount awarded for check-in", + "Maximum request starts in a rolling 60-second window. 0 means unlimited.": "Maximum request starts in a rolling 60-second window. 0 means unlimited.", "Maximum tokens including hidden reasoning tokens": "Maximum tokens including hidden reasoning tokens", "Maximum tokens per response": "Maximum tokens per response", "Maximum tokens per user": "Maximum tokens per user", @@ -3931,6 +3935,7 @@ "Rows per page": "Rows per page", "RPM": "RPM", "RPM = requests per minute, TPM = tokens per minute, RPD = requests per day. Limits apply per token group.": "RPM = requests per minute, TPM = tokens per minute, RPD = requests per day. Limits apply per token group.", + "RPM limit": "RPM limit", "RSA Private Key (Production)": "RSA Private Key (Production)", "RSA Private Key (Sandbox)": "RSA Private Key (Sandbox)", "Rule": "Rule", diff --git a/web/src/i18n/locales/fr.json b/web/src/i18n/locales/fr.json index 638cf9abadd1..a75e1b78697f 100644 --- a/web/src/i18n/locales/fr.json +++ b/web/src/i18n/locales/fr.json @@ -756,6 +756,7 @@ "Channel ID is required": "L'ID du canal est requis", "Channel key": "Clé du canal", "Channel key unlocked": "Clé de canal déverrouillée", + "Channel limits must be whole numbers from 0 to 1000000": "Les limites du canal doivent être des entiers compris entre 0 et 1 000 000", "Channel Management": "Gestion des canaux", "Channel models": "Modèles de canaux", "Channel name is required": "Le nom du canal est requis", @@ -2616,11 +2617,14 @@ "Maximum 200 characters": "Maximum 200 caractères", "Maximum 500 characters. Supports Markdown and HTML.": "Maximum 500 caractères. Prend en charge Markdown et HTML.", "Maximum check-in quota": "Quota maximum de connexion", + "Maximum concurrency": "Concurrence maximale", "Maximum custom groups per token": "Nombre maximal de groupes personnalisés par jeton", + "Maximum in-flight requests for this channel. 0 means unlimited.": "Nombre maximal de requêtes en cours pour ce canal. 0 = illimité.", "Maximum input window": "Fenêtre d'entrée maximale", "Maximum number of tokens each user can create. Default 1000. Setting too large may affect performance.": "Nombre maximum de jetons que chaque utilisateur peut créer. Par défaut 1000. Une valeur trop élevée peut affecter les performances.", "Maximum number of tokens in the response": "Nombre maximum de jetons dans la réponse", "Maximum quota amount awarded for check-in": "Montant maximum de quota attribué pour la connexion", + "Maximum request starts in a rolling 60-second window. 0 means unlimited.": "Nombre maximal de requêtes lancées sur une fenêtre glissante de 60 secondes. 0 = illimité.", "Maximum tokens including hidden reasoning tokens": "Jetons maximum, y compris les jetons de raisonnement masqués", "Maximum tokens per response": "Nombre maximal de jetons par réponse", "Maximum tokens per user": "Nombre maximum de jetons par utilisateur", @@ -3931,6 +3935,7 @@ "Rows per page": "Lignes par page", "RPM": "RPM", "RPM = requests per minute, TPM = tokens per minute, RPD = requests per day. Limits apply per token group.": "RPM = requêtes/minute, TPM = jetons/minute, RPD = requêtes/jour. Les limites s'appliquent par groupe de jetons.", + "RPM limit": "Limite RPM", "RSA Private Key (Production)": "Clé privée RSA (Production)", "RSA Private Key (Sandbox)": "Clé privée RSA (Sandbox)", "Rule": "Règle", diff --git a/web/src/i18n/locales/ja.json b/web/src/i18n/locales/ja.json index 3b9d4ce693cc..9244fe5e10b5 100644 --- a/web/src/i18n/locales/ja.json +++ b/web/src/i18n/locales/ja.json @@ -756,6 +756,7 @@ "Channel ID is required": "チャネル ID が必要です", "Channel key": "チャネルキー", "Channel key unlocked": "チャネルキーが解除されました", + "Channel limits must be whole numbers from 0 to 1000000": "チャネル制限は0から1000000までの整数で指定してください", "Channel Management": "チャネル管理", "Channel models": "チャネルモデル", "Channel name is required": "チャネル名が必要です", @@ -2616,11 +2617,14 @@ "Maximum 200 characters": "最大200文字", "Maximum 500 characters. Supports Markdown and HTML.": "最大500文字。MarkdownとHTMLをサポートしています。", "Maximum check-in quota": "最大チェックインクォータ", + "Maximum concurrency": "最大同時実行数", "Maximum custom groups per token": "トークンごとのカスタムグループ上限", + "Maximum in-flight requests for this channel. 0 means unlimited.": "このチャネルで同時処理できるリクエストの最大数。0 は無制限です。", "Maximum input window": "最大入力ウィンドウ", "Maximum number of tokens each user can create. Default 1000. Setting too large may affect performance.": "各ユーザーが作成できる最大トークン数。デフォルトは 1000。大きすぎる値はパフォーマンスに影響を与える可能性があります。", "Maximum number of tokens in the response": "レスポンスの最大トークン数", "Maximum quota amount awarded for check-in": "チェックインで付与される最大クォータ量", + "Maximum request starts in a rolling 60-second window. 0 means unlimited.": "直近60秒のローリング期間で開始できるリクエストの最大数。0 は無制限です。", "Maximum tokens including hidden reasoning tokens": "隠れ推論トークンを含む最大トークン数", "Maximum tokens per response": "1 回の応答あたりの最大トークン数", "Maximum tokens per user": "ユーザーあたりの最大トークン数", @@ -3931,6 +3935,7 @@ "Rows per page": "ページあたりの行数", "RPM": "RPM", "RPM = requests per minute, TPM = tokens per minute, RPD = requests per day. Limits apply per token group.": "RPM = 1 分あたりリクエスト数、TPM = 1 分あたりトークン数、RPD = 1 日あたりリクエスト数。制限はトークングループ単位で適用されます。", + "RPM limit": "RPM 上限", "RSA Private Key (Production)": "RSA秘密鍵（本番）", "RSA Private Key (Sandbox)": "RSA秘密鍵（サンドボックス）", "Rule": "ルール", diff --git a/web/src/i18n/locales/ru.json b/web/src/i18n/locales/ru.json index 853dc874d463..b61d12f655e9 100644 --- a/web/src/i18n/locales/ru.json +++ b/web/src/i18n/locales/ru.json @@ -756,6 +756,7 @@ "Channel ID is required": "Требуется ID канала", "Channel key": "Ключ канала", "Channel key unlocked": "Ключ канала разблокирован", + "Channel limits must be whole numbers from 0 to 1000000": "Ограничения канала должны быть целыми числами от 0 до 1000000", "Channel Management": "Управление каналами", "Channel models": "Модели каналов", "Channel name is required": "Имя канала обязательно", @@ -2616,11 +2617,14 @@ "Maximum 200 characters": "Максимум 200 символов", "Maximum 500 characters. Supports Markdown and HTML.": "Максимум 500 символов. Поддерживает Markdown и HTML.", "Maximum check-in quota": "Максимальная квота регистрации", + "Maximum concurrency": "Максимальная параллельность", "Maximum custom groups per token": "Максимум пользовательских групп на токен", + "Maximum in-flight requests for this channel. 0 means unlimited.": "Максимальное число выполняемых запросов для этого канала. 0 — без ограничений.", "Maximum input window": "Максимальное окно ввода", "Maximum number of tokens each user can create. Default 1000. Setting too large may affect performance.": "Максимальное количество токенов, которое может создать каждый пользователь. По умолчанию 1000. Слишком большое значение может повлиять на производительность.", "Maximum number of tokens in the response": "Максимальное число токенов в ответе", "Maximum quota amount awarded for check-in": "Максимальная сумма квоты, присуждаемая за регистрацию", + "Maximum request starts in a rolling 60-second window. 0 means unlimited.": "Максимальное число запусков запросов за скользящий интервал 60 секунд. 0 — без ограничений.", "Maximum tokens including hidden reasoning tokens": "Максимум токенов с учётом скрытых reasoning-токенов", "Maximum tokens per response": "Максимум токенов на ответ", "Maximum tokens per user": "Максимальное количество токенов на пользователя", @@ -3931,6 +3935,7 @@ "Rows per page": "Строк на страницу", "RPM": "RPM", "RPM = requests per minute, TPM = tokens per minute, RPD = requests per day. Limits apply per token group.": "RPM = запросов в минуту, TPM = токенов в минуту, RPD = запросов в день. Ограничения применяются к каждой группе токенов.", + "RPM limit": "Лимит RPM", "RSA Private Key (Production)": "RSA-приватный ключ (Продакшн)", "RSA Private Key (Sandbox)": "RSA-приватный ключ (Песочница)", "Rule": "Правило", diff --git a/web/src/i18n/locales/vi.json b/web/src/i18n/locales/vi.json index cdc1a12626e1..d7315abbe00a 100644 --- a/web/src/i18n/locales/vi.json +++ b/web/src/i18n/locales/vi.json @@ -756,6 +756,7 @@ "Channel ID is required": "Cần có ID kênh", "Channel key": "Khóa kênh", "Channel key unlocked": "Khóa kênh đã được mở khóa", + "Channel limits must be whole numbers from 0 to 1000000": "Giới hạn kênh phải là số nguyên từ 0 đến 1000000", "Channel Management": "Quản lý kênh", "Channel models": "Mô hình kênh", "Channel name is required": "Tên kênh là bắt buộc", @@ -2616,11 +2617,14 @@ "Maximum 200 characters": "Tối đa 200 ký tự", "Maximum 500 characters. Supports Markdown and HTML.": "Tối đa 500 ký tự. Hỗ trợ Markdown và HTML.", "Maximum check-in quota": "Hạn ngạch điểm danh tối đa", + "Maximum concurrency": "Số yêu cầu đồng thời tối đa", "Maximum custom groups per token": "Số nhóm tùy chỉnh tối đa cho mỗi token", + "Maximum in-flight requests for this channel. 0 means unlimited.": "Số yêu cầu đang xử lý tối đa cho kênh này. 0 là không giới hạn.", "Maximum input window": "Cửa sổ nhập tối đa", "Maximum number of tokens each user can create. Default 1000. Setting too large may affect performance.": "Số lượng token tối đa mỗi người dùng có thể tạo. Mặc định là 1000. Đặt quá lớn có thể ảnh hưởng đến hiệu suất.", "Maximum number of tokens in the response": "Số token tối đa trong phản hồi", "Maximum quota amount awarded for check-in": "Số lượng hạn ngạch tối đa được trao cho điểm danh", + "Maximum request starts in a rolling 60-second window. 0 means unlimited.": "Số yêu cầu bắt đầu tối đa trong cửa sổ trượt 60 giây. 0 là không giới hạn.", "Maximum tokens including hidden reasoning tokens": "Số token tối đa bao gồm token suy luận ẩn", "Maximum tokens per response": "Số token tối đa mỗi phản hồi", "Maximum tokens per user": "Số token tối đa trên mỗi người dùng", @@ -3931,6 +3935,7 @@ "Rows per page": "Số hàng trên trang", "RPM": "RPM", "RPM = requests per minute, TPM = tokens per minute, RPD = requests per day. Limits apply per token group.": "RPM = yêu cầu mỗi phút, TPM = token mỗi phút, RPD = yêu cầu mỗi ngày. Giới hạn áp dụng cho từng nhóm token.", + "RPM limit": "Giới hạn RPM", "RSA Private Key (Production)": "RSA Private Key (Sản xuất)", "RSA Private Key (Sandbox)": "Khóa riêng RSA (Sandbox)", "Rule": "Quy tắc", diff --git a/web/src/i18n/locales/zh-TW.json b/web/src/i18n/locales/zh-TW.json index f8b5fafd8d2c..9b530a363b74 100644 --- a/web/src/i18n/locales/zh-TW.json +++ b/web/src/i18n/locales/zh-TW.json @@ -756,6 +756,7 @@ "Channel ID is required": "缺少渠道 ID", "Channel key": "渠道金鑰", "Channel key unlocked": "渠道金鑰已解鎖", + "Channel limits must be whole numbers from 0 to 1000000": "渠道限制必須是 0 到 1000000 之間的整數", "Channel Management": "渠道管理", "Channel models": "渠道模型", "Channel name is required": "渠道名稱是必填的", @@ -2616,11 +2617,14 @@ "Maximum 200 characters": "最多 200 個字元", "Maximum 500 characters. Supports Markdown and HTML.": "最多 500 個字元。支援 Markdown 和 HTML。", "Maximum check-in quota": "簽到最大額度", + "Maximum concurrency": "最大並發數", "Maximum custom groups per token": "每個令牌的最大自訂分組數", + "Maximum in-flight requests for this channel. 0 means unlimited.": "此渠道允許的最大進行中請求數。0 表示不限。", "Maximum input window": "最大輸入窗口", "Maximum number of tokens each user can create. Default 1000. Setting too large may affect performance.": "每個用戶可建立的最大令牌數量。預設 1000。設定過大可能會影響效能。", "Maximum number of tokens in the response": "回應中最大 token 數", "Maximum quota amount awarded for check-in": "簽到獎勵的最大額度", + "Maximum request starts in a rolling 60-second window. 0 means unlimited.": "滾動 60 秒視窗內允許開始的最大請求數。0 表示不限。", "Maximum tokens including hidden reasoning tokens": "最大 token 數（含隱藏的推理 token）", "Maximum tokens per response": "單次回應最大 token 數", "Maximum tokens per user": "每個用戶的最大令牌數", @@ -3931,6 +3935,7 @@ "Rows per page": "每頁行數", "RPM": "RPM", "RPM = requests per minute, TPM = tokens per minute, RPD = requests per day. Limits apply per token group.": "RPM = 每分鐘請求數，TPM = 每分鐘 token 數，RPD = 每日請求數。限制按令牌分組生效。", + "RPM limit": "RPM 上限", "RSA Private Key (Production)": "RSA 私鑰（生產）", "RSA Private Key (Sandbox)": "RSA 私鑰（沙盒）", "Rule": "規則", diff --git a/web/src/i18n/locales/zh.json b/web/src/i18n/locales/zh.json index 842991e714f0..338283f48245 100644 --- a/web/src/i18n/locales/zh.json +++ b/web/src/i18n/locales/zh.json @@ -756,6 +756,7 @@ "Channel ID is required": "缺少渠道 ID", "Channel key": "渠道密钥", "Channel key unlocked": "渠道密钥已解锁", + "Channel limits must be whole numbers from 0 to 1000000": "渠道限制必须是 0 到 1000000 之间的整数", "Channel Management": "渠道管理", "Channel models": "渠道模型", "Channel name is required": "渠道名称是必填的", @@ -2616,11 +2617,14 @@ "Maximum 200 characters": "最多 200 个字符", "Maximum 500 characters. Supports Markdown and HTML.": "最多 500 个字符。支持 Markdown 和 HTML。", "Maximum check-in quota": "签到最大额度", + "Maximum concurrency": "最大并发数", "Maximum custom groups per token": "每个令牌的最大自定义分组数", + "Maximum in-flight requests for this channel. 0 means unlimited.": "此渠道允许的最大在途请求数。0 表示不限。", "Maximum input window": "最大输入窗口", "Maximum number of tokens each user can create. Default 1000. Setting too large may affect performance.": "每个用户可创建的最大令牌数量。默认 1000。设置过大可能会影响性能。", "Maximum number of tokens in the response": "响应中最大 token 数", "Maximum quota amount awarded for check-in": "签到奖励的最大额度", + "Maximum request starts in a rolling 60-second window. 0 means unlimited.": "滚动 60 秒窗口内允许开始的最大请求数。0 表示不限。", "Maximum tokens including hidden reasoning tokens": "最大 token 数（含隐藏的推理 token）", "Maximum tokens per response": "单次响应最大 token 数", "Maximum tokens per user": "每个用户的最大令牌数", @@ -3931,6 +3935,7 @@ "Rows per page": "每页行数", "RPM": "RPM", "RPM = requests per minute, TPM = tokens per minute, RPD = requests per day. Limits apply per token group.": "RPM = 每分钟请求数，TPM = 每分钟 token 数，RPD = 每日请求数。限制按令牌分组生效。", + "RPM limit": "RPM 上限", "RSA Private Key (Production)": "RSA 私钥（生产）", "RSA Private Key (Sandbox)": "RSA 私钥（沙盒）", "Rule": "规则", diff --git a/web/src/i18n/static-keys.ts b/web/src/i18n/static-keys.ts index a435b0471155..2676dcb62707 100644 --- a/web/src/i18n/static-keys.ts +++ b/web/src/i18n/static-keys.ts @@ -537,6 +537,9 @@ export const STATIC_I18N_KEYS = [ 'OpenAI Models upstream path must not contain {model}', 'OpenAI Models route is required to enable upstream model checks', + // Channel admission validation + 'Channel limits must be whole numbers from 0 to 1000000', + // Dashboard flow stages (labels/descriptions passed to t at runtime) 'User', 'Node',