Version: 0.6.7.dev.260317

 feat(agent): 新增 Token 配额门禁中间件(Redis 快照 + 封禁键 + 7 天懒重置)

- 🚪 在 `POST /api/v1/agent/chat` 挂载 `TokenQuotaGuard`,在请求进入业务逻辑前完成额度校验
-  新增 Redis 配额快照与封禁键机制:超额用户命中封禁键后可快速拦截,降低重复查库带来的开销
- 🗃️ 新增用户配额 DAO 能力:按需读取 `token_limit`、`token_usage`、`last_reset_at`,并支持基于“到期条件更新”的懒重置
- 🔄 实现 7 天懒重置策略:用户访问时若检测到配额周期已到期,则重置 `token_usage` 并清理封禁状态
- 🚫 新增超额响应码 `40051`,用于标识 `token usage exceeds limit`
This commit is contained in:
LoveLosita
2026-03-17 19:46:08 +08:00
parent 96be3e2a02
commit bc56d471a8
6 changed files with 324 additions and 3 deletions

View File

@@ -1,6 +1,7 @@
package dao
import (
"context"
"errors"
"time"
@@ -85,3 +86,47 @@ func (r *UserDAO) GetUserByID(id int) (*model.User, error) {
}
return &user, nil
}
// GetUserTokenQuotaByID 查询用户 token 配额快照(仅查询配额相关字段)。
//
// 职责边界:
// 1. 只返回 token_limit / token_usage / last_reset_at 等“额度判断必需字段”;
// 2. 不负责做超额判断与重置判断(由中间件统一决策);
// 3. 不返回密码等敏感字段,避免把无关信息带入鉴权链路。
func (r *UserDAO) GetUserTokenQuotaByID(ctx context.Context, id int) (*model.User, error) {
var user model.User
err := r.db.WithContext(ctx).
Select("id", "token_limit", "token_usage", "last_reset_at").
Where("id = ?", id).
First(&user).Error
if err != nil {
return nil, err
}
return &user, nil
}
// ResetUserTokenUsageIfDue 在“已到重置窗口”时执行懒重置。
//
// 输入输出语义:
// 1. dueBefore判定“到期可重置”的截止时间通常是 now-7d
// 2. resetAt本次重置写入的时间戳
// 3. 返回值 bool
// - true 表示本次调用实际执行了重置;
// - false 表示条件未命中(尚未到期或记录不存在)。
//
// 并发与幂等说明:
// 1. 使用条件更新WHERE last_reset_at <= dueBefore保证并发下最多一次成功重置
// 2. 重复调用是安全的,未命中条件时不会破坏现有统计。
func (r *UserDAO) ResetUserTokenUsageIfDue(ctx context.Context, id int, dueBefore time.Time, resetAt time.Time) (bool, error) {
result := r.db.WithContext(ctx).
Model(&model.User{}).
Where("id = ? AND (last_reset_at IS NULL OR last_reset_at <= ?)", id, dueBefore).
Updates(map[string]interface{}{
"token_usage": 0,
"last_reset_at": resetAt,
})
if result.Error != nil {
return false, result.Error
}
return result.RowsAffected > 0, nil
}