7 Commits
Author SHA1 Message Date
gr 5db9674490 更新 2026-07-31 16:37:23 +08:00
Your Name 3fc94c4a89 更新 2026-07-30 10:06:53 +08:00
Your Name 8f68af1c2c 更新 2026-07-28 15:04:17 +08:00
Your Name ac406a5f99 更新 2026-07-28 11:49:36 +08:00
Your Name f99a4edf83 更新 2026-07-28 09:13:07 +08:00
Your Name 153db97dc7 更新 2026-07-28 09:00:19 +08:00
Your Name 8ba13a8ff9 更新 2026-07-27 15:20:15 +08:00
44 changed files with 5090 additions and 412 deletions
+1
View File
@@ -10,6 +10,7 @@ backend/**/__pycache__/
backend/kefu.db
backend/kefu.db-journal
backend/sessions/
/douyin.zip
*.env
.env.*
Binary file not shown.
Binary file not shown.
+600 -87
View File
@@ -20,7 +20,7 @@ from settings import CORS_ORIGINS, SERVE_WEB, STATIC_DIR
from help_pages import serve_credential_tool
from web_static import mount_frontend
from pydantic import BaseModel, Field
from sqlalchemy import select, update, delete, text, func
from sqlalchemy import select, update, delete, text, func, case, or_, cast, String
from sqlalchemy.ext.asyncio import AsyncSession
from models.database import engine, Base, get_db, AsyncSessionLocal
@@ -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()
@@ -427,6 +495,22 @@ async def _conversations_from_logs(
return list(seen.values())
async def _release_db_connection(db: AsyncSession) -> None:
"""Check this session's pooled connection back in before a long await.
The whole backend shares a handful of pooled connections (five by default
on SQLite). A session left open across credential validation or worker
initialization holds one of them for tens of seconds per account, so every
unrelated request then waits out ``pool_timeout`` and the panel looks
frozen during a bulk start. Committing ends the transaction and returns
the connection; ``expire_on_commit=False`` keeps loaded attributes usable.
"""
try:
await db.commit()
except Exception:
logger.debug("Releasing the database connection failed", exc_info=True)
async def _reset_account_credentials(account_id: int, db: AsyncSession) -> Account:
"""清除账号 Cookie、IM 会话等登录数据,并停止托管。"""
await manager.stop_worker(account_id)
@@ -501,10 +585,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 +594,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 +643,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 +794,34 @@ 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."""
total_accounts: int = 0
online_accounts: int = 0
my_accounts: int = 0
my_online_accounts: int = 0
class AccountUpdate(BaseModel):
phone: Optional[str] = None
username: Optional[str] = None
@@ -761,13 +958,18 @@ def _get_account_cookie_data(account: Account) -> Optional[str]:
return read_cookie_file(account.id)
async def _build_cookie_response(account: Account, cookie_data: Optional[str] = None) -> AccountCookieResponse:
async def _build_cookie_response(
account: Account,
cookie_data: Optional[str] = None,
*,
runtime_check: bool = True,
) -> AccountCookieResponse:
cookie_data = cookie_data if cookie_data is not None else _get_account_cookie_data(account)
summary = cookie_summary(cookie_data)
im_detail = await build_cookie_credential_detail(
cookie_data,
account.im_session_data,
runtime_check=True,
runtime_check=runtime_check,
)
return AccountCookieResponse(
account_id=account.id,
@@ -1010,6 +1212,51 @@ def _account_matches_keyword(account: Account, keyword: str) -> bool:
)
@app.get(
"/api/dashboard/account-stats",
response_model=DashboardAccountStatsResponse,
)
async def get_dashboard_account_stats(
db: AsyncSession = Depends(get_db),
user: User = Depends(get_current_user),
):
"""Return public platform totals without exposing other users' accounts."""
result = await db.execute(
select(
func.count(Account.id).label("total_accounts"),
func.coalesce(
func.sum(case((Account.status == "online", 1), else_=0)),
0,
).label("online_accounts"),
func.coalesce(
func.sum(case((Account.owner_id == user.id, 1), else_=0)),
0,
).label("my_accounts"),
func.coalesce(
func.sum(
case(
(
(Account.owner_id == user.id)
& (Account.status == "online"),
1,
),
else_=0,
)
),
0,
).label("my_online_accounts"),
)
)
row = result.one()
return DashboardAccountStatsResponse(
total_accounts=int(row.total_accounts or 0),
online_accounts=int(row.online_accounts or 0),
my_accounts=int(row.my_accounts or 0),
my_online_accounts=int(row.my_online_accounts or 0),
)
@app.get("/api/accounts")
async def get_accounts(
page: Optional[int] = Query(None, ge=1),
@@ -1026,42 +1273,91 @@ async def get_accounts(
支持 q(昵称/抖音ID/手机号/账号ID 搜索)与 status 筛选。
Cookie 解析等重逻辑只对当前页执行。
"""
result = await db.execute(accounts_for_user(user))
accounts = result.scalars().all()
# 更新内存中的运行状态与数据库同步,以防异常断开
for acc in accounts:
if page is None:
result = await db.execute(accounts_for_user(user))
accounts = result.scalars().all()
# 兼容旧的全量接口行为。
for acc in accounts:
is_running = manager.is_running(acc.id)
if is_running and acc.status == "offline":
acc.status = "online"
elif not is_running and acc.status in ("online", "logging_in", "starting"):
acc.status = "offline"
return [_build_account_response(acc) for acc in accounts]
keyword = (q or "").strip().lower()
status_filter = (status or "").strip()
stmt = accounts_for_user(user)
running_account_ids = [
int(account_id)
for account_id, worker in list(manager.workers.items())
if worker and worker.is_running
]
runtime_running = Account.id.in_(running_account_ids)
effective_status = case(
(
runtime_running & (Account.status == "offline"),
"online",
),
(
(~runtime_running)
& Account.status.in_(("online", "logging_in", "starting")),
"offline",
),
else_=Account.status,
)
if keyword:
like_value = f"%{keyword}%"
search_conditions = [
func.lower(Account.username).like(like_value),
func.lower(Account.douyin_uid).like(like_value),
func.lower(Account.phone).like(like_value),
func.lower(Account.user_agent).like(like_value),
cast(Account.id, String).like(like_value),
]
matching_profiles = [
profile
for profile in list_device_profiles()
if keyword in str(profile.get("label") or "").lower()
or keyword in str(profile.get("platform") or "").lower()
]
if matching_profiles:
matching_uas = [profile["user_agent"] for profile in matching_profiles]
search_conditions.append(Account.user_agent.in_(matching_uas))
if any(profile.get("id") == "chrome_win120" for profile in matching_profiles):
search_conditions.append(Account.user_agent.is_(None))
stmt = stmt.where(or_(*search_conditions))
if status_filter and status_filter != "all":
if status_filter == "quota_disabled":
stmt = stmt.where(Account.quota_disabled.is_(True))
else:
stmt = stmt.where(
Account.quota_disabled.is_not(True),
effective_status == status_filter,
)
total_result = await db.execute(
stmt.with_only_columns(func.count(Account.id)).order_by(None)
)
total = int(total_result.scalar_one() or 0)
result = await db.execute(
stmt.order_by(Account.id.asc())
.offset((page - 1) * page_size)
.limit(page_size)
)
items = result.scalars().all()
# Only reconcile the current page. The old implementation hydrated every
# account including large Cookie/IM blobs on every refresh, which became
# visibly slow beyond a few hundred accounts.
for acc in items:
is_running = manager.is_running(acc.id)
if is_running and acc.status == "offline":
acc.status = "online"
elif not is_running and acc.status in ("online", "logging_in", "starting"):
acc.status = "offline"
if page is None:
return [_build_account_response(acc) for acc in accounts]
keyword = (q or "").strip().lower()
status_filter = (status or "").strip()
filtered = [
acc
for acc in accounts
if _account_matches_keyword(acc, keyword)
and (
not status_filter
or status_filter == "all"
or (
status_filter == "quota_disabled"
and bool(acc.quota_disabled)
)
or (
status_filter != "quota_disabled"
and not acc.quota_disabled
and acc.status == status_filter
)
)
]
total = len(filtered)
start = (page - 1) * page_size
items = filtered[start : start + page_size]
return {
"items": [_build_account_response(acc) for acc in items],
"total": total,
@@ -1069,6 +1365,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,
@@ -1248,6 +1621,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
@@ -1271,19 +1648,65 @@ 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)
@app.get("/api/reply-queues", response_model=ReplyQueueSummaryResponse)
async def get_reply_queue_summaries(
account_ids: Optional[str] = None,
db: AsyncSession = Depends(get_db),
user: User = Depends(get_current_user),
):
"""聚合当前用户可见账号的排队数量,供账号页低频轮询"""
allowed_ids = await owned_account_ids(db, user)
"""聚合可见账号的排队数量账号页可限定为当前页账号"""
requested_ids: list[int] | None = None
if account_ids is not None:
requested_ids = []
seen_ids: set[int] = set()
for raw in str(account_ids).split(","):
token = raw.strip()
if not token:
continue
try:
account_id = int(token)
except (TypeError, ValueError):
raise HTTPException(status_code=400, detail="账号编号格式错误")
if account_id > 0 and account_id not in seen_ids:
seen_ids.add(account_id)
requested_ids.append(account_id)
if len(requested_ids) > 100:
raise HTTPException(status_code=400, detail="单次最多查询 100 个账号的回复队列")
allowed_ids: set[int] | None
if requested_ids is None:
allowed_ids = await owned_account_ids(db, user)
worker_entries = list(manager.workers.items())
else:
if is_admin(user.role):
allowed_ids = set(requested_ids)
elif requested_ids:
owned_result = await db.execute(
select(Account.id).where(
Account.owner_id == user.id,
Account.id.in_(requested_ids),
)
)
allowed_ids = {int(row[0]) for row in owned_result.all()}
else:
allowed_ids = set()
worker_entries = [
(account_id, manager.workers.get(account_id))
for account_id in requested_ids
if account_id in allowed_ids and account_id in manager.workers
]
summaries: list[ReplyQueueSummaryItem] = []
total_pending = 0
for account_id, worker in list(manager.workers.items()):
for account_id, worker in worker_entries:
if allowed_ids is not None and account_id not in allowed_ids:
continue
service = worker._im_service if worker else None
@@ -1464,7 +1887,17 @@ async def get_account_cookie(
cookie_data = _get_account_cookie_data(account)
if cookie_data and purpose != "management":
await _require_desktop_login_sec_user_id(account, db)
return await _build_cookie_response(account, cookie_data)
return await _build_cookie_response(
account,
cookie_data,
# Opening the edit dialog is a read-only management action. A live
# Douyin credential probe can take many seconds and, with hundreds of
# hosted accounts, would wait behind recurring background traffic.
# The dialog only needs the locally stored credential fields; users
# can still request an authoritative online check with the explicit
# "recheck credential" action.
runtime_check=purpose != "management",
)
@app.get(
@@ -1514,6 +1947,7 @@ async def update_account_cookie(
# are committed so no worker can start in the stop/commit gap.
async with manager.preparation_lock(account_id):
account = await get_owned_account(db, user, account_id, write=True)
await _release_db_connection(db)
await batch_start_queue.cancel_account(account_id)
await manager.stop_worker(account_id)
@@ -1553,6 +1987,7 @@ async def delete_account_cookie(
# before the cleared credentials are committed.
async with manager.preparation_lock(account_id):
account = await get_owned_account(db, user, account_id, write=True)
await _release_db_connection(db)
await batch_start_queue.cancel_account(account_id)
await manager.stop_worker(account_id)
@@ -1616,6 +2051,7 @@ async def delete_account(
):
await get_owned_account(db, user, account_id, write=True)
# 停止运行中的任务
await _release_db_connection(db)
await batch_start_queue.cancel_account(account_id)
await manager.stop_worker(account_id)
clear_cookie_file(account_id)
@@ -1634,7 +2070,11 @@ async def validate_account_credential(
account = await get_owned_account(db, user, account_id)
cookie_data = _get_account_cookie_data(account)
assessment = await assess_account_credential(cookie_data, account.im_session_data)
assessment = await assess_account_credential(
cookie_data,
account.im_session_data,
startup_priority=True,
)
return CredentialValidateResponse(**assessment)
@@ -1645,6 +2085,7 @@ async def reset_account_credentials(
user: User = Depends(require_write),
):
await get_owned_account(db, user, account_id, write=True)
await _release_db_connection(db)
await batch_start_queue.cancel_account(account_id)
account = await _reset_account_credentials(account_id, db)
return {
@@ -1657,6 +2098,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)
@@ -1664,7 +2107,14 @@ async def _start_account_rpa_impl(
return {"status": "running", "message": "RPA worker is already running."}
cookie_data = _get_account_cookie_data(account)
assessment = await assess_account_credential(cookie_data, account.im_session_data)
# Credential assessment issues real network requests to Douyin, and a bulk
# start runs it for every queued account. Release the connection first.
await _release_db_connection(db)
assessment = await assess_account_credential(
cookie_data,
account.im_session_data,
startup_priority=True,
)
login_mode = requested_login_mode or assessment["login_mode"]
reset_performed = False
@@ -1686,10 +2136,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
@@ -1700,9 +2159,20 @@ async def _start_account_rpa_impl(
account.qr_code_base64 = None
account.error_message = None
await db.commit()
else:
# Nothing to persist, but the reads above may still hold a pooled
# connection, and waiting for readiness below takes seconds.
await _release_db_connection(db)
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 "启动托管失败"
@@ -1710,7 +2180,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 = "凭证已失效,已清除旧数据,正在打开浏览器重新登录..."
@@ -1720,7 +2192,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"],
@@ -1742,7 +2214,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:
@@ -1764,6 +2240,10 @@ async def start_account_rpa(
user: User = Depends(require_write),
):
account = await get_owned_account(db, user, account_id, write=True)
# Cancelling waits for an in-flight queued start, and the preparation lock
# waits for whichever start owns this account. Neither may keep a pooled
# connection checked out while it waits.
await _release_db_connection(db)
await batch_start_queue.cancel_account(account_id)
async with manager.preparation_lock(account_id):
return await _start_account_rpa_impl(account, db, body.login_mode)
@@ -1841,6 +2321,9 @@ async def stop_account_rpa(
):
account = await get_owned_account(db, user, account_id, write=True)
# Cancelling drains an in-flight queued start, which can take as long as
# the batch per-account deadline. Do not hold a pooled connection for it.
await _release_db_connection(db)
await batch_start_queue.cancel_account(account_id)
stopped = await manager.stop_worker(account_id)
# 强制将数据库中的状态重置为 offline
@@ -2050,21 +2533,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])
@@ -2078,14 +2567,25 @@ async def get_logs(
if account_id is not None:
await get_owned_account(db, user, account_id)
limit = max(1, min(int(limit or 50), 500))
stmt = (
# Fetch only primary keys while MySQL performs the cross-account sort.
# Selecting the full ORM row here includes MEDIUMTEXT/TEXT columns, which
# makes MySQL 5.7 materialize a huge on-disk temporary table for non-admin
# users. Under polling load that grew ibtmp1 by tens of gigabytes.
id_stmt = (
logs_for_user(user, account_id)
.with_only_columns(MessageLog.id)
.order_by(MessageLog.created_at.desc())
.offset(max(0, int(offset or 0)))
.limit(limit)
)
result = await db.execute(stmt)
return result.scalars().all()
ordered_ids = list((await db.execute(id_stmt)).scalars().all())
if not ordered_ids:
return []
rows = (
await db.execute(select(MessageLog).where(MessageLog.id.in_(ordered_ids)))
).scalars().all()
rows_by_id = {row.id: row for row in rows}
return [rows_by_id[row_id] for row_id in ordered_ids if row_id in rows_by_id]
@app.get("/api/received-messages", response_model=List[ReceivedMessageLogResponse])
@@ -2099,13 +2599,26 @@ async def get_received_messages(
if account_id is not None:
await get_owned_account(db, user, account_id)
limit = max(1, min(int(limit or 100), 500))
stmt = (
# Keep the global sort narrow for the same reason as /api/logs. raw_content
# can be large and must only be loaded after LIMIT has selected the IDs.
id_stmt = (
received_logs_for_user(user, account_id)
.with_only_columns(ReceivedMessageLog.id)
.order_by(ReceivedMessageLog.created_at.desc())
.limit(limit)
)
result = await db.execute(stmt)
return result.scalars().all()
ordered_ids = list((await db.execute(id_stmt)).scalars().all())
if not ordered_ids:
return []
rows = (
await db.execute(
select(ReceivedMessageLog).where(
ReceivedMessageLog.id.in_(ordered_ids)
)
)
).scalars().all()
rows_by_id = {row.id: row for row in rows}
return [rows_by_id[row_id] for row_id in ordered_ids if row_id in rows_by_id]
@app.get("/api/system-logs", response_model=List[SystemLogResponse])
+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:
+56
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,
@@ -106,6 +126,12 @@ def migrate_accounts_table(conn) -> None:
"douyin_uid",
{"default": "ALTER TABLE accounts ADD COLUMN douyin_uid VARCHAR(64)"},
)
add_index_if_missing(
conn,
"accounts",
"ix_accounts_owner_id",
("owner_id",),
)
def migrate_account_videos_table(conn) -> None:
@@ -124,6 +150,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:
+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):
"""账号额度购买订单。"""
+88 -7
View File
@@ -18,11 +18,44 @@ StartHandler = Callable[[int], Awaitable[dict[str, Any]]]
JobToken = tuple[str, int]
class _StartPreparationTimeout(Exception):
"""Internal marker for the queue's own per-account deadline."""
DEFAULT_CONCURRENCY = 6
MAX_CONCURRENCY = 32
def _configured_concurrency() -> int:
"""Admission width for account preparation.
Every account spends most of its startup waiting: for the shared network
lane, for a WebSocket handshake, for signing work in a thread. Admitting
only two at a time therefore left the network lane idle and made a fleet of
several hundred accounts take tens of minutes. Actual outbound traffic is
still capped by the traffic controller, so a wider admission window fills
the existing lane instead of adding load.
"""
try:
return max(1, min(8, int(os.getenv("KEFU_BATCH_START_CONCURRENCY", "2"))))
return max(
1,
min(
MAX_CONCURRENCY,
int(os.getenv("KEFU_BATCH_START_CONCURRENCY", str(DEFAULT_CONCURRENCY))),
),
)
except (TypeError, ValueError):
return 2
return DEFAULT_CONCURRENCY
def _configured_timeout_seconds() -> float:
try:
value = float(os.getenv("KEFU_BATCH_START_TIMEOUT_SECONDS", "90"))
except (TypeError, ValueError):
return 90.0
if value <= 0:
return 0.0
return max(5.0, min(600.0, value))
def _utc_now() -> str:
@@ -55,28 +88,40 @@ class BatchStartQueue:
handler: StartHandler,
concurrency: int | None = None,
max_batches: int = 100,
timeout_seconds: float | None = None,
) -> None:
self._handler = handler
self.concurrency = max(1, int(concurrency or _configured_concurrency()))
self.max_batches = max(10, int(max_batches or 100))
self.timeout_seconds = (
_configured_timeout_seconds()
if timeout_seconds is None
else max(0.0, float(timeout_seconds or 0.0))
)
self._queue: asyncio.Queue[JobToken] = asyncio.Queue()
self._pending_jobs: dict[int, JobToken] = {}
self._active_tasks: dict[int, tuple[JobToken, asyncio.Task]] = {}
self._batches: dict[str, _BatchRecord] = {}
self._workers: list[asyncio.Task] = []
self._worker_sequence = 0
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:
if self._stopping:
return
for index in range(self.concurrency):
# Top up to the configured width instead of only starting from
# zero. A worker that died on an unexpected error used to shrink
# the queue permanently, so later batches crawled through a single
# remaining worker with no way to recover short of a restart.
while len(self._workers) < self.concurrency:
self._worker_sequence += 1
self._workers.append(
asyncio.create_task(
self._worker(index + 1),
name=f"account-batch-start-{index + 1}",
self._worker(self._worker_sequence),
name=f"account-batch-start-{self._worker_sequence}",
)
)
@@ -157,7 +202,21 @@ class BatchStartQueue:
)
self._active_tasks[account_id] = (job_token, handler_task)
result = await handler_task
if self.timeout_seconds > 0:
try:
result = await asyncio.wait_for(
handler_task,
timeout=self.timeout_seconds,
)
except asyncio.TimeoutError as exc:
# wait_for cancels its task only when this queue's
# deadline expires. Preserve a TimeoutError raised by
# the handler itself as its real account failure.
if handler_task.cancelled():
raise _StartPreparationTimeout from exc
raise
else:
result = await handler_task
async with self._lock:
record = self._batches.get(batch_id)
if record:
@@ -171,6 +230,28 @@ class BatchStartQueue:
elapsed_seconds=round(time.monotonic() - started_at, 3),
)
record.updated_at = _utc_now()
except _StartPreparationTimeout:
elapsed = round(time.monotonic() - started_at, 3)
logger.warning(
"Batch start timed out account=%s worker=%s after %.1fs",
account_id,
worker_number,
self.timeout_seconds,
)
async with self._lock:
record = self._batches.get(batch_id)
if record:
item = record.items[account_id]
if item.get("status") != "cancelled":
item.update(
status="failed",
message=(
f"启动准备超过 {self.timeout_seconds:g} 秒,"
"已跳过并继续处理后续账号"
),
elapsed_seconds=elapsed,
)
record.updated_at = _utc_now()
except asyncio.CancelledError:
async with self._lock:
record = self._batches.get(batch_id)
+37 -18
View File
@@ -5,8 +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.http_client import DouyinImHttpClient
from rpa_engine.douyin_im.session import DouyinImSession
from utils.cookie_store import analyze_cookie
@@ -131,41 +129,57 @@ async def build_cookie_credential_detail(
async def validate_im_session(
session: DouyinImSession,
_bypass_global_limit: bool = False,
*,
startup_priority: bool = False,
) -> tuple[bool, str]:
if not _bypass_global_limit:
from rpa_engine.douyin_im.traffic_control import get_traffic_controller
controller = get_traffic_controller()
async with controller.background_slot(0, "credential validation"):
return await validate_im_session(session, _bypass_global_limit=True)
# Startup validation must not sit behind hundreds of recurring
# conversation polls. It still shares the same global concurrency
# cap, so this changes ordering without increasing bandwidth usage.
async with controller.background_slot(
0,
"credential validation",
startup=startup_priority,
):
return await validate_im_session(
session,
_bypass_global_limit=True,
startup_priority=startup_priority,
)
if not session.can_direct_im():
if not has_im_session_token(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_uiduid_tt 是加密串,
# int() 解析必然失败而回退到网络请求;该请求偶发失败会误判为“未就绪”)。
uid = session.my_uid or auth.get_uid()
uid = session.my_uid
if not uid:
# get_uid() may fall back to a synchronous HTTP request with a
# multi-second timeout. Keep that work off FastAPI's event loop
# so a manual credential recheck cannot freeze account editing or
# unrelated API requests.
uid = await asyncio.to_thread(auth.get_uid)
if not uid:
return False, "服务端未认可当前 Cookie(无法获取用户 UID)"
if not auth.is_sign_ready():
return False, "缺少 IM 签名密钥(web_protect/keys),请用浏览器登录补全"
session.my_uid = int(uid)
async with DouyinImHttpClient(session) as http:
await http.get_unread_count()
# 若已缓存到会话票据,优先校验其是否仍新鲜(最理想)。
if session.conv_meta:
ok, reason = await http.verify_messaging_capability(auth, session.my_uid)
if ok:
return True, reason
# 没有缓存会话票据是首次登录的正常情况:会话 ticket 会在发送时即时
# 创建/获取(resolve_conversation_meta),因此只要 Cookie + sessionid +
# 签名密钥(web_protect/keys) + UID 齐全,就视为可 IM 直连托管,不必再开浏览器。
return True, "IM 凭证就绪(Cookie 与签名密钥齐全,可直连托管)"
# unread_count and ticket probes were previously issued here, but
# neither result changed the final decision: unread failures become
# zero and a stale/missing ticket is resolved lazily at send time.
# Keeping those probes doubled large-batch startup traffic without
# adding an authoritative validation signal.
return True, "IM 凭证就绪(Cookie 与签名密钥齐全,可直连托管)"
except Exception as e:
logger.warning(f"IM session validation failed: {e}")
return False, f"IM 运行时验证失败: {e}"
@@ -174,6 +188,8 @@ async def validate_im_session(
async def assess_account_credential(
cookie_data: Optional[str],
im_session_data: Optional[str] = None,
*,
startup_priority: bool = False,
) -> dict:
cookie_info = analyze_cookie(cookie_data)
result = {
@@ -212,7 +228,10 @@ async def assess_account_credential(
result["should_reset"] = _should_reset_credentials(result)
return result
im_ok, im_reason = await validate_im_session(session)
im_ok, im_reason = await validate_im_session(
session,
startup_priority=startup_priority,
)
result["im_ready"] = im_ok
if im_ok:
result["can_skip_browser"] = True
+10 -2
View File
@@ -737,7 +737,7 @@ class DouyinImHttpClient:
pass
return total
async def get_conversations(self) -> list[dict]:
async def get_conversations(self, *, enrich_profiles: bool = True) -> list[dict]:
"""拉取会话列表,返回标准化会话"""
payloads = [
{"cursor": 0, "count": 50, "inbox_type": 0},
@@ -754,7 +754,12 @@ class DouyinImHttpClient:
if data is None:
data = await self._request("GET", "/v1/conversation/list", body)
if data is None:
continue
# Payload variants only help with schema compatibility. They
# cannot repair a network outage, so stop after POST + GET
# both fail instead of occupying a scarce global slot for up
# to four more full request timeouts.
logger.warning("Conversation poll transport failed; skipping payload fallbacks")
break
status_code = data.get("status_code") if isinstance(data, dict) else None
error_text = ""
@@ -814,6 +819,9 @@ class DouyinImHttpClient:
enriched: list[dict] = []
for item in conversations:
conv = enrich_conversation_item(item, my_uid)
if not enrich_profiles:
enriched.append(conv)
continue
peer_uid = str(conv.get("peer_uid") or "")
name = (conv.get("sender_name") or "").strip()
avatar = str(conv.get("sender_avatar") or "").strip()
+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 (
+16 -6
View File
@@ -87,20 +87,29 @@ class AccountReplyQueue:
details: Optional[dict[str, Any]] = None,
merge_key: str = "",
merge_keys: Optional[Iterable[str]] = None,
immediate_if_idle: bool = False,
) -> int:
"""Append one reply job and return its current 1-based queue position."""
"""Append one reply job and return its current 1-based queue position.
When ``immediate_if_idle`` is enabled, the first job in a completely
idle account queue reserves a zero-second slot. Jobs arriving behind
it still reserve the configured interval, so the normal per-account
pacing resumes from the second job onward.
"""
interval = max(0.0, float(delay_seconds or 0))
loop = asyncio.get_running_loop()
async with self._state_lock:
if not self._running or not self._task or self._task.done():
raise RuntimeError("reply queue is not running")
due_at = max(loop.time(), self._tail_due_at) + interval
queue_is_idle = self.pending_count == 0
slot_seconds = 0.0 if immediate_if_idle and queue_is_idle else interval
due_at = max(loop.time(), self._tail_due_at) + slot_seconds
self._tail_due_at = due_at
self._waiting.append(
_QueueItem(
job_id=uuid.uuid4().hex,
due_at=due_at,
slot_seconds=interval,
slot_seconds=slot_seconds,
callback=callback,
description=description,
queued_at=time.time(),
@@ -256,9 +265,10 @@ class AccountReplyQueue:
item = self._waiting.pop(selected_index)
shift_seconds = max(0.0, item.slot_seconds)
shifted_count = 0
for later in self._waiting[selected_index:]:
later.due_at -= shift_seconds
shifted_count += 1
if shift_seconds > 0:
for later in self._waiting[selected_index:]:
later.due_at -= shift_seconds
shifted_count += 1
item.due_at = asyncio.get_running_loop().time()
item.expedited = True
+558 -27
View File
@@ -1,6 +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
@@ -24,6 +29,280 @@ 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:
try:
return max(minimum, float(os.getenv(name, str(default))))
except (TypeError, ValueError):
return default
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
reconciliation when that path exists, so running it every 15 seconds for
hundreds of accounts wastes bandwidth and eventually starves new starts.
Accounts without WebSocket keep the original fast polling cadence.
"""
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:
@@ -43,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
@@ -53,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))
@@ -84,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}"
@@ -578,19 +870,32 @@ class DouyinImService:
"replies": list(replies),
},
merge_keys=queue_merge_keys,
immediate_if_idle=True,
)
scheduled_wait = 0 if position == 1 else delay_seconds
logger.info(
"Queued reply to %s for account %s: position=%s interval=%ss",
"Queued reply to %s for account %s: position=%s wait=%ss interval=%ss",
sender,
self.account_id,
position,
scheduled_wait,
delay_seconds,
)
if position == 1:
queue_detail = (
f"{sender} 是当前账号队列的首条任务,等待时间为 0 秒;"
f"后续任务仍按 {delay_seconds} 秒间隔排队。"
)
else:
queue_detail = (
f"{sender} 当前排在第 {position} 位;账号生效间隔为 {delay_seconds} 秒,"
"后续任务继续依次排队。"
)
system_logger.record(
"自动回复已进入账号队列",
detail=(
f"{sender} 当前排在第 {position} 位;账号生效间隔为 {delay_seconds} 秒,"
"账号内计时与排位独立;到点后再进入全局带宽队列逐条投递。"
f"{queue_detail} 账号内计时与排位独立;"
"发送时仍进入全局带宽队列逐条投递。"
),
level="info",
category="send",
@@ -687,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)
@@ -696,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"]
@@ -717,21 +1031,96 @@ 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:
unread_total = await http.get_unread_count()
if unread_total:
logger.info(f"IM unread total: {unread_total}")
conversations = await http.get_conversations()
await self._index_conversations(conversations)
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)
# 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
)
if unread_total:
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。
@@ -745,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)
@@ -763,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)
@@ -788,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}")
@@ -798,8 +1200,14 @@ 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}")
system_logger.record(
@@ -810,25 +1218,143 @@ class DouyinImService:
account_id=self.account_id,
)
loop_count = 0
# 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,
)
loop = asyncio.get_running_loop()
if initial_poll_succeeded:
initial_retry_interval = poll_interval
initial_retry_stagger = poll_stagger
else:
# If the authoritative first poll failed, retry on the fast HTTP
# cadence even when WebSocket connected in the meantime.
initial_retry_interval, initial_retry_stagger = (
_conversation_poll_timing(self.account_id, False)
)
next_conversation_poll_at = (
loop.time() + initial_retry_interval + initial_retry_stagger
)
logger.info(
"Conversation reconciliation account=%s interval=%.1fs stagger=%.1fs ws=%s",
self.account_id,
poll_interval,
poll_stagger,
"yes" if ws_connected else "no",
)
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:
if loop_count % 3 == 0:
await self._poll_conversations()
if loop_count % 6 == 0:
logger.info(f"IM direct tick #{loop_count} account={self.account_id}")
current_ws_connected = bool(
self._ws_client
and getattr(self._ws_client, "connected", False)
)
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,
)
candidate_poll_at = loop.time() + poll_interval + poll_stagger
# Never postpone an already scheduled reconciliation.
# In particular, reconnecting must preserve the earlier
# fallback poll that covers messages missed while offline.
next_conversation_poll_at = min(
next_conversation_poll_at,
candidate_poll_at,
)
logger.info(
"Conversation reconciliation rescheduled account=%s "
"interval=%.1fs ws=%s",
self.account_id,
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:
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 * (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(
@@ -841,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:
@@ -393,8 +393,27 @@ class TrafficController:
1.0,
)
)
self._background = asyncio.Semaphore(
_env_int("KEFU_BACKGROUND_NETWORK_CONCURRENCY", 2)
background_capacity = _env_int("KEFU_BACKGROUND_NETWORK_CONCURRENCY", 4)
self._background = asyncio.Semaphore(background_capacity)
# Recurring polls can create hundreds of waiters when many accounts
# are online. Admit at most one normal waiter to the semaphore at a
# time so startup validation can join near the front instead of being
# buried behind the entire polling backlog. The shared semaphore is
# still the single bandwidth cap; startup work does not add extra
# network concurrency.
self._background_normal_admission = asyncio.Lock()
self._background_startup_clear = asyncio.Event()
self._background_startup_clear.set()
# Starting several hundred accounts keeps startup waiters queued for
# many minutes on end. Leave one slot of the shared lane for recurring
# work so hosted accounts keep receiving messages during a bulk start
# instead of going silent until the last account is up.
self._background_startup = asyncio.Semaphore(
max(1, background_capacity - 1)
)
self._background_normal_max_defer = _env_float(
"KEFU_BACKGROUND_NORMAL_MAX_DEFER_SECONDS",
5.0,
)
self._browser = asyncio.Semaphore(
_env_int("KEFU_BROWSER_START_CONCURRENCY", 1)
@@ -410,12 +429,20 @@ class TrafficController:
)
self.background_waiting = 0
self.background_active = 0
self.background_startup_waiting = 0
self.background_startup_active = 0
self.browser_waiting = 0
self.browser_active = 0
self.media_proxy_active = 0
@asynccontextmanager
async def background_slot(self, account_id: int = 0, description: str = "request"):
async def background_slot(
self,
account_id: int = 0,
description: str = "request",
*,
startup: bool = False,
):
current_task = asyncio.current_task()
owner_task, depth = self._background_owner.get()
if owner_task is current_task and depth > 0:
@@ -429,12 +456,52 @@ class TrafficController:
started = asyncio.get_running_loop().time()
self.background_waiting += 1
try:
await self._background.acquire()
if startup:
self.background_startup_waiting += 1
self._background_startup_clear.clear()
startup_reservation = False
try:
await self._background_startup.acquire()
startup_reservation = True
await self._background.acquire()
except BaseException:
if startup_reservation:
self._background_startup.release()
raise
finally:
self.background_startup_waiting -= 1
if self.background_startup_waiting == 0:
self._background_startup_clear.set()
else:
# Only one recurring/background request may wait directly on
# the shared semaphore. A later startup request therefore
# has at most one normal request ahead of it, not hundreds.
async with self._background_normal_admission:
# Yield to pending startup work, but only while a startup
# waiter could still claim a slot, and never for longer
# than the deferral budget. A batch of several hundred
# accounts otherwise keeps startup waiters pending for the
# whole run, which stalled every recurring poll behind it.
if (
self._background_normal_max_defer > 0
and not self._background_startup_clear.is_set()
and not self._background_startup.locked()
):
try:
await asyncio.wait_for(
self._background_startup_clear.wait(),
timeout=self._background_normal_max_defer,
)
except asyncio.TimeoutError:
pass
await self._background.acquire()
except BaseException:
self.background_waiting -= 1
raise
self.background_waiting -= 1
self.background_active += 1
if startup:
self.background_startup_active += 1
token = self._background_owner.set((current_task, 1))
waited = asyncio.get_running_loop().time() - started
if waited >= 1.0:
@@ -448,6 +515,9 @@ class TrafficController:
yield
finally:
self._background_owner.reset(token)
if startup:
self.background_startup_active -= 1
self._background_startup.release()
self.background_active -= 1
self._background.release()
@@ -492,6 +562,8 @@ class TrafficController:
"background": {
"active": self.background_active,
"waiting": self.background_waiting,
"startup_active": self.background_startup_active,
"startup_waiting": self.background_startup_waiting,
},
"browser": {
"active": self.browser_active,
+362 -141
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 WebSocketwebsocket-client,与 DouYin_Spider 一致)"""
"""Async frontier-im WebSocket client with bounded message backpressure."""
def __init__(
self,
@@ -27,12 +121,16 @@ class DouyinImWsClient:
self.on_message = on_message
self.account_id = account_id
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")
@@ -45,186 +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
with self._ws_lock:
if self._ws_app:
self.connected = False
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):
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):
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:
with self._ws_lock:
if self._ws_app is ws_app:
self._ws_app = None
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
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()
+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()
+306
View File
@@ -0,0 +1,306 @@
from __future__ import annotations
import os
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
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))
import main
from models.database import Base
from models.models import Account, MessageLog
class _CountResult:
def __init__(self, count: int):
self.count = count
def scalar_one(self):
return self.count
class _RowsResult:
def __init__(self, rows):
self.rows = list(rows)
def scalars(self):
return self
def all(self):
return list(self.rows)
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"),
SimpleNamespace(id=11, status="online"),
]
db = SimpleNamespace(
execute=AsyncMock(
side_effect=[
_CountResult(392),
_RowsResult(page_rows),
]
)
)
with (
patch.object(main.manager, "is_running", side_effect=[False, True]),
patch.object(
main,
"_build_account_response",
side_effect=lambda account: {"id": account.id, "status": account.status},
) as build_response,
):
response = await main.get_accounts(
page=20,
page_size=20,
q=None,
status=None,
db=db,
user=SimpleNamespace(id=1, role="admin"),
)
self.assertEqual(db.execute.await_count, 2)
self.assertEqual(response["total"], 392)
self.assertEqual(response["page"], 20)
self.assertEqual(response["page_size"], 20)
self.assertEqual([item["id"] for item in response["items"]], [10, 11])
self.assertEqual(build_response.call_count, 2)
count_sql = str(db.execute.await_args_list[0].args[0]).upper()
page_sql = str(db.execute.await_args_list[1].args[0]).upper()
self.assertIn("COUNT", count_sql)
self.assertNotIn(" LIMIT ", count_sql)
self.assertIn(" LIMIT ", page_sql)
self.assertIn(" OFFSET ", page_sql)
async def test_status_filter_uses_effective_runtime_worker_state(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 = {
1001: SimpleNamespace(is_running=True),
1002: SimpleNamespace(is_running=False),
}
try:
async with session_factory() as db:
db.add_all(
[
Account(id=1001, status="offline"),
Account(id=1002, status="online"),
]
)
await db.commit()
with patch.object(
main,
"_build_account_response",
side_effect=lambda account: {
"id": account.id,
"status": account.status,
},
):
online = await main.get_accounts(
page=1,
page_size=20,
q=None,
status="online",
db=db,
user=SimpleNamespace(id=1, role="admin"),
)
await db.rollback()
db.expire_all()
offline = await main.get_accounts(
page=1,
page_size=20,
q=None,
status="offline",
db=db,
user=SimpleNamespace(id=1, role="admin"),
)
self.assertEqual(online["total"], 1)
self.assertEqual(online["items"], [{"id": 1001, "status": "online"}])
self.assertEqual(offline["total"], 1)
self.assertEqual(offline["items"], [{"id": 1002, "status": "offline"}])
finally:
main.manager.workers = original_workers
await engine.dispose()
if __name__ == "__main__":
unittest.main()
+263
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import os
import sys
import unittest
@@ -32,6 +33,268 @@ 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_holds_no_db_connection_while_it_waits(self):
"""A queued start must not pin one of the few pooled connections.
Credential validation and readiness waiting take seconds per account.
Holding a session open across them exhausted the pool during a bulk
start, so every unrelated request waited out ``pool_timeout``.
"""
account = SimpleNamespace(
id=505,
status="starting",
qr_code_base64=None,
error_message=None,
im_session_data="saved-session",
)
events: list[str] = []
assessment = {
"login_mode": "im_direct",
"should_reset": False,
"can_skip_browser": True,
"message": "ready",
"cookie_valid": True,
"im_ready": True,
}
async def commit():
events.append("release")
async def assess(*_args, **_kwargs):
events.append("assess")
return assessment
async def start_worker(*_args, **_kwargs):
events.append("start-worker")
return True
db = SimpleNamespace(commit=AsyncMock(side_effect=commit))
with (
patch.object(main.manager, "is_running", return_value=False),
patch.object(main.manager, "start_worker", AsyncMock(side_effect=start_worker)),
patch.object(main, "_get_account_cookie_data", return_value="{}"),
patch.object(main, "assess_account_credential", AsyncMock(side_effect=assess)),
):
await main._start_account_rpa_impl(account, db, wait_for_ready=True)
# "starting" was already persisted, so the only commits here exist to
# return the connection: one before validation, one before the wait.
self.assertEqual(events, ["release", "assess", "release", "start-worker"])
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")
+96 -2
View File
@@ -17,8 +17,18 @@ from rpa_engine import batch_start as batch_start_module
class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
def _make_queue(self, handler, *, concurrency: int = 2) -> BatchStartQueue:
queue = BatchStartQueue(handler, concurrency=concurrency)
def _make_queue(
self,
handler,
*,
concurrency: int = 2,
timeout_seconds: float | None = None,
) -> BatchStartQueue:
queue = BatchStartQueue(
handler,
concurrency=concurrency,
timeout_seconds=timeout_seconds,
)
self.addAsyncCleanup(queue.stop)
return queue
@@ -77,6 +87,42 @@ class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(maximum_active, 2)
self.assertEqual(active, 0)
async def test_default_concurrency_admits_more_than_two_accounts(self):
with patch.dict(os.environ, {}, clear=False):
os.environ.pop("KEFU_BATCH_START_CONCURRENCY", None)
queue = BatchStartQueue(lambda _account_id: None)
self.assertEqual(queue.concurrency, batch_start_module.DEFAULT_CONCURRENCY)
self.assertGreaterEqual(queue.concurrency, 4)
async def test_dead_worker_is_replaced_so_width_never_shrinks(self):
"""One crashed worker must not permanently narrow the queue.
Width used to be restored only when every worker had exited, so a
single unexpected worker death left later batches crawling through the
survivors until the process restarted.
"""
async def handler(account_id: int) -> dict:
return {"message": f"started-{account_id}"}
queue = self._make_queue(handler, concurrency=3)
first = await queue.submit([91])
await self._wait_for_complete(queue, first["batch_id"])
self.assertEqual(len(queue._workers), 3)
casualty = queue._workers[0]
casualty.cancel()
await asyncio.gather(casualty, return_exceptions=True)
second = await queue.submit([92])
await self._wait_for_complete(queue, second["batch_id"])
self.assertEqual(len(queue._workers), 3)
self.assertNotIn(casualty, queue._workers)
self.assertTrue(all(not task.done() for task in queue._workers))
names = [task.get_name() for task in queue._workers]
self.assertEqual(len(set(names)), 3)
async def test_submit_returns_while_handler_is_blocked(self):
handler_started = asyncio.Event()
release_handler = asyncio.Event()
@@ -168,6 +214,54 @@ class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(by_account[32]["status"], "submitted")
self.assertEqual(by_account[33]["status"], "submitted")
async def test_two_timeouts_release_both_workers_for_following_accounts(self):
never_release = asyncio.Event()
calls: list[int] = []
async def handler(account_id: int) -> dict:
calls.append(account_id)
if account_id in (71, 72):
await never_release.wait()
return {"message": f"started-{account_id}"}
queue = self._make_queue(
handler,
concurrency=2,
timeout_seconds=0.02,
)
with patch.object(batch_start_module.logger, "warning"):
submitted = await queue.submit([71, 72, 73, 74])
completed = await self._wait_for_complete(
queue,
submitted["batch_id"],
)
self.assertEqual(calls, [71, 72, 73, 74])
self.assertEqual(completed["failed_count"], 2)
self.assertEqual(completed["submitted_count"], 2)
by_account = {item["account_id"]: item for item in completed["items"]}
self.assertIn("已跳过并继续处理后续账号", by_account[71]["message"])
self.assertIn("已跳过并继续处理后续账号", by_account[72]["message"])
self.assertEqual(by_account[73]["status"], "submitted")
self.assertEqual(by_account[74]["status"], "submitted")
async def test_handler_timeout_error_keeps_its_original_detail(self):
async def handler(_account_id: int) -> dict:
raise asyncio.TimeoutError("upstream request timed out")
queue = self._make_queue(
handler,
concurrency=1,
timeout_seconds=10,
)
with patch.object(batch_start_module.logger, "exception"):
submitted = await queue.submit([75])
completed = await self._wait_for_complete(queue, submitted["batch_id"])
item = completed["items"][0]
self.assertEqual(item["status"], "failed")
self.assertEqual(item["message"], "upstream request timed out")
async def test_failed_account_can_be_submitted_again(self):
attempts = 0
@@ -1,10 +1,13 @@
from __future__ import annotations
import asyncio
import os
import sys
import unittest
from contextlib import asynccontextmanager
from pathlib import Path
from unittest.mock import AsyncMock
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch
BACKEND_DIR = Path(__file__).resolve().parents[1]
@@ -16,6 +19,8 @@ if str(BACKEND_DIR) not in sys.path:
from rpa_engine.douyin_im.http_client import DouyinImHttpClient
from rpa_engine.douyin_im.session import DouyinImSession
from rpa_engine.douyin_im.service import DouyinImService, _conversation_poll_timing
from rpa_engine.douyin_im import service as service_module
class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
@@ -75,6 +80,360 @@ class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
["POST", "GET"],
)
async def test_transport_outage_stops_after_one_post_and_get_pair(self):
client = self._make_client()
client._request = AsyncMock(return_value=None)
with self.assertLogs("douyin_im.http", level="WARNING"):
self.assertEqual(await client.get_conversations(), [])
self.assertEqual(client._request.await_count, 2)
self.assertEqual(
[call.args[0] for call in client._request.await_args_list],
["POST", "GET"],
)
def test_websocket_reconciliation_is_slow_and_http_fallback_stays_fast(self):
with patch.dict(
os.environ,
{
"KEFU_WS_RECONCILE_INTERVAL_SECONDS": "120",
"KEFU_HTTP_POLL_INTERVAL_SECONDS": "15",
},
):
ws_interval, ws_stagger = _conversation_poll_timing(123, True)
http_interval, http_stagger = _conversation_poll_timing(123, False)
self.assertEqual(ws_interval, 120)
self.assertEqual(http_interval, 15)
self.assertGreaterEqual(ws_stagger, 0)
self.assertLess(ws_stagger, ws_interval)
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
async def background_slot(self, *_args, **_kwargs):
yield
class _HttpClient:
def __init__(self):
self.get_conversations = AsyncMock(return_value=[])
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
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._index_conversations = AsyncMock()
with (
patch.object(service_module, "get_traffic_controller", return_value=_Controller()),
patch.object(service_module, "DouyinImHttpClient", return_value=http),
):
await service._poll_conversations()
http.get_conversations.assert_awaited_once_with(enrich_profiles=False)
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__":
unittest.main()
@@ -152,6 +152,9 @@ class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase):
[
"lock-enter",
"authorize",
# Checks the pooled connection back in before cancel/stop,
# which may wait on an in-flight start.
"commit",
"cancel",
"stop",
"write-cookie",
@@ -245,6 +248,9 @@ class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase):
[
"lock-enter",
"authorize",
# Checks the pooled connection back in before cancel/stop,
# which may wait on an in-flight start.
"commit",
"cancel",
"stop",
"clear-cookie-file",
@@ -0,0 +1,47 @@
import threading
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from rpa_engine.credential import validate_im_session
class CredentialResponsivenessTests(unittest.IsolatedAsyncioTestCase):
async def test_uid_lookup_does_not_block_event_loop(self):
event_loop_thread_id = threading.get_ident()
lookup_thread_ids = []
def get_uid():
lookup_thread_ids.append(threading.get_ident())
return 123456
session = SimpleNamespace(
my_uid=0,
can_direct_im=lambda: True,
)
auth = SimpleNamespace(
get_uid=get_uid,
is_sign_ready=lambda: True,
)
with patch(
"rpa_engine.credential.DouyinAuth.from_im_session",
return_value=auth,
):
result = await validate_im_session(
session,
_bypass_global_limit=True,
)
self.assertTrue(result[0])
self.assertEqual(len(lookup_thread_ids), 1)
self.assertNotEqual(
lookup_thread_ids[0],
event_loop_thread_id,
"the synchronous UID lookup ran on the event-loop thread",
)
self.assertEqual(session.my_uid, 123456)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,93 @@
from __future__ import annotations
import os
import sys
import unittest
from collections import namedtuple
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock
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))
import main
_AggregateRow = namedtuple(
"_AggregateRow",
(
"total_accounts",
"online_accounts",
"my_accounts",
"my_online_accounts",
),
)
class _AggregateResult:
def __init__(self, row):
self._row = row
def one(self):
return self._row
class DashboardAccountStatsTests(unittest.IsolatedAsyncioTestCase):
async def test_conditional_aggregate_maps_global_and_personal_counts(self):
row = _AggregateRow(
total_accounts=12,
online_accounts=5,
my_accounts=3,
my_online_accounts=2,
)
db = SimpleNamespace(
execute=AsyncMock(return_value=_AggregateResult(row))
)
user = SimpleNamespace(id=42, role="operator")
response = await main.get_dashboard_account_stats(db=db, user=user)
self.assertIsInstance(response, main.DashboardAccountStatsResponse)
self.assertEqual(response.total_accounts, 12)
self.assertEqual(response.online_accounts, 5)
self.assertEqual(response.my_accounts, 3)
self.assertEqual(response.my_online_accounts, 2)
db.execute.assert_awaited_once()
async def test_owner_scope_is_only_inside_personal_aggregates(self):
row = _AggregateRow(
total_accounts=8,
online_accounts=4,
my_accounts=2,
my_online_accounts=1,
)
db = SimpleNamespace(
execute=AsyncMock(return_value=_AggregateResult(row))
)
user = SimpleNamespace(id=73, role="viewer")
await main.get_dashboard_account_stats(db=db, user=user)
statement = db.execute.await_args.args[0]
sql = " ".join(str(statement).lower().split())
compiled_params = list(statement.compile().params.values())
# All roles receive the same global totals. The current user id may
# appear in CASE expressions for the two personal counters, but must
# never filter the entire aggregate query through a global WHERE.
self.assertIn("owner_id", sql)
self.assertGreaterEqual(sql.count("case when"), 3)
self.assertEqual(sql.count("accounts.owner_id"), 2)
self.assertIn(73, compiled_params)
self.assertNotIn(" where ", f" {sql} ")
self.assertEqual(db.execute.await_count, 1)
if __name__ == "__main__":
unittest.main()
@@ -171,7 +171,52 @@ class DesktopLoginSecUserIdTests(unittest.IsolatedAsyncioTestCase):
self.assertIs(result, response)
db.execute.assert_not_awaited()
build.assert_awaited_once_with(account, "cookie-json")
build.assert_awaited_once_with(
account,
"cookie-json",
runtime_check=False,
)
async def test_cookie_response_forwards_static_management_check(self):
account = self._account()
summary = {
"cookie_count": 2,
"cookie_valid": True,
"cookie_expired": False,
"reason": "ok",
"expires_at": None,
"key_names": ["sessionid"],
}
detail = {
"has_sessionid": True,
"sessionid": "sid",
"sessionid_ss": "",
"im_ready": False,
"im_status": "static",
"can_skip_browser": False,
"should_reset": False,
}
with (
patch.object(main, "cookie_summary", return_value=summary),
patch.object(
main,
"build_cookie_credential_detail",
new=AsyncMock(return_value=detail),
) as build_detail,
):
result = await main._build_cookie_response(
account,
"cookie-json",
runtime_check=False,
)
self.assertTrue(result.has_sessionid)
build_detail.assert_awaited_once_with(
"cookie-json",
None,
runtime_check=False,
)
async def test_legacy_cookie_request_is_guarded_for_existing_desktop_clients(self):
account = self._account()
+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()
+136 -4
View File
@@ -30,36 +30,127 @@ class AccountReplyQueueTests(unittest.IsolatedAsyncioTestCase):
await asyncio.sleep(0.002)
self.assertEqual(queue.pending_count, 0)
async def test_one_account_runs_three_jobs_at_successive_fifo_slots(self):
async def test_immediate_if_idle_runs_first_now_then_successive_fifo_slots(self):
queue = await self._start_queue(account_id=101)
interval = 0.05
loop = asyncio.get_running_loop()
started_at = loop.time()
calls: list[tuple[int, float]] = []
finished = asyncio.Event()
first_started = asyncio.Event()
release_first = asyncio.Event()
def callback_for(index: int):
async def callback() -> None:
calls.append((index, loop.time() - started_at))
if index == 0:
first_started.set()
await release_first.wait()
if len(calls) == 3:
finished.set()
return callback
for index in range(3):
await queue.enqueue(interval, callback_for(index), description=str(index))
await queue.enqueue(
interval,
callback_for(0),
description="0",
immediate_if_idle=True,
)
await asyncio.wait_for(first_started.wait(), timeout=0.1)
for index in (1, 2):
await queue.enqueue(
interval,
callback_for(index),
description=str(index),
immediate_if_idle=True,
)
release_first.set()
await asyncio.wait_for(finished.wait(), timeout=0.75)
await self._wait_until_idle(queue)
self.assertEqual([index for index, _ in calls], [0, 1, 2])
elapsed = [timestamp for _, timestamp in calls]
for timestamp, expected in zip(elapsed, (interval, interval * 2, interval * 3)):
for timestamp, expected in zip(elapsed, (0, interval, interval * 2)):
self.assertGreaterEqual(timestamp, expected - 0.015)
self.assertLess(timestamp, expected + 0.15)
self.assertGreaterEqual(elapsed[1] - elapsed[0], interval - 0.02)
self.assertGreaterEqual(elapsed[2] - elapsed[1], interval - 0.02)
async def test_immediate_if_idle_does_not_bypass_active_send(self):
queue = await self._start_queue(account_id=102)
interval = 1.0
first_started = asyncio.Event()
release_first = asyncio.Event()
second_started = asyncio.Event()
async def first_callback() -> None:
first_started.set()
await release_first.wait()
async def second_callback() -> None:
second_started.set()
await queue.enqueue(
interval,
first_callback,
description="first",
immediate_if_idle=True,
)
await asyncio.wait_for(first_started.wait(), timeout=0.1)
await queue.enqueue(
interval,
second_callback,
description="second",
immediate_if_idle=True,
)
snapshot = await queue.snapshot()
self.assertEqual([item["status"] for item in snapshot], ["sending", "waiting"])
self.assertEqual(snapshot[0]["interval_seconds"], 0)
self.assertEqual(snapshot[1]["interval_seconds"], int(interval))
release_first.set()
await asyncio.sleep(0.02)
self.assertFalse(second_started.is_set())
async def test_immediate_if_idle_resets_after_queue_drains(self):
queue = await self._start_queue(account_id=103)
interval = 1.0
first_finished = asyncio.Event()
second_started = asyncio.Event()
release_second = asyncio.Event()
async def first_callback() -> None:
first_finished.set()
async def second_callback() -> None:
second_started.set()
await release_second.wait()
await queue.enqueue(
interval,
first_callback,
description="first wave",
immediate_if_idle=True,
)
await asyncio.wait_for(first_finished.wait(), timeout=0.1)
await self._wait_until_idle(queue)
await queue.enqueue(
interval,
second_callback,
description="second wave",
immediate_if_idle=True,
)
await asyncio.wait_for(second_started.wait(), timeout=0.1)
snapshot = await queue.snapshot()
self.assertEqual(len(snapshot), 1)
self.assertEqual(snapshot[0]["status"], "sending")
self.assertEqual(snapshot[0]["interval_seconds"], 0)
release_second.set()
async def test_separate_account_queues_reach_first_slot_without_blocking(self):
first_queue = await self._start_queue(account_id=201)
second_queue = await self._start_queue(account_id=202)
@@ -194,6 +285,47 @@ class AccountReplyQueueTests(unittest.IsolatedAsyncioTestCase):
new_due = datetime.fromisoformat(new["scheduled_at"]).timestamp()
self.assertAlmostEqual(old_due - new_due, interval, delta=0.05)
async def test_send_now_on_zero_slot_does_not_claim_later_jobs_shifted(self):
queue = await self._start_queue(account_id=512)
interval = 1.0
async def noop() -> None:
pass
# Both enqueues complete without yielding to the consumer, preserving
# the narrow management-API window where the zero-slot first job is
# still waiting and can be selected by send-now.
await queue.enqueue(
interval,
noop,
description="immediate first",
immediate_if_idle=True,
)
await queue.enqueue(
interval,
noop,
description="scheduled second",
immediate_if_idle=True,
)
before = await queue.snapshot()
second_before = next(
item for item in before if item["description"] == "scheduled second"
)
result = await queue.send_now(before[0]["job_id"])
after = await queue.snapshot()
second_after = next(
item for item in after if item["description"] == "scheduled second"
)
self.assertEqual(result["status"], "accepted")
self.assertEqual(result["shifted_count"], 0)
second_due_before = datetime.fromisoformat(
second_before["scheduled_at"]
).timestamp()
second_due_after = datetime.fromisoformat(second_after["scheduled_at"]).timestamp()
self.assertAlmostEqual(second_due_after, second_due_before, delta=0.01)
async def test_send_now_middle_runs_first_and_only_shifts_jobs_behind_it(self):
queue = await self._start_queue(account_id=503)
interval = 0.12
+21
View File
@@ -84,6 +84,27 @@ class ReplyQueueApiTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(response.total_pending, 1)
self.assertEqual([item.account_id for item in response.items], [1])
async def test_summary_with_account_ids_only_scans_requested_page(self):
service_one = _FakeService(1, [_queue_item(1)])
service_two = _FakeService(2, [_queue_item(2, "job-2")])
service_one.get_reply_queue_snapshot = AsyncMock(return_value=service_one.items)
service_two.get_reply_queue_snapshot = AsyncMock(return_value=service_two.items)
main.manager.workers = {
1: SimpleNamespace(is_running=True, _im_service=service_one),
2: SimpleNamespace(is_running=True, _im_service=service_two),
}
response = await main.get_reply_queue_summaries(
account_ids="2",
db=object(),
user=SimpleNamespace(id=1, role="admin"),
)
service_one.get_reply_queue_snapshot.assert_not_awaited()
service_two.get_reply_queue_snapshot.assert_awaited_once()
self.assertEqual(response.total_pending, 1)
self.assertEqual([item.account_id for item in response.items], [2])
async def test_offline_account_detail_returns_empty_snapshot(self):
account = SimpleNamespace(id=1, reply_delay_seconds=0)
with (
+15 -2
View File
@@ -23,7 +23,9 @@ from rpa_engine.playwright_worker import DouyinWorker
class _RecordingQueue:
def __init__(self) -> None:
self.jobs: list[tuple[float, object, str, dict, frozenset[str]]] = []
self.jobs: list[
tuple[float, object, str, dict, frozenset[str], bool]
] = []
async def enqueue(
self,
@@ -33,11 +35,19 @@ class _RecordingQueue:
details=None,
merge_key="",
merge_keys=None,
immediate_if_idle=False,
) -> int:
keys = merge_keys if merge_keys is not None else [merge_key]
normalized_keys = frozenset(str(key) for key in keys if str(key or "").strip())
self.jobs.append(
(delay_seconds, callback, description, dict(details or {}), normalized_keys)
(
delay_seconds,
callback,
description,
dict(details or {}),
normalized_keys,
bool(immediate_if_idle),
)
)
return len(self.jobs)
@@ -75,6 +85,7 @@ class _RecordingQueue:
job[2],
merged_details,
frozenset(job[4] | incoming_keys),
job[5],
)
return {
"status": "merged",
@@ -123,6 +134,8 @@ class ReplyQueueIntegrationTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(len(service._reply_queue.jobs), 1)
self.assertEqual(match_reply.await_count, 1)
self.assertEqual(service._reply_queue.jobs[0][0], 60)
self.assertTrue(service._reply_queue.jobs[0][5])
details = service._reply_queue.jobs[0][3]
self.assertEqual(details["sender_name"], "张三")
self.assertEqual(details["conversation_id"], "conv-1")
+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=" ")
+145
View File
@@ -389,6 +389,10 @@ class GlobalSendQueueTests(unittest.IsolatedAsyncioTestCase):
class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase):
async def _wait_until(self, predicate) -> None:
while not predicate():
await asyncio.sleep(0.001)
async def test_background_slot_respects_configured_concurrency_limit(self):
with patch.dict(
os.environ,
@@ -458,6 +462,147 @@ class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(len(tasks_seen), 3)
self.assertTrue(all(task is tasks_seen[0] for task in tasks_seen))
async def test_startup_traffic_leaves_one_slot_for_recurring_work(self):
with patch.dict(
os.environ,
{"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "3"},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
release = asyncio.Event()
async def request(name: str, *, startup: bool = False) -> None:
async with controller.background_slot(1, name, startup=startup):
entered.append(name)
await release.wait()
starts = [
asyncio.create_task(request(f"startup-{index}", startup=True))
for index in range(3)
]
while len(entered) < 2:
await asyncio.sleep(0)
await asyncio.sleep(0.01)
# A flood of startups may occupy at most capacity - 1 slots, so a
# recurring poll still gets in while the fleet is coming online.
self.assertEqual(len(entered), 2)
self.assertEqual(controller.background_startup_active, 2)
poll = asyncio.create_task(request("poll"))
await asyncio.wait_for(
self._wait_until(lambda: "poll" in entered),
timeout=0.5,
)
release.set()
await asyncio.wait_for(asyncio.gather(*starts, poll), timeout=0.5)
self.assertEqual(controller.background_active, 0)
self.assertEqual(controller.background_startup_active, 0)
async def test_recurring_work_is_not_deferred_indefinitely_by_startups(self):
"""Pending startup work may delay a recurring poll, never block it.
Starting several hundred accounts keeps startup requests queued for the
whole run. Yielding to that queue without a deadline left every hosted
account silent until the last account had finished coming online.
"""
with patch.dict(
os.environ,
{
"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1",
"KEFU_BACKGROUND_NORMAL_MAX_DEFER_SECONDS": "0.02",
},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
release = asyncio.Event()
async def poll() -> None:
async with controller.background_slot(1, "conversation poll"):
entered.append("poll")
await release.wait()
# Stands in for a batch whose startup requests never stop arriving.
controller._background_startup_clear.clear()
task = asyncio.create_task(poll())
await asyncio.wait_for(
self._wait_until(lambda: entered == ["poll"]),
timeout=1.0,
)
release.set()
await asyncio.wait_for(task, timeout=0.5)
async def test_startup_request_is_not_buried_behind_normal_backlog(self):
with patch.dict(
os.environ,
{"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1"},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
releases = {
name: asyncio.Event()
for name in ("active", "normal-1", "normal-2", "startup")
}
async def request(name: str, *, startup: bool = False) -> None:
async with controller.background_slot(
1,
name,
startup=startup,
):
entered.append(name)
await releases[name].wait()
active = asyncio.create_task(request("active"))
while entered != ["active"]:
await asyncio.sleep(0)
normal_one = asyncio.create_task(request("normal-1"))
normal_two = asyncio.create_task(request("normal-2"))
# Let one normal request reach the shared semaphore while the other is
# held at normal admission, then add the priority startup request.
await asyncio.sleep(0)
await asyncio.sleep(0)
startup = asyncio.create_task(request("startup", startup=True))
await asyncio.sleep(0)
releases["active"].set()
while len(entered) < 2:
await asyncio.sleep(0)
self.assertEqual(entered[:2], ["active", "normal-1"])
releases["normal-1"].set()
while len(entered) < 3:
await asyncio.sleep(0)
self.assertEqual(entered[:3], ["active", "normal-1", "startup"])
releases["startup"].set()
while len(entered) < 4:
await asyncio.sleep(0)
releases["normal-2"].set()
await asyncio.wait_for(
asyncio.gather(active, normal_one, normal_two, startup),
timeout=0.2,
)
self.assertEqual(
entered,
["active", "normal-1", "startup", "normal-2"],
)
self.assertEqual(controller.background_active, 0)
self.assertEqual(controller.background_waiting, 0)
self.assertEqual(controller.background_startup_active, 0)
self.assertEqual(controller.background_startup_waiting, 0)
class TrafficControllerLoopIsolationTests(unittest.TestCase):
def test_get_traffic_controller_does_not_reuse_asyncio_primitives(self):
+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()
+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()
+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),
)
+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:
+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":
+1 -1
View File
@@ -63,7 +63,7 @@ const handleLogout = () => {
</a-menu-item>
<a-menu-item key="/accounts">
<template #icon><UserOutlined /></template>
<span>账号管理</span>
<span>{{ auth.isAdmin ? '账号管理' : '我的账号' }}</span>
</a-menu-item>
<a-menu-item key="/messages">
<template #icon><MessageOutlined /></template>
+84 -23
View File
@@ -1,6 +1,6 @@
<script setup>
import { ref, onMounted, onUnmounted, computed, watch } from 'vue'
import { useRouter } from 'vue-router'
import { useRoute, useRouter } from 'vue-router'
import api from '../api'
import { message } from 'ant-design-vue'
import { useAuthStore } from '../stores/auth'
@@ -39,6 +39,7 @@ import {
import { useIsMobile } from '../composables/useIsMobile'
const auth = useAuthStore()
const route = useRoute()
const router = useRouter()
const isMobile = useIsMobile()
const profileModalWidth = computed(() => (isMobile.value ? 'calc(100vw - 32px)' : 860))
@@ -75,6 +76,13 @@ let batchStatusTimer = null
let batchStatusRequestActive = false
let batchStatusGeneration = 0
const BATCH_STATUS_POLL_MS = 1500
const BATCH_STATUS_MAX_POLL_MS = 5000
const BATCH_STATUS_BACKOFF_STEP_MS = 750
let batchStatusPollDelayMs = BATCH_STATUS_POLL_MS
let batchStatusLastFinished = null
let batchStatusLastTotal = null
let batchStatusLastProcessing = null
let batchStatusLastQueued = null
const selectedIds = ref([])
const addVisible = ref(false)
const addSaving = ref(false)
@@ -710,9 +718,18 @@ const applyQueueSnapshot = (data, accountId) => {
const fetchReplyQueueSummaries = async ({ silent = true } = {}) => {
if (queueSummaryLoading.value) return
const accountIds = accounts.value
.map((account) => Number(account?.id))
.filter((accountId) => Number.isInteger(accountId) && accountId > 0)
if (!accountIds.length) {
replyQueueSummaries.value = {}
return
}
queueSummaryLoading.value = true
try {
const res = await api.get('/reply-queues')
const res = await api.get('/reply-queues', {
params: { account_ids: accountIds.join(',') }
})
const rows = Array.isArray(res.data?.items) ? res.data.items : []
const next = {}
for (const row of rows) {
@@ -983,6 +1000,10 @@ const goAccountRulesPage = (accountId) => {
}
const openAddModal = () => {
if (!auth.canWrite) {
message.warning('当前账号为只读角色,不能添加托管账号')
return
}
if (!canAddAccount.value) {
if (canPurchaseSlots.value) {
purchaseVisible.value = true
@@ -1147,6 +1168,11 @@ const stopBatchStatusPolling = () => {
batchStatusTimer = null
}
batchStatusRequestActive = false
batchStatusPollDelayMs = BATCH_STATUS_POLL_MS
batchStatusLastFinished = null
batchStatusLastTotal = null
batchStatusLastProcessing = null
batchStatusLastQueued = null
}
const finishBatchStart = async (snapshot) => {
@@ -1176,6 +1202,7 @@ const finishBatchStart = async (snapshot) => {
}
await fetchAccounts()
batchStarting.value = false
startReplyQueueSummaryPolling()
}
const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
@@ -1191,11 +1218,31 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
Math.max(0, Number(snapshot?.failed_count) || 0) +
Math.max(0, Number(snapshot?.skipped_count) || 0) +
Math.max(0, Number(snapshot?.cancelled_count) || 0)
message.loading({
content: `账号启动队列处理中:${Math.min(finished, total)}/${total}`,
key: 'batch_start',
duration: 0
})
const processing = Math.max(0, Number(snapshot?.processing_count) || 0)
const queued = Math.max(0, Number(snapshot?.queued_count) || 0)
const visibleFinished = Math.min(finished, total)
const progressChanged =
visibleFinished !== batchStatusLastFinished ||
total !== batchStatusLastTotal ||
processing !== batchStatusLastProcessing ||
queued !== batchStatusLastQueued
if (progressChanged) {
batchStatusLastFinished = visibleFinished
batchStatusLastTotal = total
batchStatusLastProcessing = processing
batchStatusLastQueued = queued
batchStatusPollDelayMs = BATCH_STATUS_POLL_MS
message.loading({
content: `账号启动队列处理中:${visibleFinished}/${total}(正在处理 ${processing},等待 ${queued},系统正错峰启动)`,
key: 'batch_start',
duration: 0
})
} else {
batchStatusPollDelayMs = Math.min(
BATCH_STATUS_MAX_POLL_MS,
batchStatusPollDelayMs + BATCH_STATUS_BACKOFF_STEP_MS
)
}
if (snapshot?.complete) {
await finishBatchStart(snapshot)
return true
@@ -1225,12 +1272,13 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
duration: 5
})
await fetchAccounts()
startReplyQueueSummaryPolling()
return
} finally {
if (generation === batchStatusGeneration) batchStatusRequestActive = false
}
if (generation === batchStatusGeneration && activeStartBatchId.value === batchId) {
batchStatusTimer = setTimeout(poll, BATCH_STATUS_POLL_MS)
batchStatusTimer = setTimeout(poll, batchStatusPollDelayMs)
}
}
@@ -1240,7 +1288,7 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
generation === batchStatusGeneration &&
activeStartBatchId.value === batchId
) {
batchStatusTimer = setTimeout(poll, BATCH_STATUS_POLL_MS)
batchStatusTimer = setTimeout(poll, batchStatusPollDelayMs)
}
})
}
@@ -1249,6 +1297,7 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
const runBatchStart = async ({ accountIds = [], allAccounts = false }) => {
if (batchStarting.value) return
batchStarting.value = true
stopReplyQueueSummaryPolling()
stopBatchStatusPolling()
const submitGeneration = batchStatusGeneration
message.loading({ content: '正在提交账号启动队列...', key: 'batch_start', duration: 0 })
@@ -1270,6 +1319,7 @@ const runBatchStart = async ({ accountIds = [], allAccounts = false }) => {
if (submitGeneration !== batchStatusGeneration) return
batchStarting.value = false
activeStartBatchId.value = null
startReplyQueueSummaryPolling()
message.error({
content: error.response?.data?.detail || error.message || '提交批量启动失败',
key: 'batch_start',
@@ -1496,7 +1546,7 @@ const formatCookieTime = (value) => {
return date.toLocaleString('zh-CN')
}
const applyCookieResponse = (data) => {
const applyCookieResponse = (data, { preserveRuntimeCredential = false } = {}) => {
editForm.value.cookie_updated_at = data.cookie_updated_at
editForm.value.cookie_count = data.cookie_count
editForm.value.cookie_valid = data.cookie_valid
@@ -1506,10 +1556,12 @@ const applyCookieResponse = (data) => {
editForm.value.has_sessionid = !!data.has_sessionid
editForm.value.sessionid = data.sessionid || ''
editForm.value.sessionid_ss = data.sessionid_ss || ''
editForm.value.im_ready = !!data.im_ready
editForm.value.im_status = data.im_status || ''
editForm.value.can_skip_browser = !!data.can_skip_browser
editForm.value.should_reset = !!data.should_reset
if (!preserveRuntimeCredential) {
editForm.value.im_ready = !!data.im_ready
editForm.value.im_status = data.im_status || ''
editForm.value.can_skip_browser = !!data.can_skip_browser
editForm.value.should_reset = !!data.should_reset
}
}
const openEditModal = async (acc) => {
@@ -1578,7 +1630,7 @@ const refreshCredential = async () => {
}
const cookieRes = await api.get(`/accounts/${editForm.value.id}/cookie?purpose=management`)
applyCookieResponse(cookieRes.data)
applyCookieResponse(cookieRes.data, { preserveRuntimeCredential: true })
message.success('凭证检测完成')
} catch (error) {
message.error(error.response?.data?.detail || '凭证检测失败')
@@ -1705,13 +1757,21 @@ const clearCookie = async () => {
}
}
onMounted(() => {
auth.fetchMe()
fetchPaymentConfig()
fetchDeviceProfiles()
fetchAccounts()
fetchRules()
onMounted(async () => {
await Promise.all([
auth.fetchMe(),
fetchPaymentConfig(),
fetchDeviceProfiles(),
fetchAccounts(),
fetchRules(),
])
startReplyQueueSummaryPolling()
if (route.query.action === 'add') {
openAddModal()
const nextQuery = { ...route.query }
delete nextQuery.action
router.replace({ path: route.path, query: nextQuery })
}
// //
//
})
@@ -1797,6 +1857,7 @@ onUnmounted(() => {
购买额度
</a-button>
<a-button
v-if="auth.canWrite"
type="primary"
class="gradient-btn"
:disabled="!canAddAccount && !canPurchaseSlots"
@@ -2052,7 +2113,7 @@ onUnmounted(() => {
<UserOutlined style="font-size: 4rem; color: var(--text-muted); margin-bottom: 16px;" />
<h3>暂无托管账号</h3>
<p style="color: var(--text-secondary); margin-bottom: 20px;">添加一个抖音账号开始自动化回复工作吧</p>
<a-button type="primary" class="gradient-btn" @click="openAddModal">
<a-button v-if="auth.canWrite" type="primary" class="gradient-btn" @click="openAddModal">
<template #icon><PlusOutlined /></template>
立即添加
</a-button>
@@ -2257,7 +2318,7 @@ onUnmounted(() => {
placeholder="0 或留空则继承系统默认"
/>
<div class="field-hint">
设置 N 秒后同一账号的待回复会话会依次排 1 条在 N 发送 2 条在 2N 后发送以此类推各账号队列互不影响0 或留空表示继承系统默认当前生效 {{ editForm.reply_delay_effective }} 0 表示不启用兜底排队立即回复
设置 N 秒后同一账号队列为空时首条回复等待 0 并立即发送后续待回复会话按 N 2N 秒依次排队各账号队列互不影响实际发送仍受全局带宽队列保护0 或留空表示继承系统默认当前生效 {{ editForm.reply_delay_effective }} 0 表示不启用兜底排队收到消息后立即回复
</div>
</a-form-item>
</a-col>
+75 -14
View File
@@ -2,18 +2,25 @@
import { ref, onMounted, onUnmounted } from 'vue'
import api from '../api'
import MessageBubble from '../components/MessageBubble.vue'
import { useAuthStore } from '../stores/auth'
import {
UserOutlined,
MessageOutlined,
CheckCircleOutlined,
ClockCircleOutlined,
ArrowRightOutlined,
ThunderboltOutlined
ThunderboltOutlined,
SettingOutlined,
PlusOutlined
} from '@ant-design/icons-vue'
const auth = useAuthStore()
const stats = ref({
totalAccounts: 0,
activeAccounts: 0,
myAccounts: 0,
myActiveAccounts: 0,
totalMessages: 0,
repliedMessages: 0,
replyRate: '0%'
@@ -24,14 +31,15 @@ const loading = ref(true)
const fetchStats = async () => {
try {
const [accountsRes, statsRes] = await Promise.all([
api.get(`/accounts`),
const [accountStatsRes, statsRes] = await Promise.all([
api.get(`/dashboard/account-stats`),
api.get(`/logs/stats`)
])
const accounts = accountsRes.data
stats.value.totalAccounts = accounts.length
stats.value.activeAccounts = accounts.filter(a => a.status === 'online').length
stats.value.totalAccounts = accountStatsRes.data.total_accounts || 0
stats.value.activeAccounts = accountStatsRes.data.online_accounts || 0
stats.value.myAccounts = accountStatsRes.data.my_accounts || 0
stats.value.myActiveAccounts = accountStatsRes.data.my_online_accounts || 0
//
stats.value.totalMessages = statsRes.data.total || 0
@@ -61,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>
@@ -84,8 +97,27 @@ onUnmounted(() => {
多账户自动回复RPA后台支持快捷扫码登录状态持久化保存以及自定义关键字规则精准答复
</p>
</div>
<div class="banner-icon">
<ThunderboltOutlined style="font-size: 4rem; color: #c084fc; opacity: 0.3;" />
<div class="banner-side">
<div v-if="auth.canWrite" class="banner-actions">
<router-link to="/accounts">
<a-button size="large">
<template #icon><UserOutlined /></template>
{{ auth.isAdmin ? '账号管理' : `我的账号(${stats.myAccounts}` }}
</a-button>
</router-link>
<router-link
v-if="auth.canWrite"
:to="{ path: '/accounts', query: { action: 'add' } }"
>
<a-button type="primary" size="large" class="gradient-btn">
<template #icon><PlusOutlined /></template>
{{ auth.isAdmin ? '添加账号' : '添加自己的账号' }}
</a-button>
</router-link>
</div>
<div class="banner-icon">
<ThunderboltOutlined style="font-size: 4rem; color: #c084fc; opacity: 0.3;" />
</div>
</div>
</div>
@@ -98,7 +130,7 @@ onUnmounted(() => {
<UserOutlined />
</div>
<div class="stat-info">
<span class="stat-label">托管账号</span>
<span class="stat-label">全平台托管账号</span>
<h2 class="stat-value">{{ stats.totalAccounts }}</h2>
</div>
</div>
@@ -111,7 +143,7 @@ onUnmounted(() => {
<CheckCircleOutlined />
</div>
<div class="stat-info">
<span class="stat-label">在线运行</span>
<span class="stat-label">全平台在线运行</span>
<h2 class="stat-value text-green">{{ stats.activeAccounts }}</h2>
</div>
</div>
@@ -199,8 +231,9 @@ onUnmounted(() => {
<div class="quick-actions-grid" style="margin-top: 20px;">
<router-link to="/accounts" class="quick-action-card">
<UserOutlined class="action-icon text-gradient" />
<span>账号配置</span>
<p>扫码登录并托管多个抖音账号</p>
<span>{{ auth.isAdmin ? '账号管理' : '我的账号' }}</span>
<p v-if="auth.isAdmin">管理全平台账号配置与运行状态</p>
<p v-else>仅查看和管理自己添加的账号当前 {{ stats.myAccounts }} </p>
</router-link>
<router-link to="/rules" class="quick-action-card">
@@ -225,6 +258,20 @@ onUnmounted(() => {
border-left: 4px solid var(--primary-color);
}
.banner-side {
display: flex;
align-items: center;
gap: 28px;
flex-shrink: 0;
}
.banner-actions {
display: flex;
gap: 12px;
flex-wrap: wrap;
justify-content: flex-end;
}
.stat-card {
display: flex;
align-items: center;
@@ -436,6 +483,20 @@ onUnmounted(() => {
display: none;
}
.banner-side,
.banner-actions {
width: 100%;
justify-content: flex-start;
}
.banner-actions > a {
flex: 1 1 180px;
}
.banner-actions :deep(.ant-btn) {
width: 100%;
}
.stat-card {
padding: 16px;
}
+19 -6
View File
@@ -39,9 +39,23 @@ const LOGS_PAGE_SIZE = 20
const hasMoreLogs = ref(true)
const loadingMore = ref(false)
let pollTimer = null
let pollStopped = false
//
let lastLogsSignature = ''
// Use a self-scheduling timeout instead of setInterval. A slow request must
// finish before the next poll is scheduled, otherwise overlapping requests can
// exhaust the database pool and amplify one slow query into dozens.
const scheduleLogPoll = () => {
if (pollStopped) return
pollTimer = setTimeout(async () => {
if (!document.hidden) {
await fetchLogs(true)
}
scheduleLogPoll()
}, POLL_INTERVAL_MS)
}
const dedupeMessages = (messages) => {
const map = new Map()
for (const item of messages) {
@@ -172,7 +186,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)
@@ -593,10 +607,8 @@ onMounted(async () => {
}
await fetchAccounts()
await fetchLogs()
pollTimer = setInterval(() => {
if (document.hidden) return
fetchLogs(true)
}, POLL_INTERVAL_MS)
pollStopped = false
scheduleLogPoll()
})
watch(
@@ -613,8 +625,9 @@ watch(
)
onUnmounted(() => {
pollStopped = true
if (pollTimer) {
clearInterval(pollTimer)
clearTimeout(pollTimer)
pollTimer = null
}
})
+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)
+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)
+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
+1 -1
View File
@@ -740,7 +740,7 @@ onMounted(() => {
/>
<div class="field-hint">
账号未单独设置时使用此间隔同一账号的第 1 条待回复在 N 发送 2 条在 2N 秒后发送以此类推各账号队列互不影响设为 0 表示不启用兜底排队收到消息后立即回复
账号未单独设置时使用此间隔同一账号队列为空时首条回复等待 0 并立即发送后续待回复会话按 N 2N 秒依次排队各账号队列互不影响实际发送仍受全局带宽队列保护设为 0 表示不启用兜底排队收到消息后立即回复
</div>
</a-form-item>
+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)