更新

backend/rpa_engine/playwright_worker.py

@@ -14,6 +14,11 @@ from playwright.async_api import async_playwright
 from models.database import AsyncSessionLocal
 from models.models import Account, AutoReplyRule, MessageLog, AccountProfileDetail, FollowWelcomeLog
 from utils.received_message_log import record_received_message
+from utils.log_limits import (
+    bound_error_log_content,
+    bound_message_log_content,
+    truncate_text,
+)
 from utils.cookie_store import get_cookie_path, read_cookie_file, analyze_cookie, merge_playwright_cookies
 from utils import system_logger
 from rpa_engine.douyin_im import DouyinImService
@@ -64,9 +69,16 @@ def format_error(exc: BaseException) -> str:
 
 
 class DouyinWorker:
-    def __init__(self, account_id: int, login_mode: str = "auto"):
+    def __init__(
+        self,
+        account_id: int,
+        login_mode: str = "auto",
+        *,
+        credential_prevalidated: bool = False,
+    ):
         self.account_id = account_id
         self.login_mode = login_mode  # auto | im_direct | browser
+        self.credential_prevalidated = bool(credential_prevalidated)
         self.browser = None
         self.context = None
         self.page = None
@@ -74,6 +86,8 @@ class DouyinWorker:
         self.is_running = False
         self.stopping = False
         self._task: asyncio.Task | None = None
+        self._startup_ready = asyncio.Event()
+        self._startup_error = ""
         self.session_dir = os.path.join(
             os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
             "sessions"
@@ -104,6 +118,15 @@ class DouyinWorker:
         self._refresh_cooldown = 90.0
         self._user_agent: str = ""
         self._sec_user_id_missing_fired = False
+        # Lightweight follow-welcome configuration. Disabled accounts refresh
+        # infrequently, so 500 idle workers do not query Account + sec_user_id
+        # every minute merely to discover that the feature is still off.
+        self._follow_config_lock = asyncio.Lock()
+        self._follow_config_loaded = False
+        self._follow_config_refresh_at = 0.0
+        self._follow_welcome_enabled = False
+        self._follow_welcome_content = ""
+        self._follow_welcome_sec_user_id = ""
 
     async def _load_user_agent(self) -> str:
         """读取账号配置的伪装设备头，用于浏览器与 IM 全链路一致。"""
@@ -111,9 +134,10 @@ class DouyinWorker:
             return self._user_agent
         db = await self.get_db()
         try:
-            result = await db.execute(select(Account).where(Account.id == self.account_id))
-            account = result.scalar_one_or_none()
-            self._user_agent = resolve_user_agent(account.user_agent if account else None)
+            result = await db.execute(
+                select(Account.user_agent).where(Account.id == self.account_id)
+            )
+            self._user_agent = resolve_user_agent(result.scalar_one_or_none())
         finally:
             await db.close()
         return self._user_agent
@@ -145,6 +169,105 @@ class DouyinWorker:
     async def get_db(self):
         return AsyncSessionLocal()
 
+    def _mark_startup_ready(self) -> None:
+        self._startup_error = ""
+        self._startup_ready.set()
+
+    def _mark_startup_failed(self, detail: str = "") -> None:
+        if self._startup_ready.is_set():
+            return
+        self._startup_error = (
+            str(detail or "").strip()
+            or "托管任务在完成初始化前已退出"
+        )
+        self._startup_ready.set()
+
+    async def wait_until_ready(self) -> None:
+        """Wait until IM startup completed, or raise its initialization error.
+
+        Batch admission can await this signal so its concurrency limit covers
+        UID/frontier/WS/first-poll initialization instead of only covering the
+        creation of a detached worker task.
+        """
+        await self._startup_ready.wait()
+        if self._startup_error:
+            raise RuntimeError(self._startup_error)
+
+    async def _refresh_follow_welcome_config(
+        self,
+        *,
+        force: bool = False,
+    ) -> tuple[bool, str, str]:
+        now = time.monotonic()
+        if (
+            not force
+            and self._follow_config_loaded
+            and now < self._follow_config_refresh_at
+        ):
+            return (
+                self._follow_welcome_enabled,
+                self._follow_welcome_content,
+                self._follow_welcome_sec_user_id,
+            )
+
+        async with self._follow_config_lock:
+            now = time.monotonic()
+            if (
+                not force
+                and self._follow_config_loaded
+                and now < self._follow_config_refresh_at
+            ):
+                return (
+                    self._follow_welcome_enabled,
+                    self._follow_welcome_content,
+                    self._follow_welcome_sec_user_id,
+                )
+
+            db = await self.get_db()
+            try:
+                row = (
+                    await db.execute(
+                        select(
+                            Account.follow_welcome_enabled,
+                            Account.follow_welcome_content,
+                            AccountProfileDetail.sec_user_id,
+                        )
+                        .outerjoin(
+                            AccountProfileDetail,
+                            AccountProfileDetail.account_id == Account.id,
+                        )
+                        .where(Account.id == self.account_id)
+                    )
+                ).first()
+            finally:
+                await db.close()
+
+            if row:
+                enabled, content, sec_user_id = row
+                self._follow_welcome_enabled = bool(enabled)
+                self._follow_welcome_content = str(content or "").strip()
+                self._follow_welcome_sec_user_id = str(sec_user_id or "").strip()
+            else:
+                self._follow_welcome_enabled = False
+                self._follow_welcome_content = ""
+                self._follow_welcome_sec_user_id = ""
+            self._follow_config_loaded = True
+            # Enabled accounts retain the old one-minute configuration
+            # responsiveness. Disabled accounts perform only one lightweight
+            # refresh every ten minutes instead of one full Account read/minute.
+            ttl = 60.0 if self._follow_welcome_enabled else 600.0
+            self._follow_config_refresh_at = now + ttl
+            return (
+                self._follow_welcome_enabled,
+                self._follow_welcome_content,
+                self._follow_welcome_sec_user_id,
+            )
+
+    def invalidate_follow_welcome_config(self) -> None:
+        """Make the next follow tick reload settings after an account edit."""
+        self._follow_config_loaded = False
+        self._follow_config_refresh_at = 0.0
+
     async def _load_sec_user_id(self) -> str:
         """Return the locally persisted Douyin sec_user_id for this account."""
         db = await self.get_db()
@@ -188,10 +311,14 @@ class DouyinWorker:
         """Resolve and persist sec_user_id once from the account's current Cookie."""
         db = await self.get_db()
         try:
-            result = await db.execute(select(Account).where(Account.id == self.account_id))
-            account = result.scalar_one_or_none()
-            cookie_data = account.cookie_data if account else None
-            user_agent = account.user_agent if account else None
+            result = await db.execute(
+                select(Account.cookie_data, Account.user_agent).where(
+                    Account.id == self.account_id
+                )
+            )
+            row = result.first()
+            cookie_data = row.cookie_data if row else None
+            user_agent = row.user_agent if row else None
         finally:
             await db.close()
 
@@ -424,10 +551,12 @@ class DouyinWorker:
         """从数据库或本地文件加载 Playwright storage_state"""
         db = await self.get_db()
         try:
-            result = await db.execute(select(Account).where(Account.id == self.account_id))
-            account = result.scalar_one_or_none()
-            if account and account.cookie_data:
-                return json.loads(account.cookie_data)
+            result = await db.execute(
+                select(Account.cookie_data).where(Account.id == self.account_id)
+            )
+            cookie_data = result.scalar_one_or_none()
+            if cookie_data:
+                return json.loads(cookie_data)
         except Exception as e:
             logger.warning(f"Failed to load cookie from database: {e}")
         finally:
@@ -478,10 +607,10 @@ class DouyinWorker:
         db = await self.get_db()
         saved_im = None
         try:
-            result = await db.execute(select(Account).where(Account.id == self.account_id))
-            account = result.scalar_one_or_none()
-            if account:
-                saved_im = account.im_session_data
+            result = await db.execute(
+                select(Account.im_session_data).where(Account.id == self.account_id)
+            )
+            saved_im = result.scalar_one_or_none()
         finally:
             await db.close()
 
@@ -560,7 +689,21 @@ class DouyinWorker:
         """Cookie 有效时跳过浏览器，直接 IM 直连托管"""
         await self._load_user_agent()
         im_session = await self._build_im_session_from_storage(storage_state)
-        ok, reason = await validate_im_session(im_session)
+        if self.credential_prevalidated:
+            # Batch preparation already performed the remote credential probe.
+            # Re-check only the immutable local requirements after rebuilding
+            # the session, avoiding a duplicate query/user request per account.
+            from rpa_engine.douyin_im.auth import DouyinAuth
+
+            auth = DouyinAuth.from_im_session(im_session)
+            ok = bool(im_session.can_direct_im() and auth.is_sign_ready())
+            reason = (
+                "IM 凭证已在启动队列中校验"
+                if ok
+                else "启动后的本地 IM 凭证不再满足直连条件"
+            )
+        else:
+            ok, reason = await validate_im_session(im_session)
         if not ok:
             logger.warning(f"IM session validation failed: {reason}")
             system_logger.record(
@@ -598,10 +741,12 @@ class DouyinWorker:
     async def _load_storage_state(self) -> dict | None:
         db = await self.get_db()
         try:
-            result = await db.execute(select(Account).where(Account.id == self.account_id))
-            account = result.scalar_one_or_none()
-            if account and account.cookie_data:
-                return json.loads(account.cookie_data)
+            result = await db.execute(
+                select(Account.cookie_data).where(Account.id == self.account_id)
+            )
+            cookie_data = result.scalar_one_or_none()
+            if cookie_data:
+                return json.loads(cookie_data)
         except Exception:
             pass
         finally:
@@ -687,11 +832,15 @@ class DouyinWorker:
         """读取账号专属排队间隔；0/NULL 均表示未设置、继承系统默认。"""
         db = await self.get_db()
         try:
-            result = await db.execute(select(Account).where(Account.id == self.account_id))
-            account = result.scalar_one_or_none()
-            if not account:
+            result = await db.execute(
+                select(Account.reply_delay_seconds).where(
+                    Account.id == self.account_id
+                )
+            )
+            reply_delay = result.scalar_one_or_none()
+            if reply_delay is None:
                 return None
-            value = max(0, int(account.reply_delay_seconds or 0))
+            value = max(0, int(reply_delay or 0))
             return value if value > 0 else None
         finally:
             await db.close()
@@ -722,11 +871,15 @@ class DouyinWorker:
         """读取该账号专属冷却秒数；返回 None 表示继承全局设置。"""
         db = await self.get_db()
         try:
-            result = await db.execute(select(Account).where(Account.id == self.account_id))
-            account = result.scalar_one_or_none()
-            if not account or account.reply_cooldown_seconds is None:
+            result = await db.execute(
+                select(Account.reply_cooldown_seconds).where(
+                    Account.id == self.account_id
+                )
+            )
+            reply_cooldown = result.scalar_one_or_none()
+            if reply_cooldown is None:
                 return None
-            return max(0, int(account.reply_cooldown_seconds))
+            return max(0, int(reply_cooldown))
         finally:
             await db.close()
 
@@ -756,6 +909,10 @@ class DouyinWorker:
 
     async def _run_im_direct_service(self, session: DouyinImSession):
         """运行 IM API + WebSocket 直连自动回复"""
+        # Cache the only account fields needed by the follow-welcome timer.
+        # Disabled accounts subsequently avoid the old full Account query on
+        # every minute tick.
+        await self._refresh_follow_welcome_config(force=True)
         reply_delay = await self.get_reply_delay()
         im_service = DouyinImService(
             session=session,
@@ -770,6 +927,9 @@ class DouyinWorker:
             follow_tick=self.follow_welcome_tick,
             # IM 登录失效（INVALID_REQUEST）时自动下线
             on_session_invalid=self.on_im_session_invalid,
+            # Batch admission waits for UID/frontier/WS/first-poll completion;
+            # it no longer releases its slot immediately after create_task().
+            on_ready=self._mark_startup_ready,
             # 实时解析冷却时间（账号专属优先，否则全局），改设置无需重启托管
             cooldown_resolver=self.resolve_cooldown_seconds,
             # 不在发送链路上自动开浏览器刷新：实测重载页面并不会重生 web_protect，
@@ -1038,15 +1198,24 @@ class DouyinWorker:
                 sender_name=sender_name,
                 sender_id=sender_id,
                 sender_avatar=sender_avatar or None,
-                message_content=message,
-                reply_content=reply,
+                message_content=bound_message_log_content(message),
+                reply_content=(
+                    bound_message_log_content(reply) if reply is not None else None
+                ),
                 status=status,
-                error_message=error,
+                error_message=(
+                    bound_error_log_content(error) if error is not None else None
+                ),
                 created_at=datetime.utcnow()
             )
             db.add(log)
             await db.commit()
-            logger.info(f"Logged message: sender={sender_name}, msg={message}, reply={reply}")
+            logger.debug(
+                "Logged message: sender=%s, msg=%s, reply=%s",
+                sender_name,
+                truncate_text(message, 300),
+                truncate_text(reply, 300) if reply is not None else None,
+            )
         except Exception as e:
             logger.error(f"Failed to log message: {e}")
             await db.rollback()
@@ -1108,23 +1277,41 @@ class DouyinWorker:
             return
         from rpa_engine.douyin_im.follower_poll import fetch_recent_followers
 
-        # sec_user_id 是托管账号的必要身份字段。这个 tick 始终由 IM 主循环调用，
-        # 因此即使关闭了关注欢迎语，也能在运行中发现字段被清空并自动退出托管。
-        sec_user_id = await self._require_sec_user_id("托管运行中")
+        # The direct-service startup normally preloads the lightweight cache.
+        # Keep a guard-first fallback for legacy/tests/partial initialization:
+        # identity safety must not depend on follow-welcome configuration.
+        guarded_sec_user_id = ""
+        if not self._follow_config_loaded:
+            guarded_sec_user_id = await self._require_sec_user_id("托管运行中")
+            if not guarded_sec_user_id:
+                return
+
+        try:
+            enabled, content, sec_user_id = (
+                await self._refresh_follow_welcome_config()
+            )
+        except Exception:
+            if not guarded_sec_user_id:
+                # Even when the optional config read fails, execute the
+                # hosting identity guard before surfacing the transient error.
+                await self._require_sec_user_id("托管运行中")
+            raise
+
+        sec_user_id = str(sec_user_id or guarded_sec_user_id or "").strip()
+        # sec_user_id remains a hosting invariant. The lightweight cached
+        # refresh detects a later database removal without making every
+        # disabled account query the database once per minute.
+        if not sec_user_id:
+            sec_user_id = await self._require_sec_user_id("托管运行中")
         if not sec_user_id:
             return
+        self._follow_welcome_sec_user_id = sec_user_id
+        if not enabled or not content:
+            return
 
-        # 1) 读账号配置 + 已处理过的粉丝集合
+        # 1) 功能已启用时才读取已处理过的粉丝集合
         db = await self.get_db()
         try:
-            acc = (
-                await db.execute(select(Account).where(Account.id == self.account_id))
-            ).scalar_one_or_none()
-            if not acc or not acc.follow_welcome_enabled:
-                return
-            content = (acc.follow_welcome_content or "").strip()
-            if not content:
-                return
             rows = (
                 await db.execute(
                     select(FollowWelcomeLog.follower_uid).where(
@@ -1230,6 +1417,8 @@ class DouyinWorker:
             return
         self.stopping = False
         self.is_running = True
+        self._startup_ready = asyncio.Event()
+        self._startup_error = ""
         task = asyncio.create_task(
             self._run_loop(),
             name=f"douyin-worker-{self.account_id}",
@@ -1253,6 +1442,7 @@ class DouyinWorker:
         """停止 RPA 任务"""
         self.stopping = True
         self.is_running = False
+        self._mark_startup_failed("托管初始化已取消")
         if self._im_service:
             await self._im_service.stop()
         task = self._task
@@ -1281,6 +1471,7 @@ class DouyinWorker:
 
             if self.login_mode == "im_direct":
                 if not storage_state:
+                    self._mark_startup_failed("未保存 Cookie，无法直连 IM")
                     await self.update_account_status(
                         "error",
                         error_msg="未保存 Cookie，无法直连 IM",
@@ -1289,6 +1480,9 @@ class DouyinWorker:
                 started, reason = await self._try_cookie_only_im_start(storage_state)
                 if started:
                     return
+                self._mark_startup_failed(
+                    reason or "凭证验证失败，无法直连 IM"
+                )
                 if self.stopping:
                     return
                 await self.update_account_status(
@@ -1317,11 +1511,13 @@ class DouyinWorker:
 
         except asyncio.CancelledError:
             logger.warning(f"Worker {self.account_id} cancelled")
+            self._mark_startup_failed("托管初始化已取消")
             if not self.stopping:
                 await self.update_account_status("offline", error_msg="RPA 任务已中断，请重新点击启动")
             raise
         except Exception as e:
             logger.exception(f"Error in RPA worker loop: {e}")
+            self._mark_startup_failed(format_error(e))
             if not self.stopping:
                 await self.update_account_status("error", error_msg=format_error(e))
                 system_logger.record(
@@ -1333,6 +1529,7 @@ class DouyinWorker:
                 )
         finally:
             self.is_running = False
+            self._mark_startup_failed()
             if not self.stopping:
                 await self.cleanup()