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
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()