diff --git a/tooling/visualization/scripts/graph-visualization.js b/tooling/visualization/scripts/graph-visualization.js
index <HASH>..<HASH> 100644
--- a/tooling/visualization/scripts/graph-visualization.js
+++ b/tooling/visualization/scripts/graph-visualization.js
@@ -93,17 +93,14 @@
     }
     
     function getNodeColor(node) {
-        var isFinalNode = node.outgoingEdges.length === 0;
-        
-        if (isFinalNode) {
-            return "#4CD964";
-        }
-        
-        if (node.isEntryNode) {
-            return "#FFCC00";
+        switch (node.type) {
+            case Styx.NodeType.Entry:
+                return "#FFCC00";
+            case Styx.NodeType.Exit:
+                return "#4CD964";
+            default:
+                return "#97C2FC";
         }
-        
-        return "#97C2FC";
     }
     
     function getEdgeColor(edge) {
Inspects the 'type' property to determine a node's color