feat: sync deployed v1.0.0-rc.38-custom.42

Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.

middleware/model_rate_limit_test.go

@@ -2,9 +2,16 @@ package middleware
 
 import (
 	"context"
+	"fmt"
+	"net/http"
+	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/gin-gonic/gin"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -32,3 +39,64 @@ func TestModelRedisRateLimitUsesUTCRegardlessOfLocalTimezone(t *testing.T) {
 	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) {
+	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) {
+				userID := 7200000 + int(modelRateLimitTestUsers.Add(1))
+				handler := memoryRateLimitHandler(60, totalLimit, 1)
+				if backend == "redis" {
+					useRateLimitMiniRedis(t)
+					handler = redisRateLimitHandler(60, totalLimit, 1)
+				}
+				router := gin.New()
+				router.GET("/:outcome", func(c *gin.Context) { c.Set("id", userID) }, 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) {
+	userID := 7200000 + int(modelRateLimitTestUsers.Add(1))
+	entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan struct{})
+	router := gin.New()
+	router.GET("/:outcome", func(c *gin.Context) { c.Set("id", userID) }, memoryRateLimitHandler(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)
+}