diff --git a/control_plane/agents/migrations/0005_progenysignal_graph_lineage.py b/control_plane/agents/migrations/0005_progenysignal_graph_lineage.py new file mode 100644 index 0000000..a766f45 --- /dev/null +++ b/control_plane/agents/migrations/0005_progenysignal_graph_lineage.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import django.db.models.deletion +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("agents", "0004_progenysignal"), + ("graph", "0003_graphapproval"), + ] + + operations = [ + migrations.AddField( + model_name="progenysignal", + name="execution_graph_version", + field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to="graph.executiongraphversion"), + ), + migrations.AddField( + model_name="progenysignal", + name="graph_node_run", + field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to="graph.graphnoderun"), + ), + migrations.AddField( + model_name="progenysignal", + name="graph_run", + field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to="graph.graphrun"), + ), + ] diff --git a/control_plane/agents/models.py b/control_plane/agents/models.py index 5f88f5d..ee3cc39 100644 --- a/control_plane/agents/models.py +++ b/control_plane/agents/models.py @@ -88,6 +88,9 @@ class ProgenySignal(TimestampedModel): task = models.ForeignKey("projects.Task", on_delete=models.SET_NULL, null=True, blank=True) milestone = models.ForeignKey("projects.Milestone", on_delete=models.SET_NULL, null=True, blank=True) agent_version = models.ForeignKey(AgentVersion, on_delete=models.SET_NULL, null=True, blank=True) + graph_run = models.ForeignKey("graph.GraphRun", on_delete=models.SET_NULL, null=True, blank=True) + graph_node_run = models.ForeignKey("graph.GraphNodeRun", on_delete=models.SET_NULL, null=True, blank=True) + execution_graph_version = models.ForeignKey("graph.ExecutionGraphVersion", on_delete=models.SET_NULL, null=True, blank=True) source = models.CharField(max_length=80) severity = models.CharField(max_length=32, default="INFO") failure_category = models.CharField(max_length=80, blank=True) diff --git a/tests/test_execution_graph_phase_a.py b/tests/test_execution_graph_phase_a.py index f0c567e..821ad37 100644 --- a/tests/test_execution_graph_phase_a.py +++ b/tests/test_execution_graph_phase_a.py @@ -2,6 +2,7 @@ from __future__ import annotations import pytest +from control_plane.agents.models import ProgenySignal from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus, GraphRun, GraphRunStatus from graph.registry import NodeHandlerRegistry, NodeResult from graph.spec import ExecutionGraphSpec, outgoing_edges @@ -67,3 +68,25 @@ def test_graph_models_persist_versioned_run_state() -> None: assert run.execution_graph_version.graph.name == "task_execution" assert run.execution_graph_version.status == ExecutionGraphVersionStatus.CHAMPION assert run.metadata["node_count"] == 12 + + +def test_progeny_signal_can_reference_graph_lineage() -> None: + spec = task_execution_graph_v1() + definition = ExecutionGraphDefinition.objects.create(name="lineage_graph", graph_type=spec.graph_type) + version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict()) + run = GraphRun.objects.create(execution_graph_version=version, current_node=spec.entry) + node_run = run.node_runs.create(node_id="coder", node_type="coder", visit_index=1) + + signal = ProgenySignal.objects.create( + source="runtime", + severity="high", + failure_category="MODEL_OUTPUT_INVALID", + summary="Malformed model JSON", + graph_run=run, + graph_node_run=node_run, + execution_graph_version=version, + ) + + assert signal.graph_run == run + assert signal.graph_node_run == node_run + assert signal.execution_graph_version == version