boss_dp/common/protocol.py

# -*- coding: utf-8 -*-
"""
共享消息协议定义。
服务器与 Worker 之间通过 WebSocket 传递 JSON 消息，每条消息包含 type 字段标识消息类型。
"""
from enum import Enum


# ────────────────────────── 消息类型 ──────────────────────────

class MsgType(str, Enum):
    """WebSocket 消息类型枚举（str 混入方便 JSON 序列化）。"""

    # Worker → Server
    REGISTER = "register"                    # 注册：上报 worker 信息与浏览器列表
    HEARTBEAT = "heartbeat"                  # 心跳
    BROWSER_LIST_UPDATE = "browser_list_update"  # 浏览器列表变更
    TASK_PROGRESS = "task_progress"          # 任务进度上报
    TASK_RESULT = "task_result"              # 任务最终结果
    TASK_STATUS_REPORT = "task_status_report"  # 任务执行状态回报（响应服务端查询）

    # Server → Worker
    REGISTER_ACK = "register_ack"            # 注册确认
    HEARTBEAT_ACK = "heartbeat_ack"          # 心跳确认
    TASK_ASSIGN = "task_assign"              # 派发任务
    TASK_CANCEL = "task_cancel"              # 取消任务
    TASK_STATUS_QUERY = "task_status_query"  # 查询任务执行状态

    # 双向
    ERROR = "error"                          # 错误消息


# ────────────────────────── 任务状态 ──────────────────────────

class TaskStatus(str, Enum):
    """任务生命周期状态。"""
    PENDING = "pending"          # 已创建，等待派发
    DISPATCHED = "dispatched"    # 已派发给 Worker
    RUNNING = "running"          # Worker 正在执行
    SUCCESS = "success"          # 执行成功
    FAILED = "failed"            # 执行失败
    CANCELLED = "cancelled"      # 已取消


# ────────────────────────── 任务类型 ──────────────────────────

class TaskType(str, Enum):
    """可扩展的任务类型。新增任务在此追加即可。"""
    BOSS_RECRUIT = "boss_recruit"     # BOSS 直聘招聘流程
    CHECK_LOGIN = "check_login"       # 检测 BOSS 账号是否已登录


# ────────────────────────── 辅助函数 ──────────────────────────

def make_msg(msg_type: MsgType, **payload) -> dict:
    """构造一条标准 WebSocket JSON 消息。"""
    return {"type": msg_type.value, **payload}