THANOS / protocol.py
st3v3x's picture
Upload 11 files
794cf9d verified
Raw History Blame Contribute Delete
3.1 kB
"""
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