1 files changed, 7 insertions, 16 deletions
diff --git a/makima/frontend/src/components/directives/DirectiveDAG.tsx b/makima/frontend/src/components/directives/DirectiveDAG.tsx
index 8c7def9..142df41 100644
--- a/makima/frontend/src/components/directives/DirectiveDAG.tsx
+++ b/makima/frontend/src/components/directives/DirectiveDAG.tsx
@@ -2,7 +2,6 @@ import { useMemo } from "react";
 import type { DirectiveStep } from "../../lib/api";
 import { StepNode } from "./StepNode";
 import {
-  OrchestratorStepNode,
   type OrchestratorStepType,
   type OrchestratorStepStatus,
 } from "./OrchestratorStepNode";
@@ -35,13 +34,6 @@ interface Layer {
   steps: DirectiveStep[];
 }
 
-/** Types that should appear before the regular DAG steps */
-const BEFORE_TYPES = new Set<OrchestratorStepType>([
-  "planning",
-  "replanning",
-  "plan-orders",
-]);
-
 function topoSort(steps: DirectiveStep[]): Layer[] {
   if (steps.length === 0) return [];
 
@@ -97,14 +89,13 @@ export function DirectiveDAG({ steps, specializedSteps, onComplete, onFail, onSk
             </div>
           )}
           <div className="flex flex-wrap gap-3 justify-center">
-            {afterSteps.map((vs) => (
-              <OrchestratorStepNode
-                key={`${vs.type}-${vs.taskId}`}
-                type={vs.type}
-                taskId={vs.taskId}
-                status={vs.status}
-                label={vs.label}
-                hasQuestions={vs.hasQuestions}
+            {layer.steps.map((step) => (
+              <StepNode
+                key={step.id}
+                step={step}
+                onComplete={onComplete ? () => onComplete(step.id) : undefined}
+                onFail={onFail ? () => onFail(step.id) : undefined}
+                onSkip={onSkip ? () => onSkip(step.id) : undefined}
               />
             ))}
           </div>