Add native execution graph runtime
This commit is contained in:
Daniel Maddern
parent
acafad9f12
commit
fde298a1b5
6 changed files with 794 additions and 1 deletions
18
graph/migrations/0002_graphnoderun_visit_index.py
Normal file
18
graph/migrations/0002_graphnoderun_visit_index.py
Normal file
|
|
@ -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"),
|
||||
),
|
||||
]
|
||||
|
|
@ -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"])]
|
||||
|
||||
|
||||
|
|
|
|||
195
graph/native_runtime.py
Normal file
195
graph/native_runtime.py
Normal file
|
|
@ -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
|
||||
353
graph/task_nodes.py
Normal file
353
graph/task_nodes.py
Normal file
|
|
@ -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
|
||||
114
tests/test_native_graph_runtime.py
Normal file
114
tests/test_native_graph_runtime.py
Normal file
|
|
@ -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"
|
||||
112
tests/test_task_execution_native_graph.py
Normal file
112
tests/test_task_execution_native_graph.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue