"""虎牙工作台:配置、商品快照、任务批次与日志。""" from fastapi import APIRouter from ._common import * router = APIRouter() @router.get("/task-types") def task_types(current: User = Depends(require_permission("huya:task"))): """返回当前规划的虎牙任务类型。""" return SUPPORTED_TASK_TYPES @router.get("/config", response_model=HuyaConfigOut) def get_config( db: Session = Depends(get_db), current: User = Depends(require_permission("huya:config")), ): """获取虎牙配置。""" config = ensure_huya_config(db) return _config_out(config) @router.put("/config", response_model=HuyaConfigOut) def update_config( req: HuyaConfigUpdate, db: Session = Depends(get_db), current: User = Depends(require_permission("huya:config")), ): """更新虎牙配置。""" config = ensure_huya_config(db) changed_fields = [] for field in HUYA_CONFIG_FIELDS: value = getattr(req, field) if value is not None: setattr(config, field, value.strip()) changed_fields.append(field) apply_huya_config_defaults(config) config.updated_at = datetime.now(UTC) if changed_fields: record_audit( db, current, action="recharge:huya:config", target="huya_config", detail={"changed_fields": changed_fields}, ) db.commit() db.refresh(config) return _config_out(config) @router.get("/workbench-accounts") def list_huya_workbench_accounts( handbook_scope: str = Query(..., pattern="^elite$|^browser_exchange$"), db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """读取当前用户指定工作台(精英宝典/浏览器兑换)的账号集合。""" ids = _workbench_account_ids(db, current, handbook_scope) configured = ( db.query(HuyaWorkbench.id) .filter( HuyaWorkbench.user_id == current.id, HuyaWorkbench.handbook_scope == handbook_scope, ) .first() is not None ) return {"account_ids": ids, "configured": configured} @router.put("/workbench-accounts") def update_huya_workbench_accounts( req: HuyaWorkbenchAccountsUpdate, db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """覆盖当前用户指定工作台(精英宝典/浏览器兑换)的账号集合,支持跨浏览器同步。""" ids = sorted(set(req.account_ids)) if any(account_id < 1 for account_id in ids): raise HTTPException(status_code=400, detail="无效的账号 ID") if ids: _require_huya_task_account_access(db, current, ids) workbench = ( db.query(HuyaWorkbench) .filter( HuyaWorkbench.user_id == current.id, HuyaWorkbench.handbook_scope == req.handbook_scope, ) .first() ) if workbench is None: db.add(HuyaWorkbench(user_id=current.id, handbook_scope=req.handbook_scope)) else: workbench.updated_at = datetime.now(UTC) db.query(HuyaWorkbenchAccount).filter( HuyaWorkbenchAccount.user_id == current.id, HuyaWorkbenchAccount.handbook_scope == req.handbook_scope, ).delete(synchronize_session=False) db.add_all( [ HuyaWorkbenchAccount( user_id=current.id, handbook_scope=req.handbook_scope, account_id=account_id, ) for account_id in ids ] ) db.commit() return {"account_ids": ids, "success": True} @router.get("/goods", response_model=list[HuyaGoodsOut]) def list_goods( db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """查看已缓存的虎牙商品快照。""" rows = db.query(HuyaGoodsSnapshot).order_by(HuyaGoodsSnapshot.id.desc()).all() latest = {} for row in rows: latest.setdefault((row.sid, row.product_id), row) return list(latest.values()) @router.get("/recharge-goods", response_model=list[HuyaRechargeGoodsOut]) def list_recharge_goods( db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """查看已缓存的虎牙充值商品快照。""" rows = ( db.query(HuyaRechargeGoodsSnapshot) .order_by(HuyaRechargeGoodsSnapshot.id.desc()) .all() ) latest = {} for row in rows: latest.setdefault((row.sid, row.spu_id), row) return list(latest.values()) @router.post("/tasks/batch") async def create_task_batch( req: HuyaTaskBatchRequest, db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """创建虎牙任务记录并启动后台执行器。""" if not req.account_ids: raise HTTPException(status_code=400, detail="请选择虎牙账号") _require_huya_task_account_access(db, current, req.account_ids) # 先清掉没有执行器的历史 running/pending,避免前端被假活跃批次锁死。 # 不清理 planned:避免与刚创建的新任务产生竞态。 cleanup_orphan_huya_tasks( db, active_batch_ids=huya_batch_registry.active_ids(), statuses=("pending", "running"), message="任务已中断(无执行器接管)", ) try: batch_id, count = create_planned_tasks( db, req.account_ids, req.task_type, current.id, req.payload, req.handbook_scope, ) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc if count == 0: raise HTTPException(status_code=400, detail="没有有效的虎牙账号") if req.task_type == "create_recharge_order": record_audit( db, current, action="recharge:huya:create", target=f"huya_batch:{batch_id}", detail={ "batch_id": batch_id, "task_type": req.task_type, "count": count, "account_count": len(req.account_ids), }, ) db.commit() log_queue = asyncio.Queue() loop = asyncio.get_running_loop() thread_db = SessionLocal() runner = HuyaBatchRunner( db=thread_db, batch_id=batch_id, task_type=req.task_type, payload=req.payload, log_queue=log_queue, loop=loop, concurrency=req.concurrency, ) huya_batch_registry.register(batch_id, log_queue, loop, runner) thread = threading.Thread(target=runner.run, daemon=True) thread.start() return {"batch_id": batch_id, "count": count, "success": True} @router.get("/tasks", response_model=list[HuyaTaskOut]) def list_tasks( batch_id: str | None = None, handbook_scope: str | None = Query(None, pattern="^(legacy|elite)$"), include_images: bool = Query( False, description="是否返回 base64 小程序码(默认否,轮询请保持 false)" ), db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """查看虎牙任务记录。""" query = _visible_huya_tasks_query(db, current) if batch_id: query = query.filter(HuyaTask.batch_id == batch_id) if handbook_scope: query = query.filter(HuyaTask.handbook_scope == handbook_scope) tasks = query.order_by(HuyaTask.id.desc()).limit(300).all() return [_task_out(task, include_images=include_images) for task in tasks] @router.post("/tasks/cleanup-orphans") def cleanup_orphan_tasks( db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """手动清理没有内存执行器接管的虎牙任务。""" cleaned = cleanup_orphan_huya_tasks( db, active_batch_ids=huya_batch_registry.active_ids(), statuses=("pending", "running"), message="任务已中断(无执行器接管)", ) return { "message": f"已清理 {cleaned} 个残留任务", "cleaned": cleaned, "success": True, } @router.get("/tasks/summary") def tasks_summary( db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """虎牙任务统计。""" return _huya_task_summary(_visible_huya_tasks_query(db, current)) @router.get("/tasks/{task_id}", response_model=HuyaTaskOut) def get_task( task_id: int, db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """获取单条虎牙任务详情(含二维码图片)。""" task = _visible_huya_tasks_query(db, current).filter(HuyaTask.id == task_id).first() if not task: raise HTTPException(status_code=404, detail="任务不存在") return _task_out(task, include_images=True) @router.post("/stop/{batch_id}") def stop_batch( batch_id: str, db: Session = Depends(get_db), current: User = Depends(require_permission("huya:task")), ): """停止正在运行的虎牙批次。""" _require_huya_batch_owner(db, current, batch_id) is_recharge_batch = ( db.query(HuyaTask.id) .filter( HuyaTask.batch_id == batch_id, HuyaTask.task_type == "create_recharge_order" ) .first() is not None ) batch = huya_batch_registry.get(batch_id) if batch: if batch.get("finished"): huya_batch_registry.pop(batch_id) cleaned = cleanup_orphan_huya_tasks( db, batch_id=batch_id, message="批次已结束", ) if cleaned: return { "message": f"批次已结束,已清理 {cleaned} 个残留任务", "success": True, } raise HTTPException(status_code=404, detail="批次已结束") batch["runner"].stop() if is_recharge_batch: record_audit( db, current, action="recharge:huya:stop", target=f"huya_batch:{batch_id}", detail={"batch_id": batch_id, "mode": "running"}, ) db.commit() return {"message": "已发送停止信号", "success": True} cleaned = cleanup_orphan_huya_tasks( db, batch_id=batch_id, message="任务已停止(批次不存在)", ) if cleaned: if is_recharge_batch: record_audit( db, current, action="recharge:huya:stop", target=f"huya_batch:{batch_id}", detail={ "batch_id": batch_id, "mode": "orphan_cleanup", "cleaned": cleaned, }, ) db.commit() return {"message": f"已清理 {cleaned} 个残留任务", "success": True} raise HTTPException(status_code=404, detail="批次不存在或已结束") @router.websocket("/ws/{batch_id}") async def ws_huya_logs(websocket: WebSocket, batch_id: str): """虎牙实时日志推送通道。""" user = authenticate_websocket(websocket) if not user: await websocket.close(code=1008, reason="未授权") return if not _has_huya_perm(user, "huya:task"): await websocket.close(code=1008, reason="无权限") return db = SessionLocal() try: _require_huya_batch_owner(db, user, batch_id) except HTTPException: await websocket.close(code=1008, reason="无权访问该任务批次") return finally: db.close() await websocket.accept() batch = huya_batch_registry.get(batch_id) if not batch: await websocket.send_json({"level": "error", "message": "批次不存在或已结束"}) await websocket.close() return 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) if msg.get("level") == "result": await asyncio.sleep(0.1) break except TimeoutError: await websocket.send_json({"level": "heartbeat", "message": ""}) except WebSocketDisconnect: pass finally: latest = huya_batch_registry.get(batch_id) if latest and latest.get("finished"): huya_batch_registry.pop(batch_id)