浅念
f938118fee
LoginLimiter._fails 与 rate_limit._buckets 只增不删——懒清理只覆盖被 回访的键，扫号换 IP/设备轮换产生的键永不被回访 → 进程长跑缓慢累积。 两处各加 sweep()（全表清「全部时间戳均已出窗」的键，窗沿 <= 与既有 保留判据 t > now - W 同沿），main.lifespan 起周期任务（300s loop） 启动创建/关停取消（cancel 后 await 落地，不留 pending task）。 Co-Authored-By: Claude Code <noreply@anthropic.com>
60 lines
2.2 KiB
Python
60 lines
2.2 KiB
Python
"""Minimal in-memory rate limiter (per (tenant, ip, device_id), rolling 1h window).
|
|
|
|
Kept aligned with the reference implementation's rate_limit.py shape. Only the chat message route uses this;
|
|
session/history endpoints are cheap and not throttled. State is process-local
|
|
and resets on restart, which is acceptable for a single-worker ai-kefu deploy.
|
|
Task 9: 桶键加租户维度——多租户同 IP(共享出口 NAT/公司网络)不互相挤占
|
|
配额;limit 经 config.engine.rate_limit_per_hour 逐租户可调(消费点
|
|
chat_rest, 默认 DEFAULT_MAX)。
|
|
|
|
Caveats:
|
|
- Single-process, in-memory. Under `uvicorn --workers >1` each worker keeps
|
|
its own `_buckets`, so the effective cap is `max_requests * workers`.
|
|
- Keyed on (tenant, ip, device_id); a client rotating the `X-Device-Id` header per
|
|
request bypasses it. This is a light safeguard against accidental GLM-cost
|
|
burn, NOT abuse-proof.
|
|
"""
|
|
import time
|
|
from collections import defaultdict
|
|
|
|
_buckets: dict[tuple[str, str, str], list[float]] = defaultdict(list)
|
|
WINDOW = 3600
|
|
DEFAULT_MAX = 30
|
|
|
|
|
|
def check(
|
|
tenant: str | None,
|
|
ip: str | None,
|
|
device_id: str | None,
|
|
*,
|
|
max_requests: int = DEFAULT_MAX,
|
|
) -> bool:
|
|
"""Return True if the request is under the cap (and record it), else False."""
|
|
key = (tenant or "", ip or "", device_id or "")
|
|
now = time.time()
|
|
_buckets[key] = [t for t in _buckets[key] if t > now - WINDOW]
|
|
if len(_buckets[key]) >= max_requests:
|
|
return False
|
|
_buckets[key].append(now)
|
|
return True
|
|
|
|
|
|
def sweep(now: float | None = None) -> int:
|
|
"""Sweep expired buckets (spec 2026-10-19 §7 memory cap): drop keys whose
|
|
timestamps have ALL left the window (boundary inclusive, matching check()'s
|
|
keep predicate ``t > now - WINDOW``); return the number removed.
|
|
|
|
Companion to check()'s per-key lazy pruning: keys never revisited (client
|
|
churning tenant/ip/device_id) would otherwise accumulate forever; called
|
|
by main.lifespan's periodic task."""
|
|
now = time.time() if now is None else now
|
|
dead = [k for k, ts in _buckets.items()
|
|
if not ts or ts[-1] <= now - WINDOW]
|
|
for k in dead:
|
|
del _buckets[k]
|
|
return len(dead)
|
|
|
|
|
|
def _reset_for_test() -> None:
|
|
_buckets.clear()
|