# -*- coding: utf-8 -*- """ 任务提交与查询 API（需要登录）。 统一任务入口：前端通过 task_type 指定任务类型（如 check_login、boss_recruit）。 """ import json import logging from datetime import datetime from typing import Optional from asgiref.sync import async_to_sync from rest_framework import status as http_status from rest_framework.decorators import api_view from common.protocol import MsgType, TaskStatus, TaskType, make_msg from server.core.response import api_success, api_error from server.models import BossAccount, TaskCreate, TaskLog, Task from server.serializers import TaskCreateSerializer from server.core.worker_manager import worker_manager from server.core.task_dispatcher import task_dispatcher logger = logging.getLogger("server.api.tasks") def _format_timestamp(ts: float) -> str: """将时间戳转换为指定格式的字符串：2026-03-01T20:19:51""" dt = datetime.fromtimestamp(ts) return dt.strftime("%Y-%m-%dT%H:%M:%S") def _task_to_dict(t) -> dict: """将任务实例（Task ORM）转为可序列化字典。""" return { "task_id": t.task_id, "task_type": str(t.task_type), "status": str(t.status), "worker_id": t.worker_id, "account_name": getattr(t, "account_name", None), "params": t.params or {}, "progress": getattr(t, "progress", None), "result": t.result, "error": t.error, "created_at": ( t.created_at.strftime("%Y-%m-%dT%H:%M:%S") if hasattr(t.created_at, "strftime") else _format_timestamp(t.created_at) ), "updated_at": ( t.updated_at.strftime("%Y-%m-%dT%H:%M:%S") if hasattr(t.updated_at, "strftime") else _format_timestamp(t.updated_at) ), } def _parse_limit(raw_value, default: int = 50, max_value: int = 200) -> int: """解析分页上限，避免非法参数导致 500。""" try: value = int(raw_value) except (TypeError, ValueError): return default if value <= 0: return default return min(value, max_value) def _parse_positive_int(raw_value, default: int = 1) -> int: """解析正整数参数。""" try: value = int(raw_value) except (TypeError, ValueError): return default if value <= 0: return default return value def _parse_task_status(raw_status: Optional[str]) -> Optional[TaskStatus]: """解析任务状态。为空时返回 None。""" if not raw_status: return None try: return TaskStatus(raw_status) except ValueError: return None def _task_log_account_name(task_log: TaskLog) -> Optional[str]: """从任务日志里提取可识别的环境名。""" result = task_log.result if isinstance(task_log.result, dict) else {} params = task_log.params if isinstance(task_log.params, dict) else {} account_name = ( result.get("browser_name") or result.get("account_name") or params.get("account_name") or params.get("browser_name") ) if not account_name: return None return str(account_name).strip() or None def _task_log_to_dict(task_log: TaskLog, account_name: Optional[str] = None) -> dict: """将 TaskLog 转为统一响应结构。""" return { "task_id": task_log.task_id, "task_type": task_log.task_type, "status": task_log.status, "worker_id": task_log.worker_id, "account_name": account_name or _task_log_account_name(task_log), "params": task_log.params if isinstance(task_log.params, dict) else {}, "progress": None, "result": task_log.result, "error": task_log.error, "created_at": task_log.created_at.strftime("%Y-%m-%dT%H:%M:%S"), "updated_at": task_log.created_at.strftime("%Y-%m-%dT%H:%M:%S"), } def _is_task_log_for_account(task_log: TaskLog, account: BossAccount) -> bool: """判断任务日志是否属于某个账号（用于账号任务列表兼容）。""" if task_log.worker_id and task_log.worker_id != account.worker_id: return False if account.current_task_id and task_log.task_id == account.current_task_id: return True matched_name = _task_log_account_name(task_log) if matched_name and matched_name == account.browser_name: return True params = task_log.params if isinstance(task_log.params, dict) else {} account_pk = str(account.pk) for key in ("id", "account_id", "boss_id"): value = params.get(key) if value is None: continue if str(value).strip() == account_pk: return True return False def _list_tasks_by_account(account: BossAccount, task_status: Optional[TaskStatus], limit: Optional[int] = 50) -> list: """ 按账号维度查询任务列表，完全基于 Task 表。 - 不再依赖内存中的 TaskInfo； - 你可以直接修改数据库中 Task.status/Task.result 等字段来影响这里的返回。 """ qs = Task.objects.filter( worker_id=account.worker_id, account_name=account.browser_name, ).order_by("-created_at") if task_status: qs = qs.filter(status=task_status.value if hasattr(task_status, "value") else str(task_status)) if limit is not None: qs = qs[:limit] return [_task_to_dict(t) for t in qs] @api_view(["GET", "POST"]) def task_list(request): """ GET -> 查询任务列表，支持以下过滤参数： - ?worker_id= : 按 Worker ID 过滤 - ?account_id= : 按账号 ID 过滤（自动查询对应的 account_name） - ?status= : 按任务状态过滤 - ?limit= : 返回条数上限，默认 50 POST -> 提交新任务（支持 JSON 和 form-data） """ if request.method == "GET": wid = request.query_params.get("worker_id") account_id = request.query_params.get("account_id") st = request.query_params.get("status") limit = _parse_limit(request.query_params.get("limit", 50)) task_status = _parse_task_status(st) if st and task_status is None: return api_error(http_status.HTTP_400_BAD_REQUEST, f"不支持的任务状态: {st}") # 如果指定了 account_id，直接返回该账号的任务列表（兼容前端“查看任务”） if account_id: try: account_id = int(account_id) account = BossAccount.objects.get(pk=account_id) except (ValueError, BossAccount.DoesNotExist): return api_error( http_status.HTTP_404_NOT_FOUND, f"未找到 id={account_id} 的账号", ) return api_success(_list_tasks_by_account(account, task_status=task_status, limit=limit)) qs = Task.objects.all().order_by("-created_at") if wid: qs = qs.filter(worker_id=wid) if task_status: qs = qs.filter(status=task_status.value if hasattr(task_status, "value") else str(task_status)) qs = qs[:limit] return api_success([_task_to_dict(t) for t in qs]) # POST: 提交新任务 data = request.data.copy() # form-data 中 params 可能是 JSON 字符串 params_raw = data.get("params", {}) if isinstance(params_raw, str): try: data["params"] = json.loads(params_raw) if params_raw.strip() else {} except (json.JSONDecodeError, ValueError): data["params"] = {} ser = TaskCreateSerializer(data=data) if not ser.is_valid(): logger.warning("任务提交参数校验失败: data=%s, errors=%s", data, ser.errors) from rest_framework.exceptions import ValidationError raise ValidationError(ser.errors) validated = ser.validated_data.copy() # boss_id 即 /api/accounts 返回的 id（账号主键），与 id/account_id 等价，均按主键查表 account_id = ( validated.pop("id", None) or validated.pop("account_id", None) or validated.pop("boss_id", None) ) if account_id is not None: try: account_id = int(account_id) except (TypeError, ValueError): account_id = None req = TaskCreate(**validated) target_worker_id = req.worker_id or "" account_name = req.account_name or "" # 指定 id / account_id / boss_id（均为账号主键）时，按主键查表 if account_id: try: account = BossAccount.objects.get(pk=account_id) except BossAccount.DoesNotExist: return api_error( http_status.HTTP_404_NOT_FOUND, f"未找到 id={account_id} 的账号", ) target_worker_id = account.worker_id account_name = account.browser_name req.worker_id = target_worker_id req.account_name = account_name if not target_worker_id and account_name: target_worker_id = worker_manager.find_worker_by_account(account_name) if not target_worker_id: return api_error( http_status.HTTP_404_NOT_FOUND, f"未找到拥有浏览器 '{account_name}' 的在线 Worker", ) req.worker_id = target_worker_id if not target_worker_id: return api_error(http_status.HTTP_400_BAD_REQUEST, "请指定 id、boss_id、worker_id 或 account_name") if not worker_manager.is_online(target_worker_id): return api_error(http_status.HTTP_503_SERVICE_UNAVAILABLE, f"Worker {target_worker_id} 不在线") req.worker_id = target_worker_id try: task = task_dispatcher.create_task(req) except ValueError as e: return api_error(http_status.HTTP_409_CONFLICT, str(e)) send_fn = worker_manager.get_send_fn(target_worker_id) if not send_fn: task.status = TaskStatus.FAILED task.error = "Worker WebSocket 连接不存在" return api_error(http_status.HTTP_503_SERVICE_UNAVAILABLE, "Worker WebSocket 连接不存在") success = async_to_sync(task_dispatcher.dispatch)(task, send_fn) if not success: return api_error(http_status.HTTP_503_SERVICE_UNAVAILABLE, f"任务派发失败: {task.error}") worker_manager.set_current_task(target_worker_id, task.task_id) # check_login 任务：关联账号的任务状态 if req.task_type == TaskType.CHECK_LOGIN and req.account_name: try: account = BossAccount.objects.filter( browser_name=req.account_name, worker_id=target_worker_id, ).first() if account: account.current_task_id = task.task_id # 派发成功后，将状态设为 DISPATCHED（不是 task.status.value，因为那可能还是 PENDING） account.current_task_status = TaskStatus.DISPATCHED.value account.save(update_fields=["current_task_id", "current_task_status"]) logger.info("账号 %s 任务关联: task_id=%s, status=%s", req.account_name, task.task_id, account.current_task_status) except Exception as e: logger.warning("关联账号任务状态失败: %s", e) return api_success(_task_to_dict(task), http_status=http_status.HTTP_201_CREATED) @api_view(["GET"]) def task_detail(request, task_id: str): """ 按 task_id 查询单个任务详情。 优先查 Task 表（活跃任务），再查 TaskLog 表（历史记录）。 """ task = Task.objects.filter(task_id=task_id).first() if task: return api_success(_task_to_dict(task)) task_log = TaskLog.objects.filter(task_id=task_id).first() if task_log: return api_success(_task_log_to_dict(task_log)) return api_error(http_status.HTTP_404_NOT_FOUND, f"任务 {task_id} 不存在") @api_view(["POST"]) def task_cancel(request, task_id: str): """ 取消任务。 - 若任务已结束（success/failed/cancelled）返回 409； - 若任务可取消：先写入 cancelled，再向 Worker 下发 TASK_CANCEL。 """ task = Task.objects.filter(task_id=task_id).first() if not task: return api_error(http_status.HTTP_404_NOT_FOUND, f"任务 {task_id} 不存在") active_status_values = { TaskStatus.PENDING.value, TaskStatus.DISPATCHED.value, TaskStatus.RUNNING.value, } if str(task.status) not in active_status_values: return api_error(http_status.HTTP_409_CONFLICT, f"任务当前状态为 {task.status}，不可取消") cancelled_task = task_dispatcher.cancel_task(task_id, error="任务已取消") if not cancelled_task: return api_error(http_status.HTTP_409_CONFLICT, "任务已结束或已被取消") if cancelled_task.worker_id: worker_manager.set_current_task(cancelled_task.worker_id, None) if cancelled_task.account_name: BossAccount.objects.filter( worker_id=cancelled_task.worker_id, browser_name=cancelled_task.account_name, current_task_id=cancelled_task.task_id, ).update(current_task_status=TaskStatus.CANCELLED.value) send_fn = worker_manager.get_send_fn(cancelled_task.worker_id) if send_fn: cancel_msg = make_msg(MsgType.TASK_CANCEL, task_id=cancelled_task.task_id) try: async_to_sync(send_fn)(cancel_msg) except Exception as e: logger.warning("向 Worker 下发任务取消失败 task_id=%s: %s", cancelled_task.task_id, e) return api_success(_task_to_dict(cancelled_task), msg="任务已取消") @api_view(["GET"]) def task_list_by_account(request, account_id: int): """ 按账号 ID 查询任务列表（不支持按 task_id 查询）。 - 兼容模式：不传 page/page_size 时，返回数组 - 分页模式：传 page 或 page_size 时，返回 {total, page, page_size, results} """ account = BossAccount.objects.filter(pk=account_id).first() if not account: return api_error(http_status.HTTP_404_NOT_FOUND, f"账号 {account_id} 不存在") st = request.query_params.get("status") task_status = _parse_task_status(st) if st and task_status is None: return api_error(http_status.HTTP_400_BAD_REQUEST, f"不支持的任务状态: {st}") # 旧接口兼容：默认返回数组，limit 控制条数 enable_paging = ("page" in request.query_params) or ("page_size" in request.query_params) if not enable_paging: limit = _parse_limit(request.query_params.get("limit", 50)) return api_success(_list_tasks_by_account(account, task_status=task_status, limit=limit)) page = _parse_positive_int(request.query_params.get("page"), default=1) page_size = _parse_limit( request.query_params.get("page_size", request.query_params.get("limit", 20)), default=20, max_value=200, ) all_items = _list_tasks_by_account(account, task_status=task_status, limit=None) total = len(all_items) start = (page - 1) * page_size end = start + page_size return api_success({ "total": total, "page": page, "page_size": page_size, "results": all_items[start:end], })