agentic-mobile-control/tests/test_workflow_runner.py

from __future__ import annotations

from core.models import Bounds, Scene, SceneElement, Task
from runtime.executor import Executor, ExecutorConfig
from runtime.planner import Planner
from runtime.task import TaskRunner, TaskRunnerConfig
from skills_learning.models import FlowStep, FlowTemplateSkill, SkillMetadata
from skills_learning.store import SkillStore
from tests.fakes import PNG_10X20
from workflow.models import (
    BranchStep,
    ConditionSpec,
    PlannedGoalStep,
    SkillInvocationStep,
    WaitForConditionStep,
    WorkflowDefinition,
    WorkflowStepResult,
)
from workflow.runner import WorkflowRunner
from workflow.store import WorkflowStore


class FakeTaskRunner:
    def __init__(
        self,
        calls: list[str],
        *,
        status: str = "completed",
        failure_reason: str | None = None,
    ) -> None:
        self.calls = calls
        self.status = status
        self.failure_reason = failure_reason

    def run(self, task: Task) -> Task:
        self.calls.append(task.goal)
        task.status = self.status  # type: ignore[assignment]
        task.failure_reason = self.failure_reason
        return task


def _scene(text: str = "Ready") -> Scene:
    return Scene(
        width=10,
        height=20,
        elements=[
            SceneElement(
                id="label",
                type="text",
                text=text,
                bounds=Bounds(1, 2, 3, 4),
            )
        ],
    )


def _skill_store() -> tuple[SkillStore, FlowTemplateSkill]:
    store = SkillStore()
    skill = store.create_version(
        FlowTemplateSkill(
            metadata=SkillMetadata(
                name="send message",
                description="Send message",
                kind="flow_template",
            ),
            steps=[FlowStep("input_text", {"text": "{message}"})],
            parameters={"message": {"type": "string"}},
        )
    )
    return store, skill


def _store(tmp_path) -> WorkflowStore:
    return WorkflowStore(tmp_path / "workflows.sqlite3")


def test_workflow_runner_linear_planned_goal_completes(tmp_path) -> None:
    calls: list[str] = []
    definition = WorkflowDefinition(
        name="linear",
        entry_step_id="first",
        steps=[
            PlannedGoalStep("first", "open app", next_step_id="second"),
            PlannedGoalStep("second", "send message"),
        ],
    )
    runner = WorkflowRunner(
        _store(tmp_path),
        task_runner_factory=lambda: FakeTaskRunner(calls),
    )

    run = runner.run(definition, "phone")

    assert run.status == "completed"
    assert calls == ["open app", "send message"]
    assert [result.step_id for result in run.step_results] == ["first", "second"]


def test_workflow_runner_failing_planned_goal_marks_run_failed(tmp_path) -> None:
    definition = WorkflowDefinition(
        name="fail",
        entry_step_id="first",
        steps=[PlannedGoalStep("first", "open app")],
    )
    runner = WorkflowRunner(
        _store(tmp_path),
        task_runner_factory=lambda: FakeTaskRunner(
            [],
            status="failed",
            failure_reason="device offline",
        ),
    )

    run = runner.run(definition, "phone")

    assert run.status == "failed"
    assert run.step_results[0].detail["failure_reason"] == "device offline"


def test_workflow_runner_skill_invocation_executes_resolved_tool_calls(tmp_path) -> None:
    skill_store, skill = _skill_store()
    calls: list[dict[str, object]] = []
    definition = WorkflowDefinition(
        name="skill",
        entry_step_id="skill",
        steps=[SkillInvocationStep("skill", skill.id, {"message": "hello"})],
    )
    runner = WorkflowRunner(
        _store(tmp_path),
        skill_store=skill_store,
        tools={"input_text": lambda **kwargs: calls.append(kwargs) or {"ok": True}},
    )

    run = runner.run(definition, "phone")

    assert run.status == "completed"
    assert calls == [{"text": "hello"}]


def test_workflow_runner_wait_step_succeeds_and_times_out(tmp_path) -> None:
    success_definition = WorkflowDefinition(
        name="wait success",
        entry_step_id="wait",
        steps=[
            WaitForConditionStep(
                "wait",
                ConditionSpec("scene_contains_text", {"text": "ready"}),
                timeout_seconds=0.1,
                poll_interval_seconds=0,
            )
        ],
    )
    success_runner = WorkflowRunner(
        _store(tmp_path / "success"),
        scene_provider=lambda: _scene("Ready"),
        sleep_func=lambda seconds: None,
    )

    success_run = success_runner.run(success_definition, "phone")

    assert success_run.status == "completed"

    timeout_definition = WorkflowDefinition(
        name="wait timeout",
        entry_step_id="wait",
        steps=[
            WaitForConditionStep(
                "wait",
                ConditionSpec("scene_contains_text", {"text": "missing"}),
                timeout_seconds=0,
                poll_interval_seconds=0,
            )
        ],
    )
    timeout_runner = WorkflowRunner(
        _store(tmp_path / "timeout"),
        scene_provider=lambda: _scene("Ready"),
        sleep_func=lambda seconds: None,
    )

    timeout_run = timeout_runner.run(timeout_definition, "phone")

    assert timeout_run.status == "failed"
    assert "timed out" in timeout_run.step_results[0].detail["reason"]


def test_workflow_runner_branch_paths_and_sequential_advancement(tmp_path) -> None:
    true_calls: list[str] = []
    definition = WorkflowDefinition(
        name="branch",
        entry_step_id="branch",
        steps=[
            BranchStep(
                "branch",
                ConditionSpec("world_variable_equals", {"name": "ready", "value": True}),
                on_true="true-step",
                on_false="false-step",
            ),
            PlannedGoalStep("true-step", "true path", next_step_id="end"),
            PlannedGoalStep("false-step", "false path", next_step_id="end"),
            WaitForConditionStep(
                "end",
                ConditionSpec("elapsed_seconds", {"seconds": 0}),
                timeout_seconds=0,
                poll_interval_seconds=0,
            ),
        ],
    )
    true_runner = WorkflowRunner(
        _store(tmp_path / "true"),
        task_runner_factory=lambda: FakeTaskRunner(true_calls),
    )

    true_run = true_runner.run(definition, "phone", {"ready": True})

    assert true_run.status == "completed"
    assert true_calls == ["true path"]
    assert [result.step_id for result in true_run.step_results] == [
        "branch",
        "true-step",
        "end",
    ]

    false_calls: list[str] = []
    false_runner = WorkflowRunner(
        _store(tmp_path / "false"),
        task_runner_factory=lambda: FakeTaskRunner(false_calls),
    )

    false_run = false_runner.run(definition, "phone", {"ready": False})

    assert false_run.status == "completed"
    assert false_calls == ["false path"]


def test_workflow_runner_resume_skips_completed_steps_after_restart(tmp_path) -> None:
    calls: list[str] = []
    db_path = tmp_path / "workflows.sqlite3"
    definition = WorkflowDefinition(
        name="resume",
        entry_step_id="one",
        steps=[
            PlannedGoalStep("one", "one"),
            PlannedGoalStep("two", "two"),
            PlannedGoalStep("three", "three"),
        ],
    )
    first_runner = WorkflowRunner(
        WorkflowStore(db_path),
        task_runner_factory=lambda: FakeTaskRunner(calls),
        step_limit=2,
    )

    partial = first_runner.run(definition, "phone")

    assert partial.status == "running"
    assert partial.current_step_id == "three"
    assert calls == ["one", "two"]

    resumed_runner = WorkflowRunner(
        WorkflowStore(db_path),
        task_runner_factory=lambda: FakeTaskRunner(calls),
    )
    resumed = resumed_runner.resume(partial.id)

    assert resumed.status == "completed"
    assert calls == ["one", "two", "three"]


def test_workflow_runner_resume_after_branch_crash_advances_to_recorded_target(
    tmp_path,
) -> None:
    """Simulates a crash between append_step_result() and update_run() for a
    BranchStep: the step's result is persisted but the checkpoint never
    happened, so current_step_id still points at the branch step itself.
    Resuming must recompute the next step from the *recorded* branch
    target instead of losing it (and must not loop forever re-processing
    the branch step).
    """
    calls: list[str] = []
    store = _store(tmp_path)
    definition = WorkflowDefinition(
        name="branch-resume",
        entry_step_id="branch",
        steps=[
            BranchStep(
                "branch",
                ConditionSpec("world_variable_equals", {"name": "ready", "value": True}),
                on_true="true-step",
                on_false="false-step",
            ),
            PlannedGoalStep("false-step", "false path"),
            PlannedGoalStep("true-step", "true path"),
        ],
    )
    store.save_definition(definition)
    run = store.create_run(definition.id, {"ready": True}, device_id="phone")
    store.append_step_result(
        run.id,
        WorkflowStepResult(
            step_id="branch",
            kind="branch",
            success=True,
            detail={"condition_result": True, "next_step_id": "true-step"},
        ),
    )

    runner = WorkflowRunner(
        store,
        task_runner_factory=lambda: FakeTaskRunner(calls),
    )
    resumed = runner.resume(run.id)

    assert resumed.status == "completed"
    assert calls == ["true path"]
    assert [result.step_id for result in resumed.step_results] == [
        "branch",
        "true-step",
    ]


def test_workflow_runner_resume_after_failed_step_crash_ends_failed(tmp_path) -> None:
    """Simulates a crash between append_step_result(success=False) and
    update_run(): the failed result is persisted but current_step_id still
    points at the failed step. Resuming must recognize the recorded
    failure and end the run as failed rather than silently continuing (or
    looping forever re-processing the same step).
    """
    store = _store(tmp_path)
    definition = WorkflowDefinition(
        name="fail-resume",
        entry_step_id="only",
        steps=[PlannedGoalStep("only", "risky")],
    )
    store.save_definition(definition)
    run = store.create_run(definition.id, {}, device_id="phone")
    store.append_step_result(
        run.id,
        WorkflowStepResult(
            step_id="only",
            kind="planned_goal",
            success=False,
            detail={"task_status": "failed", "failure_reason": "boom"},
        ),
    )

    runner = WorkflowRunner(store)
    resumed = runner.resume(run.id)

    assert resumed.status == "failed"
    assert resumed.current_step_id is None


def test_workflow_runner_failed_step_routes_to_recovery_branch(tmp_path) -> None:
    """A planned-goal step that fails still routes to its declared
    next_step_id (a branch step evaluating step_result_success). Live
    execution must reach and execute the on_false recovery path rather
    than immediately terminating the run as failed.
    """
    calls: list[str] = []

    class ScriptedTaskRunner:
        def __init__(self, calls: list[str], failing_goals: set[str]) -> None:
            self.calls = calls
            self.failing_goals = failing_goals

        def run(self, task: Task) -> Task:
            self.calls.append(task.goal)
            if task.goal in self.failing_goals:
                task.status = "failed"  # type: ignore[assignment]
                task.failure_reason = "boom"
            else:
                task.status = "completed"  # type: ignore[assignment]
                task.failure_reason = None
            return task

    definition = WorkflowDefinition(
        name="failure-recovery",
        entry_step_id="goal",
        steps=[
            PlannedGoalStep("goal", "risky goal", next_step_id="check"),
            BranchStep(
                "check",
                ConditionSpec("step_result_success", {"step_id": "goal"}),
                on_true="happy",
                on_false="recover",
            ),
            PlannedGoalStep("happy", "happy path"),
            PlannedGoalStep("recover", "recover path"),
        ],
    )
    runner = WorkflowRunner(
        _store(tmp_path),
        task_runner_factory=lambda: ScriptedTaskRunner(calls, {"risky goal"}),
    )

    run = runner.run(definition, "phone")

    assert run.status == "completed"
    assert calls == ["risky goal", "recover path"]
    assert [result.step_id for result in run.step_results] == [
        "goal",
        "check",
        "recover",
    ]
    assert run.step_results[0].success is False


def test_workflow_run_persists_variables_and_ordered_results(tmp_path) -> None:
    store = _store(tmp_path)
    definition = WorkflowDefinition(
        name="persist",
        entry_step_id="one",
        steps=[PlannedGoalStep("one", "one"), PlannedGoalStep("two", "two")],
    )
    run = WorkflowRunner(
        store,
        task_runner_factory=lambda: FakeTaskRunner([]),
    ).run(definition, "phone", {"contact": "Zhang San"})

    loaded = store.get_run(run.id)

    assert loaded is not None
    assert loaded.variables == {"contact": "Zhang San"}
    assert [result.step_id for result in loaded.step_results] == ["one", "two"]


def test_workflow_runner_branch_wait_and_skill_combination(tmp_path) -> None:
    skill_store, skill = _skill_store()
    calls: list[dict[str, object]] = []
    definition = WorkflowDefinition(
        name="combo",
        entry_step_id="branch",
        steps=[
            BranchStep(
                "branch",
                ConditionSpec("world_variable_equals", {"name": "ready", "value": True}),
                on_true="wait",
                on_false="skip",
            ),
            WaitForConditionStep(
                "wait",
                ConditionSpec("elapsed_seconds", {"seconds": 0}),
                timeout_seconds=1,
                poll_interval_seconds=0,
                next_step_id="skill",
            ),
            SkillInvocationStep(
                "skill",
                skill.id,
                {"message": "hello"},
                next_step_id="end",
            ),
            PlannedGoalStep("skip", "skip", next_step_id="end"),
            WaitForConditionStep(
                "end",
                ConditionSpec("elapsed_seconds", {"seconds": 0}),
                timeout_seconds=0,
                poll_interval_seconds=0,
            ),
        ],
    )
    runner = WorkflowRunner(
        _store(tmp_path),
        skill_store=skill_store,
        tools={"input_text": lambda **kwargs: calls.append(kwargs) or {"ok": True}},
    )

    run = runner.run(definition, "phone", {"ready": True})

    assert run.status == "completed"
    assert [result.step_id for result in run.step_results] == [
        "branch",
        "wait",
        "skill",
        "end",
    ]
    assert calls == [{"text": "hello"}]


def test_resume_completed_run_is_noop(tmp_path) -> None:
    calls: list[str] = []
    runner = WorkflowRunner(
        _store(tmp_path),
        task_runner_factory=lambda: FakeTaskRunner(calls),
    )
    definition = WorkflowDefinition(
        name="done",
        entry_step_id="one",
        steps=[PlannedGoalStep("one", "one")],
    )
    completed = runner.run(definition, "phone")

    resumed = runner.resume(completed.id)

    assert resumed == completed
    assert calls == ["one"]


def test_workflow_runner_composes_real_task_runner_contract(tmp_path) -> None:
    scene = _scene("Ready")

    def task_runner_factory() -> TaskRunner:
        return TaskRunner(
            planner=Planner(),
            executor=Executor(
                tools={"describe_screen": lambda **kwargs: scene},
                config=ExecutorConfig(max_retries=1, backoff_seconds=0),
            ),
            config=TaskRunnerConfig(max_steps=3),
            observer=lambda device_id: scene,
            screenshot_provider=lambda device_id: PNG_10X20,
        )

    definition = WorkflowDefinition(
        name="real task runner",
        entry_step_id="goal",
        steps=[PlannedGoalStep("goal", "inspect")],
    )

    run = WorkflowRunner(
        _store(tmp_path),
        task_runner_factory=task_runner_factory,
    ).run(definition, "phone")

    assert run.status == "completed"
    assert run.step_results[0].task_id is not None