Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5db9674490 | ||
|
|
3fc94c4a89 | ||
|
|
8f68af1c2c | ||
|
|
ac406a5f99 | ||
|
|
f99a4edf83 | ||
|
|
153db97dc7 | ||
|
|
8ba13a8ff9 |
@@ -10,6 +10,7 @@ backend/**/__pycache__/
|
|||||||
backend/kefu.db
|
backend/kefu.db
|
||||||
backend/kefu.db-journal
|
backend/kefu.db-journal
|
||||||
backend/sessions/
|
backend/sessions/
|
||||||
|
/douyin.zip
|
||||||
*.env
|
*.env
|
||||||
.env.*
|
.env.*
|
||||||
|
|
||||||
|
|||||||
Binary file not shown.
Binary file not shown.
+600
-87
@@ -20,7 +20,7 @@ from settings import CORS_ORIGINS, SERVE_WEB, STATIC_DIR
|
|||||||
from help_pages import serve_credential_tool
|
from help_pages import serve_credential_tool
|
||||||
from web_static import mount_frontend
|
from web_static import mount_frontend
|
||||||
from pydantic import BaseModel, Field
|
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 sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
from models.database import engine, Base, get_db, AsyncSessionLocal
|
from models.database import engine, Base, get_db, AsyncSessionLocal
|
||||||
@@ -107,7 +107,14 @@ class WorkerManager:
|
|||||||
"""Serialize credential preparation across single and batch starts."""
|
"""Serialize credential preparation across single and batch starts."""
|
||||||
return self._preparation_locks.setdefault(int(account_id), asyncio.Lock())
|
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):
|
async with self._account_lock(account_id):
|
||||||
if account_id in self.workers:
|
if account_id in self.workers:
|
||||||
worker = self.workers[account_id]
|
worker = self.workers[account_id]
|
||||||
@@ -115,11 +122,60 @@ class WorkerManager:
|
|||||||
return False
|
return False
|
||||||
del self.workers[account_id]
|
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
|
self.workers[account_id] = worker
|
||||||
await worker.start()
|
await worker.start()
|
||||||
|
|
||||||
|
if not wait_until_ready:
|
||||||
return True
|
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 def stop_worker(self, account_id: int):
|
||||||
async with self._account_lock(account_id):
|
async with self._account_lock(account_id):
|
||||||
if account_id in self.workers:
|
if account_id in self.workers:
|
||||||
@@ -331,17 +387,29 @@ async def _flush_system_logs_loop():
|
|||||||
async def _sync_legacy_cookie_files():
|
async def _sync_legacy_cookie_files():
|
||||||
"""将旧版仅保存在文件的 Cookie 同步到数据库"""
|
"""将旧版仅保存在文件的 Cookie 同步到数据库"""
|
||||||
async with AsyncSessionLocal() as db:
|
async with AsyncSessionLocal() as db:
|
||||||
result = await db.execute(select(Account))
|
# Only legacy rows without a database Cookie need filesystem work.
|
||||||
accounts = result.scalars().all()
|
# 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
|
changed = False
|
||||||
for acc in accounts:
|
for account_id in account_ids:
|
||||||
if acc.cookie_data:
|
file_data = read_cookie_file(int(account_id))
|
||||||
continue
|
|
||||||
file_data = read_cookie_file(acc.id)
|
|
||||||
if file_data:
|
if file_data:
|
||||||
acc.cookie_data = file_data
|
await db.execute(
|
||||||
acc.cookie_path = get_cookie_path(acc.id)
|
update(Account)
|
||||||
acc.cookie_updated_at = datetime.utcnow()
|
.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
|
changed = True
|
||||||
if changed:
|
if changed:
|
||||||
await db.commit()
|
await db.commit()
|
||||||
@@ -427,6 +495,22 @@ async def _conversations_from_logs(
|
|||||||
return list(seen.values())
|
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:
|
async def _reset_account_credentials(account_id: int, db: AsyncSession) -> Account:
|
||||||
"""清除账号 Cookie、IM 会话等登录数据,并停止托管。"""
|
"""清除账号 Cookie、IM 会话等登录数据,并停止托管。"""
|
||||||
await manager.stop_worker(account_id)
|
await manager.stop_worker(account_id)
|
||||||
@@ -501,10 +585,8 @@ async def _seed_admin_user():
|
|||||||
# 初始化数据库
|
# 初始化数据库
|
||||||
@app.on_event("startup")
|
@app.on_event("startup")
|
||||||
async def startup():
|
async def startup():
|
||||||
# 放大默认线程池:每个托管账号的 a_bogus/web_protect/ts_sign 签名都是阻塞的 Node
|
# a_bogus/web_protect/ts_sign 仍需在线程池中执行阻塞的 Node 调用,但账号启动、
|
||||||
# 子进程调用,经 asyncio.to_thread 跑在默认线程池里。Python 默认池大小仅
|
# 后台请求和发送通道都已有独立并发限制,因此线程池按 CPU 有界配置即可。
|
||||||
# min(32, cpu+4),在 1 核云服务器上只有 5 个线程,导致超过 5 个账号并发时第 6 个
|
|
||||||
# 账号的签名/取信息/发送会一直排队阻塞直至超时失败。这里显式放大线程池消除该瓶颈。
|
|
||||||
import concurrent.futures as _futures
|
import concurrent.futures as _futures
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -512,7 +594,11 @@ async def startup():
|
|||||||
except ValueError:
|
except ValueError:
|
||||||
_pool_size = 0
|
_pool_size = 0
|
||||||
if _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 = asyncio.get_running_loop()
|
||||||
loop.set_default_executor(
|
loop.set_default_executor(
|
||||||
_futures.ThreadPoolExecutor(
|
_futures.ThreadPoolExecutor(
|
||||||
@@ -557,12 +643,95 @@ async def startup():
|
|||||||
|
|
||||||
@app.on_event("shutdown")
|
@app.on_event("shutdown")
|
||||||
async def 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
|
# Stop queued preparations first so no new workers appear while the
|
||||||
# existing workers are being drained below.
|
# existing workers are being drained below. Its cancellation path may
|
||||||
await batch_start_queue.stop()
|
# itself wait for a half-open login/DB operation, so it needs an
|
||||||
# 停止所有正在运行的 RPA 任务
|
# independent deadline; the remaining workers are still covered by the
|
||||||
for account_id in list(manager.workers.keys()):
|
# bounded parallel stop below.
|
||||||
await manager.stop_worker(account_id)
|
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
|
from rpa_engine.douyin_im.traffic_control import shutdown_traffic_controller
|
||||||
await shutdown_traffic_controller()
|
await shutdown_traffic_controller()
|
||||||
if _system_log_flush_task:
|
if _system_log_flush_task:
|
||||||
@@ -625,6 +794,34 @@ class AccountResponse(BaseModel):
|
|||||||
from_attributes = True
|
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):
|
class AccountUpdate(BaseModel):
|
||||||
phone: Optional[str] = None
|
phone: Optional[str] = None
|
||||||
username: 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)
|
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)
|
cookie_data = cookie_data if cookie_data is not None else _get_account_cookie_data(account)
|
||||||
summary = cookie_summary(cookie_data)
|
summary = cookie_summary(cookie_data)
|
||||||
im_detail = await build_cookie_credential_detail(
|
im_detail = await build_cookie_credential_detail(
|
||||||
cookie_data,
|
cookie_data,
|
||||||
account.im_session_data,
|
account.im_session_data,
|
||||||
runtime_check=True,
|
runtime_check=runtime_check,
|
||||||
)
|
)
|
||||||
return AccountCookieResponse(
|
return AccountCookieResponse(
|
||||||
account_id=account.id,
|
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")
|
@app.get("/api/accounts")
|
||||||
async def get_accounts(
|
async def get_accounts(
|
||||||
page: Optional[int] = Query(None, ge=1),
|
page: Optional[int] = Query(None, ge=1),
|
||||||
@@ -1026,42 +1273,91 @@ async def get_accounts(
|
|||||||
支持 q(昵称/抖音ID/手机号/账号ID 搜索)与 status 筛选。
|
支持 q(昵称/抖音ID/手机号/账号ID 搜索)与 status 筛选。
|
||||||
Cookie 解析等重逻辑只对当前页执行。
|
Cookie 解析等重逻辑只对当前页执行。
|
||||||
"""
|
"""
|
||||||
result = await db.execute(accounts_for_user(user))
|
if page is None:
|
||||||
accounts = result.scalars().all()
|
result = await db.execute(accounts_for_user(user))
|
||||||
# 更新内存中的运行状态与数据库同步,以防异常断开
|
accounts = result.scalars().all()
|
||||||
for acc in accounts:
|
# 兼容旧的全量接口行为。
|
||||||
|
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)
|
is_running = manager.is_running(acc.id)
|
||||||
if is_running and acc.status == "offline":
|
if is_running and acc.status == "offline":
|
||||||
acc.status = "online"
|
acc.status = "online"
|
||||||
elif not is_running and acc.status in ("online", "logging_in", "starting"):
|
elif not is_running and acc.status in ("online", "logging_in", "starting"):
|
||||||
acc.status = "offline"
|
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 {
|
return {
|
||||||
"items": [_build_account_response(acc) for acc in items],
|
"items": [_build_account_response(acc) for acc in items],
|
||||||
"total": total,
|
"total": total,
|
||||||
@@ -1069,6 +1365,83 @@ async def get_accounts(
|
|||||||
"page_size": page_size,
|
"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)
|
@app.get("/api/accounts/{account_id}", response_model=AccountResponse)
|
||||||
async def get_account(
|
async def get_account(
|
||||||
account_id: int,
|
account_id: int,
|
||||||
@@ -1248,6 +1621,10 @@ async def update_account(
|
|||||||
user: User = Depends(require_write),
|
user: User = Depends(require_write),
|
||||||
):
|
):
|
||||||
account = await get_owned_account(db, user, account_id, write=True)
|
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:
|
if body.phone is not None:
|
||||||
account.phone = body.phone
|
account.phone = body.phone
|
||||||
@@ -1271,19 +1648,65 @@ async def update_account(
|
|||||||
account.updated_at = datetime.utcnow()
|
account.updated_at = datetime.utcnow()
|
||||||
await db.commit()
|
await db.commit()
|
||||||
await db.refresh(account)
|
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)
|
return _build_account_response(account)
|
||||||
|
|
||||||
|
|
||||||
@app.get("/api/reply-queues", response_model=ReplyQueueSummaryResponse)
|
@app.get("/api/reply-queues", response_model=ReplyQueueSummaryResponse)
|
||||||
async def get_reply_queue_summaries(
|
async def get_reply_queue_summaries(
|
||||||
|
account_ids: Optional[str] = None,
|
||||||
db: AsyncSession = Depends(get_db),
|
db: AsyncSession = Depends(get_db),
|
||||||
user: User = Depends(get_current_user),
|
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] = []
|
summaries: list[ReplyQueueSummaryItem] = []
|
||||||
total_pending = 0
|
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:
|
if allowed_ids is not None and account_id not in allowed_ids:
|
||||||
continue
|
continue
|
||||||
service = worker._im_service if worker else None
|
service = worker._im_service if worker else None
|
||||||
@@ -1464,7 +1887,17 @@ async def get_account_cookie(
|
|||||||
cookie_data = _get_account_cookie_data(account)
|
cookie_data = _get_account_cookie_data(account)
|
||||||
if cookie_data and purpose != "management":
|
if cookie_data and purpose != "management":
|
||||||
await _require_desktop_login_sec_user_id(account, db)
|
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(
|
@app.get(
|
||||||
@@ -1514,6 +1947,7 @@ async def update_account_cookie(
|
|||||||
# are committed so no worker can start in the stop/commit gap.
|
# are committed so no worker can start in the stop/commit gap.
|
||||||
async with manager.preparation_lock(account_id):
|
async with manager.preparation_lock(account_id):
|
||||||
account = await get_owned_account(db, user, account_id, write=True)
|
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 batch_start_queue.cancel_account(account_id)
|
||||||
await manager.stop_worker(account_id)
|
await manager.stop_worker(account_id)
|
||||||
|
|
||||||
@@ -1553,6 +1987,7 @@ async def delete_account_cookie(
|
|||||||
# before the cleared credentials are committed.
|
# before the cleared credentials are committed.
|
||||||
async with manager.preparation_lock(account_id):
|
async with manager.preparation_lock(account_id):
|
||||||
account = await get_owned_account(db, user, account_id, write=True)
|
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 batch_start_queue.cancel_account(account_id)
|
||||||
await manager.stop_worker(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 get_owned_account(db, user, account_id, write=True)
|
||||||
# 停止运行中的任务
|
# 停止运行中的任务
|
||||||
|
await _release_db_connection(db)
|
||||||
await batch_start_queue.cancel_account(account_id)
|
await batch_start_queue.cancel_account(account_id)
|
||||||
await manager.stop_worker(account_id)
|
await manager.stop_worker(account_id)
|
||||||
clear_cookie_file(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)
|
account = await get_owned_account(db, user, account_id)
|
||||||
|
|
||||||
cookie_data = _get_account_cookie_data(account)
|
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)
|
return CredentialValidateResponse(**assessment)
|
||||||
|
|
||||||
|
|
||||||
@@ -1645,6 +2085,7 @@ async def reset_account_credentials(
|
|||||||
user: User = Depends(require_write),
|
user: User = Depends(require_write),
|
||||||
):
|
):
|
||||||
await get_owned_account(db, user, account_id, write=True)
|
await get_owned_account(db, user, account_id, write=True)
|
||||||
|
await _release_db_connection(db)
|
||||||
await batch_start_queue.cancel_account(account_id)
|
await batch_start_queue.cancel_account(account_id)
|
||||||
account = await _reset_account_credentials(account_id, db)
|
account = await _reset_account_credentials(account_id, db)
|
||||||
return {
|
return {
|
||||||
@@ -1657,6 +2098,8 @@ async def _start_account_rpa_impl(
|
|||||||
account: Account,
|
account: Account,
|
||||||
db: AsyncSession,
|
db: AsyncSession,
|
||||||
requested_login_mode: Optional[str] = None,
|
requested_login_mode: Optional[str] = None,
|
||||||
|
*,
|
||||||
|
wait_for_ready: bool = False,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
"""Validate once, persist the starting state, then spawn one worker."""
|
"""Validate once, persist the starting state, then spawn one worker."""
|
||||||
account_id = int(account.id)
|
account_id = int(account.id)
|
||||||
@@ -1664,7 +2107,14 @@ async def _start_account_rpa_impl(
|
|||||||
return {"status": "running", "message": "RPA worker is already running."}
|
return {"status": "running", "message": "RPA worker is already running."}
|
||||||
|
|
||||||
cookie_data = _get_account_cookie_data(account)
|
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"]
|
login_mode = requested_login_mode or assessment["login_mode"]
|
||||||
reset_performed = False
|
reset_performed = False
|
||||||
|
|
||||||
@@ -1686,10 +2136,19 @@ async def _start_account_rpa_impl(
|
|||||||
detail=assessment["message"] or "凭证未通过验证,无法直连 IM",
|
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
|
# Write before spawning the task. This gives both single and batch calls
|
||||||
# an immediate authoritative status and avoids racing the worker's first
|
# an immediate authoritative status and avoids racing the worker's first
|
||||||
# database update. Batch submission normally wrote this state in bulk, so
|
# database update. Batch readiness keeps the number of accounts reaching
|
||||||
# no per-account commit is needed in that path.
|
# this commit bounded instead of letting the whole batch write at once.
|
||||||
state_changed = (
|
state_changed = (
|
||||||
account.status != "starting"
|
account.status != "starting"
|
||||||
or account.qr_code_base64 is not None
|
or account.qr_code_base64 is not None
|
||||||
@@ -1700,9 +2159,20 @@ async def _start_account_rpa_impl(
|
|||||||
account.qr_code_base64 = None
|
account.qr_code_base64 = None
|
||||||
account.error_message = None
|
account.error_message = None
|
||||||
await db.commit()
|
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:
|
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:
|
except Exception as exc:
|
||||||
account.status = "error"
|
account.status = "error"
|
||||||
account.error_message = str(exc) or "启动托管失败"
|
account.error_message = str(exc) or "启动托管失败"
|
||||||
@@ -1710,7 +2180,9 @@ async def _start_account_rpa_impl(
|
|||||||
raise
|
raise
|
||||||
|
|
||||||
if started:
|
if started:
|
||||||
if login_mode == "im_direct":
|
if wait_for_ready:
|
||||||
|
msg = "IM 托管已完成初始化"
|
||||||
|
elif login_mode == "im_direct":
|
||||||
msg = assessment["message"] or "凭证有效,正在直连 IM 托管(无需浏览器)"
|
msg = assessment["message"] or "凭证有效,正在直连 IM 托管(无需浏览器)"
|
||||||
elif reset_performed:
|
elif reset_performed:
|
||||||
msg = "凭证已失效,已清除旧数据,正在打开浏览器重新登录..."
|
msg = "凭证已失效,已清除旧数据,正在打开浏览器重新登录..."
|
||||||
@@ -1720,7 +2192,7 @@ async def _start_account_rpa_impl(
|
|||||||
msg = "未保存 Cookie,将打开浏览器扫码登录..."
|
msg = "未保存 Cookie,将打开浏览器扫码登录..."
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"status": "starting",
|
"status": "running" if wait_for_ready else "starting",
|
||||||
"login_mode": login_mode,
|
"login_mode": login_mode,
|
||||||
"cookie_valid": assessment["cookie_valid"],
|
"cookie_valid": assessment["cookie_valid"],
|
||||||
"im_ready": assessment["im_ready"],
|
"im_ready": assessment["im_ready"],
|
||||||
@@ -1742,7 +2214,11 @@ async def _start_queued_account(account_id: int) -> dict:
|
|||||||
if account.quota_disabled:
|
if account.quota_disabled:
|
||||||
raise RuntimeError("账号已停用,无法启动托管")
|
raise RuntimeError("账号已停用,无法启动托管")
|
||||||
try:
|
try:
|
||||||
return await _start_account_rpa_impl(account, db)
|
return await _start_account_rpa_impl(
|
||||||
|
account,
|
||||||
|
db,
|
||||||
|
wait_for_ready=True,
|
||||||
|
)
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
raise
|
raise
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
@@ -1764,6 +2240,10 @@ async def start_account_rpa(
|
|||||||
user: User = Depends(require_write),
|
user: User = Depends(require_write),
|
||||||
):
|
):
|
||||||
account = await get_owned_account(db, user, account_id, write=True)
|
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)
|
await batch_start_queue.cancel_account(account_id)
|
||||||
async with manager.preparation_lock(account_id):
|
async with manager.preparation_lock(account_id):
|
||||||
return await _start_account_rpa_impl(account, db, body.login_mode)
|
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)
|
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)
|
await batch_start_queue.cancel_account(account_id)
|
||||||
stopped = await manager.stop_worker(account_id)
|
stopped = await manager.stop_worker(account_id)
|
||||||
# 强制将数据库中的状态重置为 offline
|
# 强制将数据库中的状态重置为 offline
|
||||||
@@ -2050,21 +2533,27 @@ async def get_logs_stats(
|
|||||||
"""消息日志全量统计(数据库计数,不受列表 limit 限制)。"""
|
"""消息日志全量统计(数据库计数,不受列表 limit 限制)。"""
|
||||||
if account_id is not None:
|
if account_id is not None:
|
||||||
await get_owned_account(db, user, account_id)
|
await get_owned_account(db, user, account_id)
|
||||||
base = logs_for_user(user, account_id).subquery()
|
# Select only the indexed status column and calculate both counters in one
|
||||||
total = int(
|
# scan. The old implementation queried the growing log table twice on
|
||||||
(await db.execute(select(func.count()).select_from(base))).scalar() or 0
|
# every dashboard refresh.
|
||||||
|
base = (
|
||||||
|
logs_for_user(user, account_id)
|
||||||
|
.with_only_columns(MessageLog.status)
|
||||||
|
.order_by(None)
|
||||||
|
.subquery()
|
||||||
)
|
)
|
||||||
replied = int(
|
row = (
|
||||||
(
|
await db.execute(
|
||||||
await db.execute(
|
select(
|
||||||
select(func.count())
|
func.count().label("total"),
|
||||||
.select_from(base)
|
func.coalesce(
|
||||||
.where(base.c.status == "replied")
|
func.sum(case((base.c.status == "replied", 1), else_=0)),
|
||||||
)
|
0,
|
||||||
).scalar()
|
).label("replied"),
|
||||||
or 0
|
).select_from(base)
|
||||||
)
|
)
|
||||||
return {"total": total, "replied": replied}
|
).one()
|
||||||
|
return {"total": int(row.total or 0), "replied": int(row.replied or 0)}
|
||||||
|
|
||||||
|
|
||||||
@app.get("/api/logs", response_model=List[LogResponse])
|
@app.get("/api/logs", response_model=List[LogResponse])
|
||||||
@@ -2078,14 +2567,25 @@ async def get_logs(
|
|||||||
if account_id is not None:
|
if account_id is not None:
|
||||||
await get_owned_account(db, user, account_id)
|
await get_owned_account(db, user, account_id)
|
||||||
limit = max(1, min(int(limit or 50), 500))
|
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)
|
logs_for_user(user, account_id)
|
||||||
|
.with_only_columns(MessageLog.id)
|
||||||
.order_by(MessageLog.created_at.desc())
|
.order_by(MessageLog.created_at.desc())
|
||||||
.offset(max(0, int(offset or 0)))
|
.offset(max(0, int(offset or 0)))
|
||||||
.limit(limit)
|
.limit(limit)
|
||||||
)
|
)
|
||||||
result = await db.execute(stmt)
|
ordered_ids = list((await db.execute(id_stmt)).scalars().all())
|
||||||
return result.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])
|
@app.get("/api/received-messages", response_model=List[ReceivedMessageLogResponse])
|
||||||
@@ -2099,13 +2599,26 @@ async def get_received_messages(
|
|||||||
if account_id is not None:
|
if account_id is not None:
|
||||||
await get_owned_account(db, user, account_id)
|
await get_owned_account(db, user, account_id)
|
||||||
limit = max(1, min(int(limit or 100), 500))
|
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)
|
received_logs_for_user(user, account_id)
|
||||||
|
.with_only_columns(ReceivedMessageLog.id)
|
||||||
.order_by(ReceivedMessageLog.created_at.desc())
|
.order_by(ReceivedMessageLog.created_at.desc())
|
||||||
.limit(limit)
|
.limit(limit)
|
||||||
)
|
)
|
||||||
result = await db.execute(stmt)
|
ordered_ids = list((await db.execute(id_stmt)).scalars().all())
|
||||||
return result.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])
|
@app.get("/api/system-logs", response_model=List[SystemLogResponse])
|
||||||
|
|||||||
@@ -9,8 +9,10 @@ from typing import Any
|
|||||||
from urllib.parse import quote_plus
|
from urllib.parse import quote_plus
|
||||||
|
|
||||||
from dotenv import load_dotenv
|
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.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
|
BACKEND_DIR = Path(__file__).resolve().parent.parent
|
||||||
PROJECT_ROOT = BACKEND_DIR.parent
|
PROJECT_ROOT = BACKEND_DIR.parent
|
||||||
@@ -37,6 +39,10 @@ def _env_int(name: str, default: int) -> int:
|
|||||||
return default
|
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:
|
def normalize_db_type(value: str | None) -> str:
|
||||||
raw = (value or "sqlite").strip().lower()
|
raw = (value or "sqlite").strip().lower()
|
||||||
if raw in ("postgres", "pgsql"):
|
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]:
|
def engine_kwargs_for_url(url: str) -> dict[str, Any]:
|
||||||
kwargs: dict[str, Any] = {"echo": False}
|
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_pre_ping"] = True
|
||||||
kwargs["pool_recycle"] = 3600
|
kwargs["pool_recycle"] = 3600
|
||||||
# 默认连接池仅 pool_size=5 + max_overflow=10。多账号托管 + 前端并发请求时
|
# 默认连接池仅 pool_size=5 + max_overflow=10。多账号托管 + 前端并发请求时
|
||||||
@@ -155,9 +189,33 @@ def engine_kwargs_for_url(url: str) -> dict[str, Any]:
|
|||||||
return kwargs
|
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:
|
def create_database_engine(config: DatabaseConfig | None = None) -> AsyncEngine:
|
||||||
url = build_database_url(config)
|
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:
|
def mask_database_url(url: str) -> str:
|
||||||
|
|||||||
@@ -33,6 +33,26 @@ def add_column_if_missing(conn, table: str, column: str, ddl_by_dialect: dict[st
|
|||||||
conn.execute(text(ddl))
|
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:
|
def migrate_accounts_table(conn) -> None:
|
||||||
add_column_if_missing(
|
add_column_if_missing(
|
||||||
conn,
|
conn,
|
||||||
@@ -106,6 +126,12 @@ def migrate_accounts_table(conn) -> None:
|
|||||||
"douyin_uid",
|
"douyin_uid",
|
||||||
{"default": "ALTER TABLE accounts ADD COLUMN douyin_uid VARCHAR(64)"},
|
{"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:
|
def migrate_account_videos_table(conn) -> None:
|
||||||
@@ -124,6 +150,36 @@ def migrate_message_logs_table(conn) -> None:
|
|||||||
"sender_avatar",
|
"sender_avatar",
|
||||||
{"default": "ALTER TABLE message_logs ADD COLUMN sender_avatar TEXT"},
|
{"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:
|
def migrate_rules_table(conn) -> None:
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from sqlalchemy import Column, Integer, String, Boolean, DateTime, ForeignKey, Text, UniqueConstraint
|
from sqlalchemy import Column, Integer, String, Boolean, DateTime, ForeignKey, Index, Text, UniqueConstraint
|
||||||
from sqlalchemy.orm import relationship
|
from sqlalchemy.orm import relationship, validates
|
||||||
from .database import Base
|
from .database import Base
|
||||||
|
from utils.log_limits import bound_error_log_content, bound_message_log_content
|
||||||
|
|
||||||
|
|
||||||
class User(Base):
|
class User(Base):
|
||||||
@@ -204,10 +205,23 @@ class MessageLog(Base):
|
|||||||
reply_content = Column(Text, nullable=True) # 回复的消息
|
reply_content = Column(Text, nullable=True) # 回复的消息
|
||||||
status = Column(String(50), default="received") # received, replied, ignored, failed
|
status = Column(String(50), default="received") # received, replied, ignored, failed
|
||||||
error_message = Column(Text, nullable=True)
|
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")
|
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):
|
class ReceivedMessageLog(Base):
|
||||||
"""接收消息原始日志:仅记录收到的消息,内容原样保存。"""
|
"""接收消息原始日志:仅记录收到的消息,内容原样保存。"""
|
||||||
@@ -225,6 +239,14 @@ class ReceivedMessageLog(Base):
|
|||||||
raw_content = Column(Text, nullable=False)
|
raw_content = Column(Text, nullable=False)
|
||||||
created_at = Column(DateTime, default=datetime.utcnow, index=True)
|
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):
|
class FollowWelcomeLog(Base):
|
||||||
"""关注欢迎语去重表:每个账号对每个新粉丝只发送一次欢迎语(重启后仍生效)。"""
|
"""关注欢迎语去重表:每个账号对每个新粉丝只发送一次欢迎语(重启后仍生效)。"""
|
||||||
@@ -255,6 +277,10 @@ class SystemLog(Base):
|
|||||||
detail = Column(Text, nullable=True) # 详细原因
|
detail = Column(Text, nullable=True) # 详细原因
|
||||||
created_at = Column(DateTime, default=datetime.utcnow, index=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):
|
class PaymentOrder(Base):
|
||||||
"""账号额度购买订单。"""
|
"""账号额度购买订单。"""
|
||||||
|
|||||||
@@ -18,11 +18,44 @@ StartHandler = Callable[[int], Awaitable[dict[str, Any]]]
|
|||||||
JobToken = tuple[str, int]
|
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:
|
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:
|
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):
|
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:
|
def _utc_now() -> str:
|
||||||
@@ -55,28 +88,40 @@ class BatchStartQueue:
|
|||||||
handler: StartHandler,
|
handler: StartHandler,
|
||||||
concurrency: int | None = None,
|
concurrency: int | None = None,
|
||||||
max_batches: int = 100,
|
max_batches: int = 100,
|
||||||
|
timeout_seconds: float | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._handler = handler
|
self._handler = handler
|
||||||
self.concurrency = max(1, int(concurrency or _configured_concurrency()))
|
self.concurrency = max(1, int(concurrency or _configured_concurrency()))
|
||||||
self.max_batches = max(10, int(max_batches or 100))
|
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._queue: asyncio.Queue[JobToken] = asyncio.Queue()
|
||||||
self._pending_jobs: dict[int, JobToken] = {}
|
self._pending_jobs: dict[int, JobToken] = {}
|
||||||
self._active_tasks: dict[int, tuple[JobToken, asyncio.Task]] = {}
|
self._active_tasks: dict[int, tuple[JobToken, asyncio.Task]] = {}
|
||||||
self._batches: dict[str, _BatchRecord] = {}
|
self._batches: dict[str, _BatchRecord] = {}
|
||||||
self._workers: list[asyncio.Task] = []
|
self._workers: list[asyncio.Task] = []
|
||||||
|
self._worker_sequence = 0
|
||||||
self._lock = asyncio.Lock()
|
self._lock = asyncio.Lock()
|
||||||
self._stopping = False
|
self._stopping = False
|
||||||
|
|
||||||
async def _ensure_workers(self) -> None:
|
async def _ensure_workers(self) -> None:
|
||||||
async with self._lock:
|
async with self._lock:
|
||||||
self._workers = [task for task in self._workers if not task.done()]
|
self._workers = [task for task in self._workers if not task.done()]
|
||||||
if self._workers or self._stopping:
|
if self._stopping:
|
||||||
return
|
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(
|
self._workers.append(
|
||||||
asyncio.create_task(
|
asyncio.create_task(
|
||||||
self._worker(index + 1),
|
self._worker(self._worker_sequence),
|
||||||
name=f"account-batch-start-{index + 1}",
|
name=f"account-batch-start-{self._worker_sequence}",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -157,7 +202,21 @@ class BatchStartQueue:
|
|||||||
)
|
)
|
||||||
self._active_tasks[account_id] = (job_token, handler_task)
|
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:
|
async with self._lock:
|
||||||
record = self._batches.get(batch_id)
|
record = self._batches.get(batch_id)
|
||||||
if record:
|
if record:
|
||||||
@@ -171,6 +230,28 @@ class BatchStartQueue:
|
|||||||
elapsed_seconds=round(time.monotonic() - started_at, 3),
|
elapsed_seconds=round(time.monotonic() - started_at, 3),
|
||||||
)
|
)
|
||||||
record.updated_at = _utc_now()
|
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:
|
except asyncio.CancelledError:
|
||||||
async with self._lock:
|
async with self._lock:
|
||||||
record = self._batches.get(batch_id)
|
record = self._batches.get(batch_id)
|
||||||
|
|||||||
@@ -5,8 +5,6 @@ import logging
|
|||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
from rpa_engine.douyin_im.auth import DouyinAuth
|
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 rpa_engine.douyin_im.session import DouyinImSession
|
||||||
from utils.cookie_store import analyze_cookie
|
from utils.cookie_store import analyze_cookie
|
||||||
|
|
||||||
@@ -131,41 +129,57 @@ async def build_cookie_credential_detail(
|
|||||||
async def validate_im_session(
|
async def validate_im_session(
|
||||||
session: DouyinImSession,
|
session: DouyinImSession,
|
||||||
_bypass_global_limit: bool = False,
|
_bypass_global_limit: bool = False,
|
||||||
|
*,
|
||||||
|
startup_priority: bool = False,
|
||||||
) -> tuple[bool, str]:
|
) -> tuple[bool, str]:
|
||||||
if not _bypass_global_limit:
|
if not _bypass_global_limit:
|
||||||
from rpa_engine.douyin_im.traffic_control import get_traffic_controller
|
from rpa_engine.douyin_im.traffic_control import get_traffic_controller
|
||||||
|
|
||||||
controller = get_traffic_controller()
|
controller = get_traffic_controller()
|
||||||
async with controller.background_slot(0, "credential validation"):
|
# Startup validation must not sit behind hundreds of recurring
|
||||||
return await validate_im_session(session, _bypass_global_limit=True)
|
# 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 session.can_direct_im():
|
||||||
if not has_im_session_token(session):
|
if not has_im_session_token(session):
|
||||||
return False, "缺少 sessionid,无法直连 IM"
|
return False, "缺少 sessionid,无法直连 IM"
|
||||||
return False, "Cookie 不满足 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:
|
try:
|
||||||
auth = DouyinAuth.from_im_session(session)
|
auth = DouyinAuth.from_im_session(session)
|
||||||
# 优先用已持久化的 my_uid,避免每次都发起网络 query_my_uid(uid_tt 是加密串,
|
# 优先用已持久化的 my_uid,避免每次都发起网络 query_my_uid(uid_tt 是加密串,
|
||||||
# int() 解析必然失败而回退到网络请求;该请求偶发失败会误判为“未就绪”)。
|
# 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:
|
if not uid:
|
||||||
return False, "服务端未认可当前 Cookie(无法获取用户 UID)"
|
return False, "服务端未认可当前 Cookie(无法获取用户 UID)"
|
||||||
if not auth.is_sign_ready():
|
if not auth.is_sign_ready():
|
||||||
return False, "缺少 IM 签名密钥(web_protect/keys),请用浏览器登录补全"
|
return False, "缺少 IM 签名密钥(web_protect/keys),请用浏览器登录补全"
|
||||||
session.my_uid = int(uid)
|
session.my_uid = int(uid)
|
||||||
async with DouyinImHttpClient(session) as http:
|
# unread_count and ticket probes were previously issued here, but
|
||||||
await http.get_unread_count()
|
# neither result changed the final decision: unread failures become
|
||||||
# 若已缓存到会话票据,优先校验其是否仍新鲜(最理想)。
|
# zero and a stale/missing ticket is resolved lazily at send time.
|
||||||
if session.conv_meta:
|
# Keeping those probes doubled large-batch startup traffic without
|
||||||
ok, reason = await http.verify_messaging_capability(auth, session.my_uid)
|
# adding an authoritative validation signal.
|
||||||
if ok:
|
return True, "IM 凭证就绪(Cookie 与签名密钥齐全,可直连托管)"
|
||||||
return True, reason
|
|
||||||
# 没有缓存会话票据是首次登录的正常情况:会话 ticket 会在发送时即时
|
|
||||||
# 创建/获取(resolve_conversation_meta),因此只要 Cookie + sessionid +
|
|
||||||
# 签名密钥(web_protect/keys) + UID 齐全,就视为可 IM 直连托管,不必再开浏览器。
|
|
||||||
return True, "IM 凭证就绪(Cookie 与签名密钥齐全,可直连托管)"
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"IM session validation failed: {e}")
|
logger.warning(f"IM session validation failed: {e}")
|
||||||
return False, f"IM 运行时验证失败: {e}"
|
return False, f"IM 运行时验证失败: {e}"
|
||||||
@@ -174,6 +188,8 @@ async def validate_im_session(
|
|||||||
async def assess_account_credential(
|
async def assess_account_credential(
|
||||||
cookie_data: Optional[str],
|
cookie_data: Optional[str],
|
||||||
im_session_data: Optional[str] = None,
|
im_session_data: Optional[str] = None,
|
||||||
|
*,
|
||||||
|
startup_priority: bool = False,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
cookie_info = analyze_cookie(cookie_data)
|
cookie_info = analyze_cookie(cookie_data)
|
||||||
result = {
|
result = {
|
||||||
@@ -212,7 +228,10 @@ async def assess_account_credential(
|
|||||||
result["should_reset"] = _should_reset_credentials(result)
|
result["should_reset"] = _should_reset_credentials(result)
|
||||||
return 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
|
result["im_ready"] = im_ok
|
||||||
if im_ok:
|
if im_ok:
|
||||||
result["can_skip_browser"] = True
|
result["can_skip_browser"] = True
|
||||||
|
|||||||
@@ -737,7 +737,7 @@ class DouyinImHttpClient:
|
|||||||
pass
|
pass
|
||||||
return total
|
return total
|
||||||
|
|
||||||
async def get_conversations(self) -> list[dict]:
|
async def get_conversations(self, *, enrich_profiles: bool = True) -> list[dict]:
|
||||||
"""拉取会话列表,返回标准化会话"""
|
"""拉取会话列表,返回标准化会话"""
|
||||||
payloads = [
|
payloads = [
|
||||||
{"cursor": 0, "count": 50, "inbox_type": 0},
|
{"cursor": 0, "count": 50, "inbox_type": 0},
|
||||||
@@ -754,7 +754,12 @@ class DouyinImHttpClient:
|
|||||||
if data is None:
|
if data is None:
|
||||||
data = await self._request("GET", "/v1/conversation/list", body)
|
data = await self._request("GET", "/v1/conversation/list", body)
|
||||||
if data is None:
|
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
|
status_code = data.get("status_code") if isinstance(data, dict) else None
|
||||||
error_text = ""
|
error_text = ""
|
||||||
@@ -814,6 +819,9 @@ class DouyinImHttpClient:
|
|||||||
enriched: list[dict] = []
|
enriched: list[dict] = []
|
||||||
for item in conversations:
|
for item in conversations:
|
||||||
conv = enrich_conversation_item(item, my_uid)
|
conv = enrich_conversation_item(item, my_uid)
|
||||||
|
if not enrich_profiles:
|
||||||
|
enriched.append(conv)
|
||||||
|
continue
|
||||||
peer_uid = str(conv.get("peer_uid") or "")
|
peer_uid = str(conv.get("peer_uid") or "")
|
||||||
name = (conv.get("sender_name") or "").strip()
|
name = (conv.get("sender_name") or "").strip()
|
||||||
avatar = str(conv.get("sender_avatar") or "").strip()
|
avatar = str(conv.get("sender_avatar") or "").strip()
|
||||||
|
|||||||
@@ -1,7 +1,11 @@
|
|||||||
import gzip
|
import gzip
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
|
import logging.handlers
|
||||||
|
import os
|
||||||
|
import queue
|
||||||
import re
|
import re
|
||||||
|
import threading
|
||||||
from typing import Any, Optional
|
from typing import Any, Optional
|
||||||
|
|
||||||
from .message_content import (
|
from .message_content import (
|
||||||
@@ -21,7 +25,6 @@ from .message_content import (
|
|||||||
|
|
||||||
logger = logging.getLogger("douyin_im.protocol")
|
logger = logging.getLogger("douyin_im.protocol")
|
||||||
|
|
||||||
import os
|
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
|
|
||||||
@@ -42,6 +45,73 @@ def _is_control_payload(content_json: Any, msg_type: int = 0) -> bool:
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
_WS_DEBUG_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "ws_media_debug.log")
|
_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(
|
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"{datetime.now().isoformat()} type={msg_type} "
|
||||||
f"conv={conversation_id} content={content_str}{extra}\n"
|
f"conv={conversation_id} content={content_str}{extra}\n"
|
||||||
)
|
)
|
||||||
with open(_WS_DEBUG_PATH, "a", encoding="utf-8") as fh:
|
record_limit = _bounded_env_int(
|
||||||
fh.write(line)
|
"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:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@@ -153,8 +228,6 @@ def parse_ws_payload(raw: bytes | str) -> list[dict]:
|
|||||||
content_str = msg.content
|
content_str = msg.content
|
||||||
server_message_id = str(getattr(msg, "server_message_id", "") or "")
|
server_message_id = str(getattr(msg, "server_message_id", "") or "")
|
||||||
|
|
||||||
_dump_ws_message(msg_type, conversation_id, content_str, msg)
|
|
||||||
|
|
||||||
text_content = ""
|
text_content = ""
|
||||||
media_msg: dict = {}
|
media_msg: dict = {}
|
||||||
content_json: dict = {}
|
content_json: dict = {}
|
||||||
@@ -175,6 +248,12 @@ def parse_ws_payload(raw: bytes | str) -> list[dict]:
|
|||||||
)
|
)
|
||||||
return messages
|
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):
|
if _should_emit_ws_message(conversation_id, msg_type):
|
||||||
sender_uid = str(msg.sender)
|
sender_uid = str(msg.sender)
|
||||||
if media_msg and (
|
if media_msg and (
|
||||||
|
|||||||
@@ -87,20 +87,29 @@ class AccountReplyQueue:
|
|||||||
details: Optional[dict[str, Any]] = None,
|
details: Optional[dict[str, Any]] = None,
|
||||||
merge_key: str = "",
|
merge_key: str = "",
|
||||||
merge_keys: Optional[Iterable[str]] = None,
|
merge_keys: Optional[Iterable[str]] = None,
|
||||||
|
immediate_if_idle: bool = False,
|
||||||
) -> int:
|
) -> 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))
|
interval = max(0.0, float(delay_seconds or 0))
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
async with self._state_lock:
|
async with self._state_lock:
|
||||||
if not self._running or not self._task or self._task.done():
|
if not self._running or not self._task or self._task.done():
|
||||||
raise RuntimeError("reply queue is not running")
|
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._tail_due_at = due_at
|
||||||
self._waiting.append(
|
self._waiting.append(
|
||||||
_QueueItem(
|
_QueueItem(
|
||||||
job_id=uuid.uuid4().hex,
|
job_id=uuid.uuid4().hex,
|
||||||
due_at=due_at,
|
due_at=due_at,
|
||||||
slot_seconds=interval,
|
slot_seconds=slot_seconds,
|
||||||
callback=callback,
|
callback=callback,
|
||||||
description=description,
|
description=description,
|
||||||
queued_at=time.time(),
|
queued_at=time.time(),
|
||||||
@@ -256,9 +265,10 @@ class AccountReplyQueue:
|
|||||||
item = self._waiting.pop(selected_index)
|
item = self._waiting.pop(selected_index)
|
||||||
shift_seconds = max(0.0, item.slot_seconds)
|
shift_seconds = max(0.0, item.slot_seconds)
|
||||||
shifted_count = 0
|
shifted_count = 0
|
||||||
for later in self._waiting[selected_index:]:
|
if shift_seconds > 0:
|
||||||
later.due_at -= shift_seconds
|
for later in self._waiting[selected_index:]:
|
||||||
shifted_count += 1
|
later.due_at -= shift_seconds
|
||||||
|
shifted_count += 1
|
||||||
|
|
||||||
item.due_at = asyncio.get_running_loop().time()
|
item.due_at = asyncio.get_running_loop().time()
|
||||||
item.expedited = True
|
item.expedited = True
|
||||||
|
|||||||
@@ -1,6 +1,11 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
|
import inspect
|
||||||
import logging
|
import logging
|
||||||
|
import os
|
||||||
import time
|
import time
|
||||||
|
import weakref
|
||||||
|
from collections import deque
|
||||||
|
from dataclasses import dataclass, field
|
||||||
from typing import Awaitable, Callable, Optional
|
from typing import Awaitable, Callable, Optional
|
||||||
|
|
||||||
from utils import system_logger
|
from utils import system_logger
|
||||||
@@ -24,6 +29,280 @@ logger = logging.getLogger("douyin_im.service")
|
|||||||
MatchReplyFn = Callable[[str], Awaitable[Optional[list[str]]]]
|
MatchReplyFn = Callable[[str], Awaitable[Optional[list[str]]]]
|
||||||
LogFn = Callable[..., Awaitable[None]]
|
LogFn = Callable[..., Awaitable[None]]
|
||||||
ReceivedLogFn = 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:
|
class DouyinImService:
|
||||||
@@ -43,6 +322,7 @@ class DouyinImService:
|
|||||||
refresh_credentials: Optional[Callable[[], Awaitable[bool]]] = None,
|
refresh_credentials: Optional[Callable[[], Awaitable[bool]]] = None,
|
||||||
follow_tick: Optional[Callable[[], Awaitable[None]]] = None,
|
follow_tick: Optional[Callable[[], Awaitable[None]]] = None,
|
||||||
on_session_invalid: Optional[Callable[[str], Awaitable[None]]] = None,
|
on_session_invalid: Optional[Callable[[str], Awaitable[None]]] = None,
|
||||||
|
on_ready: Optional[ReadyFn] = None,
|
||||||
):
|
):
|
||||||
self.session = session
|
self.session = session
|
||||||
self.match_reply = match_reply
|
self.match_reply = match_reply
|
||||||
@@ -53,6 +333,8 @@ class DouyinImService:
|
|||||||
self.follow_tick = follow_tick
|
self.follow_tick = follow_tick
|
||||||
# 由 worker 注入:检测到 IM 登录失效(INVALID_REQUEST)时回调,用于自动下线
|
# 由 worker 注入:检测到 IM 登录失效(INVALID_REQUEST)时回调,用于自动下线
|
||||||
self.on_session_invalid = on_session_invalid
|
self.on_session_invalid = on_session_invalid
|
||||||
|
self._on_ready = on_ready
|
||||||
|
self._ready_notified = False
|
||||||
self._session_invalid_strikes = 0
|
self._session_invalid_strikes = 0
|
||||||
self._session_invalid_fired = False
|
self._session_invalid_fired = False
|
||||||
self.reply_delay_seconds = max(0, int(reply_delay_seconds or 0))
|
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._ws_client: Optional[DouyinImWsClient] = None
|
||||||
self.last_error: str = ""
|
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:
|
def _reply_key(self, conversation_key: str, content: str) -> str:
|
||||||
return f"{conversation_key}::{content}"
|
return f"{conversation_key}::{content}"
|
||||||
|
|
||||||
@@ -578,19 +870,32 @@ class DouyinImService:
|
|||||||
"replies": list(replies),
|
"replies": list(replies),
|
||||||
},
|
},
|
||||||
merge_keys=queue_merge_keys,
|
merge_keys=queue_merge_keys,
|
||||||
|
immediate_if_idle=True,
|
||||||
)
|
)
|
||||||
|
scheduled_wait = 0 if position == 1 else delay_seconds
|
||||||
logger.info(
|
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,
|
sender,
|
||||||
self.account_id,
|
self.account_id,
|
||||||
position,
|
position,
|
||||||
|
scheduled_wait,
|
||||||
delay_seconds,
|
delay_seconds,
|
||||||
)
|
)
|
||||||
|
if position == 1:
|
||||||
|
queue_detail = (
|
||||||
|
f"{sender} 是当前账号队列的首条任务,等待时间为 0 秒;"
|
||||||
|
f"后续任务仍按 {delay_seconds} 秒间隔排队。"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
queue_detail = (
|
||||||
|
f"{sender} 当前排在第 {position} 位;账号生效间隔为 {delay_seconds} 秒,"
|
||||||
|
"后续任务继续依次排队。"
|
||||||
|
)
|
||||||
system_logger.record(
|
system_logger.record(
|
||||||
"自动回复已进入账号队列",
|
"自动回复已进入账号队列",
|
||||||
detail=(
|
detail=(
|
||||||
f"{sender} 当前排在第 {position} 位;账号生效间隔为 {delay_seconds} 秒,"
|
f"{queue_detail} 账号内计时与排位独立;"
|
||||||
"账号内计时与排位独立;到点后再进入全局带宽队列逐条投递。"
|
"发送时仍进入全局带宽队列逐条投递。"
|
||||||
),
|
),
|
||||||
level="info",
|
level="info",
|
||||||
category="send",
|
category="send",
|
||||||
@@ -687,7 +992,12 @@ class DouyinImService:
|
|||||||
f"(sent={sent_any}, count={len(replies)})"
|
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)
|
my_uid = int(self.session.my_uid or 0)
|
||||||
for raw in conversations:
|
for raw in conversations:
|
||||||
conv = enrich_conversation_item(raw, my_uid)
|
conv = enrich_conversation_item(raw, my_uid)
|
||||||
@@ -696,7 +1006,11 @@ class DouyinImService:
|
|||||||
avatar = str(conv.get("sender_avatar") or "").strip()
|
avatar = str(conv.get("sender_avatar") or "").strip()
|
||||||
peer_uid = str(conv.get("peer_uid") or "")
|
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)
|
profile = await fetch_peer_profile(self.session, peer_uid, self.account_id)
|
||||||
if profile.get("nickname"):
|
if profile.get("nickname"):
|
||||||
name = profile["nickname"]
|
name = profile["nickname"]
|
||||||
@@ -717,21 +1031,96 @@ class DouyinImService:
|
|||||||
if peer_uid and name:
|
if peer_uid and name:
|
||||||
self._conv_names[peer_uid] = 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()
|
controller = get_traffic_controller()
|
||||||
async with controller.background_slot(self.account_id, "conversation poll"):
|
async with controller.background_slot(
|
||||||
async with DouyinImHttpClient(self.session, account_id=self.account_id) as http:
|
self.account_id,
|
||||||
unread_total = await http.get_unread_count()
|
"conversation poll",
|
||||||
if unread_total:
|
startup=initial,
|
||||||
logger.info(f"IM unread total: {unread_total}")
|
):
|
||||||
conversations = await http.get_conversations()
|
# Do not keep one idle connection pool per hosted account. At 500
|
||||||
await self._index_conversations(conversations)
|
# 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
|
# Message handling may wait in the global send lane. Do not keep one
|
||||||
# of the scarce background HTTP slots occupied while that happens.
|
# 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)
|
unread = int(conv.get("unread_count") or 0)
|
||||||
if unread > 0 or conv.get("content"):
|
current_preview = str(conv.get("content") or "")
|
||||||
await self._handle_incoming(conv)
|
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):
|
async def _verify_account_uid(self):
|
||||||
"""启动时用 query/user 接口核验账号真实 UID,修正采集端可能取错的 my_uid/device_id。
|
"""启动时用 query/user 接口核验账号真实 UID,修正采集端可能取错的 my_uid/device_id。
|
||||||
@@ -745,7 +1134,11 @@ class DouyinImService:
|
|||||||
from .auth import DouyinAuth
|
from .auth import DouyinAuth
|
||||||
auth = DouyinAuth.from_im_session(self.session)
|
auth = DouyinAuth.from_im_session(self.session)
|
||||||
controller = get_traffic_controller()
|
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:
|
async with DouyinImHttpClient(self.session, account_id=self.account_id) as http:
|
||||||
old = int(self.session.my_uid or 0)
|
old = int(self.session.my_uid or 0)
|
||||||
resolved = await asyncio.to_thread(http._resolve_authoritative_uid, auth)
|
resolved = await asyncio.to_thread(http._resolve_authoritative_uid, auth)
|
||||||
@@ -763,11 +1156,16 @@ class DouyinImService:
|
|||||||
async def run(self):
|
async def run(self):
|
||||||
"""主循环:WebSocket + HTTP 轮询"""
|
"""主循环:WebSocket + HTTP 轮询"""
|
||||||
self._running = True
|
self._running = True
|
||||||
|
_active_service_ids.add(self.account_id)
|
||||||
await self._reply_queue.start()
|
await self._reply_queue.start()
|
||||||
await self._verify_account_uid()
|
await self._verify_account_uid()
|
||||||
# ensure_frontier_ws 可能触发签名/HTTP(阻塞),放线程池避免多账号启动时卡死事件循环
|
# ensure_frontier_ws 可能触发签名/HTTP(阻塞),放线程池避免多账号启动时卡死事件循环
|
||||||
controller = get_traffic_controller()
|
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)
|
await asyncio.to_thread(ensure_frontier_ws, self.session)
|
||||||
has_ws = bool(self.session.frontier_ws_url())
|
has_ws = bool(self.session.frontier_ws_url())
|
||||||
cred_summary = format_session_credential_summary(self.session)
|
cred_summary = format_session_credential_summary(self.session)
|
||||||
@@ -788,7 +1186,11 @@ class DouyinImService:
|
|||||||
from .emoji_pack import ensure_emoji_map, is_fresh
|
from .emoji_pack import ensure_emoji_map, is_fresh
|
||||||
|
|
||||||
if not 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)
|
await asyncio.to_thread(ensure_emoji_map, self.session)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug(f"emoji map preload failed: {e}")
|
logger.debug(f"emoji map preload failed: {e}")
|
||||||
@@ -798,8 +1200,14 @@ class DouyinImService:
|
|||||||
)
|
)
|
||||||
await self._ws_client.start()
|
await self._ws_client.start()
|
||||||
|
|
||||||
|
initial_poll_succeeded = False
|
||||||
|
initial_unread: list[dict] = []
|
||||||
try:
|
try:
|
||||||
await self._poll_conversations()
|
initial_unread = await self._poll_conversations(
|
||||||
|
initial=True,
|
||||||
|
defer_handlers=True,
|
||||||
|
)
|
||||||
|
initial_poll_succeeded = True
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Initial conversation poll failed: {e}")
|
logger.warning(f"Initial conversation poll failed: {e}")
|
||||||
system_logger.record(
|
system_logger.record(
|
||||||
@@ -810,25 +1218,143 @@ class DouyinImService:
|
|||||||
account_id=self.account_id,
|
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:
|
while self._running:
|
||||||
# The initial poll above is authoritative. Sleep before the next
|
# The initial poll above is authoritative. Sleep before the next
|
||||||
# recurring tick so startup cannot issue two back-to-back polls.
|
# recurring tick so startup cannot issue two back-to-back polls.
|
||||||
await asyncio.sleep(5)
|
await asyncio.sleep(5)
|
||||||
if not self._running:
|
if not self._running:
|
||||||
break
|
break
|
||||||
loop_count += 1
|
|
||||||
try:
|
try:
|
||||||
if loop_count % 3 == 0:
|
current_ws_connected = bool(
|
||||||
await self._poll_conversations()
|
self._ws_client
|
||||||
if loop_count % 6 == 0:
|
and getattr(self._ws_client, "connected", False)
|
||||||
logger.info(f"IM direct tick #{loop_count} account={self.account_id}")
|
)
|
||||||
|
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 检测一次新粉丝(独立于私信轮询,失败不影响主循环)
|
# 关注欢迎语:约每 60s 检测一次新粉丝(独立于私信轮询,失败不影响主循环)
|
||||||
if self.follow_tick and loop_count % 12 == 0:
|
if self.follow_tick and loop.time() >= next_follow_tick_at:
|
||||||
try:
|
try:
|
||||||
await self.follow_tick()
|
await self.follow_tick()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"follow welcome tick error: {e}")
|
logger.error(f"follow welcome tick error: {e}")
|
||||||
|
finally:
|
||||||
|
next_follow_tick_at = loop.time() + follow_interval
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"IM poll error: {e}")
|
logger.error(f"IM poll error: {e}")
|
||||||
system_logger.record(
|
system_logger.record(
|
||||||
@@ -841,6 +1367,11 @@ class DouyinImService:
|
|||||||
|
|
||||||
async def stop(self):
|
async def stop(self):
|
||||||
self._running = False
|
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 get_traffic_controller().send_queue.cancel_account(self.account_id)
|
||||||
await self._reply_queue.stop()
|
await self._reply_queue.stop()
|
||||||
if self._ws_client:
|
if self._ws_client:
|
||||||
|
|||||||
@@ -393,8 +393,27 @@ class TrafficController:
|
|||||||
1.0,
|
1.0,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
self._background = asyncio.Semaphore(
|
background_capacity = _env_int("KEFU_BACKGROUND_NETWORK_CONCURRENCY", 4)
|
||||||
_env_int("KEFU_BACKGROUND_NETWORK_CONCURRENCY", 2)
|
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(
|
self._browser = asyncio.Semaphore(
|
||||||
_env_int("KEFU_BROWSER_START_CONCURRENCY", 1)
|
_env_int("KEFU_BROWSER_START_CONCURRENCY", 1)
|
||||||
@@ -410,12 +429,20 @@ class TrafficController:
|
|||||||
)
|
)
|
||||||
self.background_waiting = 0
|
self.background_waiting = 0
|
||||||
self.background_active = 0
|
self.background_active = 0
|
||||||
|
self.background_startup_waiting = 0
|
||||||
|
self.background_startup_active = 0
|
||||||
self.browser_waiting = 0
|
self.browser_waiting = 0
|
||||||
self.browser_active = 0
|
self.browser_active = 0
|
||||||
self.media_proxy_active = 0
|
self.media_proxy_active = 0
|
||||||
|
|
||||||
@asynccontextmanager
|
@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()
|
current_task = asyncio.current_task()
|
||||||
owner_task, depth = self._background_owner.get()
|
owner_task, depth = self._background_owner.get()
|
||||||
if owner_task is current_task and depth > 0:
|
if owner_task is current_task and depth > 0:
|
||||||
@@ -429,12 +456,52 @@ class TrafficController:
|
|||||||
started = asyncio.get_running_loop().time()
|
started = asyncio.get_running_loop().time()
|
||||||
self.background_waiting += 1
|
self.background_waiting += 1
|
||||||
try:
|
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:
|
except BaseException:
|
||||||
self.background_waiting -= 1
|
self.background_waiting -= 1
|
||||||
raise
|
raise
|
||||||
self.background_waiting -= 1
|
self.background_waiting -= 1
|
||||||
self.background_active += 1
|
self.background_active += 1
|
||||||
|
if startup:
|
||||||
|
self.background_startup_active += 1
|
||||||
token = self._background_owner.set((current_task, 1))
|
token = self._background_owner.set((current_task, 1))
|
||||||
waited = asyncio.get_running_loop().time() - started
|
waited = asyncio.get_running_loop().time() - started
|
||||||
if waited >= 1.0:
|
if waited >= 1.0:
|
||||||
@@ -448,6 +515,9 @@ class TrafficController:
|
|||||||
yield
|
yield
|
||||||
finally:
|
finally:
|
||||||
self._background_owner.reset(token)
|
self._background_owner.reset(token)
|
||||||
|
if startup:
|
||||||
|
self.background_startup_active -= 1
|
||||||
|
self._background_startup.release()
|
||||||
self.background_active -= 1
|
self.background_active -= 1
|
||||||
self._background.release()
|
self._background.release()
|
||||||
|
|
||||||
@@ -492,6 +562,8 @@ class TrafficController:
|
|||||||
"background": {
|
"background": {
|
||||||
"active": self.background_active,
|
"active": self.background_active,
|
||||||
"waiting": self.background_waiting,
|
"waiting": self.background_waiting,
|
||||||
|
"startup_active": self.background_startup_active,
|
||||||
|
"startup_waiting": self.background_startup_waiting,
|
||||||
},
|
},
|
||||||
"browser": {
|
"browser": {
|
||||||
"active": self.browser_active,
|
"active": self.browser_active,
|
||||||
|
|||||||
@@ -1,9 +1,10 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
import threading
|
import os
|
||||||
|
import weakref
|
||||||
from typing import Awaitable, Callable, Optional
|
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 utils import system_logger
|
||||||
from .protocol import parse_ws_payload
|
from .protocol import parse_ws_payload
|
||||||
@@ -13,9 +14,102 @@ logger = logging.getLogger("douyin_im.ws")
|
|||||||
|
|
||||||
MessageHandler = Callable[[dict], Awaitable[None]]
|
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:
|
class DouyinImWsClient:
|
||||||
"""直连 frontier-im WebSocket(websocket-client,与 DouYin_Spider 一致)"""
|
"""Async frontier-im WebSocket client with bounded message backpressure."""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -27,12 +121,16 @@ class DouyinImWsClient:
|
|||||||
self.on_message = on_message
|
self.on_message = on_message
|
||||||
self.account_id = account_id
|
self.account_id = account_id
|
||||||
self._running = False
|
self._running = False
|
||||||
|
self.connected = False
|
||||||
self._task: Optional[asyncio.Task] = None
|
self._task: Optional[asyncio.Task] = None
|
||||||
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
self._connection: Optional[WebSocketClientProtocol] = None
|
||||||
self._ws_app: Optional[WebSocketApp] = None
|
self._last_connection_lifetime = 0.0
|
||||||
self._ws_lock = threading.Lock()
|
self._message_queue: Optional[asyncio.Queue[dict]] = None
|
||||||
|
self._dispatcher_task: Optional[asyncio.Task] = None
|
||||||
|
|
||||||
async def start(self):
|
async def start(self):
|
||||||
|
if self._task and not self._task.done():
|
||||||
|
return
|
||||||
url = self.session.frontier_ws_url()
|
url = self.session.frontier_ws_url()
|
||||||
if not url:
|
if not url:
|
||||||
logger.warning("No frontier WebSocket URL captured; WS listener disabled")
|
logger.warning("No frontier WebSocket URL captured; WS listener disabled")
|
||||||
@@ -45,186 +143,309 @@ class DouyinImWsClient:
|
|||||||
)
|
)
|
||||||
return
|
return
|
||||||
self._running = True
|
self._running = True
|
||||||
self._loop = asyncio.get_running_loop()
|
self._ensure_dispatcher()
|
||||||
self._task = asyncio.create_task(self._run_loop(url))
|
self._task = asyncio.create_task(
|
||||||
|
self._run_loop(url),
|
||||||
|
name=f"im-ws-{self.account_id or 'na'}",
|
||||||
|
)
|
||||||
|
|
||||||
async def stop(self):
|
async def stop(self):
|
||||||
self._running = False
|
self._running = False
|
||||||
with self._ws_lock:
|
self.connected = False
|
||||||
if self._ws_app:
|
|
||||||
|
connection = self._connection
|
||||||
|
if connection is not None:
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(
|
||||||
|
connection.close(code=1000, reason="client stopping"),
|
||||||
|
timeout=_CLOSE_GRACE_SECONDS,
|
||||||
|
)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
# Shutdown iterates over every hosted account. One broken
|
||||||
|
# peer must not consume close_timeout repeatedly and turn a
|
||||||
|
# 500-account shutdown into a many-minute operation.
|
||||||
|
logger.debug("Timed out closing IM WebSocket; aborting transport")
|
||||||
try:
|
try:
|
||||||
self._ws_app.close()
|
connection.fail_connection()
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
self._ws_app = None
|
except Exception:
|
||||||
if self._task:
|
logger.debug("Failed to close IM WebSocket cleanly", exc_info=True)
|
||||||
self._task.cancel()
|
|
||||||
|
task = self._task
|
||||||
|
if task and task is not asyncio.current_task() and not task.done():
|
||||||
|
task.cancel()
|
||||||
try:
|
try:
|
||||||
await self._task
|
await task
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
pass
|
pass
|
||||||
|
if self._task is task:
|
||||||
self._task = None
|
self._task = None
|
||||||
|
self._connection = None
|
||||||
|
await self._stop_dispatcher()
|
||||||
|
|
||||||
async def _run_ws_thread(self, url: str):
|
def _record_connection_system_event(
|
||||||
"""在独立守护线程中跑 run_forever,直到连接断开/关闭。
|
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()
|
loop = asyncio.get_running_loop()
|
||||||
done = asyncio.Event()
|
state = _get_loop_state()
|
||||||
error: list[BaseException] = []
|
key = (int(self.account_id or 0), event_key)
|
||||||
|
now = loop.time()
|
||||||
def _runner():
|
last_at = state.system_log_last_at.get(key)
|
||||||
try:
|
if last_at is not None and now - last_at < _system_log_throttle_seconds():
|
||||||
self._connect_sync(url)
|
return False
|
||||||
except BaseException as e:
|
state.system_log_last_at[key] = now
|
||||||
error.append(e)
|
system_logger.record(
|
||||||
finally:
|
message,
|
||||||
try:
|
detail=detail,
|
||||||
loop.call_soon_threadsafe(done.set)
|
level=level,
|
||||||
except RuntimeError:
|
category="ws",
|
||||||
pass # 事件循环已关闭
|
account_id=self.account_id,
|
||||||
|
|
||||||
thread = threading.Thread(
|
|
||||||
target=_runner,
|
|
||||||
name=f"im-ws-{self.account_id or 'na'}",
|
|
||||||
daemon=True,
|
|
||||||
)
|
)
|
||||||
thread.start()
|
return True
|
||||||
try:
|
|
||||||
await done.wait()
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
# stop() 会 close ws_app 使 run_forever 退出,线程随之结束
|
|
||||||
raise
|
|
||||||
if error:
|
|
||||||
raise error[0]
|
|
||||||
|
|
||||||
async def _run_loop(self, url: str):
|
def _reset_connection_system_log_throttle(self) -> None:
|
||||||
retry = 0
|
loop = asyncio.get_running_loop()
|
||||||
while self._running:
|
state = _LOOP_STATES.get(loop)
|
||||||
from .frontier import ensure_frontier_ws
|
if state is None:
|
||||||
from .traffic_control import get_traffic_controller
|
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(阻塞),放线程池避免卡事件循环
|
def _ensure_dispatcher(self) -> None:
|
||||||
controller = get_traffic_controller()
|
if self._dispatcher_task and not self._dispatcher_task.done():
|
||||||
async with controller.background_slot(self.account_id or 0, "websocket prepare"):
|
return
|
||||||
await asyncio.to_thread(ensure_frontier_ws, self.session)
|
if self._message_queue is None:
|
||||||
connect_url = self.session.frontier_ws_url() or url
|
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:
|
try:
|
||||||
logger.info(f"Connecting IM WebSocket: {connect_url[:100]}...")
|
await task
|
||||||
await self._run_ws_thread(connect_url)
|
except asyncio.CancelledError:
|
||||||
retry = 0
|
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:
|
except asyncio.CancelledError:
|
||||||
break
|
break
|
||||||
except Exception as e:
|
except Exception as exc:
|
||||||
logger.warning(f"IM WebSocket error: {e}")
|
if self._running:
|
||||||
system_logger.record(
|
logger.warning("IM WebSocket error: %s", exc)
|
||||||
"实时接收连接异常",
|
self._record_connection_system_event(
|
||||||
detail=f"建立 frontier WebSocket 失败:{e}",
|
"retry",
|
||||||
level="error",
|
"实时接收连接异常",
|
||||||
category="ws",
|
detail=f"建立 frontier WebSocket 失败:{exc}",
|
||||||
account_id=self.account_id,
|
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:
|
if not self._running:
|
||||||
break
|
break
|
||||||
|
|
||||||
retry += 1
|
retry += 1
|
||||||
# Stable per-account jitter prevents every hosted account from
|
wait = _reconnect_delay(self.account_id, retry)
|
||||||
# reconnecting in the same second after a shared network outage.
|
logger.info("IM WebSocket reconnect in %.1fs...", wait)
|
||||||
jitter = ((int(self.account_id or 0) * 2654435761) % 5000) / 1000.0
|
self._record_connection_system_event(
|
||||||
wait = min(30.0, 2.0 * retry) + jitter
|
"retry",
|
||||||
logger.info(f"IM WebSocket reconnect in {wait:.1f}s...")
|
|
||||||
system_logger.record(
|
|
||||||
f"实时接收断开,{wait:.1f}s 后重连",
|
f"实时接收断开,{wait:.1f}s 后重连",
|
||||||
detail="frontier WebSocket 连接已断开,正在自动重连。",
|
detail="frontier WebSocket 连接已断开,正在自动重连。",
|
||||||
level="warning",
|
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):
|
def _connection_headers(self) -> list[tuple[str, str]]:
|
||||||
if not self._loop:
|
headers = [
|
||||||
return
|
("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):
|
async def _run_connection(self, url: str) -> None:
|
||||||
logger.info("IM WebSocket connected")
|
"""Open one connection and dispatch messages sequentially.
|
||||||
system_logger.record(
|
|
||||||
"实时接收通道已连接",
|
|
||||||
detail="frontier WebSocket 已建立,可实时接收私信。",
|
|
||||||
level="success",
|
|
||||||
category="ws",
|
|
||||||
account_id=self.account_id,
|
|
||||||
)
|
|
||||||
|
|
||||||
def on_message(_ws, message):
|
``max_queue`` bounds the library's receive buffer and ``_dispatch``
|
||||||
asyncio.run_coroutine_threadsafe(self._dispatch(message), self._loop)
|
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):
|
loop = asyncio.get_running_loop()
|
||||||
if self._running:
|
connected_at: float | None = None
|
||||||
logger.warning(f"IM WebSocket error: {error}")
|
connection: Optional[WebSocketClientProtocol] = None
|
||||||
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
|
|
||||||
try:
|
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:
|
finally:
|
||||||
with self._ws_lock:
|
if connected_at is not None:
|
||||||
if self._ws_app is ws_app:
|
self._last_connection_lifetime = max(0.0, loop.time() - connected_at)
|
||||||
self._ws_app = None
|
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):
|
async def _dispatch(self, raw):
|
||||||
|
self._ensure_dispatcher()
|
||||||
|
queue = self._message_queue
|
||||||
|
if queue is None:
|
||||||
|
return
|
||||||
if isinstance(raw, str):
|
if isinstance(raw, str):
|
||||||
payload = raw.encode("utf-8", errors="ignore")
|
payload = raw.encode("utf-8", errors="ignore")
|
||||||
else:
|
else:
|
||||||
payload = raw
|
payload = raw
|
||||||
items = parse_ws_payload(payload)
|
items = parse_ws_payload(payload)
|
||||||
for item in items:
|
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:
|
try:
|
||||||
await self.on_message(item)
|
if self._running:
|
||||||
except Exception as e:
|
# Every account owns one dispatcher, preserving its FIFO.
|
||||||
logger.debug(f"WS message handler error: {e}")
|
# 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(
|
system_logger.record(
|
||||||
"实时消息处理失败",
|
"实时消息处理失败",
|
||||||
detail=f"处理收到的私信时出错:{e}",
|
detail=f"处理收到的私信时出错:{exc}",
|
||||||
level="error",
|
level="error",
|
||||||
category="recv",
|
category="recv",
|
||||||
account_id=self.account_id,
|
account_id=self.account_id,
|
||||||
)
|
)
|
||||||
|
finally:
|
||||||
|
queue.task_done()
|
||||||
|
|||||||
@@ -14,6 +14,11 @@ from playwright.async_api import async_playwright
|
|||||||
from models.database import AsyncSessionLocal
|
from models.database import AsyncSessionLocal
|
||||||
from models.models import Account, AutoReplyRule, MessageLog, AccountProfileDetail, FollowWelcomeLog
|
from models.models import Account, AutoReplyRule, MessageLog, AccountProfileDetail, FollowWelcomeLog
|
||||||
from utils.received_message_log import record_received_message
|
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.cookie_store import get_cookie_path, read_cookie_file, analyze_cookie, merge_playwright_cookies
|
||||||
from utils import system_logger
|
from utils import system_logger
|
||||||
from rpa_engine.douyin_im import DouyinImService
|
from rpa_engine.douyin_im import DouyinImService
|
||||||
@@ -64,9 +69,16 @@ def format_error(exc: BaseException) -> str:
|
|||||||
|
|
||||||
|
|
||||||
class DouyinWorker:
|
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.account_id = account_id
|
||||||
self.login_mode = login_mode # auto | im_direct | browser
|
self.login_mode = login_mode # auto | im_direct | browser
|
||||||
|
self.credential_prevalidated = bool(credential_prevalidated)
|
||||||
self.browser = None
|
self.browser = None
|
||||||
self.context = None
|
self.context = None
|
||||||
self.page = None
|
self.page = None
|
||||||
@@ -74,6 +86,8 @@ class DouyinWorker:
|
|||||||
self.is_running = False
|
self.is_running = False
|
||||||
self.stopping = False
|
self.stopping = False
|
||||||
self._task: asyncio.Task | None = None
|
self._task: asyncio.Task | None = None
|
||||||
|
self._startup_ready = asyncio.Event()
|
||||||
|
self._startup_error = ""
|
||||||
self.session_dir = os.path.join(
|
self.session_dir = os.path.join(
|
||||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||||
"sessions"
|
"sessions"
|
||||||
@@ -104,6 +118,15 @@ class DouyinWorker:
|
|||||||
self._refresh_cooldown = 90.0
|
self._refresh_cooldown = 90.0
|
||||||
self._user_agent: str = ""
|
self._user_agent: str = ""
|
||||||
self._sec_user_id_missing_fired = False
|
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:
|
async def _load_user_agent(self) -> str:
|
||||||
"""读取账号配置的伪装设备头,用于浏览器与 IM 全链路一致。"""
|
"""读取账号配置的伪装设备头,用于浏览器与 IM 全链路一致。"""
|
||||||
@@ -111,9 +134,10 @@ class DouyinWorker:
|
|||||||
return self._user_agent
|
return self._user_agent
|
||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
try:
|
try:
|
||||||
result = await db.execute(select(Account).where(Account.id == self.account_id))
|
result = await db.execute(
|
||||||
account = result.scalar_one_or_none()
|
select(Account.user_agent).where(Account.id == self.account_id)
|
||||||
self._user_agent = resolve_user_agent(account.user_agent if account else None)
|
)
|
||||||
|
self._user_agent = resolve_user_agent(result.scalar_one_or_none())
|
||||||
finally:
|
finally:
|
||||||
await db.close()
|
await db.close()
|
||||||
return self._user_agent
|
return self._user_agent
|
||||||
@@ -145,6 +169,105 @@ class DouyinWorker:
|
|||||||
async def get_db(self):
|
async def get_db(self):
|
||||||
return AsyncSessionLocal()
|
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:
|
async def _load_sec_user_id(self) -> str:
|
||||||
"""Return the locally persisted Douyin sec_user_id for this account."""
|
"""Return the locally persisted Douyin sec_user_id for this account."""
|
||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
@@ -188,10 +311,14 @@ class DouyinWorker:
|
|||||||
"""Resolve and persist sec_user_id once from the account's current Cookie."""
|
"""Resolve and persist sec_user_id once from the account's current Cookie."""
|
||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
try:
|
try:
|
||||||
result = await db.execute(select(Account).where(Account.id == self.account_id))
|
result = await db.execute(
|
||||||
account = result.scalar_one_or_none()
|
select(Account.cookie_data, Account.user_agent).where(
|
||||||
cookie_data = account.cookie_data if account else None
|
Account.id == self.account_id
|
||||||
user_agent = account.user_agent if account else None
|
)
|
||||||
|
)
|
||||||
|
row = result.first()
|
||||||
|
cookie_data = row.cookie_data if row else None
|
||||||
|
user_agent = row.user_agent if row else None
|
||||||
finally:
|
finally:
|
||||||
await db.close()
|
await db.close()
|
||||||
|
|
||||||
@@ -424,10 +551,12 @@ class DouyinWorker:
|
|||||||
"""从数据库或本地文件加载 Playwright storage_state"""
|
"""从数据库或本地文件加载 Playwright storage_state"""
|
||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
try:
|
try:
|
||||||
result = await db.execute(select(Account).where(Account.id == self.account_id))
|
result = await db.execute(
|
||||||
account = result.scalar_one_or_none()
|
select(Account.cookie_data).where(Account.id == self.account_id)
|
||||||
if account and account.cookie_data:
|
)
|
||||||
return json.loads(account.cookie_data)
|
cookie_data = result.scalar_one_or_none()
|
||||||
|
if cookie_data:
|
||||||
|
return json.loads(cookie_data)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Failed to load cookie from database: {e}")
|
logger.warning(f"Failed to load cookie from database: {e}")
|
||||||
finally:
|
finally:
|
||||||
@@ -478,10 +607,10 @@ class DouyinWorker:
|
|||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
saved_im = None
|
saved_im = None
|
||||||
try:
|
try:
|
||||||
result = await db.execute(select(Account).where(Account.id == self.account_id))
|
result = await db.execute(
|
||||||
account = result.scalar_one_or_none()
|
select(Account.im_session_data).where(Account.id == self.account_id)
|
||||||
if account:
|
)
|
||||||
saved_im = account.im_session_data
|
saved_im = result.scalar_one_or_none()
|
||||||
finally:
|
finally:
|
||||||
await db.close()
|
await db.close()
|
||||||
|
|
||||||
@@ -560,7 +689,21 @@ class DouyinWorker:
|
|||||||
"""Cookie 有效时跳过浏览器,直接 IM 直连托管"""
|
"""Cookie 有效时跳过浏览器,直接 IM 直连托管"""
|
||||||
await self._load_user_agent()
|
await self._load_user_agent()
|
||||||
im_session = await self._build_im_session_from_storage(storage_state)
|
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:
|
if not ok:
|
||||||
logger.warning(f"IM session validation failed: {reason}")
|
logger.warning(f"IM session validation failed: {reason}")
|
||||||
system_logger.record(
|
system_logger.record(
|
||||||
@@ -598,10 +741,12 @@ class DouyinWorker:
|
|||||||
async def _load_storage_state(self) -> dict | None:
|
async def _load_storage_state(self) -> dict | None:
|
||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
try:
|
try:
|
||||||
result = await db.execute(select(Account).where(Account.id == self.account_id))
|
result = await db.execute(
|
||||||
account = result.scalar_one_or_none()
|
select(Account.cookie_data).where(Account.id == self.account_id)
|
||||||
if account and account.cookie_data:
|
)
|
||||||
return json.loads(account.cookie_data)
|
cookie_data = result.scalar_one_or_none()
|
||||||
|
if cookie_data:
|
||||||
|
return json.loads(cookie_data)
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
finally:
|
finally:
|
||||||
@@ -687,11 +832,15 @@ class DouyinWorker:
|
|||||||
"""读取账号专属排队间隔;0/NULL 均表示未设置、继承系统默认。"""
|
"""读取账号专属排队间隔;0/NULL 均表示未设置、继承系统默认。"""
|
||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
try:
|
try:
|
||||||
result = await db.execute(select(Account).where(Account.id == self.account_id))
|
result = await db.execute(
|
||||||
account = result.scalar_one_or_none()
|
select(Account.reply_delay_seconds).where(
|
||||||
if not account:
|
Account.id == self.account_id
|
||||||
|
)
|
||||||
|
)
|
||||||
|
reply_delay = result.scalar_one_or_none()
|
||||||
|
if reply_delay is None:
|
||||||
return 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
|
return value if value > 0 else None
|
||||||
finally:
|
finally:
|
||||||
await db.close()
|
await db.close()
|
||||||
@@ -722,11 +871,15 @@ class DouyinWorker:
|
|||||||
"""读取该账号专属冷却秒数;返回 None 表示继承全局设置。"""
|
"""读取该账号专属冷却秒数;返回 None 表示继承全局设置。"""
|
||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
try:
|
try:
|
||||||
result = await db.execute(select(Account).where(Account.id == self.account_id))
|
result = await db.execute(
|
||||||
account = result.scalar_one_or_none()
|
select(Account.reply_cooldown_seconds).where(
|
||||||
if not account or account.reply_cooldown_seconds is None:
|
Account.id == self.account_id
|
||||||
|
)
|
||||||
|
)
|
||||||
|
reply_cooldown = result.scalar_one_or_none()
|
||||||
|
if reply_cooldown is None:
|
||||||
return None
|
return None
|
||||||
return max(0, int(account.reply_cooldown_seconds))
|
return max(0, int(reply_cooldown))
|
||||||
finally:
|
finally:
|
||||||
await db.close()
|
await db.close()
|
||||||
|
|
||||||
@@ -756,6 +909,10 @@ class DouyinWorker:
|
|||||||
|
|
||||||
async def _run_im_direct_service(self, session: DouyinImSession):
|
async def _run_im_direct_service(self, session: DouyinImSession):
|
||||||
"""运行 IM API + WebSocket 直连自动回复"""
|
"""运行 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()
|
reply_delay = await self.get_reply_delay()
|
||||||
im_service = DouyinImService(
|
im_service = DouyinImService(
|
||||||
session=session,
|
session=session,
|
||||||
@@ -770,6 +927,9 @@ class DouyinWorker:
|
|||||||
follow_tick=self.follow_welcome_tick,
|
follow_tick=self.follow_welcome_tick,
|
||||||
# IM 登录失效(INVALID_REQUEST)时自动下线
|
# IM 登录失效(INVALID_REQUEST)时自动下线
|
||||||
on_session_invalid=self.on_im_session_invalid,
|
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,
|
cooldown_resolver=self.resolve_cooldown_seconds,
|
||||||
# 不在发送链路上自动开浏览器刷新:实测重载页面并不会重生 web_protect,
|
# 不在发送链路上自动开浏览器刷新:实测重载页面并不会重生 web_protect,
|
||||||
@@ -1038,15 +1198,24 @@ class DouyinWorker:
|
|||||||
sender_name=sender_name,
|
sender_name=sender_name,
|
||||||
sender_id=sender_id,
|
sender_id=sender_id,
|
||||||
sender_avatar=sender_avatar or None,
|
sender_avatar=sender_avatar or None,
|
||||||
message_content=message,
|
message_content=bound_message_log_content(message),
|
||||||
reply_content=reply,
|
reply_content=(
|
||||||
|
bound_message_log_content(reply) if reply is not None else None
|
||||||
|
),
|
||||||
status=status,
|
status=status,
|
||||||
error_message=error,
|
error_message=(
|
||||||
|
bound_error_log_content(error) if error is not None else None
|
||||||
|
),
|
||||||
created_at=datetime.utcnow()
|
created_at=datetime.utcnow()
|
||||||
)
|
)
|
||||||
db.add(log)
|
db.add(log)
|
||||||
await db.commit()
|
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:
|
except Exception as e:
|
||||||
logger.error(f"Failed to log message: {e}")
|
logger.error(f"Failed to log message: {e}")
|
||||||
await db.rollback()
|
await db.rollback()
|
||||||
@@ -1108,23 +1277,41 @@ class DouyinWorker:
|
|||||||
return
|
return
|
||||||
from rpa_engine.douyin_im.follower_poll import fetch_recent_followers
|
from rpa_engine.douyin_im.follower_poll import fetch_recent_followers
|
||||||
|
|
||||||
# sec_user_id 是托管账号的必要身份字段。这个 tick 始终由 IM 主循环调用,
|
# The direct-service startup normally preloads the lightweight cache.
|
||||||
# 因此即使关闭了关注欢迎语,也能在运行中发现字段被清空并自动退出托管。
|
# Keep a guard-first fallback for legacy/tests/partial initialization:
|
||||||
sec_user_id = await self._require_sec_user_id("托管运行中")
|
# 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:
|
if not sec_user_id:
|
||||||
return
|
return
|
||||||
|
self._follow_welcome_sec_user_id = sec_user_id
|
||||||
|
if not enabled or not content:
|
||||||
|
return
|
||||||
|
|
||||||
# 1) 读账号配置 + 已处理过的粉丝集合
|
# 1) 功能已启用时才读取已处理过的粉丝集合
|
||||||
db = await self.get_db()
|
db = await self.get_db()
|
||||||
try:
|
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 = (
|
rows = (
|
||||||
await db.execute(
|
await db.execute(
|
||||||
select(FollowWelcomeLog.follower_uid).where(
|
select(FollowWelcomeLog.follower_uid).where(
|
||||||
@@ -1230,6 +1417,8 @@ class DouyinWorker:
|
|||||||
return
|
return
|
||||||
self.stopping = False
|
self.stopping = False
|
||||||
self.is_running = True
|
self.is_running = True
|
||||||
|
self._startup_ready = asyncio.Event()
|
||||||
|
self._startup_error = ""
|
||||||
task = asyncio.create_task(
|
task = asyncio.create_task(
|
||||||
self._run_loop(),
|
self._run_loop(),
|
||||||
name=f"douyin-worker-{self.account_id}",
|
name=f"douyin-worker-{self.account_id}",
|
||||||
@@ -1253,6 +1442,7 @@ class DouyinWorker:
|
|||||||
"""停止 RPA 任务"""
|
"""停止 RPA 任务"""
|
||||||
self.stopping = True
|
self.stopping = True
|
||||||
self.is_running = False
|
self.is_running = False
|
||||||
|
self._mark_startup_failed("托管初始化已取消")
|
||||||
if self._im_service:
|
if self._im_service:
|
||||||
await self._im_service.stop()
|
await self._im_service.stop()
|
||||||
task = self._task
|
task = self._task
|
||||||
@@ -1281,6 +1471,7 @@ class DouyinWorker:
|
|||||||
|
|
||||||
if self.login_mode == "im_direct":
|
if self.login_mode == "im_direct":
|
||||||
if not storage_state:
|
if not storage_state:
|
||||||
|
self._mark_startup_failed("未保存 Cookie,无法直连 IM")
|
||||||
await self.update_account_status(
|
await self.update_account_status(
|
||||||
"error",
|
"error",
|
||||||
error_msg="未保存 Cookie,无法直连 IM",
|
error_msg="未保存 Cookie,无法直连 IM",
|
||||||
@@ -1289,6 +1480,9 @@ class DouyinWorker:
|
|||||||
started, reason = await self._try_cookie_only_im_start(storage_state)
|
started, reason = await self._try_cookie_only_im_start(storage_state)
|
||||||
if started:
|
if started:
|
||||||
return
|
return
|
||||||
|
self._mark_startup_failed(
|
||||||
|
reason or "凭证验证失败,无法直连 IM"
|
||||||
|
)
|
||||||
if self.stopping:
|
if self.stopping:
|
||||||
return
|
return
|
||||||
await self.update_account_status(
|
await self.update_account_status(
|
||||||
@@ -1317,11 +1511,13 @@ class DouyinWorker:
|
|||||||
|
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
logger.warning(f"Worker {self.account_id} cancelled")
|
logger.warning(f"Worker {self.account_id} cancelled")
|
||||||
|
self._mark_startup_failed("托管初始化已取消")
|
||||||
if not self.stopping:
|
if not self.stopping:
|
||||||
await self.update_account_status("offline", error_msg="RPA 任务已中断,请重新点击启动")
|
await self.update_account_status("offline", error_msg="RPA 任务已中断,请重新点击启动")
|
||||||
raise
|
raise
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.exception(f"Error in RPA worker loop: {e}")
|
logger.exception(f"Error in RPA worker loop: {e}")
|
||||||
|
self._mark_startup_failed(format_error(e))
|
||||||
if not self.stopping:
|
if not self.stopping:
|
||||||
await self.update_account_status("error", error_msg=format_error(e))
|
await self.update_account_status("error", error_msg=format_error(e))
|
||||||
system_logger.record(
|
system_logger.record(
|
||||||
@@ -1333,6 +1529,7 @@ class DouyinWorker:
|
|||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
self.is_running = False
|
self.is_running = False
|
||||||
|
self._mark_startup_failed()
|
||||||
if not self.stopping:
|
if not self.stopping:
|
||||||
await self.cleanup()
|
await self.cleanup()
|
||||||
|
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import unittest
|
import unittest
|
||||||
@@ -32,6 +33,268 @@ def _fake_db(rows):
|
|||||||
|
|
||||||
|
|
||||||
class BatchStartApiTests(unittest.IsolatedAsyncioTestCase):
|
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):
|
async def test_start_all_uses_lightweight_select_and_submits_once(self):
|
||||||
db = _fake_db([(1, False), (2, True), (3, False), (4, False)])
|
db = _fake_db([(1, False), (2, True), (3, False), (4, False)])
|
||||||
user = SimpleNamespace(id=9, role="admin")
|
user = SimpleNamespace(id=9, role="admin")
|
||||||
|
|||||||
@@ -17,8 +17,18 @@ from rpa_engine import batch_start as batch_start_module
|
|||||||
|
|
||||||
|
|
||||||
class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
|
class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
|
||||||
def _make_queue(self, handler, *, concurrency: int = 2) -> BatchStartQueue:
|
def _make_queue(
|
||||||
queue = BatchStartQueue(handler, concurrency=concurrency)
|
self,
|
||||||
|
handler,
|
||||||
|
*,
|
||||||
|
concurrency: int = 2,
|
||||||
|
timeout_seconds: float | None = None,
|
||||||
|
) -> BatchStartQueue:
|
||||||
|
queue = BatchStartQueue(
|
||||||
|
handler,
|
||||||
|
concurrency=concurrency,
|
||||||
|
timeout_seconds=timeout_seconds,
|
||||||
|
)
|
||||||
self.addAsyncCleanup(queue.stop)
|
self.addAsyncCleanup(queue.stop)
|
||||||
return queue
|
return queue
|
||||||
|
|
||||||
@@ -77,6 +87,42 @@ class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
self.assertEqual(maximum_active, 2)
|
self.assertEqual(maximum_active, 2)
|
||||||
self.assertEqual(active, 0)
|
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):
|
async def test_submit_returns_while_handler_is_blocked(self):
|
||||||
handler_started = asyncio.Event()
|
handler_started = asyncio.Event()
|
||||||
release_handler = 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[32]["status"], "submitted")
|
||||||
self.assertEqual(by_account[33]["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):
|
async def test_failed_account_can_be_submitted_again(self):
|
||||||
attempts = 0
|
attempts = 0
|
||||||
|
|
||||||
|
|||||||
@@ -1,10 +1,13 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import unittest
|
import unittest
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
from pathlib import Path
|
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]
|
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.http_client import DouyinImHttpClient
|
||||||
from rpa_engine.douyin_im.session import DouyinImSession
|
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):
|
class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
|
||||||
@@ -75,6 +80,360 @@ class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
["POST", "GET"],
|
["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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -152,6 +152,9 @@ class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
[
|
[
|
||||||
"lock-enter",
|
"lock-enter",
|
||||||
"authorize",
|
"authorize",
|
||||||
|
# Checks the pooled connection back in before cancel/stop,
|
||||||
|
# which may wait on an in-flight start.
|
||||||
|
"commit",
|
||||||
"cancel",
|
"cancel",
|
||||||
"stop",
|
"stop",
|
||||||
"write-cookie",
|
"write-cookie",
|
||||||
@@ -245,6 +248,9 @@ class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
[
|
[
|
||||||
"lock-enter",
|
"lock-enter",
|
||||||
"authorize",
|
"authorize",
|
||||||
|
# Checks the pooled connection back in before cancel/stop,
|
||||||
|
# which may wait on an in-flight start.
|
||||||
|
"commit",
|
||||||
"cancel",
|
"cancel",
|
||||||
"stop",
|
"stop",
|
||||||
"clear-cookie-file",
|
"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)
|
self.assertIs(result, response)
|
||||||
db.execute.assert_not_awaited()
|
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):
|
async def test_legacy_cookie_request_is_guarded_for_existing_desktop_clients(self):
|
||||||
account = self._account()
|
account = self._account()
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -30,36 +30,127 @@ class AccountReplyQueueTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
await asyncio.sleep(0.002)
|
await asyncio.sleep(0.002)
|
||||||
self.assertEqual(queue.pending_count, 0)
|
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)
|
queue = await self._start_queue(account_id=101)
|
||||||
interval = 0.05
|
interval = 0.05
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
started_at = loop.time()
|
started_at = loop.time()
|
||||||
calls: list[tuple[int, float]] = []
|
calls: list[tuple[int, float]] = []
|
||||||
finished = asyncio.Event()
|
finished = asyncio.Event()
|
||||||
|
first_started = asyncio.Event()
|
||||||
|
release_first = asyncio.Event()
|
||||||
|
|
||||||
def callback_for(index: int):
|
def callback_for(index: int):
|
||||||
async def callback() -> None:
|
async def callback() -> None:
|
||||||
calls.append((index, loop.time() - started_at))
|
calls.append((index, loop.time() - started_at))
|
||||||
|
if index == 0:
|
||||||
|
first_started.set()
|
||||||
|
await release_first.wait()
|
||||||
if len(calls) == 3:
|
if len(calls) == 3:
|
||||||
finished.set()
|
finished.set()
|
||||||
|
|
||||||
return callback
|
return callback
|
||||||
|
|
||||||
for index in range(3):
|
await queue.enqueue(
|
||||||
await queue.enqueue(interval, callback_for(index), description=str(index))
|
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 asyncio.wait_for(finished.wait(), timeout=0.75)
|
||||||
await self._wait_until_idle(queue)
|
await self._wait_until_idle(queue)
|
||||||
|
|
||||||
self.assertEqual([index for index, _ in calls], [0, 1, 2])
|
self.assertEqual([index for index, _ in calls], [0, 1, 2])
|
||||||
elapsed = [timestamp for _, timestamp in calls]
|
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.assertGreaterEqual(timestamp, expected - 0.015)
|
||||||
self.assertLess(timestamp, expected + 0.15)
|
self.assertLess(timestamp, expected + 0.15)
|
||||||
self.assertGreaterEqual(elapsed[1] - elapsed[0], interval - 0.02)
|
self.assertGreaterEqual(elapsed[1] - elapsed[0], interval - 0.02)
|
||||||
self.assertGreaterEqual(elapsed[2] - elapsed[1], 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):
|
async def test_separate_account_queues_reach_first_slot_without_blocking(self):
|
||||||
first_queue = await self._start_queue(account_id=201)
|
first_queue = await self._start_queue(account_id=201)
|
||||||
second_queue = await self._start_queue(account_id=202)
|
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()
|
new_due = datetime.fromisoformat(new["scheduled_at"]).timestamp()
|
||||||
self.assertAlmostEqual(old_due - new_due, interval, delta=0.05)
|
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):
|
async def test_send_now_middle_runs_first_and_only_shifts_jobs_behind_it(self):
|
||||||
queue = await self._start_queue(account_id=503)
|
queue = await self._start_queue(account_id=503)
|
||||||
interval = 0.12
|
interval = 0.12
|
||||||
|
|||||||
@@ -84,6 +84,27 @@ class ReplyQueueApiTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
self.assertEqual(response.total_pending, 1)
|
self.assertEqual(response.total_pending, 1)
|
||||||
self.assertEqual([item.account_id for item in response.items], [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):
|
async def test_offline_account_detail_returns_empty_snapshot(self):
|
||||||
account = SimpleNamespace(id=1, reply_delay_seconds=0)
|
account = SimpleNamespace(id=1, reply_delay_seconds=0)
|
||||||
with (
|
with (
|
||||||
|
|||||||
@@ -23,7 +23,9 @@ from rpa_engine.playwright_worker import DouyinWorker
|
|||||||
|
|
||||||
class _RecordingQueue:
|
class _RecordingQueue:
|
||||||
def __init__(self) -> None:
|
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(
|
async def enqueue(
|
||||||
self,
|
self,
|
||||||
@@ -33,11 +35,19 @@ class _RecordingQueue:
|
|||||||
details=None,
|
details=None,
|
||||||
merge_key="",
|
merge_key="",
|
||||||
merge_keys=None,
|
merge_keys=None,
|
||||||
|
immediate_if_idle=False,
|
||||||
) -> int:
|
) -> int:
|
||||||
keys = merge_keys if merge_keys is not None else [merge_key]
|
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())
|
normalized_keys = frozenset(str(key) for key in keys if str(key or "").strip())
|
||||||
self.jobs.append(
|
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)
|
return len(self.jobs)
|
||||||
|
|
||||||
@@ -75,6 +85,7 @@ class _RecordingQueue:
|
|||||||
job[2],
|
job[2],
|
||||||
merged_details,
|
merged_details,
|
||||||
frozenset(job[4] | incoming_keys),
|
frozenset(job[4] | incoming_keys),
|
||||||
|
job[5],
|
||||||
)
|
)
|
||||||
return {
|
return {
|
||||||
"status": "merged",
|
"status": "merged",
|
||||||
@@ -123,6 +134,8 @@ class ReplyQueueIntegrationTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
|
|
||||||
self.assertEqual(len(service._reply_queue.jobs), 1)
|
self.assertEqual(len(service._reply_queue.jobs), 1)
|
||||||
self.assertEqual(match_reply.await_count, 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]
|
details = service._reply_queue.jobs[0][3]
|
||||||
self.assertEqual(details["sender_name"], "张三")
|
self.assertEqual(details["sender_name"], "张三")
|
||||||
self.assertEqual(details["conversation_id"], "conv-1")
|
self.assertEqual(details["conversation_id"], "conv-1")
|
||||||
|
|||||||
@@ -58,6 +58,7 @@ class SecUserIdGuardTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
worker.is_running = True
|
worker.is_running = True
|
||||||
worker._im_service = SimpleNamespace(_running=True)
|
worker._im_service = SimpleNamespace(_running=True)
|
||||||
worker._require_sec_user_id = AsyncMock(return_value=False)
|
worker._require_sec_user_id = AsyncMock(return_value=False)
|
||||||
|
worker._refresh_follow_welcome_config = AsyncMock()
|
||||||
worker.get_db = AsyncMock()
|
worker.get_db = AsyncMock()
|
||||||
|
|
||||||
await worker.follow_welcome_tick()
|
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
|
# The identity guard must run before the account's follow-welcome flag
|
||||||
# is queried. Otherwise accounts with that feature disabled could stay
|
# is queried. Otherwise accounts with that feature disabled could stay
|
||||||
# hosted indefinitely without a sec_user_id.
|
# hosted indefinitely without a sec_user_id.
|
||||||
|
worker._refresh_follow_welcome_config.assert_not_awaited()
|
||||||
worker.get_db.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):
|
async def test_blank_sec_user_id_is_missing_and_stops_hosting(self):
|
||||||
worker = DouyinWorker(account_id=303, login_mode="im_direct")
|
worker = DouyinWorker(account_id=303, login_mode="im_direct")
|
||||||
worker._load_sec_user_id = AsyncMock(return_value=" ")
|
worker._load_sec_user_id = AsyncMock(return_value=" ")
|
||||||
|
|||||||
@@ -389,6 +389,10 @@ class GlobalSendQueueTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
|
|
||||||
|
|
||||||
class BackgroundTrafficLimitTests(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):
|
async def test_background_slot_respects_configured_concurrency_limit(self):
|
||||||
with patch.dict(
|
with patch.dict(
|
||||||
os.environ,
|
os.environ,
|
||||||
@@ -458,6 +462,147 @@ class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
self.assertEqual(len(tasks_seen), 3)
|
self.assertEqual(len(tasks_seen), 3)
|
||||||
self.assertTrue(all(task is tasks_seen[0] for task in tasks_seen))
|
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):
|
class TrafficControllerLoopIsolationTests(unittest.TestCase):
|
||||||
def test_get_traffic_controller_does_not_reuse_asyncio_primitives(self):
|
def test_get_traffic_controller_does_not_reuse_asyncio_primitives(self):
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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()
|
||||||
@@ -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),
|
||||||
|
)
|
||||||
@@ -8,6 +8,7 @@ from typing import Optional
|
|||||||
|
|
||||||
from models.database import AsyncSessionLocal
|
from models.database import AsyncSessionLocal
|
||||||
from models.models import ReceivedMessageLog
|
from models.models import ReceivedMessageLog
|
||||||
|
from utils.log_limits import bound_raw_message_log_content
|
||||||
|
|
||||||
logger = logging.getLogger("received_message_log")
|
logger = logging.getLogger("received_message_log")
|
||||||
|
|
||||||
@@ -23,8 +24,9 @@ async def record_received_message(
|
|||||||
message_type: Optional[int] = None,
|
message_type: Optional[int] = None,
|
||||||
server_message_id: Optional[str] = None,
|
server_message_id: Optional[str] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
# 原样落库:不做 strip / parse / serialize,空字符串也记录
|
# The live message object remains untouched for matching/replying. Only
|
||||||
store_content = raw_content if raw_content is not None else ""
|
# this diagnostic copy is bounded before persistence.
|
||||||
|
store_content = bound_raw_message_log_content(raw_content)
|
||||||
|
|
||||||
async with AsyncSessionLocal() as db:
|
async with AsyncSessionLocal() as db:
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ from collections import deque
|
|||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
|
from .log_limits import bound_system_log_detail, truncate_text
|
||||||
|
|
||||||
logger = logging.getLogger("douyin_im.system")
|
logger = logging.getLogger("douyin_im.system")
|
||||||
|
|
||||||
_VALID_LEVELS = ("info", "success", "warning", "error")
|
_VALID_LEVELS = ("info", "success", "warning", "error")
|
||||||
@@ -43,14 +45,16 @@ def record(
|
|||||||
"account_id": account_id,
|
"account_id": account_id,
|
||||||
"level": level,
|
"level": level,
|
||||||
"category": category,
|
"category": category,
|
||||||
"event": str(event or ""),
|
"event": truncate_text(event, 255),
|
||||||
"detail": str(detail or ""),
|
"detail": bound_system_log_detail(detail),
|
||||||
"created_at": datetime.utcnow().isoformat(),
|
"created_at": datetime.utcnow().isoformat(),
|
||||||
}
|
}
|
||||||
_buffer.appendleft(entry)
|
_buffer.appendleft(entry)
|
||||||
_pending.append(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":
|
if level == "error":
|
||||||
logger.error(msg)
|
logger.error(msg)
|
||||||
elif level == "warning":
|
elif level == "warning":
|
||||||
|
|||||||
@@ -63,7 +63,7 @@ const handleLogout = () => {
|
|||||||
</a-menu-item>
|
</a-menu-item>
|
||||||
<a-menu-item key="/accounts">
|
<a-menu-item key="/accounts">
|
||||||
<template #icon><UserOutlined /></template>
|
<template #icon><UserOutlined /></template>
|
||||||
<span>账号管理</span>
|
<span>{{ auth.isAdmin ? '账号管理' : '我的账号' }}</span>
|
||||||
</a-menu-item>
|
</a-menu-item>
|
||||||
<a-menu-item key="/messages">
|
<a-menu-item key="/messages">
|
||||||
<template #icon><MessageOutlined /></template>
|
<template #icon><MessageOutlined /></template>
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
<script setup>
|
<script setup>
|
||||||
import { ref, onMounted, onUnmounted, computed, watch } from 'vue'
|
import { ref, onMounted, onUnmounted, computed, watch } from 'vue'
|
||||||
import { useRouter } from 'vue-router'
|
import { useRoute, useRouter } from 'vue-router'
|
||||||
import api from '../api'
|
import api from '../api'
|
||||||
import { message } from 'ant-design-vue'
|
import { message } from 'ant-design-vue'
|
||||||
import { useAuthStore } from '../stores/auth'
|
import { useAuthStore } from '../stores/auth'
|
||||||
@@ -39,6 +39,7 @@ import {
|
|||||||
import { useIsMobile } from '../composables/useIsMobile'
|
import { useIsMobile } from '../composables/useIsMobile'
|
||||||
|
|
||||||
const auth = useAuthStore()
|
const auth = useAuthStore()
|
||||||
|
const route = useRoute()
|
||||||
const router = useRouter()
|
const router = useRouter()
|
||||||
const isMobile = useIsMobile()
|
const isMobile = useIsMobile()
|
||||||
const profileModalWidth = computed(() => (isMobile.value ? 'calc(100vw - 32px)' : 860))
|
const profileModalWidth = computed(() => (isMobile.value ? 'calc(100vw - 32px)' : 860))
|
||||||
@@ -75,6 +76,13 @@ let batchStatusTimer = null
|
|||||||
let batchStatusRequestActive = false
|
let batchStatusRequestActive = false
|
||||||
let batchStatusGeneration = 0
|
let batchStatusGeneration = 0
|
||||||
const BATCH_STATUS_POLL_MS = 1500
|
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 selectedIds = ref([])
|
||||||
const addVisible = ref(false)
|
const addVisible = ref(false)
|
||||||
const addSaving = ref(false)
|
const addSaving = ref(false)
|
||||||
@@ -710,9 +718,18 @@ const applyQueueSnapshot = (data, accountId) => {
|
|||||||
|
|
||||||
const fetchReplyQueueSummaries = async ({ silent = true } = {}) => {
|
const fetchReplyQueueSummaries = async ({ silent = true } = {}) => {
|
||||||
if (queueSummaryLoading.value) return
|
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
|
queueSummaryLoading.value = true
|
||||||
try {
|
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 rows = Array.isArray(res.data?.items) ? res.data.items : []
|
||||||
const next = {}
|
const next = {}
|
||||||
for (const row of rows) {
|
for (const row of rows) {
|
||||||
@@ -983,6 +1000,10 @@ const goAccountRulesPage = (accountId) => {
|
|||||||
}
|
}
|
||||||
|
|
||||||
const openAddModal = () => {
|
const openAddModal = () => {
|
||||||
|
if (!auth.canWrite) {
|
||||||
|
message.warning('当前账号为只读角色,不能添加托管账号')
|
||||||
|
return
|
||||||
|
}
|
||||||
if (!canAddAccount.value) {
|
if (!canAddAccount.value) {
|
||||||
if (canPurchaseSlots.value) {
|
if (canPurchaseSlots.value) {
|
||||||
purchaseVisible.value = true
|
purchaseVisible.value = true
|
||||||
@@ -1147,6 +1168,11 @@ const stopBatchStatusPolling = () => {
|
|||||||
batchStatusTimer = null
|
batchStatusTimer = null
|
||||||
}
|
}
|
||||||
batchStatusRequestActive = false
|
batchStatusRequestActive = false
|
||||||
|
batchStatusPollDelayMs = BATCH_STATUS_POLL_MS
|
||||||
|
batchStatusLastFinished = null
|
||||||
|
batchStatusLastTotal = null
|
||||||
|
batchStatusLastProcessing = null
|
||||||
|
batchStatusLastQueued = null
|
||||||
}
|
}
|
||||||
|
|
||||||
const finishBatchStart = async (snapshot) => {
|
const finishBatchStart = async (snapshot) => {
|
||||||
@@ -1176,6 +1202,7 @@ const finishBatchStart = async (snapshot) => {
|
|||||||
}
|
}
|
||||||
await fetchAccounts()
|
await fetchAccounts()
|
||||||
batchStarting.value = false
|
batchStarting.value = false
|
||||||
|
startReplyQueueSummaryPolling()
|
||||||
}
|
}
|
||||||
|
|
||||||
const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
|
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?.failed_count) || 0) +
|
||||||
Math.max(0, Number(snapshot?.skipped_count) || 0) +
|
Math.max(0, Number(snapshot?.skipped_count) || 0) +
|
||||||
Math.max(0, Number(snapshot?.cancelled_count) || 0)
|
Math.max(0, Number(snapshot?.cancelled_count) || 0)
|
||||||
message.loading({
|
const processing = Math.max(0, Number(snapshot?.processing_count) || 0)
|
||||||
content: `账号启动队列处理中:${Math.min(finished, total)}/${total}`,
|
const queued = Math.max(0, Number(snapshot?.queued_count) || 0)
|
||||||
key: 'batch_start',
|
const visibleFinished = Math.min(finished, total)
|
||||||
duration: 0
|
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) {
|
if (snapshot?.complete) {
|
||||||
await finishBatchStart(snapshot)
|
await finishBatchStart(snapshot)
|
||||||
return true
|
return true
|
||||||
@@ -1225,12 +1272,13 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
|
|||||||
duration: 5
|
duration: 5
|
||||||
})
|
})
|
||||||
await fetchAccounts()
|
await fetchAccounts()
|
||||||
|
startReplyQueueSummaryPolling()
|
||||||
return
|
return
|
||||||
} finally {
|
} finally {
|
||||||
if (generation === batchStatusGeneration) batchStatusRequestActive = false
|
if (generation === batchStatusGeneration) batchStatusRequestActive = false
|
||||||
}
|
}
|
||||||
if (generation === batchStatusGeneration && activeStartBatchId.value === batchId) {
|
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 &&
|
generation === batchStatusGeneration &&
|
||||||
activeStartBatchId.value === batchId
|
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 }) => {
|
const runBatchStart = async ({ accountIds = [], allAccounts = false }) => {
|
||||||
if (batchStarting.value) return
|
if (batchStarting.value) return
|
||||||
batchStarting.value = true
|
batchStarting.value = true
|
||||||
|
stopReplyQueueSummaryPolling()
|
||||||
stopBatchStatusPolling()
|
stopBatchStatusPolling()
|
||||||
const submitGeneration = batchStatusGeneration
|
const submitGeneration = batchStatusGeneration
|
||||||
message.loading({ content: '正在提交账号启动队列...', key: 'batch_start', duration: 0 })
|
message.loading({ content: '正在提交账号启动队列...', key: 'batch_start', duration: 0 })
|
||||||
@@ -1270,6 +1319,7 @@ const runBatchStart = async ({ accountIds = [], allAccounts = false }) => {
|
|||||||
if (submitGeneration !== batchStatusGeneration) return
|
if (submitGeneration !== batchStatusGeneration) return
|
||||||
batchStarting.value = false
|
batchStarting.value = false
|
||||||
activeStartBatchId.value = null
|
activeStartBatchId.value = null
|
||||||
|
startReplyQueueSummaryPolling()
|
||||||
message.error({
|
message.error({
|
||||||
content: error.response?.data?.detail || error.message || '提交批量启动失败',
|
content: error.response?.data?.detail || error.message || '提交批量启动失败',
|
||||||
key: 'batch_start',
|
key: 'batch_start',
|
||||||
@@ -1496,7 +1546,7 @@ const formatCookieTime = (value) => {
|
|||||||
return date.toLocaleString('zh-CN')
|
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_updated_at = data.cookie_updated_at
|
||||||
editForm.value.cookie_count = data.cookie_count
|
editForm.value.cookie_count = data.cookie_count
|
||||||
editForm.value.cookie_valid = data.cookie_valid
|
editForm.value.cookie_valid = data.cookie_valid
|
||||||
@@ -1506,10 +1556,12 @@ const applyCookieResponse = (data) => {
|
|||||||
editForm.value.has_sessionid = !!data.has_sessionid
|
editForm.value.has_sessionid = !!data.has_sessionid
|
||||||
editForm.value.sessionid = data.sessionid || ''
|
editForm.value.sessionid = data.sessionid || ''
|
||||||
editForm.value.sessionid_ss = data.sessionid_ss || ''
|
editForm.value.sessionid_ss = data.sessionid_ss || ''
|
||||||
editForm.value.im_ready = !!data.im_ready
|
if (!preserveRuntimeCredential) {
|
||||||
editForm.value.im_status = data.im_status || ''
|
editForm.value.im_ready = !!data.im_ready
|
||||||
editForm.value.can_skip_browser = !!data.can_skip_browser
|
editForm.value.im_status = data.im_status || ''
|
||||||
editForm.value.should_reset = !!data.should_reset
|
editForm.value.can_skip_browser = !!data.can_skip_browser
|
||||||
|
editForm.value.should_reset = !!data.should_reset
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
const openEditModal = async (acc) => {
|
const openEditModal = async (acc) => {
|
||||||
@@ -1578,7 +1630,7 @@ const refreshCredential = async () => {
|
|||||||
}
|
}
|
||||||
|
|
||||||
const cookieRes = await api.get(`/accounts/${editForm.value.id}/cookie?purpose=management`)
|
const cookieRes = await api.get(`/accounts/${editForm.value.id}/cookie?purpose=management`)
|
||||||
applyCookieResponse(cookieRes.data)
|
applyCookieResponse(cookieRes.data, { preserveRuntimeCredential: true })
|
||||||
message.success('凭证检测完成')
|
message.success('凭证检测完成')
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
message.error(error.response?.data?.detail || '凭证检测失败')
|
message.error(error.response?.data?.detail || '凭证检测失败')
|
||||||
@@ -1705,13 +1757,21 @@ const clearCookie = async () => {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
onMounted(() => {
|
onMounted(async () => {
|
||||||
auth.fetchMe()
|
await Promise.all([
|
||||||
fetchPaymentConfig()
|
auth.fetchMe(),
|
||||||
fetchDeviceProfiles()
|
fetchPaymentConfig(),
|
||||||
fetchAccounts()
|
fetchDeviceProfiles(),
|
||||||
fetchRules()
|
fetchAccounts(),
|
||||||
|
fetchRules(),
|
||||||
|
])
|
||||||
startReplyQueueSummaryPolling()
|
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>
|
||||||
<a-button
|
<a-button
|
||||||
|
v-if="auth.canWrite"
|
||||||
type="primary"
|
type="primary"
|
||||||
class="gradient-btn"
|
class="gradient-btn"
|
||||||
:disabled="!canAddAccount && !canPurchaseSlots"
|
:disabled="!canAddAccount && !canPurchaseSlots"
|
||||||
@@ -2052,7 +2113,7 @@ onUnmounted(() => {
|
|||||||
<UserOutlined style="font-size: 4rem; color: var(--text-muted); margin-bottom: 16px;" />
|
<UserOutlined style="font-size: 4rem; color: var(--text-muted); margin-bottom: 16px;" />
|
||||||
<h3>暂无托管账号</h3>
|
<h3>暂无托管账号</h3>
|
||||||
<p style="color: var(--text-secondary); margin-bottom: 20px;">添加一个抖音账号,开始自动化回复工作吧</p>
|
<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>
|
<template #icon><PlusOutlined /></template>
|
||||||
立即添加
|
立即添加
|
||||||
</a-button>
|
</a-button>
|
||||||
@@ -2257,7 +2318,7 @@ onUnmounted(() => {
|
|||||||
placeholder="0 或留空则继承系统默认"
|
placeholder="0 或留空则继承系统默认"
|
||||||
/>
|
/>
|
||||||
<div class="field-hint">
|
<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>
|
</div>
|
||||||
</a-form-item>
|
</a-form-item>
|
||||||
</a-col>
|
</a-col>
|
||||||
|
|||||||
@@ -2,18 +2,25 @@
|
|||||||
import { ref, onMounted, onUnmounted } from 'vue'
|
import { ref, onMounted, onUnmounted } from 'vue'
|
||||||
import api from '../api'
|
import api from '../api'
|
||||||
import MessageBubble from '../components/MessageBubble.vue'
|
import MessageBubble from '../components/MessageBubble.vue'
|
||||||
|
import { useAuthStore } from '../stores/auth'
|
||||||
import {
|
import {
|
||||||
UserOutlined,
|
UserOutlined,
|
||||||
MessageOutlined,
|
MessageOutlined,
|
||||||
CheckCircleOutlined,
|
CheckCircleOutlined,
|
||||||
ClockCircleOutlined,
|
ClockCircleOutlined,
|
||||||
ArrowRightOutlined,
|
ArrowRightOutlined,
|
||||||
ThunderboltOutlined
|
ThunderboltOutlined,
|
||||||
|
SettingOutlined,
|
||||||
|
PlusOutlined
|
||||||
} from '@ant-design/icons-vue'
|
} from '@ant-design/icons-vue'
|
||||||
|
|
||||||
|
const auth = useAuthStore()
|
||||||
|
|
||||||
const stats = ref({
|
const stats = ref({
|
||||||
totalAccounts: 0,
|
totalAccounts: 0,
|
||||||
activeAccounts: 0,
|
activeAccounts: 0,
|
||||||
|
myAccounts: 0,
|
||||||
|
myActiveAccounts: 0,
|
||||||
totalMessages: 0,
|
totalMessages: 0,
|
||||||
repliedMessages: 0,
|
repliedMessages: 0,
|
||||||
replyRate: '0%'
|
replyRate: '0%'
|
||||||
@@ -24,14 +31,15 @@ const loading = ref(true)
|
|||||||
|
|
||||||
const fetchStats = async () => {
|
const fetchStats = async () => {
|
||||||
try {
|
try {
|
||||||
const [accountsRes, statsRes] = await Promise.all([
|
const [accountStatsRes, statsRes] = await Promise.all([
|
||||||
api.get(`/accounts`),
|
api.get(`/dashboard/account-stats`),
|
||||||
api.get(`/logs/stats`)
|
api.get(`/logs/stats`)
|
||||||
])
|
])
|
||||||
|
|
||||||
const accounts = accountsRes.data
|
stats.value.totalAccounts = accountStatsRes.data.total_accounts || 0
|
||||||
stats.value.totalAccounts = accounts.length
|
stats.value.activeAccounts = accountStatsRes.data.online_accounts || 0
|
||||||
stats.value.activeAccounts = accounts.filter(a => a.status === 'online').length
|
stats.value.myAccounts = accountStatsRes.data.my_accounts || 0
|
||||||
|
stats.value.myActiveAccounts = accountStatsRes.data.my_online_accounts || 0
|
||||||
|
|
||||||
// 后端全量统计所有消息,不受列表条数限制
|
// 后端全量统计所有消息,不受列表条数限制
|
||||||
stats.value.totalMessages = statsRes.data.total || 0
|
stats.value.totalMessages = statsRes.data.total || 0
|
||||||
@@ -61,16 +69,21 @@ const fetchRecentLogs = async () => {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let statsInterval = null
|
let statsInterval = null
|
||||||
|
const refreshVisibleStats = () => {
|
||||||
|
if (document.visibilityState === 'visible') fetchStats()
|
||||||
|
}
|
||||||
|
|
||||||
onMounted(() => {
|
onMounted(() => {
|
||||||
fetchStats()
|
fetchStats()
|
||||||
fetchRecentLogs()
|
fetchRecentLogs()
|
||||||
// 统计卡片每10秒自动刷新一次
|
// 汇总统计需要扫描日志索引,分钟级刷新足够且不会持续压迫数据库。
|
||||||
statsInterval = setInterval(fetchStats, 10000)
|
statsInterval = setInterval(refreshVisibleStats, 60000)
|
||||||
|
document.addEventListener('visibilitychange', refreshVisibleStats)
|
||||||
})
|
})
|
||||||
|
|
||||||
onUnmounted(() => {
|
onUnmounted(() => {
|
||||||
if (statsInterval) clearInterval(statsInterval)
|
if (statsInterval) clearInterval(statsInterval)
|
||||||
|
document.removeEventListener('visibilitychange', refreshVisibleStats)
|
||||||
})
|
})
|
||||||
</script>
|
</script>
|
||||||
|
|
||||||
@@ -84,8 +97,27 @@ onUnmounted(() => {
|
|||||||
多账户自动回复RPA后台。支持快捷扫码登录、状态持久化保存、以及自定义关键字规则精准答复。
|
多账户自动回复RPA后台。支持快捷扫码登录、状态持久化保存、以及自定义关键字规则精准答复。
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
<div class="banner-icon">
|
<div class="banner-side">
|
||||||
<ThunderboltOutlined style="font-size: 4rem; color: #c084fc; opacity: 0.3;" />
|
<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>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
@@ -98,7 +130,7 @@ onUnmounted(() => {
|
|||||||
<UserOutlined />
|
<UserOutlined />
|
||||||
</div>
|
</div>
|
||||||
<div class="stat-info">
|
<div class="stat-info">
|
||||||
<span class="stat-label">托管账号</span>
|
<span class="stat-label">全平台托管账号</span>
|
||||||
<h2 class="stat-value">{{ stats.totalAccounts }}</h2>
|
<h2 class="stat-value">{{ stats.totalAccounts }}</h2>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
@@ -111,7 +143,7 @@ onUnmounted(() => {
|
|||||||
<CheckCircleOutlined />
|
<CheckCircleOutlined />
|
||||||
</div>
|
</div>
|
||||||
<div class="stat-info">
|
<div class="stat-info">
|
||||||
<span class="stat-label">在线运行</span>
|
<span class="stat-label">全平台在线运行</span>
|
||||||
<h2 class="stat-value text-green">{{ stats.activeAccounts }}</h2>
|
<h2 class="stat-value text-green">{{ stats.activeAccounts }}</h2>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
@@ -199,8 +231,9 @@ onUnmounted(() => {
|
|||||||
<div class="quick-actions-grid" style="margin-top: 20px;">
|
<div class="quick-actions-grid" style="margin-top: 20px;">
|
||||||
<router-link to="/accounts" class="quick-action-card">
|
<router-link to="/accounts" class="quick-action-card">
|
||||||
<UserOutlined class="action-icon text-gradient" />
|
<UserOutlined class="action-icon text-gradient" />
|
||||||
<span>账号配置</span>
|
<span>{{ auth.isAdmin ? '账号管理' : '我的账号' }}</span>
|
||||||
<p>扫码登录并托管多个抖音账号</p>
|
<p v-if="auth.isAdmin">管理全平台账号配置与运行状态</p>
|
||||||
|
<p v-else>仅查看和管理自己添加的账号,当前 {{ stats.myAccounts }} 个</p>
|
||||||
</router-link>
|
</router-link>
|
||||||
|
|
||||||
<router-link to="/rules" class="quick-action-card">
|
<router-link to="/rules" class="quick-action-card">
|
||||||
@@ -225,6 +258,20 @@ onUnmounted(() => {
|
|||||||
border-left: 4px solid var(--primary-color);
|
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 {
|
.stat-card {
|
||||||
display: flex;
|
display: flex;
|
||||||
align-items: center;
|
align-items: center;
|
||||||
@@ -436,6 +483,20 @@ onUnmounted(() => {
|
|||||||
display: none;
|
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 {
|
.stat-card {
|
||||||
padding: 16px;
|
padding: 16px;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -39,9 +39,23 @@ const LOGS_PAGE_SIZE = 20
|
|||||||
const hasMoreLogs = ref(true)
|
const hasMoreLogs = ref(true)
|
||||||
const loadingMore = ref(false)
|
const loadingMore = ref(false)
|
||||||
let pollTimer = null
|
let pollTimer = null
|
||||||
|
let pollStopped = false
|
||||||
// 上次日志数据的签名;轮询时数据没变化就跳过赋值,避免无谓的重算与重渲染
|
// 上次日志数据的签名;轮询时数据没变化就跳过赋值,避免无谓的重算与重渲染
|
||||||
let lastLogsSignature = ''
|
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 dedupeMessages = (messages) => {
|
||||||
const map = new Map()
|
const map = new Map()
|
||||||
for (const item of messages) {
|
for (const item of messages) {
|
||||||
@@ -172,7 +186,7 @@ const onConvListScroll = (e) => {
|
|||||||
|
|
||||||
const fetchAccounts = async () => {
|
const fetchAccounts = async () => {
|
||||||
try {
|
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 || []
|
accounts.value = Array.isArray(res.data) ? res.data : res.data?.items || []
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error(error)
|
console.error(error)
|
||||||
@@ -593,10 +607,8 @@ onMounted(async () => {
|
|||||||
}
|
}
|
||||||
await fetchAccounts()
|
await fetchAccounts()
|
||||||
await fetchLogs()
|
await fetchLogs()
|
||||||
pollTimer = setInterval(() => {
|
pollStopped = false
|
||||||
if (document.hidden) return
|
scheduleLogPoll()
|
||||||
fetchLogs(true)
|
|
||||||
}, POLL_INTERVAL_MS)
|
|
||||||
})
|
})
|
||||||
|
|
||||||
watch(
|
watch(
|
||||||
@@ -613,8 +625,9 @@ watch(
|
|||||||
)
|
)
|
||||||
|
|
||||||
onUnmounted(() => {
|
onUnmounted(() => {
|
||||||
|
pollStopped = true
|
||||||
if (pollTimer) {
|
if (pollTimer) {
|
||||||
clearInterval(pollTimer)
|
clearTimeout(pollTimer)
|
||||||
pollTimer = null
|
pollTimer = null
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -39,7 +39,7 @@ const formatPeerId = (conv) => {
|
|||||||
}
|
}
|
||||||
|
|
||||||
const fetchAccounts = async () => {
|
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)
|
accounts.value = res.data.filter(a => a.has_cookie)
|
||||||
if (selectedAccount.value) {
|
if (selectedAccount.value) {
|
||||||
const current = accounts.value.find(a => a.id === selectedAccount.value)
|
const current = accounts.value.find(a => a.id === selectedAccount.value)
|
||||||
|
|||||||
@@ -72,7 +72,7 @@ const fetchLogs = async () => {
|
|||||||
|
|
||||||
const fetchAccounts = async () => {
|
const fetchAccounts = async () => {
|
||||||
try {
|
try {
|
||||||
const res = await api.get('/accounts')
|
const res = await api.get('/account-options')
|
||||||
accounts.value = res.data
|
accounts.value = res.data
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error(error)
|
console.error(error)
|
||||||
|
|||||||
@@ -219,7 +219,7 @@ watch(filterAccountId, () => {
|
|||||||
|
|
||||||
const fetchAccounts = async () => {
|
const fetchAccounts = async () => {
|
||||||
try {
|
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 || []
|
accounts.value = Array.isArray(res.data) ? res.data : res.data?.items || []
|
||||||
if (!cooldownAccountId.value && accounts.value.length) {
|
if (!cooldownAccountId.value && accounts.value.length) {
|
||||||
cooldownAccountId.value = filterAccountId.value
|
cooldownAccountId.value = filterAccountId.value
|
||||||
|
|||||||
@@ -740,7 +740,7 @@ onMounted(() => {
|
|||||||
/>
|
/>
|
||||||
|
|
||||||
<div class="field-hint">
|
<div class="field-hint">
|
||||||
账号未单独设置时使用此间隔。同一账号的第 1 条待回复在 N 秒后发送,第 2 条在 2N 秒后发送,以此类推;各账号队列互不影响。设为 0 表示不启用兜底排队,收到消息后立即回复。
|
账号未单独设置时使用此间隔。同一账号队列为空时,首条回复等待 0 秒并立即发送;后续待回复会话按 N 秒、2N 秒依次排队。各账号队列互不影响,实际发送仍受全局带宽队列保护。设为 0 表示不启用兜底排队,收到消息后立即回复。
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
</a-form-item>
|
</a-form-item>
|
||||||
|
|||||||
@@ -85,7 +85,7 @@ const fetchLogs = async () => {
|
|||||||
|
|
||||||
const fetchAccounts = async () => {
|
const fetchAccounts = async () => {
|
||||||
try {
|
try {
|
||||||
const res = await api.get(`/accounts`)
|
const res = await api.get(`/account-options`)
|
||||||
accounts.value = res.data
|
accounts.value = res.data
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error(error)
|
console.error(error)
|
||||||
|
|||||||
Reference in New Issue
Block a user