更新

backend/rpa_engine/douyin_im/ws_client.py

@@ -1,4 +1,5 @@
 import asyncio
+import gzip
 import logging
 import os
 import weakref
@@ -14,6 +15,31 @@ logger = logging.getLogger("douyin_im.ws")
 
 MessageHandler = Callable[[dict], Awaitable[None]]
 
+
+def _safe_frame_metadata(payload: bytes) -> str:
+    """Return non-content protobuf metadata for early connection diagnostics."""
+    try:
+        from .static import Live_pb2, Response_pb2
+
+        frame = Live_pb2.PushFrame()
+        frame.ParseFromString(payload)
+        body = bytes(frame.payload)
+        if str(frame.payloadEncoding or "").lower() == "gzip":
+            body = gzip.decompress(body)
+        response = Response_pb2.Response()
+        response.ParseFromString(body)
+        fields = [field.name for field, _ in response.body.ListFields()]
+        message = str(response.message or response.error_desc or "")[:80]
+        return (
+            f"service={frame.service} method={frame.method} "
+            f"encoding={frame.payloadEncoding or 'none'} "
+            f"type={frame.payloadType or 'none'} payload_bytes={len(body)} "
+            f"cmd={response.cmd} body={','.join(fields) or 'none'} "
+            f"status={message or 'ok'}"
+        )
+    except Exception as exc:
+        return f"metadata_unavailable={type(exc).__name__}"
+
 # Both stages are finite.  The transport queue gives the receive coroutine a
 # small amount of breathing room, while the application queue decouples Pong /
 # frame reads from potentially slow database and reply work.  Once both fill,
@@ -127,6 +153,8 @@ class DouyinImWsClient:
         self._last_connection_lifetime = 0.0
         self._message_queue: Optional[asyncio.Queue[dict]] = None
         self._dispatcher_task: Optional[asyncio.Task] = None
+        self._received_frame_count = 0
+        self._heartbeat_ack_logged = False
 
     async def start(self):
         if self._task and not self._task.done():
@@ -331,6 +359,16 @@ class DouyinImWsClient:
             headers.append(("Cookie", cookie))
         return headers
 
+    @staticmethod
+    def _uses_browser_frontier(url: str) -> bool:
+        return "zijieapi.com" in url and "access_key=" in url
+
+    async def _run_browser_heartbeat(self, websocket) -> None:
+        """Mirror Frontier's browser SDK application-level ``hi`` heartbeat."""
+        while self._running:
+            await websocket.send("hi")
+            await asyncio.sleep(30)
+
     async def _run_connection(self, url: str) -> None:
         """Open one connection and dispatch messages sequentially.
 
@@ -343,6 +381,8 @@ class DouyinImWsClient:
         loop = asyncio.get_running_loop()
         connected_at: float | None = None
         connection: Optional[WebSocketClientProtocol] = None
+        heartbeat_task: Optional[asyncio.Task] = None
+        browser_frontier = self._uses_browser_frontier(url)
         source_ip = str(getattr(self.session, "egress_source_ip", "") or "").strip()
         connect_kwargs = {"local_addr": (source_ip, 0)} if source_ip else {}
         try:
@@ -354,7 +394,9 @@ class DouyinImWsClient:
                 user_agent_header=self.session.user_agent,
                 compression="deflate",
                 open_timeout=10,
-                ping_interval=20,
+                # The current Douyin browser Frontier SDK uses a text ``hi``
+                # heartbeat instead of RFC WebSocket ping frames.
+                ping_interval=None if browser_frontier else 20,
                 # A handler may legitimately wait up to SQLite's 30s busy
                 # timeout.  Leave enough headroom for queued work so a healthy
                 # socket isn't mistaken for a dead peer during that stall.
@@ -371,7 +413,10 @@ class DouyinImWsClient:
                 self._connection = websocket
                 connected_at = loop.time()
                 self.connected = True
-                logger.info("IM WebSocket connected")
+                logger.info(
+                    "IM WebSocket connected: subprotocol=%s",
+                    getattr(websocket, "subprotocol", None) or "none",
+                )
                 self._record_connection_system_event(
                     "connected",
                     "实时接收通道已连接",
@@ -379,10 +424,23 @@ class DouyinImWsClient:
                     level="success",
                 )
 
-                async for raw in websocket:
-                    if not self._running:
-                        break
-                    await self._dispatch(raw)
+                if browser_frontier:
+                    heartbeat_task = asyncio.create_task(
+                        self._run_browser_heartbeat(websocket),
+                        name=f"im-ws-heartbeat-{self.account_id or 'na'}",
+                    )
+                try:
+                    async for raw in websocket:
+                        if not self._running:
+                            break
+                        await self._dispatch(raw)
+                finally:
+                    if heartbeat_task and not heartbeat_task.done():
+                        heartbeat_task.cancel()
+                        try:
+                            await heartbeat_task
+                        except asyncio.CancelledError:
+                            pass
         finally:
             if connected_at is not None:
                 self._last_connection_lifetime = max(0.0, loop.time() - connected_at)
@@ -408,6 +466,11 @@ class DouyinImWsClient:
                     )
 
     async def _dispatch(self, raw):
+        if raw == "hi":
+            if not self._heartbeat_ack_logged:
+                logger.info("IM WebSocket application heartbeat acknowledged")
+                self._heartbeat_ack_logged = True
+            return
         self._ensure_dispatcher()
         queue = self._message_queue
         if queue is None:
@@ -417,6 +480,17 @@ class DouyinImWsClient:
         else:
             payload = raw
         items = parse_ws_payload(payload)
+        self._received_frame_count += 1
+        if self._received_frame_count <= 3:
+            metadata = _safe_frame_metadata(payload) if not items else "parsed-message"
+            logger.info(
+                "IM WebSocket frame received: seq=%d kind=%s bytes=%d parsed=%d %s",
+                self._received_frame_count,
+                "text" if isinstance(raw, str) else "binary",
+                len(payload),
+                len(items),
+                metadata,
+            )
         for item in items:
             if not self._running:
                 return