agent-maxim/bridge/controller/controller_loop.py

"""Controller Loop V1.

Minimal autonomous orchestration:
  goal → decision → bridge run (adapter=agent) → observe → decision → stop

Stop conditions:
  1. decision.action == "stop"
  2. iteration >= MAX_ITERATIONS
  3. bridge run status == done AND agent_final_status == done  (natural success)
  4. bridge run failed twice in a row  (repeated failure guard)
"""
from __future__ import annotations

import json
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

import adapters as adapter_registry
from controller.controller_state import ControllerState, save as save_state, state_path
from controller.decision_model import DECISION_MODE, decide
from core.run_manager import RunManager
from policy.guard import check as policy_check

MAX_ITERATIONS: int = 3
RUNS_DIR: Path = Path(__file__).resolve().parents[1] / "runs"


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


# ---------------------------------------------------------------------------
# Controller-level audit logger
# ---------------------------------------------------------------------------

class _CtrlAudit:
    def __init__(self, ctrl_id: str, runs_dir: Path) -> None:
        self._path = runs_dir / ctrl_id / "controller_audit.jsonl"
        self._ctrl_id = ctrl_id
        self._path.parent.mkdir(parents=True, exist_ok=True)

    def log(self, event: str, *, status: str = "ok", details: Optional[dict] = None) -> None:
        record = {
            "timestamp": _now(),
            "trace_id": self._ctrl_id,     # trace_id == controller_run_id
            "run_id": self._ctrl_id,
            "layer": "controller",
            "controller_run_id": self._ctrl_id,
            "event_type": event,
            "status": status,
            "details": details or {},
        }
        with self._path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(record) + "\n")


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------

def run_controller(
    goal: str,
    run_id: Optional[str] = None,
) -> ControllerState:
    """Execute the full controller loop and return final ControllerState."""
    ctrl_id = run_id or f"ctrl-{str(uuid.uuid4())[:8]}"

    if state_path(RUNS_DIR, ctrl_id).exists():
        raise ValueError(f"controller run_id '{ctrl_id}' already exists")

    state = ControllerState(
        controller_run_id=ctrl_id,
        original_goal=goal,
        decision_mode=DECISION_MODE,
    )
    audit = _CtrlAudit(ctrl_id, RUNS_DIR)
    manager = RunManager(RUNS_DIR)

    save_state(state, RUNS_DIR)
    audit.log("controller_started", details={"goal": goal, "decision_mode": DECISION_MODE,
                                              "max_iterations": MAX_ITERATIONS})

    _t_total = time.time()
    last_result: Optional[dict[str, Any]] = None
    consecutive_failures = 0
    recovery_attempted = False

    for i in range(MAX_ITERATIONS):
        state.iteration_count = i + 1
        state.touch()

        # 1. DECIDE
        decision = decide(goal, i, last_result, MAX_ITERATIONS, recovery_attempted, trace_id=ctrl_id)

        # Log recovery audit events based on decision
        if last_result is not None and (
            last_result.get("status") == "failed"
            or last_result.get("agent_final_status") == "failed"
        ) and not recovery_attempted:
            audit.log("recovery_considered", details={
                "iteration": i,
                "original_goal": goal,
                "recovery_applied": decision.recovery_applied,
                "reason": decision.reason,
            })
            if decision.recovery_applied:
                audit.log("recovery_applied", details={
                    "iteration": i,
                    "original_goal": goal,
                    "rewritten_goal": decision.goal,
                    "rule_reason": decision.reason,
                })
                state.recovery_attempted = True
                state.recovery_reason = decision.reason
                state.rewritten_goal = decision.goal
                recovery_attempted = True
                # Update running goal to rewritten goal
                goal = decision.goal
            else:
                audit.log("recovery_skipped", details={
                    "iteration": i,
                    "original_goal": goal,
                    "reason": decision.reason,
                })

        state.decisions.append({**decision.to_dict(), "iteration": i, "decided_at": _now()})
        save_state(state, RUNS_DIR)

        # Hard stop guard — layer 2: if last run succeeded, never run agent again
        terminal_enforced = False
        if (
            decision.action == "run_agent"
            and last_result is not None
            and last_result.get("status") == "done"
            and last_result.get("agent_final_status") == "done"
        ):
            decision.action = "stop"
            terminal_enforced = True

        audit.log("decision_made", details={
            "iteration": i,
            "action": decision.action,
            "goal": decision.goal,
            "reason": decision.reason,
            "mode": decision.mode,
            "parsed_success": decision.parsed_success,
            "fallback_used": decision.mode in ("stub", "ollama-fallback", "real-fallback"),
            "fallback_reason": decision.fallback_reason,
            "corrected": decision.corrected,
            "correction_reason": decision.correction_reason,
            "terminal_enforced": terminal_enforced,
        })

        if decision.action == "stop":
            # Determine final status from last result
            if last_result is None:
                state.final_status = "stopped"
            elif last_result.get("status") == "done":
                state.final_status = "done"
            else:
                state.final_status = "failed"
            break

        # 2. EXECUTE BRIDGE RUN
        bridge_run_id = f"bri-{ctrl_id}-i{i}"
        agent_mod = adapter_registry.get("agent")

        audit.log("bridge_run_started", details={
            "iteration": i,
            "bridge_run_id": bridge_run_id,
            "goal": decision.goal,
        })

        _t_bridge = time.time()
        try:
            bridge_result = manager.execute_run(
                goal=decision.goal,
                adapter_name="agent",
                adapter_fn=agent_mod.execute,   # type: ignore[attr-defined]
                guard_fn=policy_check,
                run_id=bridge_run_id,
                trace_id=ctrl_id,
            )
        except Exception as exc:  # noqa: BLE001
            error_msg = f"{type(exc).__name__}: {exc}"
            audit.log("bridge_run_error", status="failed", details={
                "iteration": i,
                "bridge_run_id": bridge_run_id,
                "error": error_msg,
            })
            bridge_result = {
                "run_id": bridge_run_id,
                "status": "failed",
                "error": error_msg,
                "agent_final_status": "failed",
                "exit_code": 1,
            }

        print(f"[TIMING] bridge.execute (iter {i}): {time.time() - _t_bridge:.3f}s")

        state.linked_bridge_run_ids.append(bridge_run_id)
        state.bridge_run_results.append({
            "iteration": i,
            "bridge_run_id": bridge_run_id,
            "status": bridge_result.get("status"),
            "agent_final_status": bridge_result.get("agent_final_status"),
            "exit_code": bridge_result.get("exit_code"),
            "error": bridge_result.get("error"),
        })
        last_result = bridge_result
        state.touch()
        save_state(state, RUNS_DIR)

        audit.log("bridge_run_finished", details={
            "iteration": i,
            "bridge_run_id": bridge_run_id,
            "status": bridge_result.get("status"),
            "agent_final_status": bridge_result.get("agent_final_status"),
            "exit_code": bridge_result.get("exit_code"),
        })

        # 3. IMMEDIATE STOP — natural success, no further decision call
        if (bridge_result.get("status") == "done"
                and bridge_result.get("agent_final_status") == "done"):
            state.final_status = "done"
            save_state(state, RUNS_DIR)
            audit.log("controller_finished", status="done",
                      details={"reason": "agent success", "iterations": state.iteration_count,
                               "linked_runs": state.linked_bridge_run_ids})
            print(f"[TIMING] run_controller total: {time.time() - _t_total:.3f}s")
            return state

        # Repeated failure guard
        consecutive_failures += 1
        if consecutive_failures >= 2:
            state.final_status = "failed"
            audit.log("controller_finished", status="failed",
                      details={"reason": "Repeated consecutive failures", "iterations": state.iteration_count})
            save_state(state, RUNS_DIR)
            print(f"[TIMING] run_controller total: {time.time() - _t_total:.3f}s")
            return state

    else:
        # MAX_ITERATIONS exhausted without explicit stop
        if last_result and last_result.get("status") == "done":
            state.final_status = "done"
        else:
            state.final_status = "failed"

    save_state(state, RUNS_DIR)
    audit.log("controller_finished", status=state.final_status,
              details={"iterations": state.iteration_count,
                       "linked_runs": state.linked_bridge_run_ids})
    print(f"[TIMING] run_controller total: {time.time() - _t_total:.3f}s")
    return state