api/job_queue.py

import os
import json
import time
import asyncio
import logging
import uuid
try:
    import resource as _resource
    _HAS_RESOURCE = True
except ImportError:
    _HAS_RESOURCE = False
from typing import Optional, List, Dict, Any
from fastapi import APIRouter, HTTPException, Request, Depends
from .auth_guard import require_role, AuthRole
from pydantic import BaseModel
from .load_balancer import balancer  # S951
import httpx

_logger = logging.getLogger("agente_ai.jq")
router = APIRouter(prefix="/jq", tags=["job_queue"])

# ── Configurazione ────────────────────────────────────────────────────────────
_SPACE_ROLE     = os.getenv("SPACE_ROLE", "unknown")
_JQ_ENABLED     = os.getenv("JQ_ENABLED", "0") == "1"
_INTERNAL_TOKEN = os.getenv("INTERNAL_TOKEN", "")

# C2-FIX: contatore reale di job attivi — incrementato/decrementato intorno a loop_inst.run()
_active_job_count: int = 0

# Redis keys (Upstash)
_K_PENDING  = "jq:pending"
_K_WAKE     = "jq:wake"
_K_CONSUMER = f"jq:consumer:{_SPACE_ROLE}"
_K_RESULT   = lambda tid: f"jq:result:{tid}"
_K_EVENTS   = lambda tid: f"jq:events:{tid}"
_K_LOAD     = lambda role: f"jq:load:{role}"

class JobPayload(BaseModel):
    taskId: str
    goal: str
    context: Optional[Dict[str, Any]] = None
    priority: int = 1

# ── Helper Redis (via HTTP REST per stabilità mobile/serverless) ──────────────
async def _rcmd(cmd: List[Any]) -> Optional[Dict[str, Any]]:
    url = os.getenv("UPSTASH_REDIS_REST_URL")
    tok = os.getenv("UPSTASH_REDIS_REST_TOKEN")
    if not url or not tok:
        return None
    try:
        async with httpx.AsyncClient() as client:
            r = await client.post(
                url,
                headers={"Authorization": f"Bearer {tok}"},
                json=cmd,
                timeout=5.0
            )
            return r.json()
    except Exception as e:
        _logger.error("[jq] redis error: %s", e)
        return None

def _redis_ok() -> bool:
    return bool(os.getenv("UPSTASH_REDIS_REST_URL") and os.getenv("UPSTASH_REDIS_REST_TOKEN"))

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

async def _lrange(key: str, start: int, end: int = -1) -> List[str]:
    res = await _rcmd(["LRANGE", key, start, end])
    return res.get("result", []) if res else []

# ── Core Logic ────────────────────────────────────────────────────────────────
async def publish_load_metrics():
    """Pubblica il carico corrente su Redis per il bilanciamento (S951).
    C2-FIX: usa _active_job_count (contatore reale) invece di asyncio.all_tasks()-5.
    """
    if not _redis_ok(): return
    # Memoria processo in MB (Linux: ru_maxrss è in kB)
    mem_mb = 0
    if _HAS_RESOURCE:
        try:
            mem_mb = _resource.getrusage(_resource.RUSAGE_SELF).ru_maxrss // 1024
        except Exception:
            mem_mb = 0
    data = {
        "role":         _SPACE_ROLE,
        "ts":           int(time.time() * 1000),
        "active_tasks": max(0, _active_job_count),  # mai negativo
        "cpu":          0,                            # placeholder — psutil non installato
        "mem":          mem_mb,
    }
    await _rcmd(["SET", _K_LOAD(_SPACE_ROLE), json.dumps(data), "EX", "60"])

async def _load_publisher_loop():
    """Loop periodico per aggiornare lo stato del nodo."""
    while True:
        try:
            await publish_load_metrics()
        except Exception as e:
            _logger.debug("[jq] load publisher error: %s", e)
        await asyncio.sleep(30)

async def _hands_consumer_loop():
    """Loop consumer reale per nodi distribuiti (S42)."""
    global _active_job_count
    _logger.info("[jq] consumer loop avviato per ruolo: %s", _SPACE_ROLE)
    _consecutive_errors = 0
    _MAX_CONSECUTIVE_ERRORS = 3  # INV-R1: dopo 3 errori infra → pausa 30s

    while True:
        try:
            # 1. Heartbeat consumer
            await _rcmd(["SET", _K_CONSUMER, "1", "EX", "15"])

            # 2. Prelievo job
            res = await _rcmd(["RPOP", _K_PENDING])
            if res and res.get("result"):
                job_raw = res["result"]
                try:
                    job_data = json.loads(job_raw)
                    # S951: Check assignment
                    assigned_to = job_data.get("assignedTo", "hands")
                    if assigned_to != _SPACE_ROLE and _SPACE_ROLE != "brain":
                        await _rcmd(["LPUSH", _K_PENDING, job_raw])
                        await asyncio.sleep(1)
                        continue
                except Exception:
                    pass
                task_id = job_data.get("taskId", str(uuid.uuid4()))
                goal = job_data.get("goal", "")
                context_raw = job_data.get("context", {})
                _logger.info("[jq] job ricevuto: %s | goal: %.80s", task_id, goal)

                if isinstance(context_raw, dict):
                    context_str = "\n".join(f"{k}: {v}" for k, v in context_raw.items() if v)
                elif isinstance(context_raw, str):
                    context_str = context_raw
                else:
                    context_str = ""

                result_payload: Dict[str, Any] = {}
                try:
                    from agents.unified_loop import UnifiedAgentLoop
                    from .state import (
                        _get_ai_client, _get_mem_manager_async,
                        _get_executor, _get_planner,
                    )
                    ai_client = _get_ai_client()
                    try:
                        from agents.critic import Critic
                        from agents.response_verifier import ResponseVerifier
                        _critic   = Critic(llm_client=ai_client)
                        _verifier = ResponseVerifier()
                    except Exception:
                        _critic   = None
                        _verifier = None

                    _steps: list = []

                    async def _on_step(step: dict) -> None:
                        _steps.append({
                            "action": step.get("action", ""),
                            "output": str(step.get("output", ""))[:200],
                        })

                    loop_inst = UnifiedAgentLoop(
                        llm_client=ai_client,
                        critic=_critic,
                        verifier=_verifier,
                        memory=await _get_mem_manager_async(),
                        executor=_get_executor(),
                        planner=_get_planner(),
                    )

                    _priority = int(job_data.get("priority", 1))
                    _max_steps = max(5, min(4 + _priority * 2, 16))

                    # C2-FIX: incrementa contatore PRIMA di run(), decrementa in finally
                    _active_job_count += 1
                    try:
                        raw_result = await loop_inst.run(
                            goal=goal,
                            context=context_str,
                            max_steps=_max_steps,
                            on_step=_on_step,
                            session_id=task_id,
                        )
                    finally:
                        _active_job_count = max(0, _active_job_count - 1)

                    output  = raw_result.get("output", "") if isinstance(raw_result, dict) else str(raw_result)
                    success = bool(raw_result.get("success", False)) if isinstance(raw_result, dict) else bool(output)
                    engine  = raw_result.get("engine", "unknown") if isinstance(raw_result, dict) else "unknown"

                    result_payload = {
                        "taskId":  task_id,
                        "status":  "completed" if success else "failed",
                        "worker":  _SPACE_ROLE,
                        "output":  output[:4000],
                        "engine":  engine,
                        "success": success,
                        "steps":   len(_steps),
                        "ts":      int(time.time() * 1000),
                    }
                    _consecutive_errors = 0
                    _logger.info("[jq] job completato: %s | engine: %s | success: %s", task_id, engine, success)

                except (ImportError, ModuleNotFoundError) as imp_err:
                    _logger.warning("[jq] agents.unified_loop non disponibile su %s: %s", _SPACE_ROLE, imp_err)
                    result_payload = {
                        "taskId": task_id,
                        "status": "unavailable",
                        "worker": _SPACE_ROLE,
                        "error":  f"UnifiedAgentLoop non disponibile su nodo {_SPACE_ROLE}: {imp_err}",
                        "ts":     int(time.time() * 1000),
                    }

                except Exception as exec_err:
                    _consecutive_errors += 1
                    _logger.error("[jq] job execution error task=%s: %s", task_id, exec_err)
                    result_payload = {
                        "taskId": task_id,
                        "status": "error",
                        "worker": _SPACE_ROLE,
                        "error":  str(exec_err)[:500],
                        "ts":     int(time.time() * 1000),
                    }
                    if _consecutive_errors >= _MAX_CONSECUTIVE_ERRORS:
                        _logger.error(
                            "[jq] INV-R1: %d errori consecutivi — pausa 30s (nodo: %s)",
                            _consecutive_errors, _SPACE_ROLE,
                        )
                        await asyncio.sleep(30)

                # 3. Salva risultato su Redis (TTL 1h) — Tool Success Contract S429
                await _rcmd(["SET", _K_RESULT(task_id), json.dumps(result_payload), "EX", "3600"])
                _logger.info("[jq] risultato Redis: %s → %s", task_id, result_payload.get("status"))

            else:
                _consecutive_errors = 0
                await asyncio.sleep(2)

        except Exception as e:
            _logger.error("[jq] consumer loop error: %s", e)
            await asyncio.sleep(5)

async def start_job_queue_consumer() -> None:
    """Punto di ingresso per main.py _on_startup()."""
    if not _redis_ok():
        _logger.warning("[jq] Redis non configurato — job queue disabilitato")
        return
    _bg_tasks.append(asyncio.create_task(_load_publisher_loop()))
    if _SPACE_ROLE in ("hands", "memory", "audit", "unknown"):
        _bg_tasks.append(asyncio.create_task(_hands_consumer_loop()))
    else:
        _logger.info("[jq] SPACE_ROLE=%s — consumer non avviato (solo load publisher)", _SPACE_ROLE)

# ── FastAPI endpoints ──────────────────────────────────────────────────────────
@router.get("/status")
async def jq_status(
    role: AuthRole = Depends(require_role(AuthRole.MACHINE)),
):
    return {
        "space_role":       _SPACE_ROLE,
        "jq_enabled":       _JQ_ENABLED,
        "redis_configured": _redis_ok(),
        "active_tasks":     max(0, _active_job_count),
        "ts":               int(time.time() * 1000),
    }

@router.get("/load/{role}")
async def jq_load(
    role: str,
    auth_role: AuthRole = Depends(require_role(AuthRole.MACHINE)),
):
    if role not in ("brain", "hands", "memory", "audit"):
        raise HTTPException(400, "role non valido")
    res = await _rcmd(["GET", _K_LOAD(role)])
    if not res or not res.get("result"):
        raise HTTPException(404, f"Metriche {role} non disponibili")
    try:
        return json.loads(res["result"])
    except json.JSONDecodeError as _je:
        raise HTTPException(500, f"Metriche Redis corrotte per {role}: {_je}")

@router.post("/submit")
async def jq_submit(
    job: JobPayload,
    role: AuthRole = Depends(require_role(AuthRole.MACHINE)),
):
    """S42 — Grid Orchestrator: sottomissione job reale su Redis."""
    if not _redis_ok():
        raise HTTPException(503, "Job queue non disponibile: Redis non configurato")

    task_id = job.taskId if job.taskId else str(uuid.uuid4())

    # S951: Dynamic Load Balancing
    target_role = job.context.get('preferred_role', 'hands') if job.context else 'hands'
    assigned_role = await balancer.get_best_node(target_role)

    job_payload = {
        'taskId':      task_id,
        'goal':        job.goal,
        'context':     job.context or {},
        'priority':    job.priority,
        'submittedAt': int(time.time() * 1000),
        'submittedBy': _SPACE_ROLE,
        'assignedTo':  assigned_role,
    }
    res = await _rcmd(["LPUSH", _K_PENDING, json.dumps(job_payload)])
    if res is None:
        raise HTTPException(503, "Errore Redis durante la sottomissione del job")

    queue_len = int(res.get("result", 0)) if res else 0
    await _rcmd(["SET", _K_WAKE, "1", "EX", "10"])
    _logger.info("[jq] job sottomesso: %s (coda: %d)", task_id, queue_len)
    return {
        "taskId":      task_id,
        "status":      "queued",
        "queueLength": queue_len,
        "ts":          int(time.time() * 1000),
    }

@router.get("/result/{task_id}")
async def jq_result(
    task_id: str,
    role: AuthRole = Depends(require_role(AuthRole.MACHINE)),
):
    """S429 — Tool Success Contract: recupero risultato job per taskId."""
    if not _redis_ok():
        raise HTTPException(503, "Job queue non disponibile: Redis non configurato")

    res = await _rcmd(["GET", _K_RESULT(task_id)])
    if not res or not res.get("result"):
        return {"taskId": task_id, "status": "pending", "ts": int(time.time() * 1000)}
    try:
        return json.loads(res["result"])
    except json.JSONDecodeError as _je:
        _logger.error("[jq] risultato Redis corrotto task=%s: %s", task_id, _je)
        raise HTTPException(500, f"Risultato corrotto in Redis per {task_id}: {_je}")

# Lista dei background tasks (popolata da start_job_queue_consumer)
_bg_tasks: List[asyncio.Task] = []