Просмотр исходного кода

Merge branch 'feat/relay-capture'

Relay 请求抓包日志 + 用户模型 RPM 限流
master
fengsilin 1 день назад
Родитель
Сommit
9ee5f535c8
22 измененных файлов: 1915 добавлений и 9 удалений
  1. +24
    -4
      common/limiter/limiter.go
  2. +18
    -0
      common/limiter/lua/rate_refund.lua
  3. +12
    -0
      common/rate-limit.go
  4. +2
    -0
      controller/user.go
  5. +74
    -0
      controller/user_rate_limit.go
  6. +152
    -0
      controller/user_rate_limit_test.go
  7. +5
    -0
      main.go
  8. +177
    -0
      middleware/relay_capture.go
  9. +402
    -0
      middleware/relay_capture_integration_test.go
  10. +194
    -0
      middleware/relay_capture_test.go
  11. +199
    -0
      middleware/user-model-rate-limit.go
  12. +216
    -0
      middleware/user_model_rate_limit_test.go
  13. +1
    -0
      model/main.go
  14. +8
    -5
      model/user.go
  15. +43
    -0
      model/user_capture_test.go
  16. +54
    -0
      model/user_model_rate_limit.go
  17. +110
    -0
      model/user_model_rate_limit_test.go
  18. +4
    -0
      router/api-router.go
  19. +4
    -0
      router/relay-router.go
  20. +18
    -0
      web/src/components/table/users/modals/EditUserModal.jsx
  21. +192
    -0
      web/src/components/table/users/modals/UserModelRateLimitSection.jsx
  22. +6
    -0
      web/src/helpers/api.js

+ 24
- 4
common/limiter/limiter.go Просмотреть файл

@@ -13,9 +13,13 @@ import (
//go:embed lua/rate_limit.lua
var rateLimitScript string

//go:embed lua/rate_refund.lua
var rateRefundScript string

type RedisLimiter struct {
client *redis.Client
limitScriptSHA string
client *redis.Client
limitScriptSHA string
refundScriptSHA string
}

var (
@@ -30,9 +34,14 @@ func New(ctx context.Context, r *redis.Client) *RedisLimiter {
if err != nil {
common.SysLog(fmt.Sprintf("Failed to load rate limit script: %v", err))
}
refundSHA, err := r.ScriptLoad(ctx, rateRefundScript).Result()
if err != nil {
common.SysLog(fmt.Sprintf("Failed to load rate refund script: %v", err))
}
instance = &RedisLimiter{
client: r,
limitScriptSHA: limitSHA,
client: r,
limitScriptSHA: limitSHA,
refundScriptSHA: refundSHA,
}
})

@@ -68,6 +77,17 @@ func (rl *RedisLimiter) Allow(ctx context.Context, key string, opts ...Option) (
return result == 1, nil
}

func (rl *RedisLimiter) Refund(ctx context.Context, key string, requested, capacity int64) error {
_, err := rl.client.EvalSha(
ctx,
rl.refundScriptSHA,
[]string{key},
requested,
capacity,
).Int()
return err
}

// Config 配置选项模式
type Config struct {
Capacity int64


+ 18
- 0
common/limiter/lua/rate_refund.lua Просмотреть файл

@@ -0,0 +1,18 @@
-- 令牌桶退款脚本
-- KEYS[1]: 限流器唯一标识
-- ARGV[1]: 退还的令牌数
-- ARGV[2]: 桶容量(上限)
local key = KEYS[1]
local refund = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])

local bucket = redis.call('HMGET', key, 'tokens', 'last_time')
local tokens = tonumber(bucket[1])

if not tokens then
return 0
end

tokens = math.min(capacity, tokens + refund)
redis.call('HSET', key, 'tokens', tokens)
return 1

+ 12
- 0
common/rate-limit.go Просмотреть файл

@@ -68,3 +68,15 @@ func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration in
}
return true
}

// Refund 移除 key 最近一次请求记录(用于请求失败时退还配额)
func (l *InMemoryRateLimiter) Refund(key string) bool {
l.mutex.Lock()
defer l.mutex.Unlock()
queue, ok := l.store[key]
if !ok || len(*queue) == 0 {
return false
}
*queue = (*queue)[:len(*queue)-1]
return true
}

+ 2
- 0
controller/user.go Просмотреть файл

@@ -14,6 +14,7 @@ import (
"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"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/service/region_sync"
@@ -578,6 +579,7 @@ func UpdateUser(c *gin.Context) {
common.ApiError(c, err)
return
}
middleware.SetCaptureEnabled(int64(updatedUser.Id), updatedUser.CaptureRelay)
if originUser.Quota != updatedUser.Quota {
model.RecordLog(originUser.Id, model.LogTypeManage, fmt.Sprintf("管理员将用户额度从 %s修改为 %s", logger.LogQuota(originUser.Quota), logger.LogQuota(updatedUser.Quota)))
}


+ 74
- 0
controller/user_rate_limit.go Просмотреть файл

@@ -0,0 +1,74 @@
package controller

import (
"strconv"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
)

// GetUserRateLimits GET /api/user/:id/rate-limits
// 查询指定用户的所有模型 RPM 限制配置(管理员权限)
func GetUserRateLimits(c *gin.Context) {
userId, err := strconv.Atoi(c.Param("id"))
if err != nil {
common.ApiErrorMsg(c, "invalid user id")
return
}

list, err := model.GetUserModelRateLimits(userId)
if err != nil {
common.ApiError(c, err)
return
}

// 保证返回空列表而非 null
if list == nil {
list = []model.UserModelRateLimit{}
}

c.JSON(200, gin.H{
"success": true,
"message": "",
"data": list,
})
}

// SetUserRateLimits PUT /api/user/:id/rate-limits
// 覆盖式写入指定用户的所有模型 RPM 限制配置(管理员权限)
func SetUserRateLimits(c *gin.Context) {
userId, err := strconv.Atoi(c.Param("id"))
if err != nil {
common.ApiErrorMsg(c, "invalid user id")
return
}

var items []model.UserModelRateLimit
if err := c.ShouldBindJSON(&items); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}

for _, item := range items {
if item.Rpm < 0 {
common.ApiErrorMsg(c, "rpm must be >= 0")
return
}
}

if err := model.SetUserModelRateLimits(userId, items); err != nil {
common.ApiError(c, err)
return
}

// 删除 Redis 缓存
if common.RedisEnabled {
_ = common.RedisDel("user_model_rate_limit:" + strconv.Itoa(userId))
}

c.JSON(200, gin.H{
"success": true,
"message": "",
})
}

+ 152
- 0
controller/user_rate_limit_test.go Просмотреть файл

@@ -0,0 +1,152 @@
package controller

import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/glebarez/sqlite"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)

func setupUserRateLimitControllerDB(t *testing.T) {
t.Helper()
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, _ := db.DB()
sqlDB.SetMaxOpenConns(1)

origDB := model.DB
model.DB = db
common.RedisEnabled = false

require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{}))

t.Cleanup(func() {
model.DB = origDB
sqlDB.Close()
})
}

func setupUserRateLimitRouter() *gin.Engine {
gin.SetMode(gin.TestMode)
r := gin.New()
g := r.Group("/api/user")
{
g.GET("/:id/rate-limits", GetUserRateLimits)
g.PUT("/:id/rate-limits", SetUserRateLimits)
}
return r
}

func TestGetUserRateLimits_Empty(t *testing.T) {
setupUserRateLimitControllerDB(t)
router := setupUserRateLimitRouter()

w := httptest.NewRecorder()
req, _ := http.NewRequest(http.MethodGet, "/api/user/1/rate-limits", nil)
router.ServeHTTP(w, req)

assert.Equal(t, http.StatusOK, w.Code)
var resp map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
assert.True(t, resp["success"].(bool))
// data 应为空列表(非 nil)
data, ok := resp["data"]
require.True(t, ok, "response should have data key")
assert.Empty(t, data)
}

func TestSetUserRateLimits_OK(t *testing.T) {
setupUserRateLimitControllerDB(t)
router := setupUserRateLimitRouter()

items := []model.UserModelRateLimit{
{Model: "gpt-4", Rpm: 60},
}
body, _ := json.Marshal(items)

w := httptest.NewRecorder()
req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(w, req)

assert.Equal(t, http.StatusOK, w.Code)
var resp map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
assert.True(t, resp["success"].(bool))

// 验证数据确实写入
list, err := model.GetUserModelRateLimits(1)
require.NoError(t, err)
require.Len(t, list, 1)
assert.Equal(t, "gpt-4", list[0].Model)
assert.Equal(t, 60, list[0].Rpm)
}

func TestSetUserRateLimits_NegativeRpm(t *testing.T) {
setupUserRateLimitControllerDB(t)
router := setupUserRateLimitRouter()

items := []model.UserModelRateLimit{
{Model: "gpt-4", Rpm: -1},
}
body, _ := json.Marshal(items)

w := httptest.NewRecorder()
req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(w, req)

assert.Equal(t, http.StatusOK, w.Code)
var resp map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
assert.False(t, resp["success"].(bool))
}

func TestGetUserRateLimits_InvalidId(t *testing.T) {
setupUserRateLimitControllerDB(t)
router := setupUserRateLimitRouter()

w := httptest.NewRecorder()
req, _ := http.NewRequest(http.MethodGet, "/api/user/notanid/rate-limits", nil)
router.ServeHTTP(w, req)

assert.Equal(t, http.StatusOK, w.Code)
var resp map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
assert.False(t, resp["success"].(bool))
}

func TestSetUserRateLimits_EmptySlice(t *testing.T) {
setupUserRateLimitControllerDB(t)
router := setupUserRateLimitRouter()

// 先写一条
_ = model.SetUserModelRateLimits(2, []model.UserModelRateLimit{
{UserId: 2, Model: "claude-3-opus", Rpm: 10},
})

// 用空 slice 清空
body, _ := json.Marshal([]model.UserModelRateLimit{})
w := httptest.NewRecorder()
req, _ := http.NewRequest(http.MethodPut, "/api/user/2/rate-limits", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(w, req)

assert.Equal(t, http.StatusOK, w.Code)
var resp map[string]interface{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
assert.True(t, resp["success"].(bool))

list, err := model.GetUserModelRateLimits(2)
require.NoError(t, err)
assert.Empty(t, list)
}

+ 5
- 0
main.go Просмотреть файл

@@ -293,6 +293,11 @@ func InitResources() error {
// 初始化模型
model.GetPricing()

// 加载抓包用户缓存
if err := middleware.LoadCaptureEnabledUsers(); err != nil {
common.SysError("failed to load capture enabled users: " + err.Error())
}

// Initialize SQL Database
err = model.InitLogDB()
if err != nil {


+ 177
- 0
middleware/relay_capture.go Просмотреть файл

@@ -0,0 +1,177 @@
package middleware

import (
"fmt"
"os"
"path/filepath"
"sync"
"time"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
)

var captureEnabledUsers sync.Map // key: int64, value: struct{}

func isCaptureEnabled(userID int64) bool {
_, ok := captureEnabledUsers.Load(userID)
return ok
}

func SetCaptureEnabled(userID int64, enabled bool) {
if enabled {
captureEnabledUsers.Store(userID, struct{}{})
} else {
captureEnabledUsers.Delete(userID)
}
}

// LoadCaptureEnabledUsers 从数据库加载 capture_relay=true 的用户到内存缓存
func LoadCaptureEnabledUsers() error {
var ids []int64
if err := model.DB.Model(&model.User{}).
Where("capture_relay = ?", true).
Pluck("id", &ids).Error; err != nil {
return err
}
for _, id := range ids {
captureEnabledUsers.Store(id, struct{}{})
}
return nil
}

// GetCaptureEnabledUsersCount 仅用于测试
func GetCaptureEnabledUsersCount() int {
count := 0
captureEnabledUsers.Range(func(key, value interface{}) bool {
count++
return true
})
return count
}

// --- captureResponseWriter ---

type captureResponseWriter struct {
gin.ResponseWriter
ch chan<- []byte
}

func (cw *captureResponseWriter) Write(b []byte) (int, error) {
n, err := cw.ResponseWriter.Write(b)
if n > 0 {
buf := make([]byte, n)
copy(buf, b[:n])
cw.ch <- buf
}
return n, err
}

func (cw *captureResponseWriter) WriteString(s string) (int, error) {
n, err := cw.ResponseWriter.WriteString(s)
if n > 0 {
buf := make([]byte, n)
copy(buf, s[:n])
cw.ch <- buf
}
return n, err
}

// --- startCaptureWriter ---

// startCaptureWriter 启动单个 writer goroutine,从 chIn 顺序写入文件
func startCaptureWriter(f *os.File, requestID string, start time.Time) (in chan<- []byte, done <-chan struct{}) {
chIn := make(chan []byte, 128)
doneCh := make(chan struct{})

go func() {
defer close(doneCh)
defer f.Close()
for chunk := range chIn {
f.Write(chunk)
}
fmt.Fprintf(f, "\n=== END duration_ms=%d ===\n", time.Since(start).Milliseconds())
}()

return chIn, doneCh
}

// --- 辅助函数 ---

func createCaptureFile(requestID string) (*os.File, error) {
dir := filepath.Join("./data", "relay-capture", time.Now().Format("2006-01-02"))
if err := os.MkdirAll(dir, 0755); err != nil {
return nil, err
}
return os.OpenFile(
filepath.Join(dir, requestID+".log"),
os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644,
)
}

func writeRequestBlock(c *gin.Context, f *os.File, userID int64) {
fmt.Fprintf(f, "=== REQUEST %s ===\n", time.Now().UTC().Format(time.RFC3339Nano))
fmt.Fprintf(f, "%s %s\n", c.Request.Method, c.Request.URL.Path)
fmt.Fprintf(f, "user_id: %d\n", userID)

for key, values := range c.Request.Header {
for _, v := range values {
fmt.Fprintf(f, "%s: %s\n", key, v)
}
}
fmt.Fprintf(f, "\n")

// 读取请求体
if c.Request.Body != nil {
if storage, err := common.GetBodyStorage(c); err == nil {
if body, err := storage.Bytes(); err == nil && len(body) > 0 {
f.Write(body)
}
}
}
fmt.Fprintf(f, "\n")
}

// RelayCaptureMiddleware 抓包中间件
// 对开启 capture_relay 的用户,将请求体和响应体写入本地文件
func RelayCaptureMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
userID := int64(c.GetInt("id"))
if !isCaptureEnabled(userID) {
c.Next()
return
}

start := time.Now()
requestID := c.GetString(common.RequestIdKey)

f, err := createCaptureFile(requestID)
if err != nil {
c.Next()
return
}

// 同步写 REQUEST 块
writeRequestBlock(c, f, userID)
fmt.Fprintf(f, "\n=== RESPONSE ===\n")

// 启动 writer goroutine
chIn, doneCh := startCaptureWriter(f, requestID, start)

c.Writer = &captureResponseWriter{
ResponseWriter: c.Writer,
ch: chIn,
}

c.Next()

// 关闭 chIn,等待 writer 完成(带超时保护)
close(chIn)
select {
case <-doneCh:
case <-time.After(30 * time.Second):
common.SysError("relay capture: timeout waiting for writer goroutine, request_id=" + requestID)
}
}
}

+ 402
- 0
middleware/relay_capture_integration_test.go Просмотреть файл

@@ -0,0 +1,402 @@
package middleware

import (
"bytes"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"sync/atomic"
"testing"
"time"

"github.com/QuantumNous/new-api/common"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func setupCaptureTestRouter() *gin.Engine {
gin.SetMode(gin.TestMode)
return gin.New()
}

// captureTestCaptureDir 集成测试共享的捕获文件目录
const captureTestDir = "./data/relay-capture"

// --- 非流式请求测试 ---

func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) {
clearCaptureCache()

SetCaptureEnabled(42, true)
defer SetCaptureEnabled(42, false)

router := setupCaptureTestRouter()
router.Use(func(c *gin.Context) {
c.Set("id", 42) // must be int, not int64 — matches production TokenAuth behavior
c.Set(common.RequestIdKey, "test-non-stream-"+time.Now().Format("150405"))
c.Next()
})
router.Use(RelayCaptureMiddleware())
router.POST("/v1/chat/completions", func(c *gin.Context) {
c.JSON(200, gin.H{
"id": "chatcmpl-test",
"object": "chat.completion",
"choices": []gin.H{{"message": gin.H{"role": "assistant", "content": "Hello!"}}},
})
})

body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}`
req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body)))
req.Header.Set("Authorization", "Bearer sk-test")
req.Header.Set("Content-Type", "application/json")

rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)

assert.Equal(t, 200, rec.Code)

today := time.Now().Format("2006-01-02")
dateDir := filepath.Join(captureTestDir, today)
files, err := os.ReadDir(dateDir)
require.NoError(t, err)
require.True(t, len(files) > 0, "should have capture log files")

logPath := filepath.Join(dateDir, files[len(files)-1].Name())
content, err := os.ReadFile(logPath)
require.NoError(t, err)
contentStr := string(content)

t.Logf("Capture file content:\n%s", contentStr)

assert.Contains(t, contentStr, "=== REQUEST")
assert.Contains(t, contentStr, "POST /v1/chat/completions")
assert.Contains(t, contentStr, "user_id: 42")
assert.Contains(t, contentStr, "Authorization: Bearer sk-test")
assert.Contains(t, contentStr, "=== RESPONSE")
assert.Contains(t, contentStr, "chatcmpl-test")
assert.Contains(t, contentStr, "=== END")
assert.Contains(t, contentStr, "duration_ms=")
}

// --- 流式 SSE 请求测试 ---

func TestRelayCapture_Streaming_EnabledUser(t *testing.T) {
clearCaptureCache()
t.Cleanup(func() { os.RemoveAll(captureTestDir) })

SetCaptureEnabled(42, true)
defer SetCaptureEnabled(42, false)

router := setupCaptureTestRouter()
router.Use(func(c *gin.Context) {
c.Set("id", 42)
c.Set(common.RequestIdKey, "test-stream-"+time.Now().Format("150405"))
c.Next()
})
router.Use(RelayCaptureMiddleware())
router.POST("/v1/chat/completions", func(c *gin.Context) {
c.Header("Content-Type", "text/event-stream")
c.Header("Cache-Control", "no-cache")
c.Header("Connection", "keep-alive")

flusher := c.Writer.(http.Flusher)

chunks := []string{
"data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n",
"data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"world!\"}}]}\n\n",
"data: [DONE]\n\n",
}

for _, chunk := range chunks {
c.Writer.Write([]byte(chunk))
flusher.Flush()
time.Sleep(5 * time.Millisecond)
}
})

body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}],"stream":true}`
req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body)))
req.Header.Set("Authorization", "Bearer sk-test")
req.Header.Set("Content-Type", "application/json")

rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)

assert.Equal(t, 200, rec.Code)

today := time.Now().Format("2006-01-02")
dateDir := filepath.Join(captureTestDir, today)
files, err := os.ReadDir(dateDir)
require.NoError(t, err)
require.True(t, len(files) > 0, "should have capture log files")

logPath := filepath.Join(dateDir, files[len(files)-1].Name())
content, err := os.ReadFile(logPath)
require.NoError(t, err)
contentStr := string(content)

t.Logf("Stream capture file content:\n%s", contentStr)

assert.Contains(t, contentStr, "=== REQUEST")
assert.Contains(t, contentStr, "POST /v1/chat/completions")
assert.Contains(t, contentStr, "user_id: 42")
assert.Contains(t, contentStr, "=== RESPONSE")
assert.Contains(t, contentStr, "Hello ")
assert.Contains(t, contentStr, "world!")
assert.Contains(t, contentStr, "[DONE]")
assert.Contains(t, contentStr, "=== END")
assert.Contains(t, contentStr, "duration_ms=")
}

// --- 未开启抓包用户测试 ---

func TestRelayCapture_DisabledUser(t *testing.T) {
clearCaptureCache()

router := setupCaptureTestRouter()
router.Use(func(c *gin.Context) {
c.Set("id", 99)
c.Set(common.RequestIdKey, "test-disabled-"+time.Now().Format("150405"))
c.Next()
})
router.Use(RelayCaptureMiddleware())
router.POST("/v1/chat/completions", func(c *gin.Context) {
c.JSON(200, gin.H{"result": "ok"})
})

req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`)))
req.Header.Set("Authorization", "Bearer sk-test")

rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)

assert.Equal(t, 200, rec.Code)

today := time.Now().Format("2006-01-02")
dateDir := filepath.Join(captureTestDir, today)
if _, err := os.Stat(dateDir); err == nil {
files, _ := os.ReadDir(dateDir)
for _, f := range files {
assert.NotContains(t, f.Name(), "test-disabled", "disabled user should not generate capture files")
}
}
}

// --- 大量数据流测试 ---

func TestRelayCapture_LargeStream(t *testing.T) {
clearCaptureCache()
t.Cleanup(func() { os.RemoveAll(captureTestDir) })

SetCaptureEnabled(42, true)
defer SetCaptureEnabled(42, false)

router := setupCaptureTestRouter()
router.Use(func(c *gin.Context) {
c.Set("id", 42)
c.Set(common.RequestIdKey, "test-large-"+time.Now().Format("150405"))
c.Next()
})
router.Use(RelayCaptureMiddleware())
router.POST("/v1/chat/completions", func(c *gin.Context) {
c.Header("Content-Type", "text/event-stream")
flusher := c.Writer.(http.Flusher)

for i := 0; i < 50; i++ {
chunk := fmt.Sprintf("data: {\"id\":\"%d\",\"choices\":[{\"delta\":{\"content\":\"chunk%d \"}}]}\n\n", i, i)
c.Writer.Write([]byte(chunk))
flusher.Flush()
}
c.Writer.Write([]byte("data: [DONE]\n\n"))
flusher.Flush()
})

req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4","stream":true}`)))
req.Header.Set("Authorization", "Bearer sk-test")

rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)

assert.Equal(t, 200, rec.Code)

today := time.Now().Format("2006-01-02")
dateDir := filepath.Join(captureTestDir, today)
files, err := os.ReadDir(dateDir)
require.NoError(t, err)
require.True(t, len(files) > 0)

logPath := filepath.Join(dateDir, files[len(files)-1].Name())
content, err := os.ReadFile(logPath)
require.NoError(t, err)
contentStr := string(content)

assert.Contains(t, contentStr, "chunk0")
assert.Contains(t, contentStr, "chunk49")
assert.Contains(t, contentStr, "[DONE]")
}

// --- GET 请求测试 ---

func TestRelayCapture_GetRequest(t *testing.T) {
clearCaptureCache()
t.Cleanup(func() { os.RemoveAll(captureTestDir) })

SetCaptureEnabled(42, true)
defer SetCaptureEnabled(42, false)

router := setupCaptureTestRouter()
router.Use(func(c *gin.Context) {
c.Set("id", 42)
c.Set(common.RequestIdKey, "test-get-"+time.Now().Format("150405"))
c.Next()
})
router.Use(RelayCaptureMiddleware())
router.GET("/v1/models", func(c *gin.Context) {
c.JSON(200, gin.H{
"data": []gin.H{{"id": "gpt-4", "object": "model"}},
})
})

req, _ := http.NewRequest("GET", "/v1/models", nil)
req.Header.Set("Authorization", "Bearer sk-test")

rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)

assert.Equal(t, 200, rec.Code)

today := time.Now().Format("2006-01-02")
dateDir := filepath.Join(captureTestDir, today)
files, err := os.ReadDir(dateDir)
require.NoError(t, err)
require.True(t, len(files) > 0)

logPath := filepath.Join(dateDir, files[len(files)-1].Name())
content, err := os.ReadFile(logPath)
require.NoError(t, err)
contentStr := string(content)

assert.Contains(t, contentStr, "=== REQUEST")
assert.Contains(t, contentStr, "GET /v1/models")
assert.Contains(t, contentStr, "=== RESPONSE")
assert.Contains(t, contentStr, "gpt-4")
assert.Contains(t, contentStr, "=== END")
}

// --- 零 ID 用户测试 ---

func TestRelayCapture_ZeroUserID(t *testing.T) {
clearCaptureCache()

router := setupCaptureTestRouter()
router.Use(func(c *gin.Context) {
c.Set(common.RequestIdKey, "test-zero-"+time.Now().Format("150405"))
c.Next()
})
router.Use(RelayCaptureMiddleware())
router.GET("/v1/models", func(c *gin.Context) {
c.JSON(200, gin.H{"data": []gin.H{}})
})

req, _ := http.NewRequest("GET", "/v1/models", nil)
rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)

assert.Equal(t, 200, rec.Code)
}

// --- 验证响应不被篡改 ---

func TestRelayCapture_ResponseUnchanged(t *testing.T) {
clearCaptureCache()
t.Cleanup(func() { os.RemoveAll(captureTestDir) })

SetCaptureEnabled(42, true)
defer SetCaptureEnabled(42, false)

router := setupCaptureTestRouter()
router.Use(func(c *gin.Context) {
c.Set("id", 42)
c.Set(common.RequestIdKey, "test-unchanged-"+time.Now().Format("150405"))
c.Next()
})
router.Use(RelayCaptureMiddleware())
router.POST("/v1/chat/completions", func(c *gin.Context) {
c.JSON(200, gin.H{
"id": "chatcmpl-123",
"choices": []gin.H{{"message": gin.H{"content": "exact response"}}},
})
})

req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4"}`)))
req.Header.Set("Content-Type", "application/json")

rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)

assert.Equal(t, 200, rec.Code)
assert.Contains(t, rec.Body.String(), "exact response")
assert.Contains(t, rec.Body.String(), "chatcmpl-123")
}

// --- 并发安全测试 ---

func TestRelayCapture_ConcurrentRequests(t *testing.T) {
clearCaptureCache()
t.Cleanup(func() { os.RemoveAll(captureTestDir) })

SetCaptureEnabled(42, true)
defer SetCaptureEnabled(42, false)

router := setupCaptureTestRouter()
var reqCounter int64 = 0
router.Use(func(c *gin.Context) {
c.Set("id", 42)
n := atomic.AddInt64(&reqCounter, 1)
c.Set(common.RequestIdKey, fmt.Sprintf("concurrent-%d-%d", n, time.Now().UnixNano()))
c.Next()
})
router.Use(RelayCaptureMiddleware())
router.POST("/v1/chat/completions", func(c *gin.Context) {
c.JSON(200, gin.H{"status": "ok"})
})

done := make(chan bool, 10)
for i := 0; i < 10; i++ {
go func(idx int) {
req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(fmt.Sprintf(`{"req":%d}`, idx))))
req.Header.Set("Content-Type", "application/json")
rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)
assert.Equal(t, 200, rec.Code)
done <- true
}(i)
}

for i := 0; i < 10; i++ {
select {
case <-done:
case <-time.After(10 * time.Second):
t.Fatal("timeout waiting for concurrent request")
}
}

time.Sleep(100 * time.Millisecond)
today := time.Now().Format("2006-01-02")
dateDir := filepath.Join(captureTestDir, today)
files, err := os.ReadDir(dateDir)
require.NoError(t, err)
assert.Equal(t, 10, len(files), "should have 10 capture log files")

for _, f := range files {
content, err := os.ReadFile(filepath.Join(dateDir, f.Name()))
require.NoError(t, err)
contentStr := string(content)
assert.Contains(t, contentStr, "=== REQUEST")
assert.Contains(t, contentStr, "=== RESPONSE")
assert.Contains(t, contentStr, "=== END")
}
}

+ 194
- 0
middleware/relay_capture_test.go Просмотреть файл

@@ -0,0 +1,194 @@
package middleware

import (
"fmt"
"io"
"net/http/httptest"
"os"
"path/filepath"
"testing"
"time"

"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func clearCaptureCache() {
captureEnabledUsers.Range(func(key, value interface{}) bool {
captureEnabledUsers.Delete(key)
return true
})
}

func TestCaptureEnabled_NotFound(t *testing.T) {
clearCaptureCache()
assert.False(t, isCaptureEnabled(1))
}

func TestSetCaptureEnabled_Enable(t *testing.T) {
clearCaptureCache()

SetCaptureEnabled(1, true)
assert.True(t, isCaptureEnabled(1))

SetCaptureEnabled(1, false)
}

func TestSetCaptureEnabled_Disable(t *testing.T) {
clearCaptureCache()

SetCaptureEnabled(1, true)
SetCaptureEnabled(1, false)
assert.False(t, isCaptureEnabled(1))
}

func TestSetCaptureEnabled_DifferentUser(t *testing.T) {
clearCaptureCache()

SetCaptureEnabled(1, true)
assert.False(t, isCaptureEnabled(2))

SetCaptureEnabled(1, false)
}

func TestGetCaptureEnabledUsersCount(t *testing.T) {
clearCaptureCache()

assert.Equal(t, 0, GetCaptureEnabledUsersCount())

SetCaptureEnabled(1, true)
SetCaptureEnabled(2, true)
assert.Equal(t, 2, GetCaptureEnabledUsersCount())

SetCaptureEnabled(1, false)
assert.Equal(t, 1, GetCaptureEnabledUsersCount())

SetCaptureEnabled(2, false)
}

// --- captureResponseWriter 测试 ---

func TestCaptureResponseWriter_Write(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
ch := make(chan []byte, 16)

cw := &captureResponseWriter{
ResponseWriter: c.Writer,
ch: ch,
}

n, err := cw.Write([]byte("hello"))
assert.NoError(t, err)
assert.Equal(t, 5, n)
assert.Equal(t, "hello", rec.Body.String())

select {
case data := <-ch:
assert.Equal(t, []byte("hello"), data)
default:
t.Fatal("expected data in channel")
}
}

func TestCaptureResponseWriter_WriteString(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
ch := make(chan []byte, 16)

cw := &captureResponseWriter{
ResponseWriter: c.Writer,
ch: ch,
}

n, err := cw.WriteString("test-string")
assert.NoError(t, err)
assert.Equal(t, len("test-string"), n)
assert.Equal(t, "test-string", rec.Body.String())

select {
case data := <-ch:
assert.Equal(t, []byte("test-string"), data)
default:
t.Fatal("expected data in channel")
}
}

// --- startCaptureWriter 测试 ---

func TestStartCaptureWriter_WritesAllChunks(t *testing.T) {
f, err := os.CreateTemp("", "test-capture-*.log")
require.NoError(t, err)
filePath := f.Name()
defer os.Remove(filePath)

start := time.Now()
chIn, doneCh := startCaptureWriter(f, "test-req", start)

// 发送 130 个 chunk(大于 channel buffer 128)
for i := 0; i < 130; i++ {
chIn <- []byte(fmt.Sprintf("data: %d\n", i))
}
close(chIn)

select {
case <-doneCh:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for doneCh")
}

content, err := os.ReadFile(filePath)
require.NoError(t, err)
contentStr := string(content)
assert.Contains(t, contentStr, "data: 0")
assert.Contains(t, contentStr, "data: 19")
assert.Contains(t, contentStr, "=== END")
assert.Contains(t, contentStr, "duration_ms=")
}

// --- 辅助函数测试 ---

func TestCreateCaptureFile(t *testing.T) {
f, err := createCaptureFile("test-req-123")
require.NoError(t, err)
defer os.RemoveAll(filepath.Dir(f.Name()))

assert.Contains(t, f.Name(), "relay-capture")
assert.Contains(t, f.Name(), "test-req-123.log")

dir := filepath.Dir(f.Name())
info, err := os.Stat(dir)
require.NoError(t, err)
assert.True(t, info.IsDir())

f.Close()
}

func TestWriteRequestBlock(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)

c.Set("id", int64(42))
c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
c.Request.Header.Set("Authorization", "Bearer sk-test")
c.Request.Header.Set("Content-Type", "application/json")

f, err := os.CreateTemp("", "test-req-block-*.log")
require.NoError(t, err)
filePath := f.Name()
defer os.Remove(filePath)

writeRequestBlock(c, f, 42)

f.Seek(0, io.SeekStart)
content, err := io.ReadAll(f)
require.NoError(t, err)
contentStr := string(content)
assert.Contains(t, contentStr, "=== REQUEST")
assert.Contains(t, contentStr, "POST /v1/chat/completions")
assert.Contains(t, contentStr, "user_id: 42")
}

+ 199
- 0
middleware/user-model-rate-limit.go Просмотреть файл

@@ -0,0 +1,199 @@
package middleware

import (
"context"
"fmt"
"net/http"
"strconv"
"sync"
"time"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/common/limiter"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting"

"github.com/gin-gonic/gin"
)

const (
UserModelRateLimitMark = "UMRL"
)

var userModelRateLimiter common.InMemoryRateLimiter

// userModelRpmCache 非 Redis 模式下的进程内缓存,避免每次请求打 DB
var userModelRpmCache sync.Map // key: "userId:model" -> rpmCacheEntry

type rpmCacheEntry struct {
rpm int
expiresAt int64
}

const rpmCacheTTLSeconds = 300 // 5 分钟,与 Redis 路径一致

// loadUserModelRpm 查询指定用户+模型的 RPM 限制。
// Redis 可用时使用 Hash 缓存批量回填;Redis 不可用时直接查 DB。
// 返回 0 表示不限制(无配置)。
func loadUserModelRpm(userId int, modelName string) int {
if common.RedisEnabled {
return loadUserModelRpmFromRedis(userId, modelName)
}
return loadUserModelRpmFromMemory(userId, modelName)
}

func loadUserModelRpmFromMemory(userId int, modelName string) int {
cacheKey := fmt.Sprintf("%d:%s", userId, modelName)
now := time.Now().Unix()

if v, ok := userModelRpmCache.Load(cacheKey); ok {
entry := v.(rpmCacheEntry)
if now < entry.expiresAt {
return entry.rpm
}
userModelRpmCache.Delete(cacheKey)
}

rpm, _ := model.GetUserModelRpm(userId, modelName)
userModelRpmCache.Store(cacheKey, rpmCacheEntry{rpm: rpm, expiresAt: now + rpmCacheTTLSeconds})
return rpm
}

func loadUserModelRpmFromRedis(userId int, modelName string) int {
ctx := context.Background()
rdb := common.RDB
hashKey := fmt.Sprintf("user_model_rate_limit:%d", userId)

// 单次 HMGet 同时查目标模型和哨兵字段
vals, err := rdb.HMGet(ctx, hashKey, modelName, "__loaded").Result()
if err == nil && len(vals) == 2 {
// vals[0] = 模型 RPM,vals[1] = 哨兵
if vals[0] != nil {
if s, ok := vals[0].(string); ok {
if rpm, parseErr := strconv.Atoi(s); parseErr == nil {
return rpm
}
}
}
// 哨兵存在说明已加载过,该模型无配置
if vals[1] != nil {
return 0
}
}

// 首次查询:从 DB 加载并批量回填(含哨兵字段,防止未配置用户每次打 DB)
limits, dbErr := model.GetUserModelRateLimits(userId)
if dbErr != nil {
return 0
}

pipe := rdb.Pipeline()
for _, l := range limits {
pipe.HSet(ctx, hashKey, l.Model, l.Rpm)
}
pipe.HSet(ctx, hashKey, "__loaded", "1")
pipe.Expire(ctx, hashKey, 5*time.Minute)
if _, err := pipe.Exec(ctx); err != nil {
common.SysError(fmt.Sprintf("user model rate limit: failed to cache rpm for user %d: %v", userId, err))
rdb.Del(ctx, hashKey)
}

for _, l := range limits {
if l.Model == modelName {
return l.Rpm
}
}
return 0
}

// UserModelRateLimit 用户-模型维度的 RPM 速率限制中间件。
// 依赖 context 中的 "id"(userId int)和 "original_model"(modelName string)。
func UserModelRateLimit() gin.HandlerFunc {
duration := time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute
userModelRateLimiter.Init(duration)

return func(c *gin.Context) {
// 1. 全局开关
if !setting.ModelRequestRateLimitEnabled {
c.Next()
return
}

// 2. 读取 userId / modelName
userId := c.GetInt("id")
modelName := c.GetString("original_model")
if userId == 0 || modelName == "" {
c.Next()
return
}

// 3. 查 RPM 配置
rpm := loadUserModelRpm(userId, modelName)
if rpm == 0 {
c.Next()
return
}

duration := time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute

if common.RedisEnabled {
userModelRedisRateLimit(c, userId, modelName, rpm, duration)
} else {
userModelMemoryRateLimit(c, userId, modelName, rpm, duration)
}
}
}

// userModelRedisRateLimit Redis 令牌桶限流
func userModelRedisRateLimit(c *gin.Context, userId int, modelName string, rpm int, duration time.Duration) {
ctx := context.Background()
rdb := common.RDB
key := fmt.Sprintf("rateLimit:%s:%d:%s", UserModelRateLimitMark, userId, modelName)
durationSeconds := int64(duration.Seconds())

tb := limiter.New(ctx, rdb)
allowed, err := tb.Allow(ctx, key,
limiter.WithCapacity(int64(rpm)*durationSeconds),
limiter.WithRate(int64(rpm)),
limiter.WithRequested(durationSeconds),
)
if err != nil {
// Redis 异常时降级到内存限流
userModelMemoryRateLimit(c, userId, modelName, rpm, duration)
return
}

if !allowed {
abortWithOpenAiMessage(c, http.StatusTooManyRequests,
fmt.Sprintf("用户模型速率限制:每 %d 分钟最多 %d 次请求 (%s)",
setting.ModelRequestRateLimitDurationMinutes, rpm, modelName))
return
}

c.Next()

// 请求失败时退还令牌
if c.Writer.Status() >= 400 {
_ = tb.Refund(ctx, key, durationSeconds, int64(rpm)*durationSeconds)
}
}

// userModelMemoryRateLimit 内存滑动窗口限流(参考 memoryRateLimitHandler)
func userModelMemoryRateLimit(c *gin.Context, userId int, modelName string, rpm int, duration time.Duration) {
key := fmt.Sprintf("%s%d:%s", UserModelRateLimitMark, userId, modelName)
durationSeconds := int64(duration.Seconds())

if !userModelRateLimiter.Request(key, rpm, durationSeconds) {
abortWithOpenAiMessage(c, http.StatusTooManyRequests,
fmt.Sprintf("用户模型速率限制:每 %d 分钟最多 %d 次请求 (%s)",
setting.ModelRequestRateLimitDurationMinutes, rpm, modelName))
return
}

c.Next()

// 请求失败时退还配额
if c.Writer.Status() >= 400 {
userModelRateLimiter.Refund(key)
}
}

+ 216
- 0
middleware/user_model_rate_limit_test.go Просмотреть файл

@@ -0,0 +1,216 @@
package middleware

import (
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)

func setupTestDB(t *testing.T) {
t.Helper()
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, _ := db.DB()
sqlDB.SetMaxOpenConns(1)

origDB := model.DB
model.DB = db
require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{}))

t.Cleanup(func() {
model.DB = origDB
sqlDB.Close()
})
}

// buildTestRouter 创建用于测试的 gin Engine,前置中间件注入 userId 和 modelName
func buildTestRouter(userId int, modelName string) *gin.Engine {
gin.SetMode(gin.TestMode)
r := gin.New()
r.POST("/test",
func(c *gin.Context) {
c.Set("id", userId)
c.Set("original_model", modelName)
c.Next()
},
UserModelRateLimit(),
func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
},
)
return r
}

func doRequest(r *gin.Engine) *httptest.ResponseRecorder {
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/test", nil)
r.ServeHTTP(w, req)
return w
}

// TestUserModelRateLimit_Disabled:全局开关关闭时即使有限流配置也应放行
func TestUserModelRateLimit_Disabled(t *testing.T) {
setupTestDB(t)
common.RedisEnabled = false
setting.ModelRequestRateLimitEnabled = false

// 插入 rpm=1 的严格限制
require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
UserId: 42, Model: "gpt-4", Rpm: 1,
}).Error)

r := buildTestRouter(42, "gpt-4")

// 连续发两次请求,都应该 200
w1 := doRequest(r)
assert.Equal(t, http.StatusOK, w1.Code)
w2 := doRequest(r)
assert.Equal(t, http.StatusOK, w2.Code)
}

// TestUserModelRateLimit_NoModel:original_model 为空时放行
func TestUserModelRateLimit_NoModel(t *testing.T) {
setupTestDB(t)
common.RedisEnabled = false
setting.ModelRequestRateLimitEnabled = true

r := buildTestRouter(42, "")

w := doRequest(r)
assert.Equal(t, http.StatusOK, w.Code)
}

// TestUserModelRateLimit_NoConfig:DB 中无该用户+模型的配置时放行
func TestUserModelRateLimit_NoConfig(t *testing.T) {
setupTestDB(t)
common.RedisEnabled = false
setting.ModelRequestRateLimitEnabled = true
setting.ModelRequestRateLimitDurationMinutes = 1

// DB 中没有任何配置
r := buildTestRouter(99, "gpt-4")

w := doRequest(r)
assert.Equal(t, http.StatusOK, w.Code)
}

// TestUserModelRateLimit_UnderLimit:rpm=2,发 1 次请求,放行
func TestUserModelRateLimit_UnderLimit(t *testing.T) {
setupTestDB(t)
common.RedisEnabled = false
setting.ModelRequestRateLimitEnabled = true
setting.ModelRequestRateLimitDurationMinutes = 1

require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
UserId: 10, Model: "claude-3", Rpm: 2,
}).Error)

r := buildTestRouter(10, "claude-3")

w := doRequest(r)
assert.Equal(t, http.StatusOK, w.Code)
}

// TestUserModelRateLimit_OverLimit:rpm=2,连续发 3 次请求,第 3 次被 429
func TestUserModelRateLimit_OverLimit(t *testing.T) {
setupTestDB(t)
common.RedisEnabled = false
setting.ModelRequestRateLimitEnabled = true
setting.ModelRequestRateLimitDurationMinutes = 1

require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
UserId: 20, Model: "gpt-4o", Rpm: 2,
}).Error)

r := buildTestRouter(20, "gpt-4o")

w1 := doRequest(r)
assert.Equal(t, http.StatusOK, w1.Code, "第 1 次请求应放行")

w2 := doRequest(r)
assert.Equal(t, http.StatusOK, w2.Code, "第 2 次请求应放行")

w3 := doRequest(r)
assert.Equal(t, http.StatusTooManyRequests, w3.Code, "第 3 次请求应被限流")
}

// TestUserModelRateLimit_ConcurrentMemory:内存路径并发安全验证。
// RPM=2,并发 20 个请求,预期最多 2 个 200,其余 429。
func TestUserModelRateLimit_ConcurrentMemory(t *testing.T) {
setupTestDB(t)
common.RedisEnabled = false
setting.ModelRequestRateLimitEnabled = true
setting.ModelRequestRateLimitDurationMinutes = 1

require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
UserId: 30, Model: "deepseek-r1", Rpm: 2,
}).Error)

r := buildTestRouter(30, "deepseek-r1")

var successCount int64
var wg sync.WaitGroup
totalRequests := 20

for i := 0; i < totalRequests; i++ {
wg.Add(1)
go func() {
defer wg.Done()
w := doRequest(r)
if w.Code == http.StatusOK {
atomic.AddInt64(&successCount, 1)
}
}()
}
wg.Wait()

sc := int(atomic.LoadInt64(&successCount))
assert.LessOrEqual(t, sc, 2, "并发请求中最多 %d 个应成功,实际 %d 个", 2, sc)
assert.GreaterOrEqual(t, sc, 1, "至少应有 1 个请求成功")
}

// TestUserModelRateLimit_UserIsolation:不同用户之间的限流互不影响
func TestUserModelRateLimit_UserIsolation(t *testing.T) {
setupTestDB(t)
common.RedisEnabled = false
setting.ModelRequestRateLimitEnabled = true
setting.ModelRequestRateLimitDurationMinutes = 1

// 两个用户都设置 RPM=2
require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
UserId: 100, Model: "gpt-4", Rpm: 2,
}).Error)
require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
UserId: 101, Model: "gpt-4", Rpm: 2,
}).Error)

r1 := buildTestRouter(100, "gpt-4")
r2 := buildTestRouter(101, "gpt-4")

// 用户 100 用完配额
w1 := doRequest(r1)
assert.Equal(t, http.StatusOK, w1.Code)
w2 := doRequest(r1)
assert.Equal(t, http.StatusOK, w2.Code)
w3 := doRequest(r1)
assert.Equal(t, http.StatusTooManyRequests, w3.Code)

// 用户 101 应该不受影响
w4 := doRequest(r2)
assert.Equal(t, http.StatusOK, w4.Code)
w5 := doRequest(r2)
assert.Equal(t, http.StatusOK, w5.Code)
w6 := doRequest(r2)
assert.Equal(t, http.StatusTooManyRequests, w6.Code)
}

+ 1
- 0
model/main.go Просмотреть файл

@@ -290,6 +290,7 @@ func migrateDB() error {
&PendingSyncRecord{},
&QuotaSyncLog{},
&EmailQuotaRule{},
&UserModelRateLimit{},
)
if err != nil {
return err


+ 8
- 5
model/user.go Просмотреть файл

@@ -65,6 +65,8 @@ type User struct {
RemoteUserId int `json:"remote_user_id" gorm:"type:int;default:0;column:remote_user_id"`
SyncedQuota int `json:"synced_quota" gorm:"type:int;default:0;column:synced_quota"`
LastSyncAt int64 `json:"last_sync_at" gorm:"type:bigint;default:0;column:last_sync_at"`
// Relay 抓包开关
CaptureRelay bool `json:"capture_relay" gorm:"default:false"`
}

// IsSyncedUser 判断是否为国内同步用户
@@ -580,11 +582,12 @@ func (user *User) Edit(updatePassword bool) error {

newUser := *user
updates := map[string]interface{}{
"username": newUser.Username,
"display_name": newUser.DisplayName,
"group": newUser.Group,
"quota": newUser.Quota,
"remark": newUser.Remark,
"username": newUser.Username,
"display_name": newUser.DisplayName,
"group": newUser.Group,
"quota": newUser.Quota,
"remark": newUser.Remark,
"capture_relay": newUser.CaptureRelay,
}
if updatePassword {
updates["password"] = newUser.Password


+ 43
- 0
model/user_capture_test.go Просмотреть файл

@@ -0,0 +1,43 @@
package model

import (
"testing"

"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)

func TestUser_CaptureRelayField_DB(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
defer func() {
sqlDB, _ := db.DB()
sqlDB.Close()
}()

err = db.AutoMigrate(&User{})
require.NoError(t, err)

user := &User{
Username: "capture-test-user",
Password: "12345678",
DisplayName: "capture test",
CaptureRelay: false,
}
err = db.Create(user).Error
require.NoError(t, err)

var retrieved User
err = db.First(&retrieved, user.Id).Error
require.NoError(t, err)
assert.False(t, retrieved.CaptureRelay)

// 测试设置为 true
err = db.Model(&retrieved).Update("capture_relay", true).Error
require.NoError(t, err)
err = db.First(&retrieved, user.Id).Error
require.NoError(t, err)
assert.True(t, retrieved.CaptureRelay)
}

+ 54
- 0
model/user_model_rate_limit.go Просмотреть файл

@@ -0,0 +1,54 @@
package model

import (
"github.com/QuantumNous/new-api/common"
"gorm.io/gorm"
)

type UserModelRateLimit struct {
Id int `json:"id" gorm:"primaryKey"`
UserId int `json:"user_id" gorm:"uniqueIndex:idx_user_model_rate_limit"`
Model string `json:"model" gorm:"size:255;uniqueIndex:idx_user_model_rate_limit"`
Rpm int `json:"rpm"`
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at"`
}

// GetUserModelRateLimits 查询指定用户的所有 RPM 配置
func GetUserModelRateLimits(userId int) ([]UserModelRateLimit, error) {
var list []UserModelRateLimit
err := DB.Where("user_id = ?", userId).Find(&list).Error
return list, err
}

// SetUserModelRateLimits 覆盖式写入:事务中先硬删再批量插入。
// 空 slice 时仅删除该用户所有配置。
func SetUserModelRateLimits(userId int, items []UserModelRateLimit) error {
return DB.Transaction(func(tx *gorm.DB) error {
// 先硬删该用户所有配置
if err := tx.Where("user_id = ?", userId).Delete(&UserModelRateLimit{}).Error; err != nil {
return err
}
if len(items) == 0 {
return nil
}
now := common.GetTimestamp()
for i := range items {
items[i].UserId = userId
items[i].Id = 0 // 让数据库自增
items[i].CreatedAt = now
items[i].UpdatedAt = now
}
return tx.Create(&items).Error
})
}

// GetUserModelRpm 查单个用户+模型的 RPM,不命中返回 (0, false)
func GetUserModelRpm(userId int, model string) (int, bool) {
var item UserModelRateLimit
err := DB.Where("user_id = ? AND model = ?", userId, model).First(&item).Error
if err != nil {
return 0, false
}
return item.Rpm, true
}

+ 110
- 0
model/user_model_rate_limit_test.go Просмотреть файл

@@ -0,0 +1,110 @@
package model

import (
"testing"

"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)

func setupUserModelRateLimitDB(t *testing.T) {
t.Helper()
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, _ := db.DB()
sqlDB.SetMaxOpenConns(1)

origDB := DB
DB = db

require.NoError(t, db.AutoMigrate(&UserModelRateLimit{}))

t.Cleanup(func() {
DB = origDB
sqlDB.Close()
})
}

func TestGetUserModelRateLimits_Empty(t *testing.T) {
setupUserModelRateLimitDB(t)

list, err := GetUserModelRateLimits(1)
require.NoError(t, err)
assert.Empty(t, list)
}

func TestSetUserModelRateLimits_CreateAndOverwrite(t *testing.T) {
setupUserModelRateLimitDB(t)

// 首次写入 2 条
items := []UserModelRateLimit{
{UserId: 1, Model: "gpt-4", Rpm: 60},
{UserId: 1, Model: "gpt-3.5-turbo", Rpm: 120},
}
err := SetUserModelRateLimits(1, items)
require.NoError(t, err)

list, err := GetUserModelRateLimits(1)
require.NoError(t, err)
assert.Len(t, list, 2)

// 再覆盖写入 1 条(只保留 gpt-4,rpm 改为 30)
items2 := []UserModelRateLimit{
{UserId: 1, Model: "gpt-4", Rpm: 30},
}
err = SetUserModelRateLimits(1, items2)
require.NoError(t, err)

list2, err := GetUserModelRateLimits(1)
require.NoError(t, err)
assert.Len(t, list2, 1)
assert.Equal(t, "gpt-4", list2[0].Model)
assert.Equal(t, 30, list2[0].Rpm)
}

func TestSetUserModelRateLimits_EmptySlice(t *testing.T) {
setupUserModelRateLimitDB(t)

// 先写 1 条
items := []UserModelRateLimit{
{UserId: 2, Model: "claude-3-opus", Rpm: 10},
}
err := SetUserModelRateLimits(2, items)
require.NoError(t, err)

list, err := GetUserModelRateLimits(2)
require.NoError(t, err)
assert.Len(t, list, 1)

// 再用空 slice 清除
err = SetUserModelRateLimits(2, []UserModelRateLimit{})
require.NoError(t, err)

list2, err := GetUserModelRateLimits(2)
require.NoError(t, err)
assert.Empty(t, list2)
}

func TestGetUserModelRpm_Found(t *testing.T) {
setupUserModelRateLimitDB(t)

items := []UserModelRateLimit{
{UserId: 3, Model: "gemini-pro", Rpm: 20},
}
err := SetUserModelRateLimits(3, items)
require.NoError(t, err)

rpm, found := GetUserModelRpm(3, "gemini-pro")
assert.True(t, found)
assert.Equal(t, 20, rpm)
}

func TestGetUserModelRpm_NotFound(t *testing.T) {
setupUserModelRateLimitDB(t)

rpm, found := GetUserModelRpm(999, "nonexistent-model")
assert.False(t, found)
assert.Equal(t, 0, rpm)
}

+ 4
- 0
router/api-router.go Просмотреть файл

@@ -142,6 +142,10 @@ func SetApiRouter(router *gin.Engine) {
// Admin 2FA routes
adminRoute.GET("/2fa/stats", controller.Admin2FAStats)
adminRoute.DELETE("/:id/2fa", controller.AdminDisable2FA)

// User model RPM rate limit routes
adminRoute.GET("/:id/rate-limits", controller.GetUserRateLimits)
adminRoute.PUT("/:id/rate-limits", controller.SetUserRateLimits)
}
}



+ 4
- 0
router/relay-router.go Просмотреть файл

@@ -66,6 +66,7 @@ func SetRelayRouter(router *gin.Engine) {
relayV1Router.Use(middleware.SystemPerformanceCheck())
relayV1Router.Use(middleware.TokenAuth())
relayV1Router.Use(middleware.ModelRequestRateLimit())
relayV1Router.Use(middleware.RelayCaptureMiddleware())
{
// WebSocket 路由(统一到 Relay)
wsRouter := relayV1Router.Group("")
@@ -78,6 +79,7 @@ func SetRelayRouter(router *gin.Engine) {
//http router
httpRouter := relayV1Router.Group("")
httpRouter.Use(middleware.Distribute())
httpRouter.Use(middleware.UserModelRateLimit())

// claude related routes
httpRouter.POST("/messages", func(c *gin.Context) {
@@ -182,7 +184,9 @@ func SetRelayRouter(router *gin.Engine) {
relayGeminiRouter.Use(middleware.SystemPerformanceCheck())
relayGeminiRouter.Use(middleware.TokenAuth())
relayGeminiRouter.Use(middleware.ModelRequestRateLimit())
relayGeminiRouter.Use(middleware.RelayCaptureMiddleware())
relayGeminiRouter.Use(middleware.Distribute())
relayGeminiRouter.Use(middleware.UserModelRateLimit())
{
// Gemini API 路径格式: /v1beta/models/{model_name}:{action}
relayGeminiRouter.POST("/models/*path", func(c *gin.Context) {


+ 18
- 0
web/src/components/table/users/modals/EditUserModal.jsx Просмотреть файл

@@ -57,6 +57,7 @@ import {
} from '@douyinfe/semi-icons';
import UserBindingManagementModal from './UserBindingManagementModal';
import UserRatioSection from './UserRatioSection';
import UserModelRateLimitSection from './UserModelRateLimitSection';

const { Text, Title } = Typography;

@@ -88,6 +89,7 @@ const EditUserModal = (props) => {
quota: 0,
group: 'default',
remark: '',
capture_relay: false,
});

const fetchGroups = async () => {
@@ -267,6 +269,16 @@ const EditUserModal = (props) => {
showClear
/>
</Col>

{userId && (
<Col span={24}>
<Form.Switch
field='capture_relay'
label={t('Relay 抓包')}
helpText={t('开启后该用户所有 relay 请求的请求体和响应体将保存到服务器本地文件,仅用于调试')}
/>
</Col>
)}
</Row>
</Card>

@@ -335,6 +347,12 @@ const EditUserModal = (props) => {
<UserRatioSection userId={userId} />
</Card>
)}
{/* 模型速率限制 */}
{userId && (
<Card className='!rounded-2xl shadow-sm border-0'>
<UserModelRateLimitSection userId={userId} />
</Card>
)}
{/* 绑定信息入口 */}
{userId && (
<Card className='!rounded-2xl shadow-sm border-0'>


+ 192
- 0
web/src/components/table/users/modals/UserModelRateLimitSection.jsx Просмотреть файл

@@ -0,0 +1,192 @@
import React, { useEffect, useState, useRef } from 'react';
import { useTranslation } from 'react-i18next';
import { userRateLimitApi, showError, showSuccess } from '../../../../helpers';
import {
Button,
Table,
Modal,
Form,
Avatar,
Typography,
Popconfirm,
} from '@douyinfe/semi-ui';
import { IconDelete, IconPlus, IconClock } from '@douyinfe/semi-icons';

const { Text } = Typography;

const UserModelRateLimitSection = ({ userId }) => {
const { t } = useTranslation();
const [limits, setLimits] = useState([]);
const [loading, setLoading] = useState(false);
const [addModalVisible, setAddModalVisible] = useState(false);
const formApiRef = useRef(null);

const loadLimits = async () => {
if (!userId) return;
setLoading(true);
try {
const res = await userRateLimitApi.get(userId);
const { success, data, message } = res.data;
if (success) {
setLimits(data || []);
} else {
showError(message);
}
} catch (e) {
showError(e.message);
}
setLoading(false);
};

useEffect(() => {
loadLimits();
}, [userId]);

const saveLimits = async (newLimits) => {
try {
const res = await userRateLimitApi.set(userId, newLimits);
const { success, message } = res.data;
if (success) {
showSuccess(t('保存成功'));
return true;
} else {
showError(message);
return false;
}
} catch (e) {
showError(e.message);
return false;
}
};

const handleAdd = async (values) => {
const rpm = parseInt(values.rpm) || 0;
const newItem = { model: values.model, rpm };
const filtered = limits.filter((item) => item.model !== values.model);
const newLimits = [...filtered, newItem];
const ok = await saveLimits(newLimits);
if (ok) {
setLimits(newLimits);
setAddModalVisible(false);
formApiRef.current?.reset();
}
};

const handleDelete = async (model) => {
const newLimits = limits.filter((item) => item.model !== model);
const ok = await saveLimits(newLimits);
if (ok) {
setLimits(newLimits);
}
};

const columns = [
{
title: t('模型'),
dataIndex: 'model',
key: 'model',
},
{
title: t('RPM 限制'),
dataIndex: 'rpm',
key: 'rpm',
render: (rpm) =>
rpm === 0 ? (
<Text type='tertiary'>{t('不限制')}</Text>
) : (
<Text>{rpm}</Text>
),
},
{
title: t('操作'),
key: 'action',
render: (_, record) => (
<Popconfirm
title={t('确认删除该限制?')}
onConfirm={() => handleDelete(record.model)}
>
<Button
type='danger'
size='small'
icon={<IconDelete />}
>
{t('删除')}
</Button>
</Popconfirm>
),
},
];

return (
<>
<div className='flex items-center justify-between mb-2'>
<div className='flex items-center'>
<Avatar size='small' color='cyan' className='mr-2 shadow-md'>
<IconClock size={16} />
</Avatar>
<div>
<Text className='text-lg font-medium'>
{t('模型速率限制(RPM)')}
</Text>
<div className='text-xs text-gray-600'>
{t('为该用户配置各模型每分钟请求数上限,0 表示不限制')}
</div>
</div>
</div>
<Button
size='small'
icon={<IconPlus />}
onClick={() => setAddModalVisible(true)}
>
{t('添加限制')}
</Button>
</div>

<Table
columns={columns}
dataSource={limits}
loading={loading}
rowKey='model'
size='small'
pagination={false}
empty={t('暂无速率限制')}
/>

<Modal
title={t('添加模型速率限制')}
visible={addModalVisible}
onOk={() => formApiRef.current?.submitForm()}
onCancel={() => {
setAddModalVisible(false);
formApiRef.current?.reset();
}}
>
<Form
getFormApi={(api) => (formApiRef.current = api)}
onSubmit={handleAdd}
>
<Form.Input
field='model'
label={t('模型名称')}
placeholder={t('请输入精确的模型名,如 gpt-4o')}
rules={[{ required: true, message: t('请输入模型名称') }]}
style={{ width: '100%' }}
/>
<Form.InputNumber
field='rpm'
label={t('RPM 限制')}
placeholder='60'
initValue={60}
min={0}
step={10}
extraText={t('0 表示不限制')}
rules={[{ required: true, message: t('请输入 RPM 限制') }]}
style={{ width: '100%' }}
/>
</Form>
</Modal>
</>
);
};

export default UserModelRateLimitSection;

+ 6
- 0
web/src/helpers/api.js Просмотреть файл

@@ -418,3 +418,9 @@ export const pricingTagApi = {
update: (id, data) => API.put(`/api/pricing_tag/${id}`, data),
delete: (id) => API.delete(`/api/pricing_tag/${id}`),
};

// 用户模型限流 API
export const userRateLimitApi = {
get: (userId) => API.get(`/api/user/${userId}/rate-limits`),
set: (userId, items) => API.put(`/api/user/${userId}/rate-limits`, items),
};

Загрузка…
Отмена
Сохранить