diff --git a/frontend/src/lib/components/FirstStepInputs.svelte b/frontend/src/lib/components/FirstStepInputs.svelte
index 775dccdda8..68bc0afe8c 100644
--- a/frontend/src/lib/components/FirstStepInputs.svelte
+++ b/frontend/src/lib/components/FirstStepInputs.svelte
@@ -28,7 +28,7 @@
 		}
 	}
 	$effect(() => {
-		flowStore.val && flowStateStore && untrack(() => loadSchema())
+		flowStore.val && flowStateStore.val && untrack(() => loadSchema())
 	})
 
 	function handleClick() {
diff --git a/frontend/src/lib/components/FlowBuilder.svelte b/frontend/src/lib/components/FlowBuilder.svelte
index 78a0363012..0573c67c40 100644
--- a/frontend/src/lib/components/FlowBuilder.svelte
+++ b/frontend/src/lib/components/FlowBuilder.svelte
@@ -1157,7 +1157,7 @@
 				</div>
 			</div>
 			<!-- metadata -->
-			{#if flowStateStore}
+			{#if flowStateStore.val}
 				<FlowEditor
 					bind:this={flowEditor}
 					{disabledFlowInputs}
diff --git a/frontend/src/lib/components/ModulePreviewForm.svelte b/frontend/src/lib/components/ModulePreviewForm.svelte
index 952892dc1c..3de3d199d8 100644
--- a/frontend/src/lib/components/ModulePreviewForm.svelte
+++ b/frontend/src/lib/components/ModulePreviewForm.svelte
@@ -101,7 +101,7 @@
 	let args = $state(<Record<string, any>>{})
 
 	onMount(() => {
-		testSteps?.updateStepArgs(mod.id, flowStateStore, flowStore?.val, previewArgs?.val)
+		testSteps?.updateStepArgs(mod.id, flowStateStore.val, flowStore?.val, previewArgs?.val)
 		args = testSteps?.getStepArgs(mod.id) ?? { value: {} }
 	})
 </script>
diff --git a/frontend/src/lib/components/ModuleTest.svelte b/frontend/src/lib/components/ModuleTest.svelte
index aacd0b10a6..d382e2580c 100644
--- a/frontend/src/lib/components/ModuleTest.svelte
+++ b/frontend/src/lib/components/ModuleTest.svelte
@@ -35,7 +35,8 @@
 	}
 
 	export function loadArgsAndRunTest() {
-		testSteps?.updateStepArgs(mod.id, flowStateStore, flowStore?.val, previewArgs?.val)
+		console.log('loadArgsAndRunTest', flowStateStore.val)
+		testSteps?.updateStepArgs(mod.id, flowStateStore.val, flowStore?.val, previewArgs?.val)
 		runTest(testSteps.getStepArgs(mod.id)?.value)
 	}
 
diff --git a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
index 29e4e1bef7..ed6b8f019c 100644
--- a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
+++ b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts
@@ -936,7 +936,7 @@ class AIChatManager {
 			const module = getModule(id)
 
 			if (module && module.value.type === 'rawscript') {
-				const moduleState: FlowModuleState | undefined = flowStateStore.val[module.id]
+				const moduleState: FlowModuleState | undefined = flowStateStore[module.id]
 
 				const editorRelated =
 					currentEditor && currentEditor.type === 'script' && currentEditor.stepId === module.id
diff --git a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
index c0688bafd8..7c9d76cafd 100644
--- a/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
+++ b/frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte
@@ -514,7 +514,7 @@
 		const cleanup = aiChatManager.listenForSelectedIdChanges(
 			$selectedId,
 			flowStore.val,
-			flowStateStore,
+			flowStateStore.val,
 			$currentEditor
 		)
 		return cleanup
diff --git a/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte b/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
index 763de8e3fb..668fcf2d54 100644
--- a/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
+++ b/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
@@ -78,7 +78,7 @@
 	}
 
 	$effect(() => {
-		computeMissingInputWarnings(flowStore, flowStateStore, flowInputsStore)
+		computeMissingInputWarnings(flowStore, flowStateStore.val, flowInputsStore)
 	})
 </script>
 
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index 308324c4a3..40528b4dbe 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -250,7 +250,7 @@
 
 	let stepPropPicker = $derived(
 		$executionCount != undefined && failureModule
-			? getFailureStepPropPicker(flowStateStore, flowStore.val, previewArgs.val)
+			? getFailureStepPropPicker(flowStateStore.val, flowStore.val, previewArgs.val)
 			: getStepPropPicker(
 					flowStateStore.val,
 					parentModule,
diff --git a/frontend/src/lib/components/flows/flowStateUtils.svelte.ts b/frontend/src/lib/components/flows/flowStateUtils.svelte.ts
index ced1907a2b..63739f2c41 100644
--- a/frontend/src/lib/components/flows/flowStateUtils.svelte.ts
+++ b/frontend/src/lib/components/flows/flowStateUtils.svelte.ts
@@ -260,7 +260,7 @@ export async function createScriptFromInlineScript(
 	return pickScript(availablePath, flowModule.summary ?? '', flowModule.id, hash)
 }
 
-export function deleteFlowStateById(id: string, flowStateStore: FlowState) {
+export function deleteFlowStateById(id: string, flowStateStore: StateStore<FlowState>) {
 	delete flowStateStore.val[id]
 }
 
@@ -295,7 +295,7 @@ export function sliceModules(
 
 export async function insertNewPreprocessorModule(
 	flowStore: StateStore<ExtendedOpenFlow>,
-	flowStateStore: FlowState,
+	flowStateStore: StateStore<FlowState>,
 	inlineScript?: {
 		language: RawScript['language']
 	},
@@ -325,7 +325,7 @@ export async function insertNewPreprocessorModule(
 
 export async function insertNewFailureModule(
 	flowStore: StateStore<ExtendedOpenFlow>,
-	flowStateStore: FlowState,
+	flowStateStore: StateStore<FlowState>,
 	inlineScript?: {
 		language: RawScript['language']
 		subkind: 'pgsql' | 'flow'
diff --git a/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte b/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
index ab8e565bce..27b59ef366 100644
--- a/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
+++ b/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
@@ -2,7 +2,7 @@
 	import { preventDefault, stopPropagation } from 'svelte/legacy'
 
 	import Popover from '$lib/components/Popover.svelte'
-	import { classNames } from '$lib/utils'
+	import { classNames, type StateStore } from '$lib/utils'
 	import {
 		AlertTriangle,
 		Bed,
@@ -158,7 +158,7 @@
 		id: string | undefined,
 		pickableIds: Record<string, any> | undefined,
 		flowPropPickerConfig: any | undefined,
-		flowStateStore: FlowState | undefined
+		flowStateStore: StateStore<FlowState> | undefined
 	) {
 		if (!id) return
 		connectingData =
diff --git a/frontend/src/lib/components/flows/propPicker/InputPickerInner.svelte b/frontend/src/lib/components/flows/propPicker/InputPickerInner.svelte
index 816804b007..11473aca9c 100644
--- a/frontend/src/lib/components/flows/propPicker/InputPickerInner.svelte
+++ b/frontend/src/lib/components/flows/propPicker/InputPickerInner.svelte
@@ -45,7 +45,7 @@
 			{#if testSteps?.isArgManuallySet(id, key)}
 				<button
 					onclick={() => {
-						testSteps?.evalArg(id, key, flowStateStore, flowStore?.val, previewArgs?.val)
+						testSteps?.evalArg(id, key, flowStateStore?.val, flowStore?.val, previewArgs?.val)
 					}}
 					title="Re-evaluate input"
 					class="-my-1 ml-0.5 hover:text-primary dark:hover:text-primary dark:text-gray-500 text-gray-300"
diff --git a/frontend/src/lib/components/stepHistoryLoader.svelte.ts b/frontend/src/lib/components/stepHistoryLoader.svelte.ts
index fe0c5f1bad..c82bf78b94 100644
--- a/frontend/src/lib/components/stepHistoryLoader.svelte.ts
+++ b/frontend/src/lib/components/stepHistoryLoader.svelte.ts
@@ -4,6 +4,7 @@ import { JobService, type Flow, type FlowModule } from '$lib/gen'
 import { dfs } from './flows/dfs'
 import { getContext, setContext } from 'svelte'
 import pLimit from 'p-limit'
+import type { StateStore } from '$lib/utils'
 
 export type stepState = {
 	initial?: boolean
@@ -75,7 +76,7 @@ export class StepHistoryLoader {
 
 	async loadIndividualStepsStates(
 		flow: Flow,
-		flowStateStore: FlowState,
+		flowStateStore: StateStore<FlowState>,
 		workspaceId: string,
 		initialPath: string,
 		path: string
diff --git a/frontend/src/lib/components/tutorials/FlowBuilderTutorialBranchOne.svelte b/frontend/src/lib/components/tutorials/FlowBuilderTutorialBranchOne.svelte
index bbb88cb226..d1c7bbe066 100644
--- a/frontend/src/lib/components/tutorials/FlowBuilderTutorialBranchOne.svelte
+++ b/frontend/src/lib/components/tutorials/FlowBuilderTutorialBranchOne.svelte
@@ -23,7 +23,7 @@
 	on:error
 	on:skipAll
 	getSteps={(driver, options) => {
-		const id = nextId(flowStateStore, flowStore.val)
+		const id = nextId(flowStateStore.val, flowStore.val)
 		const index = options?.indexToInsertAt ?? flowStore.val.value.modules.length
 
 		const steps = [
diff --git a/frontend/src/lib/components/tutorials/FlowBuilderTutorialForLoop.svelte b/frontend/src/lib/components/tutorials/FlowBuilderTutorialForLoop.svelte
index c5cad5df36..f744dcf08d 100644
--- a/frontend/src/lib/components/tutorials/FlowBuilderTutorialForLoop.svelte
+++ b/frontend/src/lib/components/tutorials/FlowBuilderTutorialForLoop.svelte
@@ -26,7 +26,7 @@
 	on:error
 	on:skipAll
 	getSteps={(driver, options) => {
-		const id = nextId(flowStateStore, flowStore.val)
+		const id = nextId(flowStateStore.val, flowStore.val)
 		const index = options?.indexToInsertAt ?? flowStore.val.value.modules.length
 
 		let tempId = ''
@@ -109,7 +109,7 @@
 						'We can refer to the result of previous steps using the results object: results.a or use static values like [1,2,3] in this case.',
 					onNextClick: () => {
 						updateFlowModuleById(flowStore.val, id, (module) => {
-							const newId = nextId(flowStateStore, flowStore.val)
+							const newId = nextId(flowStateStore.val, flowStore.val)
 							tempId = newId
 
 							if (module.value.type === 'forloopflow') {