diff --git a/frontend/src/lib/components/ArgInput.svelte b/frontend/src/lib/components/ArgInput.svelte
index 1f5bbd4745..140b4f0dfb 100644
--- a/frontend/src/lib/components/ArgInput.svelte
+++ b/frontend/src/lib/components/ArgInput.svelte
@@ -122,6 +122,7 @@
 		workspace?: string | undefined
 		actions?: import('svelte').Snippet
 		innerBottomSnippet?: import('svelte').Snippet
+		fieldHeaderActions?: import('svelte').Snippet
 	}
 
 	let {
@@ -179,7 +180,8 @@
 		computeS3ForceViewerPolicies = undefined,
 		workspace = undefined,
 		actions,
-		innerBottomSnippet
+		innerBottomSnippet,
+		fieldHeaderActions
 	}: Props = $props()
 
 	$effect(() => {
@@ -585,19 +587,22 @@
 		</div>
 	{/if}
 	{#if displayHeader}
-		<FieldHeader
-			prettify={prettifyHeader}
-			label={title && !emptyString(title) ? title : label}
-			{disabled}
-			{required}
-			{type}
-			{contentEncoding}
-			{format}
-			{simpleTooltip}
-			{lightHeader}
-			{displayType}
-			labelClass={css?.label?.class}
-		/>
+		<div class="flex">
+			<FieldHeader
+				prettify={prettifyHeader}
+				label={title && !emptyString(title) ? title : label}
+				{disabled}
+				{required}
+				{type}
+				{contentEncoding}
+				{format}
+				{simpleTooltip}
+				{lightHeader}
+				{displayType}
+				labelClass={css?.label?.class}
+			/>
+			{@render fieldHeaderActions?.()}
+		</div>
 	{/if}
 
 	<div class="flex space-x-1">
diff --git a/frontend/src/lib/components/ModulePreviewForm.svelte b/frontend/src/lib/components/ModulePreviewForm.svelte
index 9c21d1e173..0fc226a3e3 100644
--- a/frontend/src/lib/components/ModulePreviewForm.svelte
+++ b/frontend/src/lib/components/ModulePreviewForm.svelte
@@ -150,20 +150,23 @@
 								nullable={schema.properties[argName].nullable}
 								title={schema.properties[argName].title}
 								placeholder={schema.properties[argName].placeholder}
-							/>
-						{/if}
-						{#if stepsInputArgs?.isArgManuallySet(mod.id, argName)}
-							<div class="pt-6 mt-0.5">
-								<Button
-									on:click={() => {
-										plugIt(argName)
-									}}
-									size="sm"
-									variant="border"
-									color="light"
-									title="Re-evaluate input step"><RefreshCw size={14} /></Button
-								>
-							</div>
+							>
+								{#snippet fieldHeaderActions()}
+									{#if stepsInputArgs?.isArgManuallySet(mod.id, argName)}
+										<div class="ml-auto">
+											<Button
+												on:click={() => {
+													plugIt(argName)
+												}}
+												size="xs2"
+												variant="contained"
+												color="light"
+												title="Re-evaluate input step"><RefreshCw size={12} /></Button
+											>
+										</div>
+									{/if}
+								{/snippet}
+							</ArgInput>
 						{/if}
 					</div>
 				{/if}