agent-maxim/agent/core/agent_loop.py

"""Main agent execution loop: select → resolve → execute → retry? → update → persist → repeat."""
from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path

from core.classifier import FailureClass, classify
from core.resolver import has_templates, resolve_payload
from core.retry_policy import should_retry
from core.state import AgentState, RunStatus, StepResult, StepStatus
from planner.step_selector import select_next_step
from runtime.executor import Executor
from telemetry.audit_logger import AuditLogger

MAX_STEPS = 20


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def run_loop(
    state: AgentState,
    executor: Executor,
    logger: AuditLogger,
    runs_dir: Path,
) -> AgentState:
    """Main loop: select → resolve → execute → retry-or-continue → save → repeat.

    Terminates when:
    - no pending steps remain (DONE)
    - template resolution fails and no retry allowed (FAILED)
    - a step fails and no retry allowed (FAILED)
    - a 'finish' step completes (DONE)
    - MAX_STEPS is exceeded (FAILED)
    """
    from core.controller import save_state  # avoid circular import

    for _ in range(MAX_STEPS):
        step = select_next_step(state.plan)
        if step is None:
            state.status = RunStatus.DONE
            break

        # Track attempt number (0 = first time through)
        is_retry = step.attempts > 0
        step.attempts += 1
        current_attempt = step.attempts

        # pending → running
        step.status = StepStatus.RUNNING
        state.current_step_id = step.step_id
        state.updated_at = _now()

        if is_retry:
            logger.log(
                "retry_attempt_started",
                step_id=step.step_id,
                status="running",
                details={"kind": step.kind, "attempt": current_attempt},
            )
        else:
            logger.log(
                "step_started",
                step_id=step.step_id,
                status="running",
                details={"kind": step.kind, "description": step.description},
            )

        # --- Template resolution (before executor) ---
        if has_templates(step.payload):
            resolved_payload, ok, err = resolve_payload(step.payload, state)
            logger.log(
                "resolution_performed",
                step_id=step.step_id,
                status="ok" if ok else "failed",
                details={"templates_found": True, "success": ok, "error": err if not ok else None},
            )
            if not ok:
                result = StepResult(
                    step_id=step.step_id,
                    status=StepStatus.FAILED,
                    error=f"Template resolution failed: {err}",
                    started_at=_now(),
                    finished_at=_now(),
                    attempt=current_attempt,
                    failure_class=FailureClass.TEMPLATE_ERROR.value,
                )
                step.status = StepStatus.FAILED
                state.history.append(result)
                state.updated_at = _now()
                logger.log(
                    "retry_attempt_finished" if is_retry else "step_finished",
                    step_id=step.step_id,
                    status="failed",
                    details={"error": result.error, "attempt": current_attempt},
                )
                logger.log(
                    "step_failed",
                    step_id=step.step_id,
                    status="failed",
                    details={
                        "failure_class": FailureClass.TEMPLATE_ERROR.value,
                        "attempt": current_attempt,
                        "retry_scheduled": False,
                        "reason": "template_error is never retried",
                    },
                )
                state.status = RunStatus.FAILED
                state.current_step_id = None
                save_state(state, runs_dir, logger)
                break
            step.payload = resolved_payload
        else:
            if step.kind not in ("finish",):
                logger.log(
                    "resolution_performed",
                    step_id=step.step_id,
                    status="ok",
                    details={"templates_found": False, "success": True},
                )

        # --- Execute ---
        started = _now()
        result = executor.execute(step)
        result.started_at = started
        result.finished_at = _now()
        result.attempt = current_attempt

        logger.log(
            "retry_attempt_finished" if is_retry else "step_finished",
            step_id=step.step_id,
            status=result.status.value,
            details={
                "output_text": (result.output_text or "")[:300],
                "error": result.error,
                "attempt": current_attempt,
            },
        )

        # --- Retry / fail decision ---
        if result.status == StepStatus.FAILED:
            failure_class = classify(step.kind, result.error or "")
            result.failure_class = failure_class.value
            retry_ok = should_retry(failure_class, step.attempts)

            logger.log(
                "step_failed",
                step_id=step.step_id,
                status="failed",
                details={
                    "failure_class": failure_class.value,
                    "attempt": current_attempt,
                    "retry_scheduled": retry_ok,
                },
            )

            if retry_ok:
                step.status = StepStatus.PENDING  # reset for retry
                state.history.append(result)
                state.updated_at = _now()
                logger.log(
                    "retry_scheduled",
                    step_id=step.step_id,
                    status="ok",
                    details={
                        "attempt_just_failed": current_attempt,
                        "next_attempt": current_attempt + 1,
                        "failure_class": failure_class.value,
                    },
                )
                save_state(state, runs_dir, logger)
                # Loop continues — select_next_step will pick up PENDING step again
            else:
                step.status = StepStatus.FAILED
                state.history.append(result)
                state.updated_at = _now()
                state.status = RunStatus.FAILED
                state.current_step_id = None
                save_state(state, runs_dir, logger)
                break

        else:
            # Step succeeded
            step.status = StepStatus.DONE
            state.history.append(result)
            state.updated_at = _now()
            save_state(state, runs_dir, logger)

            if step.kind == "finish":
                state.status = RunStatus.DONE
                state.current_step_id = None
                break

    else:
        state.status = RunStatus.FAILED
        state.current_step_id = None
        logger.log(
            "run_aborted",
            status="failed",
            details={"reason": f"Exceeded MAX_STEPS ({MAX_STEPS})"},
        )

    return state