| @@ -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 | |||
| @@ -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 | |||
| @@ -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 | |||
| } | |||
| @@ -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))) | |||
| } | |||
| @@ -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": "", | |||
| }) | |||
| } | |||
| @@ -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) | |||
| } | |||
| @@ -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 { | |||
| @@ -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) | |||
| } | |||
| } | |||
| } | |||
| @@ -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") | |||
| } | |||
| } | |||
| @@ -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") | |||
| } | |||
| @@ -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) | |||
| } | |||
| } | |||
| @@ -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) | |||
| } | |||
| @@ -290,6 +290,7 @@ func migrateDB() error { | |||
| &PendingSyncRecord{}, | |||
| &QuotaSyncLog{}, | |||
| &EmailQuotaRule{}, | |||
| &UserModelRateLimit{}, | |||
| ) | |||
| if err != nil { | |||
| return err | |||
| @@ -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 | |||
| @@ -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) | |||
| } | |||
| @@ -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 | |||
| } | |||
| @@ -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) | |||
| } | |||
| @@ -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) | |||
| } | |||
| } | |||
| @@ -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) { | |||
| @@ -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'> | |||
| @@ -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; | |||
| @@ -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), | |||
| }; | |||