fix: exclude failed requests from memory rate limit success count - #6528
fix: exclude failed requests from memory rate limit success count#6528feitianbubu wants to merge 1 commit into
Conversation
WalkthroughThe in-memory rate limiter adds a non-counting eligibility check. Model middleware uses it before downstream execution and records only successful requests afterward. Tests verify failed responses do not consume the success quota and successful requests eventually receive HTTP 429. ChangesModel rate-limit success accounting
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant memoryRateLimitHandler
participant InMemoryRateLimiter
participant DownstreamHandler
Client->>memoryRateLimitHandler: Send model request
memoryRateLimitHandler->>InMemoryRateLimiter: CanRequest success key
InMemoryRateLimiter-->>memoryRateLimitHandler: Return eligibility
memoryRateLimitHandler->>DownstreamHandler: c.Next()
DownstreamHandler-->>memoryRateLimitHandler: Return response status
memoryRateLimitHandler->>InMemoryRateLimiter: Request success key for successful response
memoryRateLimitHandler-->>Client: Return response or HTTP 429
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with 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.
Inline comments:
In `@middleware/model_rate_limit_test.go`:
- Around line 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.
- Around line 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.
In `@middleware/model-rate-limit.go`:
- Around line 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.
- Around line 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.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 85969bbe-c029-428e-b22b-cda7c70bb582
📒 Files selected for processing (3)
common/rate-limit.gomiddleware/model-rate-limit.gomiddleware/model_rate_limit_test.go
| // 成功数限制只统计成功请求，失败请求不占配额；拒绝时返回 JSON 错误体而非空 429。 | ||
| func TestModelMemoryRateLimitSuccessCountIgnoresFailedRequests(t *testing.T) { | ||
| gin.SetMode(gin.TestMode) | ||
|
|
||
| // 限流器是进程级全局且无重置接口，user id 每次唯一才能保证 -count=2 复跑 | ||
| userID := int(time.Now().UnixNano() % 1_000_000_000) |
There was a problem hiding this comment.
🩺 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
| 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") | ||
| } |
There was a problem hiding this comment.
🎯 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.
| // 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. 如果请求成功，记录到成功请求计数中 |
There was a problem hiding this comment.
🩺 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 middlewareRepository: 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 middlewareRepository: 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 -nRepository: 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)
PYRepository: 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)
PYRepository: 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.
1 participant
📝 变更描述 / Description
最大成功请求数将失败也计入了总数修法：
成功请求数abortWithOpenAiMessage，文案与 Redis 一致。🚀 变更类型 / Type of change
🔗 关联任务 / Related Issue
✅ 提交前检查项 / Checklist
Bug fix，我已提交或关联对应 Issue，且不会将设计取舍、预期不一致或理解偏差直接归类为 bug。📸 运行证明 / Proof of Work
设置:
成功数限制(不含失败):
总请求数限制:
