Nexus/server/api/agent.py

"""Nexus — Agent API Routes (Heartbeat receiver + exec endpoint)
Presentation layer — Agent servers call these endpoints to report status and receive commands.

Heartbeat flow:
  1. Agent sends heartbeat → write to Redis (real-time data for frontend)
  2. Alert detection: CPU/mem/disk > threshold → WebSocket push + Telegram
  3. Recovery detection: previously alerted metrics now normal → WebSocket + Telegram
  4. Background task: Redis→MySQL batch flush every 10min
"""

import json
import logging
from datetime import datetime, timezone
from typing import Optional

from fastapi import APIRouter, Depends, Header, HTTPException

from server.api.dependencies import get_server_service
from server.api.schemas import AgentHeartbeat, AgentExec
from server.api.websocket import broadcast_alert, broadcast_recovery
from server.application.services.server_service import ServerService
from server.config import settings
from server.infrastructure.redis.client import get_redis, get_redis_sync
from server.infrastructure.telegram import send_telegram_system_alert

logger = logging.getLogger("nexus.agent")

router = APIRouter(prefix="/api/agent", tags=["agent"])

# Redis key patterns
REDIS_KEY_PREFIX = "heartbeat:"
REDIS_KEY_EXPIRE = 600  # 10 min TTL — matches flush interval
REDIS_ALERT_KEY_PREFIX = "alerts:"


def _verify_api_key(x_api_key: str = Header(...)):
    """Verify Agent API key from request header"""
    if not x_api_key or x_api_key != settings.API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")
    return x_api_key


def _detect_alerts(system_info: dict, thresholds: dict) -> list:
    """Detect metrics exceeding alert thresholds

    Returns list of (alert_type, value) tuples.
    E.g. [("cpu", 85.3), ("mem", 92.1)]
    """
    alerts = []
    cpu = system_info.get("cpu_usage") or system_info.get("cpu")
    mem = system_info.get("mem_usage") or system_info.get("mem")
    disk = system_info.get("disk_usage") or system_info.get("disk")

    if cpu is not None and float(cpu) > thresholds.get("cpu", settings.CPU_ALERT_THRESHOLD):
        alerts.append(("cpu", float(cpu)))
    if mem is not None and float(mem) > thresholds.get("mem", settings.MEM_ALERT_THRESHOLD):
        alerts.append(("mem", float(mem)))
    if disk is not None and float(disk) > thresholds.get("disk", settings.DISK_ALERT_THRESHOLD):
        alerts.append(("disk", float(disk)))

    return alerts


def _detect_recovery(system_info: dict, prev_alerts: set, thresholds: dict) -> list:
    """Detect previously alerted metrics that have returned to normal

    Returns list of (metric, current_value) tuples.
    """
    recoveries = []
    for prev in prev_alerts:
        parts = prev.split(":")
        if len(parts) != 2:
            continue
        metric, prev_val = parts[0], float(parts[1])
        current = system_info.get(f"{metric}_usage") or system_info.get(metric)

        if current is not None and float(current) < thresholds.get(metric, 80):
            recoveries.append((metric, float(current)))

    return recoveries


@router.post("/heartbeat", response_model=dict)
async def receive_heartbeat(
    payload: AgentHeartbeat,
    api_key: str = Depends(_verify_api_key),
    service: ServerService = Depends(get_server_service),
):
    """Receive Agent heartbeat data

    Flow:
    1. Write heartbeat to Redis (frontend reads Redis for live status)
    2. Detect alerts (CPU/mem/disk > threshold) → WebSocket + Telegram push
    3. Detect recoveries (previously alerted metrics now normal) → push
    4. Also update MySQL via ServerService (for persistent record)
    """
    server_id = payload.server_id
    is_online = payload.is_online
    system_info = payload.system_info or {}
    agent_version = payload.agent_version or ""

    # ── 1. Write to Redis (real-time data) ──
    redis = get_redis()
    try:
        redis_key = f"{REDIS_KEY_PREFIX}{server_id}"
        await redis.hset(redis_key, mapping={
            "is_online": str(is_online),
            "system_info": json.dumps(system_info),
            "last_heartbeat": datetime.now(timezone.utc).isoformat(),
            "agent_version": agent_version,
        })
        # Set TTL = flush interval × 1.5 (prevent stale online status after key expires)
        await redis.expire(redis_key, REDIS_KEY_EXPIRE)
    except Exception as e:
        logger.error(f"Redis heartbeat write failed for server {server_id}: {e}")
        # Redis write failed — still continue with MySQL update

    # ── 1b. Look up server name for alert broadcasts ──
    server_name = f"server-{server_id}"
    try:
        server = await service.get_server(server_id)
        if server:
            server_name = server.name
    except Exception:
        pass

    # ── 2. Alert detection ──
    alert_thresholds = {
        "cpu": settings.CPU_ALERT_THRESHOLD,
        "mem": settings.MEM_ALERT_THRESHOLD,
        "disk": settings.DISK_ALERT_THRESHOLD,
    }

    if system_info:
        alerts = _detect_alerts(system_info, alert_thresholds)
        for alert_type, alert_value in alerts:
            logger.warning(f"Alert: server {server_id} {alert_type}={alert_value}%")
            await broadcast_alert(server_id, alert_type, alert_value, server_name)

            # Track active alert in Redis
            try:
                redis = get_redis()
                await redis.sadd(f"{REDIS_ALERT_KEY_PREFIX}{server_id}", f"{alert_type}:{alert_value}")
                await redis.expire(f"{REDIS_ALERT_KEY_PREFIX}{server_id}", 3600)  # 1h TTL
            except Exception:
                pass  # Non-critical

    # ── 3. Recovery detection ──
    if system_info:
        try:
            redis = get_redis()
            prev_alerts = await redis.smembers(f"{REDIS_ALERT_KEY_PREFIX}{server_id}")
            if prev_alerts:
                recoveries = _detect_recovery(system_info, prev_alerts, alert_thresholds)
                for metric, value in recoveries:
                    logger.info(f"Recovery: server {server_id} {metric}={value}%")
                    await broadcast_recovery(server_id, metric, value, server_name)

                    # Remove recovered alert from Redis tracking
                    for prev in prev_alerts:
                        if prev.startswith(f"{metric}:"):
                            await redis.srem(f"{REDIS_ALERT_KEY_PREFIX}{server_id}", prev)
        except Exception as e:
            logger.error(f"Recovery detection failed for server {server_id}: {e}")

    # ── 4. Also update MySQL (persistent record for non-real-time queries) ──
    await service.update_heartbeat(server_id, payload.model_dump())

    return {"status": "ok", "server_id": server_id}


@router.post("/exec", response_model=dict)
async def agent_exec(
    payload: AgentExec,
    api_key: str = Depends(_verify_api_key),
):
    """Execute a shell command on this Agent server

    This endpoint is called by Nexus to dispatch commands to remote Agents.
    Each Agent server runs its own instance of this endpoint.
    """
    command = payload.command
    timeout = payload.timeout
    if payload.sudo:
        command = f"sudo {command}"

    import asyncio
    try:
        proc = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
        exit_code = proc.returncode or 0

        stdout_text = stdout.decode()[:10000]
        stderr_text = stderr.decode()[:10000]

        return {
            "status": "success" if exit_code == 0 else "failed",
            "stdout": stdout_text,
            "stderr": stderr_text,
            "exit_code": exit_code,
        }
    except asyncio.TimeoutError:
        proc.kill()
        return {
            "status": "timeout",
            "stdout": "",
            "stderr": f"Command timed out after {timeout}s",
            "exit_code": -1,
        }
    except Exception as e:
        return {
            "status": "error",
            "stdout": "",
            "stderr": str(e),
            "exit_code": -1,
        }