3 files changed, 27 insertions, 4 deletions

diff --git a/editor/plugins/visual_shader_editor_plugin.cpp b/editor/plugins/visual_shader_editor_plugin.cpp
index f5ae3c8bf5..9810fbee2c 100644
--- a/editor/plugins/visual_shader_editor_plugin.cpp
+++ b/editor/plugins/visual_shader_editor_plugin.cpp
@@ -93,6 +93,7 @@ void VisualShaderEditor::edit(VisualShader *p_visual_shader) {
 			edit_type = edit_type_standart;
 			particles_mode = false;
 		}
+		visual_shader->set_shader_type(get_current_shader_type());
 	} else {
 		if (visual_shader.is_valid()) {
 			if (visual_shader->is_connected("changed", callable_mp(this, &VisualShaderEditor::_update_preview))) {
@@ -540,6 +541,7 @@ void VisualShaderEditor::_update_graph() {
 		String expression = "";
 
 		GraphNode *node = memnew(GraphNode);
+		visual_shader->set_graph_node(type, nodes[n_i], node);
 
 		if (is_group) {
 			size = group_node->get_size();
@@ -1518,14 +1520,10 @@ VisualShaderNode *VisualShaderEditor::_add_node(int p_idx, int p_op_idx) {
 void VisualShaderEditor::_node_dragged(const Vector2 &p_from, const Vector2 &p_to, int p_node) {
 	VisualShader::Type type = get_current_shader_type();
 
-	updating = true;
 	undo_redo->create_action(TTR("Node Moved"));
 	undo_redo->add_do_method(visual_shader.ptr(), "set_node_position", type, p_node, p_to);
 	undo_redo->add_undo_method(visual_shader.ptr(), "set_node_position", type, p_node, p_from);
-	undo_redo->add_do_method(this, "_update_graph");
-	undo_redo->add_undo_method(this, "_update_graph");
 	undo_redo->commit_action();
-	updating = false;
 }
 
 void VisualShaderEditor::_connection_request(const String &p_from, int p_from_index, const String &p_to, int p_to_index) {
@@ -2055,6 +2053,7 @@ void VisualShaderEditor::_delete_nodes() {
 }
 
 void VisualShaderEditor::_mode_selected(int p_id) {
+	visual_shader->set_shader_type(VisualShader::Type(p_id));
 	_update_options_menu();
 	_update_graph();
 }
diff --git a/scene/resources/visual_shader.cpp b/scene/resources/visual_shader.cpp
index b851b4c4dc..394133dd47 100644
--- a/scene/resources/visual_shader.cpp
+++ b/scene/resources/visual_shader.cpp
@@ -317,6 +317,17 @@ VisualShaderNodeCustom::VisualShaderNodeCustom() {
 
 /////////////////////////////////////////////////////////
 
+void VisualShader::set_graph_node(Type p_type, int p_id, GraphNode *p_graph_node) {
+	ERR_FAIL_INDEX(p_type, TYPE_MAX);
+	Graph *g = &graph[p_type];
+	ERR_FAIL_COND(!g->nodes.has(p_id));
+	g->nodes[p_id].graph_node = p_graph_node;
+}
+
+void VisualShader::set_shader_type(Type p_type) {
+	current_type = p_type;
+}
+
 void VisualShader::set_version(const String &p_version) {
 	version = p_version;
 }
@@ -400,6 +411,11 @@ void VisualShader::set_node_position(Type p_type, int p_id, const Vector2 &p_pos
 	Graph *g = &graph[p_type];
 	ERR_FAIL_COND(!g->nodes.has(p_id));
 	g->nodes[p_id].position = p_position;
+	if (current_type == p_type) {
+		if (g->nodes[p_id].graph_node != nullptr) {
+			g->nodes[p_id].graph_node->set_offset(p_position);
+		}
+	}
 }
 
 Vector2 VisualShader::get_node_position(Type p_type, int p_id) const {
diff --git a/scene/resources/visual_shader.h b/scene/resources/visual_shader.h
index 3f5266142a..d45029ec03 100644
--- a/scene/resources/visual_shader.h
+++ b/scene/resources/visual_shader.h
@@ -33,6 +33,7 @@
 
 #include "core/string_builder.h"
 #include "scene/gui/control.h"
+#include "scene/gui/graph_edit.h"
 #include "scene/resources/shader.h"
 
 class VisualShaderNodeUniform;
@@ -69,10 +70,13 @@ public:
 	};
 
 private:
+	Type current_type;
+
 	struct Node {
 		Ref<VisualShaderNode> node;
 		Vector2 position;
 		List<int> prev_connected_nodes;
+		GraphNode *graph_node;
 	};
 
 	struct Graph {
@@ -124,6 +128,10 @@ protected:
 	bool _get(const StringName &p_name, Variant &r_ret) const;
 	void _get_property_list(List<PropertyInfo> *p_list) const;
 
+public: // internal methods
+	void set_graph_node(Type p_type, int p_id, GraphNode *p_graph_node);
+	void set_shader_type(Type p_type);
+
 public:
 	void set_version(const String &p_version);
 	String get_version() const;