Lifecycle de Execução

import asyncio
import redis.asyncio as aioredis
from contextlib import asynccontextmanager

HEARTBEAT_TTL      = 30   # segundos — TTL da chave no Redis
HEARTBEAT_INTERVAL = 10   # segundos — frequência de renovação (razão 3x com TTL)


class AgentHeartbeat:
    def __init__(self, redis: aioredis.Redis, thread_id: str, task_id: str):
        self.redis     = redis
        self.thread_id = thread_id
        self.task_id   = task_id  # ECS task ID — útil para correlacionar logs
        self._key      = f"heartbeat:{thread_id}"

    async def _beat(self):
        # SETEX: SET + EXPIRE atômico.
        # Valor = task_id — permite saber qual ECS task estava rodando o thread.
        await self.redis.setex(self._key, HEARTBEAT_TTL, self.task_id)

    @asynccontextmanager
    async def running(self):
        async def _loop():
            while True:
                await self._beat()
                await asyncio.sleep(HEARTBEAT_INTERVAL)

        task = asyncio.create_task(_loop())
        try:
            await self._beat()  # bate imediatamente ao entrar — não espera 10s
            yield
        finally:
            task.cancel()
            # Deleção explícita: OrphanScanner não vai detectar este thread como órfão.
            # Sem isso, o TTL de 30s causaria um falso positivo para threads normais.
            await self.redis.delete(self._key)


async def run_agent(thread_id: str, input: dict, ecs_task_id: str):
    await redis.sadd("threads:in_progress", thread_id)

    heartbeat = AgentHeartbeat(redis, thread_id, ecs_task_id)
    async with heartbeat.running():
        result = await graph.ainvoke(
            input, config={"configurable": {"thread_id": thread_id}},
        )

    await redis.srem("threads:in_progress", thread_id)
    return result

class OrphanScanner:
    def __init__(self, redis):
        self.redis = redis

    async def scan_once(self) -> list[str]:
        thread_ids: set[str] = await self.redis.smembers("threads:in_progress")
        if not thread_ids:
            return []

        orphans = []
        for thread_id in thread_ids:
            alive = await self.redis.exists(f"heartbeat:{thread_id}")
            if not alive:
                orphans.append(thread_id)
        return orphans

    async def enqueue_orphans(self, orphans: list[str], guard: QueueGuard) -> None:
        for thread_id in orphans:
            # Verifica se já está na fila — o scanner pode rodar enquanto
            # um ciclo anterior ainda está sendo processado.
            already = await self.redis.lpos("recovery:queue", thread_id)
            if already is not None:
                continue

            await self.redis.srem("threads:in_progress", thread_id)
            enqueued = await guard.safe_enqueue(thread_id)
            if enqueued:
                logger.warning("Órfão detectado e enfileirado: %s", thread_id)

    async def run(self, guard: QueueGuard) -> None:
        logger.info("OrphanScanner iniciado (intervalo=20s)")
        while True:
            try:
                orphans = await self.scan_once()
                if orphans:
                    await self.enqueue_orphans(orphans, guard)
            except Exception:
                logger.exception("Erro no OrphanScanner")
            await asyncio.sleep(20)

class RecoveryConsumer:
    def __init__(self, redis, graph, dlq: "DLQManager"):
        self.redis = redis
        self.graph = graph
        self.dlq   = dlq

    async def consume_one(self, thread_id: str) -> None:
        config = {"configurable": {"thread_id": thread_id}}
        try:
            state = await self.graph.aget_state(config)

            if not state:
                # Checkpoint não existe — pode ter expirado ou nunca ter sido salvo.
                # Reenfileirar não vai resolver: vai para DLQ diretamente.
                await self.dlq.send_to_dlq(
                    thread_id, reason="no_checkpoint",
                    error="Checkpoint ausente ou expirado",
                )
                return

            if not state.next:
                # Estado terminal — grafo já terminou.
                # Provavelmente falso positivo do scanner. Limpa sem reexecutar.
                logger.info("Thread %s em estado terminal. Falso positivo.", thread_id)
                await self.redis.delete(f"recovery:attempts:{thread_id}")
                return

            logger.info("Retomando %s nos nós: %s", thread_id, list(state.next))
            await asyncio.wait_for(
                self.graph.ainvoke(None, config=config),
                timeout=120,
            )
            await self.redis.delete(f"recovery:attempts:{thread_id}")
            logger.info("Thread %s retomado com sucesso.", thread_id)

        except asyncio.TimeoutError:
            await self.dlq.handle_failure(
                thread_id, reason="timeout", error="Execução excedeu 120s")
        except Exception as exc:
            await self.dlq.handle_failure(
                thread_id, reason="exception", error=str(exc))

    async def run(self) -> None:
        while True:
            try:
                # BLPOP bloqueia até item disponível — sem polling busy.
                # timeout=5: retorna None após 5s sem item — evita bloqueio eterno.
                result = await self.redis.blpop("recovery:queue", timeout=5)
                if result:
                    _, thread_id = result
                    await self.consume_one(thread_id)
            except Exception:
                logger.exception("Erro no RecoveryConsumer")
                await asyncio.sleep(1)

from datetime import datetime, timezone

MAX_ATTEMPTS = int(os.getenv("RECOVERY_MAX_ATTEMPTS", "3"))
ATTEMPTS_TTL = 60 * 60 * 24  # 24h — reseta automaticamente


class DLQManager:
    def __init__(self, redis):
        self.redis = redis

    async def get_attempts(self, thread_id: str) -> int:
        val = await self.redis.get(f"recovery:attempts:{thread_id}")
        return int(val) if val else 0

    async def increment_attempts(self, thread_id: str) -> int:
        key      = f"recovery:attempts:{thread_id}"
        attempts = await self.redis.incr(key)
        if attempts == 1:
            # Define TTL apenas na primeira falha — INCR não aceita EX.
            await self.redis.expire(key, ATTEMPTS_TTL)
        return attempts

    async def send_to_dlq(self, thread_id: str, reason: str, error: str = "") -> None:
        attempts = await self.get_attempts(thread_id)
        meta = {
            "thread_id": thread_id,
            "reason":    reason,
            "error":     error[:500],
            "attempts":  str(attempts),
            "failed_at": datetime.now(timezone.utc).isoformat(),
        }
        pipe = self.redis.pipeline()
        pipe.rpush("recovery:dlq", thread_id)
        pipe.hset(f"recovery:dlq:meta:{thread_id}", mapping=meta)
        pipe.delete(f"recovery:attempts:{thread_id}")
        await pipe.execute()
        logger.error("Thread %s → DLQ após %d tentativas. reason=%s",
                     thread_id, attempts, reason)

    async def handle_failure(self, thread_id: str, reason: str, error: str = "") -> None:
        attempts = await self.increment_attempts(thread_id)
        if attempts >= MAX_ATTEMPTS:
            await self.send_to_dlq(thread_id, reason=reason, error=error)
        else:
            logger.warning("Thread %s falhou (%d/%d). Reenfileirando.",
                           thread_id, attempts, MAX_ATTEMPTS)
            await self.redis.rpush("recovery:queue", thread_id)

    async def list_dlq(self) -> list[dict]:
        ids    = await self.redis.lrange("recovery:dlq", 0, -1)
        result = []
        for tid in ids:
            meta = await self.redis.hgetall(f"recovery:dlq:meta:{tid}")
            result.append(meta or {"thread_id": tid})
        return result

async def requeue_from_dlq(self, thread_id: str) -> bool:
    pos = await self.redis.lpos("recovery:dlq", thread_id)
    if pos is None:
        logger.warning("Thread %s não encontrado na DLQ.", thread_id)
        return False

    # Pipeline: remove da DLQ, limpa meta, reseta contador, reenfileira.
    # Tudo atomicamente — não deixa estado parcial.
    pipe = self.redis.pipeline()
    pipe.lrem("recovery:dlq", 1, thread_id)
    pipe.delete(f"recovery:dlq:meta:{thread_id}")
    pipe.delete(f"recovery:attempts:{thread_id}")
    pipe.rpush("recovery:queue", thread_id)
    await pipe.execute()
    logger.info("Thread %s removido da DLQ e reenfileirado.", thread_id)
    return True

from contextlib import asynccontextmanager

_active_tasks: set[asyncio.Task] = set()
_shutting_down = False


@asynccontextmanager
async def lifespan(app):
    yield
    global _shutting_down
    _shutting_down = True
    if _active_tasks:
        logger.info("Aguardando %d task(s) ativas...", len(_active_tasks))
        # timeout=25 — stop_timeout do ECS task definition deve ser 120s,
        # deixando margem para overhead de shutdown após as tasks terminarem.
        await asyncio.wait(_active_tasks, timeout=25.0)

app = FastAPI(lifespan=lifespan)


@app.post("/agents/invoke")
async def invoke(body: InvokeRequest):
    if _shutting_down:
        # 503 durante shutdown — ALB roteia para outro target saudável.
        raise HTTPException(status_code=503, detail="Serviço em manutenção.")
    task = asyncio.current_task()
    _active_tasks.add(task)
    try:
        return await run_agent(body.thread_id, body.input)
    finally:
        _active_tasks.discard(task)

import pytest
from unittest.mock import AsyncMock, MagicMock
import fakeredis.aioredis as fakeredis
from recovery_worker import OrphanScanner, RecoveryConsumer, DLQManager, QueueGuard

@pytest.fixture
def redis(): return fakeredis.FakeRedis(decode_responses=True)

@pytest.fixture
def dlq(redis): return DLQManager(redis)


@pytest.mark.asyncio
async def test_scanner_detecta_orfao(redis):
    # Thread em andamento sem heartbeat → deve ser detectado como órfão
    await redis.sadd("threads:in_progress", "t-001")
    orphans = await OrphanScanner(redis).scan_once()
    assert "t-001" in orphans


@pytest.mark.asyncio
async def test_scanner_ignora_thread_vivo(redis):
    # Thread com heartbeat ativo → não é órfão
    await redis.sadd("threads:in_progress", "t-002")
    await redis.setex("heartbeat:t-002", 30, "task-xyz")
    orphans = await OrphanScanner(redis).scan_once()
    assert "t-002" not in orphans


@pytest.mark.asyncio
async def test_dlq_apos_max_attempts(redis, dlq):
    for i in range(3):  # MAX_ATTEMPTS = 3
        await dlq.handle_failure("t-003", reason="exception", error=f"err {i}")

    items = await redis.lrange("recovery:dlq", 0, -1)
    assert "t-003" in items
    meta  = await redis.hgetall("recovery:dlq:meta:t-003")
    assert meta["reason"] == "exception"
    assert meta["attempts"] == "3"


@pytest.mark.asyncio
async def test_consumer_limpa_attempts_em_sucesso(redis, dlq):
    # Simula uma falha anterior — tentativa 1
    await dlq.increment_attempts("t-004")

    mock_graph = MagicMock()
    mock_graph.aget_state = AsyncMock(return_value=MagicMock(next=("node_b",)))
    mock_graph.ainvoke    = AsyncMock(return_value={"ok": True})

    await RecoveryConsumer(redis, mock_graph, dlq).consume_one("t-004")

    # Após sucesso: contador deve ser zerado, ainvoke chamado com None
    assert await dlq.get_attempts("t-004") == 0
    mock_graph.ainvoke.assert_called_once_with(
        None, config={"configurable": {"thread_id": "t-004"}}
    )