feat: WS 每 IP 并发连接上限(默认10,(租户,IP) 键,finally 释放)
Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
@@ -5,6 +5,8 @@ from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||||
|
||||
from app.engine.models import ConversationSource, ConversationStatus
|
||||
from app.engine.routes.runtime import get_runtime
|
||||
from app.engine.services.conn_limiter import DEFAULT_WS_CONN_LIMIT
|
||||
from app.engine.services.conn_limiter import limiter as conn_limiter
|
||||
from app.engine.services.conversation import ConversationManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -13,21 +15,32 @@ router = APIRouter()
|
||||
|
||||
@router.websocket("/ws/chat")
|
||||
async def ws_chat(websocket: WebSocket):
|
||||
await websocket.accept()
|
||||
# Task 8: 每 (租户, IP) 并发连接上限——必须在 accept 前拒绝(close 代
|
||||
# accept,客户端握手即断)。IP 提取与下方 client_info 同一来源(XFF 优先,
|
||||
# scope client 兜底),握手期 headers/scope 已就绪、不依赖首条 init 消息;
|
||||
# 客户端自报的 client_info.ip 仍按原逻辑只进会话记录,不参与限流。
|
||||
runtime = get_runtime(websocket)
|
||||
client_ip = websocket.client.host if websocket.client else ""
|
||||
forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip()
|
||||
if forwarded:
|
||||
client_ip = forwarded
|
||||
|
||||
limit = runtime.config.data["engine"].get(
|
||||
"ws_conn_limit", DEFAULT_WS_CONN_LIMIT)
|
||||
if not conn_limiter.try_acquire(runtime.config.slug, client_ip, limit=limit):
|
||||
await websocket.close(code=1013, reason="too many connections")
|
||||
return
|
||||
|
||||
session_id = None
|
||||
|
||||
try:
|
||||
await websocket.accept()
|
||||
# Task 5b: 租户运行时经 state/默认租户兜底解析(过渡期无中间件);
|
||||
# manager 由装配线从 runtime 造(db 会话 + runtime 引用)
|
||||
runtime = get_runtime(websocket)
|
||||
async with runtime.session_factory() as db:
|
||||
manager = ConversationManager(db, runtime=runtime)
|
||||
|
||||
# Collect client info
|
||||
client_ip = websocket.client.host if websocket.client else ""
|
||||
forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip()
|
||||
if forwarded:
|
||||
client_ip = forwarded
|
||||
# Collect client info(client_ip/forwarded 已在握手期提取,复用)
|
||||
user_agent = websocket.headers.get("user-agent", "")
|
||||
client_info = {"ip": client_ip, "user_agent": user_agent}
|
||||
|
||||
@@ -107,3 +120,7 @@ async def ws_chat(websocket: WebSocket):
|
||||
logger.info("Client disconnected: %s", session_id)
|
||||
except Exception:
|
||||
logger.exception("WebSocket error for session %s", session_id)
|
||||
finally:
|
||||
# Task 8: 正常关闭与异常断开都归还名额——acquire 成功后所有路径
|
||||
# 必经此处,不泄漏(accept 失败/循环中抛错同样覆盖)
|
||||
conn_limiter.release(runtime.config.slug, client_ip)
|
||||
|
||||
Reference in New Issue
Block a user