优化 ui

web/backend/routers/login.py

@@ -15,12 +15,12 @@ from ..services.login_service import LoginBatchRunner
 
 router = APIRouter(prefix="/api/login", tags=["登录任务"])
 
-# 运行中的批次: batch_id -> {runner, log_queue, loop}
+# 运行中的批次: batch_id -> {log_queue, loop, runner}
 _active_batches: dict[str, dict] = {}
 
 
 @router.post("/batch")
-def create_batch(
+async def create_batch(
     req: LoginBatchRequest,
     db: Session = Depends(get_db),
     current: User = Depends(require_permission("login:batch")),
@@ -46,6 +46,10 @@ def create_batch(
     if not valid_ids:
         raise HTTPException(status_code=403, detail="没有可登录的账号")
 
+    # 在主事件循环中创建 log_queue，传给后台线程
+    log_queue = asyncio.Queue()
+    loop = asyncio.get_running_loop()
+
     # 创建执行器（用独立的 DB 会话，因为在线程中运行）
     thread_db = SessionLocal()
     runner = LoginBatchRunner(
@@ -55,10 +59,19 @@ def create_batch(
         creator_role=current.role,
         max_geetest_retries=req.max_geetest_retries,
         proxy_config=proxy,
+        log_queue=log_queue,
+        loop=loop,
     )
 
     batch_id = runner.batch_id
 
+    # 先注册到全局，再启动线程，确保 WebSocket 连接时能找到
+    _active_batches[batch_id] = {
+        "log_queue": log_queue,
+        "loop": loop,
+        "runner": runner,
+    }
+
     # 启动线程
     thread = threading.Thread(target=runner.run, daemon=True)
     thread.start()
@@ -114,21 +127,24 @@ async def ws_login_logs(websocket: WebSocket, batch_id: str):
     """WebSocket 推送登录实时日志。"""
     await websocket.accept()
 
-    log_queue = asyncio.Queue()
-    loop = asyncio.get_running_loop()
+    # 从已注册的批次中获取 log_queue（由 create_batch 创建）
+    batch = _active_batches.get(batch_id)
+    if not batch:
+        await websocket.send_json({"level": "error", "message": "批次不存在或已结束"})
+        await websocket.close()
+        return
 
-    # 查找已运行的批次，或等待新批次
-    # 简化：直接把 log_queue 注册到全局，前端创建批次后连 ws
-    _active_batches[batch_id] = {
-        "log_queue": log_queue,
-        "loop": loop,
-    }
+    log_queue: asyncio.Queue = batch["log_queue"]
 
     try:
         while True:
             try:
                 msg = await asyncio.wait_for(log_queue.get(), timeout=30)
                 await websocket.send_json(msg)
+                # 收到 result 表示任务结束
+                if msg.get("level") == "result":
+                    await asyncio.sleep(0.1)
+                    break
             except asyncio.TimeoutError:
                 await websocket.send_json({"level": "heartbeat", "message": ""})
     except WebSocketDisconnect: