From 48a9058cb2d5db4ad243313c70b65311884cd2b0 Mon Sep 17 00:00:00 2001
From: claude_dev <claude-dev@sanguo.local>
Date: Mon, 6 Jul 2026 18:15:22 +0800
Subject: [PATCH] =?UTF-8?q?feat(api):=20WS=20=E8=BF=9E=E6=8E=A5=E6=B1=A0?=
 =?UTF-8?q?=20ConnectionManager?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 sanguo_api/ws.py     | 28 ++++++++++++++++++++++++++++
 tests/api/test_ws.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 sanguo_api/ws.py
 create mode 100644 tests/api/test_ws.py

diff --git a/sanguo_api/ws.py b/sanguo_api/ws.py
new file mode 100644
index 0000000..8c96423
--- /dev/null
+++ b/sanguo_api/ws.py
@@ -0,0 +1,28 @@
+# sanguo_api/ws.py
+"""WS 连接池：task_id → 订阅者集合。简单广播，不做重连/心跳。"""
+from fastapi import WebSocket
+
+
+class ConnectionManager:
+    def __init__(self):
+        self._connections: dict[str, set[WebSocket]] = {}
+
+    def connect(self, task_id: str, ws: WebSocket):
+        self._connections.setdefault(task_id, set()).add(ws)
+
+    def disconnect(self, task_id: str, ws: WebSocket):
+        conns = self._connections.get(task_id)
+        if conns:
+            conns.discard(ws)
+            if not conns:
+                del self._connections[task_id]
+
+    async def broadcast(self, task_id: str, msg: dict):
+        for ws in list(self._connections.get(task_id, [])):
+            try:
+                await ws.send_json(msg)
+            except Exception:
+                self.disconnect(task_id, ws)
+
+
+manager = ConnectionManager()
diff --git a/tests/api/test_ws.py b/tests/api/test_ws.py
new file mode 100644
index 0000000..6c310ff
--- /dev/null
+++ b/tests/api/test_ws.py
@@ -0,0 +1,30 @@
+# tests/api/test_ws.py
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+
+
+@pytest.mark.asyncio
+async def test_connect_and_broadcast():
+    from sanguo_api.ws import ConnectionManager
+    mgr = ConnectionManager()
+    ws = AsyncMock()
+    mgr.connect("t1", ws)
+    assert "t1" in mgr._connections
+    await mgr.broadcast("t1", {"status": "running"})
+    ws.send_json.assert_called_with({"status": "running"})
+
+
+def test_disconnect_removes_ws():
+    from sanguo_api.ws import ConnectionManager
+    mgr = ConnectionManager()
+    ws = MagicMock()
+    mgr.connect("t1", ws)
+    mgr.disconnect("t1", ws)
+    assert ws not in mgr._connections.get("t1", set())
+
+
+@pytest.mark.asyncio
+async def test_broadcast_no_subscribers_no_error():
+    from sanguo_api.ws import ConnectionManager
+    mgr = ConnectionManager()
+    await mgr.broadcast("nope", {"x": 1})  # 不抛