yml2213
d11fa1032d
一、失败入账号一生（此前完全缺失）
生产实测：`login_ok` 1,456 条、`relogin_failed` 73 条，但**没有 login_failed**，
而 login_tasks 里 attempt_status='failed' 有 247 条 —— 这 247 次失败在账号时间线上
一条痕迹都没有，「这账号为什么一直登录不上」在账号页上无从回答。
- 新增事件类型 login_failed，在 `_mark_attempt_failed` 里写入
- **被停止 / 被中断不写**：分型把「已停止/中断/服务重启」归到 interrupted，
那不是账号的问题 —— 否则运营按一次停止就能给几十个账号刷满假失败
- 事件带分型（captcha/credential/email_code/email_quota/network），
同时落进 detail.kind，以后可直接按「验证码失败」筛账号
- 幂等键 login_task:{id}:failed:{ts}，与成功路径同构
二、失败分型下沉成共用 service
`classify_failure` 原在 automation/jobs/ck_relogin.py，手动登录也要用。
挪到 services/login_failure.py（job 改为 import，`ck_relogin.classify_failure`
这个名字保留，旧调用方与用例不受影响），并提供型别中文标签与
`is_account_fault` 判据 —— 与仓库既有的「service 不该依赖 job」分层一致。
三、按状态导出 CSV
- 新增 GET /api/login/tasks/export：不分页，导出当前筛选的**全部**匹配记录；
列：账号/状态/失败原因/失败分型/批次/最近尝试时间(UTC)/标签/归属客服
- 权限与筛选与列表**完全共用**，客服只导出自己名下账号
- 带 UTF-8 BOM（不加 Excel 打开中文表头是乱码），导出动作写审计
四、修一个口径不一致的真 bug
页面「失败 N」= failed + error 之和，但筛选只认 failed ——
所以截图里「失败 389」点进去只有 386 条（差的就是那几条 error）；
「清空失败」同样漏删 error。现在把状态组收敛成一份 `_apply_status_filter`，
列表 / 导出 / 批量删除三处共用。
测试：684 passed, 6 skipped（本次新增 13 个）；三条关键断言做了反向验证
（去掉 stopped 判断 / 失败组去掉 error / 去掉 BOM，对应用例均变红）。
590 lines
21 KiB
Python
590 lines
21 KiB
Python
"""登录任务路由 + WebSocket 实时日志"""
|
||
|
||
import asyncio
|
||
import csv
|
||
import io
|
||
import threading
|
||
from datetime import UTC, datetime
|
||
|
||
from fastapi import (
|
||
APIRouter,
|
||
Depends,
|
||
HTTPException,
|
||
WebSocket,
|
||
WebSocketDisconnect,
|
||
)
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlalchemy import and_, func, or_
|
||
from sqlalchemy.orm import Session, load_only
|
||
|
||
from ..database import SessionLocal, get_db
|
||
from ..deps import authenticate_websocket, get_current_user, require_permission
|
||
from ..models import Account, LoginTask, User
|
||
from ..models import ProxyConfig as ProxyConfigModel
|
||
from ..permissions import get_user_permissions, user_has_permission
|
||
from ..schemas import LoginBatchRequest, LoginTaskOut
|
||
from ..services.audit_service import record_audit
|
||
from ..services.login_failure import classify_failure, failure_kind_label
|
||
from ..services.login_service import LoginBatchRunner, batch_registry
|
||
|
||
router = APIRouter(prefix="/api/login", tags=["登录任务"])
|
||
|
||
|
||
def _effective_attempt_status():
|
||
"""页面展示用的状态：优先最近一次登录尝试，没有尝试时回退到 CK 状态。
|
||
|
||
``attempt_*`` 才是「这次登录怎么样了」，``status`` 只表示这条记录是否
|
||
持有可用 CK —— 两者分开，停止/失败/重启都不会让账号从列表里消失。
|
||
"""
|
||
return func.coalesce(
|
||
func.nullif(LoginTask.attempt_status, ""),
|
||
LoginTask.status,
|
||
)
|
||
|
||
|
||
#: 页面上「失败」是一个**组**：库里 `failed` 与 `error` 分两列存，但用户眼里都是失败。
|
||
#: 汇总口径早就是两者相加（见 `_login_task_summary`），筛选却只认 `failed` ——
|
||
#: 于是「失败 389」点进去只剩 386 条（差的就是那几条 `error`）。
|
||
#: 这里统一成组，列表 / 导出 / 批量删除共用一份，避免三处各写一遍再漂移。
|
||
STATUS_GROUPS: dict[str, tuple[str, ...]] = {
|
||
"failed": ("failed", "error"),
|
||
}
|
||
|
||
#: 导出 CSV 用的中文状态。与前端 `LoginTasksPage.STATUS_LABELS` 同源同义；
|
||
#: 后端要自己出一份是因为 CSV 在服务端生成，拿不到前端的表。
|
||
LOGIN_STATUS_LABELS: dict[str, str] = {
|
||
"pending": "等待中",
|
||
"running": "登录中",
|
||
"success": "成功",
|
||
"failed": "失败",
|
||
"error": "异常",
|
||
"stopped": "已停止",
|
||
"account_cancelled": "账号已注销",
|
||
"password_wrong": "账号密码错误",
|
||
"account_unverified": "账号未认证",
|
||
"account_verified": "账号已认证",
|
||
"account_auth_unknown": "认证状态未知",
|
||
}
|
||
|
||
|
||
def _status_values(status: str) -> tuple[str, ...]:
|
||
"""把页面上的状态名展开成库里真正的取值集合。"""
|
||
return STATUS_GROUPS.get(status, (status,))
|
||
|
||
|
||
def _apply_status_filter(query, status: str | None):
|
||
"""按「最近一次尝试的状态」筛选（与列表序列化口径一致）。
|
||
|
||
优先看 `attempt_status`，为空时回退到 `status` —— 两者分开是刻意的：
|
||
`status` 只表示这条记录是否持有可用 CK（见 `_effective_attempt_status`）。
|
||
"""
|
||
if not status:
|
||
return query
|
||
wanted = _status_values(status)
|
||
return query.filter(
|
||
or_(
|
||
and_(LoginTask.attempt_status != "", LoginTask.attempt_status.in_(wanted)),
|
||
and_(LoginTask.attempt_status == "", LoginTask.status.in_(wanted)),
|
||
)
|
||
)
|
||
|
||
|
||
def _login_task_summary(query):
|
||
"""按状态汇总登录任务，避免概览页拉完整任务列表。"""
|
||
rows = (
|
||
query.order_by(None)
|
||
.with_entities(_effective_attempt_status(), func.count(LoginTask.id))
|
||
.group_by(_effective_attempt_status())
|
||
.all()
|
||
)
|
||
status_counts = {status or "": count for status, count in rows}
|
||
return {
|
||
"total": sum(status_counts.values()),
|
||
"success": status_counts.get("success", 0),
|
||
"failed": sum(status_counts.get(status, 0) for status in ("failed", "error")),
|
||
"status_counts": status_counts,
|
||
}
|
||
|
||
|
||
@router.post("/batch")
|
||
async def create_batch(
|
||
req: LoginBatchRequest,
|
||
db: Session = Depends(get_db),
|
||
current: User = Depends(require_permission("login:batch")),
|
||
):
|
||
"""创建批量登录或账号检测任务。"""
|
||
if not req.account_ids:
|
||
raise HTTPException(status_code=400, detail="请选择账号")
|
||
|
||
# 读取代理配置
|
||
proxy = db.query(ProxyConfigModel).first()
|
||
|
||
# 权限过滤账号
|
||
valid_ids = []
|
||
seen_ids = set()
|
||
for aid in req.account_ids:
|
||
if aid in seen_ids:
|
||
continue
|
||
seen_ids.add(aid)
|
||
acc = db.query(Account).filter(Account.id == aid).first()
|
||
if not acc:
|
||
continue
|
||
if (
|
||
not user_has_permission(current, "login:view_all")
|
||
and acc.assigned_to != current.id
|
||
):
|
||
continue
|
||
valid_ids.append(aid)
|
||
|
||
if not valid_ids:
|
||
action_name = "检测" if req.mode == "check" else "登录"
|
||
raise HTTPException(status_code=403, detail=f"没有可{action_name}的账号")
|
||
|
||
# 在主事件循环中创建 log_queue，传给后台线程
|
||
log_queue = asyncio.Queue()
|
||
loop = asyncio.get_running_loop()
|
||
|
||
# 创建执行器（用独立的 DB 会话，因为在线程中运行）
|
||
thread_db = SessionLocal()
|
||
runner = LoginBatchRunner(
|
||
db=thread_db,
|
||
account_ids=valid_ids,
|
||
created_by=current.id,
|
||
creator_permissions=get_user_permissions(current),
|
||
max_login_retries=req.max_login_retries,
|
||
max_total_time=req.max_total_time,
|
||
proxy_config=proxy,
|
||
log_queue=log_queue,
|
||
loop=loop,
|
||
concurrency=req.concurrency,
|
||
api_strategy=req.api_strategy,
|
||
mode=req.mode,
|
||
)
|
||
|
||
batch_id = runner.batch_id
|
||
|
||
# 先注册到批次注册表，再启动线程，确保 WebSocket 连接时能找到
|
||
batch_registry.register(batch_id, log_queue, loop, runner)
|
||
|
||
# 启动线程
|
||
thread = threading.Thread(target=runner.run, daemon=True)
|
||
thread.start()
|
||
|
||
# 审计：批次是谁在什么时候发起、多少个账号、什么模式。
|
||
# 自动重登与人工登录共用同一条链路，没有这条记录就无法追溯来源。
|
||
try:
|
||
record_audit(
|
||
db,
|
||
current,
|
||
action="login:batch",
|
||
target=batch_id,
|
||
detail={
|
||
"mode": req.mode,
|
||
"requested": len(req.account_ids),
|
||
"accepted": len(valid_ids),
|
||
"concurrency": req.concurrency,
|
||
"api_strategy": req.api_strategy,
|
||
},
|
||
)
|
||
db.commit()
|
||
except Exception: # noqa: BLE001 审计失败不应影响已启动的批次
|
||
db.rollback()
|
||
|
||
return {"batch_id": batch_id, "count": len(valid_ids), "success": True}
|
||
|
||
|
||
def _visible_tasks_query(db: Session, current: User, batch_id: str | None):
|
||
"""按权限与批次过滤登录任务。客服只能看自己名下账号的任务。
|
||
|
||
这里已经**不需要 defer cookie** —— 登录记录不再持有凭证（2026-10-04 CK 解耦）。
|
||
"""
|
||
query = db.query(LoginTask).join(Account, LoginTask.account_id == Account.id)
|
||
|
||
if not user_has_permission(current, "login:view_all"):
|
||
query = query.filter(Account.assigned_to == current.id)
|
||
|
||
if batch_id:
|
||
# 记录持有的 CK 属于产出它的批次，最近一次尝试属于新批次，两者都要算。
|
||
query = query.filter(
|
||
or_(LoginTask.batch_id == batch_id, LoginTask.attempt_batch_id == batch_id)
|
||
)
|
||
return query
|
||
|
||
|
||
def _serialize_tasks(db: Session, rows) -> list[LoginTaskOut]:
|
||
"""把任务行拼上账号名后序列化。"""
|
||
account_ids = [t.account_id for t in rows]
|
||
accounts_map = {}
|
||
if account_ids:
|
||
accs = (
|
||
db.query(Account)
|
||
.options(load_only(Account.id, Account.username))
|
||
.filter(Account.id.in_(account_ids))
|
||
.all()
|
||
)
|
||
accounts_map = {a.id: a.username for a in accs}
|
||
|
||
result = []
|
||
for t in rows:
|
||
attempt_status = t.attempt_status or ""
|
||
result.append(
|
||
LoginTaskOut(
|
||
id=t.id,
|
||
batch_id=t.attempt_batch_id or t.batch_id,
|
||
account_id=t.account_id,
|
||
account_username=accounts_map.get(t.account_id, ""),
|
||
status=attempt_status or t.status,
|
||
cookie="",
|
||
message=(t.attempt_message if attempt_status else "")
|
||
or t.message
|
||
or "",
|
||
created_by=t.created_by,
|
||
created_at=t.created_at,
|
||
finished_at=t.attempt_at or t.finished_at,
|
||
)
|
||
)
|
||
return result
|
||
|
||
|
||
@router.get("/tasks")
|
||
def list_tasks(
|
||
batch_id: str | None = None,
|
||
# 默认值保持 None（而不是 Query(None)）：这个接口会被测试与内部代码直接
|
||
# 调用，用 Query() 做默认值时拿到的会是那个 Query 对象而非 None。
|
||
page: int | None = None,
|
||
page_size: int = 50,
|
||
status: str | None = None,
|
||
db: Session = Depends(get_db),
|
||
current: User = Depends(get_current_user),
|
||
):
|
||
"""查看登录任务列表。
|
||
|
||
**不传 `page` 时返回裸数组**（上限 200），沿用老调用方（Cookie 页等）的
|
||
契约；传了 `page` 才返回 ``{items, total, page, page_size}``，供表格分页。
|
||
|
||
传 `page` 时同时返回 `total`，是为了让前端能「全选所有筛选结果」——
|
||
否则跨页选择无从下手，而任务量早就过了单页能装下的规模。
|
||
"""
|
||
query = _apply_status_filter(_visible_tasks_query(db, current, batch_id), status)
|
||
|
||
ordered = query.order_by(LoginTask.id.desc())
|
||
|
||
if page is None:
|
||
return _serialize_tasks(db, ordered.limit(200).all())
|
||
|
||
page = max(1, int(page))
|
||
page_size = max(1, min(int(page_size), 500))
|
||
total = ordered.count()
|
||
rows = ordered.offset((page - 1) * page_size).limit(page_size).all()
|
||
return {
|
||
"items": _serialize_tasks(db, rows),
|
||
"total": total,
|
||
"page": page,
|
||
"page_size": page_size,
|
||
}
|
||
|
||
|
||
def _fmt_dt(value) -> str:
|
||
"""导出用的时间格式。
|
||
|
||
统一输出 **UTC**，并在表头写明 —— 服务端猜本地时区迟早会错，
|
||
而错的时间在人工筛选时是看不出来的。
|
||
"""
|
||
if value is None:
|
||
return ""
|
||
return value.strftime("%Y-%m-%d %H:%M:%S")
|
||
|
||
|
||
@router.get("/tasks/export")
|
||
def export_tasks(
|
||
batch_id: str | None = None,
|
||
status: str | None = None,
|
||
db: Session = Depends(get_db),
|
||
current: User = Depends(get_current_user),
|
||
):
|
||
"""按当前筛选导出登录任务 CSV（账号名 / 状态 / 失败原因 / 分型 / 批次 / 时间 / 标签 / 归属客服）。
|
||
|
||
几个刻意的选择：
|
||
|
||
- **不分页**：导出的用途就是「拿去 Excel 人工筛」，只给一页没有意义；
|
||
- **筛选与权限跟列表完全共用**（`_visible_tasks_query` + `_apply_status_filter`），
|
||
所以「客服只能导出自己名下账号」这条不需要在这里重写一遍；
|
||
- **带 UTF-8 BOM**：不加的话 Excel 打开中文表头是乱码；
|
||
- 失败分型复用登录失败那套（`login_failure.classify_failure`），
|
||
与账号时间线上写的是同一份判据，不会出现「页面上算 A、导出算 B」。
|
||
"""
|
||
query = _apply_status_filter(_visible_tasks_query(db, current, batch_id), status)
|
||
rows = query.order_by(LoginTask.id.desc()).all()
|
||
|
||
account_ids = [t.account_id for t in rows]
|
||
info: dict[int, tuple[str, str, str]] = {}
|
||
if account_ids:
|
||
for acc_id, username, tag, owner in (
|
||
db.query(Account.id, Account.username, Account.tag, User.username)
|
||
.outerjoin(User, Account.assigned_to == User.id)
|
||
.filter(Account.id.in_(account_ids))
|
||
.all()
|
||
):
|
||
info[acc_id] = (username or "", tag or "", owner or "")
|
||
|
||
output = io.StringIO()
|
||
writer = csv.writer(output)
|
||
writer.writerow(
|
||
[
|
||
"账号",
|
||
"状态",
|
||
"失败原因",
|
||
"失败分型",
|
||
"批次",
|
||
"最近尝试时间(UTC)",
|
||
"标签",
|
||
"归属客服",
|
||
]
|
||
)
|
||
for task in rows:
|
||
attempt_status = task.attempt_status or ""
|
||
effective = attempt_status or task.status or ""
|
||
message = (task.attempt_message if attempt_status else "") or task.message or ""
|
||
kind = classify_failure(message) if effective in {"failed", "error"} else ""
|
||
username, tag, owner = info.get(task.account_id, ("", "", ""))
|
||
writer.writerow(
|
||
[
|
||
username,
|
||
LOGIN_STATUS_LABELS.get(effective, effective),
|
||
message,
|
||
failure_kind_label(kind) if kind else "",
|
||
task.attempt_batch_id or task.batch_id or "",
|
||
_fmt_dt(task.attempt_at or task.finished_at or task.created_at),
|
||
tag,
|
||
owner,
|
||
]
|
||
)
|
||
|
||
content = "\ufeff" + output.getvalue()
|
||
stamp = datetime.now(UTC).strftime("%Y%m%d_%H%M")
|
||
filename = f"login_tasks_{status or 'all'}_{stamp}.csv"
|
||
record_audit(
|
||
db,
|
||
current,
|
||
action="login:tasks:export",
|
||
target=f"status={status or 'all'}",
|
||
detail={"rows": len(rows), "batch_id": batch_id or ""},
|
||
)
|
||
db.commit()
|
||
return StreamingResponse(
|
||
iter([content]),
|
||
media_type="text/csv; charset=utf-8",
|
||
headers={"Content-Disposition": f"attachment; filename={filename}"},
|
||
)
|
||
|
||
|
||
@router.get("/tasks/summary")
|
||
def tasks_summary(
|
||
db: Session = Depends(get_db),
|
||
current: User = Depends(get_current_user),
|
||
):
|
||
"""登录任务统计。"""
|
||
query = db.query(LoginTask).join(Account, LoginTask.account_id == Account.id)
|
||
if not user_has_permission(current, "login:view_all"):
|
||
query = query.filter(Account.assigned_to == current.id)
|
||
return _login_task_summary(query)
|
||
|
||
|
||
def _in_flight_clause():
|
||
"""不可删除的记录：正被批次处理的**在途尝试**。
|
||
|
||
⚠️ **CK 自 2026-10-04 起不再是护栏理由** —— CK 已迁到 `accounts.cookie`，
|
||
删登录历史不再影响任何账号的登录态（这正是解耦的目的）。但「在途」必须保留：
|
||
删掉会让批次线程找不到记录；停批次请走 ``/stop/{batch_id}``。
|
||
"""
|
||
in_flight = ("pending", "running")
|
||
return or_(
|
||
and_(
|
||
LoginTask.attempt_status != "",
|
||
LoginTask.attempt_status.in_(in_flight),
|
||
),
|
||
and_(LoginTask.attempt_status == "", LoginTask.status.in_(in_flight)),
|
||
)
|
||
|
||
|
||
def _deletable_clause():
|
||
"""可安全删除的记录：不在途。"""
|
||
return ~_in_flight_clause()
|
||
|
||
|
||
def _is_in_flight(task: LoginTask) -> bool:
|
||
"""单条记录是否在途（与 `_in_flight_clause` 同口径）。"""
|
||
in_flight = ("pending", "running")
|
||
if task.attempt_status:
|
||
return task.attempt_status in in_flight
|
||
return (task.status or "") in in_flight
|
||
|
||
|
||
@router.delete("/tasks/{task_id}")
|
||
def delete_task(
|
||
task_id: int,
|
||
db: Session = Depends(get_db),
|
||
current: User = Depends(require_permission("login:batch")),
|
||
):
|
||
"""删除单个登录任务（**在途的记录不允许删除**；删历史不影响 CK）。"""
|
||
task = db.query(LoginTask).filter(LoginTask.id == task_id).first()
|
||
if not task:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
if _is_in_flight(task):
|
||
raise HTTPException(
|
||
status_code=409,
|
||
detail="该记录正在被批次处理，删除会让批次失去目标；请先停止批次",
|
||
)
|
||
db.delete(task)
|
||
db.commit()
|
||
return {"message": "已删除", "success": True}
|
||
|
||
|
||
@router.delete("/tasks")
|
||
def delete_tasks(
|
||
task_ids: str = "",
|
||
batch_id: str | None = None,
|
||
status: str | None = None,
|
||
all_matching: bool = False,
|
||
db: Session = Depends(get_db),
|
||
current: User = Depends(require_permission("login:batch")),
|
||
):
|
||
"""批量删除登录任务。三种模式，互斥：
|
||
|
||
1. ``task_ids``：删指定 id（表格勾选）；
|
||
2. ``all_matching=true``：删当前筛选条件下的**全部**（跨页全选用）；
|
||
3. 两者都没给 → 400，绝不默认全删。
|
||
|
||
两条护栏（都会在响应里报告跳过数量）：
|
||
|
||
- **在途（pending/running）的记录不删**：删掉会让批次线程找不到记录；
|
||
停批次请走 ``/stop/{batch_id}``。
|
||
- ~~持有 CK 的记录不删~~ —— **这条护栏已随 CK 解耦（2026-10-04）取消**：
|
||
CK 现在挂在账号上，删登录历史不影响登录态。历史终于可以自由清理。
|
||
"""
|
||
# 护栏：不删在途任务，避免批次线程与数据库状态打架。
|
||
not_in_flight = _deletable_clause()
|
||
|
||
if task_ids:
|
||
ids = [int(x) for x in task_ids.split(",") if x.strip().isdigit()]
|
||
if not ids:
|
||
raise HTTPException(status_code=400, detail="无效的任务ID")
|
||
elif all_matching:
|
||
# 带 join 的 query 不能直接 bulk delete（SQLAlchemy 会拒绝），
|
||
# 所以先按权限与筛选条件查出 id，再按主键删。
|
||
# 状态筛选与列表/导出共用一份 —— 「清空失败」以前只删 `failed`、
|
||
# 漏掉 `error`，删完与页面上那个「失败 N」的数还对不上。
|
||
query = _apply_status_filter(
|
||
_visible_tasks_query(db, current, batch_id), status
|
||
)
|
||
ids = [row[0] for row in query.with_entities(LoginTask.id).all()]
|
||
else:
|
||
raise HTTPException(status_code=400, detail="请指定任务ID或 all_matching=true")
|
||
|
||
if not ids:
|
||
return {
|
||
"message": "没有符合条件的任务",
|
||
"deleted": 0,
|
||
"skipped": 0,
|
||
"success": True,
|
||
}
|
||
|
||
rows = (
|
||
db.query(LoginTask.id, not_in_flight.label("ok_to_delete"))
|
||
.filter(LoginTask.id.in_(ids))
|
||
.all()
|
||
)
|
||
target_ids = [r[0] for r in rows if r[1]]
|
||
skipped = len(rows) - len(target_ids)
|
||
|
||
if not target_ids:
|
||
return {
|
||
"message": f"没有可删除的记录（{skipped} 条在途，已跳过）",
|
||
"deleted": 0,
|
||
"skipped": skipped,
|
||
"success": True,
|
||
}
|
||
|
||
deleted = (
|
||
db.query(LoginTask)
|
||
.filter(LoginTask.id.in_(target_ids))
|
||
.delete(synchronize_session=False)
|
||
)
|
||
db.commit()
|
||
|
||
message = f"已删除 {deleted} 个任务"
|
||
if skipped:
|
||
message += f"（{skipped} 条在途，已跳过）"
|
||
return {"message": message, "deleted": deleted, "skipped": skipped, "success": True}
|
||
|
||
|
||
@router.post("/stop/{batch_id}")
|
||
def stop_batch(
|
||
batch_id: str,
|
||
current: User = Depends(require_permission("login:batch")),
|
||
):
|
||
batch = batch_registry.get(batch_id)
|
||
if batch:
|
||
if batch.get("finished"):
|
||
# 已结束但 WS 尚未清理（或从未有 WS 订阅）：回收注册表条目。
|
||
batch_registry.pop(batch_id)
|
||
raise HTTPException(status_code=404, detail="批次已结束")
|
||
batch["runner"].stop()
|
||
return {"message": "已发送停止信号", "success": True}
|
||
raise HTTPException(status_code=404, detail="批次不存在或已结束")
|
||
|
||
|
||
@router.websocket("/ws/login/{batch_id}")
|
||
async def ws_login_logs(websocket: WebSocket, batch_id: str):
|
||
"""WebSocket 推送登录实时日志（需认证）。"""
|
||
# 认证：从 cookie 或 token query param 验证用户身份
|
||
user = authenticate_websocket(websocket)
|
||
if not user:
|
||
await websocket.close(code=1008, reason="未授权")
|
||
return
|
||
|
||
# 归属校验：非 view_all 只能订阅自己创建的批次（与虎牙批次一致）。
|
||
if not user_has_permission(user, "login:view_all"):
|
||
check_db = SessionLocal()
|
||
try:
|
||
owned = (
|
||
check_db.query(LoginTask.id)
|
||
.filter(LoginTask.batch_id == batch_id, LoginTask.created_by == user.id)
|
||
.first()
|
||
)
|
||
finally:
|
||
check_db.close()
|
||
if not owned:
|
||
await websocket.close(code=1008, reason="无权访问该任务批次")
|
||
return
|
||
|
||
await websocket.accept()
|
||
|
||
# 从已注册的批次中获取 log_queue（由 create_batch 创建）
|
||
batch = 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)
|
||
# 收到 result 表示任务结束
|
||
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 = batch_registry.get(batch_id)
|
||
if latest and latest.get("finished"):
|
||
batch_registry.pop(batch_id)
|