forja/tests/unit/test_runtime_orchestrator.py

"""Tests del AgentOrchestrator: invoke, pausa HITL, resume (aprobar/rechazar) y persistencia."""

from datetime import UTC, datetime
from pathlib import Path
from uuid import uuid4

import pytest

from agentforge_core.domain.agent import AgentDefinition, LLMConfig
from agentforge_core.domain.policy import PolicyDefinition, PolicyValidator
from agentforge_core.guardrails.guardrails_ai import GuardrailsAIEngine
from agentforge_core.llm.mock import MockProvider
from agentforge_core.runtime.orchestrator import AgentOrchestrator


def _agent(threshold: int = 4) -> AgentDefinition:
    return AgentDefinition(
        name="incident_analyzer",
        version="v1",
        owner="Juan",
        purpose="Análisis de incidentes",
        state="active",
        guardrails=["default"],
        llm=LLMConfig(provider="mock"),
        system_prompt="Eres un analista. Responde SIEMPRE con JSON.",
        output_schema={"type": "object"},
        risk_threshold_for_hitl=threshold,
        updated_at=datetime.now(UTC),
    )


def _policy(input_validators: list[PolicyValidator] | None = None) -> PolicyDefinition:
    return PolicyDefinition(
        name="p",
        version="v1",
        description="t",
        input_validators=input_validators or [],
        output_validators=[],
    )


@pytest.fixture
def orchestrator(tmp_path: Path) -> AgentOrchestrator:
    return AgentOrchestrator(
        provider=MockProvider(), engine=GuardrailsAIEngine(), data_dir=tmp_path
    )


async def test_invoke_camino_feliz(orchestrator: AgentOrchestrator) -> None:
    ex = await orchestrator.invoke(
        agent_def=_agent(), policy=_policy(), user_input="degradación MOS pool SBC"
    )
    assert ex.status == "completed"
    assert ex.finished_at is not None
    assert ex.needs_human_for is None
    assert ex.proposed_actions and ex.proposed_actions[0].action == "scale_out_sbc_pool"
    assert ex.final_output is not None and ex.final_output["approved_actions"]
    assert [s.step for s in ex.decision_path] == [
        "validate_input",
        "llm_reason",
        "validate_output",
        "propose_actions",
        "approve_gate",
        "finalize",
    ]


async def test_snapshot_relee_estado_sin_avanzar(orchestrator: AgentOrchestrator) -> None:
    agent = _agent()
    paused = await orchestrator.invoke(
        agent_def=agent, policy=_policy(), user_input="caída registros sip"
    )
    again = await orchestrator.snapshot(agent_def=agent, policy=_policy(), trace_id=paused.trace_id)
    assert again is not None
    assert again.status == "awaiting_approval"
    assert again.trace_id == paused.trace_id
    # Idempotente: leer no añade pasos al decision_path.
    assert [s.step for s in again.decision_path] == [s.step for s in paused.decision_path]


async def test_snapshot_trace_id_desconocido_devuelve_none(orchestrator: AgentOrchestrator) -> None:
    assert await orchestrator.snapshot(agent_def=_agent(), policy=_policy(), trace_id=uuid4()) is None


async def test_invoke_pausa_hitl_y_resume_aprueba(orchestrator: AgentOrchestrator) -> None:
    agent = _agent()
    paused = await orchestrator.invoke(
        agent_def=agent, policy=_policy(), user_input="caída registros sip"
    )
    assert paused.status == "awaiting_approval"
    assert paused.finished_at is None
    assert {a.id for a in paused.needs_human_for or []} == {"act-1"}  # act-2 risk=3 < 4

    resumed = await orchestrator.resume(
        agent_def=agent,
        policy=_policy(),
        trace_id=paused.trace_id,
        decision={"approved_action_ids": ["act-1"]},
    )
    assert resumed.status == "completed"
    assert resumed.trace_id == paused.trace_id
    assert resumed.final_output is not None
    assert {a["id"] for a in resumed.final_output["approved_actions"]} == {"act-1"}


async def test_resume_rechazo_marca_failed(orchestrator: AgentOrchestrator) -> None:
    agent = _agent()
    paused = await orchestrator.invoke(
        agent_def=agent, policy=_policy(), user_input="caída registros sip"
    )
    resumed = await orchestrator.resume(
        agent_def=agent,
        policy=_policy(),
        trace_id=paused.trace_id,
        decision={"rejected": True},
    )
    assert resumed.status == "failed"
    assert resumed.error == "rejected_by_human"


async def test_invoke_bloqueado_por_guardrail(orchestrator: AgentOrchestrator) -> None:
    policy = _policy(
        [
            PolicyValidator(
                type="detect_pii",
                config={"entities": ["EMAIL_ADDRESS"], "severity_on_match": "block"},
            )
        ]
    )
    ex = await orchestrator.invoke(
        agent_def=_agent(), policy=policy, user_input="manda correo a juan@example.com"
    )
    assert ex.status == "blocked_by_guardrail"
    assert ex.violations and any(v.blocked for v in ex.violations)
    assert ex.final_output is None


async def test_estado_persiste_entre_instancias(tmp_path: Path) -> None:
    agent = _agent()
    o1 = AgentOrchestrator(
        provider=MockProvider(), engine=GuardrailsAIEngine(), data_dir=tmp_path
    )
    paused = await o1.invoke(agent_def=agent, policy=_policy(), user_input="caída registros sip")
    assert paused.status == "awaiting_approval"

    # Nueva instancia (simula reinicio del proceso) apuntando al mismo data_dir.
    o2 = AgentOrchestrator(
        provider=MockProvider(), engine=GuardrailsAIEngine(), data_dir=tmp_path
    )
    resumed = await o2.resume(
        agent_def=agent,
        policy=_policy(),
        trace_id=paused.trace_id,
        decision={"approved_action_ids": ["act-1"]},
    )
    assert resumed.status == "completed"