Artifex/graph/native_runtime.py

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any
from uuid import UUID

from asgiref.sync import sync_to_async
from django.utils import timezone

from control_plane.events.bus import EventBus
from graph.models import GraphApprovalStatus, GraphEdgeTraversal, GraphNodeRun, GraphNodeRunStatus, GraphRun, GraphRunStatus
from graph.registry import NodeHandlerRegistry, NodeResult
from graph.runtime import GraphRuntime
from graph.spec import ExecutionGraphSpec, outgoing_edges


@dataclass
class GraphExecutionContext:
    graph_run: GraphRun
    spec: ExecutionGraphSpec
    metadata: dict[str, Any] = field(default_factory=dict)
    signal_payload: dict[str, Any] = field(default_factory=dict)


class NativeGraphRuntime(GraphRuntime):
    def __init__(self, registry: NodeHandlerRegistry, bus: EventBus | None = None) -> None:
        self.registry = registry
        self.bus = bus or EventBus()

    async def start(self, project_id: UUID | None = None, *, graph_run: GraphRun | None = None) -> str:
        if graph_run is None:
            raise ValueError("NativeGraphRuntime.start requires a persisted GraphRun")
        if graph_run.status == GraphRunStatus.PENDING:
            graph_run.status = GraphRunStatus.RUNNING
            graph_run.started_at = timezone.now()
            graph_run.current_node = graph_run.current_node or graph_run.execution_graph_version.graph_spec["entry"]
            graph_run.save(update_fields=["status", "started_at", "current_node", "updated_at"])
            self.bus.publish("GRAPH_RUN_STARTED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id})
        return str(graph_run.id)

    async def pause(self, run_id: str) -> None:
        graph_run = GraphRun.objects.get(id=run_id)
        graph_run.status = GraphRunStatus.PAUSED
        graph_run.save(update_fields=["status", "updated_at"])
        self.bus.publish("GRAPH_RUN_PAUSED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id})

    async def resume(self, run_id: str) -> None:
        graph_run = GraphRun.objects.get(id=run_id)
        if graph_run.status == GraphRunStatus.PAUSED:
            graph_run.status = GraphRunStatus.RUNNING
            graph_run.save(update_fields=["status", "updated_at"])
            self.bus.publish("GRAPH_RUN_RESUMED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id})
        self.run_until_terminal_or_paused(graph_run)

    async def cancel(self, run_id: str) -> None:
        graph_run = GraphRun.objects.get(id=run_id)
        graph_run.status = GraphRunStatus.CANCELLED
        graph_run.completed_at = timezone.now()
        graph_run.save(update_fields=["status", "completed_at", "updated_at"])
        self.bus.publish("GRAPH_RUN_CANCELLED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id})

    async def signal(self, run_id: str, event: dict[str, Any]) -> None:
        await sync_to_async(self.signal_now)(run_id, event)

    def signal_now(self, run_id: str, event: dict[str, Any]) -> None:
        graph_run = GraphRun.objects.get(id=run_id)
        metadata = dict(graph_run.metadata)
        metadata["last_signal"] = self._bounded(event)
        if event.get("action") == "approve" and graph_run.status == GraphRunStatus.PAUSED:
            for approval in graph_run.approvals.filter(status=GraphApprovalStatus.PENDING):
                approval.status = GraphApprovalStatus.APPROVED
                approval.decided_by = str(event.get("actor", "system"))
                approval.decided_at = timezone.now()
                approval.save(update_fields=["status", "decided_by", "decided_at", "updated_at"])
            graph_run.status = GraphRunStatus.RUNNING
        graph_run.metadata = metadata
        graph_run.save(update_fields=["metadata", "status", "updated_at"])
        self.bus.publish("GRAPH_RUN_SIGNALED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id, "signal": self._bounded(event)})

    def run_until_terminal_or_paused(self, graph_run: GraphRun, *, interrupt_after: str | None = None) -> GraphRun:
        graph_run.refresh_from_db()
        spec = ExecutionGraphSpec.from_dict(graph_run.execution_graph_version.graph_spec)
        spec.validate()
        if graph_run.status == GraphRunStatus.PENDING:
            graph_run.status = GraphRunStatus.RUNNING
            graph_run.started_at = timezone.now()
            graph_run.current_node = graph_run.current_node or spec.entry
            graph_run.save(update_fields=["status", "started_at", "current_node", "updated_at"])
            self.bus.publish("GRAPH_RUN_STARTED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id})
        while graph_run.status == GraphRunStatus.RUNNING:
            node_id = graph_run.current_node or spec.entry
            if node_id in spec.terminal_nodes:
                self._finish_terminal(graph_run, node_id)
                break
            node_spec = spec.nodes[node_id]
            context = GraphExecutionContext(graph_run=graph_run, spec=spec, metadata=dict(graph_run.metadata))
            node_run, created = self._node_run(graph_run, node_id, node_spec.node_type)
            if not created and node_run.status == GraphNodeRunStatus.COMPLETE:
                result = NodeResult("COMPLETE", str(node_run.output_metadata.get("edge_result", "success")), node_run.output_metadata, node_run.telemetry)
            else:
                result = self._execute_node(context, node_run)
            if interrupt_after == node_id:
                graph_run.metadata = {**dict(graph_run.metadata), "interrupted_after_node": node_id}
                graph_run.save(update_fields=["metadata", "updated_at"])
                graph_run.refresh_from_db()
                return graph_run
            if result.status == "PAUSED":
                graph_run.status = GraphRunStatus.PAUSED
                graph_run.failure_reason = result.pause_reason
                graph_run.save(update_fields=["status", "failure_reason", "updated_at"])
                self.bus.publish("GRAPH_RUN_PAUSED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id, "reason": result.pause_reason})
                break
            if result.status == "FAILED":
                graph_run.status = GraphRunStatus.FAILED
                graph_run.completed_at = timezone.now()
                graph_run.failure_reason = str((result.failure_evidence or {}).get("reason", "node failed"))
                graph_run.save(update_fields=["status", "completed_at", "failure_reason", "updated_at"])
                self.bus.publish("GRAPH_RUN_FAILED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id, "node_id": node_id})
                break
            next_node = self._select_next(spec, node_id, result.edge_result)
            if next_node is None:
                graph_run.status = GraphRunStatus.FAILED
                graph_run.completed_at = timezone.now()
                graph_run.failure_reason = f"No edge from {node_id} for {result.edge_result}"
                graph_run.save(update_fields=["status", "completed_at", "failure_reason", "updated_at"])
                self.bus.publish("GRAPH_RUN_FAILED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id, "reason": graph_run.failure_reason})
                break
            GraphEdgeTraversal.objects.create(graph_run=graph_run, source_node=node_id, target_node=next_node, condition=result.edge_result, result=result.edge_result)
            self.bus.publish("GRAPH_EDGE_TRAVERSED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id, "source": node_id, "target": next_node, "condition": result.edge_result})
            graph_run.current_node = next_node
            graph_run.metadata = self._merge_summary(graph_run.metadata, result)
            graph_run.metadata["last_node_id"] = node_id
            graph_run.metadata["last_edge_result"] = result.edge_result
            graph_run.metadata["last_output"] = self._bounded(result.output_metadata or {})
            graph_run.metadata.pop("interrupted_after_node", None)
            graph_run.save(update_fields=["current_node", "metadata", "updated_at"])
        graph_run.refresh_from_db()
        return graph_run

    def _node_run(self, graph_run: GraphRun, node_id: str, node_type: str) -> tuple[GraphNodeRun, bool]:
        last = graph_run.node_runs.filter(node_id=node_id).order_by("-visit_index").first()
        if last is not None and last.status in {GraphNodeRunStatus.PENDING, GraphNodeRunStatus.RUNNING, GraphNodeRunStatus.PAUSED}:
            return last, False
        if last is not None and last.status == GraphNodeRunStatus.COMPLETE and graph_run.current_node == node_id and graph_run.metadata.get("interrupted_after_node") == node_id:
            return last, False
        visit_index = 1 if last is None else last.visit_index + 1
        return GraphNodeRun.objects.create(graph_run=graph_run, node_id=node_id, node_type=node_type, visit_index=visit_index), True

    def _execute_node(self, context: GraphExecutionContext, node_run: GraphNodeRun) -> NodeResult:
        handler = self.registry.get(node_run.node_type)
        started = time.monotonic()
        node_run.status = GraphNodeRunStatus.RUNNING
        node_run.started_at = timezone.now()
        node_run.save(update_fields=["status", "started_at", "updated_at"])
        self.bus.publish("GRAPH_NODE_STARTED", project=context.graph_run.project, task=context.graph_run.task, payload={"graph_run_id": context.graph_run.id, "node_id": node_run.node_id})
        try:
            result = handler.run(context)
        except Exception as exc:
            result = NodeResult("FAILED", "failure", failure_evidence={"reason": str(exc)})
        node_run.status = result.status if result.status in GraphNodeRunStatus.values else GraphNodeRunStatus.FAILED
        node_run.completed_at = timezone.now()
        output = dict(result.output_metadata or {})
        output["edge_result"] = result.edge_result
        node_run.output_metadata = self._bounded(output)
        telemetry = dict(result.telemetry or {})
        telemetry["duration_ms"] = int((time.monotonic() - started) * 1000)
        node_run.telemetry = self._bounded(telemetry)
        node_run.failure_evidence = self._bounded(result.failure_evidence or {})
        node_run.save(update_fields=["status", "completed_at", "output_metadata", "telemetry", "failure_evidence", "updated_at"])
        event = "GRAPH_NODE_COMPLETED" if result.status == "COMPLETE" else "GRAPH_NODE_FAILED"
        self.bus.publish(event, project=context.graph_run.project, task=context.graph_run.task, payload={"graph_run_id": context.graph_run.id, "node_id": node_run.node_id, "status": result.status})
        return result

    def _select_next(self, spec: ExecutionGraphSpec, node_id: str, edge_result: str) -> str | None:
        edges = outgoing_edges(spec, node_id)
        for edge in edges:
            if edge.condition == edge_result:
                return edge.target
        for edge in edges:
            if not edge.condition:
                return edge.target
        return None

    def _finish_terminal(self, graph_run: GraphRun, node_id: str) -> None:
        graph_run.status = GraphRunStatus.FAILED if node_id == "fail" else GraphRunStatus.COMPLETE
        graph_run.completed_at = timezone.now()
        graph_run.save(update_fields=["status", "completed_at", "updated_at"])
        event = "GRAPH_RUN_COMPLETED" if graph_run.status == GraphRunStatus.COMPLETE else "GRAPH_RUN_FAILED"
        self.bus.publish(event, project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id, "terminal_node": node_id})

    def _merge_summary(self, metadata: dict[str, Any], result: NodeResult) -> dict[str, Any]:
        summary = dict(metadata)
        telemetry = dict(summary.get("telemetry", {}))
        for key, value in (result.telemetry or {}).items():
            if isinstance(value, int | float):
                telemetry[key] = telemetry.get(key, 0) + value
        summary["telemetry"] = telemetry
        return summary

    def _bounded(self, value: Any) -> Any:
        text = str(value)
        if len(text) > 20000:
            return {"truncated": True, "excerpt": text[:20000]}
        return value