From 224bad1fc23630bd14c1789a6dd693179aa702c9 Mon Sep 17 00:00:00 2001
From: Guilhem <guilhemlemouel@gmail.com>
Date: Tue, 1 Jul 2025 22:41:26 +0100
Subject: [PATCH] fix(frontend): improve step job load (#6109)

* fix(frontend): improve step job load

* nit
---
 .../ModulePreviewResultViewer.svelte          | 24 +++---------
 frontend/src/lib/components/ModuleTest.svelte |  7 +++-
 .../flows/content/FlowModuleComponent.svelte  |  6 ++-
 .../flows/map/FlowModuleSchemaItem.svelte     | 17 +++------
 .../flows/propPicker/OutputPickerInner.svelte | 38 ++++++++++++-------
 5 files changed, 46 insertions(+), 46 deletions(-)

diff --git a/frontend/src/lib/components/ModulePreviewResultViewer.svelte b/frontend/src/lib/components/ModulePreviewResultViewer.svelte
index 3bbe230375..4d4d6a10bf 100644
--- a/frontend/src/lib/components/ModulePreviewResultViewer.svelte
+++ b/frontend/src/lib/components/ModulePreviewResultViewer.svelte
@@ -24,7 +24,7 @@
 		disableMock?: boolean
 		disableHistory?: boolean
 		onUpdateMock?: (mock: { enabled: boolean; return_value?: unknown }) => void
-		loadingHistory?: boolean
+		loadingJob?: boolean
 	}
 
 	let {
@@ -40,26 +40,15 @@
 		disableMock = false,
 		disableHistory = false,
 		onUpdateMock,
-		loadingHistory = false
+		loadingJob = false
 	}: Props = $props()
 
 	const { testSteps } = getContext<FlowEditorContext>('FlowEditorContext')
 
 	let selectedJob: Job | undefined = $state(undefined)
-	let fetchingLastJob = false
 	let preview: 'mock' | 'job' | undefined = $state(undefined)
 	let jobProgressReset: () => void = $state(() => {})
 
-	let nlastJob = $derived.by(() => {
-		if (testJob && testJob.type === 'CompletedJob') {
-			return { ...testJob, preview: true }
-		}
-		if (lastJob) {
-			return { ...lastJob, preview: false }
-		}
-		return undefined
-	})
-
 	let forceJson = $state(false)
 
 	const logJob = $derived(testJob ?? selectedJob)
@@ -77,7 +66,8 @@
 		{/if}
 
 		<OutputPickerInner
-			lastJob={nlastJob}
+			{lastJob}
+			{testJob}
 			fullResult
 			moduleId={mod.id}
 			closeOnOutsideClick={true}
@@ -86,7 +76,7 @@
 			mock={mod.mock}
 			bind:forceJson
 			bind:selectedJob
-			isLoading={(testIsLoading && !scriptProgress) || fetchingLastJob || loadingHistory}
+			isLoading={testIsLoading || loadingJob}
 			bind:preview
 			path={`path` in mod.value ? mod.value.path : ''}
 			{loopStatus}
@@ -116,9 +106,7 @@
 				duration={logJob?.['duration_ms']}
 				mem={logJob?.['mem_peak']}
 				content={logJob?.logs}
-				isLoading={(testIsLoading && logJob?.['running'] == false) ||
-					fetchingLastJob ||
-					loadingHistory}
+				isLoading={(testIsLoading && logJob?.['running'] == false) || loadingJob}
 				tag={logJob?.tag}
 			/>
 		{/if}
diff --git a/frontend/src/lib/components/ModuleTest.svelte b/frontend/src/lib/components/ModuleTest.svelte
index 7f3213ce32..f12e791003 100644
--- a/frontend/src/lib/components/ModuleTest.svelte
+++ b/frontend/src/lib/components/ModuleTest.svelte
@@ -1,7 +1,7 @@
 <script lang="ts" module>
 	type testModuleState = {
 		loading: boolean
-		cancel?: () => void
+		cancel?: () => Promise<void>
 	}
 
 	let testModulesState = $state<Record<string, testModuleState>>({})
@@ -52,7 +52,10 @@
 		// Not defined if JobProgressBar not loaded
 		if (jobProgressReset) jobProgressReset()
 
-		testModulesState[mod.id].cancel = testJobLoader?.cancelJob
+		testModulesState[mod.id].cancel = async () => {
+			await testJobLoader?.cancelJob()
+			testJob = undefined
+		}
 
 		const val = mod.value
 		// let jobId: string | undefined = undefined
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index 9d5dad9154..88cdbd90f3 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -249,7 +249,9 @@
 		$selectedId && untrack(() => onSelectedIdChange())
 	})
 	$effect(() => {
-		if ($workspaceStore && $pathStore && flowModule?.id && $flowStateStore) {
+		if (testJob && testJob.type === 'CompletedJob') {
+			lastJob = $state.snapshot(testJob)
+		} else if ($workspaceStore && $pathStore && flowModule?.id && $flowStateStore) {
 			untrack(() => getLastJob())
 		}
 	})
@@ -858,7 +860,7 @@
 											{testIsLoading}
 											disableMock={preprocessorModule || failureModule}
 											disableHistory={failureModule}
-											loadingHistory={stepHistoryLoader?.stepStates[flowModule.id]?.loadingJobs}
+											loadingJob={stepHistoryLoader?.stepStates[flowModule.id]?.loadingJobs}
 										/>
 									</Pane>
 								{/if}
diff --git a/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte b/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
index b24e8bb240..87e41fb1d8 100644
--- a/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
+++ b/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
@@ -175,17 +175,11 @@
 	}
 
 	$effect(() => {
-		flowStateStore && $flowStateStore && untrack(() => updateLastJob($flowStateStore))
-	})
-
-	let nlastJob = $derived.by(() => {
-		if (testJob) {
-			return { ...testJob, preview: true }
+		if (testJob && testJob.type === 'CompletedJob') {
+			lastJob = $state.snapshot(testJob)
+		} else if (flowStateStore && $flowStateStore) {
+			untrack(() => updateLastJob($flowStateStore))
 		}
-		if (lastJob) {
-			return { ...lastJob, preview: false }
-		}
-		return undefined
 	})
 
 	let isConnectingCandidate = $derived(
@@ -444,7 +438,8 @@
 						prefix={'results'}
 						connectingData={isConnecting ? connectingData : undefined}
 						{mock}
-						lastJob={nlastJob}
+						{lastJob}
+						{testJob}
 						moduleId={id}
 						onSelect={selectConnection}
 						{onUpdateMock}
diff --git a/frontend/src/lib/components/flows/propPicker/OutputPickerInner.svelte b/frontend/src/lib/components/flows/propPicker/OutputPickerInner.svelte
index c8534d0a00..9786e8b2db 100644
--- a/frontend/src/lib/components/flows/propPicker/OutputPickerInner.svelte
+++ b/frontend/src/lib/components/flows/propPicker/OutputPickerInner.svelte
@@ -56,7 +56,8 @@
 		rightMargin?: boolean
 		disableMock?: boolean
 		disableHistory?: boolean
-		lastJob?: SelectedJob
+		lastJob?: Job
+		testJob?: Job
 		derivedHistoryOpen?: boolean // derived from historyOpen
 		historyOffset?: any
 		clazz?: string
@@ -71,6 +72,7 @@
 
 	let {
 		lastJob = undefined,
+		testJob = undefined,
 		prefix = '',
 		allowCopy = false,
 		connectingData = undefined,
@@ -130,25 +132,25 @@
 		}
 	}
 
-	function selectJob(job: SelectedJob | undefined) {
-		if (job && 'result' in job) {
+	function selectJob(nJob: SelectedJob | undefined) {
+		if (nJob && 'result' in nJob) {
+			selectedJob = nJob
+		} else if (job && 'result' in job) {
 			selectedJob = job
-		} else if (lastJob && 'result' in lastJob) {
-			selectedJob = lastJob
 		} else {
 			selectedJob = undefined
 		}
 	}
 
 	$effect(() => {
-		if (!lastJob || !('result' in lastJob)) {
+		if (!job || !('result' in job)) {
 			return
 		}
-		selectJob(lastJob)
+		selectJob(job)
 
-		if (lastJob.preview && mock?.enabled) {
+		if (job.preview && mock?.enabled) {
 			preview = 'job'
-			lastJob.preview = false
+			job.preview = false
 		}
 	})
 
@@ -254,6 +256,16 @@
 		}
 	})
 
+	let job = $derived.by(() => {
+		if (testJob) {
+			return { ...testJob, preview: testJob.type === 'CompletedJob' }
+		}
+		if (lastJob) {
+			return { ...lastJob, preview: false }
+		}
+		return undefined
+	})
+
 	let popoverHeight = $derived(customHeight ?? (clientHeight > 0 ? clientHeight : 0))
 
 	const isLoadingAndNotMock = $derived(isLoading && !mock?.enabled)
@@ -324,7 +336,7 @@
 											togglePreview('mock')
 											return
 										}
-										if (detail.id === lastJob?.id && !mock?.enabled) {
+										if (detail.id === job?.id && !mock?.enabled) {
 											togglePreview(undefined)
 											return
 										}
@@ -363,7 +375,7 @@
 					<div
 						class={twMerge(
 							'w-grow min-w-0 flex gap-1 items-center h-[27px] rounded-md  group',
-							preview || selectedJob?.id !== lastJob?.id ? 'p-[2px] bg-surface-secondary' : ''
+							preview || selectedJob?.id !== job?.id ? 'p-[2px] bg-surface-secondary' : ''
 						)}
 					>
 						{#if loopStatus?.type === 'self'}
@@ -399,7 +411,7 @@
 								job={selectedJob}
 								class={twMerge(
 									'min-w-16 text-secondary',
-									preview || selectedJob?.id !== lastJob?.id ? 'bg-surface shadow-sm h-[23px]' : ''
+									preview || selectedJob?.id !== job?.id ? 'bg-surface shadow-sm h-[23px]' : ''
 								)}
 							/>
 						{/if}
@@ -686,7 +698,7 @@
 						pureViewer={false}
 					/>
 				{/if}
-			{:else if !lastJob}
+			{:else if !job}
 				<div class="flex flex-col items-center justify-center h-full">
 					<p class="text-xs text-secondary">
 						Test this step to see results{#if !disableMock}