diff --git a/backend/main.py b/backend/main.py index 2958c5c..e553199 100644 --- a/backend/main.py +++ b/backend/main.py @@ -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 diff --git a/backend/rpa_engine/batch_start.py b/backend/rpa_engine/batch_start.py index 2bb685e..92c79b3 100644 --- a/backend/rpa_engine/batch_start.py +++ b/backend/rpa_engine/batch_start.py @@ -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) diff --git a/backend/rpa_engine/credential.py b/backend/rpa_engine/credential.py index 413ec92..0c181ba 100644 --- a/backend/rpa_engine/credential.py +++ b/backend/rpa_engine/credential.py @@ -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 diff --git a/backend/rpa_engine/douyin_im/http_client.py b/backend/rpa_engine/douyin_im/http_client.py index ea261a6..178a027 100644 --- a/backend/rpa_engine/douyin_im/http_client.py +++ b/backend/rpa_engine/douyin_im/http_client.py @@ -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() diff --git a/backend/rpa_engine/douyin_im/service.py b/backend/rpa_engine/douyin_im/service.py index 9e06e44..70e815c 100644 --- a/backend/rpa_engine/douyin_im/service.py +++ b/backend/rpa_engine/douyin_im/service.py @@ -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 检测一次新粉丝(独立于私信轮询,失败不影响主循环) diff --git a/backend/rpa_engine/douyin_im/traffic_control.py b/backend/rpa_engine/douyin_im/traffic_control.py index ae9d561..15c9399 100644 --- a/backend/rpa_engine/douyin_im/traffic_control.py +++ b/backend/rpa_engine/douyin_im/traffic_control.py @@ -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, diff --git a/backend/rpa_engine/douyin_im/ws_client.py b/backend/rpa_engine/douyin_im/ws_client.py index 4874697..35e341d 100644 --- a/backend/rpa_engine/douyin_im/ws_client.py +++ b/backend/rpa_engine/douyin_im/ws_client.py @@ -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 diff --git a/backend/tests/test_account_pagination.py b/backend/tests/test_account_pagination.py new file mode 100644 index 0000000..d51ccc9 --- /dev/null +++ b/backend/tests/test_account_pagination.py @@ -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() diff --git a/backend/tests/test_batch_start_queue.py b/backend/tests/test_batch_start_queue.py index 28ad83b..0043330 100644 --- a/backend/tests/test_batch_start_queue.py +++ b/backend/tests/test_batch_start_queue.py @@ -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 diff --git a/backend/tests/test_conversation_poll_bandwidth.py b/backend/tests/test_conversation_poll_bandwidth.py index e21cfe1..9253abe 100644 --- a/backend/tests/test_conversation_poll_bandwidth.py +++ b/backend/tests/test_conversation_poll_bandwidth.py @@ -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() diff --git a/backend/tests/test_reply_queue_api.py b/backend/tests/test_reply_queue_api.py index 49b439e..1b5ff53 100644 --- a/backend/tests/test_reply_queue_api.py +++ b/backend/tests/test_reply_queue_api.py @@ -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 ( diff --git a/backend/tests/test_traffic_control.py b/backend/tests/test_traffic_control.py index eab9bce..0179f8f 100644 --- a/backend/tests/test_traffic_control.py +++ b/backend/tests/test_traffic_control.py @@ -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): diff --git a/douyin.zip b/douyin.zip new file mode 100644 index 0000000..dc4820e Binary files /dev/null and b/douyin.zip differ diff --git a/frontend/src/views/Accounts.vue b/frontend/src/views/Accounts.vue index 5550a4e..e5fd62c 100644 --- a/frontend/src/views/Accounts.vue +++ b/frontend/src/views/Accounts.vue @@ -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',