Artifex/tests/test_native_graph_runtime.py

from __future__ import annotations

from control_plane.agents.models import ProgenySignal
from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus, GraphEdgeTraversal, GraphRun, GraphRunStatus
from graph.native_runtime import NativeGraphRuntime
from graph.registry import NodeHandlerRegistry, NodeResult
from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec


class FixedNode:
    idempotent = True
    replay_safe = True
    destructive = False

    def __init__(self, node_type: str, edge_result: str = "success") -> None:
        self.node_type = node_type
        self.edge_result = edge_result
        self.calls = 0

    def run(self, context: object) -> NodeResult:
        self.calls += 1
        return NodeResult("COMPLETE", self.edge_result, {"calls": self.calls}, {"model_requests": 1})


def persist_spec(spec: ExecutionGraphSpec) -> GraphRun:
    definition = ExecutionGraphDefinition.objects.create(name=spec.name, graph_type=spec.graph_type)
    version = ExecutionGraphVersion.objects.create(graph=definition, version=spec.version, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
    return GraphRun.objects.create(execution_graph_version=version, current_node=spec.entry)


def test_native_graph_runtime_executes_conditional_edges_and_persists_history() -> None:
    spec = ExecutionGraphSpec(
        name="fixture_graph",
        version=1,
        graph_type="FIXTURE",
        entry="start",
        nodes={
            "start": GraphNodeSpec("start", "start"),
            "success": GraphNodeSpec("success", "success"),
            "fail": GraphNodeSpec("fail", "fail"),
        },
        edges=[GraphEdgeSpec("start", "success", "ok"), GraphEdgeSpec("start", "fail", "bad")],
        terminal_nodes=["success", "fail"],
    )
    graph_run = persist_spec(spec)
    registry = NodeHandlerRegistry()
    start = FixedNode("start", "ok")
    registry.register(start)

    result = NativeGraphRuntime(registry).run_until_terminal_or_paused(graph_run)

    assert result.status == GraphRunStatus.COMPLETE
    assert result.current_node == "success"
    assert start.calls == 1
    assert list(GraphEdgeTraversal.objects.filter(graph_run=graph_run).values_list("source_node", "target_node", "condition")) == [("start", "success", "ok")]
    assert graph_run.node_runs.get(node_id="start").telemetry["duration_ms"] >= 0


def test_native_graph_runtime_resume_does_not_repeat_completed_current_node() -> None:
    spec = ExecutionGraphSpec(
        name="resume_graph",
        version=1,
        graph_type="FIXTURE",
        entry="first",
        nodes={
            "first": GraphNodeSpec("first", "first"),
            "second": GraphNodeSpec("second", "second"),
            "done": GraphNodeSpec("done", "done"),
        },
        edges=[GraphEdgeSpec("first", "second", "success"), GraphEdgeSpec("second", "done", "success")],
        terminal_nodes=["done"],
    )
    graph_run = persist_spec(spec)
    registry = NodeHandlerRegistry()
    first = FixedNode("first")
    second = FixedNode("second")
    registry.register(first)
    registry.register(second)
    runtime = NativeGraphRuntime(registry)

    runtime.run_until_terminal_or_paused(graph_run, interrupt_after="first")
    graph_run.refresh_from_db()
    assert graph_run.current_node == "first"

    runtime.run_until_terminal_or_paused(graph_run)

    graph_run.refresh_from_db()
    assert graph_run.status == GraphRunStatus.COMPLETE
    assert first.calls == 1
    assert second.calls == 1


def test_native_graph_runtime_records_paused_runs() -> None:
    class PauseNode(FixedNode):
        def run(self, context: object) -> NodeResult:
            self.calls += 1
            return NodeResult("PAUSED", "awaiting", pause_reason="AWAITING_APPROVAL")

    spec = ExecutionGraphSpec(
        name="pause_graph",
        version=1,
        graph_type="FIXTURE",
        entry="approval",
        nodes={"approval": GraphNodeSpec("approval", "approval"), "done": GraphNodeSpec("done", "done")},
        edges=[GraphEdgeSpec("approval", "done", "approved")],
        terminal_nodes=["done"],
    )
    graph_run = persist_spec(spec)
    registry = NodeHandlerRegistry()
    registry.register(PauseNode("approval"))

    result = NativeGraphRuntime(registry).run_until_terminal_or_paused(graph_run)

    assert result.status == GraphRunStatus.PAUSED
    assert result.failure_reason == "AWAITING_APPROVAL"


def test_native_graph_runtime_signal_created_for_missing_edge() -> None:
    spec = ExecutionGraphSpec(
        name="missing_edge_graph",
        version=1,
        graph_type="FIXTURE",
        entry="start",
        nodes={"start": GraphNodeSpec("start", "start"), "done": GraphNodeSpec("done", "done")},
        edges=[GraphEdgeSpec("start", "done", "expected")],
        terminal_nodes=["done"],
    )
    graph_run = persist_spec(spec)
    registry = NodeHandlerRegistry()
    registry.register(FixedNode("start", "unexpected"))

    result = NativeGraphRuntime(registry).run_until_terminal_or_paused(graph_run)

    assert result.status == GraphRunStatus.FAILED
    signal = ProgenySignal.objects.get(source="graph_runtime", failure_category="GRAPH_RUNTIME_ERROR")
    assert signal.graph_run == graph_run
    assert signal.graph_node_run.node_id == "start"
    assert signal.evidence["edge_result"] == "unexpected"