更新

backend/rpa_engine/douyin_im/protocol.py

@@ -1,7 +1,11 @@
 import gzip
 import json
 import logging
+import logging.handlers
+import os
+import queue
 import re
+import threading
 from typing import Any, Optional
 
 from .message_content import (
@@ -21,7 +25,6 @@ from .message_content import (
 
 logger = logging.getLogger("douyin_im.protocol")
 
-import os
 from datetime import datetime
 
 
@@ -42,6 +45,73 @@ def _is_control_payload(content_json: Any, msg_type: int = 0) -> bool:
     return False
 
 _WS_DEBUG_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "ws_media_debug.log")
+_WS_DEBUG_WRITER_LOCK = threading.Lock()
+_WS_DEBUG_LOGGER: Optional[logging.Logger] = None
+
+
+def _bounded_env_int(name: str, default: int, minimum: int, maximum: int) -> int:
+    try:
+        value = int(os.getenv(name, str(default)) or default)
+    except (TypeError, ValueError):
+        value = default
+    return max(minimum, min(maximum, value))
+
+
+class _DroppingQueueHandler(logging.handlers.QueueHandler):
+    """Never let optional diagnostics block the IM event loop."""
+
+    def enqueue(self, record) -> None:
+        try:
+            self.queue.put_nowait(record)
+        except queue.Full:
+            # Debug output is intentionally lossy under pressure.  Receiving
+            # and replying to messages must always take precedence.
+            return
+
+
+def _get_ws_debug_logger() -> logging.Logger:
+    global _WS_DEBUG_LOGGER
+    if _WS_DEBUG_LOGGER is not None:
+        return _WS_DEBUG_LOGGER
+    with _WS_DEBUG_WRITER_LOCK:
+        if _WS_DEBUG_LOGGER is not None:
+            return _WS_DEBUG_LOGGER
+
+        max_bytes = _bounded_env_int(
+            "KEFU_WS_DEBUG_MAX_BYTES", 10 * 1024 * 1024, 1024 * 1024, 100 * 1024 * 1024
+        )
+        backup_count = _bounded_env_int(
+            "KEFU_WS_DEBUG_BACKUP_COUNT", 2, 1, 10
+        )
+        queue_size = _bounded_env_int(
+            "KEFU_WS_DEBUG_QUEUE_SIZE", 1000, 100, 10000
+        )
+        records: queue.Queue = queue.Queue(maxsize=queue_size)
+        rotating = logging.handlers.RotatingFileHandler(
+            _WS_DEBUG_PATH,
+            maxBytes=max_bytes,
+            backupCount=backup_count,
+            encoding="utf-8",
+            delay=True,
+        )
+        rotating.setFormatter(logging.Formatter("%(message)s"))
+        listener = logging.handlers.QueueListener(
+            records,
+            rotating,
+            respect_handler_level=True,
+        )
+        listener.start()
+
+        debug_logger = logging.getLogger("douyin_im.ws_raw_debug")
+        debug_logger.handlers.clear()
+        debug_logger.addHandler(_DroppingQueueHandler(records))
+        debug_logger.setLevel(logging.INFO)
+        debug_logger.propagate = False
+        # Keep strong references for the lifetime of the logger/listener.
+        debug_logger._kefu_queue_listener = listener  # type: ignore[attr-defined]
+        debug_logger._kefu_rotating_handler = rotating  # type: ignore[attr-defined]
+        _WS_DEBUG_LOGGER = debug_logger
+        return debug_logger
 
 
 def _should_emit_ws_message(
@@ -84,8 +154,13 @@ def _dump_ws_message(msg_type: int, conversation_id: str, content_str: str, msg:
             f"{datetime.now().isoformat()} type={msg_type} "
             f"conv={conversation_id} content={content_str}{extra}\n"
         )
-        with open(_WS_DEBUG_PATH, "a", encoding="utf-8") as fh:
-            fh.write(line)
+        record_limit = _bounded_env_int(
+            "KEFU_WS_DEBUG_RECORD_MAX_CHARS", 16384, 1024, 262144
+        )
+        if len(line) > record_limit:
+            marker = "...[单条调试记录过长，已截断]\n"
+            line = line[: max(0, record_limit - len(marker))] + marker
+        _get_ws_debug_logger().info(line.rstrip("\n"))
     except Exception:
         pass
 
@@ -153,8 +228,6 @@ def parse_ws_payload(raw: bytes | str) -> list[dict]:
                         content_str = msg.content
                         server_message_id = str(getattr(msg, "server_message_id", "") or "")
 
-                        _dump_ws_message(msg_type, conversation_id, content_str, msg)
-
                         text_content = ""
                         media_msg: dict = {}
                         content_json: dict = {}
@@ -175,6 +248,12 @@ def parse_ws_payload(raw: bytes | str) -> list[dict]:
                             )
                             return messages
 
+                        # Raw diagnostics are optional and intentionally run
+                        # only after control/status frames have been filtered.
+                        # The writer itself is queued and rotating, so it can
+                        # never block message parsing or grow without bound.
+                        _dump_ws_message(msg_type, conversation_id, content_str, msg)
+
                         if _should_emit_ws_message(conversation_id, msg_type):
                             sender_uid = str(msg.sender)
                             if media_msg and (