diff --git a/frontend/src/lib/components/FlowBuilder.svelte b/frontend/src/lib/components/FlowBuilder.svelte
index e751d2463e..e330012172 100644
--- a/frontend/src/lib/components/FlowBuilder.svelte
+++ b/frontend/src/lib/components/FlowBuilder.svelte
@@ -137,8 +137,7 @@
 					'state',
 					encodeState({
 						flow: $flowStore,
-						selectedId: $selectedIdStore,
-						flowStateStore: $flowStateStore
+						selectedId: $selectedIdStore
 					})
 				),
 			500
@@ -175,7 +174,6 @@
 	}
 
 	onMount(() => {
-		console.log(selectedId)
 		selectedId && select(selectedId)
 	})
 
diff --git a/frontend/src/lib/components/FlowStatusViewer.svelte b/frontend/src/lib/components/FlowStatusViewer.svelte
index 7cba20f27c..07eb1a74b0 100644
--- a/frontend/src/lib/components/FlowStatusViewer.svelte
+++ b/frontend/src/lib/components/FlowStatusViewer.svelte
@@ -112,7 +112,7 @@
 								if (flowState) {
 									if (
 										!flowState[flowJobIds.moduleId].previewResult ||
-										!Array.isArray(flowState[flowJobIds.moduleId].previewResult)
+										!Array.isArray(flowState[flowJobIds.moduleId]?.previewResult)
 									) {
 										flowState[flowJobIds.moduleId].previewResult = []
 									}
@@ -163,7 +163,7 @@
 									: undefined}
 								on:jobsLoaded={(e) => {
 									if (mod.id && (mod.flow_jobs ?? []).length == 0) {
-										if (flowState) {
+										if (flowState && flowState[mod.id]) {
 											flowState[mod.id].previewResult = e.detail.result
 											flowState[mod.id].previewArgs = e.detail.args
 										}
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index 5a4914c0c1..c75bce373b 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -84,7 +84,11 @@
 		})
 
 		if (JSON.stringify(schema) !== JSON.stringify($flowStateStore[flowModule.id]?.schema)) {
-			$flowStateStore[flowModule.id].schema = schema
+			if (!$flowStateStore[flowModule.id]) {
+				$flowStateStore[flowModule.id] = { schema }
+			} else {
+				$flowStateStore[flowModule.id].schema = schema
+			}
 		}
 	}
 
@@ -174,7 +178,7 @@
 									<Kbd>Ctrl/Cmd</Kbd> + <Kbd>S</Kbd>
 								</Tooltip>Inputs</Tab
 							>
-							<Tab value="test">Test</Tab>
+							<Tab value="test">Test this step</Tab>
 							<Tab value="retries">Retries</Tab>
 							{#if !$selectedId.includes('failure')}
 								<Tab value="early-stop">Early Stop</Tab>
@@ -189,7 +193,7 @@
 										pickableProperties={stepPropPicker.pickableProperties}
 									>
 										<SchemaForm
-											schema={$flowStateStore[$selectedId].schema}
+											schema={$flowStateStore[$selectedId]?.schema ?? {}}
 											inputTransform={true}
 											importPath={$selectedId}
 											bind:args={flowModule.value.input_transforms}
@@ -201,7 +205,7 @@
 								<ModulePreview
 									bind:this={modulePreview}
 									mod={flowModule}
-									schema={$flowStateStore[$selectedId].schema}
+									schema={$flowStateStore[$selectedId]?.schema ?? {}}
 								/>
 							{:else if selected === 'retries'}
 								<FlowRetries bind:flowModule class="px-4 pb-4 h-full overflow-auto" />
diff --git a/frontend/src/lib/components/flows/content/FlowModuleSuspend.svelte b/frontend/src/lib/components/flows/content/FlowModuleSuspend.svelte
index a6bcc869e6..a235ccd463 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleSuspend.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleSuspend.svelte
@@ -25,7 +25,7 @@
 	let editor: SimpleEditor | undefined = undefined
 
 	const pickableProperties = {
-		result: $flowStateStore[$selectedId].previewResult
+		result: $flowStateStore[$selectedId]?.previewResult ?? {}
 	}
 
 	$: isSuspendEnabled = Boolean(flowModule.suspend)
diff --git a/frontend/src/lib/components/flows/flowState.ts b/frontend/src/lib/components/flows/flowState.ts
index 430a2eceea..e8f26ba3e8 100644
--- a/frontend/src/lib/components/flows/flowState.ts
+++ b/frontend/src/lib/components/flows/flowState.ts
@@ -19,8 +19,8 @@ export type FlowState = Record<string, FlowModuleState>
  */
 export const flowStateStore = writable<FlowState>({})
 
-export async function initFlowState(flow: Flow, flowState: FlowState | undefined) {
-	const modulesState: FlowState = flowState ?? {}
+export async function initFlowState(flow: Flow) {
+	const modulesState: FlowState = {}
 
 	await mapFlowModules(flow.value.modules, modulesState)
 
@@ -28,6 +28,7 @@ export async function initFlowState(flow: Flow, flowState: FlowState | undefined
 		? await loadFlowModuleState(flow.value.failure_module)
 		: emptyFlowModuleState()
 
+	console.log(modulesState)
 	flowStateStore.set({
 		...modulesState,
 		failure: failureModule
@@ -63,6 +64,7 @@ async function mapFlowModule(flowModule: FlowModule, modulesState: FlowState) {
 		const flowModuleState = await loadFlowModuleState(flowModule)
 		modulesState[flowModule.id] = flowModuleState
 	}
+
 }
 
 async function mapFlowModules(flowModules: FlowModule[], modulesState: FlowState) {
diff --git a/frontend/src/lib/components/flows/flowStateUtils.ts b/frontend/src/lib/components/flows/flowStateUtils.ts
index 2484457cb0..2108f44172 100644
--- a/frontend/src/lib/components/flows/flowStateUtils.ts
+++ b/frontend/src/lib/components/flows/flowStateUtils.ts
@@ -14,7 +14,7 @@ import {
 	NEVER_TESTED_THIS_FAR,
 	numberToChars
 } from './utils'
-import { Mutex, Semaphore, withTimeout } from 'async-mutex';
+import { Mutex } from 'async-mutex';
 
 export async function loadFlowModuleState(flowModule: FlowModule): Promise<FlowModuleState> {
 	try {
diff --git a/frontend/src/lib/components/flows/flowStore.ts b/frontend/src/lib/components/flows/flowStore.ts
index fe9c16689f..6c74463841 100644
--- a/frontend/src/lib/components/flows/flowStore.ts
+++ b/frontend/src/lib/components/flows/flowStore.ts
@@ -51,7 +51,7 @@ export function dfs(modules: FlowModule[], previewOrder: boolean = false): strin
 
 export const flowIds = derived(flowStore, flow => dfs(flow.value.modules))
 
-export async function initFlow(flow: Flow, flowState: FlowState | undefined) {
+export async function initFlow(flow: Flow) {
 
 	let counter = 40
 	for (const mod of flow.value.modules) {
@@ -68,7 +68,7 @@ export async function initFlow(flow: Flow, flowState: FlowState | undefined) {
 		}
 	}
 
-	await initFlowState(flow, flowState)
+	await initFlowState(flow)
 
 	flowStore.set(flow)
 
diff --git a/frontend/src/lib/components/flows/header/FlowImportExportMenu.svelte b/frontend/src/lib/components/flows/header/FlowImportExportMenu.svelte
index d675fef763..08b1f3e758 100644
--- a/frontend/src/lib/components/flows/header/FlowImportExportMenu.svelte
+++ b/frontend/src/lib/components/flows/header/FlowImportExportMenu.svelte
@@ -20,7 +20,7 @@
 	function importJson() {
 		Object.assign($flowStore, JSON.parse(pendingJson))
 
-		initFlow($flowStore, undefined)
+		initFlow($flowStore)
 		sendUserToast('OpenFlow imported from JSON')
 		jsonSetterDrawer.toggleDrawer()
 	}
diff --git a/frontend/src/lib/components/flows/previousResults.ts b/frontend/src/lib/components/flows/previousResults.ts
index 242ea93636..d25a7da789 100644
--- a/frontend/src/lib/components/flows/previousResults.ts
+++ b/frontend/src/lib/components/flows/previousResults.ts
@@ -110,7 +110,7 @@ function getFlowInput(
 				} else {
 					return {
 						...parentFlowInput,
-						previous_result: flattenPreviousResult(flowState[parentPreviousModuleId].previewResult)
+						previous_result: flattenPreviousResult(flowState[parentPreviousModuleId]?.previewResult ?? {})
 					}
 				}
 			}
@@ -131,7 +131,7 @@ export function getStepPropPicker(
 	const flowInput = getFlowInput(dfs(parentModule?.id, flow), flowState, args, flow.schema)
 
 	const previousResults = previousModule
-		? flowState[previousModule.id].previewResult
+		? flowState[previousModule.id]?.previewResult
 		: flattenPreviousResult(flowInput)
 
 	const pickableProperties = {
diff --git a/frontend/src/lib/components/scripts/CreateActions.svelte b/frontend/src/lib/components/scripts/CreateActions.svelte
index 10d3f7f0cd..0a4c71d091 100644
--- a/frontend/src/lib/components/scripts/CreateActions.svelte
+++ b/frontend/src/lib/components/scripts/CreateActions.svelte
@@ -44,7 +44,7 @@
 	function importJson() {
 		Object.assign($flowStore, JSON.parse(pendingJson))
 
-		initFlow($flowStore, undefined)
+		initFlow($flowStore)
 		sendUserToast('OpenFlow imported from JSON')
 		drawers.json?.toggleDrawer()
 	}
diff --git a/frontend/src/routes/flows/add.svelte b/frontend/src/routes/flows/add.svelte
index 2c097567b7..3b72dd6bf8 100644
--- a/frontend/src/routes/flows/add.svelte
+++ b/frontend/src/routes/flows/add.svelte
@@ -50,7 +50,7 @@
 			$page.url.searchParams.delete('hub')
 		}
 		selectedId = state?.selectedId
-		await initFlow(flow, state?.flowState)
+		await initFlow(flow)
 	}
 
 	loadFlow()
diff --git a/frontend/src/routes/flows/edit/[...path].svelte b/frontend/src/routes/flows/edit/[...path].svelte
index d9c20bd79a..a4ac322820 100644
--- a/frontend/src/routes/flows/edit/[...path].svelte
+++ b/frontend/src/routes/flows/edit/[...path].svelte
@@ -34,7 +34,7 @@
 
 	let initialPath: string = ''
 
-	initFlow(flow, undefined)
+	initFlow(flow)
 
 	async function loadFlow(): Promise<void> {
 		flow =
@@ -46,7 +46,7 @@
 				  })
 		initialPath = flow.path
 
-		await initFlow(flow, stateLoadedFromUrl?.flowState)
+		await initFlow(flow)
 		selectedId = stateLoadedFromUrl?.selectedId
 		$dirtyStore = false
 	}