package middleware import ( "context" "fmt" "net/http" "strconv" "sync/atomic" "testing" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/i18n" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestModelRedisRateLimitUsesUTCRegardlessOfLocalTimezone(t *testing.T) { redisServer, redisClient := useRateLimitMiniRedis(t) previousLocation := time.Local time.Local = time.FixedZone("test-utc-plus-eight", 8*60*60) t.Cleanup(func() { time.Local = previousLocation }) ctx := context.Background() recordKey := "rateLimit:model-utc-record" recordRedisRequest(ctx, redisClient, recordKey, 2) recorded, err := redisClient.LIndex(ctx, recordKey, 0).Result() require.NoError(t, err) recordedAt, err := time.Parse(modelRateLimitTimeFormat, recorded) require.NoError(t, err) assert.WithinDuration(t, time.Now().UTC(), recordedAt, 2*time.Second) checkKey := "rateLimit:model-utc-check" withinWindow := time.Now().UTC().Add(-30 * time.Second).Format(modelRateLimitTimeFormat) _, err = redisServer.Push(checkKey, withinWindow, withinWindow) require.NoError(t, err) allowed, err := checkRedisRateLimit(ctx, redisClient, checkKey, 2, 60) require.NoError(t, err) assert.False(t, allowed, "an existing UTC timestamp inside the window must remain limited on a non-UTC host") } var modelRateLimitTestUsers atomic.Int64 func TestModelRateLimitStreamFailuresDoNotConsumeSuccessLimit(t *testing.T) { require.NoError(t, i18n.Init()) for _, backend := range []string{"memory", "redis"} { for _, totalLimit := range []int{0, 2} { t.Run(fmt.Sprintf("%s/total=%d", backend, totalLimit), func(t *testing.T) { subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1))) handler := memoryRateLimitHandler(subject, 60, totalLimit, 1) if backend == "redis" { useRateLimitMiniRedis(t) handler = redisRateLimitHandler(subject, 60, totalLimit, 1) } router := gin.New() router.GET("/:outcome", handler, func(c *gin.Context) { status := relaycommon.NewStreamStatus() if c.Param("outcome") == "failed" { status.MarkFailed("server_error", "", 0) } else { status.MarkCompleted() } common.SetContextKey(c, constant.ContextKeyResponseStreamStatus, status) c.Status(http.StatusOK) }) assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code) if totalLimit > 0 { assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code) } else { assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code) } assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code) }) } } } func TestModelMemoryRateLimitReservesConcurrentSuccessAdmission(t *testing.T) { require.NoError(t, i18n.Init()) subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1))) entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan struct{}) router := gin.New() router.GET("/:outcome", memoryRateLimitHandler(subject, 60, 0, 1), func(c *gin.Context) { if c.Param("outcome") == "failed" { close(entered) <-release status := relaycommon.NewStreamStatus() status.MarkFailed("server_error", "", 0) common.SetContextKey(c, constant.ContextKeyResponseStreamStatus, status) } c.Status(http.StatusOK) }) go func() { defer close(finished) assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code) }() <-entered assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code) close(release) <-finished assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code) assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code) } // The memory limiter answers like the Redis limiter: an OpenAI error body // naming the limit that was reached, not an empty 429. func TestModelMemoryRateLimitRejectionsReturnOpenAIError(t *testing.T) { require.NoError(t, i18n.Init()) for _, tc := range []struct { name string totalLimit, successLimit int messageKey string }{ {name: "total", totalLimit: 1, messageKey: i18n.MsgRateLimitTotalReached}, {name: "success", successLimit: 1, messageKey: i18n.MsgRateLimitReached}, } { t.Run(tc.name, func(t *testing.T) { subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1))) router := gin.New() router.GET("/", func(c *gin.Context) { c.Set(common.RequestIdKey, "req-limited") }, memoryRateLimitHandler(subject, 60, tc.totalLimit, tc.successLimit), func(c *gin.Context) { c.Status(http.StatusOK) }) require.Equal(t, http.StatusOK, performRateLimitRequest(router, "/", "127.0.0.1:1000").Code) resp := performRateLimitRequest(router, "/", "127.0.0.1:1000") assert.Equal(t, http.StatusTooManyRequests, resp.Code) message := i18n.Translate("", tc.messageKey, map[string]any{"Minutes": setting.ModelRequestRateLimitDurationMinutes, "Max": 1}) want, err := common.Marshal(gin.H{"error": gin.H{ "message": common.MessageWithRequestId(message, "req-limited"), "type": "new_api_error", "code": "", }}) require.NoError(t, err) assert.JSONEq(t, string(want), resp.Body.String()) }) } } // A group with its own limit counts each user's requests on its own; groups // without one keep sharing the user's global limit. func TestModelRateLimitCountsConfiguredGroupsSeparately(t *testing.T) { require.NoError(t, i18n.Init()) previousEnabled, previousMinutes := setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes previousTotal, previousSuccess := setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount previousGroups, previousRedisEnabled := setting.ModelRequestRateLimitGroup, common.RedisEnabled t.Cleanup(func() { setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes = previousEnabled, previousMinutes setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount = previousTotal, previousSuccess setting.ModelRequestRateLimitGroup, common.RedisEnabled = previousGroups, previousRedisEnabled }) setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes = true, 1 common.RedisEnabled = false // Groups a and b have their own limits; c and d have none and share the global one. steps := []struct { group string want int }{ {"a", http.StatusOK}, {"a", http.StatusTooManyRequests}, {"b", http.StatusOK}, {"b", http.StatusTooManyRequests}, {"c", http.StatusOK}, {"d", http.StatusTooManyRequests}, } // The Redis token bucket binds to the first Redis client it sees and another // test already drives it, so the total limit runs on the memory backend only. for _, tc := range []struct { backend, limit string counts [2]int // [total, success] for groups a and b and for the global limit }{ {"memory", "success", [2]int{0, 1}}, {"redis", "success", [2]int{0, 1}}, {"memory", "total", [2]int{1, 1000}}, } { t.Run(tc.backend+"/"+tc.limit, func(t *testing.T) { if tc.backend == "redis" { useRateLimitMiniRedis(t) } setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount = tc.counts[0], tc.counts[1] setting.ModelRequestRateLimitGroup = map[string][2]int{"a": tc.counts, "b": tc.counts} userID := 7200000 + int(modelRateLimitTestUsers.Add(1)) router := gin.New() router.GET("/:group", func(c *gin.Context) { c.Set("id", userID) common.SetContextKey(c, constant.ContextKeyTokenGroup, c.Param("group")) }, ModelRequestRateLimit(), func(c *gin.Context) { c.Status(http.StatusOK) }) for _, step := range steps { assert.Equal(t, step.want, performRateLimitRequest(router, "/"+step.group, "127.0.0.1:1000").Code, "group %s", step.group) } }) } }