ZYT/dy
Template
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

更新

Browse Source
This commit is contained in:
Your Name
2026-07-28 15:04:17 +08:00
parent ac406a5f99
commit 8f68af1c2c
27 changed files with 3442 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -10,6 +10,7 @@ backend/**/__pycache__/
backend/kefu.db
backend/kefu.db-journal
backend/sessions/
/douyin.zip
*.env
.env.*
backend/main.py
+329 -41
View File
@@ -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])
backend/models/db_config.py
+61 -3
View File
@@ -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:
backend/models/db_migrate.py
+50
View File
@@ -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:
backend/models/models.py
+29 -3
View File
@@ -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):
"""账号额度购买订单。"""
backend/rpa_engine/credential.py
+3 -2
View File
@@ -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 是加密串，
backend/rpa_engine/douyin_im/protocol.py
+84 -5
View File
@@ -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 (
backend/rpa_engine/douyin_im/service.py
+455 -30
View File
@@ -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:
backend/rpa_engine/douyin_im/ws_client.py
+359 -143
View File
@@ -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()
backend/rpa_engine/playwright_worker.py
+242 -45
View File
@@ -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()
backend/tests/test_account_pagination.py
+156 -2
View File
@@ -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"),
backend/tests/test_batch_start_api.py
+214
View File
@@ -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")
backend/tests/test_conversation_poll_bandwidth.py
+292 -1
View File
@@ -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__":
backend/tests/test_credential_responsiveness.py
+3 -9
View File
@@ -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,
backend/tests/test_io_safety.py
+281
View File
@@ -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()
backend/tests/test_sec_user_id_guard.py
+17
View File
@@ -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=" ")
backend/tests/test_worker_scale_controls.py
+191
View File
@@ -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()
backend/tests/test_ws_client_scaling.py
+551
View File
@@ -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()
backend/utils/log_limits.py
+101
View File
@@ -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),
)
backend/utils/received_message_log.py
+4 -2
View File
@@ -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:
backend/utils/system_logger.py
+7 -3
View File
@@ -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":
frontend/src/views/Dashboard.vue
+7 -2
View File
@@ -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)
})
</script>
frontend/src/views/Logs.vue
+1 -1
View File
@@ -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)
frontend/src/views/Messages.vue
+1 -1
View File
@@ -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)
frontend/src/views/ReceivedMessages.vue
+1 -1
View File
@@ -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)
frontend/src/views/Rules.vue
+1 -1
View File
@@ -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
frontend/src/views/SystemLogs.vue
+1 -1
View File
@@ -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)