agent-maxim/bridge/core/run_manager.py

"""Run lifecycle management for Bridge V1.2.

One-shot flow inside execute_run():
  created → running → done | failed

All state transitions are persisted immediately.
Agent-adapter metadata (agent_run_id, agent_final_status, agent_state_path,
exit_code) is extracted from AdapterResult.metadata and stored in state.json.
"""
from __future__ import annotations

import inspect
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from adapters.base import AdapterResult
from core.audit import BridgeAuditLogger
from core.instruction_builder import build as build_instruction
from core.state_store import StateStore
from models.schemas import RunStatus


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


class RunManager:
    def __init__(self, runs_dir: Path) -> None:
        self._runs_dir = runs_dir
        self._store = StateStore(runs_dir)
        self._cancel_flags: set[str] = set()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def execute_run(
        self,
        goal: str,
        adapter_name: str,
        adapter_fn,   # callable(prompt, **kwargs) -> AdapterResult
        guard_fn,     # callable(prompt) -> (bool, str)
        run_id: Optional[str] = None,
        trace_id: Optional[str] = None,
    ) -> dict[str, Any]:
        """Full one-shot lifecycle: created → running → done|failed."""
        rid = run_id or str(uuid.uuid4())[:8]
        if self._store.exists(rid):
            raise ValueError(f"run_id '{rid}' already exists")

        logger = BridgeAuditLogger(rid, self._runs_dir, trace_id=trace_id)
        now = _now()

        # 1. CREATED
        state: dict[str, Any] = {
            "run_id": rid,
            "trace_id": trace_id,
            "goal": goal,
            "status": RunStatus.CREATED.value,
            "adapter": adapter_name,
            "adapter_mode": None,
            "instruction": None,
            "response": None,
            "error": None,
            "exit_code": None,
            "agent_run_id": None,
            "agent_final_status": None,
            "agent_state_path": None,
            "created_at": now,
            "updated_at": now,
        }
        self._store.save(rid, state)
        logger.log("run_created", details={"goal": goal, "adapter": adapter_name})

        # 2. BUILD INSTRUCTION
        instruction = build_instruction(goal)
        state["instruction"] = instruction
        state["updated_at"] = _now()
        self._store.save(rid, state)
        logger.log("instruction_generated", details={"instruction_len": len(instruction)})

        # 3. POLICY CHECK
        ok, reason = guard_fn(instruction)
        if not ok:
            return self._fail(state, logger, f"Policy blocked: {reason}")

        # 4. RUNNING
        state["status"] = RunStatus.RUNNING.value
        state["updated_at"] = _now()
        self._store.save(rid, state)
        logger.log("task_sent", details={
            "adapter": adapter_name,
            "instruction_len": len(instruction),
        })

        # 5. EXECUTE — pass bridge_run_id, trace_id, and is_cancelled callback
        sig = inspect.signature(adapter_fn)
        kwargs: dict[str, Any] = {}
        if "bridge_run_id" in sig.parameters:
            kwargs["bridge_run_id"] = rid
        if "trace_id" in sig.parameters:
            kwargs["trace_id"] = trace_id
        if "is_cancelled" in sig.parameters:
            kwargs["is_cancelled"] = lambda: rid in self._cancel_flags
        result: AdapterResult = adapter_fn(instruction, **kwargs)

        # 6. PERSIST ADAPTER METADATA
        state["adapter_mode"] = result.mode
        state["exit_code"] = result.exit_code
        if result.metadata:
            state["agent_run_id"] = result.metadata.get("agent_run_id")
            state["agent_final_status"] = result.metadata.get("agent_final_status")
            state["agent_state_path"] = result.metadata.get("agent_state_path")

        logger.log(
            "response_received",
            status="ok" if result.ok else "failed",
            details={
                "adapter": adapter_name,
                "mode": result.mode,
                "exit_code": result.exit_code,
                "response_len": len(result.response) if result.response else 0,
                "error": result.error,
                "agent_run_id": result.metadata.get("agent_run_id"),
            },
        )

        # 7. CLEANUP: Remove cancel flag after task completes (success, failed, or cancelled)
        self.clear_cancel_flag(rid)

        # 8. DONE or FAILED or CANCELLED
        if result.ok:
            state["response"] = result.response
            state["status"] = RunStatus.DONE.value
        elif result.exit_code == 130 or (result.error and "user_cancelled" in result.error):
            # Special exit code for cancelled runs
            state["response"] = result.response  # may have partial output
            state["error"] = result.error or "user_cancelled"
            state["status"] = "cancelled"
        else:
            state["response"] = result.response  # may have partial output
            state["error"] = result.error
            state["status"] = RunStatus.FAILED.value

        state["updated_at"] = _now()
        self._store.save(rid, state)
        logger.log(
            "run_finished",
            status=state["status"],
            details={
                "adapter_mode": result.mode,
                "exit_code": result.exit_code,
                "agent_run_id": result.metadata.get("agent_run_id"),
                "agent_final_status": result.metadata.get("agent_final_status"),
            },
        )

        return state

    def get_run(self, run_id: str) -> dict[str, Any] | None:
        return self._store.load(run_id)

    def cancel_run(self, run_id: str) -> None:
        """Mark a run as cancelled by adding it to in-memory cancel flags."""
        self._cancel_flags.add(run_id)

    def is_cancelled(self, run_id: str) -> bool:
        """Check if a run is marked as cancelled (for internal use by adapters)."""
        return run_id in self._cancel_flags

    def clear_cancel_flag(self, run_id: str) -> None:
        """Remove cancel flag after task completion (for internal use)."""
        self._cancel_flags.discard(run_id)

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _fail(
        self,
        state: dict[str, Any],
        logger: BridgeAuditLogger,
        error: str,
    ) -> dict[str, Any]:
        state["error"] = error
        state["status"] = RunStatus.FAILED.value
        state["updated_at"] = _now()
        self._store.save(state["run_id"], state)
        logger.log("run_finished", status="failed", details={"error": error})
        return state