diff --git a/frontend/src/lib/components/FlowGraphViewerStep.svelte b/frontend/src/lib/components/FlowGraphViewerStep.svelte
index 7a13353227..f756132721 100644
--- a/frontend/src/lib/components/FlowGraphViewerStep.svelte
+++ b/frontend/src/lib/components/FlowGraphViewerStep.svelte
@@ -178,7 +178,7 @@
 					<HighlightCode
 						language={stepDetail.value.language}
 						code={stepDetail.value.content}
-						class="whitespace-pre-wrap"
+						className="whitespace-pre-wrap"
 					/>
 				</div>
 				<h3 class="mb-2 mt-4">Lockfile</h3>
diff --git a/frontend/src/lib/components/HighlightCode.svelte b/frontend/src/lib/components/HighlightCode.svelte
index aebd2bf0f1..63e1d8a999 100644
--- a/frontend/src/lib/components/HighlightCode.svelte
+++ b/frontend/src/lib/components/HighlightCode.svelte
@@ -21,10 +21,27 @@
 	import HighlightTheme from './HighlightTheme.svelte'
 	import type { LanguageType } from 'svelte-highlight/languages'
 
-	export let code: string = ''
-	export let language: Script['language'] | 'bunnative' | 'frontend' | undefined
-	export let highlightLanguage: LanguageType<string> | undefined = undefined
-	export let lines = false
+	interface Props {
+		code?: string
+		language: Script['language'] | 'bunnative' | 'frontend' | undefined
+		highlightLanguage?: LanguageType<string> | undefined
+		lines?: boolean
+		className?: string
+		onApplyCode?: () => void
+		showApplyButton?: boolean
+		applyButtonIcon?: typeof ClipboardCopy
+	}
+
+	let {
+		code = '',
+		language,
+		highlightLanguage = undefined,
+		lines = false,
+		className = '',
+		onApplyCode = undefined,
+		showApplyButton = false,
+		applyButtonIcon = undefined
+	}: Props = $props()
 
 	function getLang(lang: Script['language'] | 'bunnative' | 'frontend' | undefined) {
 		switch (lang) {
@@ -71,16 +88,16 @@
 			case 'ansible':
 				return yaml
 			case 'java':
-				return java;
+				return java
 			case 'ruby':
-				return ruby;
-			// for related places search: ADD_NEW_LANG 
+				return ruby
+			// for related places search: ADD_NEW_LANG
 			default:
 				return typescript
 		}
 	}
 
-	$: lang = highlightLanguage ?? getLang(language)
+	const lang = $derived(highlightLanguage ?? getLang(language))
 </script>
 
 <HighlightTheme />
@@ -88,25 +105,39 @@
 <div class="relative">
 	<Button
 		wrapperClasses="absolute top-2 right-2 z-20"
-		on:click={() => copyToClipboard(code)}
+		onclick={() => copyToClipboard(code)}
 		color="light"
 		size="xs2"
 		startIcon={{
 			icon: ClipboardCopy
 		}}
 		iconOnly
+		title="Copy to clipboard"
 	/>
+	{#if showApplyButton}
+		<Button
+			wrapperClasses="absolute top-2 right-10 z-20"
+			onclick={onApplyCode}
+			color="light"
+			size="xs2"
+			startIcon={{
+				icon: applyButtonIcon
+			}}
+			iconOnly
+			title="Apply code"
+		/>
+	{/if}
 	<div class="overflow-x-auto">
 		{#if code?.length < 10000}
 			{#if !lines}
-				<Highlight class="nowrap {$$props.class}" language={lang} {code} />
+				<Highlight class="nowrap {className}" language={lang} {code} />
 			{:else}
-				<Highlight class="nowrap {$$props.class}" language={lang} {code} let:highlighted>
+				<Highlight class="nowrap {className}" language={lang} {code} let:highlighted>
 					<LineNumbers {highlighted} />
 				</Highlight>
 			{/if}
 		{:else}
-			<pre class="overflow-auto max-h-screen text-xs {$$props.class}"
+			<pre class="overflow-auto max-h-screen text-xs {className}"
 				><code class="language-{language}">{code}</code></pre
 			>
 		{/if}
diff --git a/frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte b/frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte
index 7e12a4526e..4b24221c6a 100644
--- a/frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte
+++ b/frontend/src/lib/components/copilot/chat/ContextElementBadge.svelte
@@ -83,7 +83,7 @@
 				<HighlightCode
 					language={contextElement.lang}
 					code={contextElement.content}
-					class="w-full p-2 "
+					className="w-full p-2 "
 				/>
 			</div>
 		{:else if contextElement.type === 'flow_module'}
@@ -92,7 +92,7 @@
 					<HighlightCode
 						language={contextElement.value.language}
 						code={contextElement.value.content}
-						class="w-full p-2 "
+						className="w-full p-2 "
 					/>
 				</div>
 			{:else}
diff --git a/frontend/src/lib/components/copilot/chat/script/CodeDisplay.svelte b/frontend/src/lib/components/copilot/chat/script/CodeDisplay.svelte
index 018fec2bb6..d2f94a1998 100644
--- a/frontend/src/lib/components/copilot/chat/script/CodeDisplay.svelte
+++ b/frontend/src/lib/components/copilot/chat/script/CodeDisplay.svelte
@@ -14,6 +14,8 @@
 		typescript,
 		yaml
 	} from 'svelte-highlight/languages'
+	import { aiChatManager, AIMode } from '../AIChatManager.svelte'
+	import { Check, Play } from 'lucide-svelte'
 
 	const astNode = getAstNode()
 
@@ -79,15 +81,38 @@
 	let language = $derived(
 		(astNode.current.children?.[0]?.properties?.class as string | undefined)?.split('-')[1]
 	)
+
+	// Check if the apply button should be shown
+	let showApplyButton = $derived.by(() => {
+		if (
+			aiChatManager.mode !== AIMode.SCRIPT ||
+			!aiChatManager.scriptEditorApplyCode ||
+			code === aiChatManager.scriptEditorOptions?.code
+		) {
+			return false
+		}
+		return true
+	})
+
+	function handleApplyCode() {
+		if (code && aiChatManager.scriptEditorApplyCode) {
+			aiChatManager.scriptEditorApplyCode(code, { mode: 'apply' })
+		}
+	}
 </script>
 
-<div
-	class="flex flex-col not-prose relative w-full border border-gray-300 dark:border-gray-600 rounded-lg overflow-hidden"
->
-	<HighlightCode
-		class="p-1"
-		code={code ?? ''}
-		highlightLanguage={SMART_LANG_TO_HIGHLIGHT_LANG[getSmartLang(language as string)]}
-		language={undefined}
-	/>
+<div class="flex flex-col gap-0.5 rounded-lg relative not-prose">
+	<div
+		class="relative w-full border border-gray-300 dark:border-gray-600 rounded-lg overflow-hidden"
+	>
+		<HighlightCode
+			className="p-1"
+			code={code ?? ''}
+			highlightLanguage={SMART_LANG_TO_HIGHLIGHT_LANG[getSmartLang(language as string)]}
+			language={undefined}
+			onApplyCode={handleApplyCode}
+			{showApplyButton}
+			applyButtonIcon={aiChatManager.pendingNewCode ? Check : Play}
+		/>
+	</div>
 </div>
diff --git a/frontend/src/lib/components/copilot/chat/script/core.ts b/frontend/src/lib/components/copilot/chat/script/core.ts
index 4380a37f4a..ef6e167a84 100644
--- a/frontend/src/lib/components/copilot/chat/script/core.ts
+++ b/frontend/src/lib/components/copilot/chat/script/core.ts
@@ -31,7 +31,12 @@ const DOCS_CONTEXT_PERCENTAGE = 1
 // percentage of the context window for types of npm packages
 const TYPES_CONTEXT_PERCENTAGE = 1
 // good providers for diff-based edit
-const DIFF_BASED_EDIT_PROVIDERS: AIProvider[] = ['openai', 'anthropic', 'googleai', 'azure_openai']
+export const DIFF_BASED_EDIT_PROVIDERS: AIProvider[] = [
+	'openai',
+	'anthropic',
+	'googleai',
+	'azure_openai'
+]
 
 export function formatResourceTypes(
 	allResourceTypes: ResourceType[],
@@ -218,9 +223,9 @@ export function getLangContext(
 		(isPreprocessor
 			? TS_PREPROCESSOR_INSTRUCTION
 			: TS_RESOURCE_TYPE_SYSTEM +
-			(allowResourcesFetch
-				? `To query the RT namespace, you can use the \`search_resource_types\` tool.\n`
-				: '')) + TS_WINDMILL_CLIENT_CONTEXT
+				(allowResourcesFetch
+					? `To query the RT namespace, you can use the \`search_resource_types\` tool.\n`
+					: '')) + TS_WINDMILL_CLIENT_CONTEXT
 
 	const mainFunctionName = isPreprocessor ? 'preprocessor' : 'main'
 
@@ -248,7 +253,7 @@ export function getLangContext(
 				(isPreprocessor
 					? PYTHON_PREPROCESSOR_INSTRUCTION
 					: PYTHON_RESOURCE_TYPE_SYSTEM +
-					`${allowResourcesFetch ? `\nTo query the available resource types, you can use the \`search_resource_types\` tool.` : ''}`) +
+						`${allowResourcesFetch ? `\nTo query the available resource types, you can use the \`search_resource_types\` tool.` : ''}`) +
 				PYTHON_WINDMILL_CLIENT_CONTEXT
 			)
 		case 'php':
@@ -775,7 +780,7 @@ export async function fetchNpmPackageTypes(
 						typeDefinitions.set(path, code)
 					}
 				},
-				localFile: () => { }
+				localFile: () => {}
 			}
 		})
 
diff --git a/frontend/src/routes/(root)/(logged)/scripts/get/[...hash]/+page.svelte b/frontend/src/routes/(root)/(logged)/scripts/get/[...hash]/+page.svelte
index 3cf7672bd5..3dc0905a08 100644
--- a/frontend/src/routes/(root)/(logged)/scripts/get/[...hash]/+page.svelte
+++ b/frontend/src/routes/(root)/(logged)/scripts/get/[...hash]/+page.svelte
@@ -822,7 +822,7 @@
 									<HighlightCode
 										language={script.language}
 										code={script.content}
-										class="whitespace-pre-wrap"
+										className="whitespace-pre-wrap"
 									/>
 								</div>
 							</TabContent>