Skip to content

fix: exclude failed requests from memory rate limit success count #6528

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
feitianbubu wants to merge 1 commit into
base: main
Choose a base branch
from
+51 −9
Open

fix: exclude failed requests from memory rate limit success count #6528

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions common/rate-limit.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,18 @@ func (l *InMemoryRateLimiter) clearExpiredItems() {
}
}

// CanRequest reports whether Request would allow key right now, without
// recording anything. Parameter duration's unit is seconds.
func (l *InMemoryRateLimiter) CanRequest(key string, maxRequestNum int, duration int64) bool {
l.mutex.Lock()
defer l.mutex.Unlock()
queue, ok := l.store[key]
if !ok || len(*queue) < maxRequestNum {
return true
}
return time.Now().Unix()-(*queue)[0] >= duration
}

// Request parameter duration's unit is seconds
func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration int64) bool {
l.mutex.Lock()
Expand Down
14 changes: 5 additions & 9 deletions middleware/model-rate-limit.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,24 +140,20 @@ func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int)

// 1. 检查总请求数限制（当totalMaxCount为0时跳过）
if totalMaxCount > 0 && !inMemoryRateLimiter.Request(totalKey, totalMaxCount, duration) {
c.Status(http.StatusTooManyRequests)
c.Abort()
abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到总请求数限制：%d分钟内最多请求%d次，包括失败次数，请检查您的请求是否正确", setting.ModelRequestRateLimitDurationMinutes, totalMaxCount))
return
}

// 2. 检查成功请求数限制
// 使用一个临时key来检查限制，这样可以避免实际记录
checkKey := successKey + "_check"
if !inMemoryRateLimiter.Request(checkKey, successMaxCount, duration) {
c.Status(http.StatusTooManyRequests)
c.Abort()
// 2. 检查成功请求数限制：只检查不计数，成功后才在步骤4计入
if !inMemoryRateLimiter.CanRequest(successKey, successMaxCount, duration) {
abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到请求数限制：%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount))
return
}

// 3. 处理请求
c.Next()

// 4. 如果请求成功，记录到实际的成功请求计数中
// 4. 如果请求成功，记录到成功请求计数中
Comment on lines +147 to +156

@coderabbitai coderabbitai Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Make success-limit admission atomic with success accounting.

CanRequest releases the limiter lock before c.Next(), so concurrent requests can all pass the check while capacity is available. They are then all admitted, even though only some later Request calls can be recorded; the return value from that recording is ignored. Use reservation/commit/rollback semantics or track in-flight reservations so successful traffic cannot exceed the configured quota. The current sequential regression test will not catch this race.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@middleware/model-rate-limit.go` around lines 147 - 156, Update the
success-limit flow around CanRequest and the subsequent accounting after
c.Next() to reserve capacity atomically before admitting a request, then commit
the reservation only when the request succeeds and roll it back otherwise.
Ensure concurrent requests cannot collectively exceed successMaxCount, and
handle the recording operation’s result rather than ignoring it; extend coverage
with a concurrent regression test for this race.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "Files matching middleware/model-rate-limit.go:"
fd -a 'model-rate-limit\.go$' . || true

file="$(fd 'model-rate-limit\.go$' . | head -n 1 || true)"
if [ -z "${file:-}" ]; then
  echo "Target file not found"
  exit 1
fi

echo "File: $file"
wc -l "$file"
echo "--- outline ---"
ast-grep outline "$file" --view compact || true
echo "--- relevant lines 1-220 ---"
sed -n '1,220p' "$file" | cat -n

echo "--- related successMaxCount occurrences ---"
rg -n "successMaxCount|ModelRequestRateLimitDurationMinutes|ModelSuccessRateLimit|rate.*limit|InMemoryRateLimiter|Request\\(|CanRequest" "$file"

Repository: QuantumNous/new-api

Length of output: 9826

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "locating limiter package files:"
fd -a '.*\.go$' common/limiter || true

echo "--- common/limiter files sizes ---"
while IFS= read -r f; do wc -l "$f"; done < <(fd '.*\.go$' common/limiter)

echo "--- common/limiter outline/contents ---"
while IFS= read -r f; do
  echo "===== $f ====="
  sed -n '1,260p' "$f" | cat -n
done < <(fd '.*\.go$' common/limiter)

echo "--- inMemoryRateLimiter declarations/usages ---"
rg -n "inMemoryRateLimiter|CanRequest|func .*Request|type .*Limiter" common middleware

Repository: QuantumNous/new-api

Length of output: 321

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "--- common/limiter/limiter.go ---"
wc -l common/limiter/limiter.go
sed -n '1,260p' common/limiter/limiter.go | cat -n

echo "--- inMemoryRateLimiter declarations/usages ---"
rg -n "inMemoryRateLimiter|CanRequest|func .*Request|type .*Limiter" common middleware

Repository: QuantumNous/new-api

Length of output: 7278

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "--- common/rate-limit.go ---"
wc -l common/rate-limit.go
sed -n '1,140p' common/rate-limit.go | cat -n

echo "--- model_rate_limit_test.go relevant lines ---"
sed -n '1,140p' middleware/model_rate_limit_test.go | cat -n

Repository: QuantumNous/new-api

Length of output: 5615

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
def can_request(store, key, max_request_num, now, duration):
    queue = store.get(key)
    if queue is None or len(queue) < max_request_num:
        return True
    return (now - queue[0]) >= duration

def request(store, key, max_request_num, now, duration):
    now = now
    queue = store.get(key)
    if queue is not None:
        if len(queue) < max_request_num:
            store[key] = queue + [now]
            return True
        else:
            if now - queue[0] >= duration:
                store[key] = queue[1:] + [now]
                return True
            else:
                return False
    else:
        store[key] = [now]
        return True

def memory_handler(store, key_prefix, user_id, downstream_statuses, duration):
    store[:] = {}
    for status in downstream_statuses:
        success_key = key_prefix + str(user_id)
        allowed = can_request(store, success_key, 0, 100, duration)
        if not allowed:
            yield "abort", 429
            continue
        # process
        if status < 400:
            request(store, success_key, 0, 100, duration)
        yield "processed", status, store.get(success_key)

statuses = [200, 200, 200]
events = list(memory_handler([], "model_success_", 1, statuses, 60))
for e in events:
    print(e)
PY

Repository: QuantumNous/new-api

Length of output: 370

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
def can_request(store, key, max_request_num, now, duration):
    queue = store.get(key)
    if queue is None or len(queue) < max_request_num:
        return True
    return (now - queue[0]) >= duration

def request(store, key, max_request_num, now, duration):
    queue = store.get(key)
    if queue is not None:
        if len(queue) < max_request_num:
            store[key] = queue + [now]
            return True
        else:
            if now - queue[0] >= duration:
                store[key] = queue[1:] + [now]
                return True
            else:
                return False
    else:
        store[key] = [now]
        return True

def memory_handler(store, key_prefix, user_id, downstream_statuses, duration):
    for status in downstream_statuses:
        success_key = key_prefix + str(user_id)
        allowed = can_request(store, success_key, 0, 100, duration)
        if not allowed:
            yield "abort", 429
            continue
        if status < 400:
            request(store, success_key, 0, 100, duration, )
        yield "processed_after", status, store.get(success_key)

statuses = [200, 200, 200]
events = list(memory_handler({}, "model_success_", 1, statuses, 60))
for e in events:
    print(e)
PY

Repository: QuantumNous/new-api

Length of output: 219

Guard the in-memory success limiter when successMaxCount is zero.

In-memory handling mirrors the Redis zero-value contract for total limits, but the success path calls CanRequest and Request with successMaxCount == 0, so a single successful request records an entry and subsequent successful requests are rejected until expiry. Skip both calls when successMaxCount == 0, or document and enforce a 0-as-1 request limit consistently.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@middleware/model-rate-limit.go` around lines 147 - 156, Update the
success-rate-limit flow around CanRequest and the subsequent success-count
Request call to bypass both in-memory limiter operations when successMaxCount is
zero, preserving the Redis zero-value contract and allowing successful requests
without creating limiter entries.

if c.Writer.Status() < 400 {
inMemoryRateLimiter.Request(successKey, successMaxCount, duration)
}
Expand Down
34 changes: 34 additions & 0 deletions middleware/model_rate_limit_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,12 @@ package middleware

import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"

"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
Expand Down Expand Up @@ -32,3 +35,34 @@ 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")
}

// 成功数限制只统计成功请求，失败请求不占配额；拒绝时返回 JSON 错误体而非空 429。
func TestModelMemoryRateLimitSuccessCountIgnoresFailedRequests(t *testing.T) {
gin.SetMode(gin.TestMode)

// 限流器是进程级全局且无重置接口，user id 每次唯一才能保证 -count=2 复跑
userID := int(time.Now().UnixNano() % 1_000_000_000)
Comment on lines +39 to +44

@coderabbitai coderabbitai Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Use deterministic isolation for the global limiter state.

A wall-clock-derived ID can collide with another test or repeated execution, reusing stale process-global quota state. Reset/inject the limiter for the test, or allocate IDs from a reserved namespace using an atomic counter.

As per coding guidelines, backend tests should prefer deterministic inputs and explicit state.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@middleware/model_rate_limit_test.go` around lines 39 - 44, Replace the
wall-clock-derived userID in
TestModelMemoryRateLimitSuccessCountIgnoresFailedRequests with deterministic
isolation: reset or inject the process-global limiter for the test, or generate
IDs through an atomic counter in a reserved namespace. Ensure repeated and
parallel test runs cannot reuse stale quota state.

Source: Coding guidelines

downstreamStatus := http.StatusOK
router := gin.New()
router.GET("/limited", func(c *gin.Context) {
c.Set("id", userID)
}, memoryRateLimitHandler(60, 0, 2), func(c *gin.Context) {
c.Status(downstreamStatus)
})

do := func(status int) *httptest.ResponseRecorder {
downstreamStatus = status
return performRateLimitRequest(router, "/limited", "192.0.2.70:12345")
}

for range 5 {
assert.Equal(t, http.StatusInternalServerError, do(http.StatusInternalServerError).Code, "failed requests must not consume the success quota")
}

require.Equal(t, http.StatusOK, do(http.StatusOK).Code)
require.Equal(t, http.StatusOK, do(http.StatusOK).Code)

limited := do(http.StatusOK)
require.Equal(t, http.StatusTooManyRequests, limited.Code)
assert.Contains(t, limited.Body.String(), "您已达到请求数限制", "429 must carry the same error message as the Redis path")
}
Comment on lines +65 to +68

@coderabbitai coderabbitai Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Assert the JSON response contract, not only the message text.

assert.Contains would pass for a plain-text 429 containing the same phrase. Also assert an application/json content type and validate the expected JSON field/body shape.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@middleware/model_rate_limit_test.go` around lines 65 - 68, Update the 429
response assertions in the do test to verify the JSON contract rather than only
searching response text: assert an application/json content type and decode or
inspect the body to confirm the expected JSON field and value contain the
rate-limit message.