150 lines
7.2 KiB
Python
150 lines
7.2 KiB
Python
from __future__ import annotations
|
|
|
|
from django.utils import timezone
|
|
|
|
from agents.steward import StewardService
|
|
from graph.models import GraphApproval, GraphApprovalStatus
|
|
from graph.native_runtime import GraphExecutionContext
|
|
from graph.registry import NodeHandlerRegistry, NodeResult
|
|
from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec
|
|
|
|
|
|
def steward_run_graph_v1() -> ExecutionGraphSpec:
|
|
spec = ExecutionGraphSpec(
|
|
name="steward_run",
|
|
version=1,
|
|
graph_type="STEWARD_RUN",
|
|
entry="prepare",
|
|
nodes={
|
|
"prepare": GraphNodeSpec("prepare", "steward_prepare"),
|
|
"collect_signals": GraphNodeSpec("collect_signals", "steward_collect_signals"),
|
|
"run_checks": GraphNodeSpec("run_checks", "steward_run_checks"),
|
|
"normalize_findings": GraphNodeSpec("normalize_findings", "steward_normalize_findings"),
|
|
"classify": GraphNodeSpec("classify", "steward_classify"),
|
|
"route": GraphNodeSpec("route", "steward_route"),
|
|
"await_approval": GraphNodeSpec("await_approval", "steward_await_approval"),
|
|
"summarize": GraphNodeSpec("summarize", "steward_summarize"),
|
|
"complete": GraphNodeSpec("complete", "complete", {"terminal": True}),
|
|
},
|
|
edges=[
|
|
GraphEdgeSpec("prepare", "collect_signals", "success"),
|
|
GraphEdgeSpec("collect_signals", "run_checks", "success"),
|
|
GraphEdgeSpec("run_checks", "normalize_findings", "success"),
|
|
GraphEdgeSpec("normalize_findings", "classify", "success"),
|
|
GraphEdgeSpec("classify", "route", "success"),
|
|
GraphEdgeSpec("route", "await_approval", "approval_required"),
|
|
GraphEdgeSpec("route", "summarize", "success"),
|
|
GraphEdgeSpec("await_approval", "summarize", "approved"),
|
|
GraphEdgeSpec("await_approval", "summarize", "rejected"),
|
|
GraphEdgeSpec("summarize", "complete", "success"),
|
|
],
|
|
terminal_nodes=["complete"],
|
|
metadata={"description": "Steward V1 governance workflow; detection, classification, routing, no code mutation."},
|
|
)
|
|
spec.validate()
|
|
return spec
|
|
|
|
|
|
class StewardNode:
|
|
idempotent = True
|
|
replay_safe = True
|
|
destructive = False
|
|
|
|
def __init__(self, service: StewardService, node_type: str) -> None:
|
|
self.service = service
|
|
self.node_type = node_type
|
|
|
|
def steward_run_id(self, context: GraphExecutionContext) -> str:
|
|
return str(context.graph_run.metadata["steward_run_id"])
|
|
|
|
|
|
class StewardPrepareNode(StewardNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
return NodeResult("COMPLETE", "success", {"steward_run_id": self.steward_run_id(context)})
|
|
|
|
|
|
class StewardCollectSignalsNode(StewardNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
return NodeResult("COMPLETE", "success")
|
|
|
|
|
|
class StewardRunChecksNode(StewardNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
from control_plane.projects.models import StewardRun
|
|
|
|
steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
|
|
checks = self.service.run_checks(steward_run)
|
|
return NodeResult("COMPLETE", "success", {"check_count": len(checks)})
|
|
|
|
|
|
class StewardNormalizeFindingsNode(StewardNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
from control_plane.projects.models import StewardRun
|
|
|
|
steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
|
|
findings = self.service.normalize_findings(steward_run)
|
|
return NodeResult("COMPLETE", "success", {"finding_count": len(findings)})
|
|
|
|
|
|
class StewardClassifyNode(StewardNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
from control_plane.projects.models import StewardRun
|
|
|
|
steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
|
|
findings = self.service.classify_findings(steward_run)
|
|
return NodeResult("COMPLETE", "success", {"classified_count": len(findings)})
|
|
|
|
|
|
class StewardRouteNode(StewardNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
from control_plane.projects.models import StewardRun
|
|
|
|
steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
|
|
actions = self.service.route_findings(steward_run)
|
|
approval_required = any(action.requires_approval for action in actions)
|
|
return NodeResult("COMPLETE", "approval_required" if approval_required else "success", {"action_count": len(actions), "approval_required": approval_required})
|
|
|
|
|
|
class StewardApprovalNode(StewardNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
node_run = context.graph_run.node_runs.filter(node_id=context.graph_run.current_node).order_by("-visit_index").first()
|
|
if GraphApproval.objects.filter(graph_run=context.graph_run, status=GraphApprovalStatus.APPROVED).exists():
|
|
from control_plane.projects.models import StewardRun
|
|
|
|
steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
|
|
for action in steward_run.findings.filter(actions__requires_approval=True).values_list("actions__id", flat=True).distinct():
|
|
from control_plane.projects.models import StewardAction
|
|
|
|
steward_action = StewardAction.objects.get(id=action)
|
|
if steward_action.task_id is None and steward_action.action_type == "REPAIR":
|
|
steward_action.task = self.service._create_repair_task(steward_action.finding)
|
|
steward_action.status = "ROUTED"
|
|
steward_action.approved_at = timezone.now()
|
|
steward_action.save(update_fields=["task", "status", "approved_at", "updated_at"])
|
|
return NodeResult("COMPLETE", "approved")
|
|
if GraphApproval.objects.filter(graph_run=context.graph_run, status=GraphApprovalStatus.REJECTED).exists():
|
|
return NodeResult("COMPLETE", "rejected")
|
|
GraphApproval.objects.get_or_create(graph_run=context.graph_run, node_run=node_run, reason="AWAITING_STEWARD_ROUTING_APPROVAL")
|
|
return NodeResult("PAUSED", "awaiting", pause_reason="AWAITING_STEWARD_ROUTING_APPROVAL")
|
|
|
|
|
|
class StewardSummarizeNode(StewardNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
from control_plane.projects.models import StewardRun
|
|
|
|
steward_run = StewardRun.objects.get(id=self.steward_run_id(context))
|
|
self.service.complete_run(steward_run)
|
|
return NodeResult("COMPLETE", "success", {"summary": steward_run.summary})
|
|
|
|
|
|
def steward_registry(service: StewardService) -> NodeHandlerRegistry:
|
|
registry = NodeHandlerRegistry()
|
|
registry.register(StewardPrepareNode(service, "steward_prepare"))
|
|
registry.register(StewardCollectSignalsNode(service, "steward_collect_signals"))
|
|
registry.register(StewardRunChecksNode(service, "steward_run_checks"))
|
|
registry.register(StewardNormalizeFindingsNode(service, "steward_normalize_findings"))
|
|
registry.register(StewardClassifyNode(service, "steward_classify"))
|
|
registry.register(StewardRouteNode(service, "steward_route"))
|
|
registry.register(StewardApprovalNode(service, "steward_await_approval"))
|
|
registry.register(StewardSummarizeNode(service, "steward_summarize"))
|
|
return registry
|