diff --git a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
index 0dd3bfa20b..5c54e31dcb 100644
--- a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
+++ b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
@@ -115,6 +115,19 @@
 			if (snapshot) {
 				flowStore.val = snapshot
 				refreshStateStore(flowStore)
+
+				if ($currentEditor) {
+					const module = getModule($currentEditor.stepId, snapshot)
+					if (module) {
+						if ($currentEditor.type === 'script' && module.value.type === 'rawscript') {
+							$currentEditor.editor.setCode(module.value.content)
+						} else if ($currentEditor.type === 'iterator' && module.value.type === 'forloopflow') {
+							$currentEditor.editor.setCode(
+								module.value.iterator.type === 'javascript' ? module.value.iterator.expr : ''
+							)
+						}
+					}
+				}
 			}
 		},
 		showModuleDiff(id: string) {