agentic-mobile-control/workflow/runner.py

from __future__ import annotations

from collections.abc import Callable
from datetime import datetime
from types import SimpleNamespace
from typing import Any
from time import sleep

from core.models import Scene, Task
from runtime.task import StopReason, TaskRunner, is_cancellation_reason
from skills_learning.store import SkillStore, get_default_store
from workflow.conditions import (
    ConditionEvaluator,
    UnknownConditionKindError,
    evaluate_condition,
)
from workflow.config import WorkflowConfig, load_config
from workflow.models import (
    BranchStep,
    PlannedGoalStep,
    SkillInvocationStep,
    WaitForConditionStep,
    WorkflowDefinition,
    WorkflowRun,
    WorkflowStep,
    WorkflowStepResult,
)
from workflow.skill_exec import SkillExecutionError, run_flow_template_skill
from workflow.store import WorkflowStore

TaskRunnerFactory = Callable[[], TaskRunner]
SceneProvider = Callable[[], Scene | None]
WorldStateProvider = Callable[[], object | None]
SleepFunc = Callable[[float], None]
StopRequested = Callable[[], bool]

TERMINAL_STATUSES = {"completed", "failed", "cancelled"}


class WorkflowRunner:
    def __init__(
        self,
        store: WorkflowStore | None = None,
        *,
        task_runner_factory: TaskRunnerFactory | None = None,
        skill_store: SkillStore | None = None,
        tools: dict[str, Callable[..., Any]] | None = None,
        condition_registry: dict[str, ConditionEvaluator] | None = None,
        scene_provider: SceneProvider | None = None,
        world_state_provider: WorldStateProvider | None = None,
        sleep_func: SleepFunc = sleep,
        config: WorkflowConfig | None = None,
        step_limit: int | None = None,
    ) -> None:
        self.config = config or load_config()
        self.store = store or WorkflowStore(self.config.db_path)
        self.task_runner_factory = task_runner_factory or (lambda: TaskRunner())
        self.skill_store = skill_store or get_default_store()
        self.tools = tools
        self.condition_registry = condition_registry
        self.scene_provider = scene_provider or (lambda: None)
        self.world_state_provider = world_state_provider
        self.sleep_func = sleep_func
        self.step_limit = step_limit

    def run(
        self,
        definition: WorkflowDefinition,
        device_id: str,
        initial_variables: dict[str, Any] | None = None,
        *,
        should_stop: StopRequested | None = None,
        stop_reason: StopReason | None = None,
    ) -> WorkflowRun:
        self.store.save_definition(definition)
        run = self.store.create_run(
            definition.id,
            initial_variables or {},
            device_id=device_id,
        )
        return self._drive(definition, run, should_stop=should_stop, stop_reason=stop_reason)

    def resume(
        self,
        run_id: str,
        *,
        should_stop: StopRequested | None = None,
        stop_reason: StopReason | None = None,
    ) -> WorkflowRun:
        run = self.store.get_run(run_id)
        if run is None:
            raise KeyError(f"unknown workflow run {run_id}")
        if run.status in TERMINAL_STATUSES:
            return run
        definition = self.store.get_definition(run.definition_id)
        if definition is None:
            raise KeyError(f"unknown workflow definition {run.definition_id}")
        return self._drive(definition, run, should_stop=should_stop, stop_reason=stop_reason)

    def _stop_status(self, stop_reason: StopReason | None) -> str:
        """Resolve the terminal status for a should_stop-triggered stop.

        No `stop_reason` preserves the pre-existing default of `cancelled` for
        any stop; a supplied reason distinguishes an explicit cancellation
        from other stop conditions (e.g. lost lease), which resolve to `failed`.
        """
        if stop_reason is None:
            return "cancelled"
        return "cancelled" if is_cancellation_reason(stop_reason()) else "failed"

    def _drive(
        self,
        definition: WorkflowDefinition,
        run: WorkflowRun,
        *,
        should_stop: StopRequested | None = None,
        stop_reason: StopReason | None = None,
    ) -> WorkflowRun:
        executed = 0
        while run.status not in TERMINAL_STATUSES and run.current_step_id:
            if should_stop is not None and should_stop():
                return self._checkpoint(
                    run, run.current_step_id, self._stop_status(stop_reason)
                )
            if self.step_limit is not None and executed >= self.step_limit:
                return run
            step = definition.step_by_id(run.current_step_id)
            recorded = _recorded_result(run, step.step_id)
            if recorded is not None:
                branch_next_step_id = None
                if isinstance(step, BranchStep):
                    branch_next_step_id = recorded.detail.get("next_step_id")
                next_status, next_step_id = self._resolve_outcome(
                    definition, step, recorded.success, branch_next_step_id
                )
                run = self._checkpoint(run, next_step_id, next_status)
                continue

            result, branch_next_step_id = self._execute_step(
                definition,
                run,
                step,
                should_stop=should_stop,
                stop_reason=stop_reason,
            )
            if should_stop is not None and should_stop():
                self.store.append_step_result(run.id, result)
                return self._checkpoint(
                    run, run.current_step_id, self._stop_status(stop_reason)
                )
            next_status, next_step_id = self._resolve_outcome(
                definition, step, result.success, branch_next_step_id
            )

            self.store.append_step_result(run.id, result)
            run = self._checkpoint(
                run,
                next_step_id,
                next_status,
            )
            executed += 1
        return run

    def _resolve_outcome(
        self,
        definition: WorkflowDefinition,
        step: WorkflowStep,
        success: bool,
        branch_next_step_id: str | None,
    ) -> tuple[str, str | None]:
        """Compute the (status, next_step_id) for a step's outcome.

        Shared by live execution and resume so a failed step's routing
        (e.g. to a branch step evaluating the failure) can never drift
        between the two code paths. The run only becomes terminal when
        there is truly no next step to route to.
        """
        next_step_id = self._next_step_id(definition, step, branch_next_step_id)
        if next_step_id is None:
            return ("completed" if success else "failed"), None
        return "running", next_step_id

    def _checkpoint(
        self,
        run: WorkflowRun,
        current_step_id: str | None,
        status: str,
    ) -> WorkflowRun:
        self.store.update_run(
            run.id,
            status=status,  # type: ignore[arg-type]
            current_step_id=current_step_id,
            variables=run.variables,
        )
        updated = self.store.get_run(run.id)
        if updated is None:
            raise KeyError(f"unknown workflow run {run.id}")
        return updated

    def _execute_step(
        self,
        definition: WorkflowDefinition,
        run: WorkflowRun,
        step: WorkflowStep,
        *,
        should_stop: StopRequested | None = None,
        stop_reason: StopReason | None = None,
    ) -> tuple[WorkflowStepResult, str | None]:
        if isinstance(step, PlannedGoalStep):
            return self._execute_planned_goal_step(
                run,
                step,
                should_stop=should_stop,
                stop_reason=stop_reason,
            ), None
        if isinstance(step, SkillInvocationStep):
            return self._execute_skill_invocation_step(step), None
        if isinstance(step, WaitForConditionStep):
            return self._execute_wait_step(
                run,
                step,
                should_stop=should_stop,
            ), None
        if isinstance(step, BranchStep):
            return self._execute_branch_step(run, step)
        return (
            WorkflowStepResult(
                step_id=getattr(step, "step_id", "unknown"),
                kind="unknown",
                success=False,
                detail={"reason": "unknown workflow step type"},
            ),
            None,
        )

    def _execute_planned_goal_step(
        self,
        run: WorkflowRun,
        step: PlannedGoalStep,
        *,
        should_stop: StopRequested | None = None,
        stop_reason: StopReason | None = None,
    ) -> WorkflowStepResult:
        task = Task(goal=step.goal, device_id=run.device_id or "")
        task_runner = self.task_runner_factory()
        if should_stop is None:
            result_task = task_runner.run(task)
        elif stop_reason is None:
            result_task = task_runner.run(task, should_stop=should_stop)
        else:
            result_task = task_runner.run(
                task, should_stop=should_stop, stop_reason=stop_reason
            )
        success = result_task.status == "completed"
        return WorkflowStepResult(
            step_id=step.step_id,
            kind=step.kind,
            success=success,
            detail={
                "task_status": result_task.status,
                "failure_reason": result_task.failure_reason,
            },
            task_id=result_task.id,
        )

    def _execute_skill_invocation_step(
        self,
        step: SkillInvocationStep,
    ) -> WorkflowStepResult:
        skill = self.skill_store.get_by_id(step.skill_id)
        if skill is None:
            return WorkflowStepResult(
                step_id=step.step_id,
                kind=step.kind,
                success=False,
                detail={"reason": f"unknown skill {step.skill_id}"},
            )
        try:
            results = run_flow_template_skill(skill, step.args, tools=self.tools)
        except SkillExecutionError as exc:
            return WorkflowStepResult(
                step_id=step.step_id,
                kind=step.kind,
                success=False,
                detail={"reason": str(exc)},
            )
        success = all(result.success for result in results)
        return WorkflowStepResult(
            step_id=step.step_id,
            kind=step.kind,
            success=success,
            detail={
                "step_results": [result.to_dict() for result in results],
            },
        )

    def _execute_wait_step(
        self,
        run: WorkflowRun,
        step: WaitForConditionStep,
        *,
        should_stop: StopRequested | None = None,
    ) -> WorkflowStepResult:
        started_at = datetime.now().astimezone()
        timeout_seconds = step.timeout_seconds
        poll_interval_seconds = step.poll_interval_seconds
        while True:
            if should_stop is not None and should_stop():
                return WorkflowStepResult(
                    step_id=step.step_id,
                    kind=step.kind,
                    success=False,
                    detail={"reason": "execution interrupted"},
                )
            try:
                if self._condition_is_true(run, step.condition, started_at=started_at):
                    return WorkflowStepResult(
                        step_id=step.step_id,
                        kind=step.kind,
                        success=True,
                        detail={"condition": step.condition.to_dict()},
                    )
            except UnknownConditionKindError as exc:
                return WorkflowStepResult(
                    step_id=step.step_id,
                    kind=step.kind,
                    success=False,
                    detail={"reason": f"unknown condition kind: {exc}"},
                )

            elapsed = (datetime.now(started_at.tzinfo) - started_at).total_seconds()
            if elapsed >= timeout_seconds:
                return WorkflowStepResult(
                    step_id=step.step_id,
                    kind=step.kind,
                    success=False,
                    detail={"reason": "condition timed out"},
                )
            self.sleep_func(poll_interval_seconds)

    def _execute_branch_step(
        self,
        run: WorkflowRun,
        step: BranchStep,
    ) -> tuple[WorkflowStepResult, str | None]:
        try:
            matched = self._condition_is_true(run, step.condition)
        except UnknownConditionKindError as exc:
            return (
                WorkflowStepResult(
                    step_id=step.step_id,
                    kind=step.kind,
                    success=False,
                    detail={"reason": f"unknown condition kind: {exc}"},
                ),
                None,
            )
        target = step.on_true if matched else step.on_false
        return (
            WorkflowStepResult(
                step_id=step.step_id,
                kind=step.kind,
                success=True,
                detail={"condition_result": matched, "next_step_id": target},
            ),
            target,
        )

    def _condition_is_true(
        self,
        run: WorkflowRun,
        condition,
        *,
        started_at: datetime | None = None,
    ) -> bool:
        return evaluate_condition(
            condition,
            scene=self.scene_provider(),
            world_state=self._world_state(run),
            started_at=started_at,
            step_results=run.step_results,
            registry=self.condition_registry,
        )

    def _world_state(self, run: WorkflowRun) -> object:
        if self.world_state_provider is not None:
            world_state = self.world_state_provider()
            if world_state is not None:
                return world_state
        return SimpleNamespace(variables=dict(run.variables))

    def _next_step_id(
        self,
        definition: WorkflowDefinition,
        step: WorkflowStep,
        branch_next_step_id: str | None,
    ) -> str | None:
        if branch_next_step_id is not None:
            return branch_next_step_id
        explicit = getattr(step, "next_step_id", None)
        if explicit is not None:
            return explicit
        if isinstance(step, BranchStep):
            return None
        return definition.next_step_id_after(step.step_id)


def _recorded_result(run: WorkflowRun, step_id: str) -> WorkflowStepResult | None:
    for result in reversed(run.step_results):
        if result.step_id == step_id:
            return result
    return None