mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat(rate-limit): 实现分层 RPM 限速，支持系统默认/用户/独立Key三级配置
- 新增用户级 rate_limit 字段，支持系统默认/用户自定义/不限制三种模式 - 独立 Key 的 rate_limit 语义调整：null=跟随系统默认，0=不限制，>0=自定义 - 实现 UserRpmLimiter 基于 Redis sliding window 的 RPM 限速引擎 - Pipeline 请求流程集成用户级 RPM 检查 - 管理后台和用户面板新增 RPM 限速配置与实时状态查看 - 系统设置新增全局默认 RPM 配置项 - 迁移脚本回填现有 API Key 的 rate_limit 默认值 - 新增用户/Key RPM 状态监控 API 和前端展示 Closes #231 Co-authored-by: LewisPen <LewisPen@nyadoo.com>
This commit is contained in:
fawney19
@@ -18,6 +18,7 @@ from src.api.base.adapter import ApiMode
|
||||
from src.api.base.pipeline import ApiRequestPipeline
|
||||
from src.core.enums import UserRole
|
||||
from src.core.modules.hooks import AUTH_TOKEN_PREFIX_AUTHENTICATORS
|
||||
from src.services.rate_limit.user_rpm_limiter import RpmCheckResult
|
||||
|
||||
|
||||
class TestPipelineBalanceCalculation:
|
||||
@@ -499,6 +500,166 @@ class TestPipelineAuthentication:
|
||||
assert "锁定" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
class TestPipelineUserRateLimit:
|
||||
@pytest.fixture
|
||||
def pipeline(self) -> ApiRequestPipeline:
|
||||
return ApiRequestPipeline()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_user_rate_limit_uses_system_default_for_user_scope(
|
||||
self, pipeline: ApiRequestPipeline, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request = MagicMock()
|
||||
request.state = MagicMock()
|
||||
db = MagicMock()
|
||||
user = MagicMock(id="user-1", rate_limit=None)
|
||||
api_key = MagicMock(id="key-1", is_standalone=False, rate_limit=0)
|
||||
|
||||
limiter = MagicMock()
|
||||
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:1"
|
||||
limiter.get_key_rpm_key.return_value = "rpm:key:key-1:1"
|
||||
limiter.check_and_consume = AsyncMock(return_value=RpmCheckResult(allowed=True))
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.base.pipeline.get_user_rpm_limiter",
|
||||
AsyncMock(return_value=limiter),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.base.pipeline.SystemConfigService.get_config",
|
||||
lambda *_a, **_k: 60,
|
||||
)
|
||||
|
||||
await pipeline._check_user_rate_limit(request, db, user, api_key)
|
||||
|
||||
limiter.check_and_consume.assert_awaited_once_with(
|
||||
user_rpm_key="rpm:user:user-1:1",
|
||||
user_rpm_limit=60,
|
||||
key_rpm_key="rpm:key:key-1:1",
|
||||
key_rpm_limit=0,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_user_rate_limit_returns_429_with_scope_header(
|
||||
self, pipeline: ApiRequestPipeline, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request = MagicMock()
|
||||
request.state = MagicMock()
|
||||
db = MagicMock()
|
||||
user = MagicMock(id="user-1", rate_limit=100)
|
||||
api_key = MagicMock(id="key-1", is_standalone=False, rate_limit=10)
|
||||
|
||||
limiter = MagicMock()
|
||||
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:1"
|
||||
limiter.get_key_rpm_key.return_value = "rpm:key:key-1:1"
|
||||
limiter.get_retry_after.return_value = 17
|
||||
limiter.check_and_consume = AsyncMock(
|
||||
return_value=RpmCheckResult(
|
||||
allowed=False,
|
||||
scope="key",
|
||||
limit=10,
|
||||
remaining=0,
|
||||
retry_after=17,
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.base.pipeline.get_user_rpm_limiter",
|
||||
AsyncMock(return_value=limiter),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.base.pipeline.SystemConfigService.get_config",
|
||||
lambda *_a, **_k: 60,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pipeline._check_user_rate_limit(request, db, user, api_key)
|
||||
|
||||
assert exc_info.value.status_code == 429
|
||||
assert exc_info.value.headers == {
|
||||
"Retry-After": "17",
|
||||
"X-RateLimit-Limit": "10",
|
||||
"X-RateLimit-Remaining": "0",
|
||||
"X-RateLimit-Scope": "key",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_user_rate_limit_uses_system_default_for_standalone_key(
|
||||
self, pipeline: ApiRequestPipeline, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request = MagicMock()
|
||||
request.state = MagicMock()
|
||||
db = MagicMock()
|
||||
user = MagicMock(id="user-1", rate_limit=999)
|
||||
api_key = MagicMock(id="standalone-1", is_standalone=True, rate_limit=None)
|
||||
|
||||
limiter = MagicMock()
|
||||
limiter.get_standalone_rpm_key.return_value = "rpm:ukey:standalone-1:1"
|
||||
limiter.get_key_rpm_key.return_value = "rpm:key:standalone-1:1"
|
||||
limiter.check_and_consume = AsyncMock(return_value=RpmCheckResult(allowed=True))
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.base.pipeline.get_user_rpm_limiter",
|
||||
AsyncMock(return_value=limiter),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.base.pipeline.SystemConfigService.get_config",
|
||||
lambda *_a, **_k: 60,
|
||||
)
|
||||
|
||||
await pipeline._check_user_rate_limit(request, db, user, api_key)
|
||||
|
||||
limiter.check_and_consume.assert_awaited_once_with(
|
||||
user_rpm_key="rpm:ukey:standalone-1:1",
|
||||
user_rpm_limit=60,
|
||||
key_rpm_key="rpm:key:standalone-1:1",
|
||||
key_rpm_limit=0,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_user_rate_limit_returns_429_with_user_scope_header(
|
||||
self, pipeline: ApiRequestPipeline, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request = MagicMock()
|
||||
request.state = MagicMock()
|
||||
db = MagicMock()
|
||||
user = MagicMock(id="user-1", rate_limit=3)
|
||||
api_key = MagicMock(id="key-1", is_standalone=False, rate_limit=10)
|
||||
|
||||
limiter = MagicMock()
|
||||
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:1"
|
||||
limiter.get_key_rpm_key.return_value = "rpm:key:key-1:1"
|
||||
limiter.get_retry_after.return_value = 23
|
||||
limiter.check_and_consume = AsyncMock(
|
||||
return_value=RpmCheckResult(
|
||||
allowed=False,
|
||||
scope="user",
|
||||
limit=3,
|
||||
remaining=0,
|
||||
retry_after=23,
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.base.pipeline.get_user_rpm_limiter",
|
||||
AsyncMock(return_value=limiter),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.base.pipeline.SystemConfigService.get_config",
|
||||
lambda *_a, **_k: 60,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pipeline._check_user_rate_limit(request, db, user, api_key)
|
||||
|
||||
assert exc_info.value.status_code == 429
|
||||
assert exc_info.value.headers == {
|
||||
"Retry-After": "23",
|
||||
"X-RateLimit-Limit": "3",
|
||||
"X-RateLimit-Remaining": "0",
|
||||
"X-RateLimit-Scope": "user",
|
||||
}
|
||||
|
||||
|
||||
class TestPipelineTokenPrefixAuth:
|
||||
"""Tests token-prefix auth isolation."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user