File size: 3,101 Bytes
3216f7c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """
Definición de protocolo WebSocket y serialización de mensajes para comunicación con el Manager.
100% idéntico a internal/protocol/messages.go.
"""
import json
import base64
from typing import Any, Tuple, Dict, List, Optional
from dataclasses import dataclass, asdict
from config import logger
# Tipos de mensajes
TYPE_HANDSHAKE = "handshake"
TYPE_HANDSHAKE_ACK = "handshake_ack"
TYPE_PING = "ping"
TYPE_PONG = "pong"
TYPE_TASK_ASSIGN = "task_assign"
TYPE_TASK_PROGRESS = "task_progress"
TYPE_TASK_COMPLETE = "task_complete"
TYPE_TASK_FAILED = "task_failed"
TYPE_TASK_CANCEL = "task_cancel"
@dataclass
class VideoSettings:
resolution: str = ""
crf: str = ""
audio_bitrate: str = ""
fps: str = ""
preset: str = ""
codec: str = ""
watermark: str = ""
watermark_pos: str = ""
watermark_image_file_id: str = ""
format: str = ""
thumbnail_file_id: str = ""
@dataclass
class TaskState:
task_id: str
status: str # "converting", "completed"
upload_url: str = ""
@dataclass
class Handshake:
worker_id: str
secret: str
cpu_cores: int
max_jobs: int
active_tasks: List[Dict[str, Any]]
@dataclass
class TaskProgress:
task_id: str
percent: float
speed: str
time: str
eta: str
@dataclass
class TaskComplete:
task_id: str
original_size: int
compressed_size: int
duration_sec: float
@dataclass
class TaskFailed:
task_id: str
error: str
@dataclass
class TaskCancel:
task_id: str
def build_ws_packet(msg_type: str, payload_obj: Any = None) -> str:
"""
Empaqueta el mensaje en formato WebSocket esperado por el Manager en Go:
type WSMessage struct {
Type MessageType `json:"type"`
Payload []byte `json:"payload,omitempty"`
}
En Go, un campo `[]byte` se codifica/decodifica como una cadena Base64.
"""
if payload_obj is not None:
if hasattr(payload_obj, "__dataclass_fields__"):
payload_data = asdict(payload_obj)
elif isinstance(payload_obj, dict):
payload_data = payload_obj
else:
payload_data = payload_obj
raw_json = json.dumps(payload_data).encode("utf-8")
b64_payload = base64.b64encode(raw_json).decode("ascii")
return json.dumps({"type": msg_type, "payload": b64_payload})
return json.dumps({"type": msg_type})
def parse_ws_packet(raw_text: str) -> Tuple[str, Any]:
"""
Decodifica el mensaje JSON entrante soportando tanto payload en Base64
como payload embebido directamente como diccionario/objeto.
"""
try:
data = json.loads(raw_text)
msg_type = data.get("type", "")
payload = data.get("payload")
if isinstance(payload, str):
try:
decoded_bytes = base64.b64decode(payload)
payload = json.loads(decoded_bytes.decode("utf-8"))
except Exception:
pass
return msg_type, payload
except Exception as e:
logger.error(f"Error decodificando paquete WS: {e}")
return "", None
|