from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from app.database import get_db from app.models.process import Worker from app.models.user import AdminUser, User from app.routers.auth import get_current_user from app.routers.admin import get_current_admin from app.utils.notifications import create_admin_notification from pydantic import BaseModel from typing import List, Optional from datetime import datetime import uuid from app.websocket import ws_manager router = APIRouter(prefix="/workers", tags=["Workers"]) class WorkerRegister(BaseModel): hostname: str name: Optional[str] = None status: Optional[str] = "online" class WorkerResponse(BaseModel): id: int hostname: str name: Optional[str] status: str last_seen: datetime class Config: orm_mode = True class WorkerUpdate(BaseModel): name: Optional[str] = None status: Optional[str] = None @router.post("/register", response_model=WorkerResponse) def register_worker(data: WorkerRegister, db: Session = Depends(get_db)): worker = db.query(Worker).filter(Worker.hostname == data.hostname).first() is_new = worker is None previous_status = None if is_new else worker.status if is_new: worker = Worker( hostname=data.hostname, name=data.name or data.hostname, status=data.status, last_seen=datetime.utcnow() ) db.add(worker) else: worker.status = data.status worker.last_seen = datetime.utcnow() if data.name: worker.name = data.name db.commit() db.refresh(worker) if is_new: create_admin_notification( db, "Worker Registrado", f"El equipo worker '{worker.hostname}' se ha registrado en la red." ) elif previous_status == "offline" and worker.status in ["online", "busy"]: create_admin_notification( db, "Worker Conectado", f"El equipo worker '{worker.name or worker.hostname}' se ha conectado (estado: {worker.status})." ) ws_manager.broadcast_sync({"type": "refresh", "tab": "workers"}) return worker @router.get("", response_model=List[WorkerResponse]) def get_workers(db: Session = Depends(get_db), current_admin: AdminUser = Depends(get_current_admin)): if not current_admin.perm_workers: raise HTTPException(status_code=403, detail="No tiene permisos para gestionar workers") # Marcar automáticamente como 'offline' los que no han reportado en 2 minutos workers = db.query(Worker).order_by(Worker.id).all() now = datetime.utcnow() updated = False for w in workers: if w.status in ["online", "busy"] and (now - w.last_seen).total_seconds() > 120: w.status = "offline" updated = True if updated: db.commit() workers = db.query(Worker).order_by(Worker.id).all() return workers @router.put("/{worker_id}", response_model=WorkerResponse) def update_worker(worker_id: int, data: WorkerUpdate, db: Session = Depends(get_db), current_admin: AdminUser = Depends(get_current_admin)): if not current_admin.perm_workers: raise HTTPException(status_code=403, detail="No tiene permisos para gestionar workers") worker = db.query(Worker).filter(Worker.id == worker_id).first() if not worker: raise HTTPException(status_code=404, detail="Worker no encontrado") if data.name is not None: worker.name = data.name if data.status is not None: worker.status = data.status if data.status == "suspended": worker.is_active = False worker.command = "stop" try: from tasks import celery_app celery_app.control.cancel_consumer('celery', destination=[f'celery@{worker.hostname}']) except Exception as e: print(f"[Worker] Error cancelando consumidor: {e}") elif data.status == "online": worker.is_active = True worker.command = "start" try: from tasks import celery_app celery_app.control.add_consumer('celery', destination=[f'celery@{worker.hostname}']) except Exception as e: print(f"[Worker] Error agregando consumidor: {e}") db.commit() db.refresh(worker) # Registrar en notificaciones de admin action_desc = f"actualizado (estado: {worker.status})" if data.status == "suspended": action_desc = "suspendido" elif data.status == "online": action_desc = "activado" create_admin_notification( db, title=f"Worker {action_desc}", message=f"El Worker '{worker.name}' (ID: {worker.id}, Host: {worker.hostname}) ha sido {action_desc} por el administrador '{current_admin.username}'.", category="worker" ) ws_manager.broadcast_sync({"type": "refresh", "tab": "workers"}) return worker @router.delete("/{worker_id}") def delete_worker(worker_id: int, db: Session = Depends(get_db), current_admin: AdminUser = Depends(get_current_admin)): if not current_admin.perm_workers: raise HTTPException(status_code=403, detail="No tiene permisos para gestionar workers") worker = db.query(Worker).filter(Worker.id == worker_id).first() if not worker: raise HTTPException(status_code=404, detail="Worker no encontrado") db.delete(worker) db.commit() create_admin_notification( db, title="Worker eliminado", message=f"El Worker '{worker.name}' (ID: {worker_id}, Host: {worker.hostname}) ha sido eliminado permanentemente por el administrador '{current_admin.username}'.", category="worker" ) ws_manager.broadcast_sync({"type": "refresh", "tab": "workers"}) return {"message": "Worker eliminado exitosamente"} # ========================================== # BYOK / LOCAL WORKERS ENDPOINTS # ========================================== @router.get("/my-worker") def get_my_worker(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)): worker = db.query(Worker).filter(Worker.user_id == current_user.id).first() if not worker: # Generar token unico de worker y crearlo token = f"melodix_wk_{uuid.uuid4().hex}" worker = Worker( hostname="unregistered", name=f"Worker de {current_user.username}", status="offline", last_seen=datetime.utcnow(), user_id=current_user.id, token=token, command="none", is_active=False ) db.add(worker) db.commit() db.refresh(worker) return { "id": worker.id, "name": worker.name, "status": worker.status, "last_seen": worker.last_seen, "token": worker.token, "is_active": worker.is_active, "command": worker.command, "hostname": worker.hostname } @router.post("/my-worker/toggle") def toggle_my_worker(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)): worker = db.query(Worker).filter(Worker.user_id == current_user.id).first() if not worker: raise HTTPException(status_code=404, detail="Worker no encontrado") worker.is_active = not worker.is_active db.commit() db.refresh(worker) return {"is_active": worker.is_active, "status": worker.status} @router.post("/my-worker/command") def send_worker_command(command: str, db: Session = Depends(get_db), current_user: User = Depends(get_current_user)): if command not in ["start", "stop"]: raise HTTPException(status_code=400, detail="Comando invalido. Use 'start' o 'stop'") worker = db.query(Worker).filter(Worker.user_id == current_user.id).first() if not worker: raise HTTPException(status_code=404, detail="Worker no encontrado") worker.command = command db.commit() return {"message": f"Comando '{command}' enviado al worker"} @router.get("/agent/poll") def agent_poll(token: str, status: str, hostname: str, db: Session = Depends(get_db)): worker = db.query(Worker).filter(Worker.token == token).first() if not worker: raise HTTPException(status_code=404, detail="Token de Worker no registrado") worker.last_seen = datetime.utcnow() # Si el administrador suspendió el worker, mantener el estado "suspended" y detener el engine if worker.status == "suspended": current_command = "stop" if status == "busy" else worker.command if worker.command != "none": worker.command = "none" db.commit() return { "command": current_command, "is_active": False, "user_id": worker.user_id, "status": "suspended" } # De lo contrario, actualizar estado según el reporte del agente previous_status = worker.status if status == "offline": worker.status = "offline" elif status == "online": if worker.status != "busy": worker.status = "online" else: worker.status = status if hostname: worker.hostname = hostname current_command = worker.command if worker.command != "none": worker.command = "none" db.commit() if previous_status == "offline" and worker.status in ["online", "busy"]: create_admin_notification( db, "Worker Conectado", f"El worker local/BYOK '{worker.name or worker.hostname}' se ha conectado." ) return { "command": current_command, "is_active": worker.is_active, "user_id": worker.user_id, "status": worker.status } @router.get("/agent/poll-public") def agent_poll_public(hostname: str, status: str, db: Session = Depends(get_db)): worker = db.query(Worker).filter(Worker.hostname == hostname, Worker.user_id == None).first() is_new = worker is None if not worker: # Registrar de forma automática si no existe en la base de datos worker = Worker( hostname=hostname, name=hostname, status="online", last_seen=datetime.utcnow(), command="none", is_active=True ) db.add(worker) db.commit() db.refresh(worker) # Enviar notificación de nuevo worker create_admin_notification( db, "Worker Registrado", f"El equipo worker '{worker.hostname}' se ha autoregistrado vía poll público." ) worker.last_seen = datetime.utcnow() # Si el administrador suspendió el worker, mantener el estado "suspended" if worker.status == "suspended": current_command = "stop" if status == "busy" else worker.command if worker.command != "none": worker.command = "none" db.commit() return { "command": current_command, "status": "suspended" } # De lo contrario, actualizar estado según el reporte del agente previous_status = worker.status if status == "offline": worker.status = "offline" elif status == "online": if worker.status != "busy": worker.status = "online" else: worker.status = status current_command = worker.command if worker.command != "none": worker.command = "none" db.commit() if not is_new and previous_status == "offline" and worker.status in ["online", "busy"]: create_admin_notification( db, "Worker Conectado", f"El equipo worker '{worker.name or worker.hostname}' se ha conectado." ) return { "command": current_command, "status": worker.status }