feat: WS 每 IP 并发连接上限（默认10，(租户,IP) 键，finally 释放）

Co-Authored-By: Claude Code <noreply@anthropic.com>

app/engine/routes/chat.py

@@ -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()
-    session_id = None
-
-    try:
-        # Task 5b: 租户运行时经 state/默认租户兜底解析(过渡期无中间件);
-        # manager 由装配线从 runtime 造(db 会话 + runtime 引用)
+    # Task 8: 每 (租户, IP) 并发连接上限——必须在 accept 前拒绝(close 代
+    # accept,客户端握手即断)。IP 提取与下方 client_info 同一来源(XFF 优先,
+    # scope client 兜底),握手期 headers/scope 已就绪、不依赖首条 init 消息;
+    # 客户端自报的 client_info.ip 仍按原逻辑只进会话记录,不参与限流。
     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
+
+    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 引用)
+        async with runtime.session_factory() as db:
+            manager = ConversationManager(db, runtime=runtime)
+
+            # 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)

app/engine/services/conn_limiter.py

@@ -0,0 +1,57 @@
+"""WS 每 IP 并发连接上限(spec §6.1 P2)——进程内 (租户, IP) 计数闸。
+
+动机: 平台化后单 IP 可对所有租户域名各开一圈 WS 耗尽单进程资源;
+REST 消息接口已有的 rate_limit.py 管不到长连接握手。key 是
+(tenant_slug, ip) 二元组——各租户名额独立,不共享预算。
+
+约束与取舍(与 rate_limit.py 同一族):
+- 单进程内存态, ``--workers 1`` 约束下准确; 多 worker 各自一份计数,
+  实际上限 = limit * workers。
+- 纯同步 dict 操作、方法内无 await——单事件循环下 check-then-set
+  原子,无需锁;不得在多线程里使用。
+- 计数归零即删 key, 长期运行不随历史 (租户, IP) 对积累内存;
+  release 未持有的 key 是安全 no-op(不出负数)。
+- limit 可在 try_acquire 时按租户配置覆盖: 模块单例服务全部租户,而
+  ws_conn_limit 是 engine 组的租户级配置——按次传入让各租户差异实时
+  生效(含 registry_reload 后), 又不必维护"每租户一个 limiter 实例"
+  的第二层簿记(实例缓存还有配置改了不生效的滞留问题)。
+"""
+from typing import Optional
+
+DEFAULT_WS_CONN_LIMIT = 10
+
+
+class ConnectionLimiter:
+    """每 (tenant, ip) 并发计数器: 满则 try_acquire 拒绝, release 归零清键。"""
+
+    def __init__(self, limit: int = DEFAULT_WS_CONN_LIMIT) -> None:
+        self._limit = limit
+        self._counts: dict[tuple[str, str], int] = {}
+
+    def try_acquire(self, tenant: str, ip: str, *,
+                    limit: Optional[int] = None) -> bool:
+        """未满则记账并返回 True; 已满返回 False(不记账)。"""
+        cap = self._limit if limit is None else limit
+        key = (tenant, ip)
+        n = self._counts.get(key, 0)
+        if n >= cap:
+            return False
+        self._counts[key] = n + 1
+        return True
+
+    def release(self, tenant: str, ip: str) -> None:
+        """归还一个名额; 计数归零删 key; 未持有的 key 是安全 no-op。"""
+        key = (tenant, ip)
+        n = self._counts.get(key, 0)
+        if n <= 1:
+            self._counts.pop(key, None)
+        else:
+            self._counts[key] = n - 1
+
+
+# app 级单例: chat.py 的 WS 路由共用这一份进程内计数(--workers 1 约束)
+limiter = ConnectionLimiter()
+
+
+def _reset_for_test() -> None:
+    limiter._counts.clear()

app/platform/defaults.py

@@ -16,7 +16,8 @@ DEFAULT_CONFIG: dict = {
     "push": {"enabled": False, "url": "", "api_key": "", "prefix": "", "timeout": 10,
              "max_attempts": 3, "retry_backoff_base": 1.0},
     "ai": {"provider": "glm", "api_key": None, "base_url": None, "model": "glm-4-flash"},
-    "engine": {"context_seed_keys": {}, "rate_limit_per_hour": 30},
+    "engine": {"context_seed_keys": {}, "rate_limit_per_hour": 30,
+               "ws_conn_limit": 10},
     "misc": {"name_to_region_extras": {}},
 }
 

tests/conftest.py

@@ -64,6 +64,16 @@ def _reset_rate_limit():
     _reset_for_test()
 
 
+@pytest.fixture(autouse=True)
+def _reset_conn_limiter():
+    """Task 8: 同上——WS 连接限流器的模块级单例计数前后复位,跨测试不泄漏
+    (TestClient 所有连接同落 ip="testclient",不复位会互相顶名额)。"""
+    from app.engine.services.conn_limiter import _reset_for_test
+    _reset_for_test()
+    yield
+    _reset_for_test()
+
+
 @pytest.fixture
 def jiancai_runtime():
     """直连测试用的 jiancai TenantRuntime(Task 6: manager 的 runtime 必填)。

tests/test_conn_limiter.py

@@ -0,0 +1,80 @@
+# Task 8: WS 每 IP 并发连接上限(spec §6.1 P2)。
+#
+# 两层测试:
+# - 纯单元: ConnectionLimiter 算法(名额/键隔离/归零清 key/no-op release);
+# - 接入: /ws/chat 握手期按 (租户, IP) 限流——第 limit+1 条并发连接被
+#   1013 拒(close 代 accept),连接关闭后名额必须释放。
+import json
+from contextlib import ExitStack
+
+import pytest
+from starlette.websockets import WebSocketDisconnect
+
+from app.engine.services.conn_limiter import ConnectionLimiter
+
+# ---- 纯单元: 算法 -----------------------------------------------------------
+
+
+def test_limit_per_ip():
+    lim = ConnectionLimiter(limit=2)
+    assert lim.try_acquire("t1", "1.1.1.1") and lim.try_acquire("t1", "1.1.1.1")
+    assert not lim.try_acquire("t1", "1.1.1.1")          # 第 3 个拒绝
+    lim.release("t1", "1.1.1.1")
+    assert lim.try_acquire("t1", "1.1.1.1")              # 释放后可再进
+    assert lim.try_acquire("t2", "1.1.1.1")              # 其他租户不受影响
+    assert lim.try_acquire("t1", "2.2.2.2")              # 其他 IP 不受影响
+
+
+def test_count_zero_cleans_key():
+    # 计数归零删 key: 长期运行下内部 dict 不随历史 (租户,IP) 对无限增长
+    lim = ConnectionLimiter(limit=3)
+    lim.try_acquire("t", "1.1.1.1")
+    lim.try_acquire("t", "1.1.1.1")
+    lim.release("t", "1.1.1.1")
+    assert ("t", "1.1.1.1") in lim._counts               # 还剩 1 个,键保留
+    lim.release("t", "1.1.1.1")
+    assert lim._counts == {}                              # 归零即清空
+
+
+def test_release_unheld_key_is_safe_noop():
+    lim = ConnectionLimiter(limit=2)
+    lim.release("nobody", "9.9.9.9")                      # 从未 acquire 过
+    lim.try_acquire("t", "1.1.1.1")
+    lim.release("t", "1.1.1.1")
+    lim.release("t", "1.1.1.1")                           # 双重 release
+    assert lim._counts == {}                              # 不出负计数
+    assert lim.try_acquire("t", "1.1.1.1")                # 状态仍可用
+
+
+# ---- 接入: /ws/chat 握手期限流 ----------------------------------------------
+# TestClient 的 WS scope client 恒为 ("testclient", 50000) 且测试不带 XFF,
+# 故所有连接同落 ip="testclient";app_client 默认 jiancai 域——同租户同 IP,
+# 正好是被限的键。默认 ws_conn_limit=10(defaults.py engine 组)。
+
+
+def test_ws_over_limit_rejected_with_1013(app_client):
+    """填满 10 个名额 → 第 11 条握手即被拒(1013, 未 accept)。"""
+    with ExitStack() as stack:
+        conns = [
+            stack.enter_context(app_client.websocket_connect("/ws/chat"))
+            for _ in range(10)
+        ]
+        # 前 10 条真实可用(走完 init),不是恰好没崩
+        conns[0].send_text(json.dumps({"source": "standalone"}))
+        assert conns[0].receive_json()["type"] == "init"
+
+        with pytest.raises(WebSocketDisconnect) as ei:
+            with app_client.websocket_connect("/ws/chat"):
+                pass  # pragma: no cover — 未 accept 即 close,进不来
+        assert ei.value.code == 1013
+        assert ei.value.reason == "too many connections"
+
+
+def test_ws_conn_release_allows_reentry(app_client):
+    with ExitStack() as stack:
+        for _ in range(10):
+            stack.enter_context(app_client.websocket_connect("/ws/chat"))
+    # ExitStack 退出 = 10 条全部关闭,服务端 finally 已 release
+    with app_client.websocket_connect("/ws/chat") as ws:
+        ws.send_text(json.dumps({"source": "standalone"}))
+        assert ws.receive_json()["type"] == "init"