From 8f68af1c2c299611d52513b57dcae8688eca1a62 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 28 Jul 2026 15:04:17 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 1 + backend/main.py | 370 ++++++++++-- backend/models/db_config.py | 64 +- backend/models/db_migrate.py | 50 ++ backend/models/models.py | 32 +- backend/rpa_engine/credential.py | 5 +- backend/rpa_engine/douyin_im/protocol.py | 89 ++- backend/rpa_engine/douyin_im/service.py | 485 ++++++++++++++- backend/rpa_engine/douyin_im/ws_client.py | 502 +++++++++++----- backend/rpa_engine/playwright_worker.py | 287 +++++++-- backend/tests/test_account_pagination.py | 158 ++++- backend/tests/test_batch_start_api.py | 214 +++++++ .../tests/test_conversation_poll_bandwidth.py | 293 +++++++++- .../tests/test_credential_responsiveness.py | 12 +- backend/tests/test_io_safety.py | 281 +++++++++ backend/tests/test_sec_user_id_guard.py | 17 + backend/tests/test_worker_scale_controls.py | 191 ++++++ backend/tests/test_ws_client_scaling.py | 551 ++++++++++++++++++ backend/utils/log_limits.py | 101 ++++ backend/utils/received_message_log.py | 6 +- backend/utils/system_logger.py | 10 +- frontend/src/views/Dashboard.vue | 9 +- frontend/src/views/Logs.vue | 2 +- frontend/src/views/Messages.vue | 2 +- frontend/src/views/ReceivedMessages.vue | 2 +- frontend/src/views/Rules.vue | 2 +- frontend/src/views/SystemLogs.vue | 2 +- 27 files changed, 3442 insertions(+), 296 deletions(-) create mode 100644 backend/tests/test_io_safety.py create mode 100644 backend/tests/test_worker_scale_controls.py create mode 100644 backend/tests/test_ws_client_scaling.py create mode 100644 backend/utils/log_limits.py diff --git a/.gitignore b/.gitignore index 8cb4f3e..52dfd73 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ backend/**/__pycache__/ backend/kefu.db backend/kefu.db-journal backend/sessions/ +/douyin.zip *.env .env.* diff --git a/backend/main.py b/backend/main.py index 2753d77..566fa50 100644 --- a/backend/main.py +++ b/backend/main.py @@ -107,7 +107,14 @@ class WorkerManager: """Serialize credential preparation across single and batch starts.""" return self._preparation_locks.setdefault(int(account_id), asyncio.Lock()) - async def start_worker(self, account_id: int, login_mode: str = "auto"): + async def start_worker( + self, + account_id: int, + login_mode: str = "auto", + *, + wait_until_ready: bool = False, + credential_prevalidated: bool = False, + ): async with self._account_lock(account_id): if account_id in self.workers: worker = self.workers[account_id] @@ -115,11 +122,60 @@ class WorkerManager: return False del self.workers[account_id] - worker = DouyinWorker(account_id, login_mode=login_mode) + worker = DouyinWorker( + account_id, + login_mode=login_mode, + credential_prevalidated=credential_prevalidated, + ) self.workers[account_id] = worker await worker.start() + + if not wait_until_ready: return True + try: + await worker.wait_until_ready() + except asyncio.CancelledError: + # Batch timeout/cancellation must not leave a detached worker + # continuing to initialize after its queue slot was released. + async with self._account_lock(account_id): + if self.workers.get(account_id) is worker: + try: + await worker.stop() + except Exception: + logger.exception( + "Failed to stop cancelled startup for account %s", + account_id, + ) + finally: + if self.workers.get(account_id) is worker: + self.workers.pop(account_id, None) + raise + except Exception: + # A normal initialization failure marks the account error inside + # the worker. Give that task a short chance to finish its status + # write before removing it; cancelling immediately would overwrite + # the useful error with an offline state. + task = getattr(worker, "_task", None) + if task and not task.done(): + try: + await asyncio.wait_for(asyncio.shield(task), timeout=5.0) + except (asyncio.TimeoutError, asyncio.CancelledError): + try: + await worker.stop() + except Exception: + logger.exception( + "Failed to stop unsuccessful startup for account %s", + account_id, + ) + except Exception: + pass + async with self._account_lock(account_id): + if self.workers.get(account_id) is worker: + self.workers.pop(account_id, None) + raise + return True + async def stop_worker(self, account_id: int): async with self._account_lock(account_id): if account_id in self.workers: @@ -331,17 +387,29 @@ async def _flush_system_logs_loop(): async def _sync_legacy_cookie_files(): """将旧版仅保存在文件的 Cookie 同步到数据库""" async with AsyncSessionLocal() as db: - result = await db.execute(select(Account)) - accounts = result.scalars().all() + # Only legacy rows without a database Cookie need filesystem work. + # Selecting full Account entities used to hydrate every large + # cookie_data / im_session_data value on each process start, which is + # especially expensive with hundreds of hosted accounts. + result = await db.execute( + select(Account.id).where( + Account.cookie_data.is_(None) | (Account.cookie_data == "") + ) + ) + account_ids = list(result.scalars().all()) changed = False - for acc in accounts: - if acc.cookie_data: - continue - file_data = read_cookie_file(acc.id) + for account_id in account_ids: + file_data = read_cookie_file(int(account_id)) if file_data: - acc.cookie_data = file_data - acc.cookie_path = get_cookie_path(acc.id) - acc.cookie_updated_at = datetime.utcnow() + await db.execute( + update(Account) + .where(Account.id == int(account_id)) + .values( + cookie_data=file_data, + cookie_path=get_cookie_path(int(account_id)), + cookie_updated_at=datetime.utcnow(), + ) + ) changed = True if changed: await db.commit() @@ -501,10 +569,8 @@ async def _seed_admin_user(): # 初始化数据库 @app.on_event("startup") async def startup(): - # 放大默认线程池：每个托管账号的 a_bogus/web_protect/ts_sign 签名都是阻塞的 Node - # 子进程调用，经 asyncio.to_thread 跑在默认线程池里。Python 默认池大小仅 - # min(32, cpu+4)，在 1 核云服务器上只有 5 个线程，导致超过 5 个账号并发时第 6 个 - # 账号的签名/取信息/发送会一直排队阻塞直至超时失败。这里显式放大线程池消除该瓶颈。 + # a_bogus/web_protect/ts_sign 仍需在线程池中执行阻塞的 Node 调用，但账号启动、 + # 后台请求和发送通道都已有独立并发限制，因此线程池按 CPU 有界配置即可。 import concurrent.futures as _futures try: @@ -512,7 +578,11 @@ async def startup(): except ValueError: _pool_size = 0 if _pool_size <= 0: - _pool_size = max(64, ((os.cpu_count() or 1) * 8)) + # WebSocket connections are fully asynchronous now. The executor is + # only for short signing / compatibility calls, whose network lanes + # are already bounded. Keeping 64 threads on a 2-core host increases + # context switching and swap pressure without adding throughput. + _pool_size = max(8, min(32, (os.cpu_count() or 1) * 4)) loop = asyncio.get_running_loop() loop.set_default_executor( _futures.ThreadPoolExecutor( @@ -557,12 +627,95 @@ async def startup(): @app.on_event("shutdown") async def shutdown(): + # Stop accounts concurrently with global deadlines. Sequentially waiting + # for hundreds of WebSocket close handshakes can otherwise turn a normal + # deployment restart into a many-minute outage. + try: + stop_concurrency = max( + 1, + min(64, int(os.getenv("KEFU_SHUTDOWN_CONCURRENCY", "32") or 32)), + ) + except ValueError: + stop_concurrency = 32 + try: + shutdown_timeout = max( + 5.0, + min( + 180.0, + float(os.getenv("KEFU_SHUTDOWN_TIMEOUT_SECONDS", "60") or 60), + ), + ) + except ValueError: + shutdown_timeout = 60.0 + + try: + batch_stop_timeout = max( + 1.0, + min( + 30.0, + float(os.getenv("KEFU_BATCH_STOP_TIMEOUT_SECONDS", "10") or 10), + ), + ) + except ValueError: + batch_stop_timeout = 10.0 + # Stop queued preparations first so no new workers appear while the - # existing workers are being drained below. - await batch_start_queue.stop() - # 停止所有正在运行的 RPA 任务 - for account_id in list(manager.workers.keys()): - await manager.stop_worker(account_id) + # existing workers are being drained below. Its cancellation path may + # itself wait for a half-open login/DB operation, so it needs an + # independent deadline; the remaining workers are still covered by the + # bounded parallel stop below. + try: + await asyncio.wait_for( + batch_start_queue.stop(), + timeout=batch_stop_timeout, + ) + except asyncio.TimeoutError: + logger.warning( + "Batch-start queue shutdown exceeded %.1fs; continuing with worker drain", + batch_stop_timeout, + ) + + stop_gate = asyncio.Semaphore(stop_concurrency) + + async def _stop_account(account_id: int) -> None: + async with stop_gate: + try: + await manager.stop_worker(account_id) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Failed to stop account %s during shutdown", account_id) + + stop_tasks = [ + asyncio.create_task( + _stop_account(account_id), + name=f"shutdown-account-{account_id}", + ) + for account_id in list(manager.workers.keys()) + ] + if stop_tasks: + try: + await asyncio.wait_for( + asyncio.gather(*stop_tasks), + timeout=shutdown_timeout, + ) + except asyncio.TimeoutError: + logger.warning( + "Account shutdown exceeded %.1fs; cancelling remaining tasks", + shutdown_timeout, + ) + for task in stop_tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*stop_tasks, return_exceptions=True) + from rpa_engine.douyin_im.service import _shutdown_initial_unread_dispatcher + try: + await asyncio.wait_for( + _shutdown_initial_unread_dispatcher(), + timeout=5.0, + ) + except asyncio.TimeoutError: + logger.warning("Initial-unread dispatcher shutdown exceeded 5s") from rpa_engine.douyin_im.traffic_control import shutdown_traffic_controller await shutdown_traffic_controller() if _system_log_flush_task: @@ -625,6 +778,25 @@ class AccountResponse(BaseModel): from_attributes = True +class AccountOptionResponse(BaseModel): + """Small account payload used by selectors on non-account pages. + + Keeping this separate from ``AccountResponse`` prevents account dropdowns + from loading and parsing every account's Cookie and IM session blobs. + """ + + id: int + username: Optional[str] = None + avatar_url: Optional[str] = None + douyin_uid: Optional[str] = None + phone: Optional[str] = None + status: str + has_cookie: bool = False + reply_cooldown_seconds: Optional[int] = None + reply_cooldown_effective: int = 0 + quota_disabled: bool = False + + class DashboardAccountStatsResponse(BaseModel): """Safe account totals shown to every authenticated dashboard user.""" @@ -1177,6 +1349,83 @@ async def get_accounts( "page_size": page_size, } + +@app.get("/api/account-options", response_model=List[AccountOptionResponse]) +async def get_account_options( + db: AsyncSession = Depends(get_db), + user: User = Depends(get_current_user), +): + """Return only fields required by account dropdowns. + + The legacy unpaginated ``/api/accounts`` endpoint hydrates large Cookie and + IM session columns and then parses Cookie JSON for every account. With + hundreds of accounts that makes simply opening Messages, Rules or Logs + expensive. This query deliberately selects only short scalar fields. + """ + + has_db_cookie = case( + ( + Account.cookie_data.is_not(None) + & (Account.cookie_data != ""), + True, + ), + else_=False, + ).label("has_db_cookie") + stmt = select( + Account.id, + Account.username, + Account.avatar_url, + Account.douyin_uid, + Account.phone, + Account.status, + Account.cookie_path, + has_db_cookie, + Account.reply_cooldown_seconds, + Account.quota_disabled, + ) + if not is_admin(user.role): + stmt = stmt.where(Account.owner_id == user.id) + + rows = (await db.execute(stmt.order_by(Account.id.asc()))).all() + global_cooldown = _global_cooldown_seconds() + options: list[AccountOptionResponse] = [] + for row in rows: + account_id = int(row.id) + is_running = manager.is_running(account_id) + runtime_status = str(row.status or "offline") + if is_running and runtime_status == "offline": + runtime_status = "online" + elif not is_running and runtime_status in ("online", "logging_in", "starting"): + runtime_status = "offline" + + has_cookie = bool(row.has_db_cookie) + if not has_cookie and row.cookie_path: + has_cookie = os.path.exists(str(row.cookie_path)) + + cooldown_override = row.reply_cooldown_seconds + options.append( + AccountOptionResponse( + id=account_id, + username=row.username, + avatar_url=row.avatar_url, + douyin_uid=row.douyin_uid, + phone=row.phone, + status=runtime_status, + has_cookie=has_cookie, + reply_cooldown_seconds=( + int(cooldown_override) if cooldown_override is not None else None + ), + reply_cooldown_effective=( + max(0, int(cooldown_override)) + if cooldown_override is not None + else global_cooldown + ), + quota_disabled=bool(row.quota_disabled), + ) + ) + return options + + @app.get("/api/accounts/{account_id}", response_model=AccountResponse) async def get_account( account_id: int, @@ -1356,6 +1605,10 @@ async def update_account( user: User = Depends(require_write), ): account = await get_owned_account(db, user, account_id, write=True) + follow_config_changed = bool( + {"follow_welcome_enabled", "follow_welcome_content"} + & set(body.model_fields_set) + ) if body.phone is not None: account.phone = body.phone @@ -1379,6 +1632,11 @@ async def update_account( account.updated_at = datetime.utcnow() await db.commit() await db.refresh(account) + if follow_config_changed: + worker = manager.workers.get(account_id) + invalidate = getattr(worker, "invalidate_follow_welcome_config", None) + if callable(invalidate): + invalidate() return _build_account_response(account) @@ -1820,6 +2078,8 @@ async def _start_account_rpa_impl( account: Account, db: AsyncSession, requested_login_mode: Optional[str] = None, + *, + wait_for_ready: bool = False, ) -> dict: """Validate once, persist the starting state, then spawn one worker.""" account_id = int(account.id) @@ -1853,10 +2113,19 @@ async def _start_account_rpa_impl( detail=assessment["message"] or "凭证未通过验证，无法直连 IM", ) + if wait_for_ready and login_mode != "im_direct": + # A bulk operation cannot complete an interactive QR/browser login. + # Launching hundreds of browser tasks would only move the backlog out + # of the queue and recreate the original server stall. The single + # account start endpoint remains unchanged for interactive login. + raise RuntimeError( + "该账号需要手动浏览器登录，已跳过批量启动，请单独启动" + ) + # Write before spawning the task. This gives both single and batch calls # an immediate authoritative status and avoids racing the worker's first - # database update. Batch submission normally wrote this state in bulk, so - # no per-account commit is needed in that path. + # database update. Batch readiness keeps the number of accounts reaching + # this commit bounded instead of letting the whole batch write at once. state_changed = ( account.status != "starting" or account.qr_code_base64 is not None @@ -1869,7 +2138,14 @@ async def _start_account_rpa_impl( await db.commit() try: - started = await manager.start_worker(account_id, login_mode=login_mode) + started = await manager.start_worker( + account_id, + login_mode=login_mode, + wait_until_ready=wait_for_ready, + credential_prevalidated=bool( + login_mode == "im_direct" and assessment["can_skip_browser"] + ), + ) except Exception as exc: account.status = "error" account.error_message = str(exc) or "启动托管失败" @@ -1877,7 +2153,9 @@ async def _start_account_rpa_impl( raise if started: - if login_mode == "im_direct": + if wait_for_ready: + msg = "IM 托管已完成初始化" + elif login_mode == "im_direct": msg = assessment["message"] or "凭证有效，正在直连 IM 托管（无需浏览器）" elif reset_performed: msg = "凭证已失效，已清除旧数据，正在打开浏览器重新登录..." @@ -1887,7 +2165,7 @@ async def _start_account_rpa_impl( msg = "未保存 Cookie，将打开浏览器扫码登录..." return { - "status": "starting", + "status": "running" if wait_for_ready else "starting", "login_mode": login_mode, "cookie_valid": assessment["cookie_valid"], "im_ready": assessment["im_ready"], @@ -1909,7 +2187,11 @@ async def _start_queued_account(account_id: int) -> dict: if account.quota_disabled: raise RuntimeError("账号已停用，无法启动托管") try: - return await _start_account_rpa_impl(account, db) + return await _start_account_rpa_impl( + account, + db, + wait_for_ready=True, + ) except asyncio.CancelledError: raise except Exception as exc: @@ -2217,21 +2499,27 @@ async def get_logs_stats( """消息日志全量统计（数据库计数，不受列表 limit 限制）。""" if account_id is not None: await get_owned_account(db, user, account_id) - base = logs_for_user(user, account_id).subquery() - total = int( - (await db.execute(select(func.count()).select_from(base))).scalar() or 0 + # Select only the indexed status column and calculate both counters in one + # scan. The old implementation queried the growing log table twice on + # every dashboard refresh. + base = ( + logs_for_user(user, account_id) + .with_only_columns(MessageLog.status) + .order_by(None) + .subquery() ) - replied = int( - ( - await db.execute( - select(func.count()) - .select_from(base) - .where(base.c.status == "replied") - ) - ).scalar() - or 0 - ) - return {"total": total, "replied": replied} + row = ( + await db.execute( + select( + func.count().label("total"), + func.coalesce( + func.sum(case((base.c.status == "replied", 1), else_=0)), + 0, + ).label("replied"), + ).select_from(base) + ) + ).one() + return {"total": int(row.total or 0), "replied": int(row.replied or 0)} @app.get("/api/logs", response_model=List[LogResponse]) diff --git a/backend/models/db_config.py b/backend/models/db_config.py index dc2c308..af803be 100644 --- a/backend/models/db_config.py +++ b/backend/models/db_config.py @@ -9,8 +9,10 @@ from typing import Any from urllib.parse import quote_plus from dotenv import load_dotenv -from sqlalchemy import text +from sqlalchemy import event, text from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from sqlalchemy.engine import make_url +from sqlalchemy.pool import AsyncAdaptedQueuePool, StaticPool BACKEND_DIR = Path(__file__).resolve().parent.parent PROJECT_ROOT = BACKEND_DIR.parent @@ -37,6 +39,10 @@ def _env_int(name: str, default: int) -> int: return default +def _bounded_env_int(name: str, default: int, minimum: int, maximum: int) -> int: + return max(minimum, min(maximum, _env_int(name, default))) + + def normalize_db_type(value: str | None) -> str: raw = (value or "sqlite").strip().lower() if raw in ("postgres", "pgsql"): @@ -136,7 +142,35 @@ def build_database_url(config: DatabaseConfig | None = None) -> str: def engine_kwargs_for_url(url: str) -> dict[str, Any]: kwargs: dict[str, Any] = {"echo": False} - if not url.startswith("sqlite"): + if url.startswith("sqlite"): + # SQLAlchemy 2.0.30 defaults file-backed aiosqlite to NullPool. At + # hundreds of hosted accounts that creates and tears down an aiosqlite + # worker thread for every short query. Reuse a small bounded pool; + # SQLite still serializes writers, so a large pool only adds lock + # contention and does not improve throughput. + busy_timeout_ms = _bounded_env_int( + "KEFU_SQLITE_BUSY_TIMEOUT_MS", 30_000, 1_000, 120_000 + ) + kwargs["connect_args"] = {"timeout": busy_timeout_ms / 1000.0} + sqlite_database = make_url(url).database + is_memory_database = ( + not sqlite_database + or sqlite_database == ":memory:" + or "mode=memory" in url.lower() + ) + if is_memory_database: + # Every connection to an in-memory SQLite URL otherwise receives a + # different database. StaticPool preserves the single shared + # connection expected by tests and utility callers. + kwargs["poolclass"] = StaticPool + else: + kwargs["poolclass"] = AsyncAdaptedQueuePool + kwargs["pool_size"] = _bounded_env_int( + "KEFU_SQLITE_POOL_SIZE", 5, 1, 10 + ) + kwargs["max_overflow"] = 0 + kwargs["pool_timeout"] = max(5.0, busy_timeout_ms / 1000.0) + else: kwargs["pool_pre_ping"] = True kwargs["pool_recycle"] = 3600 # 默认连接池仅 pool_size=5 + max_overflow=10。多账号托管 + 前端并发请求时 @@ -155,9 +189,33 @@ def engine_kwargs_for_url(url: str) -> dict[str, Any]: return kwargs +def _configure_sqlite_connection(dbapi_connection, _connection_record) -> None: + """Apply process-wide SQLite settings to every pooled connection. + + WAL lets readers continue while the single writer commits. NORMAL avoids + a full disk sync for every small log/status transaction while retaining + WAL crash consistency. busy_timeout turns transient writer contention + into bounded waiting instead of immediate ``database is locked`` errors. + """ + busy_timeout_ms = _bounded_env_int( + "KEFU_SQLITE_BUSY_TIMEOUT_MS", 30_000, 1_000, 120_000 + ) + cursor = dbapi_connection.cursor() + try: + cursor.execute(f"PRAGMA busy_timeout={busy_timeout_ms}") + cursor.execute("PRAGMA journal_mode=WAL") + cursor.execute("PRAGMA synchronous=NORMAL") + cursor.execute("PRAGMA foreign_keys=ON") + finally: + cursor.close() + + def create_database_engine(config: DatabaseConfig | None = None) -> AsyncEngine: url = build_database_url(config) - return create_async_engine(url, **engine_kwargs_for_url(url)) + engine = create_async_engine(url, **engine_kwargs_for_url(url)) + if url.startswith("sqlite"): + event.listen(engine.sync_engine, "connect", _configure_sqlite_connection) + return engine def mask_database_url(url: str) -> str: diff --git a/backend/models/db_migrate.py b/backend/models/db_migrate.py index 4c2eb81..34a66e8 100644 --- a/backend/models/db_migrate.py +++ b/backend/models/db_migrate.py @@ -33,6 +33,26 @@ def add_column_if_missing(conn, table: str, column: str, ddl_by_dialect: dict[st conn.execute(text(ddl)) +def add_index_if_missing( + conn, + table: str, + index_name: str, + columns: tuple[str, ...], +) -> None: + """Create one portable index without relying on dialect-specific IF NOT EXISTS.""" + try: + insp = inspect(conn) + if not insp.has_table(table): + return + existing = {item.get("name") for item in insp.get_indexes(table)} + except Exception: + return + if index_name in existing: + return + safe_columns = ", ".join(columns) + conn.execute(text(f"CREATE INDEX {index_name} ON {table} ({safe_columns})")) + + def migrate_accounts_table(conn) -> None: add_column_if_missing( conn, @@ -124,6 +144,36 @@ def migrate_message_logs_table(conn) -> None: "sender_avatar", {"default": "ALTER TABLE message_logs ADD COLUMN sender_avatar TEXT"}, ) + add_index_if_missing( + conn, + "message_logs", + "ix_message_logs_account_created_at", + ("account_id", "created_at"), + ) + add_index_if_missing( + conn, + "message_logs", + "ix_message_logs_created_at", + ("created_at",), + ) + add_index_if_missing( + conn, + "message_logs", + "ix_message_logs_status_account_id", + ("status", "account_id"), + ) + add_index_if_missing( + conn, + "received_message_logs", + "ix_received_message_logs_account_created_at", + ("account_id", "created_at"), + ) + add_index_if_missing( + conn, + "system_logs", + "ix_system_logs_account_created_at", + ("account_id", "created_at"), + ) def migrate_rules_table(conn) -> None: diff --git a/backend/models/models.py b/backend/models/models.py index 36a2e36..793c1a3 100644 --- a/backend/models/models.py +++ b/backend/models/models.py @@ -1,7 +1,8 @@ from datetime import datetime -from sqlalchemy import Column, Integer, String, Boolean, DateTime, ForeignKey, Text, UniqueConstraint -from sqlalchemy.orm import relationship +from sqlalchemy import Column, Integer, String, Boolean, DateTime, ForeignKey, Index, Text, UniqueConstraint +from sqlalchemy.orm import relationship, validates from .database import Base +from utils.log_limits import bound_error_log_content, bound_message_log_content class User(Base): @@ -204,10 +205,23 @@ class MessageLog(Base): reply_content = Column(Text, nullable=True) # 回复的消息 status = Column(String(50), default="received") # received, replied, ignored, failed error_message = Column(Text, nullable=True) - created_at = Column(DateTime, default=datetime.utcnow) + created_at = Column(DateTime, default=datetime.utcnow, index=True) account = relationship("Account", back_populates="logs") + __table_args__ = ( + Index("ix_message_logs_account_created_at", "account_id", "created_at"), + Index("ix_message_logs_status_account_id", "status", "account_id"), + ) + + @validates("message_content", "reply_content") + def _bound_message_content(self, _key, value): + return bound_message_log_content(value) if value is not None else None + + @validates("error_message") + def _bound_error_content(self, _key, value): + return bound_error_log_content(value) if value is not None else None + class ReceivedMessageLog(Base): """接收消息原始日志：仅记录收到的消息，内容原样保存。""" @@ -225,6 +239,14 @@ class ReceivedMessageLog(Base): raw_content = Column(Text, nullable=False) created_at = Column(DateTime, default=datetime.utcnow, index=True) + __table_args__ = ( + Index( + "ix_received_message_logs_account_created_at", + "account_id", + "created_at", + ), + ) + class FollowWelcomeLog(Base): """关注欢迎语去重表：每个账号对每个新粉丝只发送一次欢迎语（重启后仍生效）。""" @@ -255,6 +277,10 @@ class SystemLog(Base): detail = Column(Text, nullable=True) # 详细原因 created_at = Column(DateTime, default=datetime.utcnow, index=True) + __table_args__ = ( + Index("ix_system_logs_account_created_at", "account_id", "created_at"), + ) + class PaymentOrder(Base): """账号额度购买订单。""" diff --git a/backend/rpa_engine/credential.py b/backend/rpa_engine/credential.py index 191783b..62674db 100644 --- a/backend/rpa_engine/credential.py +++ b/backend/rpa_engine/credential.py @@ -5,7 +5,6 @@ import logging from typing import Optional from rpa_engine.douyin_im.auth import DouyinAuth -from rpa_engine.douyin_im.frontier import ensure_frontier_ws from rpa_engine.douyin_im.session import DouyinImSession from utils.cookie_store import analyze_cookie @@ -156,7 +155,9 @@ async def validate_im_session( return False, "缺少 sessionid，无法直连 IM" return False, "Cookie 不满足 IM 直连条件" - await asyncio.to_thread(ensure_frontier_ws, session) + # Frontier discovery belongs to the worker startup lifecycle. Running it + # here populated only this temporary assessment session, so a bulk start + # immediately repeated the same signing / query work for every account. try: auth = DouyinAuth.from_im_session(session) # 优先用已持久化的 my_uid，避免每次都发起网络 query_my_uid（uid_tt 是加密串， diff --git a/backend/rpa_engine/douyin_im/protocol.py b/backend/rpa_engine/douyin_im/protocol.py index 6e255c3..0cbcd31 100644 --- a/backend/rpa_engine/douyin_im/protocol.py +++ b/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 ( diff --git a/backend/rpa_engine/douyin_im/service.py b/backend/rpa_engine/douyin_im/service.py index 70e815c..4a4a4f7 100644 --- a/backend/rpa_engine/douyin_im/service.py +++ b/backend/rpa_engine/douyin_im/service.py @@ -1,7 +1,11 @@ import asyncio +import inspect import logging import os import time +import weakref +from collections import deque +from dataclasses import dataclass, field from typing import Awaitable, Callable, Optional from utils import system_logger @@ -25,6 +29,14 @@ logger = logging.getLogger("douyin_im.service") MatchReplyFn = Callable[[str], Awaitable[Optional[list[str]]]] LogFn = Callable[..., Awaitable[None]] ReceivedLogFn = Callable[..., Awaitable[None]] +ReadyFn = Callable[[], Optional[Awaitable[None]]] + + +# These sets live on the FastAPI event-loop thread. They let every account +# derive a polling interval from the current population instead of assuming it +# is the only hosted account on the server. +_active_service_ids: set[int] = set() +_ws_service_ids: set[int] = set() def _env_poll_seconds(name: str, default: float, minimum: float = 5.0) -> float: @@ -34,7 +46,11 @@ def _env_poll_seconds(name: str, default: float, minimum: float = 5.0) -> float: return default -def _conversation_poll_timing(account_id: int, has_ws: bool) -> tuple[float, float]: +def _conversation_poll_timing( + account_id: int, + has_ws: bool, + population: Optional[int] = None, +) -> tuple[float, float]: """Return the reconciliation interval and a stable per-account stagger. WebSocket is the real-time receive path. HTTP polling is only a safety @@ -42,15 +58,253 @@ def _conversation_poll_timing(account_id: int, has_ws: bool) -> tuple[float, flo hundreds of accounts wastes bandwidth and eventually starves new starts. Accounts without WebSocket keep the original fast polling cadence. """ - interval = _env_poll_seconds( + base_interval = _env_poll_seconds( "KEFU_WS_RECONCILE_INTERVAL_SECONDS" if has_ws else "KEFU_HTTP_POLL_INTERVAL_SECONDS", 120.0 if has_ws else 15.0, ) + if population is None: + if has_ws: + population = len(_ws_service_ids) + else: + population = len(_active_service_ids - _ws_service_ids) + population = max(1, int(population or 0)) + # Reserve a bounded request budget for each class. With 500 connected + # accounts and the default 1 req/s budget, reconciliation automatically + # stretches to 500 seconds instead of permanently saturating the two + # shared background-network slots. Disconnected accounts use a separate + # budget so they cannot create a 15-second retry storm after an outage. + budget_name = "KEFU_WS_POLL_BUDGET_RPS" if has_ws else "KEFU_HTTP_POLL_BUDGET_RPS" + try: + budget_rps = max(0.05, float(os.getenv(budget_name, "1.0"))) + except (TypeError, ValueError): + budget_rps = 1.0 + interval = max(base_interval, population / budget_rps) spread_ms = max(1, int(interval * 1000)) stagger = ((int(account_id or 0) * 2654435761) % spread_ms) / 1000.0 return interval, stagger +def _initial_unread_concurrency() -> int: + try: + return max( + 1, + min(8, int(os.getenv("KEFU_INITIAL_UNREAD_CONCURRENCY", "2"))), + ) + except (TypeError, ValueError): + return 2 + + +@dataclass(eq=False) +class _InitialUnreadJob: + service: "DouyinImService" + messages: deque[dict] = field(default_factory=deque) + cancelled: bool = False + active_task: Optional[asyncio.Task] = None + + +class _InitialUnreadDispatcher: + """Process startup unread snapshots with a fixed process-wide worker set. + + One queued object represents one account and is requeued after each + message. This keeps account-local FIFO while preventing a single account + from monopolizing both consumers. Only active consumers create handler + tasks, so a 500-account start does not create 500 waiting tasks. + """ + + def __init__(self, concurrency: Optional[int] = None) -> None: + self.concurrency = max( + 1, + int(concurrency or _initial_unread_concurrency()), + ) + self._queue: asyncio.Queue[_InitialUnreadJob] = asyncio.Queue() + self._jobs: dict[int, _InitialUnreadJob] = {} + self._workers: list[asyncio.Task] = [] + self._lock = asyncio.Lock() + self._stopping = False + + async def _ensure_workers(self) -> None: + async with self._lock: + self._workers = [task for task in self._workers if not task.done()] + if self._workers or self._stopping: + return + for index in range(self.concurrency): + self._workers.append( + asyncio.create_task( + self._worker(index + 1), + name=f"initial-unread-worker-{index + 1}", + ) + ) + + async def submit( + self, + service: "DouyinImService", + messages: list[dict], + ) -> None: + pending = [message for message in messages if isinstance(message, dict)] + if not pending or not service._running: + return + await self._ensure_workers() + async with self._lock: + if self._stopping or not service._running: + return + service_key = id(service) + existing = self._jobs.get(service_key) + if existing and not existing.cancelled: + existing.messages.extend(pending) + return + job = _InitialUnreadJob( + service=service, + messages=deque(pending), + ) + self._jobs[service_key] = job + self._queue.put_nowait(job) + logger.debug( + "Initial unread queued account=%s count=%s", + service.account_id, + len(pending), + ) + + async def cancel(self, service: "DouyinImService") -> None: + operation: Optional[asyncio.Task] = None + current_task = asyncio.current_task() + async with self._lock: + job = self._jobs.pop(id(service), None) + if job is None: + return + job.cancelled = True + job.messages.clear() + operation = job.active_task + if ( + operation + and operation is not current_task + and not operation.done() + ): + operation.cancel() + if ( + operation + and operation is not current_task + and not operation.done() + ): + await asyncio.gather(operation, return_exceptions=True) + + async def _worker(self, worker_number: int) -> None: + while True: + job = await self._queue.get() + operation: Optional[asyncio.Task] = None + try: + async with self._lock: + if ( + job.cancelled + or not job.service._running + or not job.messages + ): + self._jobs.pop(id(job.service), None) + else: + message = job.messages.popleft() + operation = asyncio.create_task( + job.service._handle_incoming(message), + name=( + "initial-unread-account-" + f"{job.service.account_id}" + ), + ) + job.active_task = operation + + if operation is not None: + try: + await operation + except asyncio.CancelledError: + # Cancelling one account cancels only its bounded child + # operation; cancelling this worker still shuts it down. + if asyncio.current_task().cancelling(): + raise + except Exception as exc: + logger.error( + "Initial unread handling failed account=%s " + "worker=%s: %s", + job.service.account_id, + worker_number, + exc, + ) + + async with self._lock: + if job.active_task is operation: + job.active_task = None + if ( + job.cancelled + or not job.service._running + or not job.messages + ): + self._jobs.pop(id(job.service), None) + else: + # Round-robin across accounts while retaining FIFO + # within this account's initial unread snapshot. + self._queue.put_nowait(job) + except asyncio.CancelledError: + if operation and not operation.done(): + operation.cancel() + await asyncio.gather(operation, return_exceptions=True) + raise + finally: + self._queue.task_done() + + async def join(self) -> None: + await self._queue.join() + + async def stop(self) -> None: + async with self._lock: + if self._stopping: + workers = list(self._workers) + operations = [] + else: + self._stopping = True + workers = list(self._workers) + operations = [ + job.active_task + for job in self._jobs.values() + if job.active_task and not job.active_task.done() + ] + for job in self._jobs.values(): + job.cancelled = True + job.messages.clear() + self._jobs.clear() + for operation in operations: + operation.cancel() + for worker in workers: + worker.cancel() + if operations: + await asyncio.gather(*operations, return_exceptions=True) + if workers: + await asyncio.gather(*workers, return_exceptions=True) + while True: + try: + self._queue.get_nowait() + except asyncio.QueueEmpty: + break + else: + self._queue.task_done() + self._workers.clear() + + +_INITIAL_UNREAD_DISPATCHERS = weakref.WeakKeyDictionary() + + +def _get_initial_unread_dispatcher() -> _InitialUnreadDispatcher: + loop = asyncio.get_running_loop() + dispatcher = _INITIAL_UNREAD_DISPATCHERS.get(loop) + if dispatcher is None: + dispatcher = _InitialUnreadDispatcher() + _INITIAL_UNREAD_DISPATCHERS[loop] = dispatcher + return dispatcher + + +async def _shutdown_initial_unread_dispatcher() -> None: + loop = asyncio.get_running_loop() + dispatcher = _INITIAL_UNREAD_DISPATCHERS.pop(loop, None) + if dispatcher is not None: + await dispatcher.stop() + + class DouyinImService: """抖音 IM 直连服务：WebSocket 实时监听 + HTTP 轮询 + 自动回复""" @@ -68,6 +322,7 @@ class DouyinImService: refresh_credentials: Optional[Callable[[], Awaitable[bool]]] = None, follow_tick: Optional[Callable[[], Awaitable[None]]] = None, on_session_invalid: Optional[Callable[[str], Awaitable[None]]] = None, + on_ready: Optional[ReadyFn] = None, ): self.session = session self.match_reply = match_reply @@ -78,6 +333,8 @@ class DouyinImService: self.follow_tick = follow_tick # 由 worker 注入：检测到 IM 登录失效（INVALID_REQUEST）时回调，用于自动下线 self.on_session_invalid = on_session_invalid + self._on_ready = on_ready + self._ready_notified = False self._session_invalid_strikes = 0 self._session_invalid_fired = False self.reply_delay_seconds = max(0, int(reply_delay_seconds or 0)) @@ -109,6 +366,16 @@ class DouyinImService: self._ws_client: Optional[DouyinImWsClient] = None self.last_error: str = "" + async def _notify_ready(self) -> None: + if self._ready_notified: + return + self._ready_notified = True + if not self._on_ready: + return + result = self._on_ready() + if inspect.isawaitable(result): + await result + def _reply_key(self, conversation_key: str, content: str) -> str: return f"{conversation_key}::{content}" @@ -725,7 +992,12 @@ class DouyinImService: f"(sent={sent_any}, count={len(replies)})" ) - async def _index_conversations(self, conversations: list[dict]): + async def _index_conversations( + self, + conversations: list[dict], + *, + enrich_profiles: bool = True, + ): my_uid = int(self.session.my_uid or 0) for raw in conversations: conv = enrich_conversation_item(raw, my_uid) @@ -734,7 +1006,11 @@ class DouyinImService: avatar = str(conv.get("sender_avatar") or "").strip() peer_uid = str(conv.get("peer_uid") or "") - if peer_uid and (is_generic_peer_name(name, peer_uid) or not avatar): + if ( + enrich_profiles + and peer_uid + and (is_generic_peer_name(name, peer_uid) or not avatar) + ): profile = await fetch_peer_profile(self.session, peer_uid, self.account_id) if profile.get("nickname"): name = profile["nickname"] @@ -755,16 +1031,67 @@ class DouyinImService: if peer_uid and name: self._conv_names[peer_uid] = name - async def _poll_conversations(self): + async def _poll_conversations( + self, + *, + initial: bool = False, + defer_handlers: bool = False, + ) -> list[dict]: controller = get_traffic_controller() - async with controller.background_slot(self.account_id, "conversation poll"): - async with DouyinImHttpClient(self.session, account_id=self.account_id) as http: + async with controller.background_slot( + self.account_id, + "conversation poll", + startup=initial, + ): + # Do not keep one idle connection pool per hosted account. At 500 + # accounts that would retain hundreds of sockets between sparse + # reconciliations. Capacity-aware scheduling makes construction + # infrequent, while the context manager releases the socket as + # soon as this account's poll finishes. + async with DouyinImHttpClient( + self.session, + account_id=self.account_id, + ) as http: conversations = await http.get_conversations(enrich_profiles=False) - # Profile enrichment may involve several slow third-party requests. - # Run it after releasing the conversation-list slot; each individual - # lookup re-enters the shared controller and yields fairly to startup - # validation and other accounts between profiles. - await self._index_conversations(conversations) + + # Capture the previous preview before _index_conversations overwrites + # _conv_meta. A conversation-list preview is not inherently a new + # message: after startup we only act when it is unread or has actually + # changed since the last cache snapshot. This prevents the first + # scheduled reconciliation from replying to every historical preview. + previous_previews: list[tuple[bool, str]] = [] + previous_by_peer = { + str(meta.get("peer_uid") or ""): str(meta.get("content") or "") + for meta in self._conv_meta.values() + if str(meta.get("peer_uid") or "") + } + for conv in conversations: + conv_id = str(conv.get("conversation_id") or "").strip() + peer_uid = str( + conv.get("peer_uid") or conv.get("sender_id") or "" + ).strip() + sender_name = str(conv.get("sender_name") or "").strip() + prior: Optional[str] = None + known = False + if conv_id and conv_id in self._conv_meta: + known = True + prior = str(self._conv_meta[conv_id].get("content") or "") + elif sender_name and sender_name in self._conv_previews: + known = True + prior = str(self._conv_previews.get(sender_name) or "") + elif peer_uid and peer_uid in previous_by_peer: + known = True + prior = previous_by_peer[peer_uid] + previous_previews.append((known, prior or "")) + + # Never enrich every row in a reconciliation snapshot. At 500 accounts + # that could turn one list request into tens of thousands of profile + # calls. _handle_incoming resolves the peer lazily only for an unread + # or genuinely changed conversation after this lightweight cache pass. + await self._index_conversations( + conversations, + enrich_profiles=False, + ) unread_total = sum( max(0, int(item.get("unread_count") or 0)) for item in conversations @@ -773,10 +1100,27 @@ class DouyinImService: logger.info(f"IM unread total: {unread_total}") # Message handling may wait in the global send lane. Do not keep one # of the scarce background HTTP slots occupied while that happens. - for conv in conversations: + deferred: list[dict] = [] + for conv, (preview_known, previous_preview) in zip( + conversations, + previous_previews, + ): unread = int(conv.get("unread_count") or 0) - if unread > 0 or conv.get("content"): - await self._handle_incoming(conv) + current_preview = str(conv.get("content") or "") + preview_changed = bool( + preview_known + and current_preview + and current_preview != previous_preview + ) + # Existing previews are useful cache data, but on the authoritative + # startup snapshot they are not proof of a newly received message. + # The same is true on later polls unless the cached preview changed. + if unread > 0 or (not initial and preview_changed): + if defer_handlers: + deferred.append(conv) + else: + await self._handle_incoming(conv) + return deferred async def _verify_account_uid(self): """启动时用 query/user 接口核验账号真实 UID，修正采集端可能取错的 my_uid/device_id。 @@ -790,7 +1134,11 @@ class DouyinImService: from .auth import DouyinAuth auth = DouyinAuth.from_im_session(self.session) controller = get_traffic_controller() - async with controller.background_slot(self.account_id, "account UID verify"): + async with controller.background_slot( + self.account_id, + "account UID verify", + startup=True, + ): async with DouyinImHttpClient(self.session, account_id=self.account_id) as http: old = int(self.session.my_uid or 0) resolved = await asyncio.to_thread(http._resolve_authoritative_uid, auth) @@ -808,11 +1156,16 @@ class DouyinImService: async def run(self): """主循环：WebSocket + HTTP 轮询""" self._running = True + _active_service_ids.add(self.account_id) await self._reply_queue.start() await self._verify_account_uid() # ensure_frontier_ws 可能触发签名/HTTP（阻塞），放线程池避免多账号启动时卡死事件循环 controller = get_traffic_controller() - async with controller.background_slot(self.account_id, "frontier discovery"): + async with controller.background_slot( + self.account_id, + "frontier discovery", + startup=True, + ): await asyncio.to_thread(ensure_frontier_ws, self.session) has_ws = bool(self.session.frontier_ws_url()) cred_summary = format_session_credential_summary(self.session) @@ -833,7 +1186,11 @@ class DouyinImService: from .emoji_pack import ensure_emoji_map, is_fresh if not is_fresh(): - async with controller.background_slot(self.account_id, "emoji preload"): + async with controller.background_slot( + self.account_id, + "emoji preload", + startup=True, + ): await asyncio.to_thread(ensure_emoji_map, self.session) except Exception as e: logger.debug(f"emoji map preload failed: {e}") @@ -844,8 +1201,12 @@ class DouyinImService: await self._ws_client.start() initial_poll_succeeded = False + initial_unread: list[dict] = [] try: - await self._poll_conversations() + initial_unread = await self._poll_conversations( + initial=True, + defer_handlers=True, + ) initial_poll_succeeded = True except Exception as e: logger.warning(f"Initial conversation poll failed: {e}") @@ -857,9 +1218,25 @@ class DouyinImService: account_id=self.account_id, ) - ws_connected = bool( - self._ws_client and getattr(self._ws_client, "connected", False) - ) + # Batch readiness covers transport discovery plus the authoritative + # list/cache snapshot, not potentially slow reply generation, logging, + # or outbound sends for pre-existing unread conversations. + await self._notify_ready() + if initial_unread: + await _get_initial_unread_dispatcher().submit( + self, + initial_unread, + ) + + # A discovered WS URL is enough to start on the low-frequency + # reconciliation schedule. If the socket does not actually open, the + # first health tick switches the account to fallback mode. This avoids + # a second fast HTTP poll just because the handshake needed a moment. + ws_connected = bool(has_ws) + if ws_connected: + _ws_service_ids.add(self.account_id) + else: + _ws_service_ids.discard(self.account_id) poll_interval, poll_stagger = _conversation_poll_timing( self.account_id, ws_connected, @@ -884,15 +1261,23 @@ class DouyinImService: poll_stagger, "yes" if ws_connected else "no", ) - - loop_count = 0 + poll_failures = 0 + follow_interval = _env_poll_seconds( + "KEFU_FOLLOW_POLL_INTERVAL_SECONDS", + 60.0, + minimum=30.0, + ) + follow_stagger = ( + (int(self.account_id or 0) * 2654435761) + % max(1, int(follow_interval * 1000)) + ) / 1000.0 + next_follow_tick_at = loop.time() + follow_interval + follow_stagger while self._running: # The initial poll above is authoritative. Sleep before the next # recurring tick so startup cannot issue two back-to-back polls. await asyncio.sleep(5) if not self._running: break - loop_count += 1 try: current_ws_connected = bool( self._ws_client @@ -900,6 +1285,10 @@ class DouyinImService: ) if current_ws_connected != ws_connected: ws_connected = current_ws_connected + if ws_connected: + _ws_service_ids.add(self.account_id) + else: + _ws_service_ids.discard(self.account_id) poll_interval, poll_stagger = _conversation_poll_timing( self.account_id, ws_connected, @@ -919,22 +1308,53 @@ class DouyinImService: poll_interval, "yes" if ws_connected else "no", ) + else: + # Large batches change the active population while older + # services are already running. Future polls adopt the + # widened capacity interval without moving a poll that is + # already scheduled. + capacity_interval, _ = _conversation_poll_timing( + self.account_id, + ws_connected, + ) + if abs(capacity_interval - poll_interval) >= 1.0: + poll_interval = capacity_interval if loop.time() >= next_conversation_poll_at: try: - await self._poll_conversations() + max_poll_seconds = _env_poll_seconds( + "KEFU_CONVERSATION_POLL_DEADLINE_SECONDS", + 30.0, + ) + await asyncio.wait_for( + self._poll_conversations(), + timeout=max_poll_seconds, + ) + poll_failures = 0 + except asyncio.TimeoutError: + poll_failures = min(4, poll_failures + 1) + logger.debug( + "Dropping stale conversation poll account=%s after %.1fs", + self.account_id, + max_poll_seconds, + ) + except Exception: + poll_failures = min(4, poll_failures + 1) + raise finally: # Advance on both success and failure. Otherwise a # past deadline retries every five-second loop tick # during an outage and amplifies traffic. - next_conversation_poll_at = loop.time() + poll_interval - if loop_count % 6 == 0: - logger.info(f"IM direct tick #{loop_count} account={self.account_id}") + next_conversation_poll_at = ( + loop.time() + poll_interval * (2 ** poll_failures) + ) # 关注欢迎语：约每 60s 检测一次新粉丝（独立于私信轮询，失败不影响主循环） - if self.follow_tick and loop_count % 12 == 0: + if self.follow_tick and loop.time() >= next_follow_tick_at: try: await self.follow_tick() except Exception as e: logger.error(f"follow welcome tick error: {e}") + finally: + next_follow_tick_at = loop.time() + follow_interval except Exception as e: logger.error(f"IM poll error: {e}") system_logger.record( @@ -947,6 +1367,11 @@ class DouyinImService: async def stop(self): self._running = False + _active_service_ids.discard(self.account_id) + _ws_service_ids.discard(self.account_id) + dispatcher = _INITIAL_UNREAD_DISPATCHERS.get(asyncio.get_running_loop()) + if dispatcher is not None: + await dispatcher.cancel(self) await get_traffic_controller().send_queue.cancel_account(self.account_id) await self._reply_queue.stop() if self._ws_client: diff --git a/backend/rpa_engine/douyin_im/ws_client.py b/backend/rpa_engine/douyin_im/ws_client.py index 35e341d..a28183c 100644 --- a/backend/rpa_engine/douyin_im/ws_client.py +++ b/backend/rpa_engine/douyin_im/ws_client.py @@ -1,9 +1,10 @@ import asyncio import logging -import threading +import os +import weakref from typing import Awaitable, Callable, Optional -from websocket import WebSocketApp +from websockets.legacy.client import WebSocketClientProtocol, connect as websocket_connect from utils import system_logger from .protocol import parse_ws_payload @@ -13,9 +14,102 @@ logger = logging.getLogger("douyin_im.ws") MessageHandler = Callable[[dict], Awaitable[None]] +# 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, +# backpressure intentionally reaches TCP instead of allocating more tasks. +_TRANSPORT_MAX_QUEUE = 4 +_APPLICATION_QUEUE_SIZE = 8 +_INCOMING_MAX_SIZE = 2**20 +_STABLE_CONNECTION_SECONDS = 60.0 +_MAX_RECONNECT_BASE_SECONDS = 60.0 +_CLOSE_GRACE_SECONDS = 2.0 +_PING_TIMEOUT_SECONDS = 120.0 +_HANDLER_CONCURRENCY_ENV = "KEFU_WS_HANDLER_CONCURRENCY" +_SYSTEM_LOG_THROTTLE_ENV = "KEFU_WS_SYSTEM_LOG_THROTTLE_SECONDS" + + +def _env_int_clamped(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)) + + +def _env_float_clamped( + name: str, + default: float, + minimum: float, + maximum: float, +) -> float: + try: + value = float(os.getenv(name, str(default)) or default) + except (TypeError, ValueError): + value = default + return max(minimum, min(maximum, value)) + + +def _handler_concurrency_limit() -> int: + # SQLite serializes writes. Eight allows unrelated parsing / reads to + # progress without letting an accidental value such as 500 recreate the + # original event-loop and database stampede. + return _env_int_clamped(_HANDLER_CONCURRENCY_ENV, 8, 1, 32) + + +def _system_log_throttle_seconds() -> float: + return _env_float_clamped(_SYSTEM_LOG_THROTTLE_ENV, 300.0, 10.0, 3600.0) + + +class _LoopWsState: + """Shared limits for all WS clients owned by one asyncio event loop.""" + + def __init__(self) -> None: + self.handler_slots = asyncio.Semaphore(_handler_concurrency_limit()) + self.system_log_last_at: dict[tuple[int, str], float] = {} + + +# asyncio synchronization primitives belong to their creating event loop. +# Keeping one weakly-keyed state per loop gives production a process-wide +# limit while keeping isolated test loops and uncommon threaded loops safe. +_LOOP_STATES: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _LoopWsState]" = ( + weakref.WeakKeyDictionary() +) + + +def _get_loop_state() -> _LoopWsState: + loop = asyncio.get_running_loop() + state = _LOOP_STATES.get(loop) + if state is None: + state = _LoopWsState() + _LOOP_STATES[loop] = state + return state + + +def _reconnect_delay(account_id: int | None, retry: int) -> float: + """Return exponential backoff with stable, account-specific full jitter. + + A connection that is accepted and immediately closed is still a failed + attempt. The old client reset its retry counter whenever ``run_forever`` + returned normally, which kept those accounts reconnecting every 2-7s. + This delay reaches a 60-90s range after repeated short-lived connections. + """ + + attempt = max(1, int(retry or 1)) + base = min(_MAX_RECONNECT_BASE_SECONDS, float(2 ** min(attempt, 6))) + # Spread later retries across half of the base interval. Keep at least a + # five-second spread on early retries so a shared outage doesn't reconnect + # every account in the same instant. + spread = max(5.0, base / 2.0) + # Keep one account's fraction stable across attempts. This preserves the + # exponential ordering while different accounts remain spread apart. + seed = (int(account_id or 0) * 2654435761) & 0xFFFFFFFF + fraction = (seed % 10000) / 10000.0 + return base + (spread * fraction) + class DouyinImWsClient: - """直连 frontier-im WebSocket（websocket-client，与 DouYin_Spider 一致）""" + """Async frontier-im WebSocket client with bounded message backpressure.""" def __init__( self, @@ -29,11 +123,14 @@ class DouyinImWsClient: self._running = False self.connected = False self._task: Optional[asyncio.Task] = None - self._loop: Optional[asyncio.AbstractEventLoop] = None - self._ws_app: Optional[WebSocketApp] = None - self._ws_lock = threading.Lock() + self._connection: Optional[WebSocketClientProtocol] = None + self._last_connection_lifetime = 0.0 + self._message_queue: Optional[asyncio.Queue[dict]] = None + self._dispatcher_task: Optional[asyncio.Task] = None async def start(self): + if self._task and not self._task.done(): + return url = self.session.frontier_ws_url() if not url: logger.warning("No frontier WebSocket URL captured; WS listener disabled") @@ -46,190 +143,309 @@ class DouyinImWsClient: ) return self._running = True - self._loop = asyncio.get_running_loop() - self._task = asyncio.create_task(self._run_loop(url)) + self._ensure_dispatcher() + self._task = asyncio.create_task( + self._run_loop(url), + name=f"im-ws-{self.account_id or 'na'}", + ) async def stop(self): self._running = False self.connected = False - with self._ws_lock: - if self._ws_app: + + connection = self._connection + if connection is not None: + try: + await asyncio.wait_for( + connection.close(code=1000, reason="client stopping"), + timeout=_CLOSE_GRACE_SECONDS, + ) + except asyncio.TimeoutError: + # Shutdown iterates over every hosted account. One broken + # peer must not consume close_timeout repeatedly and turn a + # 500-account shutdown into a many-minute operation. + logger.debug("Timed out closing IM WebSocket; aborting transport") try: - self._ws_app.close() + connection.fail_connection() except Exception: pass - self._ws_app = None - if self._task: - self._task.cancel() + except Exception: + logger.debug("Failed to close IM WebSocket cleanly", exc_info=True) + + task = self._task + if task and task is not asyncio.current_task() and not task.done(): + task.cancel() try: - await self._task + await task except asyncio.CancelledError: pass + if self._task is task: self._task = None + self._connection = None + await self._stop_dispatcher() - async def _run_ws_thread(self, url: str): - """在独立守护线程中跑 run_forever，直到连接断开/关闭。 + def _record_connection_system_event( + self, + event_key: str, + message: str, + *, + detail: str, + level: str, + ) -> bool: + """Persist at most one repeated lifecycle event per account/window.""" - 不能用共享默认线程池（run_in_executor(None)/asyncio.to_thread）： - WS 长连接会永久占用一个池线程，账号数超过池大小（默认 64）后， - 所有账号的签名/轮询任务被饿死，表现为“启动几十个账号后全部卡死超时”。 - """ loop = asyncio.get_running_loop() - done = asyncio.Event() - error: list[BaseException] = [] - - def _runner(): - try: - self._connect_sync(url) - except BaseException as e: - error.append(e) - finally: - try: - loop.call_soon_threadsafe(done.set) - except RuntimeError: - pass # 事件循环已关闭 - - thread = threading.Thread( - target=_runner, - name=f"im-ws-{self.account_id or 'na'}", - daemon=True, + state = _get_loop_state() + key = (int(self.account_id or 0), event_key) + now = loop.time() + last_at = state.system_log_last_at.get(key) + if last_at is not None and now - last_at < _system_log_throttle_seconds(): + return False + state.system_log_last_at[key] = now + system_logger.record( + message, + detail=detail, + level=level, + category="ws", + account_id=self.account_id, ) - thread.start() - try: - await done.wait() - except asyncio.CancelledError: - # stop() 会 close ws_app 使 run_forever 退出，线程随之结束 - raise - if error: - raise error[0] + return True - async def _run_loop(self, url: str): - retry = 0 - while self._running: - from .frontier import ensure_frontier_ws - from .traffic_control import get_traffic_controller + def _reset_connection_system_log_throttle(self) -> None: + loop = asyncio.get_running_loop() + state = _LOOP_STATES.get(loop) + if state is None: + return + account_key = int(self.account_id or 0) + state.system_log_last_at.pop((account_key, "connected"), None) + state.system_log_last_at.pop((account_key, "retry"), None) - # ensure_frontier_ws 可能触发签名/HTTP（阻塞），放线程池避免卡事件循环 - controller = get_traffic_controller() - async with controller.background_slot(self.account_id or 0, "websocket prepare"): - await asyncio.to_thread(ensure_frontier_ws, self.session) - connect_url = self.session.frontier_ws_url() or url + def _ensure_dispatcher(self) -> None: + if self._dispatcher_task and not self._dispatcher_task.done(): + return + if self._message_queue is None: + self._message_queue = asyncio.Queue(maxsize=_APPLICATION_QUEUE_SIZE) + self._dispatcher_task = asyncio.create_task( + self._dispatch_loop(), + name=f"im-ws-dispatch-{self.account_id or 'na'}", + ) + + async def _stop_dispatcher(self) -> None: + task = self._dispatcher_task + self._dispatcher_task = None + if task and task is not asyncio.current_task() and not task.done(): + task.cancel() try: - logger.info(f"Connecting IM WebSocket: {connect_url[:100]}...") - await self._run_ws_thread(connect_url) - retry = 0 + await task + except asyncio.CancelledError: + pass + + queue = self._message_queue + self._message_queue = None + if queue is not None: + # Dropped messages must decrement the unfinished counter so tests, + # diagnostics, and a later restart can never hang on queue.join(). + while True: + try: + queue.get_nowait() + except asyncio.QueueEmpty: + break + else: + queue.task_done() + + async def _prepare_url(self, fallback_url: str) -> str: + from .frontier import ensure_frontier_ws + from .traffic_control import get_traffic_controller + + # Frontier discovery can perform synchronous signing / HTTP work. It + # remains in the shared background lane and off the FastAPI event loop. + controller = get_traffic_controller() + async with controller.background_slot( + self.account_id or 0, + "websocket prepare", + ): + await asyncio.to_thread(ensure_frontier_ws, self.session) + return self.session.frontier_ws_url() or fallback_url + + async def _run_loop(self, initial_url: str): + retry = 0 + first_attempt = True + while self._running: + self._last_connection_lifetime = 0.0 + try: + # Startup validation already prepared the captured URL. Avoid + # repeating signing / frontier discovery for all 500 accounts + # on their first connect; refresh only after a disconnect. + if first_attempt and initial_url: + connect_url = initial_url + else: + connect_url = await self._prepare_url(initial_url) + first_attempt = False + if not connect_url: + raise RuntimeError("frontier WebSocket URL is unavailable") + logger.info("Connecting IM WebSocket: %s...", connect_url[:100]) + await self._run_connection(connect_url) except asyncio.CancelledError: break - except Exception as e: - logger.warning(f"IM WebSocket error: {e}") - system_logger.record( - "实时接收连接异常", - detail=f"建立 frontier WebSocket 失败：{e}", - level="error", - category="ws", - account_id=self.account_id, - ) + except Exception as exc: + if self._running: + logger.warning("IM WebSocket error: %s", exc) + self._record_connection_system_event( + "retry", + "实时接收连接异常", + detail=f"建立 frontier WebSocket 失败：{exc}", + level="error", + ) + # Only a genuinely stable connection earns a retry reset. A + # successful handshake followed by an immediate normal close must + # continue exponential backoff rather than reconnect forever at + # the first delay. + if self._last_connection_lifetime >= _STABLE_CONNECTION_SECONDS: + retry = 0 if not self._running: break + retry += 1 - # Stable per-account jitter prevents every hosted account from - # reconnecting in the same second after a shared network outage. - jitter = ((int(self.account_id or 0) * 2654435761) % 5000) / 1000.0 - wait = min(30.0, 2.0 * retry) + jitter - logger.info(f"IM WebSocket reconnect in {wait:.1f}s...") - system_logger.record( + wait = _reconnect_delay(self.account_id, retry) + logger.info("IM WebSocket reconnect in %.1fs...", wait) + self._record_connection_system_event( + "retry", f"实时接收断开，{wait:.1f}s 后重连", detail="frontier WebSocket 连接已断开，正在自动重连。", level="warning", - category="ws", - account_id=self.account_id, ) - await asyncio.sleep(wait) + try: + await asyncio.sleep(wait) + except asyncio.CancelledError: + break - def _connect_sync(self, url: str): - if not self._loop: - return + def _connection_headers(self) -> list[tuple[str, str]]: + headers = [ + ("Pragma", "no-cache"), + ("Cache-Control", "no-cache"), + ("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8"), + ] + cookie = self.session.cookie_header() + if cookie: + headers.append(("Cookie", cookie)) + return headers - def on_open(_ws): - self.connected = True - logger.info("IM WebSocket connected") - system_logger.record( - "实时接收通道已连接", - detail="frontier WebSocket 已建立，可实时接收私信。", - level="success", - category="ws", - account_id=self.account_id, - ) + async def _run_connection(self, url: str) -> None: + """Open one connection and dispatch messages sequentially. - def on_message(_ws, message): - asyncio.run_coroutine_threadsafe(self._dispatch(message), self._loop) + ``max_queue`` bounds the library's receive buffer and ``_dispatch`` + feeds one lifecycle-owned, bounded application queue. This receive + loop therefore remains responsive to control frames during ordinary + database stalls without creating one task per incoming frame. + """ - def on_error(_ws, error): - if self._running: - logger.warning(f"IM WebSocket error: {error}") - system_logger.record( - "实时接收通道报错", - detail=f"{error}", - level="error", - category="ws", - account_id=self.account_id, - ) - - def on_close(_ws, code, msg): - self.connected = False - logger.info(f"IM WebSocket closed: code={code}, msg={msg}") - if self._running: - system_logger.record( - "实时接收通道关闭", - detail=f"code={code}, msg={msg}", - level="warning", - category="ws", - account_id=self.account_id, - ) - - headers = { - "User-Agent": self.session.user_agent, - "Pragma": "no-cache", - "Cache-Control": "no-cache", - "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8", - "Sec-WebSocket-Protocol": "binary, base64, pbbp2", - "Sec-WebSocket-Extensions": "permessage-deflate; client_max_window_bits", - } - ws_app = WebSocketApp( - url, - header=headers, - cookie=self.session.cookie_header(), - on_open=on_open, - on_message=on_message, - on_error=on_error, - on_close=on_close, - ) - with self._ws_lock: - self._ws_app = ws_app + loop = asyncio.get_running_loop() + connected_at: float | None = None + connection: Optional[WebSocketClientProtocol] = None try: - ws_app.run_forever(origin="https://www.douyin.com", ping_interval=20, ping_timeout=10) + async with websocket_connect( + url, + origin="https://www.douyin.com", + subprotocols=["binary", "base64", "pbbp2"], + extra_headers=self._connection_headers(), + user_agent_header=self.session.user_agent, + compression="deflate", + open_timeout=10, + ping_interval=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. + ping_timeout=_PING_TIMEOUT_SECONDS, + close_timeout=3, + # Frontier frames contain metadata and media URLs rather than + # media bytes. A finite frame limit plus a finite queue makes + # receive memory genuinely bounded across hundreds of peers. + max_size=_INCOMING_MAX_SIZE, + max_queue=_TRANSPORT_MAX_QUEUE, + ) as websocket: + connection = websocket + self._connection = websocket + connected_at = loop.time() + self.connected = True + logger.info("IM WebSocket connected") + self._record_connection_system_event( + "connected", + "实时接收通道已连接", + detail="frontier WebSocket 已建立，可实时接收私信。", + level="success", + ) + + async for raw in websocket: + if not self._running: + break + await self._dispatch(raw) finally: + if connected_at is not None: + self._last_connection_lifetime = max(0.0, loop.time() - connected_at) + if self._last_connection_lifetime >= _STABLE_CONNECTION_SECONDS: + # A genuinely healthy session starts a new lifecycle. Its + # next outage should be visible immediately rather than + # hidden by an old retry window. + self._reset_connection_system_log_throttle() self.connected = False - with self._ws_lock: - if self._ws_app is ws_app: - self._ws_app = None + if self._connection is connection: + self._connection = None + + if connection is not None: + code = connection.close_code + reason = connection.close_reason + logger.info("IM WebSocket closed: code=%s, msg=%s", code, reason) + if self._running: + self._record_connection_system_event( + "retry", + "实时接收通道关闭", + detail=f"code={code}, msg={reason}", + level="warning", + ) async def _dispatch(self, raw): + self._ensure_dispatcher() + queue = self._message_queue + if queue is None: + return if isinstance(raw, str): payload = raw.encode("utf-8", errors="ignore") else: payload = raw items = parse_ws_payload(payload) for item in items: + if not self._running: + return + await queue.put(item) + + async def _dispatch_loop(self) -> None: + queue = self._message_queue + if queue is None: + return + handler_slots = _get_loop_state().handler_slots + while True: + item = await queue.get() try: - await self.on_message(item) - except Exception as e: - logger.debug(f"WS message handler error: {e}") + if self._running: + # Every account owns one dispatcher, preserving its FIFO. + # The shared semaphore prevents 500 dispatchers from + # entering SQLite / reply work at the same instant. A + # dispatcher waiting here is directly cancellable by + # stop(); no detached per-message task is created. + async with handler_slots: + if self._running: + await self.on_message(item) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.debug("WS message handler error: %s", exc) system_logger.record( "实时消息处理失败", - detail=f"处理收到的私信时出错：{e}", + detail=f"处理收到的私信时出错：{exc}", level="error", category="recv", account_id=self.account_id, ) + finally: + queue.task_done() diff --git a/backend/rpa_engine/playwright_worker.py b/backend/rpa_engine/playwright_worker.py index 32737c5..d2d0b76 100644 --- a/backend/rpa_engine/playwright_worker.py +++ b/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() diff --git a/backend/tests/test_account_pagination.py b/backend/tests/test_account_pagination.py index d51ccc9..a65837d 100644 --- a/backend/tests/test_account_pagination.py +++ b/backend/tests/test_account_pagination.py @@ -5,7 +5,7 @@ import sys import unittest from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker @@ -20,7 +20,7 @@ if str(BACKEND_DIR) not in sys.path: import main from models.database import Base -from models.models import Account +from models.models import Account, MessageLog class _CountResult: @@ -43,6 +43,160 @@ class _RowsResult: class AccountPaginationTests(unittest.IsolatedAsyncioTestCase): + async def test_legacy_cookie_sync_only_reads_accounts_missing_db_cookie(self): + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + session_factory = sessionmaker( + engine, + class_=AsyncSession, + expire_on_commit=False, + ) + try: + async with session_factory() as db: + db.add_all( + [ + Account( + id=2301, + username="already-in-db", + cookie_data='{"cookies": [{"value": "large"}]}', + im_session_data="x" * 100_000, + ), + Account(id=2302, username="legacy-file", cookie_data=None), + ] + ) + await db.commit() + + def read_cookie(account_id: int): + self.assertEqual(account_id, 2302) + return '{"cookies": [{"value": "migrated"}]}' + + with ( + patch.object(main, "AsyncSessionLocal", session_factory), + patch.object(main, "read_cookie_file", side_effect=read_cookie) as read_file, + patch.object(main, "get_cookie_path", return_value="legacy-2302.json"), + ): + await main._sync_legacy_cookie_files() + + read_file.assert_called_once_with(2302) + async with session_factory() as db: + migrated = await db.get(Account, 2302) + self.assertEqual( + migrated.cookie_data, + '{"cookies": [{"value": "migrated"}]}', + ) + self.assertEqual(migrated.cookie_path, "legacy-2302.json") + finally: + await engine.dispose() + + async def test_account_edit_invalidates_running_follow_config_cache(self): + account = SimpleNamespace(id=2201) + db = SimpleNamespace(commit=AsyncMock(), refresh=AsyncMock()) + worker = SimpleNamespace(invalidate_follow_welcome_config=MagicMock()) + original_workers = main.manager.workers + main.manager.workers = {2201: worker} + try: + with ( + patch.object(main, "get_owned_account", AsyncMock(return_value=account)), + patch.object(main, "_build_account_response", return_value={"id": 2201}), + ): + response = await main.update_account( + account_id=2201, + body=main.AccountUpdate(follow_welcome_enabled=True), + db=db, + user=SimpleNamespace(id=7, role="operator"), + ) + + self.assertEqual(response, {"id": 2201}) + worker.invalidate_follow_welcome_config.assert_called_once_with() + finally: + main.manager.workers = original_workers + + async def test_log_stats_uses_one_aggregate_and_respects_ownership(self): + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + session_factory = sessionmaker( + engine, + class_=AsyncSession, + expire_on_commit=False, + ) + try: + async with session_factory() as db: + db.add_all( + [ + Account(id=2101, owner_id=7, status="offline"), + Account(id=2102, owner_id=8, status="offline"), + MessageLog(account_id=2101, status="received"), + MessageLog(account_id=2101, status="replied"), + MessageLog(account_id=2102, status="replied"), + ] + ) + await db.commit() + + with patch.object(db, "execute", wraps=db.execute) as execute: + stats = await main.get_logs_stats( + account_id=None, + db=db, + user=SimpleNamespace(id=7, role="user"), + ) + + self.assertEqual(stats, {"total": 2, "replied": 1}) + self.assertEqual(execute.await_count, 1) + finally: + await engine.dispose() + + async def test_account_options_returns_lightweight_runtime_fields(self): + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + session_factory = sessionmaker( + engine, + class_=AsyncSession, + expire_on_commit=False, + ) + original_workers = main.manager.workers + main.manager.workers = {2001: SimpleNamespace(is_running=True)} + try: + async with session_factory() as db: + db.add_all( + [ + Account( + id=2001, + owner_id=7, + username="owned", + status="offline", + cookie_data='{"cookies": []}', + im_session_data="x" * 100_000, + reply_cooldown_seconds=12, + ), + Account( + id=2002, + owner_id=8, + username="other", + status="online", + cookie_data='{"cookies": []}', + ), + ] + ) + await db.commit() + + options = await main.get_account_options( + db=db, + user=SimpleNamespace(id=7, role="user"), + ) + + self.assertEqual(len(options), 1) + self.assertEqual(options[0].id, 2001) + self.assertEqual(options[0].status, "online") + self.assertTrue(options[0].has_cookie) + self.assertEqual(options[0].reply_cooldown_seconds, 12) + self.assertEqual(options[0].reply_cooldown_effective, 12) + self.assertFalse(hasattr(options[0], "im_session_data")) + finally: + main.manager.workers = original_workers + await engine.dispose() + async def test_paginated_list_counts_then_loads_only_current_page(self): page_rows = [ SimpleNamespace(id=10, status="offline"), diff --git a/backend/tests/test_batch_start_api.py b/backend/tests/test_batch_start_api.py index 0e81ef9..d1b61de 100644 --- a/backend/tests/test_batch_start_api.py +++ b/backend/tests/test_batch_start_api.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import os import sys import unittest @@ -32,6 +33,219 @@ def _fake_db(rows): class BatchStartApiTests(unittest.IsolatedAsyncioTestCase): + async def test_shutdown_continues_when_batch_queue_cleanup_times_out(self): + original_workers = main.manager.workers + original_flush_task = main._system_log_flush_task + main.manager.workers = {} + main._system_log_flush_task = None + stop_started = asyncio.Event() + + async def blocked_batch_stop(): + stop_started.set() + await asyncio.Event().wait() + + try: + with ( + patch.dict( + os.environ, + {"KEFU_BATCH_STOP_TIMEOUT_SECONDS": "1"}, + ), + patch.object( + main.batch_start_queue, + "stop", + AsyncMock(side_effect=blocked_batch_stop), + ), + patch( + "rpa_engine.douyin_im.traffic_control.shutdown_traffic_controller", + AsyncMock(), + ) as stop_traffic, + ): + await asyncio.wait_for(main.shutdown(), timeout=2.0) + + self.assertTrue(stop_started.is_set()) + stop_traffic.assert_awaited_once_with() + finally: + main.manager.workers = original_workers + main._system_log_flush_task = original_flush_task + + async def test_shutdown_stops_many_accounts_with_bounded_parallelism(self): + active = 0 + maximum_active = 0 + stopped: list[int] = [] + original_workers = main.manager.workers + original_flush_task = main._system_log_flush_task + main.manager.workers = { + account_id: SimpleNamespace(is_running=True) + for account_id in range(601, 613) + } + main._system_log_flush_task = None + + async def stop_worker(account_id: int): + nonlocal active, maximum_active + active += 1 + maximum_active = max(maximum_active, active) + try: + await asyncio.sleep(0.005) + stopped.append(account_id) + main.manager.workers.pop(account_id, None) + return True + finally: + active -= 1 + + try: + with ( + patch.dict( + os.environ, + { + "KEFU_SHUTDOWN_CONCURRENCY": "3", + "KEFU_SHUTDOWN_TIMEOUT_SECONDS": "5", + }, + ), + patch.object(main.batch_start_queue, "stop", AsyncMock()), + patch.object(main.manager, "stop_worker", AsyncMock(side_effect=stop_worker)), + patch( + "rpa_engine.douyin_im.traffic_control.shutdown_traffic_controller", + AsyncMock(), + ), + ): + await main.shutdown() + + self.assertEqual(len(stopped), 12) + self.assertEqual(maximum_active, 3) + finally: + main.manager.workers = original_workers + main._system_log_flush_task = original_flush_task + + async def test_worker_manager_waits_for_full_ready_and_reuses_validation(self): + worker = SimpleNamespace( + is_running=True, + start=AsyncMock(), + wait_until_ready=AsyncMock(), + ) + manager = main.WorkerManager() + + with patch.object(main, "DouyinWorker", return_value=worker) as worker_factory: + started = await manager.start_worker( + 501, + login_mode="im_direct", + wait_until_ready=True, + credential_prevalidated=True, + ) + + self.assertTrue(started) + worker_factory.assert_called_once_with( + 501, + login_mode="im_direct", + credential_prevalidated=True, + ) + worker.start.assert_awaited_once_with() + worker.wait_until_ready.assert_awaited_once_with() + self.assertIs(manager.workers[501], worker) + + async def test_cancelled_ready_wait_stops_and_removes_detached_worker(self): + wait_started = asyncio.Event() + waiting = asyncio.Event() + + async def wait_forever(): + wait_started.set() + await waiting.wait() + + worker = SimpleNamespace( + is_running=True, + start=AsyncMock(), + wait_until_ready=AsyncMock(side_effect=wait_forever), + stop=AsyncMock(), + ) + manager = main.WorkerManager() + + with patch.object(main, "DouyinWorker", return_value=worker): + task = asyncio.create_task( + manager.start_worker( + 502, + login_mode="im_direct", + wait_until_ready=True, + credential_prevalidated=True, + ) + ) + await asyncio.wait_for(wait_started.wait(), timeout=0.2) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + worker.stop.assert_awaited_once_with() + self.assertNotIn(502, manager.workers) + + async def test_batch_start_waits_for_ready_and_skips_duplicate_validation(self): + account = SimpleNamespace( + id=503, + status="offline", + qr_code_base64=None, + error_message=None, + im_session_data="saved-session", + ) + db = SimpleNamespace(commit=AsyncMock()) + assessment = { + "login_mode": "im_direct", + "should_reset": False, + "can_skip_browser": True, + "message": "ready", + "cookie_valid": True, + "im_ready": True, + } + + with ( + patch.object(main.manager, "is_running", return_value=False), + patch.object(main.manager, "start_worker", AsyncMock(return_value=True)) as start, + patch.object(main, "_get_account_cookie_data", return_value="{}"), + patch.object(main, "assess_account_credential", AsyncMock(return_value=assessment)), + ): + result = await main._start_account_rpa_impl( + account, + db, + wait_for_ready=True, + ) + + start.assert_awaited_once_with( + 503, + login_mode="im_direct", + wait_until_ready=True, + credential_prevalidated=True, + ) + self.assertEqual(result["status"], "running") + + async def test_batch_start_does_not_launch_interactive_browser_login(self): + account = SimpleNamespace( + id=504, + status="offline", + qr_code_base64=None, + error_message=None, + im_session_data=None, + ) + db = SimpleNamespace(commit=AsyncMock()) + assessment = { + "login_mode": "browser", + "should_reset": False, + "can_skip_browser": False, + "message": "login required", + "cookie_valid": False, + "im_ready": False, + } + + with ( + patch.object(main.manager, "is_running", return_value=False), + patch.object(main.manager, "start_worker", AsyncMock()) as start, + patch.object(main, "_get_account_cookie_data", return_value=None), + patch.object(main, "assess_account_credential", AsyncMock(return_value=assessment)), + ): + with self.assertRaisesRegex(RuntimeError, "批量启动"): + await main._start_account_rpa_impl( + account, + db, + wait_for_ready=True, + ) + + start.assert_not_awaited() + async def test_start_all_uses_lightweight_select_and_submits_once(self): db = _fake_db([(1, False), (2, True), (3, False), (4, False)]) user = SimpleNamespace(id=9, role="admin") diff --git a/backend/tests/test_conversation_poll_bandwidth.py b/backend/tests/test_conversation_poll_bandwidth.py index 9253abe..fe272fc 100644 --- a/backend/tests/test_conversation_poll_bandwidth.py +++ b/backend/tests/test_conversation_poll_bandwidth.py @@ -1,10 +1,12 @@ from __future__ import annotations +import asyncio import os import sys import unittest from contextlib import asynccontextmanager from pathlib import Path +from types import SimpleNamespace from unittest.mock import AsyncMock, patch @@ -109,6 +111,42 @@ class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase): self.assertGreaterEqual(http_stagger, 0) self.assertLess(http_stagger, http_interval) + def test_poll_interval_expands_to_the_configured_population_budget(self): + with patch.dict( + os.environ, + { + "KEFU_WS_RECONCILE_INTERVAL_SECONDS": "120", + "KEFU_HTTP_POLL_INTERVAL_SECONDS": "15", + "KEFU_WS_POLL_BUDGET_RPS": "1", + "KEFU_HTTP_POLL_BUDGET_RPS": "1", + }, + ): + ws_interval, _ = _conversation_poll_timing( + 123, + True, + population=500, + ) + http_interval, _ = _conversation_poll_timing( + 123, + False, + population=500, + ) + + self.assertEqual(ws_interval, 500) + self.assertEqual(http_interval, 500) + + def test_initial_unread_concurrency_is_configurable_and_bounded(self): + with patch.dict( + os.environ, + {"KEFU_INITIAL_UNREAD_CONCURRENCY": "4"}, + ): + self.assertEqual(service_module._initial_unread_concurrency(), 4) + with patch.dict( + os.environ, + {"KEFU_INITIAL_UNREAD_CONCURRENCY": "999"}, + ): + self.assertEqual(service_module._initial_unread_concurrency(), 8) + async def test_service_poll_uses_one_conversation_request_without_unread_probe(self): class _Controller: @asynccontextmanager @@ -141,7 +179,260 @@ class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase): await service._poll_conversations() http.get_conversations.assert_awaited_once_with(enrich_profiles=False) - service._index_conversations.assert_awaited_once_with([]) + service._index_conversations.assert_awaited_once_with( + [], + enrich_profiles=False, + ) + + async def test_poll_only_handles_unread_or_a_genuinely_changed_preview(self): + class _Controller: + def __init__(self): + self.startup_flags = [] + + @asynccontextmanager + async def background_slot(self, *_args, **kwargs): + self.startup_flags.append(bool(kwargs.get("startup"))) + yield + + snapshots = [ + [ + { + "conversation_id": "0:1:10001:20001", + "peer_uid": "20001", + "sender_name": "历史会话", + "sender_avatar": "https://example.test/a.png", + "content": "历史消息", + "unread_count": 0, + }, + { + "conversation_id": "0:1:10001:20002", + "peer_uid": "20002", + "sender_name": "未读会话", + "sender_avatar": "https://example.test/b.png", + "content": "新消息", + "unread_count": 1, + }, + ], + [ + { + "conversation_id": "0:1:10001:20001", + "peer_uid": "20001", + "sender_name": "历史会话", + "sender_avatar": "https://example.test/a.png", + "content": "历史消息", + "unread_count": 0, + }, + { + "conversation_id": "0:1:10001:20002", + "peer_uid": "20002", + "sender_name": "未读会话", + "sender_avatar": "https://example.test/b.png", + "content": "新消息", + "unread_count": 0, + }, + ], + [ + { + "conversation_id": "0:1:10001:20001", + "peer_uid": "20001", + "sender_name": "历史会话", + "sender_avatar": "https://example.test/a.png", + "content": "真正发生变化", + "unread_count": 0, + }, + { + "conversation_id": "0:1:10001:20002", + "peer_uid": "20002", + "sender_name": "未读会话", + "sender_avatar": "https://example.test/b.png", + "content": "新消息", + "unread_count": 0, + }, + ], + ] + + class _HttpClient: + def __init__(self): + self.get_conversations = AsyncMock(side_effect=snapshots) + self.enter_count = 0 + self.exit_count = 0 + + async def __aenter__(self): + self.enter_count += 1 + return self + + async def __aexit__(self, *_args): + self.exit_count += 1 + return False + + http = _HttpClient() + service = DouyinImService( + session=DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001), + match_reply=AsyncMock(return_value=None), + log_fn=AsyncMock(), + account_id=9, + ) + service._handle_incoming = AsyncMock() + controller = _Controller() + + with ( + patch.object( + service_module, + "get_traffic_controller", + return_value=controller, + ), + patch.object(service_module, "DouyinImHttpClient", return_value=http) as factory, + ): + # Startup indexes both previews, but only the unread conversation + # is allowed to enter the reply path. + await service._poll_conversations(initial=True) + self.assertEqual(service._handle_incoming.await_count, 1) + self.assertEqual( + service._handle_incoming.await_args.args[0]["peer_uid"], + "20002", + ) + + # The first normal reconciliation sees the exact same previews; + # it must not merely defer a historical-message reply explosion. + service._handle_incoming.reset_mock() + await service._poll_conversations() + service._handle_incoming.assert_not_awaited() + + # A real preview transition is processed even if unread_count is + # unavailable/zero on the upstream response. + await service._poll_conversations() + service._handle_incoming.assert_awaited_once() + self.assertEqual( + service._handle_incoming.await_args.args[0]["peer_uid"], + "20001", + ) + + self.assertEqual(factory.call_count, 3) + self.assertEqual(http.enter_count, 3) + self.assertEqual(http.exit_count, 3) + self.assertEqual(controller.startup_flags, [True, False, False]) + + async def test_ready_is_not_blocked_by_slow_initial_unread_handler(self): + events: list[str] = [] + ready = asyncio.Event() + handler_started = asyncio.Event() + handler_cancelled = asyncio.Event() + never_release = asyncio.Event() + + def on_ready(): + events.append("ready") + ready.set() + + async def slow_handler(_message): + events.append("handler") + handler_started.set() + try: + await never_release.wait() + except asyncio.CancelledError: + handler_cancelled.set() + raise + + class _WsClient: + connected = False + + def __init__(self, *_args, **_kwargs): + self.start = AsyncMock() + self.stop = AsyncMock() + + service = DouyinImService( + session=DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001), + match_reply=AsyncMock(return_value=None), + log_fn=AsyncMock(), + account_id=901, + on_ready=on_ready, + ) + service._verify_account_uid = AsyncMock() + service._poll_conversations = AsyncMock( + return_value=[ + { + "conversation_id": "0:1:10001:29001", + "content": "startup unread", + "unread_count": 1, + } + ] + ) + service._handle_incoming = AsyncMock(side_effect=slow_handler) + service._reply_queue.start = AsyncMock() + service._reply_queue.stop = AsyncMock() + + with ( + patch.object(service_module, "DouyinImWsClient", _WsClient), + patch.object(service_module, "ensure_frontier_ws"), + patch.object(service_module.system_logger, "record"), + patch( + "rpa_engine.douyin_im.emoji_pack.is_fresh", + return_value=True, + ), + ): + run_task = asyncio.create_task(service.run()) + try: + # If initial unread were still processed inline, this wait + # would time out because slow_handler never completes. + await asyncio.wait_for(ready.wait(), timeout=0.3) + await asyncio.wait_for(handler_started.wait(), timeout=0.3) + self.assertEqual(events[:2], ["ready", "handler"]) + service._poll_conversations.assert_awaited_once_with( + initial=True, + defer_handlers=True, + ) + + # Account stop cancels its active deferred handler instead of + # leaving work detached from the service lifecycle. + await asyncio.wait_for(service.stop(), timeout=0.5) + self.assertTrue(handler_cancelled.is_set()) + finally: + run_task.cancel() + await asyncio.gather(run_task, return_exceptions=True) + await service_module._shutdown_initial_unread_dispatcher() + + async def test_initial_unread_dispatcher_has_process_wide_concurrency_limit(self): + dispatcher = service_module._InitialUnreadDispatcher(concurrency=2) + active = 0 + maximum_active = 0 + processed: list[int] = [] + two_started = asyncio.Event() + release = asyncio.Event() + + def make_service(account_id: int): + async def handle(_message): + nonlocal active, maximum_active + active += 1 + maximum_active = max(maximum_active, active) + if active == 2: + two_started.set() + try: + await release.wait() + processed.append(account_id) + finally: + active -= 1 + + return SimpleNamespace( + account_id=account_id, + _running=True, + _handle_incoming=handle, + ) + + services = [make_service(index + 1) for index in range(6)] + try: + for service in services: + await dispatcher.submit(service, [{"unread_count": 1}]) + + await asyncio.wait_for(two_started.wait(), timeout=0.3) + await asyncio.sleep(0.03) + self.assertEqual(maximum_active, 2) + self.assertEqual(processed, []) + + release.set() + await asyncio.wait_for(dispatcher.join(), timeout=0.5) + self.assertEqual(maximum_active, 2) + self.assertCountEqual(processed, range(1, 7)) + finally: + await dispatcher.stop() if __name__ == "__main__": diff --git a/backend/tests/test_credential_responsiveness.py b/backend/tests/test_credential_responsiveness.py index d46d125..8c2810c 100644 --- a/backend/tests/test_credential_responsiveness.py +++ b/backend/tests/test_credential_responsiveness.py @@ -24,15 +24,9 @@ class CredentialResponsivenessTests(unittest.IsolatedAsyncioTestCase): is_sign_ready=lambda: True, ) - with ( - patch( - "rpa_engine.credential.ensure_frontier_ws", - return_value=None, - ), - patch( - "rpa_engine.credential.DouyinAuth.from_im_session", - return_value=auth, - ), + with patch( + "rpa_engine.credential.DouyinAuth.from_im_session", + return_value=auth, ): result = await validate_im_session( session, diff --git a/backend/tests/test_io_safety.py b/backend/tests/test_io_safety.py new file mode 100644 index 0000000..103d36c --- /dev/null +++ b/backend/tests/test_io_safety.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +import json +import logging.handlers +import os +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from sqlalchemy import text +from sqlalchemy.pool import StaticPool + + +BACKEND_DIR = Path(__file__).resolve().parents[1] +if str(BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(BACKEND_DIR)) + +from models.db_config import DatabaseConfig, create_database_engine, engine_kwargs_for_url +from models.db_migrate import migrate_message_logs_table +from models.models import MessageLog +from rpa_engine.douyin_im import protocol +from rpa_engine.douyin_im.static import Live_pb2, Response_pb2 +from utils import system_logger +from utils.log_limits import ( + TRUNCATION_MARKER, + bound_message_log_content, + bound_raw_message_log_content, +) + + +class LogLimitTests(unittest.TestCase): + def tearDown(self): + system_logger.clear() + + def test_system_log_caps_persisted_and_console_detail(self): + with patch.dict(os.environ, {"KEFU_SYSTEM_LOG_MAX_CHARS": "512"}): + with self.assertLogs("douyin_im.system", level="INFO") as captured: + entry = system_logger.record("event", "x" * 5000) + + self.assertLessEqual(len(entry["detail"]), 512) + self.assertIn(TRUNCATION_MARKER.strip(), entry["detail"]) + self.assertLess(len(captured.output[0]), 700) + + def test_oversized_media_log_remains_valid_compact_json(self): + payload = json.dumps( + { + "type": "sticker", + "url": "https://example.invalid/sticker.webp", + "text": "x" * 20000, + "unused_blob": "y" * 20000, + }, + ensure_ascii=False, + ) + with patch.dict(os.environ, {"KEFU_MESSAGE_LOG_MAX_CHARS": "4096"}): + bounded = bound_message_log_content(payload) + + decoded = json.loads(bounded) + self.assertEqual(decoded["type"], "sticker") + self.assertEqual(decoded["url"], "https://example.invalid/sticker.webp") + self.assertTrue(decoded["_log_truncated"]) + self.assertNotIn("unused_blob", decoded) + self.assertLessEqual(len(bounded), 4096) + + def test_message_model_validator_caps_all_insert_paths(self): + with patch.dict(os.environ, {"KEFU_MESSAGE_LOG_MAX_CHARS": "2048"}): + row = MessageLog(message_content="m" * 10000, reply_content="r" * 10000) + + self.assertLessEqual(len(row.message_content), 2048) + self.assertLessEqual(len(row.reply_content), 2048) + + def test_raw_message_log_is_bounded(self): + with patch.dict(os.environ, {"KEFU_RAW_MESSAGE_LOG_MAX_CHARS": "4096"}): + bounded = bound_raw_message_log_content("z" * 20000) + self.assertLessEqual(len(bounded), 4096) + self.assertIn(TRUNCATION_MARKER.strip(), bounded) + + +class SqliteIoTests(unittest.IsolatedAsyncioTestCase): + async def test_short_memory_url_uses_one_static_connection(self): + kwargs = engine_kwargs_for_url("sqlite+aiosqlite://") + self.assertIs(kwargs["poolclass"], StaticPool) + self.assertNotIn("pool_size", kwargs) + + engine = create_database_engine( + DatabaseConfig( + db_type="sqlite", + database_url="sqlite+aiosqlite://", + ) + ) + try: + async with engine.begin() as conn: + await conn.execute(text("CREATE TABLE memory_probe (id INTEGER)")) + async with engine.begin() as conn: + await conn.execute(text("INSERT INTO memory_probe VALUES (1)")) + count = ( + await conn.execute(text("SELECT count(*) FROM memory_probe")) + ).scalar_one() + self.assertEqual(count, 1) + finally: + await engine.dispose() + + async def test_file_sqlite_uses_bounded_pool_and_wal_pragmas(self): + kwargs = engine_kwargs_for_url("sqlite+aiosqlite:///example.db") + self.assertEqual(kwargs["pool_size"], 5) + self.assertEqual(kwargs["max_overflow"], 0) + self.assertEqual(kwargs["connect_args"]["timeout"], 30.0) + + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "io.db" + engine = create_database_engine( + DatabaseConfig(db_type="sqlite", db_path=str(db_path)) + ) + try: + async with engine.connect() as conn: + journal_mode = (await conn.execute(text("PRAGMA journal_mode"))).scalar_one() + synchronous = (await conn.execute(text("PRAGMA synchronous"))).scalar_one() + busy_timeout = (await conn.execute(text("PRAGMA busy_timeout"))).scalar_one() + self.assertEqual(str(journal_mode).lower(), "wal") + self.assertEqual(synchronous, 1) + self.assertEqual(busy_timeout, 30000) + finally: + await engine.dispose() + + async def test_existing_log_tables_receive_composite_indexes(self): + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "migration.db" + engine = create_database_engine( + DatabaseConfig(db_type="sqlite", db_path=str(db_path)) + ) + try: + async with engine.begin() as conn: + await conn.execute( + text( + "CREATE TABLE message_logs (" + "id INTEGER PRIMARY KEY, account_id INTEGER, " + "created_at DATETIME, sender_avatar TEXT, status VARCHAR(50))" + ) + ) + await conn.execute( + text( + "CREATE TABLE received_message_logs (" + "id INTEGER PRIMARY KEY, account_id INTEGER, " + "created_at DATETIME)" + ) + ) + await conn.execute( + text( + "CREATE TABLE system_logs (" + "id INTEGER PRIMARY KEY, account_id INTEGER, " + "created_at DATETIME)" + ) + ) + await conn.run_sync(migrate_message_logs_table) + + async with engine.connect() as conn: + message_indexes = { + row[1] + for row in (await conn.execute(text("PRAGMA index_list(message_logs)"))).all() + } + received_indexes = { + row[1] + for row in ( + await conn.execute(text("PRAGMA index_list(received_message_logs)")) + ).all() + } + system_indexes = { + row[1] + for row in (await conn.execute(text("PRAGMA index_list(system_logs)"))).all() + } + latest_plan = " ".join( + str(row[-1]) + for row in ( + await conn.execute( + text( + "EXPLAIN QUERY PLAN SELECT * FROM message_logs " + "ORDER BY created_at DESC LIMIT 50" + ) + ) + ).all() + ) + status_plan = " ".join( + str(row[-1]) + for row in ( + await conn.execute( + text( + "EXPLAIN QUERY PLAN SELECT count(*) FROM message_logs " + "WHERE status = 'replied'" + ) + ) + ).all() + ) + account_plan = " ".join( + str(row[-1]) + for row in ( + await conn.execute( + text( + "EXPLAIN QUERY PLAN SELECT * FROM message_logs " + "WHERE account_id = 1 ORDER BY created_at DESC LIMIT 50" + ) + ) + ).all() + ) + system_plan = " ".join( + str(row[-1]) + for row in ( + await conn.execute( + text( + "EXPLAIN QUERY PLAN SELECT * FROM system_logs " + "WHERE account_id = 1 ORDER BY created_at DESC LIMIT 50" + ) + ) + ).all() + ) + self.assertIn("ix_message_logs_account_created_at", message_indexes) + self.assertIn("ix_message_logs_created_at", message_indexes) + self.assertIn("ix_message_logs_status_account_id", message_indexes) + self.assertIn( + "ix_received_message_logs_account_created_at", + received_indexes, + ) + self.assertIn("ix_system_logs_account_created_at", system_indexes) + self.assertIn("ix_message_logs_created_at", latest_plan) + self.assertIn("ix_message_logs_status_account_id", status_plan) + self.assertIn("ix_message_logs_account_created_at", account_plan) + self.assertIn("ix_system_logs_account_created_at", system_plan) + finally: + await engine.dispose() + + +class WebSocketDebugTests(unittest.TestCase): + def _frame(self, *, message_type: int, content: str) -> bytes: + response = Response_pb2.Response() + message = response.body.new_message_notify.message + message.conversation_id = "0:1:200:100" + message.server_message_id = 123 + message.message_type = message_type + message.sender = 200 + message.content = content + frame = Live_pb2.PushFrame() + frame.payloadType = "pb" + frame.payload = response.SerializeToString() + return frame.SerializeToString() + + def test_control_frame_is_filtered_before_debug_writer(self): + with patch.object(protocol, "_dump_ws_message") as dump: + result = protocol.parse_ws_payload( + self._frame(message_type=50001, content='{"command_type":6}') + ) + self.assertEqual(result, []) + dump.assert_not_called() + + def test_debug_writer_uses_non_blocking_rotating_queue(self): + # Inspect construction without writing chat data to the repository. + with tempfile.TemporaryDirectory() as temp_dir: + old_path = protocol._WS_DEBUG_PATH + protocol._WS_DEBUG_PATH = str(Path(temp_dir) / "ws.log") + protocol._WS_DEBUG_LOGGER = None + try: + with patch.dict(os.environ, {"KEFU_WS_DEBUG": "1"}): + protocol._dump_ws_message(1, "conv", "hello") + logger = protocol._WS_DEBUG_LOGGER + self.assertIsNotNone(logger) + self.assertIsInstance(logger.handlers[0], logging.handlers.QueueHandler) + self.assertIsInstance( + logger._kefu_rotating_handler, + logging.handlers.RotatingFileHandler, + ) + finally: + logger = protocol._WS_DEBUG_LOGGER + if logger is not None: + logger._kefu_queue_listener.stop() + logger._kefu_rotating_handler.close() + logger.handlers.clear() + protocol._WS_DEBUG_LOGGER = None + protocol._WS_DEBUG_PATH = old_path + + +if __name__ == "__main__": + unittest.main() diff --git a/backend/tests/test_sec_user_id_guard.py b/backend/tests/test_sec_user_id_guard.py index 989a1e9..314507b 100644 --- a/backend/tests/test_sec_user_id_guard.py +++ b/backend/tests/test_sec_user_id_guard.py @@ -58,6 +58,7 @@ class SecUserIdGuardTests(unittest.IsolatedAsyncioTestCase): worker.is_running = True worker._im_service = SimpleNamespace(_running=True) worker._require_sec_user_id = AsyncMock(return_value=False) + worker._refresh_follow_welcome_config = AsyncMock() worker.get_db = AsyncMock() await worker.follow_welcome_tick() @@ -68,8 +69,24 @@ class SecUserIdGuardTests(unittest.IsolatedAsyncioTestCase): # The identity guard must run before the account's follow-welcome flag # is queried. Otherwise accounts with that feature disabled could stay # hosted indefinitely without a sec_user_id. + worker._refresh_follow_welcome_config.assert_not_awaited() worker.get_db.assert_not_awaited() + async def test_cached_disabled_follow_setting_still_guards_missing_identity(self): + worker = DouyinWorker(account_id=311, login_mode="im_direct") + worker.is_running = True + worker._im_service = SimpleNamespace(_running=True) + worker._follow_config_loaded = True + worker._refresh_follow_welcome_config = AsyncMock( + return_value=(False, "", "") + ) + worker._require_sec_user_id = AsyncMock(return_value="") + + await worker.follow_welcome_tick() + + worker._refresh_follow_welcome_config.assert_awaited_once_with() + worker._require_sec_user_id.assert_awaited_once_with("托管运行中") + async def test_blank_sec_user_id_is_missing_and_stops_hosting(self): worker = DouyinWorker(account_id=303, login_mode="im_direct") worker._load_sec_user_id = AsyncMock(return_value=" ") diff --git a/backend/tests/test_worker_scale_controls.py b/backend/tests/test_worker_scale_controls.py new file mode 100644 index 0000000..d22dc32 --- /dev/null +++ b/backend/tests/test_worker_scale_controls.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +import json +import os +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + + +BACKEND_DIR = Path(__file__).resolve().parents[1] +os.environ["KEFU_DB_TYPE"] = "sqlite" +os.environ["KEFU_DATABASE_URL"] = "" +os.environ["KEFU_DB_PATH"] = str(BACKEND_DIR / "kefu.db") +if str(BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(BACKEND_DIR)) + +from rpa_engine.douyin_im.session import DouyinImSession +from rpa_engine.playwright_worker import DouyinWorker + + +class WorkerScaleControlTests(unittest.IsolatedAsyncioTestCase): + async def test_storage_load_selects_only_cookie_column(self): + class _ScalarResult: + def scalar_one_or_none(self): + return '{"cookies": [{"name": "sessionid", "value": "ok"}]}' + + db = SimpleNamespace( + execute=AsyncMock(return_value=_ScalarResult()), + close=AsyncMock(), + ) + worker = DouyinWorker(account_id=499) + worker.get_db = AsyncMock(return_value=db) + + storage = await worker._load_storage_state() + + self.assertEqual(storage["cookies"][0]["value"], "ok") + statement = db.execute.await_args.args[0] + selected_names = [ + item.get("name") for item in statement.column_descriptions + ] + self.assertEqual(selected_names, ["cookie_data"]) + self.assertNotIn("im_session_data", str(statement).lower()) + + async def test_direct_service_marks_worker_ready_after_initialization(self): + worker = DouyinWorker(account_id=500, login_mode="im_direct") + worker._refresh_follow_welcome_config = AsyncMock( + return_value=(False, "", "sec-user") + ) + worker.get_reply_delay = AsyncMock(return_value=None) + fake_service = SimpleNamespace(run=AsyncMock(), stop=AsyncMock()) + + async def complete_initialization(): + service_factory.call_args.kwargs["on_ready"]() + + fake_service.run.side_effect = complete_initialization + session = DouyinImSession( + cookies={"sessionid": "session"}, + my_uid=0, + ) + + with patch( + "rpa_engine.playwright_worker.DouyinImService", + return_value=fake_service, + ) as service_factory: + await worker._run_im_direct_service(session) + + await worker.wait_until_ready() + worker._refresh_follow_welcome_config.assert_awaited_once_with(force=True) + fake_service.run.assert_awaited_once_with() + fake_service.stop.assert_awaited_once_with() + + async def test_direct_start_failure_releases_readiness_waiter(self): + worker = DouyinWorker(account_id=501, login_mode="im_direct") + worker._load_storage_state = AsyncMock(return_value=None) + worker.update_account_status = AsyncMock() + worker.cleanup = AsyncMock() + + await worker.start() + + with self.assertRaisesRegex(RuntimeError, "未保存 Cookie"): + await worker.wait_until_ready() + + if worker._task: + await worker._task + worker.update_account_status.assert_awaited_once_with( + "error", + error_msg="未保存 Cookie，无法直连 IM", + ) + + async def test_prevalidated_start_skips_duplicate_remote_validation(self): + worker = DouyinWorker( + account_id=502, + login_mode="im_direct", + credential_prevalidated=True, + ) + session = DouyinImSession( + cookies={"sessionid": "session"}, + my_uid=10001, + keys_str=json.dumps({"ec_privateKey": "private"}), + web_protect_str=json.dumps( + { + "ticket": "ticket", + "ts_sign": "sign", + "client_cert": "certificate", + } + ), + ) + worker._load_user_agent = AsyncMock(return_value="test-agent") + worker._build_im_session_from_storage = AsyncMock(return_value=session) + worker._require_sec_user_id = AsyncMock(return_value="sec-user") + worker._persist_im_session = AsyncMock() + worker._run_im_direct_service = AsyncMock() + + with patch( + "rpa_engine.playwright_worker.validate_im_session", + new_callable=AsyncMock, + ) as validate: + started, reason = await worker._try_cookie_only_im_start( + {"cookies": []} + ) + + self.assertTrue(started) + self.assertEqual(reason, "") + validate.assert_not_awaited() + worker._require_sec_user_id.assert_awaited_once() + worker._run_im_direct_service.assert_awaited_once_with(session) + + async def test_disabled_follow_welcome_uses_cached_lightweight_config(self): + class _Result: + def first(self): + return False, "", "sec-user" + + db = SimpleNamespace( + execute=AsyncMock(return_value=_Result()), + close=AsyncMock(), + ) + worker = DouyinWorker(account_id=503) + worker.get_db = AsyncMock(return_value=db) + + first = await worker._refresh_follow_welcome_config() + second = await worker._refresh_follow_welcome_config() + + self.assertEqual(first, (False, "", "sec-user")) + self.assertEqual(second, first) + worker.get_db.assert_awaited_once_with() + db.execute.assert_awaited_once() + db.close.assert_awaited_once_with() + + # The account PUT endpoint calls this synchronous hook so enabling the + # feature does not wait for the disabled-account ten-minute TTL. + worker.invalidate_follow_welcome_config() + await worker._refresh_follow_welcome_config() + self.assertEqual(worker.get_db.await_count, 2) + self.assertEqual(db.execute.await_count, 2) + + async def test_disabled_follow_tick_does_not_read_follower_log(self): + worker = DouyinWorker(account_id=504) + worker._im_service = SimpleNamespace(session=object()) + worker._follow_config_loaded = True + worker._refresh_follow_welcome_config = AsyncMock( + return_value=(False, "", "sec-user") + ) + worker.get_db = AsyncMock() + worker._require_sec_user_id = AsyncMock() + + await worker.follow_welcome_tick() + + worker._refresh_follow_welcome_config.assert_awaited_once_with() + worker.get_db.assert_not_awaited() + worker._require_sec_user_id.assert_not_awaited() + + async def test_missing_cached_sec_user_id_stops_hosting(self): + worker = DouyinWorker(account_id=505) + worker._im_service = SimpleNamespace(session=object()) + worker._follow_config_loaded = True + worker._refresh_follow_welcome_config = AsyncMock( + return_value=(False, "", "") + ) + worker._require_sec_user_id = AsyncMock(return_value="") + + await worker.follow_welcome_tick() + + worker._require_sec_user_id.assert_awaited_once_with( + "托管运行中" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backend/tests/test_ws_client_scaling.py b/backend/tests/test_ws_client_scaling.py new file mode 100644 index 0000000..c5f1177 --- /dev/null +++ b/backend/tests/test_ws_client_scaling.py @@ -0,0 +1,551 @@ +from __future__ import annotations + +import asyncio +import os +import sys +import threading +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, Mock, patch + +from websockets.legacy.server import serve + + +BACKEND_DIR = Path(__file__).resolve().parents[1] +if str(BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(BACKEND_DIR)) + +from rpa_engine.douyin_im import ws_client as ws_module +from rpa_engine.douyin_im.session import DouyinImSession +from rpa_engine.douyin_im.ws_client import DouyinImWsClient, _reconnect_delay + + +TEST_WS_URL = "wss://frontier-im.douyin.com/ws/v2?token=test-token-value" + + +class _FakeWebSocket: + def __init__(self, frames=()): + self.frames = list(frames) + self.next_calls = 0 + self.close_code = 1000 + self.close_reason = "test complete" + self.close_calls: list[tuple[int, str]] = [] + self.fail_calls = 0 + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + def __aiter__(self): + return self + + async def __anext__(self): + self.next_calls += 1 + if not self.frames: + raise StopAsyncIteration + return self.frames.pop(0) + + async def close(self, code=1000, reason=""): + self.close_calls.append((code, reason)) + + def fail_connection(self): + self.fail_calls += 1 + + +class WebSocketScalingTests(unittest.IsolatedAsyncioTestCase): + def _make_client(self, handler=None, account_id=23): + session = DouyinImSession( + cookies={"sessionid": "session-value", "sid_tt": "sid-value"}, + ws_urls=[TEST_WS_URL], + user_agent="test-agent/1.0", + ) + return DouyinImWsClient( + session, + handler or AsyncMock(), + account_id=account_id, + ) + + async def test_async_connection_preserves_handshake_and_ping_options(self): + received: list[bytes] = [] + + async def handler(item): + received.append(item["payload"]) + + client = self._make_client(handler) + client._running = True + fake_ws = _FakeWebSocket([b"binary-frame", "text-frame"]) + connect_mock = Mock(return_value=fake_ws) + + with ( + patch.object(ws_module, "websocket_connect", connect_mock), + patch.object( + ws_module, + "parse_ws_payload", + side_effect=lambda payload: [{"payload": payload}], + ), + patch.object(ws_module.system_logger, "record"), + ): + await client._run_connection(TEST_WS_URL) + queue = client._message_queue + self.assertIsNotNone(queue) + await asyncio.wait_for(queue.join(), timeout=0.5) + await client.stop() + + self.assertEqual(received, [b"binary-frame", b"text-frame"]) + kwargs = connect_mock.call_args.kwargs + self.assertEqual(kwargs["origin"], "https://www.douyin.com") + self.assertEqual(kwargs["subprotocols"], ["binary", "base64", "pbbp2"]) + self.assertEqual(kwargs["user_agent_header"], "test-agent/1.0") + self.assertEqual(kwargs["ping_interval"], 20) + self.assertEqual(kwargs["ping_timeout"], ws_module._PING_TIMEOUT_SECONDS) + self.assertEqual(kwargs["max_queue"], ws_module._TRANSPORT_MAX_QUEUE) + self.assertEqual(kwargs["max_size"], ws_module._INCOMING_MAX_SIZE) + headers = dict(kwargs["extra_headers"]) + self.assertEqual(headers["Cookie"], "sessionid=session-value; sid_tt=sid-value") + self.assertNotIn("Sec-WebSocket-Protocol", headers) + self.assertFalse(client.connected) + self.assertIsNone(client._connection) + + async def test_starting_500_clients_does_not_create_os_threads(self): + parked = asyncio.Event() + + async def parked_run_loop(_client, _url): + await parked.wait() + + clients = [self._make_client(account_id=index + 1) for index in range(500)] + before_threads = threading.active_count() + with patch.object(DouyinImWsClient, "_run_loop", parked_run_loop): + await asyncio.gather(*(client.start() for client in clients)) + await asyncio.sleep(0) + self.assertEqual(threading.active_count(), before_threads) + self.assertEqual(sum(client._task is not None for client in clients), 500) + self.assertEqual( + sum(client._dispatcher_task is not None for client in clients), + 500, + ) + await asyncio.gather(*(client.stop() for client in clients)) + + self.assertTrue(all(client._task is None for client in clients)) + self.assertTrue(all(client._dispatcher_task is None for client in clients)) + + async def test_global_handler_concurrency_is_shared_across_clients(self): + limit = 3 + release_handlers = asyncio.Event() + limit_reached = asyncio.Event() + active = 0 + maximum_active = 0 + started = 0 + completed = 0 + + async def handler(_item): + nonlocal active, maximum_active, started, completed + active += 1 + started += 1 + maximum_active = max(maximum_active, active) + if started == limit: + limit_reached.set() + try: + await release_handlers.wait() + finally: + active -= 1 + completed += 1 + + clients = [ + self._make_client(handler, account_id=index + 1000) + for index in range(12) + ] + queues = [] + with patch.dict( + os.environ, + {ws_module._HANDLER_CONCURRENCY_ENV: str(limit)}, + ): + try: + for index, client in enumerate(clients): + client._running = True + client._ensure_dispatcher() + queue = client._message_queue + self.assertIsNotNone(queue) + queues.append(queue) + await queue.put({"index": index}) + + await asyncio.wait_for(limit_reached.wait(), timeout=0.5) + # Give every other account a chance to contend for the same + # process/event-loop-wide semaphore. + await asyncio.sleep(0) + self.assertEqual(started, limit) + self.assertEqual(maximum_active, limit) + + release_handlers.set() + await asyncio.wait_for( + asyncio.gather(*(queue.join() for queue in queues)), + timeout=1.0, + ) + finally: + release_handlers.set() + await asyncio.gather(*(client.stop() for client in clients)) + + self.assertEqual(completed, len(clients)) + self.assertEqual(maximum_active, limit) + + async def test_global_handler_limit_preserves_single_client_fifo(self): + received: list[int] = [] + + async def handler(item): + await asyncio.sleep(0) + received.append(item["sequence"]) + + client = self._make_client(handler, account_id=2001) + client._running = True + client._ensure_dispatcher() + queue = client._message_queue + self.assertIsNotNone(queue) + for sequence in range(20): + await queue.put({"sequence": sequence}) + await asyncio.wait_for(queue.join(), timeout=0.5) + await client.stop() + + self.assertEqual(received, list(range(20))) + + async def test_stop_cancels_dispatcher_waiting_for_global_handler_slot(self): + holder_started = asyncio.Event() + release_holder = asyncio.Event() + waiter_handler = AsyncMock() + + async def holder_handler(_item): + holder_started.set() + await release_holder.wait() + + holder = self._make_client(holder_handler, account_id=3001) + waiter = self._make_client(waiter_handler, account_id=3002) + with patch.dict( + os.environ, + {ws_module._HANDLER_CONCURRENCY_ENV: "1"}, + ): + holder._running = True + holder._ensure_dispatcher() + holder_queue = holder._message_queue + self.assertIsNotNone(holder_queue) + await holder_queue.put({"id": "holder"}) + await asyncio.wait_for(holder_started.wait(), timeout=0.5) + + waiter._running = True + waiter._ensure_dispatcher() + waiter_queue = waiter._message_queue + self.assertIsNotNone(waiter_queue) + await waiter_queue.put({"id": "waiter"}) + + state = ws_module._get_loop_state() + + async def wait_until_slot_has_waiter(): + while not state.handler_slots._waiters: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_until_slot_has_waiter(), timeout=0.5) + joined = asyncio.create_task(waiter_queue.join()) + await asyncio.wait_for(waiter.stop(), timeout=0.5) + await asyncio.wait_for(joined, timeout=0.5) + + waiter_handler.assert_not_awaited() + self.assertTrue(waiter_queue.empty()) + + release_holder.set() + await asyncio.wait_for(holder_queue.join(), timeout=0.5) + await holder.stop() + + async def test_bounded_dispatch_queue_keeps_receiver_responsive(self): + first_handler_started = asyncio.Event() + release_first_handler = asyncio.Event() + received: list[bytes] = [] + + async def handler(item): + received.append(item["payload"]) + if len(received) == 1: + first_handler_started.set() + await release_first_handler.wait() + + client = self._make_client(handler) + client._running = True + fake_ws = _FakeWebSocket([b"first", b"second", b"third", b"fourth"]) + + async def wait_until_third_frame_is_read(): + while fake_ws.next_calls < 3: + await asyncio.sleep(0) + + with ( + patch.object(ws_module, "_APPLICATION_QUEUE_SIZE", 1), + patch.object(ws_module, "websocket_connect", return_value=fake_ws), + patch.object( + ws_module, + "parse_ws_payload", + side_effect=lambda payload: [{"payload": payload}], + ), + patch.object(ws_module.system_logger, "record"), + ): + task = asyncio.create_task(client._run_connection(TEST_WS_URL)) + await asyncio.wait_for(first_handler_started.wait(), timeout=0.5) + await asyncio.wait_for(wait_until_third_frame_is_read(), timeout=0.5) + + # The receiver keeps consuming while the business handler is + # blocked, but stops after the bounded application queue fills. + # It has pulled the third frame and is blocked enqueueing it; the + # fourth frame hasn't been requested and memory remains bounded. + self.assertEqual(fake_ws.next_calls, 3) + self.assertEqual(received, [b"first"]) + self.assertEqual(client._message_queue.qsize(), 1) + self.assertIsNotNone(client._dispatcher_task) + + release_first_handler.set() + await asyncio.wait_for(task, timeout=0.5) + queue = client._message_queue + self.assertIsNotNone(queue) + await asyncio.wait_for(queue.join(), timeout=0.5) + await client.stop() + + self.assertEqual(received, [b"first", b"second", b"third", b"fourth"]) + self.assertEqual(fake_ws.next_calls, 5) # four frames + end-of-stream + + async def test_real_async_handshake_receives_binary_frame(self): + received: list[bytes] = [] + request: dict[str, str | None] = {} + + async def handler(item): + received.append(item["payload"]) + + async def server_handler(websocket, _path): + request["origin"] = websocket.request_headers.get("Origin") + request["cookie"] = websocket.request_headers.get("Cookie") + request["user_agent"] = websocket.request_headers.get("User-Agent") + request["subprotocol"] = websocket.subprotocol + await websocket.send(b"protobuf-frame") + await websocket.close(code=1000, reason="test complete") + + client = self._make_client(handler) + client._running = True + with ( + patch.object( + ws_module, + "parse_ws_payload", + side_effect=lambda payload: [{"payload": payload}], + ), + patch.object(ws_module.system_logger, "record"), + ): + async with serve( + server_handler, + "127.0.0.1", + 0, + origins=["https://www.douyin.com"], + subprotocols=["pbbp2"], + ) as server: + port = server.sockets[0].getsockname()[1] + await asyncio.wait_for( + client._run_connection(f"ws://127.0.0.1:{port}"), + timeout=1.0, + ) + queue = client._message_queue + self.assertIsNotNone(queue) + await asyncio.wait_for(queue.join(), timeout=0.5) + await client.stop() + + self.assertEqual(received, [b"protobuf-frame"]) + self.assertEqual(request["origin"], "https://www.douyin.com") + self.assertEqual(request["cookie"], "sessionid=session-value; sid_tt=sid-value") + self.assertEqual(request["user_agent"], "test-agent/1.0") + self.assertEqual(request["subprotocol"], "pbbp2") + + async def test_stop_cancels_slow_handler_and_drains_pending_messages(self): + handler_started = asyncio.Event() + handler_cancelled = asyncio.Event() + never_release = asyncio.Event() + + async def handler(_item): + handler_started.set() + try: + await never_release.wait() + except asyncio.CancelledError: + handler_cancelled.set() + raise + + client = self._make_client(handler) + client._running = True + client._ensure_dispatcher() + queue = client._message_queue + self.assertIsNotNone(queue) + await queue.put({"id": 1}) + await queue.put({"id": 2}) + await asyncio.wait_for(handler_started.wait(), timeout=0.5) + joined = asyncio.create_task(queue.join()) + + await asyncio.wait_for(client.stop(), timeout=0.5) + await asyncio.wait_for(joined, timeout=0.5) + + self.assertTrue(handler_cancelled.is_set()) + self.assertTrue(queue.empty()) + self.assertIsNone(client._message_queue) + self.assertIsNone(client._dispatcher_task) + + async def test_first_connection_uses_captured_url_without_refresh(self): + client = self._make_client(account_id=24) + client._running = True + refreshed_url = TEST_WS_URL + "&refreshed=1" + client._prepare_url = AsyncMock(return_value=refreshed_url) + connected_urls: list[str] = [] + + async def connection(url): + connected_urls.append(url) + client._last_connection_lifetime = 1.0 + if len(connected_urls) == 2: + client._running = False + + client._run_connection = connection + + with ( + patch.object(ws_module, "_reconnect_delay", return_value=0.0), + patch.object(ws_module.system_logger, "record"), + ): + await asyncio.wait_for(client._run_loop(TEST_WS_URL), timeout=0.5) + + self.assertEqual(connected_urls, [TEST_WS_URL, refreshed_url]) + client._prepare_url.assert_awaited_once_with(TEST_WS_URL) + + async def test_stop_closes_connection_and_cancels_receive_task(self): + client = self._make_client() + client._running = True + client.connected = True + fake_ws = _FakeWebSocket() + client._connection = fake_ws + task = asyncio.create_task(asyncio.sleep(30)) + client._task = task + + await client.stop() + + self.assertEqual(fake_ws.close_calls, [(1000, "client stopping")]) + self.assertTrue(task.cancelled()) + self.assertFalse(client.connected) + self.assertIsNone(client._connection) + self.assertIsNone(client._task) + + async def test_stop_aborts_connection_when_close_handshake_stalls(self): + client = self._make_client() + client._running = True + client.connected = True + fake_ws = _FakeWebSocket() + + async def stalled_close(code=1000, reason=""): + fake_ws.close_calls.append((code, reason)) + await asyncio.Event().wait() + + fake_ws.close = stalled_close + client._connection = fake_ws + task = asyncio.create_task(asyncio.sleep(30)) + client._task = task + + with patch.object(ws_module, "_CLOSE_GRACE_SECONDS", 0.01): + await asyncio.wait_for(client.stop(), timeout=0.2) + + self.assertEqual(fake_ws.fail_calls, 1) + self.assertTrue(task.cancelled()) + self.assertIsNone(client._connection) + + async def test_short_normal_closes_continue_exponential_retry(self): + client = self._make_client(account_id=41) + client._running = True + client._prepare_url = AsyncMock(return_value=TEST_WS_URL) + attempts = 0 + + async def short_connection(_url): + nonlocal attempts + attempts += 1 + client._last_connection_lifetime = 1.0 + if attempts == 4: + client._running = False + + client._run_connection = short_connection + + with ( + patch.object(ws_module, "_reconnect_delay", return_value=0.0) as delay, + patch.object(ws_module.system_logger, "record"), + ): + await asyncio.wait_for(client._run_loop(TEST_WS_URL), timeout=0.5) + + self.assertEqual([call.args[1] for call in delay.call_args_list], [1, 2, 3]) + + async def test_stable_connection_resets_retry_counter(self): + client = self._make_client(account_id=42) + client._running = True + client._prepare_url = AsyncMock(return_value=TEST_WS_URL) + lifetimes = [1.0, 1.0, ws_module._STABLE_CONNECTION_SECONDS + 1.0, 1.0] + attempts = 0 + + async def connection(_url): + nonlocal attempts + client._last_connection_lifetime = lifetimes[attempts] + attempts += 1 + if attempts == len(lifetimes): + client._running = False + + client._run_connection = connection + + with ( + patch.object(ws_module, "_reconnect_delay", return_value=0.0) as delay, + patch.object(ws_module.system_logger, "record"), + ): + await asyncio.wait_for(client._run_loop(TEST_WS_URL), timeout=0.5) + + self.assertEqual([call.args[1] for call in delay.call_args_list], [1, 2, 1]) + + async def test_repeated_connection_failures_throttle_system_logs_per_account(self): + client = self._make_client(account_id=4041) + client._running = True + client._prepare_url = AsyncMock(return_value=TEST_WS_URL) + attempts = 0 + + async def failing_connection(_url): + nonlocal attempts + attempts += 1 + if attempts == 6: + client._running = False + raise ConnectionError("frontier unavailable") + + client._run_connection = failing_connection + + with ( + patch.object(ws_module, "_reconnect_delay", return_value=0.0), + patch.object( + ws_module, + "_system_log_throttle_seconds", + return_value=300.0, + ), + patch.object(ws_module.system_logger, "record") as system_record, + patch.object(ws_module.logger, "warning") as ordinary_warning, + ): + await asyncio.wait_for(client._run_loop(TEST_WS_URL), timeout=0.5) + + # Ordinary diagnostics remain available for every live failure, while + # the database-backed system log receives one row for the repeated + # failure/reconnect cycle of this account. + self.assertEqual(attempts, 6) + self.assertEqual(ordinary_warning.call_count, 5) + self.assertEqual(system_record.call_count, 1) + self.assertEqual(system_record.call_args.kwargs["account_id"], 4041) + + def test_handler_concurrency_environment_value_is_safely_clamped(self): + with patch.dict(os.environ, {ws_module._HANDLER_CONCURRENCY_ENV: "0"}): + self.assertEqual(ws_module._handler_concurrency_limit(), 1) + with patch.dict(os.environ, {ws_module._HANDLER_CONCURRENCY_ENV: "500"}): + self.assertEqual(ws_module._handler_concurrency_limit(), 32) + with patch.dict(os.environ, {ws_module._HANDLER_CONCURRENCY_ENV: "invalid"}): + self.assertEqual(ws_module._handler_concurrency_limit(), 8) + + def test_reconnect_delay_is_bounded_and_spread_by_account(self): + waits = [_reconnect_delay(10, retry) for retry in range(1, 10)] + + self.assertGreater(waits[1], waits[0]) + self.assertGreater(waits[2], waits[1]) + self.assertTrue(all(2.0 <= wait < 90.0 for wait in waits)) + self.assertNotEqual(_reconnect_delay(10, 8), _reconnect_delay(11, 8)) + + +if __name__ == "__main__": + unittest.main() diff --git a/backend/utils/log_limits.py b/backend/utils/log_limits.py new file mode 100644 index 0000000..8414640 --- /dev/null +++ b/backend/utils/log_limits.py @@ -0,0 +1,101 @@ +"""Bound diagnostic payloads without changing live message processing.""" + +from __future__ import annotations + +import json +import os +from typing import Any + + +TRUNCATION_MARKER = "\n...[日志内容过长，已截断]" +_MESSAGE_FIELDS = ( + "type", + "url", + "uri", + "text", + "name", + "width", + "height", + "duration", + "sticker_id", + "mime_type", +) + + +def _env_limit(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)) + + +def truncate_text(value: Any, limit: int) -> str: + text = "" if value is None else str(value) + if len(text) <= limit: + return text + keep = max(0, limit - len(TRUNCATION_MARKER)) + return text[:keep] + TRUNCATION_MARKER + + +def bound_system_log_detail(value: Any) -> str: + return truncate_text( + value, + _env_limit("KEFU_SYSTEM_LOG_MAX_CHARS", 4096, 512, 65536), + ) + + +def _compact_media_message(value: Any) -> Any: + if not isinstance(value, dict): + return truncate_text(value, 2048) + compact: dict[str, Any] = {} + for key in _MESSAGE_FIELDS: + if key not in value: + continue + item = value[key] + compact[key] = truncate_text(item, 2048) if isinstance(item, str) else item + compact["_log_truncated"] = True + return compact + + +def bound_message_log_content(value: Any) -> str: + """Keep a valid compact media JSON payload when a log entry is oversized.""" + limit = _env_limit("KEFU_MESSAGE_LOG_MAX_CHARS", 16384, 2048, 262144) + text = "" if value is None else str(value) + if len(text) <= limit: + return text + + try: + parsed = json.loads(text) + except (TypeError, ValueError, json.JSONDecodeError): + return truncate_text(text, limit) + + if isinstance(parsed, dict) and parsed.get("type"): + compact = _compact_media_message(parsed) + elif isinstance(parsed, dict) and isinstance(parsed.get("messages"), list): + compact = { + "messages": [ + _compact_media_message(item) + for item in parsed["messages"][:20] + ], + "_log_truncated": True, + } + else: + return truncate_text(text, limit) + + encoded = json.dumps(compact, ensure_ascii=False, separators=(",", ":")) + return encoded if len(encoded) <= limit else truncate_text(encoded, limit) + + +def bound_raw_message_log_content(value: Any) -> str: + return truncate_text( + value, + _env_limit("KEFU_RAW_MESSAGE_LOG_MAX_CHARS", 32768, 4096, 262144), + ) + + +def bound_error_log_content(value: Any) -> str: + return truncate_text( + value, + _env_limit("KEFU_ERROR_LOG_MAX_CHARS", 4096, 512, 65536), + ) diff --git a/backend/utils/received_message_log.py b/backend/utils/received_message_log.py index 5eccf6f..5b9f4a6 100644 --- a/backend/utils/received_message_log.py +++ b/backend/utils/received_message_log.py @@ -8,6 +8,7 @@ from typing import Optional from models.database import AsyncSessionLocal from models.models import ReceivedMessageLog +from utils.log_limits import bound_raw_message_log_content logger = logging.getLogger("received_message_log") @@ -23,8 +24,9 @@ async def record_received_message( message_type: Optional[int] = None, server_message_id: Optional[str] = None, ) -> None: - # 原样落库：不做 strip / parse / serialize，空字符串也记录 - store_content = raw_content if raw_content is not None else "" + # The live message object remains untouched for matching/replying. Only + # this diagnostic copy is bounded before persistence. + store_content = bound_raw_message_log_content(raw_content) async with AsyncSessionLocal() as db: try: diff --git a/backend/utils/system_logger.py b/backend/utils/system_logger.py index 6cc0f2b..7542a3d 100644 --- a/backend/utils/system_logger.py +++ b/backend/utils/system_logger.py @@ -15,6 +15,8 @@ from collections import deque from datetime import datetime from typing import Optional +from .log_limits import bound_system_log_detail, truncate_text + logger = logging.getLogger("douyin_im.system") _VALID_LEVELS = ("info", "success", "warning", "error") @@ -43,14 +45,16 @@ def record( "account_id": account_id, "level": level, "category": category, - "event": str(event or ""), - "detail": str(detail or ""), + "event": truncate_text(event, 255), + "detail": bound_system_log_detail(detail), "created_at": datetime.utcnow().isoformat(), } _buffer.appendleft(entry) _pending.append(entry) - msg = f"[{category}] {event}" + (f" | {detail}" if detail else "") + msg = f"[{category}] {entry['event']}" + ( + f" | {entry['detail']}" if entry["detail"] else "" + ) if level == "error": logger.error(msg) elif level == "warning": diff --git a/frontend/src/views/Dashboard.vue b/frontend/src/views/Dashboard.vue index ffbee03..788c65a 100644 --- a/frontend/src/views/Dashboard.vue +++ b/frontend/src/views/Dashboard.vue @@ -69,16 +69,21 @@ const fetchRecentLogs = async () => { } let statsInterval = null +const refreshVisibleStats = () => { + if (document.visibilityState === 'visible') fetchStats() +} onMounted(() => { fetchStats() fetchRecentLogs() - // 统计卡片每10秒自动刷新一次 - statsInterval = setInterval(fetchStats, 10000) + // 汇总统计需要扫描日志索引，分钟级刷新足够且不会持续压迫数据库。 + statsInterval = setInterval(refreshVisibleStats, 60000) + document.addEventListener('visibilitychange', refreshVisibleStats) }) onUnmounted(() => { if (statsInterval) clearInterval(statsInterval) + document.removeEventListener('visibilitychange', refreshVisibleStats) }) diff --git a/frontend/src/views/Logs.vue b/frontend/src/views/Logs.vue index ca51650..7ed090b 100644 --- a/frontend/src/views/Logs.vue +++ b/frontend/src/views/Logs.vue @@ -172,7 +172,7 @@ const onConvListScroll = (e) => { const fetchAccounts = async () => { try { - const res = await api.get('/accounts') + const res = await api.get('/account-options') accounts.value = Array.isArray(res.data) ? res.data : res.data?.items || [] } catch (error) { console.error(error) diff --git a/frontend/src/views/Messages.vue b/frontend/src/views/Messages.vue index 5d31659..3732581 100644 --- a/frontend/src/views/Messages.vue +++ b/frontend/src/views/Messages.vue @@ -39,7 +39,7 @@ const formatPeerId = (conv) => { } const fetchAccounts = async () => { - const res = await api.get(`/accounts`) + const res = await api.get(`/account-options`) accounts.value = res.data.filter(a => a.has_cookie) if (selectedAccount.value) { const current = accounts.value.find(a => a.id === selectedAccount.value) diff --git a/frontend/src/views/ReceivedMessages.vue b/frontend/src/views/ReceivedMessages.vue index 8b72f6e..840e981 100644 --- a/frontend/src/views/ReceivedMessages.vue +++ b/frontend/src/views/ReceivedMessages.vue @@ -72,7 +72,7 @@ const fetchLogs = async () => { const fetchAccounts = async () => { try { - const res = await api.get('/accounts') + const res = await api.get('/account-options') accounts.value = res.data } catch (error) { console.error(error) diff --git a/frontend/src/views/Rules.vue b/frontend/src/views/Rules.vue index d5c8466..3c286dd 100644 --- a/frontend/src/views/Rules.vue +++ b/frontend/src/views/Rules.vue @@ -219,7 +219,7 @@ watch(filterAccountId, () => { const fetchAccounts = async () => { try { - const res = await api.get('/accounts') + const res = await api.get('/account-options') accounts.value = Array.isArray(res.data) ? res.data : res.data?.items || [] if (!cooldownAccountId.value && accounts.value.length) { cooldownAccountId.value = filterAccountId.value diff --git a/frontend/src/views/SystemLogs.vue b/frontend/src/views/SystemLogs.vue index 6eda345..6fe7891 100644 --- a/frontend/src/views/SystemLogs.vue +++ b/frontend/src/views/SystemLogs.vue @@ -85,7 +85,7 @@ const fetchLogs = async () => { const fetchAccounts = async () => { try { - const res = await api.get(`/accounts`) + const res = await api.get(`/account-options`) accounts.value = res.data } catch (error) { console.error(error)