api/job_queue.py

"""
backend/api/job_queue.py — Redis coordination BRAIN↔HANDS (S-DUAL-2)

Questo modulo implementa tre livelli di coordinamento via Upstash Redis:

  1. LOAD EXCHANGE — entrambi gli Space pubblicano metriche ogni 30s su Redis.
     Il CF router legge via /api/jq/load/{role} senza HTTP probe costosi.
     Chiavi: jq:load:brain  jq:load:hands  (TTL 90s)

  2. WAKE-UP — BRAIN segnala HANDS di scaldarsi prima che il circuito si chiuda.
     HANDS consumer legge i segnali e chiama /health su se stesso.
     Chiave: jq:wake  (LIST, RPOP, TTL 30s per elemento)

  3. TASK DELEGATION — BRAIN accoda task, HANDS consuma ed esegue.
     Chiave: jq:tasks:pending  (LIST, LPUSH/RPOP)
     Chiave: jq:result:{taskId}  (STRING, TTL 300s)
     Chiave: jq:events:{taskId}  (LIST, TTL 300s)
     Chiave: jq:consumer:alive   (STRING, TTL 30s — heartbeat HANDS consumer)

  Abilitazione:
     JQ_ENABLED=1 + UPSTASH_REDIS_REST_URL + UPSTASH_REDIS_REST_TOKEN + SPACE_ROLE
     Se JQ_ENABLED non è impostato → modulo in standby (load publish attivo, queue no)

  Endpoints:
     GET  /api/jq/status              — diagnostica queue
     GET  /api/jq/load/{role}         — metriche load da Redis (brain|hands)
     POST /api/jq/wake                — segnala wake-up HANDS (solo BRAIN)
     POST /api/jq/submit              — sottometti job a HANDS via Redis (solo BRAIN)
     GET  /api/jq/result/{taskId}     — leggi risultato job da Redis
     GET  /api/jq/events/{taskId}     — leggi eventi SSE da Redis
"""
import os, asyncio, json, time, uuid, logging
from fastapi import APIRouter, Request, HTTPException
from pydantic import BaseModel

_logger = logging.getLogger("api.job_queue")

router = APIRouter(prefix="/api/jq", tags=["job-queue"])

# ── Config ─────────────────────────────────────────────────────────────────────
_SPACE_ROLE  = os.getenv("SPACE_ROLE", "unknown")       # brain | hands | unknown
_JQ_ENABLED  = os.getenv("JQ_ENABLED", "0").strip() == "1"
_LOAD_TTL    = 90    # s — TTL metriche load su Redis
_RESULT_TTL  = 300   # s — TTL risultato job su Redis
_EVENTS_TTL  = 300   # s — TTL lista eventi SSE su Redis
_WAKE_TTL    = 30    # s — TTL singolo wake signal
_CONSUMER_HB_TTL = 30  # s — TTL heartbeat consumer HANDS

# ── Redis keys ─────────────────────────────────────────────────────────────────
_K_LOAD       = lambda role: f"jq:load:{role}"         # STRING — metriche load
_K_WAKE       = "jq:wake"                               # LIST   — wake signals
_K_PENDING    = "jq:tasks:pending"                      # LIST   — job queue
_K_RESULT     = lambda tid: f"jq:result:{tid}"          # STRING — risultato job
_K_EVENTS     = lambda tid: f"jq:events:{tid}"          # LIST   — eventi SSE
_K_CONSUMER   = "jq:consumer:alive"                     # STRING — HB consumer

# ── Bootstrap check ────────────────────────────────────────────────────────────
def _redis_ok() -> bool:
    """Verifica che Redis sia configurato (non controlla la connettività)."""
    return bool(
        os.getenv("UPSTASH_REDIS_REST_URL") and
        os.getenv("UPSTASH_REDIS_REST_TOKEN")
    )


async def _rcmd(command: list, timeout: float = 3.0) -> dict | None:
    """Esegue comando Redis via modulo centralizzato backend/redis.py."""
    try:
        from redis import redis_cmd
        return await redis_cmd(command, timeout=timeout)
    except Exception as exc:
        _logger.debug("[jq] redis_cmd error: %s", exc)
        return None


async def _rpush(key: str, value: str, ttl: int | None = None) -> bool:
    """LPUSH key value + EXPIRE key ttl (se ttl != None)."""
    r = await _rcmd(["LPUSH", key, value])
    if r and ttl:
        await _rcmd(["EXPIRE", key, ttl])
    return r is not None


async def _rpop(key: str) -> str | None:
    """RPOP key. Ritorna None se lista vuota o errore."""
    r = await _rcmd(["RPOP", key])
    if r and r.get("result") is not None:
        return r["result"]
    return None


async def _llen(key: str) -> int:
    """LLEN key."""
    r = await _rcmd(["LLEN", key])
    return int(r.get("result", 0)) if r else 0


async def _lrange(key: str, start: int = 0, stop: int = -1) -> list[str]:
    """LRANGE key start stop."""
    r = await _rcmd(["LRANGE", key, start, stop])
    return r.get("result", []) if r else []


# ── Load metric exchange ───────────────────────────────────────────────────────

async def publish_load_metrics(role: str | None = None) -> bool:
    """
    Pubblica le metriche di carico di questo Space su Redis.
    Chiamato ogni 30s dal background loop _load_publisher_loop().
    Complementa /api/health/load (HTTP) con uno snapshot Redis consultabile
    senza HTTP call dall'altro Space.
    """
    if not _redis_ok():
        return False
    _role = role or _SPACE_ROLE
    try:
        from api.priority import get_load_metrics as _glm
        metrics = _glm()
    except Exception:
        metrics = {}
    try:
        from api.state import _agent_tasks
        active = sum(1 for t in _agent_tasks.values() if t.get("status") in ("RUNNING", "running"))
    except Exception:
        active = 0

    payload = json.dumps({
        "space_role":         _role,
        "active_agent_tasks": active,
        "realtime_active":    metrics.get("realtime_active", 0),
        "background_active":  metrics.get("background_active", 0),
        "realtime_available": metrics.get("realtime_available", 6),
        "consumer_enabled":   _JQ_ENABLED,
        "ts":                 int(time.time() * 1000),
    })
    r = await _rcmd(["SET", _K_LOAD(_role), payload, "EX", _LOAD_TTL])
    return r is not None


async def get_remote_load(role: str) -> dict | None:
    """Legge le metriche load dell'altro Space da Redis. Ritorna None se stale."""
    if not _redis_ok():
        return None
    r = await _rcmd(["GET", _K_LOAD(role)])
    if not r or r.get("result") is None:
        return None
    try:
        return json.loads(r["result"])
    except Exception:
        return None


# ── Wake-up signaling ──────────────────────────────────────────────────────────

async def publish_wake_signal(reason: str = "circuit_open") -> bool:
    """
    BRAIN pubblica un segnale di wake-up per HANDS.
    HANDS consumer legge il segnale e chiama /health su se stesso
    per prevenire il cold start di HF Space free tier.
    """
    if not _redis_ok():
        return False
    payload = json.dumps({"reason": reason, "ts": int(time.time() * 1000), "from": _SPACE_ROLE})
    return await _rpush(_K_WAKE, payload, ttl=_WAKE_TTL)


async def _consume_wake_signals() -> int:
    """HANDS: legge e processa tutti i wake signal in coda. Ritorna il count."""
    count = 0
    while True:
        raw = await _rpop(_K_WAKE)
        if raw is None:
            break
        try:
            sig = json.loads(raw)
            _logger.info("[jq] wake signal from %s — reason: %s", sig.get("from"), sig.get("reason"))
        except Exception:
            pass
        count += 1
    return count


# ── Task delegation ────────────────────────────────────────────────────────────

class JobPayload(BaseModel):
    goal:       str
    session_id: str = ""
    context:    dict = {}
    priority:   str = "realtime"  # realtime | background
    max_steps:  int = 20
    task_id:    str = ""  # se vuoto → generato da BRAIN


async def submit_job(job: JobPayload) -> dict:
    """
    BRAIN: accoda un task per HANDS via Redis.
    Ritorna {taskId, status, stream_url, queue_depth}.
    
    Il client usa stream_url per connettersi direttamente a HANDS (via CF routing)
    e ricevere gli eventi SSE una volta che HANDS ha consumato il job.
    """
    if not _redis_ok():
        raise HTTPException(503, "Redis non configurato — job queue non disponibile")

    task_id = job.task_id or str(uuid.uuid4())
    payload = json.dumps({
        "taskId":     task_id,
        "goal":       job.goal,
        "session_id": job.session_id,
        "context":    job.context,
        "priority":   job.priority,
        "max_steps":  job.max_steps,
        "submitted_at": time.time(),
        "submitted_by": _SPACE_ROLE,
    })

    ok = await _rpush(_K_PENDING, payload)
    if not ok:
        raise HTTPException(503, "Impossibile accodare il task su Redis")

    depth = await _llen(_K_PENDING)
    _logger.info("[jq] job queued taskId=%s depth=%d", task_id, depth)

    return {
        "taskId":     task_id,
        "status":     "queued",
        "queue_depth": depth,
        "stream_url": f"/api/agent/tasks/{task_id}/stream",
    }


async def _execute_queued_job(job: dict) -> None:
    """
    HANDS: esegue un job prelevato dalla coda Redis.
    Crea il task nel registro locale e avvia unified_loop.
    Gli eventi SSE vengono scritti sia nel registro locale (per streaming diretto)
    sia su Redis (per relay da BRAIN).
    """
    task_id = job.get("taskId", str(uuid.uuid4()))
    goal    = job.get("goal", "")
    _logger.info("[jq] executing queued job taskId=%s goal=%.60s", task_id, goal)

    try:
        from api.state import _agent_tasks
        _agent_tasks[task_id] = {
            "status":     "QUEUED",
            "goal":       goal,
            "session_id": job.get("session_id", ""),
            "created_at": time.time() * 1000,
            "source":     "jq",
        }

        # Pubblica evento di start su Redis
        await _rcmd(["LPUSH", _K_EVENTS(task_id), json.dumps({
            "type": "task_queued", "taskId": task_id, "ts": int(time.time() * 1000)
        })])
        await _rcmd(["EXPIRE", _K_EVENTS(task_id), _EVENTS_TTL])

        # Lancia il loop tramite agent.py create_agent_task
        try:
            from api.agent import _create_task_internal
            await _create_task_internal(task_id=task_id, goal=goal, job=job)
        except (ImportError, AttributeError):
            # Fallback: usa unified_loop direttamente
            from agents.unified_loop import UnifiedLoop
            loop = UnifiedLoop()
            result = await loop.run(
                goal=goal,
                context=json.dumps(job.get("context", {})),
                max_steps=job.get("max_steps", 20),
                session_id=job.get("session_id", ""),
            )
            # Pubblica risultato
            await _rcmd(["SET", _K_RESULT(task_id), json.dumps({
                "taskId":       task_id,
                "status":       "success" if result.get("success") else "error",
                "output":       result.get("output", ""),
                "error":        result.get("error"),
                "completed_at": time.time(),
            }), "EX", _RESULT_TTL])
            _agent_tasks[task_id]["status"] = "SUCCESS" if result.get("success") else "ERROR"

    except Exception as exc:
        _logger.error("[jq] job execution failed taskId=%s: %s", task_id, exc, exc_info=True)
        await _rcmd(["SET", _K_RESULT(task_id), json.dumps({
            "taskId":       task_id,
            "status":       "error",
            "error":        str(exc),
            "completed_at": time.time(),
        }), "EX", _RESULT_TTL])
        try:
            from api.state import _agent_tasks
            _agent_tasks[task_id]["status"] = "ERROR"
        except Exception:
            pass


# ── Background loops ───────────────────────────────────────────────────────────

async def _load_publisher_loop() -> None:
    """Background loop: pubblica metriche load ogni 30s su Redis."""
    _logger.info("[jq] load publisher avviato (role=%s)", _SPACE_ROLE)
    while True:
        await asyncio.sleep(30)
        try:
            await publish_load_metrics()
        except Exception as exc:
            _logger.debug("[jq] load publish error: %s", exc)


async def _hands_consumer_loop() -> None:
    """
    Background loop HANDS: ogni 1s legge dalla queue Redis.
    - Consuma wake signals (log + no-op — siamo già svegli)
    - Consuma job dalla coda e li esegue in background
    - Pubblica heartbeat consumer ogni 10s
    """
    _logger.info("[jq] HANDS job consumer avviato")
    _hb_tick = 0
    while True:
        await asyncio.sleep(1)
        try:
            # Wake signals — siamo già svegli ma logghiamo
            await _consume_wake_signals()

            # Heartbeat consumer ogni ~10s
            _hb_tick += 1
            if _hb_tick % 10 == 0:
                await _rcmd(["SET", _K_CONSUMER, str(int(time.time())), "EX", _CONSUMER_HB_TTL])

            if not _JQ_ENABLED:
                continue  # load publisher attivo, job consumer no

            # Preleva job dalla coda
            raw = await _rpop(_K_PENDING)
            if raw is None:
                continue

            try:
                job = json.loads(raw)
            except Exception:
                _logger.warning("[jq] invalid job payload — skipping")
                continue

            # Esegui in background — non blocca il loop consumer
            t = asyncio.create_task(_execute_queued_job(job))
            t.add_done_callback(lambda task: (
                _logger.error("[jq] job task crashed: %s", task.exception(), exc_info=task.exception())
                if not task.cancelled() and task.exception() else None
            ))

        except Exception as exc:
            _logger.debug("[jq] consumer tick error: %s", exc)


async def start_job_queue_consumer() -> None:
    """
    Punto di ingresso per main.py _on_startup().
    Avvia:
    - _load_publisher_loop() (sempre, su tutti gli Space)
    - _hands_consumer_loop() (solo se SPACE_ROLE=hands o unknown)
    """
    if not _redis_ok():
        _logger.warning("[jq] Redis non configurato — job queue disabilitato")
        return

    # Load publisher su tutti gli Space
    asyncio.create_task(_load_publisher_loop())

    # Consumer solo su HANDS (o se role non impostato per compatibilità)
    if _SPACE_ROLE in ("hands", "unknown"):
        asyncio.create_task(_hands_consumer_loop())
    else:
        _logger.info("[jq] SPACE_ROLE=%s — consumer non avviato (solo load publisher)", _SPACE_ROLE)

    # Pubblica subito le metriche al boot
    await publish_load_metrics()


# ── FastAPI endpoints ──────────────────────────────────────────────────────────

@router.get("/status")
async def jq_status():
    """Diagnostica completa della job queue."""
    _redis_configured = _redis_ok()
    result = {
        "space_role":       _SPACE_ROLE,
        "jq_enabled":       _JQ_ENABLED,
        "redis_configured": _redis_configured,
        "ts":               int(time.time() * 1000),
    }
    if _redis_configured:
        result["queue_depth"]     = await _llen(_K_PENDING)
        result["wake_pending"]    = await _llen(_K_WAKE)
        _hb = await _rcmd(["GET", _K_CONSUMER])
        result["consumer_alive"]  = bool(_hb and _hb.get("result"))
        result["brain_load"]      = await get_remote_load("brain")
        result["hands_load"]      = await get_remote_load("hands")
    return result


@router.get("/load/{role}")
async def jq_load(role: str):
    """Legge le metriche load di uno Space da Redis. role: brain | hands"""
    if role not in ("brain", "hands"):
        raise HTTPException(400, "role deve essere 'brain' o 'hands'")
    data = await get_remote_load(role)
    if data is None:
        raise HTTPException(404, f"Metriche {role} non disponibili (stale o Redis non configurato)")
    return data


@router.post("/wake")
async def jq_wake(request: Request):
    """BRAIN: invia segnale wake-up a HANDS. Solo da BRAIN o con X-Internal-Token."""
    _tok = os.getenv("INTERNAL_TOKEN", "")
    if _tok and request.headers.get("X-Internal-Token", "") != _tok:
        raise HTTPException(401, "Unauthorized")
    ok = await publish_wake_signal(reason="manual_wake")
    return {"sent": ok, "ts": int(time.time() * 1000)}


@router.post("/submit")
async def jq_submit(job: JobPayload, request: Request):
    """
    BRAIN: sottomette un job a HANDS via Redis.
    Richiede X-Internal-Token.
    Ritorna {taskId, status, stream_url} — il client usa stream_url per SSE.
    """
    _tok = os.getenv("INTERNAL_TOKEN", "")
    if _tok and request.headers.get("X-Internal-Token", "") != _tok:
        raise HTTPException(401, "Unauthorized")
    if not _JQ_ENABLED:
        raise HTTPException(503, "JQ_ENABLED non impostato — job queue disabilitato")
    return await submit_job(job)


@router.get("/result/{task_id}")
async def jq_result(task_id: str):
    """Legge il risultato di un job delegato da Redis. Disponibile per ~5 min post-completamento."""
    if not _redis_ok():
        raise HTTPException(503, "Redis non configurato")
    r = await _rcmd(["GET", _K_RESULT(task_id)])
    if not r or r.get("result") is None:
        raise HTTPException(404, f"Risultato per {task_id} non trovato (non ancora pronto o scaduto)")
    try:
        return json.loads(r["result"])
    except Exception:
        raise HTTPException(500, "Risultato malformato in Redis")


@router.get("/events/{task_id}")
async def jq_events(task_id: str, from_idx: int = 0):
    """
    Legge gli eventi SSE di un task da Redis (per relay da BRAIN).
    from_idx: indice da cui iniziare (0 = tutti).
    """
    if not _redis_ok():
        raise HTTPException(503, "Redis non configurato")
    events_raw = await _lrange(_K_EVENTS(task_id), from_idx)
    events = []
    for e in events_raw:
        try:
            events.append(json.loads(e))
        except Exception:
            events.append({"raw": e})
    return {"taskId": task_id, "events": events, "count": len(events), "from_idx": from_idx}