diff --git a/frontend/src/lib/components/ArgInfo.svelte b/frontend/src/lib/components/ArgInfo.svelte
index a7128d682a..e8a1850144 100644
--- a/frontend/src/lib/components/ArgInfo.svelte
+++ b/frontend/src/lib/components/ArgInfo.svelte
@@ -24,8 +24,12 @@
 </script>
 
 <Drawer bind:this={jsonViewer} size="800px">
-	<DrawerContent title="See JSON" on:close={jsonViewer.toggleDrawer}>
-		<Highlight language={json} code={JSON.stringify(jsonViewerContent, null, 4)} />
+	<DrawerContent title="Argument Details" on:close={jsonViewer.toggleDrawer}>
+		{#if isString(jsonViewerContent)}
+			<pre>{jsonViewerContent}</pre>
+		{:else}
+			<ObjectViewer pureViewer json={jsonViewerContent} />
+		{/if}
 	</DrawerContent>
 </Drawer>
 
diff --git a/frontend/src/lib/components/FlowViewer.svelte b/frontend/src/lib/components/FlowViewer.svelte
index 979b721792..8536ec1a59 100644
--- a/frontend/src/lib/components/FlowViewer.svelte
+++ b/frontend/src/lib/components/FlowViewer.svelte
@@ -2,7 +2,7 @@
 	import Highlight from 'svelte-highlight'
 	import json from 'svelte-highlight/languages/json'
 	import { faClipboard } from '@fortawesome/free-solid-svg-icons'
-	import type { FlowModule, FlowValue } from '$lib/gen'
+	import type { FlowValue } from '$lib/gen'
 	import { Tab, Tabs, TabContent, Button } from './common'
 	import SchemaViewer from './SchemaViewer.svelte'
 	import FieldHeader from './FieldHeader.svelte'
diff --git a/frontend/src/lib/components/ModuleStatus.svelte b/frontend/src/lib/components/ModuleStatus.svelte
index 96e0244275..2f2960834c 100644
--- a/frontend/src/lib/components/ModuleStatus.svelte
+++ b/frontend/src/lib/components/ModuleStatus.svelte
@@ -30,6 +30,8 @@
 	</span>
 {:else if type == FlowStatusModule.type.SUCCESS}
 	<Badge color="green">Success</Badge>
+{:else if type == FlowStatusModule.type.FAILURE}
+	<Badge color="red">Failure</Badge>
 {/if}
 
 <style>
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index 0018723f4e..50de26a13b 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -12,7 +12,7 @@
 	import FlowCard from '../common/FlowCard.svelte'
 	import FlowModuleHeader from './FlowModuleHeader.svelte'
 	import { flowStateStore } from '../flowState'
-	import { scriptLangToEditorLang } from '$lib/utils'
+	import { schemaToObject, scriptLangToEditorLang } from '$lib/utils'
 	import PropPickerWrapper from '../propPicker/PropPickerWrapper.svelte'
 	import { afterUpdate, getContext, setContext } from 'svelte'
 	import type { FlowEditorContext } from '../types'
@@ -59,7 +59,7 @@
 	$: stepPropPicker = failureModule
 		? {
 				pickableProperties: {
-					flow_input: $previewArgs,
+					flow_input: schemaToObject($flowStore.schema, $previewArgs),
 					priorIds: {},
 					previousId: undefined,
 					hasResume: false
@@ -204,7 +204,10 @@
 						<div class="h-[calc(100%-32px)]">
 							{#if selected === 'inputs'}
 								<div class="h-full overflow-auto">
-									<PropPickerWrapper pickableProperties={stepPropPicker.pickableProperties}>
+									<PropPickerWrapper
+										pickableProperties={stepPropPicker.pickableProperties}
+										error={failureModule}
+									>
 										<SchemaForm
 											schema={$flowStateStore[$selectedId]?.schema ?? {}}
 											inputTransform={true}
diff --git a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
index a37b29acf8..99d49570e6 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
@@ -94,7 +94,12 @@
 			/>
 		{/if}
 	{:else if flowModule.value.type === 'rawscript' || flowModule.value.type === 'script' || flowModule.value.type === 'flow'}
-		<FlowModuleComponent bind:flowModule {parentModule} {previousModule} />
+		<FlowModuleComponent
+			bind:flowModule
+			{parentModule}
+			{previousModule}
+			failureModule={$selectedId === 'failure'}
+		/>
 	{/if}
 {:else if flowModule.value.type === 'forloopflow'}
 	{#each flowModule.value.modules as submodule, index (index)}
diff --git a/frontend/src/lib/components/graph/FlowGraph.svelte b/frontend/src/lib/components/graph/FlowGraph.svelte
index a2ff2f393b..e3754f5f84 100644
--- a/frontend/src/lib/components/graph/FlowGraph.svelte
+++ b/frontend/src/lib/components/graph/FlowGraph.svelte
@@ -34,19 +34,22 @@
 	let edges: Edge[] = []
 	let width: number, height: number
 
+	let loadedFlows: Record<string, FlowModule[]> = {}
+
 	let dispatch = createEventDispatcher()
 
 	$: {
 		width && height && minHeight && selectedNode && flowModuleStates
-		if (modules) {
-			idGenerator = createIdGenerator()
-			createGraph(modules, failureModule)
-		} else {
-			nodes = edges = []
-		}
+		createGraph()
 	}
 
-	function createGraph(modules: FlowModule[], failureModule?: FlowModule) {
+	async function createGraph() {
+		if (modules) {
+			idGenerator = createIdGenerator()
+		} else {
+			nodes = edges = []
+			return
+		}
 		nestedNodes = nodes = []
 
 		nestedNodes.push(createVirtualNode(getParentIds(), 'Flow start'))
diff --git a/frontend/src/lib/utils.ts b/frontend/src/lib/utils.ts
index fe51147939..5c81cf6da1 100644
--- a/frontend/src/lib/utils.ts
+++ b/frontend/src/lib/utils.ts
@@ -366,7 +366,6 @@ export function schemaToTsType(schema: Schema): string {
 
 export function schemaToObject(schema: Schema, args: Record<string, any>): Object {
 	const object = {}
-
 	if (!schema || !schema.properties) {
 		return object
 	}
@@ -375,7 +374,6 @@ export function schemaToObject(schema: Schema, args: Record<string, any>): Objec
 	propKeys.forEach((key: string) => {
 		object[key] = args[key] ?? null
 	})
-
 	return object
 }