From fde298a1b55ebce1e3f2ed5a3aef69962b1a1418 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Sat, 15 Aug 2026 17:05:01 +0700 Subject: [PATCH] Add native execution graph runtime --- .../0002_graphnoderun_visit_index.py | 18 + graph/models.py | 3 +- graph/native_runtime.py | 195 ++++++++++ graph/task_nodes.py | 353 ++++++++++++++++++ tests/test_native_graph_runtime.py | 114 ++++++ tests/test_task_execution_native_graph.py | 112 ++++++ 6 files changed, 794 insertions(+), 1 deletion(-) create mode 100644 graph/migrations/0002_graphnoderun_visit_index.py create mode 100644 graph/native_runtime.py create mode 100644 graph/task_nodes.py create mode 100644 tests/test_native_graph_runtime.py create mode 100644 tests/test_task_execution_native_graph.py diff --git a/graph/migrations/0002_graphnoderun_visit_index.py b/graph/migrations/0002_graphnoderun_visit_index.py new file mode 100644 index 0000000..8d13184 --- /dev/null +++ b/graph/migrations/0002_graphnoderun_visit_index.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("graph", "0001_initial"), + ] + + operations = [ + migrations.RemoveConstraint(model_name="graphnoderun", name="unique_graph_node_run"), + migrations.AddField(model_name="graphnoderun", name="visit_index", field=models.PositiveIntegerField(default=1)), + migrations.AddConstraint( + model_name="graphnoderun", + constraint=models.UniqueConstraint(fields=("graph_run", "node_id", "visit_index"), name="unique_graph_node_run_visit"), + ), + ] diff --git a/graph/models.py b/graph/models.py index 7edf5ad..e1212ed 100644 --- a/graph/models.py +++ b/graph/models.py @@ -79,6 +79,7 @@ class GraphNodeRun(models.Model): graph_run = models.ForeignKey(GraphRun, on_delete=models.CASCADE, related_name="node_runs") node_id = models.CharField(max_length=120) node_type = models.CharField(max_length=120) + visit_index = models.PositiveIntegerField(default=1) status = models.CharField(max_length=32, choices=GraphNodeRunStatus.choices, default=GraphNodeRunStatus.PENDING) started_at = models.DateTimeField(null=True, blank=True) completed_at = models.DateTimeField(null=True, blank=True) @@ -92,7 +93,7 @@ class GraphNodeRun(models.Model): updated_at = models.DateTimeField(auto_now=True) class Meta: - constraints = [models.UniqueConstraint(fields=["graph_run", "node_id"], name="unique_graph_node_run")] + constraints = [models.UniqueConstraint(fields=["graph_run", "node_id", "visit_index"], name="unique_graph_node_run_visit")] indexes = [models.Index(fields=["node_id", "status"])] diff --git a/graph/native_runtime.py b/graph/native_runtime.py new file mode 100644 index 0000000..8d68452 --- /dev/null +++ b/graph/native_runtime.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Any +from uuid import UUID + +from django.utils import timezone + +from control_plane.events.bus import EventBus +from graph.models import 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: + 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: + 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: + 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 diff --git a/graph/task_nodes.py b/graph/task_nodes.py new file mode 100644 index 0000000..efd0e0e --- /dev/null +++ b/graph/task_nodes.py @@ -0,0 +1,353 @@ +from __future__ import annotations + +from pathlib import Path + +from agents.coder import Coder +from agents.judge import Judge +from agents.reviewer import Reviewer +from control_plane.agents.models import AgentRole, AgentVersion +from control_plane.events.bus import EventBus +from control_plane.events.models import EventType +from control_plane.projects.models import CommitRecord, Task, TaskAttempt, TaskStatus, Worktree +from control_plane.verification.models import TestRun, VerificationResult +from graph.native_runtime import GraphExecutionContext +from graph.registry import NodeHandlerRegistry, NodeResult +from knowledge.context_builder import WorkerContextBuilder +from model_router.router import ModelRouter +from tools.capabilities import Capability +from tools.runtime import WorktreeTools +from tools.test_runner import DeterministicTestRunner +from workspace.worktrees import WorktreeManager + + +class TaskExecutionServices: + def __init__( + self, + router: ModelRouter, + *, + bus: EventBus | None = None, + test_command: list[str] | None = None, + ) -> None: + self.router = router + self.bus = bus or EventBus() + self.context_builder = WorkerContextBuilder() + self.worktrees = WorktreeManager() + self.coder = Coder(router) + self.reviewer = Reviewer() + self.judge = Judge() + self.tests = DeterministicTestRunner() + self.test_command = test_command or ["python", "-m", "pytest"] + + def champion(self, role: AgentRole) -> AgentVersion: + return AgentVersion.objects.select_related("agent").get(agent__role=role, promotion_status="CHAMPION") + + def tools(self, worktree: Worktree) -> WorktreeTools: + return WorktreeTools( + Path(worktree.worktree_path), + { + Capability.READ_REPOSITORY, + Capability.INVESTIGATE_WORKTREE, + Capability.WRITE_WORKTREE, + Capability.RUN_TESTS, + Capability.COMMIT_CHANGES, + }, + ) + + def worktree(self, task: Task) -> Worktree: + try: + return task.worktree + except Worktree.DoesNotExist: + if not task.project.repository_path: + raise RuntimeError("Task project has no repository_path") + return self.worktrees.create_for_task(task, Path(task.project.repository_path)) + + +class TaskNode: + idempotent = True + replay_safe = "checkpointed" + destructive = False + + def __init__(self, services: TaskExecutionServices, node_type: str) -> None: + self.services = services + self.node_type = node_type + + def task(self, context: GraphExecutionContext) -> Task: + if context.graph_run.task_id is None: + raise RuntimeError("Task execution graph run requires task") + return Task.objects.select_related("project", "milestone", "feature").get(id=context.graph_run.task_id) + + def metadata(self, context: GraphExecutionContext) -> dict[str, object]: + return dict(context.graph_run.metadata) + + def save_metadata(self, context: GraphExecutionContext, metadata: dict[str, object]) -> None: + context.graph_run.metadata = metadata + context.graph_run.save(update_fields=["metadata", "updated_at"]) + + +class ClaimTaskNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "claim_task") + self.idempotent = True + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + return NodeResult("COMPLETE", "success", {"task_id": str(task.id), "task_status": task.status}) + + +class PrepareWorktreeNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "prepare_worktree") + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + worktree = self.services.worktree(task) + metadata = self.metadata(context) + metadata["worktree_id"] = str(worktree.id) + metadata["worktree_path"] = worktree.worktree_path + self.save_metadata(context, metadata) + return NodeResult("COMPLETE", "success", {"worktree_id": str(worktree.id), "worktree_path": worktree.worktree_path}) + + +class BuildContextNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "build_context") + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + worktree = self.services.worktree(task) + metadata = self.metadata(context) + attempt_id = metadata.get("current_attempt_id") + if attempt_id: + attempt = TaskAttempt.objects.get(id=attempt_id) + return NodeResult("COMPLETE", "success", {"attempt_id": str(attempt.id), "attempt_number": attempt.attempt_number}) + coder_version = self.services.champion(AgentRole.CODER) + attempt = TaskAttempt.objects.create(task=task, attempt_number=task.retry_count + 1, coder=coder_version, status="RUNNING") + task_context = self.services.context_builder.build_for_task(task, Path(worktree.worktree_path)) + task_context["previous_attempts"] = list( + task.attempts.exclude(id=attempt.id).order_by("attempt_number").values("attempt_number", "status", "coder_result", "review_findings", "judge_findings") + ) + attempt.context_snapshot = task_context + attempt.save(update_fields=["context_snapshot", "updated_at"]) + metadata["current_attempt_id"] = str(attempt.id) + self.save_metadata(context, metadata) + return NodeResult("COMPLETE", "success", {"attempt_id": str(attempt.id), "attempt_number": attempt.attempt_number}) + + +class CoderNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "coder") + self.idempotent = False + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + worktree = self.services.worktree(task) + metadata = self.metadata(context) + attempt = TaskAttempt.objects.get(id=metadata["current_attempt_id"]) + coder_version = attempt.coder + result = self.services.coder.execute(attempt.context_snapshot, self.services.tools(worktree), project=task.project, agent_version=coder_version) + attempt.coder_result = {"status": result.status, "summary": result.summary, "changed_files": result.changed_files, "metadata": result.metadata} + attempt.save(update_fields=["coder_result", "updated_at"]) + metadata["last_failure_reason"] = "coder_failed" + metadata["last_failure_findings"] = [result.summary] + self.save_metadata(context, metadata) + return NodeResult("COMPLETE", "success" if result.status == "COMPLETE" else "failure", {"coder_status": result.status, "summary": result.summary}, result.metadata.get("telemetry", {}) if isinstance(result.metadata, dict) else {}) + + +class RunTestsNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "run_tests") + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + worktree = self.services.worktree(task) + metadata = self.metadata(context) + attempt = TaskAttempt.objects.get(id=metadata["current_attempt_id"]) + test_run = self.services.tests.run(task.project, task, Path(worktree.worktree_path), self.services.test_command) + metadata["test_run_id"] = str(test_run.id) + self.save_metadata(context, metadata) + if test_run.status != "PASS": + self._attach_test_failure_evidence(attempt, test_run) + self.services.bus.publish(EventType.TEST_FAILED, project=task.project, task=task, payload={"test_run_id": str(test_run.id)}) + return NodeResult("COMPLETE", "complete", {"test_run_id": str(test_run.id), "test_status": test_run.status}, {"test_status": test_run.status}) + + def _attach_test_failure_evidence(self, attempt: TaskAttempt, test_run: TestRun) -> None: + content = test_run.output_artifact.content if test_run.output_artifact else {} + stdout = str(content.get("stdout", "")) if isinstance(content, dict) else "" + stderr = str(content.get("stderr", "")) if isinstance(content, dict) else "" + coder_result = dict(attempt.coder_result or {}) + attempt_metadata = dict(coder_result.get("metadata", {})) if isinstance(coder_result.get("metadata", {}), dict) else {} + attempt_metadata["test_failure_evidence"] = {"test_run_id": str(test_run.id), "status": test_run.status, "stdout_excerpt": stdout[-12000:], "stderr_excerpt": stderr[-4000:]} + coder_result["metadata"] = attempt_metadata + attempt.coder_result = coder_result + attempt.save(update_fields=["coder_result", "updated_at"]) + + +class ReviewNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "review") + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + metadata = self.metadata(context) + attempt = TaskAttempt.objects.get(id=metadata["current_attempt_id"]) + reviewer_version = self.services.champion(AgentRole.REVIEWER) + worktree = self.services.worktree(task) + tools = self.services.tools(worktree) + tools.git(["add", "-N", "."]) + diff = tools.diff() + test_run = TestRun.objects.get(id=metadata["test_run_id"]) + review = self.services.reviewer.review(task, reviewer_version, diff, test_run.status) + attempt.review_findings = review.findings + attempt.save(update_fields=["review_findings", "updated_at"]) + metadata["review_id"] = str(review.id) + metadata["last_failure_reason"] = "review_failed" + metadata["last_failure_findings"] = review.findings + self.save_metadata(context, metadata) + if review.status != "PASS": + self.services.bus.publish(EventType.REVIEW_FAILED, project=task.project, task=task, payload={"review_id": str(review.id), "findings": review.findings}) + return NodeResult("COMPLETE", review.status, {"review_id": str(review.id), "review_status": review.status, "findings": review.findings}, {"review_status": review.status}) + + +class JudgeNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "judge") + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + metadata = self.metadata(context) + attempt = TaskAttempt.objects.get(id=metadata["current_attempt_id"]) + judge_version = self.services.champion(AgentRole.PROJECT_JUDGE) + test_run = TestRun.objects.get(id=metadata["test_run_id"]) + worktree = self.services.worktree(task) + tools = self.services.tools(worktree) + tools.git(["add", "-N", "."]) + diff = tools.diff() + verification = self.services.judge.judge(task.project, task, judge_version, diff, test_run.status) + attempt.judge_findings = verification.evidence + attempt.save(update_fields=["judge_findings", "updated_at"]) + metadata["verification_id"] = str(verification.id) + metadata["last_failure_reason"] = "judge_failed" + metadata["last_failure_findings"] = verification.evidence + self.save_metadata(context, metadata) + edge = "PASS" if verification.result == VerificationResult.PASS else "FAIL" + return NodeResult("COMPLETE", edge, {"verification_id": str(verification.id), "result": verification.result, "evidence": verification.evidence}, {"judge_result": verification.result}) + + +class RetryOrFailNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "retry_or_fail") + self.idempotent = False + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + metadata = self.metadata(context) + attempt = TaskAttempt.objects.get(id=metadata["current_attempt_id"]) + effective_max_retries = min(task.max_retries, 2) + attempt.status = "REWORK_REQUIRED" if task.retry_count < effective_max_retries else "FAILED" + attempt.save(update_fields=["status", "updated_at"]) + task.retry_count += 1 + reason = str(metadata.get("last_failure_reason", "task_failed")) + findings = metadata.get("last_failure_findings", []) + classification = self._classify_failure(reason, findings) + metadata.pop("current_attempt_id", None) + self.save_metadata(context, metadata) + if task.retry_count <= effective_max_retries: + task.status = TaskStatus.RUNNING + task.save(update_fields=["retry_count", "status", "updated_at"]) + self.services.bus.publish(EventType.TASK_FAILED, project=task.project, task=task, payload={"reason": reason, "will_retry": True, "classification": classification, "findings": findings}) + return NodeResult("COMPLETE", "retry_available", {"retry_count": task.retry_count, "classification": classification}) + task.status = TaskStatus.FAILED + task.save(update_fields=["retry_count", "status", "updated_at"]) + self.services.bus.publish("TASK_RETRY_EXHAUSTED", project=task.project, task=task, payload={"reason": reason, "classification": classification, "findings": findings}) + self.services.bus.publish(EventType.TASK_FAILED, project=task.project, task=task, payload={"reason": reason, "will_retry": False, "classification": classification, "findings": findings}) + return NodeResult("COMPLETE", "retry_exhausted", {"retry_count": task.retry_count, "classification": classification}) + + def _classify_failure(self, reason: str, findings: object) -> str: + text = f"{reason} {findings}".lower() + if "unsupported operation" in text or "missing capability" in text: + return "missing_capability" + if "context" in text or "migration" in text: + return "context_problem" + if "timeout" in text or "provider" in text: + return "environment_problem" + if "malformed json" in text or "model" in text: + return "model_problem" + if "ambiguous" in text: + return "intent_ambiguity" + if reason == "review_failed" or reason == "judge_failed": + return "replan" + return "split_task" + + +class CommitNode(TaskNode): + destructive = True + + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "commit") + self.idempotent = False + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + existing = CommitRecord.objects.filter(task=task).first() + if existing is not None: + return NodeResult("COMPLETE", "success", {"commit_id": str(existing.id), "sha": existing.sha, "deduplicated": True}) + metadata = self.metadata(context) + attempt = TaskAttempt.objects.get(id=metadata["current_attempt_id"]) + worktree = self.services.worktree(task) + test_run = TestRun.objects.get(id=metadata["test_run_id"]) + sha = self.services.tools(worktree).commit_all(f"Artifex task: {task.goal[:80]}") + commit = CommitRecord.objects.create( + project=task.project, + task=task, + worktree=worktree, + coder=attempt.coder, + reviewer=self.services.champion(AgentRole.REVIEWER), + judge=self.services.champion(AgentRole.PROJECT_JUDGE), + test_run=test_run, + review_id=metadata.get("review_id"), + verification_id=metadata.get("verification_id"), + sha=sha, + branch_name=worktree.branch_name, + message=f"Artifex task: {task.goal[:80]}", + ) + attempt.status = "COMPLETE" + attempt.save(update_fields=["status", "updated_at"]) + task.status = TaskStatus.COMPLETE + task.save(update_fields=["status", "updated_at"]) + self.services.bus.publish(EventType.COMMIT_CREATED, project=task.project, task=task, payload={"commit_id": str(commit.id), "sha": sha}) + self.services.bus.publish(EventType.TASK_COMPLETED, project=task.project, task=task, payload={"task_id": str(task.id)}) + metadata["commit_id"] = str(commit.id) + metadata["commit_sha"] = sha + self.save_metadata(context, metadata) + return NodeResult("COMPLETE", "success", {"commit_id": str(commit.id), "sha": sha}) + + +class CleanupNode(TaskNode): + def __init__(self, services: TaskExecutionServices) -> None: + super().__init__(services, "cleanup") + + def run(self, context: GraphExecutionContext) -> NodeResult: + task = self.task(context) + if task.status == TaskStatus.COMPLETE: + self.services.worktrees.validate_clean_worktree(task.worktree) + self.services.worktrees.cleanup(task.worktree) + return NodeResult("COMPLETE", "success", {"cleaned": True}) + return NodeResult("COMPLETE", "failed", {"cleaned": False}) + + +def task_execution_registry(services: TaskExecutionServices) -> NodeHandlerRegistry: + registry = NodeHandlerRegistry() + for handler in [ + ClaimTaskNode(services), + PrepareWorktreeNode(services), + BuildContextNode(services), + CoderNode(services), + RunTestsNode(services), + ReviewNode(services), + JudgeNode(services), + CommitNode(services), + RetryOrFailNode(services), + CleanupNode(services), + ]: + registry.register(handler) + return registry diff --git a/tests/test_native_graph_runtime.py b/tests/test_native_graph_runtime.py new file mode 100644 index 0000000..83b5b84 --- /dev/null +++ b/tests/test_native_graph_runtime.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +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" diff --git a/tests/test_task_execution_native_graph.py b/tests/test_task_execution_native_graph.py new file mode 100644 index 0000000..839c4b1 --- /dev/null +++ b/tests/test_task_execution_native_graph.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from pathlib import Path + +from agents.providers import DeterministicCodingProvider +from control_plane.agents.management.commands.seed_core_agents import Command as SeedAgentsCommand +from control_plane.events.models import Event, EventType +from control_plane.projects.models import CommitRecord, Project, ProjectPlan, Milestone, Task, TaskStatus +from control_plane.verification.models import Review, TestRun, Verification, VerificationResult +from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus, GraphEdgeTraversal, GraphRun, GraphRunStatus +from graph.native_runtime import NativeGraphRuntime +from graph.task_execution import task_execution_graph_v1 +from graph.task_nodes import TaskExecutionServices, task_execution_registry +from model_router.router import ModelRouter +from tests.test_m2_autonomous_loop import create_disposable_django_repo + + +def create_task(repository_path: Path, goal: str, acceptance: list[str], *, max_retries: int = 2) -> Task: + project = Project.objects.create(name=f"Graph Project {goal[:12]}", goal=goal, repository_path=str(repository_path)) + plan = ProjectPlan.objects.create(project=project, version=1, goal=goal) + milestone = Milestone.objects.create(project=project, plan=plan, key="G1", title="Graph", goal="Execution graph") + return Task.objects.create( + project=project, + milestone=milestone, + task_type="implementation", + status=TaskStatus.RUNNING, + goal=goal, + acceptance_criteria=acceptance, + max_retries=max_retries, + ) + + +def graph_run_for_task(task: Task) -> GraphRun: + spec = task_execution_graph_v1() + definition, _ = ExecutionGraphDefinition.objects.get_or_create( + name=spec.name, + defaults={"graph_type": spec.graph_type, "description": "Task execution graph"}, + ) + version, _ = ExecutionGraphVersion.objects.get_or_create( + graph=definition, + version=spec.version, + defaults={"status": ExecutionGraphVersionStatus.CHAMPION, "graph_spec": spec.to_dict()}, + ) + return GraphRun.objects.create( + execution_graph_version=version, + project=task.project, + milestone=task.milestone, + feature=task.feature, + task=task, + current_node=spec.entry, + ) + + +def run_graph(task: Task, *, interrupt_after: str | None = None) -> GraphRun: + SeedAgentsCommand().handle() + services = TaskExecutionServices(ModelRouter({"qwen": DeterministicCodingProvider()}), test_command=["python", "manage.py", "test"]) + runtime = NativeGraphRuntime(task_execution_registry(services)) + return runtime.run_until_terminal_or_paused(graph_run_for_task(task), interrupt_after=interrupt_after) + + +def test_native_task_execution_graph_success_matches_loop_semantics(tmp_path: Path) -> None: + repo = create_disposable_django_repo(tmp_path) + task = create_task(repo, "Add a /health endpoint returning JSON {\"status\": \"ok\"} and add tests.", ["/health returns JSON ok", "tests pass"]) + + graph_run = run_graph(task) + + task.refresh_from_db() + assert graph_run.status == GraphRunStatus.COMPLETE + assert task.status == TaskStatus.COMPLETE + assert CommitRecord.objects.filter(task=task).count() == 1 + assert TestRun.objects.get(task=task).status == "PASS" + assert Review.objects.get(task=task).status == "PASS" + assert Verification.objects.get(task=task).result == VerificationResult.PASS + assert task.worktree.status == "CLEANED" + assert GraphEdgeTraversal.objects.filter(graph_run=graph_run, source_node="judge", target_node="commit", condition="PASS").exists() + + +def test_native_task_execution_graph_retry_exhaustion_matches_loop_semantics(tmp_path: Path) -> None: + repo = create_disposable_django_repo(tmp_path) + task = create_task(repo, "FORCE_BAD_IMPLEMENTATION Add a /health endpoint returning JSON {\"status\": \"ok\"} and add tests.", ["/health returns JSON ok", "tests pass"], max_retries=1) + + graph_run = run_graph(task) + + task.refresh_from_db() + assert graph_run.status == GraphRunStatus.FAILED + assert task.status == TaskStatus.FAILED + assert task.retry_count == 2 + assert CommitRecord.objects.filter(task=task).count() == 0 + assert Event.objects.filter(task=task, event_type=EventType.TASK_FAILED).exists() + assert Event.objects.filter(task=task, event_type="TASK_RETRY_EXHAUSTED").exists() + + +def test_native_task_execution_graph_resume_after_coder_prevents_duplicate_commit(tmp_path: Path) -> None: + repo = create_disposable_django_repo(tmp_path) + task = create_task(repo, "Add a /health endpoint returning JSON {\"status\": \"ok\"} and add tests.", ["/health returns JSON ok", "tests pass"]) + SeedAgentsCommand().handle() + graph_run = graph_run_for_task(task) + services = TaskExecutionServices(ModelRouter({"qwen": DeterministicCodingProvider()}), test_command=["python", "manage.py", "test"]) + runtime = NativeGraphRuntime(task_execution_registry(services)) + + runtime.run_until_terminal_or_paused(graph_run, interrupt_after="coder") + graph_run.refresh_from_db() + assert graph_run.current_node == "coder" + assert task.attempts.count() == 1 + + runtime.run_until_terminal_or_paused(graph_run) + runtime.run_until_terminal_or_paused(graph_run) + + task.refresh_from_db() + assert task.status == TaskStatus.COMPLETE + assert CommitRecord.objects.filter(task=task).count() == 1 + assert graph_run.node_runs.filter(node_id="coder").count() == 1