diff --git a/src/AnimGraphEditor.cc b/src/AnimGraphEditor.cc
index eba40a4..1e61a66 100644
--- a/src/AnimGraphEditor.cc
+++ b/src/AnimGraphEditor.cc
@@ -161,9 +161,9 @@ void AnimGraphEditorUpdate() {
 
     // Header
     ImNodes::BeginNodeTitleBar();
-    if (i == 0) {
+    if (&node_resource == &graph_resource.getGraphOutputNode()) {
       ImGui::TextUnformatted("Graph Outputs");
-    } else if (i == 1) {
+    } else if (&node_resource == &graph_resource.getGraphInputNode()) {
       ImGui::TextUnformatted("Graph Inputs");
     } else {
       ImGui::TextUnformatted(node_resource.m_type_name.c_str());
diff --git a/src/AnimGraphResource.h b/src/AnimGraphResource.h
index 18471ec..b93ddbc 100644
--- a/src/AnimGraphResource.h
+++ b/src/AnimGraphResource.h
@@ -513,6 +513,19 @@ struct AnimGraphResource {
   AnimNodeResource& getGraphOutputNode() { return m_nodes[0]; }
   AnimNodeResource& getGraphInputNode() { return m_nodes[1]; }
 
+  const AnimNodeResource& getGraphOutputNode() const { return m_nodes[0]; }
+  const AnimNodeResource& getGraphInputNode() const { return m_nodes[1]; }
+
+  size_t getNodeIndex (const AnimNodeResource& node_resource) const {
+    for (size_t i = 0, n = m_nodes.size(); i < n; i++) {
+      if (&m_nodes[i] == &node_resource) {
+        return i;
+      }
+    }
+
+    return -1;
+  }
+
   size_t addNode(AnimNodeResource node_resource) {
     m_nodes.push_back(node_resource);
     return m_nodes.size() - 1;