diff --git a/controller/channel.go b/controller/channel.go index 21735170..d951a20a 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -510,6 +510,20 @@ func validateChannel(channel *model.Channel, isAdd bool) error { } } + // 校验 quota_reset_period + if channel.QuotaResetPeriod != "" { + normalized := model.NormalizeChannelResetPeriod(channel.QuotaResetPeriod) + if normalized == model.ChannelResetNever && channel.QuotaResetPeriod != "never" && channel.QuotaResetPeriod != "" { + return fmt.Errorf("invalid quota_reset_period: %s", channel.QuotaResetPeriod) + } + channel.QuotaResetPeriod = normalized + } + + // custom 模式必须有 customSeconds > 0 + if model.NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == model.ChannelResetCustom && channel.QuotaResetCustomSeconds <= 0 { + return fmt.Errorf("quota_reset_custom_seconds must be > 0 when quota_reset_period is 'custom'") + } + return nil } @@ -602,6 +616,7 @@ func AddChannel(c *gin.Context) { } addChannelRequest.Channel.CreatedTime = common.GetTimestamp() + model.InitChannelNextResetTime(addChannelRequest.Channel) keys := make([]string, 0) switch addChannelRequest.Mode { case "multi_to_single": @@ -974,6 +989,7 @@ func UpdateChannel(c *gin.Context) { // 覆盖模式:直接使用新密钥(默认行为,不需要特殊处理) } } + model.InitChannelNextResetTime(&channel.Channel) err = channel.Update() if err != nil { common.ApiError(c, err) diff --git a/docs/agent-specs/2026-05-19-channel-subscription-quota-execution.md b/docs/agent-specs/2026-05-19-channel-subscription-quota-execution.md new file mode 100644 index 00000000..fc0062dd --- /dev/null +++ b/docs/agent-specs/2026-05-19-channel-subscription-quota-execution.md @@ -0,0 +1,524 @@ +# 渠道订阅类型额度实现方案 + +**版本:** v2.0 +**日期:** 2026-05-19 +**状态:** 执行版 + +--- + +## 1. 目标与范围 + +参照 `model/subscription.go` 的订阅额度重置机制,为 Channel 增加**额度上限 + 周期重置**功能。详细设计决策见审阅版。 + +--- + +## 2. 参考实现索引 + +| 功能 | 文件:行号 | +|------|-----------| +| 重置周期常量 | `model/subscription.go:26-33` | +| `NormalizeResetPeriod` | `model/subscription.go:300-307` | +| `calcNextResetTime` | `model/subscription.go:309-348` | +| `maybeResetUserSubscriptionWithPlanTx` | `model/subscription.go:933-967` | +| `ResetDueSubscriptions` | `model/subscription.go:1100-1140` | +| 定时任务 | `service/subscription_reset_task.go:29-93` | +| 定时任务启动 | `main.go:120` | +| 渠道额度扣减 | `model/channel.go:824-837` | +| 请求完成后扣减调用点 | `service/text_quota.go:370` | +| 渠道校验 | `controller/channel.go:457-514` | +| 渠道缓存 | `model/channel_cache.go:193-205` | + +--- + +## 3. Channel 模型新增字段(修改 `model/channel.go`) + +```go +// 新增字段,接在 UsedQuota 之后 +QuotaLimit int64 `json:"quota_limit" gorm:"bigint;default:0"` +QuotaResetPeriod string `json:"quota_reset_period" gorm:"type:varchar(16);default:'never'"` +QuotaResetCustomSeconds int `json:"quota_reset_custom_seconds" gorm:"type:int;default:0"` +LastResetTime int64 `json:"last_reset_time" gorm:"bigint;default:0"` +NextResetTime int64 `json:"next_reset_time" gorm:"bigint;default:0"` +``` + +**语义**: +- `QuotaLimit = 0` → 不启用额度管控(原有行为) +- `QuotaLimit > 0` + `QuotaResetPeriod != 'never'` → 启用额度管控 + 周期重置 +- `QuotaLimit > 0` + `QuotaResetPeriod = 'never'` → 启用额度管控,永不重置 +- `UsedQuota` 沿用现有字段,重置时清零 + +--- + +## 4. 新增 `model/channel_quota.go` + +### 4.1 重置周期常量 + +```go +const ( + ChannelResetNever = "never" + ChannelResetDaily = "daily" + ChannelResetWeekly = "weekly" + ChannelResetMonthly = "monthly" + ChannelResetCustom = "custom" +) +``` + +与 `subscription.go` 的常量值一致,命名前缀改为 `Channel` 避免冲突。 + +### 4.2 NormalizeChannelResetPeriod + +```go +func NormalizeChannelResetPeriod(period string) string { + switch strings.TrimSpace(period) { + case ChannelResetDaily, ChannelResetWeekly, ChannelResetMonthly, ChannelResetCustom: + return strings.TrimSpace(period) + default: + return ChannelResetNever + } +} +``` + +### 4.3 calcChannelNextResetTime + +与 `subscription.go:calcNextResetTime` 逻辑一致,但接收 `period` 和 `customSeconds` 参数而非 `*SubscriptionPlan`: + +```go +func calcChannelNextResetTime(base time.Time, period string, customSeconds int) int64 { + period = NormalizeChannelResetPeriod(period) + if period == ChannelResetNever { + return 0 + } + var next time.Time + switch period { + case ChannelResetDaily: + next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()). + AddDate(0, 0, 1) + case ChannelResetWeekly: + weekday := int(base.Weekday()) + if weekday == 0 { + weekday = 7 + } + daysUntil := 8 - weekday + next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()). + AddDate(0, 0, daysUntil) + case ChannelResetMonthly: + next = time.Date(base.Year(), base.Month(), 1, 0, 0, 0, 0, base.Location()). + AddDate(0, 1, 0) + case ChannelResetCustom: + if customSeconds <= 0 { + return 0 + } + next = base.Add(time.Duration(customSeconds) * time.Second) + default: + return 0 + } + return next.Unix() +} +``` + +与 `calcNextResetTime` 的区别:无 `endUnix` 参数(渠道无到期时间)。 + +### 4.4 maybeResetChannelQuotaTx + +与 `subscription.go:maybeResetUserSubscriptionWithPlanTx` 逻辑一致,针对 Channel: + +```go +func maybeResetChannelQuotaTx(tx *gorm.DB, channel *Channel, now int64) error { + if tx == nil || channel == nil { + return errors.New("invalid reset args") + } + if channel.NextResetTime > 0 && channel.NextResetTime > now { + return nil + } + if NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == ChannelResetNever { + return nil + } + baseUnix := channel.LastResetTime + if baseUnix <= 0 { + baseUnix = channel.CreatedTime + } + base := time.Unix(baseUnix, 0) + next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) + + advanced := false + for next > 0 && next <= now { + advanced = true + base = time.Unix(next, 0) + next = calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) + } + + if !advanced { + if channel.NextResetTime == 0 && next > 0 { + return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ + "next_reset_time": next, + "last_reset_time": base.Unix(), + }).Error + } + return nil + } + + // 重置额度,只用 Updates 不用 Save,避免覆盖并发更新 + return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ + "used_quota": 0, + "last_reset_time": base.Unix(), + "next_reset_time": next, + }).Error +} +``` + +**与订阅实现的区别**:用 `tx.Model(&Channel{}).Where("id = ?").Updates(...)` 替代 `tx.Save(channel)`,只更新重置相关字段,避免覆盖 `UsedQuota` 的并发扣减。 + +### 4.5 ResetDueChannels + +与 `subscription.go:ResetDueSubscriptions` 逻辑一致: + +```go +func ResetDueChannels(limit int) (int, error) { + if limit <= 0 { + limit = 200 + } + now := GetDBTimestamp() + var channels []Channel + if err := DB.Where("next_reset_time > 0 AND next_reset_time <= ? AND status = ? AND quota_limit > 0", + now, common.ChannelStatusEnabled). + Order("next_reset_time asc"). + Limit(limit). + Find(&channels).Error; err != nil { + return 0, err + } + if len(channels) == 0 { + return 0, nil + } + resetCount := 0 + for _, ch := range channels { + chCopy := ch + err := DB.Transaction(func(tx *gorm.DB) error { + var locked Channel + if err := tx.Set("gorm:query_option", "FOR UPDATE"). + Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", chCopy.Id, now). + First(&locked).Error; err != nil { + return nil + } + if err := maybeResetChannelQuotaTx(tx, &locked, now); err != nil { + return err + } + resetCount++ + return nil + }) + if err != nil { + return resetCount, err + } + } + return resetCount, nil +} +``` + +**与订阅实现的区别**: +- 查询条件增加 `AND quota_limit > 0`,跳过未启用额度管控的渠道 +- 无需加载关联的 Plan(渠道的重置参数直接在 Channel 上) + +### 4.6 CheckChannelQuota + +```go +func CheckChannelQuota(channelId int) error { + var channel Channel + if err := DB.Select("id, used_quota, quota_limit, quota_reset_period, next_reset_time"). + Where("id = ?", channelId).First(&channel).Error; err != nil { + return err + } + // 不启用额度管控 + if channel.QuotaLimit <= 0 { + return nil + } + // 懒重置:在额度检查之前 + if channel.NextResetTime > 0 { + now := GetDBTimestamp() + if channel.NextResetTime <= now { + err := DB.Transaction(func(tx *gorm.DB) error { + var locked Channel + if err := tx.Set("gorm:query_option", "FOR UPDATE"). + Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", channelId, now). + First(&locked).Error; err != nil { + return nil // 已被其他事务处理 + } + return maybeResetChannelQuotaTx(tx, &locked, now) + }) + if err != nil { + return err + } + // 重新读取重置后的值 + if err := DB.Select("used_quota, quota_limit"). + Where("id = ?", channelId).First(&channel).Error; err != nil { + return err + } + } + } + if channel.UsedQuota >= channel.QuotaLimit { + return fmt.Errorf("channel quota exceeded (used=%d, limit=%d)", channel.UsedQuota, channel.QuotaLimit) + } + return nil +} +``` + +**设计要点**: +- 直接查 DB,不依赖缓存(缓存 `UsedQuota` 可能过期) +- 懒重置在额度检查之前执行(修复原方案的问题 1) +- 重置后重新读取确保拿到最新值 + +### 4.7 InitChannelNextResetTime + +创建/更新渠道时初始化 `NextResetTime`: + +```go +func InitChannelNextResetTime(channel *Channel) { + if channel.QuotaLimit <= 0 { + channel.NextResetTime = 0 + channel.LastResetTime = 0 + return + } + period := NormalizeChannelResetPeriod(channel.QuotaResetPeriod) + if period == ChannelResetNever { + channel.NextResetTime = 0 + channel.LastResetTime = 0 + return + } + baseUnix := channel.LastResetTime + if baseUnix <= 0 { + baseUnix = channel.CreatedTime + } + if baseUnix <= 0 { + baseUnix = GetDBTimestamp() + } + base := time.Unix(baseUnix, 0) + next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) + channel.NextResetTime = next + if channel.LastResetTime <= 0 { + channel.LastResetTime = baseUnix + } +} +``` + +--- + +## 5. 新增 `service/channel_quota_reset_task.go` + +与 `service/subscription_reset_task.go` 结构一致,仅替换调用: + +```go +package service + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + + "github.com/bytedance/gopkg/util/gopool" +) + +const ( + channelQuotaResetTickInterval = 1 * time.Minute + channelQuotaResetBatchSize = 300 +) + +var ( + channelQuotaResetOnce sync.Once + channelQuotaResetRunning atomic.Bool +) + +func StartChannelQuotaResetTask() { + channelQuotaResetOnce.Do(func() { + if !common.IsMasterNode { + return + } + gopool.Go(func() { + logger.LogInfo(context.Background(), fmt.Sprintf("channel quota reset task started: tick=%s", channelQuotaResetTickInterval)) + ticker := time.NewTicker(channelQuotaResetTickInterval) + defer ticker.Stop() + + runChannelQuotaResetOnce() + for range ticker.C { + runChannelQuotaResetOnce() + } + }) + }) +} + +func runChannelQuotaResetOnce() { + if !channelQuotaResetRunning.CompareAndSwap(false, true) { + return + } + defer channelQuotaResetRunning.Store(false) + + ctx := context.Background() + totalReset := 0 + for { + n, err := model.ResetDueChannels(channelQuotaResetBatchSize) + if err != nil { + logger.LogWarn(ctx, fmt.Sprintf("channel quota reset task failed: %v", err)) + return + } + if n == 0 { + break + } + totalReset += n + if n < channelQuotaResetBatchSize { + break + } + } + if common.DebugEnabled && totalReset > 0 { + logger.LogDebug(ctx, "channel quota maintenance: reset_count=%d", totalReset) + } +} +``` + +--- + +## 6. 修改 `service/text_quota.go` + +在 `PostTextConsumeQuota` 中,请求扣减前增加额度检查: + +```go +// 原代码 (line 369-370): +model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota) +model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota) + +// 修改后: +model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota) +// 渠道额度管控:检查额度是否超限 +if err := model.CheckChannelQuota(relayInfo.ChannelId); err != nil { + // 额度已超限,仍扣减已消耗的额度,但记录警告 + logger.LogWarn(ctx, fmt.Sprintf("channel quota exceeded: channel_id=%d, %v", relayInfo.ChannelId, err)) +} +model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota) +``` + +**注意**:额度检查放在**请求完成后扣减前**,而非请求前。原因: +1. 请求前的 `estimatedQuota` 不可靠(流式请求无法预估) +2. 请求已完成,必须扣减实际消耗 +3. 超限后下次请求会被 `distributor` 拒绝(见第 7 节) + +这样设计:**当前请求正常完成并扣费,但下次请求会被拒绝**。这是"实时扣减"模式的自然语义——不做预扣,只做事后拦截。 + +--- + +## 7. 修改 `middleware/distributor.go` + +在渠道选中后、`c.Next()` 前,检查渠道额度: + +```go +// 原代码 (line 158-160): +common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now()) +SetupContextForSelectedChannel(c, channel, modelRequest.Model) +c.Next() + +// 修改后: +common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now()) +SetupContextForSelectedChannel(c, channel, modelRequest.Model) +// 渠道额度管控:请求前检查 +if err := model.CheckChannelQuota(channel.Id); err != nil { + abortWithOpenAiMessage(c, http.StatusTooManyRequests, err.Error()) + return +} +c.Next() +``` + +**此检查是真正的拦截点**: +- `CheckChannelQuota` 内部会先执行懒重置,再检查 `UsedQuota >= QuotaLimit` +- 额度超限时直接返回 429,不转发请求到上游 +- `text_quota.go` 中的检查是兜底(防止并发场景下漏过) + +--- + +## 8. 修改 `controller/channel.go` + +### 8.1 validateChannel 增加校验 + +```go +// 在 validateChannel 末尾 return nil 之前添加: + +// 校验 quota_reset_period +if channel.QuotaResetPeriod != "" { + normalized := model.NormalizeChannelResetPeriod(channel.QuotaResetPeriod) + if normalized == model.ChannelResetNever && channel.QuotaResetPeriod != "never" && channel.QuotaResetPeriod != "" { + return fmt.Errorf("invalid quota_reset_period: %s", channel.QuotaResetPeriod) + } + channel.QuotaResetPeriod = normalized +} + +// custom 模式必须有 customSeconds > 0 +if model.NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == model.ChannelResetCustom && channel.QuotaResetCustomSeconds <= 0 { + return fmt.Errorf("quota_reset_custom_seconds must be > 0 when quota_reset_period is 'custom'") +} +``` + +### 8.2 创建/更新渠道时初始化 NextResetTime + +在 `AddChannel` 和 `UpdateChannel` 中,保存前调用: + +```go +model.InitChannelNextResetTime(channel) +``` + +这确保 `NextResetTime` 不为 0(修复原方案的问题 2),定时任务能正确选中需要重置的渠道。 + +--- + +## 9. 修改 `main.go` + +在 `StartSubscriptionQuotaResetTask()` 后添加: + +```go +service.StartChannelQuotaResetTask() +``` + +--- + +## 10. 数据库迁移 + +GORM `AutoMigrate` 自动处理字段新增。`Channel{}` 已在 `model/main.go` 的 AutoMigrate 列表中,新增字段会自动添加。 + +无需手动 SQL。 + +--- + +## 11. 文件清单 + +### 11.1 新增 + +| 文件 | 说明 | +|------|------| +| `model/channel_quota.go` | 常量、重置时间计算、重置事务、ResetDueChannels、CheckChannelQuota、InitChannelNextResetTime | +| `service/channel_quota_reset_task.go` | 定时任务 | + +### 11.2 修改 + +| 文件 | 修改内容 | +|------|----------| +| `model/channel.go` | 新增 5 个字段 | +| `service/text_quota.go` | 扣减前额度检查(兜底) | +| `middleware/distributor.go` | 请求前额度检查(主拦截点) | +| `controller/channel.go` | 校验 + 初始化 NextResetTime | +| `main.go` | 启动定时任务 | + +--- + +## 12. 前端(后续迭代) + +本次实现不含前端修改。前端渠道编辑表单增加 `QuotaLimit` 和 `QuotaResetPeriod` 字段、列表展示额度使用率,列为后续迭代。 + +--- + +## 13. 测试要点 + +1. **重置周期正确性**:daily/weekly/monthly/custom,验证 `NextResetTime` 计算正确 +2. **懒重置**:定时任务未触发时,请求前懒检查能补救 +3. **额度拦截**:`UsedQuota >= QuotaLimit` 时请求返回 429 +4. **兼容性**:未设置 `QuotaLimit` 的历史渠道行为不变 +5. **NextResetTime 初始化**:创建/更新渠道后 `NextResetTime` 正确 +6. **三库兼容**:SQLite / MySQL / PostgreSQL 均通过 diff --git a/docs/agent-specs/2026-05-19-channel-subscription-quota-review.md b/docs/agent-specs/2026-05-19-channel-subscription-quota-review.md new file mode 100644 index 00000000..11931409 --- /dev/null +++ b/docs/agent-specs/2026-05-19-channel-subscription-quota-review.md @@ -0,0 +1,258 @@ +# 渠道订阅类型额度设计文档 + +**版本:** v2.0 +**日期:** 2026-05-19 +**状态:** 审阅版 + +--- + +## 1. 背景与问题 + +### 1.1 现状 + +Channel 模型仅有两种额度相关字段: + +| 字段 | 语义 | 来源 | +|------|------|------| +| `Balance` | 上游 API 实际余额(USD) | 上游 API 查询 | +| `UsedQuota` | 累计消耗量 | 本地扣减,永不重置 | + +两者都无法满足"固定额度上限 + 周期重置"的运营需求。 + +### 1.2 需求场景 + +运营者希望对渠道配额管控: + +- 渠道 A:每天最多 1 万 quota,超限拒绝请求 +- 渠道 B:每月最多 50 万 quota,超限拒绝请求 +- 渠道 C:不设上限,仅记录消耗(与现有 `UsedQuota` 等价) + +### 1.3 已有参考实现 + +`model/subscription.go` 中已实现完整的用户订阅额度重置机制: + +```go +// 订阅套餐可配置重置周期 +QuotaResetPeriod // daily/weekly/monthly/custom/never + +// 用户订阅记录 +AmountTotal // 总额度 +AmountUsed // 已消耗 +LastResetTime // 上次重置 +NextResetTime // 下次重置 +``` + +渠道额度直接复用此设计模式,但消费模式不同(见 3.1)。 + +--- + +## 2. 目标与非目标 + +### 2.1 目标 + +- 渠道可配置额度上限(`QuotaLimit`) +- 渠道可配置重置周期(`QuotaResetPeriod`:daily / weekly / monthly / custom / never) +- 周期到达时自动重置 `UsedQuota` 为 0 +- 额度耗尽时拒绝请求(返回 429) + +### 2.2 非目标 + +- 不支持预扣 + 结算模式(与订阅的 `PreConsume` 不同) +- 不支持按渠道内单个 key 分别管控 +- 不支持按模型分组分别设置额度 +- 不支持额度耗尽自动降级到备用渠道(仅拒绝) +- 不支持额度使用率告警通知 + +--- + +## 3. 核心决策 + +### 3.1 消费模式:实时扣减(非预扣 + 结算) + +订阅采用"预扣 + 结算"双阶段,因为订阅钱包需要精确余额。 + +渠道采用**实时扣减**,更简单: + +1. 请求前检查:`UsedQuota >= QuotaLimit` → 拒绝 +2. 请求完成:`UsedQuota += actualQuota` + +不做预估算额。`UsedQuota >= QuotaLimit` 时拒绝,意味着当前周期内已消耗量达到上限。由于请求完成才扣减,最后一个被放行的请求可能实际消耗后超过 `QuotaLimit`,但这是可接受的——与 `UpdateChannelUsedQuota` 的原子递增语义一致。 + +### 3.2 重置时机:定时任务 + 懒检查兜底 + +与订阅一致: + +```go +// 定时任务(每分钟) +ResetDueChannels() + +// 请求前懒检查(定时任务未及时触发时补救) +maybeResetChannelQuotaTx() +``` + +**关键**:懒检查必须在额度检查之前执行,否则重置周期刚到时旧 `UsedQuota` 仍会拒绝合法请求。 + +### 3.3 并发安全:FOR UPDATE 行锁 + Updates 精确更新 + +```go +tx.Set("gorm:query_option", "FOR UPDATE"). + Where("id = ? AND next_reset_time <= ?", channelId, now). + First(&locked) +``` + +重置时只更新相关字段,不用 `Save` 避免覆盖并发更新: + +```go +tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ + "used_quota": 0, + "last_reset_time": base.Unix(), + "next_reset_time": next, +}) +``` + +### 3.4 额度耗尽行为:返回 429 + +额度耗尽时返回 HTTP 429,不自动降级。降级涉及渠道选择逻辑改动,复杂度高,列为非目标。 + +### 3.5 Balance 与 QuotaLimit 的关系 + +两者独立共存: + +| 字段 | 语义 | 来源 | +|------|------|------| +| `Balance` | 上游 API 实际余额 | 上游 API 查询 | +| `UsedQuota` | 本地累计消耗 | 本地扣减 | +| `QuotaLimit` | 运营者设定的额度上限 | 配置 | + +`Balance` 是上游真实余额,`QuotaLimit` 是本地配额管控,两者互不影响。运营者可同时使用或只选其一。 + +### 3.6 多 Key 模式 + +`UsedQuota` 是整个渠道的累计,不区分单个 key。多 key 模式下仅全局管控,与现有行为一致。 + +--- + +## 4. 影响范围 + +### 4.1 新增文件 + +| 文件 | 说明 | +|------|------| +| `model/channel_quota.go` | 渠道额度核心逻辑(重置时间计算、重置事务、ResetDueChannels、CheckChannelQuota) | +| `service/channel_quota_reset_task.go` | 渠道额度重置定时任务 | + +### 4.2 修改文件 + +| 文件 | 修改内容 | +|------|----------| +| `model/channel.go` | 新增 5 个字段 | +| `service/text_quota.go` | 请求前懒重置 + 额度检查 | +| `controller/channel.go` | 校验新字段 + 创建/更新时初始化 NextResetTime | +| `main.go` | 启动定时任务 | + +### 4.3 数据库字段变更 + +```sql +ALTER TABLE channels ADD COLUMN quota_limit BIGINT DEFAULT 0; +ALTER TABLE channels ADD COLUMN quota_reset_period VARCHAR(16) DEFAULT 'never'; +ALTER TABLE channels ADD COLUMN quota_reset_custom_seconds INT DEFAULT 0; +ALTER TABLE channels ADD COLUMN last_reset_time BIGINT DEFAULT 0; +ALTER TABLE channels ADD COLUMN next_reset_time BIGINT DEFAULT 0; +``` + +`used_quota` 已存在,沿用。现有渠道 `quota_reset_period = 'never'` 即为原有行为。 + +--- + +## 5. 数据流 + +```mermaid +sequenceDiagram + participant Req as 请求 + participant Relay as Relay 层 (text_quota.go) + participant Model as channel_quota.go + participant DB as Database + + Req->>Relay: 请求完成,准备扣减 + Relay->>Model: 懒重置检查 (channelId) + Model->>DB: SELECT ... FOR UPDATE + alt NextResetTime <= now + Model->>DB: UPDATE used_quota=0, last/next_reset_time + end + Relay->>Model: 额度检查 (channelId) + alt UsedQuota >= QuotaLimit (且 QuotaLimit > 0) + Relay-->>Req: 429 Channel quota exceeded + else + Relay->>DB: UPDATE used_quota += actual + end +``` + +```mermaid +sequenceDiagram + participant Cron as 定时任务 (每分钟) + participant Model as channel_quota.go + participant DB as Database + + Cron->>Model: ResetDueChannels(300) + Model->>DB: SELECT next_reset_time <= now AND quota_limit > 0 + loop 每条渠道 + Model->>DB: FOR UPDATE 锁行 + Model->>DB: UPDATE used_quota=0, last/next_reset_time + end +``` + +--- + +## 6. 风险与兼容性 + +| 风险 | 影响 | 缓解措施 | +|------|------|----------| +| `quota_reset_period = 'never'` 时 `UsedQuota` 行为不变 | 无 | 默认值即为此 | +| 历史渠道 `next_reset_time = 0` | 无 | 定时任务跳过,懒检查也跳过 | +| 重置时刻有进行中请求 | 低 | `FOR UPDATE` 锁保证串行 | +| 最后一个放行请求可能超用 | 低 | 可接受,与现有行为一致 | +| SQLite 下 `FOR UPDATE` 无效 | 低 | SQLite 自带数据库级锁,事务串行化 | +| 缓存中 `UsedQuota` 可能过期 | 低 | 额度检查直接查 DB,不依赖缓存 | + +--- + +## 7. 验收标准 + +### 7.1 功能验收 + +- [ ] 新增/编辑渠道时可配置 `quota_limit` 和 `quota_reset_period` +- [ ] `QuotaLimit > 0` 且 `UsedQuota >= QuotaLimit` 时请求返回 429 +- [ ] `QuotaLimit = 0` 时不做额度检查(原有行为) +- [ ] 每日重置:次日 00:00 `UsedQuota` 清零 +- [ ] 每月重置:每月 1 日 00:00 清零 +- [ ] 自定义秒数重置:按配置精确重置 +- [ ] 永不重置:`UsedQuota` 持续累计 +- [ ] 创建/更新渠道时正确初始化 `NextResetTime` + +### 7.2 兼容性验收 + +- [ ] 未设置 `QuotaLimit` 的历史渠道行为不变 +- [ ] GORM AutoMigrate 正确添加新字段 +- [ ] SQLite / MySQL / PostgreSQL 均兼容 + +--- + +## 8. 回滚与降级 + +### 8.1 回滚 + +部署后如需回滚:新字段由 GORM AutoMigrate 添加,旧代码忽略未知字段,向前兼容。回滚后定时任务不运行,`UsedQuota` 恢复持续累计行为。 + +### 8.2 降级 + +关闭单个渠道的额度管控:设 `QuotaLimit = 0` + `QuotaResetPeriod = 'never'`,等价于原有行为。 + +--- + +## 9. 参考实现 + +- `model/subscription.go:300-307` — `NormalizeResetPeriod` +- `model/subscription.go:309-348` — `calcNextResetTime` +- `model/subscription.go:933-967` — `maybeResetUserSubscriptionWithPlanTx` +- `model/subscription.go:1100-1140` — `ResetDueSubscriptions` +- `service/subscription_reset_task.go` — 定时任务 diff --git a/main.go b/main.go index 3361b8ce..016a3fe4 100644 --- a/main.go +++ b/main.go @@ -118,6 +118,8 @@ func main() { // Subscription quota reset task (daily/weekly/monthly/custom) service.StartSubscriptionQuotaResetTask() + // Channel quota reset task (daily/weekly/monthly/custom) + service.StartChannelQuotaResetTask() // Wire task polling adaptor factory (breaks service -> relay import cycle) service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor { diff --git a/middleware/distributor.go b/middleware/distributor.go index 2263fae3..5635b801 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -157,6 +157,12 @@ func Distribute() func(c *gin.Context) { } common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now()) SetupContextForSelectedChannel(c, channel, modelRequest.Model) + if channel != nil { + if err := model.CheckChannelQuota(channel.Id); err != nil { + abortWithOpenAiMessage(c, http.StatusTooManyRequests, err.Error()) + return + } + } c.Next() if channel != nil && c.Writer != nil && c.Writer.Status() < http.StatusBadRequest { service.RecordChannelAffinity(c, channel.Id) diff --git a/model/channel.go b/model/channel.go index 370a99bf..9c03d8b3 100644 --- a/model/channel.go +++ b/model/channel.go @@ -20,25 +20,30 @@ import ( ) type Channel struct { - Id int `json:"id"` - Type int `json:"type" gorm:"default:0"` - Key string `json:"key" gorm:"not null"` - OpenAIOrganization *string `json:"openai_organization"` - TestModel *string `json:"test_model"` - Status int `json:"status" gorm:"default:1"` - Name string `json:"name" gorm:"index"` - Weight *uint `json:"weight" gorm:"default:0"` - CreatedTime int64 `json:"created_time" gorm:"bigint"` - TestTime int64 `json:"test_time" gorm:"bigint"` - ResponseTime int `json:"response_time"` // in milliseconds - BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"` - Other string `json:"other"` - Balance float64 `json:"balance"` // in USD - BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"` - Models string `json:"models"` - Group string `json:"group" gorm:"type:varchar(64);default:'default'"` - UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"` - ModelMapping *string `json:"model_mapping" gorm:"type:text"` + Id int `json:"id"` + Type int `json:"type" gorm:"default:0"` + Key string `json:"key" gorm:"not null"` + OpenAIOrganization *string `json:"openai_organization"` + TestModel *string `json:"test_model"` + Status int `json:"status" gorm:"default:1"` + Name string `json:"name" gorm:"index"` + Weight *uint `json:"weight" gorm:"default:0"` + CreatedTime int64 `json:"created_time" gorm:"bigint"` + TestTime int64 `json:"test_time" gorm:"bigint"` + ResponseTime int `json:"response_time"` // in milliseconds + BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"` + Other string `json:"other"` + Balance float64 `json:"balance"` // in USD + BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"` + Models string `json:"models"` + Group string `json:"group" gorm:"type:varchar(64);default:'default'"` + UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"` + QuotaLimit int64 `json:"quota_limit" gorm:"bigint;default:0"` + QuotaResetPeriod string `json:"quota_reset_period" gorm:"type:varchar(16);default:'never'"` + QuotaResetCustomSeconds int `json:"quota_reset_custom_seconds" gorm:"type:int;default:0"` + LastResetTime int64 `json:"last_reset_time" gorm:"bigint;default:0"` + NextResetTime int64 `json:"next_reset_time" gorm:"bigint;default:0"` + ModelMapping *string `json:"model_mapping" gorm:"type:text"` //MaxInputTokens *int `json:"max_input_tokens" gorm:"default:0"` StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"` Priority *int64 `json:"priority" gorm:"bigint;default:0"` diff --git a/model/channel_quota.go b/model/channel_quota.go new file mode 100644 index 00000000..b4cfa7c9 --- /dev/null +++ b/model/channel_quota.go @@ -0,0 +1,200 @@ +package model + +import ( + "errors" + "fmt" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" +) + +const ( + ChannelResetNever = "never" + ChannelResetDaily = "daily" + ChannelResetWeekly = "weekly" + ChannelResetMonthly = "monthly" + ChannelResetCustom = "custom" +) + +func NormalizeChannelResetPeriod(period string) string { + switch strings.TrimSpace(period) { + case ChannelResetDaily, ChannelResetWeekly, ChannelResetMonthly, ChannelResetCustom: + return strings.TrimSpace(period) + default: + return ChannelResetNever + } +} + +func calcChannelNextResetTime(base time.Time, period string, customSeconds int) int64 { + period = NormalizeChannelResetPeriod(period) + if period == ChannelResetNever { + return 0 + } + var next time.Time + switch period { + case ChannelResetDaily: + next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()). + AddDate(0, 0, 1) + case ChannelResetWeekly: + weekday := int(base.Weekday()) + if weekday == 0 { + weekday = 7 + } + daysUntil := 8 - weekday + next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()). + AddDate(0, 0, daysUntil) + case ChannelResetMonthly: + next = time.Date(base.Year(), base.Month(), 1, 0, 0, 0, 0, base.Location()). + AddDate(0, 1, 0) + case ChannelResetCustom: + if customSeconds <= 0 { + return 0 + } + next = base.Add(time.Duration(customSeconds) * time.Second) + default: + return 0 + } + return next.Unix() +} + +func maybeResetChannelQuotaTx(tx *gorm.DB, channel *Channel, now int64) error { + if tx == nil || channel == nil { + return errors.New("invalid reset args") + } + if channel.NextResetTime > 0 && channel.NextResetTime > now { + return nil + } + if NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == ChannelResetNever { + return nil + } + baseUnix := channel.LastResetTime + if baseUnix <= 0 { + baseUnix = channel.CreatedTime + } + base := time.Unix(baseUnix, 0) + next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) + + advanced := false + for next > 0 && next <= now { + advanced = true + base = time.Unix(next, 0) + next = calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) + } + + if !advanced { + if channel.NextResetTime == 0 && next > 0 { + return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ + "next_reset_time": next, + "last_reset_time": base.Unix(), + }).Error + } + return nil + } + + return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ + "used_quota": 0, + "last_reset_time": base.Unix(), + "next_reset_time": next, + }).Error +} + +func ResetDueChannels(limit int) (int, error) { + if limit <= 0 { + limit = 200 + } + now := GetDBTimestamp() + var channels []Channel + if err := DB.Where("next_reset_time > 0 AND next_reset_time <= ? AND status = ? AND quota_limit > 0", + now, common.ChannelStatusEnabled). + Order("next_reset_time asc"). + Limit(limit). + Find(&channels).Error; err != nil { + return 0, err + } + if len(channels) == 0 { + return 0, nil + } + resetCount := 0 + for _, ch := range channels { + chCopy := ch + err := DB.Transaction(func(tx *gorm.DB) error { + var locked Channel + if err := tx.Set("gorm:query_option", "FOR UPDATE"). + Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", chCopy.Id, now). + First(&locked).Error; err != nil { + return nil + } + if err := maybeResetChannelQuotaTx(tx, &locked, now); err != nil { + return err + } + resetCount++ + return nil + }) + if err != nil { + return resetCount, err + } + } + return resetCount, nil +} + +func CheckChannelQuota(channelId int) error { + var channel Channel + if err := DB.Select("id, used_quota, quota_limit, quota_reset_period, next_reset_time"). + Where("id = ?", channelId).First(&channel).Error; err != nil { + return err + } + if channel.QuotaLimit <= 0 { + return nil + } + if channel.NextResetTime > 0 { + now := GetDBTimestamp() + if channel.NextResetTime <= now { + _ = DB.Transaction(func(tx *gorm.DB) error { + var locked Channel + if err := tx.Set("gorm:query_option", "FOR UPDATE"). + Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", channelId, now). + First(&locked).Error; err != nil { + return nil + } + return maybeResetChannelQuotaTx(tx, &locked, now) + }) + if err := DB.Select("used_quota, quota_limit"). + Where("id = ?", channelId).First(&channel).Error; err != nil { + return err + } + } + } + if channel.UsedQuota >= channel.QuotaLimit { + return fmt.Errorf("channel quota exceeded (used=%d, limit=%d)", channel.UsedQuota, channel.QuotaLimit) + } + return nil +} + +func InitChannelNextResetTime(channel *Channel) { + if channel.QuotaLimit <= 0 { + channel.NextResetTime = 0 + channel.LastResetTime = 0 + return + } + period := NormalizeChannelResetPeriod(channel.QuotaResetPeriod) + if period == ChannelResetNever { + channel.NextResetTime = 0 + channel.LastResetTime = 0 + return + } + baseUnix := channel.LastResetTime + if baseUnix <= 0 { + baseUnix = channel.CreatedTime + } + if baseUnix <= 0 { + baseUnix = GetDBTimestamp() + } + base := time.Unix(baseUnix, 0) + next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) + channel.NextResetTime = next + if channel.LastResetTime <= 0 { + channel.LastResetTime = baseUnix + } +} diff --git a/service/channel_quota_reset_task.go b/service/channel_quota_reset_task.go new file mode 100644 index 00000000..d5ab2761 --- /dev/null +++ b/service/channel_quota_reset_task.go @@ -0,0 +1,70 @@ +package service + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + + "github.com/bytedance/gopkg/util/gopool" +) + +const ( + channelQuotaResetTickInterval = 1 * time.Minute + channelQuotaResetBatchSize = 300 +) + +var ( + channelQuotaResetOnce sync.Once + channelQuotaResetRunning atomic.Bool +) + +func StartChannelQuotaResetTask() { + channelQuotaResetOnce.Do(func() { + if !common.IsMasterNode { + return + } + gopool.Go(func() { + logger.LogInfo(context.Background(), fmt.Sprintf("channel quota reset task started: tick=%s", channelQuotaResetTickInterval)) + ticker := time.NewTicker(channelQuotaResetTickInterval) + defer ticker.Stop() + + runChannelQuotaResetOnce() + for range ticker.C { + runChannelQuotaResetOnce() + } + }) + }) +} + +func runChannelQuotaResetOnce() { + if !channelQuotaResetRunning.CompareAndSwap(false, true) { + return + } + defer channelQuotaResetRunning.Store(false) + + ctx := context.Background() + totalReset := 0 + for { + n, err := model.ResetDueChannels(channelQuotaResetBatchSize) + if err != nil { + logger.LogWarn(ctx, fmt.Sprintf("channel quota reset task failed: %v", err)) + return + } + if n == 0 { + break + } + totalReset += n + if n < channelQuotaResetBatchSize { + break + } + } + if common.DebugEnabled && totalReset > 0 { + logger.LogDebug(ctx, "channel quota maintenance: reset_count=%d", totalReset) + } +} diff --git a/service/text_quota.go b/service/text_quota.go index 3f344dc3..c2e9b0c1 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -368,6 +368,9 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us } else { model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota) model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota) + if err := model.CheckChannelQuota(relayInfo.ChannelId); err != nil { + logger.LogWarn(ctx, fmt.Sprintf("channel quota exceeded: channel_id=%d, %v", relayInfo.ChannelId, err)) + } } if err := SettleBilling(ctx, relayInfo, summary.Quota); err != nil {