diff --git a/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte b/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
index b8e612b01c..f6f6be464f 100644
--- a/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
+++ b/frontend/src/lib/components/flows/map/FlowModuleSchemaItem.svelte
@@ -15,8 +15,7 @@
 		Square,
 		SkipForward,
 		Voicemail,
-		X,
-		Plug
+		X
 	} from 'lucide-svelte'
 	import { createEventDispatcher, getContext } from 'svelte'
 	import { fade } from 'svelte/transition'
@@ -25,12 +24,12 @@
 	import { twMerge } from 'tailwind-merge'
 	import IdEditorInput from '$lib/components/IdEditorInput.svelte'
 	import { dfs } from '../dfs'
-	import { Drawer, Popup } from '$lib/components/common'
+	import { Drawer } from '$lib/components/common'
 	import DrawerContent from '$lib/components/common/drawer/DrawerContent.svelte'
 	import { getDependeeAndDependentComponents } from '../flowExplorer'
 	import { replaceId } from '../flowStore'
-	import type { PropPickerWrapperContext } from '$lib/components/prop_picker'
-	import ObjectViewer from '$lib/components/propertyPicker/ObjectViewer.svelte'
+	import FlowPropPicker from '$lib/components/flows/propPicker/FlowPropPicker.svelte'
+	import type { PropPickerConfig } from '$lib/components/prop_picker'
 	export let selected: boolean = false
 	export let deletable: boolean = false
 	export let retry: boolean = false
@@ -48,6 +47,8 @@
 	export let concurrency: boolean = false
 	export let retries: number | undefined = undefined
 	export let warningMessage: string | undefined = undefined
+	export let propPickerConfig: PropPickerConfig | undefined = undefined
+	export let pickableIds: Record<string, any> | undefined = undefined
 
 	const { flowInputsStore } = getContext<{ flowInputsStore: Writable<FlowInput | undefined> }>(
 		'FlowGraphContext'
@@ -56,9 +57,6 @@
 	const flowEditorContext = getContext<FlowEditorContext>('FlowEditorContext')
 	const dispatch = createEventDispatcher()
 
-	const { propPickerConfig, pickableIds } =
-		getContext<PropPickerWrapperContext>('PropPickerWrapper')
-
 	const { currentStepStore: copilotCurrentStepStore } =
 		getContext<FlowCopilotContext | undefined>('FlowCopilotContext') || {}
 
@@ -271,34 +269,16 @@ hover:border-blue-700 hover:!visible {hover ? '' : '!hidden'}"
 			{/if}
 		</div>
 	</div>
-	{#if $propPickerConfig && $pickableIds && id && Object.keys($pickableIds).includes(id)}
+	{#if propPickerConfig && pickableIds && id && Object.keys(pickableIds).includes(id)}
 		<div
-			class="absolute -top-[10px] right-[85px]"
+			class="absolute -top-[12px] right-[90px]"
 			on:click|preventDefault|stopPropagation={(e) => {
 				e.preventDefault()
 				e.stopPropagation()
 			}}
 			data-prop-picker
 		>
-			<Popup floatingConfig={{ strategy: 'fixed', placement: 'bottom-start' }}>
-				<svelte:fragment slot="button">
-					<button
-						class="rounded-full h-[20px] w-[20px] trash center-center outline-[1px] outline dark:outline-gray-500 outline-gray-300 duration-150 bg-blue-500 text-white hover:bg-blue-500/80 hover:text-white"
-					>
-						<Plug class="mx-[3px]" size={12} strokeWidth={2} />
-					</button>
-				</svelte:fragment>
-				<ObjectViewer
-					json={{ [id]: $pickableIds[id] }}
-					topBrackets={false}
-					pureViewer={false}
-					prefix="results"
-					on:select={(e) => {
-						$propPickerConfig?.onSelect(e.detail)
-						propPickerConfig.set(undefined)
-					}}
-				/>
-			</Popup>
+			<FlowPropPicker {propPickerConfig} {pickableIds} {id} />
 		</div>
 	{/if}
 	{#if deletable}
@@ -315,7 +295,7 @@ hover:border-blue-700 hover:!visible {hover ? '' : '!hidden'}"
 
 		{#if id !== 'preprocessor'}
 			<button
-				class="absolute -top-[10px] right-[60px] rounded-full h-[20px] w-[20px] trash center-center text-secondary
+				class="absolute -top-[10px] right-[60px] rounded-full h-[20px] w-[20px] center-center text-secondary
 outline-[1px] outline dark:outline-gray-500 outline-gray-300 bg-surface duration-150 hover:bg-blue-400 hover:text-white
  {hover ? '' : '!hidden'}"
 				on:click|preventDefault|stopPropagation={(event) => dispatch('move')}
diff --git a/frontend/src/lib/components/flows/map/MapItem.svelte b/frontend/src/lib/components/flows/map/MapItem.svelte
index 5114758357..bbbcf02b66 100644
--- a/frontend/src/lib/components/flows/map/MapItem.svelte
+++ b/frontend/src/lib/components/flows/map/MapItem.svelte
@@ -11,6 +11,7 @@
 	import { msToSec } from '$lib/utils'
 	import BarsStaggered from '$lib/components/icons/BarsStaggered.svelte'
 	import FlowJobsMenu from './FlowJobsMenu.svelte'
+	import type { PropPickerWrapperContext } from '$lib/components/prop_picker'
 
 	export let mod: FlowModule
 	export let insertable: boolean
@@ -57,6 +58,9 @@
 	function onDelete(event: CustomEvent<MouseEvent>) {
 		dispatch('delete', event)
 	}
+
+	const { propPickerConfig, pickableIds } =
+		getContext<PropPickerWrapperContext>('PropPickerWrapper')
 </script>
 
 {#if mod}
@@ -96,6 +100,8 @@
 		<div class={moving == mod.id ? 'opacity-50' : ''}>
 			{#if mod.value.type === 'forloopflow' || mod.value.type === 'whileloopflow'}
 				<FlowModuleSchemaItem
+					propPickerConfig={$propPickerConfig}
+					pickableIds={$pickableIds}
 					deletable={insertable}
 					label={`${
 						mod.summary || (mod.value.type == 'forloopflow' ? 'For loop' : 'While loop')
@@ -163,6 +169,8 @@
 					{...itemProps}
 					modType={mod.value.type}
 					{bgColor}
+					propPickerConfig={$propPickerConfig}
+					pickableIds={$pickableIds}
 					label={mod.summary ||
 						(mod.id === 'preprocessor'
 							? 'Preprocessor'
diff --git a/frontend/src/lib/components/flows/map/VirtualItem.svelte b/frontend/src/lib/components/flows/map/VirtualItem.svelte
index c69e743d7f..39907a5935 100644
--- a/frontend/src/lib/components/flows/map/VirtualItem.svelte
+++ b/frontend/src/lib/components/flows/map/VirtualItem.svelte
@@ -4,6 +4,8 @@
 	import { classNames } from '$lib/utils'
 	import { createEventDispatcher, getContext } from 'svelte'
 	import type { FlowCopilotContext } from '$lib/components/copilot/flow'
+	import type { PropPickerConfig } from '$lib/components/prop_picker'
+	import FlowPropPicker from '$lib/components/flows/propPicker/FlowPropPicker.svelte'
 
 	export let label: string | undefined = undefined
 	export let bgColor: string = ''
@@ -14,6 +16,8 @@
 	export let borderColor: string | undefined = undefined
 	export let hideId: boolean = false
 	export let preLabel: string | undefined = undefined
+	export let propPickerConfig: PropPickerConfig | undefined = undefined
+	export let pickableIds: Record<string, any> | undefined = undefined
 
 	const dispatch = createEventDispatcher<{
 		insert: {
@@ -76,4 +80,9 @@
 			{/if}
 		</div>
 	</div>
+	{#if propPickerConfig && pickableIds && id}
+		<div class="absolute -top-[12px] right-[90px]">
+			<FlowPropPicker {propPickerConfig} {pickableIds} {id} />
+		</div>
+	{/if}
 </div>
diff --git a/frontend/src/lib/components/flows/propPicker/FlowPropPicker.svelte b/frontend/src/lib/components/flows/propPicker/FlowPropPicker.svelte
new file mode 100644
index 0000000000..b6a1e1afb5
--- /dev/null
+++ b/frontend/src/lib/components/flows/propPicker/FlowPropPicker.svelte
@@ -0,0 +1,45 @@
+<script lang="ts">
+	import ObjectViewer from '$lib/components/propertyPicker/ObjectViewer.svelte'
+	import AnimatedButton from '$lib/components/common/button/AnimatedButton.svelte'
+	import { Popup } from '$lib/components/common'
+	import { Plug } from 'lucide-svelte'
+	import { twMerge } from 'tailwind-merge'
+	import type { PropPickerConfig } from '$lib/components/prop_picker'
+
+	export let propPickerConfig: PropPickerConfig | undefined = undefined
+	export let pickableIds: Record<string, any> | undefined = undefined
+	export let id: string | undefined = undefined
+</script>
+
+{#if id}
+	<AnimatedButton
+		animate={propPickerConfig?.insertionMode === 'connect'}
+		wrapperClasses=" h-fit w-fit center-center"
+		baseRadius="9999px"
+	>
+		<Popup floatingConfig={{ strategy: 'fixed', placement: 'bottom-start' }}>
+			<svelte:fragment slot="button">
+				<button
+					class={twMerge(
+						'rounded-full trash center-center h-[20px] w-[20px]',
+						propPickerConfig?.insertionMode !== 'connect'
+							? 'outline-[1px] outline dark:outline-gray-500 outline-gray-300 duration-150 bg-blue-500 hover:bg-blue-500/80 hover:text-white text-secondary-inverse'
+							: 'bg-surface text-blue-500'
+					)}
+				>
+					<Plug class="mx-[3px]" size={12} strokeWidth={2} />
+				</button>
+			</svelte:fragment>
+			<ObjectViewer
+				json={{ [id]: pickableIds?.[id] }}
+				topBrackets={false}
+				pureViewer={false}
+				prefix="results"
+				on:select={(e) => {
+					propPickerConfig?.onSelect(e.detail)
+					propPickerConfig = undefined
+				}}
+			/>
+		</Popup>
+	</AnimatedButton>
+{/if}
diff --git a/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
index 869fb8c00b..2aff56cbff 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
@@ -7,6 +7,7 @@
 	import { getContext } from 'svelte'
 	import type { Writable } from 'svelte/store'
 	import InsertModuleButton from '$lib/components/flows/map/InsertModuleButton.svelte'
+	import type { PropPickerWrapperContext } from '$lib/components/prop_picker'
 
 	export let data: {
 		hasPreprocessor: boolean
@@ -23,6 +24,9 @@
 	const { selectedId } = getContext<{
 		selectedId: Writable<string | undefined>
 	}>('FlowGraphContext')
+
+	const { propPickerConfig, pickableIds } =
+		getContext<PropPickerWrapperContext>('PropPickerWrapper')
 </script>
 
 <NodeWrapper let:darkMode>
@@ -64,5 +68,7 @@
 		on:select={(e) => {
 			data.eventHandlers?.select(e.detail)
 		}}
+		propPickerConfig={$propPickerConfig}
+		pickableIds={$pickableIds}
 	/>
 </NodeWrapper>