newtw

core/websocket_client_v2.py

@@ -15,7 +15,7 @@ class QingjianAPIClient:
     重构后的轻简API客户端 (协议全复刻版)
     """
     
-    def __init__(self, uri=None, enable_agent: bool = True):
+    def __init__(self, uri=None, enable_agent: bool = True, worker_id: int = -1, worker_count: int = 1):
         from config.config import QINGJIAN_WS_URI
         self.uri = uri or QINGJIAN_WS_URI
         self.websocket = None
@@ -23,6 +23,12 @@ class QingjianAPIClient:
         self.logger = logger
         self.enable_agent = enable_agent
         
+        # 多进程分片逻辑
+        self.worker_id = worker_id
+        self.worker_count = worker_count
+        if self.worker_id >= 0:
+            logger.info(f"[WebSocket] 启用分片模式: Worker {self.worker_id}/{self.worker_count}")
+
         # 初始化新架构总指挥部
         self.orchestrator = init_orchestrator(ws_client=self)
         logger.info("[WebSocket] 新架构 Orchestrator 已就绪。")
@@ -36,13 +42,35 @@ class QingjianAPIClient:
     async def receive_messages(self):
         await receive_messages_flow(self)
 
+    def _should_handle(self, customer_id: str) -> bool:
+        """分片判定：这个客户归我管吗？"""
+        if self.worker_id < 0 or self.worker_count <= 1:
+            return True
+        
+        # 如果没有 customer_id，为了安全起见，只让 Worker 0 处理
+        if not customer_id:
+            return self.worker_id == 0
+            
+        import hashlib
+        # 使用稳定的哈希算法分配客户
+        hash_val = int(hashlib.md5(str(customer_id).encode("utf-8")).hexdigest(), 16)
+        return (hash_val % self.worker_count) == self.worker_id
+
     async def handle_message(self, message):
         """收到消息处理"""
         try:
             data = json.loads(message)
+            # 预提取客户ID用于分片判定
+            customer_id = str(data.get("cy_id") or data.get("from_id") or "")
+            if not self._should_handle(customer_id):
+                return
+
             await self.orchestrator.on_raw_message_received(platform="qianniu", raw_data=data)
         except Exception as e:
-            logger.error(f"[WebSocket] 处理消息异常: {e}")
+            raw_preview = str(message).replace("\n", "\\n")
+            if len(raw_preview) > 300:
+                raw_preview = raw_preview[:300] + "..."
+            logger.error(f"[WebSocket] 处理消息异常: {e} raw={raw_preview}")
 
     async def send(self, customer_id: str, acc_id: str, acc_type: str, content: str, msg_type: int = 0):
         """