This commit is contained in:
Your Name
2026-07-28 09:00:19 +08:00
parent 8ba13a8ff9
commit 153db97dc7
14 changed files with 793 additions and 75 deletions
+135 -37
View File
@@ -20,7 +20,7 @@ from settings import CORS_ORIGINS, SERVE_WEB, STATIC_DIR
from help_pages import serve_credential_tool
from web_static import mount_frontend
from pydantic import BaseModel, Field
from sqlalchemy import select, update, delete, text, func, case
from sqlalchemy import select, update, delete, text, func, case, or_, cast, String
from sqlalchemy.ext.asyncio import AsyncSession
from models.database import engine, Base, get_db, AsyncSessionLocal
@@ -1080,42 +1080,91 @@ async def get_accounts(
支持 q(昵称/抖音ID/手机号/账号ID 搜索)与 status 筛选。
Cookie 解析等重逻辑只对当前页执行。
"""
result = await db.execute(accounts_for_user(user))
accounts = result.scalars().all()
# 更新内存中的运行状态与数据库同步,以防异常断开
for acc in accounts:
if page is None:
result = await db.execute(accounts_for_user(user))
accounts = result.scalars().all()
# 兼容旧的全量接口行为。
for acc in accounts:
is_running = manager.is_running(acc.id)
if is_running and acc.status == "offline":
acc.status = "online"
elif not is_running and acc.status in ("online", "logging_in", "starting"):
acc.status = "offline"
return [_build_account_response(acc) for acc in accounts]
keyword = (q or "").strip().lower()
status_filter = (status or "").strip()
stmt = accounts_for_user(user)
running_account_ids = [
int(account_id)
for account_id, worker in list(manager.workers.items())
if worker and worker.is_running
]
runtime_running = Account.id.in_(running_account_ids)
effective_status = case(
(
runtime_running & (Account.status == "offline"),
"online",
),
(
(~runtime_running)
& Account.status.in_(("online", "logging_in", "starting")),
"offline",
),
else_=Account.status,
)
if keyword:
like_value = f"%{keyword}%"
search_conditions = [
func.lower(Account.username).like(like_value),
func.lower(Account.douyin_uid).like(like_value),
func.lower(Account.phone).like(like_value),
func.lower(Account.user_agent).like(like_value),
cast(Account.id, String).like(like_value),
]
matching_profiles = [
profile
for profile in list_device_profiles()
if keyword in str(profile.get("label") or "").lower()
or keyword in str(profile.get("platform") or "").lower()
]
if matching_profiles:
matching_uas = [profile["user_agent"] for profile in matching_profiles]
search_conditions.append(Account.user_agent.in_(matching_uas))
if any(profile.get("id") == "chrome_win120" for profile in matching_profiles):
search_conditions.append(Account.user_agent.is_(None))
stmt = stmt.where(or_(*search_conditions))
if status_filter and status_filter != "all":
if status_filter == "quota_disabled":
stmt = stmt.where(Account.quota_disabled.is_(True))
else:
stmt = stmt.where(
Account.quota_disabled.is_not(True),
effective_status == status_filter,
)
total_result = await db.execute(
stmt.with_only_columns(func.count(Account.id)).order_by(None)
)
total = int(total_result.scalar_one() or 0)
result = await db.execute(
stmt.order_by(Account.id.asc())
.offset((page - 1) * page_size)
.limit(page_size)
)
items = result.scalars().all()
# Only reconcile the current page. The old implementation hydrated every
# account including large Cookie/IM blobs on every refresh, which became
# visibly slow beyond a few hundred accounts.
for acc in items:
is_running = manager.is_running(acc.id)
if is_running and acc.status == "offline":
acc.status = "online"
elif not is_running and acc.status in ("online", "logging_in", "starting"):
acc.status = "offline"
if page is None:
return [_build_account_response(acc) for acc in accounts]
keyword = (q or "").strip().lower()
status_filter = (status or "").strip()
filtered = [
acc
for acc in accounts
if _account_matches_keyword(acc, keyword)
and (
not status_filter
or status_filter == "all"
or (
status_filter == "quota_disabled"
and bool(acc.quota_disabled)
)
or (
status_filter != "quota_disabled"
and not acc.quota_disabled
and acc.status == status_filter
)
)
]
total = len(filtered)
start = (page - 1) * page_size
items = filtered[start : start + page_size]
return {
"items": [_build_account_response(acc) for acc in items],
"total": total,
@@ -1330,14 +1379,55 @@ async def update_account(
@app.get("/api/reply-queues", response_model=ReplyQueueSummaryResponse)
async def get_reply_queue_summaries(
account_ids: Optional[str] = None,
db: AsyncSession = Depends(get_db),
user: User = Depends(get_current_user),
):
"""聚合当前用户可见账号的排队数量,供账号页低频轮询"""
allowed_ids = await owned_account_ids(db, user)
"""聚合可见账号的排队数量账号页可限定为当前页账号"""
requested_ids: list[int] | None = None
if account_ids is not None:
requested_ids = []
seen_ids: set[int] = set()
for raw in str(account_ids).split(","):
token = raw.strip()
if not token:
continue
try:
account_id = int(token)
except (TypeError, ValueError):
raise HTTPException(status_code=400, detail="账号编号格式错误")
if account_id > 0 and account_id not in seen_ids:
seen_ids.add(account_id)
requested_ids.append(account_id)
if len(requested_ids) > 100:
raise HTTPException(status_code=400, detail="单次最多查询 100 个账号的回复队列")
allowed_ids: set[int] | None
if requested_ids is None:
allowed_ids = await owned_account_ids(db, user)
worker_entries = list(manager.workers.items())
else:
if is_admin(user.role):
allowed_ids = set(requested_ids)
elif requested_ids:
owned_result = await db.execute(
select(Account.id).where(
Account.owner_id == user.id,
Account.id.in_(requested_ids),
)
)
allowed_ids = {int(row[0]) for row in owned_result.all()}
else:
allowed_ids = set()
worker_entries = [
(account_id, manager.workers.get(account_id))
for account_id in requested_ids
if account_id in allowed_ids and account_id in manager.workers
]
summaries: list[ReplyQueueSummaryItem] = []
total_pending = 0
for account_id, worker in list(manager.workers.items()):
for account_id, worker in worker_entries:
if allowed_ids is not None and account_id not in allowed_ids:
continue
service = worker._im_service if worker else None
@@ -1688,7 +1778,11 @@ async def validate_account_credential(
account = await get_owned_account(db, user, account_id)
cookie_data = _get_account_cookie_data(account)
assessment = await assess_account_credential(cookie_data, account.im_session_data)
assessment = await assess_account_credential(
cookie_data,
account.im_session_data,
startup_priority=True,
)
return CredentialValidateResponse(**assessment)
@@ -1718,7 +1812,11 @@ async def _start_account_rpa_impl(
return {"status": "running", "message": "RPA worker is already running."}
cookie_data = _get_account_cookie_data(account)
assessment = await assess_account_credential(cookie_data, account.im_session_data)
assessment = await assess_account_credential(
cookie_data,
account.im_session_data,
startup_priority=True,
)
login_mode = requested_login_mode or assessment["login_mode"]
reset_performed = False
+57 -1
View File
@@ -18,6 +18,10 @@ StartHandler = Callable[[int], Awaitable[dict[str, Any]]]
JobToken = tuple[str, int]
class _StartPreparationTimeout(Exception):
"""Internal marker for the queue's own per-account deadline."""
def _configured_concurrency() -> int:
try:
return max(1, min(8, int(os.getenv("KEFU_BATCH_START_CONCURRENCY", "2"))))
@@ -25,6 +29,16 @@ def _configured_concurrency() -> int:
return 2
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:
return datetime.now(timezone.utc).isoformat()
@@ -55,10 +69,16 @@ class BatchStartQueue:
handler: StartHandler,
concurrency: int | None = None,
max_batches: int = 100,
timeout_seconds: float | None = None,
) -> None:
self._handler = handler
self.concurrency = max(1, int(concurrency or _configured_concurrency()))
self.max_batches = max(10, int(max_batches or 100))
self.timeout_seconds = (
_configured_timeout_seconds()
if timeout_seconds is None
else max(0.0, float(timeout_seconds or 0.0))
)
self._queue: asyncio.Queue[JobToken] = asyncio.Queue()
self._pending_jobs: dict[int, JobToken] = {}
self._active_tasks: dict[int, tuple[JobToken, asyncio.Task]] = {}
@@ -157,7 +177,21 @@ class BatchStartQueue:
)
self._active_tasks[account_id] = (job_token, handler_task)
result = await handler_task
if self.timeout_seconds > 0:
try:
result = await asyncio.wait_for(
handler_task,
timeout=self.timeout_seconds,
)
except asyncio.TimeoutError as exc:
# wait_for cancels its task only when this queue's
# deadline expires. Preserve a TimeoutError raised by
# the handler itself as its real account failure.
if handler_task.cancelled():
raise _StartPreparationTimeout from exc
raise
else:
result = await handler_task
async with self._lock:
record = self._batches.get(batch_id)
if record:
@@ -171,6 +205,28 @@ class BatchStartQueue:
elapsed_seconds=round(time.monotonic() - started_at, 3),
)
record.updated_at = _utc_now()
except _StartPreparationTimeout:
elapsed = round(time.monotonic() - started_at, 3)
logger.warning(
"Batch start timed out account=%s worker=%s after %.1fs",
account_id,
worker_number,
self.timeout_seconds,
)
async with self._lock:
record = self._batches.get(batch_id)
if record:
item = record.items[account_id]
if item.get("status") != "cancelled":
item.update(
status="failed",
message=(
f"启动准备超过 {self.timeout_seconds:g} 秒,"
"已跳过并继续处理后续账号"
),
elapsed_seconds=elapsed,
)
record.updated_at = _utc_now()
except asyncio.CancelledError:
async with self._lock:
record = self._batches.get(batch_id)
+27 -15
View File
@@ -6,7 +6,6 @@ from typing import Optional
from rpa_engine.douyin_im.auth import DouyinAuth
from rpa_engine.douyin_im.frontier import ensure_frontier_ws
from rpa_engine.douyin_im.http_client import DouyinImHttpClient
from rpa_engine.douyin_im.session import DouyinImSession
from utils.cookie_store import analyze_cookie
@@ -131,13 +130,26 @@ async def build_cookie_credential_detail(
async def validate_im_session(
session: DouyinImSession,
_bypass_global_limit: bool = False,
*,
startup_priority: bool = False,
) -> tuple[bool, str]:
if not _bypass_global_limit:
from rpa_engine.douyin_im.traffic_control import get_traffic_controller
controller = get_traffic_controller()
async with controller.background_slot(0, "credential validation"):
return await validate_im_session(session, _bypass_global_limit=True)
# Startup validation must not sit behind hundreds of recurring
# conversation polls. It still shares the same global concurrency
# cap, so this changes ordering without increasing bandwidth usage.
async with controller.background_slot(
0,
"credential validation",
startup=startup_priority,
):
return await validate_im_session(
session,
_bypass_global_limit=True,
startup_priority=startup_priority,
)
if not session.can_direct_im():
if not has_im_session_token(session):
@@ -155,17 +167,12 @@ async def validate_im_session(
if not auth.is_sign_ready():
return False, "缺少 IM 签名密钥(web_protect/keys),请用浏览器登录补全"
session.my_uid = int(uid)
async with DouyinImHttpClient(session) as http:
await http.get_unread_count()
# 若已缓存到会话票据,优先校验其是否仍新鲜(最理想)。
if session.conv_meta:
ok, reason = await http.verify_messaging_capability(auth, session.my_uid)
if ok:
return True, reason
# 没有缓存会话票据是首次登录的正常情况:会话 ticket 会在发送时即时
# 创建/获取(resolve_conversation_meta),因此只要 Cookie + sessionid +
# 签名密钥(web_protect/keys) + UID 齐全,就视为可 IM 直连托管,不必再开浏览器。
return True, "IM 凭证就绪(Cookie 与签名密钥齐全,可直连托管)"
# unread_count and ticket probes were previously issued here, but
# neither result changed the final decision: unread failures become
# zero and a stale/missing ticket is resolved lazily at send time.
# Keeping those probes doubled large-batch startup traffic without
# adding an authoritative validation signal.
return True, "IM 凭证就绪(Cookie 与签名密钥齐全,可直连托管)"
except Exception as e:
logger.warning(f"IM session validation failed: {e}")
return False, f"IM 运行时验证失败: {e}"
@@ -174,6 +181,8 @@ async def validate_im_session(
async def assess_account_credential(
cookie_data: Optional[str],
im_session_data: Optional[str] = None,
*,
startup_priority: bool = False,
) -> dict:
cookie_info = analyze_cookie(cookie_data)
result = {
@@ -212,7 +221,10 @@ async def assess_account_credential(
result["should_reset"] = _should_reset_credentials(result)
return result
im_ok, im_reason = await validate_im_session(session)
im_ok, im_reason = await validate_im_session(
session,
startup_priority=startup_priority,
)
result["im_ready"] = im_ok
if im_ok:
result["can_skip_browser"] = True
+10 -2
View File
@@ -737,7 +737,7 @@ class DouyinImHttpClient:
pass
return total
async def get_conversations(self) -> list[dict]:
async def get_conversations(self, *, enrich_profiles: bool = True) -> list[dict]:
"""拉取会话列表,返回标准化会话"""
payloads = [
{"cursor": 0, "count": 50, "inbox_type": 0},
@@ -754,7 +754,12 @@ class DouyinImHttpClient:
if data is None:
data = await self._request("GET", "/v1/conversation/list", body)
if data is None:
continue
# Payload variants only help with schema compatibility. They
# cannot repair a network outage, so stop after POST + GET
# both fail instead of occupying a scarce global slot for up
# to four more full request timeouts.
logger.warning("Conversation poll transport failed; skipping payload fallbacks")
break
status_code = data.get("status_code") if isinstance(data, dict) else None
error_text = ""
@@ -814,6 +819,9 @@ class DouyinImHttpClient:
enriched: list[dict] = []
for item in conversations:
conv = enrich_conversation_item(item, my_uid)
if not enrich_profiles:
enriched.append(conv)
continue
peer_uid = str(conv.get("peer_uid") or "")
name = (conv.get("sender_name") or "").strip()
avatar = str(conv.get("sender_avatar") or "").strip()
+100 -7
View File
@@ -1,5 +1,6 @@
import asyncio
import logging
import os
import time
from typing import Awaitable, Callable, Optional
@@ -26,6 +27,30 @@ LogFn = Callable[..., Awaitable[None]]
ReceivedLogFn = Callable[..., Awaitable[None]]
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) -> 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.
"""
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,
)
spread_ms = max(1, int(interval * 1000))
stagger = ((int(account_id or 0) * 2654435761) % spread_ms) / 1000.0
return interval, stagger
class DouyinImService:
"""抖音 IM 直连服务:WebSocket 实时监听 + HTTP 轮询 + 自动回复"""
@@ -734,11 +759,18 @@ class DouyinImService:
controller = get_traffic_controller()
async with controller.background_slot(self.account_id, "conversation poll"):
async with DouyinImHttpClient(self.session, account_id=self.account_id) as http:
unread_total = await http.get_unread_count()
if unread_total:
logger.info(f"IM unread total: {unread_total}")
conversations = await http.get_conversations()
await self._index_conversations(conversations)
conversations = await http.get_conversations(enrich_profiles=False)
# Profile enrichment may involve several slow third-party requests.
# Run it after releasing the conversation-list slot; each individual
# lookup re-enters the shared controller and yields fairly to startup
# validation and other accounts between profiles.
await self._index_conversations(conversations)
unread_total = sum(
max(0, int(item.get("unread_count") or 0))
for item in conversations
)
if unread_total:
logger.info(f"IM unread total: {unread_total}")
# Message handling may wait in the global send lane. Do not keep one
# of the scarce background HTTP slots occupied while that happens.
for conv in conversations:
@@ -811,8 +843,10 @@ class DouyinImService:
)
await self._ws_client.start()
initial_poll_succeeded = False
try:
await self._poll_conversations()
initial_poll_succeeded = True
except Exception as e:
logger.warning(f"Initial conversation poll failed: {e}")
system_logger.record(
@@ -823,6 +857,34 @@ class DouyinImService:
account_id=self.account_id,
)
ws_connected = bool(
self._ws_client and getattr(self._ws_client, "connected", False)
)
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",
)
loop_count = 0
while self._running:
# The initial poll above is authoritative. Sleep before the next
@@ -832,8 +894,39 @@ class DouyinImService:
break
loop_count += 1
try:
if loop_count % 3 == 0:
await self._poll_conversations()
current_ws_connected = bool(
self._ws_client
and getattr(self._ws_client, "connected", False)
)
if current_ws_connected != ws_connected:
ws_connected = current_ws_connected
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",
)
if loop.time() >= next_conversation_poll_at:
try:
await self._poll_conversations()
finally:
# Advance on both success and failure. Otherwise a
# past deadline retries every five-second loop tick
# during an outage and amplifies traffic.
next_conversation_poll_at = loop.time() + poll_interval
if loop_count % 6 == 0:
logger.info(f"IM direct tick #{loop_count} account={self.account_id}")
# 关注欢迎语:约每 60s 检测一次新粉丝(独立于私信轮询,失败不影响主循环)
@@ -396,6 +396,15 @@ class TrafficController:
self._background = asyncio.Semaphore(
_env_int("KEFU_BACKGROUND_NETWORK_CONCURRENCY", 2)
)
# 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()
self._browser = asyncio.Semaphore(
_env_int("KEFU_BROWSER_START_CONCURRENCY", 1)
)
@@ -410,12 +419,20 @@ class TrafficController:
)
self.background_waiting = 0
self.background_active = 0
self.background_startup_waiting = 0
self.background_startup_active = 0
self.browser_waiting = 0
self.browser_active = 0
self.media_proxy_active = 0
@asynccontextmanager
async def background_slot(self, account_id: int = 0, description: str = "request"):
async def background_slot(
self,
account_id: int = 0,
description: str = "request",
*,
startup: bool = False,
):
current_task = asyncio.current_task()
owner_task, depth = self._background_owner.get()
if owner_task is current_task and depth > 0:
@@ -429,12 +446,29 @@ class TrafficController:
started = asyncio.get_running_loop().time()
self.background_waiting += 1
try:
await self._background.acquire()
if startup:
self.background_startup_waiting += 1
self._background_startup_clear.clear()
try:
await self._background.acquire()
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:
await self._background_startup_clear.wait()
await self._background.acquire()
except BaseException:
self.background_waiting -= 1
raise
self.background_waiting -= 1
self.background_active += 1
if startup:
self.background_startup_active += 1
token = self._background_owner.set((current_task, 1))
waited = asyncio.get_running_loop().time() - started
if waited >= 1.0:
@@ -448,6 +482,8 @@ class TrafficController:
yield
finally:
self._background_owner.reset(token)
if startup:
self.background_startup_active -= 1
self.background_active -= 1
self._background.release()
@@ -492,6 +528,8 @@ class TrafficController:
"background": {
"active": self.background_active,
"waiting": self.background_waiting,
"startup_active": self.background_startup_active,
"startup_waiting": self.background_startup_waiting,
},
"browser": {
"active": self.browser_active,
@@ -27,6 +27,7 @@ class DouyinImWsClient:
self.on_message = on_message
self.account_id = account_id
self._running = False
self.connected = False
self._task: Optional[asyncio.Task] = None
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._ws_app: Optional[WebSocketApp] = None
@@ -50,6 +51,7 @@ class DouyinImWsClient:
async def stop(self):
self._running = False
self.connected = False
with self._ws_lock:
if self._ws_app:
try:
@@ -150,6 +152,7 @@ class DouyinImWsClient:
return
def on_open(_ws):
self.connected = True
logger.info("IM WebSocket connected")
system_logger.record(
"实时接收通道已连接",
@@ -174,6 +177,7 @@ class DouyinImWsClient:
)
def on_close(_ws, code, msg):
self.connected = False
logger.info(f"IM WebSocket closed: code={code}, msg={msg}")
if self._running:
system_logger.record(
@@ -206,6 +210,7 @@ class DouyinImWsClient:
try:
ws_app.run_forever(origin="https://www.douyin.com", ping_interval=20, ping_timeout=10)
finally:
self.connected = False
with self._ws_lock:
if self._ws_app is ws_app:
self._ws_app = None
+152
View File
@@ -0,0 +1,152 @@
from __future__ import annotations
import os
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, 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
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_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()
+60 -2
View File
@@ -17,8 +17,18 @@ from rpa_engine import batch_start as batch_start_module
class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
def _make_queue(self, handler, *, concurrency: int = 2) -> BatchStartQueue:
queue = BatchStartQueue(handler, concurrency=concurrency)
def _make_queue(
self,
handler,
*,
concurrency: int = 2,
timeout_seconds: float | None = None,
) -> BatchStartQueue:
queue = BatchStartQueue(
handler,
concurrency=concurrency,
timeout_seconds=timeout_seconds,
)
self.addAsyncCleanup(queue.stop)
return queue
@@ -168,6 +178,54 @@ class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(by_account[32]["status"], "submitted")
self.assertEqual(by_account[33]["status"], "submitted")
async def test_two_timeouts_release_both_workers_for_following_accounts(self):
never_release = asyncio.Event()
calls: list[int] = []
async def handler(account_id: int) -> dict:
calls.append(account_id)
if account_id in (71, 72):
await never_release.wait()
return {"message": f"started-{account_id}"}
queue = self._make_queue(
handler,
concurrency=2,
timeout_seconds=0.02,
)
with patch.object(batch_start_module.logger, "warning"):
submitted = await queue.submit([71, 72, 73, 74])
completed = await self._wait_for_complete(
queue,
submitted["batch_id"],
)
self.assertEqual(calls, [71, 72, 73, 74])
self.assertEqual(completed["failed_count"], 2)
self.assertEqual(completed["submitted_count"], 2)
by_account = {item["account_id"]: item for item in completed["items"]}
self.assertIn("已跳过并继续处理后续账号", by_account[71]["message"])
self.assertIn("已跳过并继续处理后续账号", by_account[72]["message"])
self.assertEqual(by_account[73]["status"], "submitted")
self.assertEqual(by_account[74]["status"], "submitted")
async def test_handler_timeout_error_keeps_its_original_detail(self):
async def handler(_account_id: int) -> dict:
raise asyncio.TimeoutError("upstream request timed out")
queue = self._make_queue(
handler,
concurrency=1,
timeout_seconds=10,
)
with patch.object(batch_start_module.logger, "exception"):
submitted = await queue.submit([75])
completed = await self._wait_for_complete(queue, submitted["batch_id"])
item = completed["items"][0]
self.assertEqual(item["status"], "failed")
self.assertEqual(item["message"], "upstream request timed out")
async def test_failed_account_can_be_submitted_again(self):
attempts = 0
@@ -3,8 +3,9 @@ from __future__ import annotations
import os
import sys
import unittest
from contextlib import asynccontextmanager
from pathlib import Path
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, patch
BACKEND_DIR = Path(__file__).resolve().parents[1]
@@ -16,6 +17,8 @@ if str(BACKEND_DIR) not in sys.path:
from rpa_engine.douyin_im.http_client import DouyinImHttpClient
from rpa_engine.douyin_im.session import DouyinImSession
from rpa_engine.douyin_im.service import DouyinImService, _conversation_poll_timing
from rpa_engine.douyin_im import service as service_module
class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
@@ -75,6 +78,71 @@ class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
["POST", "GET"],
)
async def test_transport_outage_stops_after_one_post_and_get_pair(self):
client = self._make_client()
client._request = AsyncMock(return_value=None)
with self.assertLogs("douyin_im.http", level="WARNING"):
self.assertEqual(await client.get_conversations(), [])
self.assertEqual(client._request.await_count, 2)
self.assertEqual(
[call.args[0] for call in client._request.await_args_list],
["POST", "GET"],
)
def test_websocket_reconciliation_is_slow_and_http_fallback_stays_fast(self):
with patch.dict(
os.environ,
{
"KEFU_WS_RECONCILE_INTERVAL_SECONDS": "120",
"KEFU_HTTP_POLL_INTERVAL_SECONDS": "15",
},
):
ws_interval, ws_stagger = _conversation_poll_timing(123, True)
http_interval, http_stagger = _conversation_poll_timing(123, False)
self.assertEqual(ws_interval, 120)
self.assertEqual(http_interval, 15)
self.assertGreaterEqual(ws_stagger, 0)
self.assertLess(ws_stagger, ws_interval)
self.assertGreaterEqual(http_stagger, 0)
self.assertLess(http_stagger, http_interval)
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([])
if __name__ == "__main__":
unittest.main()
+21
View File
@@ -84,6 +84,27 @@ class ReplyQueueApiTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(response.total_pending, 1)
self.assertEqual([item.account_id for item in response.items], [1])
async def test_summary_with_account_ids_only_scans_requested_page(self):
service_one = _FakeService(1, [_queue_item(1)])
service_two = _FakeService(2, [_queue_item(2, "job-2")])
service_one.get_reply_queue_snapshot = AsyncMock(return_value=service_one.items)
service_two.get_reply_queue_snapshot = AsyncMock(return_value=service_two.items)
main.manager.workers = {
1: SimpleNamespace(is_running=True, _im_service=service_one),
2: SimpleNamespace(is_running=True, _im_service=service_two),
}
response = await main.get_reply_queue_summaries(
account_ids="2",
db=object(),
user=SimpleNamespace(id=1, role="admin"),
)
service_one.get_reply_queue_snapshot.assert_not_awaited()
service_two.get_reply_queue_snapshot.assert_awaited_once()
self.assertEqual(response.total_pending, 1)
self.assertEqual([item.account_id for item in response.items], [2])
async def test_offline_account_detail_returns_empty_snapshot(self):
account = SimpleNamespace(id=1, reply_delay_seconds=0)
with (
+64
View File
@@ -458,6 +458,70 @@ class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(len(tasks_seen), 3)
self.assertTrue(all(task is tasks_seen[0] for task in tasks_seen))
async def test_startup_request_is_not_buried_behind_normal_backlog(self):
with patch.dict(
os.environ,
{"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1"},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
releases = {
name: asyncio.Event()
for name in ("active", "normal-1", "normal-2", "startup")
}
async def request(name: str, *, startup: bool = False) -> None:
async with controller.background_slot(
1,
name,
startup=startup,
):
entered.append(name)
await releases[name].wait()
active = asyncio.create_task(request("active"))
while entered != ["active"]:
await asyncio.sleep(0)
normal_one = asyncio.create_task(request("normal-1"))
normal_two = asyncio.create_task(request("normal-2"))
# Let one normal request reach the shared semaphore while the other is
# held at normal admission, then add the priority startup request.
await asyncio.sleep(0)
await asyncio.sleep(0)
startup = asyncio.create_task(request("startup", startup=True))
await asyncio.sleep(0)
releases["active"].set()
while len(entered) < 2:
await asyncio.sleep(0)
self.assertEqual(entered[:2], ["active", "normal-1"])
releases["normal-1"].set()
while len(entered) < 3:
await asyncio.sleep(0)
self.assertEqual(entered[:3], ["active", "normal-1", "startup"])
releases["startup"].set()
while len(entered) < 4:
await asyncio.sleep(0)
releases["normal-2"].set()
await asyncio.wait_for(
asyncio.gather(active, normal_one, normal_two, startup),
timeout=0.2,
)
self.assertEqual(
entered,
["active", "normal-1", "startup", "normal-2"],
)
self.assertEqual(controller.background_active, 0)
self.assertEqual(controller.background_waiting, 0)
self.assertEqual(controller.background_startup_active, 0)
self.assertEqual(controller.background_startup_waiting, 0)
class TrafficControllerLoopIsolationTests(unittest.TestCase):
def test_get_traffic_controller_does_not_reuse_asyncio_primitives(self):
BIN
View File
Binary file not shown.
+53 -8
View File
@@ -76,6 +76,13 @@ let batchStatusTimer = null
let batchStatusRequestActive = false
let batchStatusGeneration = 0
const BATCH_STATUS_POLL_MS = 1500
const BATCH_STATUS_MAX_POLL_MS = 5000
const BATCH_STATUS_BACKOFF_STEP_MS = 750
let batchStatusPollDelayMs = BATCH_STATUS_POLL_MS
let batchStatusLastFinished = null
let batchStatusLastTotal = null
let batchStatusLastProcessing = null
let batchStatusLastQueued = null
const selectedIds = ref([])
const addVisible = ref(false)
const addSaving = ref(false)
@@ -711,9 +718,18 @@ const applyQueueSnapshot = (data, accountId) => {
const fetchReplyQueueSummaries = async ({ silent = true } = {}) => {
if (queueSummaryLoading.value) return
const accountIds = accounts.value
.map((account) => Number(account?.id))
.filter((accountId) => Number.isInteger(accountId) && accountId > 0)
if (!accountIds.length) {
replyQueueSummaries.value = {}
return
}
queueSummaryLoading.value = true
try {
const res = await api.get('/reply-queues')
const res = await api.get('/reply-queues', {
params: { account_ids: accountIds.join(',') }
})
const rows = Array.isArray(res.data?.items) ? res.data.items : []
const next = {}
for (const row of rows) {
@@ -1152,6 +1168,11 @@ const stopBatchStatusPolling = () => {
batchStatusTimer = null
}
batchStatusRequestActive = false
batchStatusPollDelayMs = BATCH_STATUS_POLL_MS
batchStatusLastFinished = null
batchStatusLastTotal = null
batchStatusLastProcessing = null
batchStatusLastQueued = null
}
const finishBatchStart = async (snapshot) => {
@@ -1181,6 +1202,7 @@ const finishBatchStart = async (snapshot) => {
}
await fetchAccounts()
batchStarting.value = false
startReplyQueueSummaryPolling()
}
const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
@@ -1196,11 +1218,31 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
Math.max(0, Number(snapshot?.failed_count) || 0) +
Math.max(0, Number(snapshot?.skipped_count) || 0) +
Math.max(0, Number(snapshot?.cancelled_count) || 0)
message.loading({
content: `账号启动队列处理中:${Math.min(finished, total)}/${total}`,
key: 'batch_start',
duration: 0
})
const processing = Math.max(0, Number(snapshot?.processing_count) || 0)
const queued = Math.max(0, Number(snapshot?.queued_count) || 0)
const visibleFinished = Math.min(finished, total)
const progressChanged =
visibleFinished !== batchStatusLastFinished ||
total !== batchStatusLastTotal ||
processing !== batchStatusLastProcessing ||
queued !== batchStatusLastQueued
if (progressChanged) {
batchStatusLastFinished = visibleFinished
batchStatusLastTotal = total
batchStatusLastProcessing = processing
batchStatusLastQueued = queued
batchStatusPollDelayMs = BATCH_STATUS_POLL_MS
message.loading({
content: `账号启动队列处理中:${visibleFinished}/${total}(正在处理 ${processing},等待 ${queued},系统正错峰启动)`,
key: 'batch_start',
duration: 0
})
} else {
batchStatusPollDelayMs = Math.min(
BATCH_STATUS_MAX_POLL_MS,
batchStatusPollDelayMs + BATCH_STATUS_BACKOFF_STEP_MS
)
}
if (snapshot?.complete) {
await finishBatchStart(snapshot)
return true
@@ -1230,12 +1272,13 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
duration: 5
})
await fetchAccounts()
startReplyQueueSummaryPolling()
return
} finally {
if (generation === batchStatusGeneration) batchStatusRequestActive = false
}
if (generation === batchStatusGeneration && activeStartBatchId.value === batchId) {
batchStatusTimer = setTimeout(poll, BATCH_STATUS_POLL_MS)
batchStatusTimer = setTimeout(poll, batchStatusPollDelayMs)
}
}
@@ -1245,7 +1288,7 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
generation === batchStatusGeneration &&
activeStartBatchId.value === batchId
) {
batchStatusTimer = setTimeout(poll, BATCH_STATUS_POLL_MS)
batchStatusTimer = setTimeout(poll, batchStatusPollDelayMs)
}
})
}
@@ -1254,6 +1297,7 @@ const pollBatchStartStatus = (batchId, initialSnapshot = null) => {
const runBatchStart = async ({ accountIds = [], allAccounts = false }) => {
if (batchStarting.value) return
batchStarting.value = true
stopReplyQueueSummaryPolling()
stopBatchStatusPolling()
const submitGeneration = batchStatusGeneration
message.loading({ content: '正在提交账号启动队列...', key: 'batch_start', duration: 0 })
@@ -1275,6 +1319,7 @@ const runBatchStart = async ({ accountIds = [], allAccounts = false }) => {
if (submitGeneration !== batchStatusGeneration) return
batchStarting.value = false
activeStartBatchId.value = null
startReplyQueueSummaryPolling()
message.error({
content: error.response?.data?.detail || error.message || '提交批量启动失败',
key: 'batch_start',