diff --git a/Examples/Nodes/MathNode.cs b/Examples/Nodes/MathNode.cs
index 6b98c73..ae1e84b 100644
--- a/Examples/Nodes/MathNode.cs
+++ b/Examples/Nodes/MathNode.cs
@@ -1,12 +1,12 @@
-using UnityEngine;
-
-[System.Serializable]
+[System.Serializable]
 public class MathNode : Node {
     [Input] public float a;
     [Input] public float b;
     [Output] public float result;
-    public enum ValueType { Float, Int }
     public enum MathType { Add, Subtract, Multiply, Divide}
-    public ValueType valueType = ValueType.Float;
     public MathType mathType = MathType.Add;
+
+    protected override void Init() {
+        name = "Math";
+    }
 }
diff --git a/Scripts/Editor/NodeEditorAction.cs b/Scripts/Editor/NodeEditorAction.cs
index e27e518..b360063 100644
--- a/Scripts/Editor/NodeEditorAction.cs
+++ b/Scripts/Editor/NodeEditorAction.cs
@@ -131,7 +131,7 @@ public partial class NodeEditorWindow {
             if (!_portConnectionPoints.ContainsKey(draggedOutput)) return;
             Vector2 from = draggedOutput.node.position.position + _portConnectionPoints[draggedOutput].center;
             Vector2 to = draggedOutputTarget != null ? draggedOutputTarget.node.position.position + portConnectionPoints[draggedOutputTarget].center : WindowToGridPosition(Event.current.mousePosition);
-            DrawConnection(from, to);
+            DrawConnection(from, to, NodeEditorUtilities.GetTypeColor(draggedOutput.type));
         }
     }
 
diff --git a/Scripts/Editor/NodeEditorGUI.cs b/Scripts/Editor/NodeEditorGUI.cs
index 67a6574..008e6e6 100644
--- a/Scripts/Editor/NodeEditorGUI.cs
+++ b/Scripts/Editor/NodeEditorGUI.cs
@@ -19,10 +19,6 @@ public partial class NodeEditorWindow {
         GUI.matrix = m;
     }
 
-    public static void DrawConnection(Vector2 from, Vector2 to, Color col) {
-        Handles.DrawBezier(from, to, from, to, col, new Texture2D(2, 2), 2);
-    }
-
     public static void BeginZoomed(Rect rect, float zoom) {
         GUI.EndClip();
             
@@ -101,7 +97,8 @@ public partial class NodeEditorWindow {
         contextMenu.DropDown(new Rect(Event.current.mousePosition, Vector2.zero));
     }
 
-    public void DrawConnection(Vector2 startPoint, Vector2 endPoint) {
+    /// <summary> Draw a bezier from startpoint to endpoint, both in grid coordinates </summary>
+    public void DrawConnection(Vector2 startPoint, Vector2 endPoint, Color col) {
         startPoint = GridToWindowPosition(startPoint);
         endPoint = GridToWindowPosition(endPoint);
 
@@ -113,8 +110,11 @@ public partial class NodeEditorWindow {
         if (startPoint.x > endPoint.x) endTangent.x = Mathf.LerpUnclamped(endPoint.x, startPoint.x, -0.7f);
         else endTangent.x = Mathf.LerpUnclamped(endPoint.x, startPoint.x, 0.7f);
 
-        Handles.DrawBezier(startPoint, endPoint, startTangent, endTangent, Color.gray, null, 4);
-        Handles.DrawBezier(startPoint, endPoint, startTangent, endTangent, Color.white, null, 2);
+        Color prevCol = GUI.color;
+        Color edgeCol = new Color(0.1f, 0.1f, 0.1f, col.a);
+        Handles.DrawBezier(startPoint, endPoint, startTangent, endTangent, edgeCol, null, 4);
+        Handles.DrawBezier(startPoint, endPoint, startTangent, endTangent, col, null, 2);
+        GUI.color = prevCol;
     }
 
     public void DrawConnections() {
@@ -128,7 +128,7 @@ public partial class NodeEditorWindow {
                 for (int k = 0; k < output.ConnectionCount; k++) {
                     NodePort input = output.GetConnection(k);
                     Vector2 to = input.node.position.position + _portConnectionPoints[input].center;
-                    DrawConnection(from, to);
+                    DrawConnection(from, to, NodeEditorUtilities.GetTypeColor(output.type));
                 }
             }
         }