diff --git a/frontend/src/lib/components/Dev.svelte b/frontend/src/lib/components/Dev.svelte
index 3cd95d7c55..b4570469ad 100644
--- a/frontend/src/lib/components/Dev.svelte
+++ b/frontend/src/lib/components/Dev.svelte
@@ -453,11 +453,13 @@
 	type LastEditFlow = {
 		flow: OpenFlow
 		uriPath: string
+		path: string
 	}
 	let lastUriPath: string | undefined = undefined
 	async function replaceFlow(lastEdit: LastEditFlow) {
 		mode = 'flow'
 		lastUriPath = lastEdit.uriPath
+		pathStore.set(lastEdit.path)
 		// sendUserToast(JSON.stringify(lastEdit.flow), true)
 		// return
 		try {
@@ -469,6 +471,10 @@
 					lastEdit.flow.value = { modules: [] }
 				}
 				flowStore.val = lastEdit.flow
+				try {
+					let ids = dfs(flowStore.val.value.modules ?? [], (m) => m.id)
+					flowStateStore.val = Object.fromEntries(ids.map((k) => [k, {}]))
+				} catch (e) {}
 				inferModuleArgs($selectedIdStore)
 			}
 		} catch (e) {
@@ -486,8 +492,9 @@
 	const selectedIdStore = writable('settings-metadata')
 	const triggersCount = writable<TriggersCount | undefined>(undefined)
 	const modulesTestStates = new ModulesTestStates((moduleId) => {
+		// console.log('FOO')
 		// Update the derived store with test job states
-		showJobStatus = false
+		// showJobStatus = false
 	})
 	const outputPickerOpenFns: Record<string, () => void> = $state({})
 
@@ -497,18 +504,21 @@
 		showCaptureHint: writable(undefined),
 		triggersState: new Triggers()
 	})
+
+	let pathStore = writable('')
+	let initialPathStore = writable('')
 	setContext<FlowEditorContext>('FlowEditorContext', {
 		selectedId: selectedIdStore,
 		previewArgs: previewArgsStore,
 		scriptEditorDrawer,
 		moving,
 		history,
-		pathStore: writable(''),
+		pathStore: pathStore,
 		flowStateStore,
 		flowStore,
 		stepsInputArgs,
 		saveDraft: () => {},
-		initialPathStore: writable(''),
+		initialPathStore,
 		fakeInitialPath: '',
 		flowInputsStore: writable<FlowInput>({}),
 		customUi: {},
@@ -791,8 +801,8 @@
 					<FlowPreviewButtons
 						bind:this={flowPreviewButtons}
 						{onJobDone}
+						bind:localModuleStates
 						onRunPreview={() => {
-							localModuleStates = {}
 							showJobStatus = true
 						}}
 					/>
@@ -807,7 +817,7 @@
 								disableTutorials
 								smallErrorHandler={true}
 								disableStaticInputs
-								{localModuleStates}
+								localModuleStates={showJobStatus ? localModuleStates : {}}
 								onTestUpTo={flowPreviewButtons?.testUpTo}
 								testModuleStates={modulesTestStates}
 								isOwner={flowPreviewContent?.getIsOwner?.()}
diff --git a/frontend/src/lib/components/ModuleTest.svelte b/frontend/src/lib/components/ModuleTest.svelte
index df22ebfd72..db2c6082c7 100644
--- a/frontend/src/lib/components/ModuleTest.svelte
+++ b/frontend/src/lib/components/ModuleTest.svelte
@@ -34,12 +34,13 @@
 	let stepHistoryLoader = getStepHistoryLoaderContext()
 
 	export function runTestWithStepArgs() {
-		runTest(stepsInputArgs.getStepArgs(mod.id))
+		const args = stepsInputArgs.getStepArgs(mod.id)
+		runTest(args)
 	}
 
 	export function loadArgsAndRunTest() {
 		stepsInputArgs?.updateStepArgs(mod.id, flowStateStore.val, flowStore?.val, previewArgs?.val)
-		runTest(stepsInputArgs.getStepArgs(mod.id))
+		runTestWithStepArgs()
 	}
 
 	export async function runTest(args: any) {
diff --git a/frontend/src/lib/components/graph/FlowGraphV2.svelte b/frontend/src/lib/components/graph/FlowGraphV2.svelte
index 2daa40d275..3b95378497 100644
--- a/frontend/src/lib/components/graph/FlowGraphV2.svelte
+++ b/frontend/src/lib/components/graph/FlowGraphV2.svelte
@@ -545,7 +545,6 @@
 {#if insertable}
 	<FlowYamlEditor bind:drawer={yamlEditorDrawer} />
 {/if}
-
 <div
 	style={`height: ${height}px; max-height: ${maxHeight}px;`}
 	class="overflow-clip"