diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index ca69ce38d..31ed330a7 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -18,12 +18,17 @@
 	const dispatch = createEventDispatcher();
 
 	export let id = '';
+
 	export let save = false;
+	export let run = true;
 
 	export let token;
 	export let lang = '';
 	export let code = '';
-	export let allow_execution = true;
+
+	export let className = 'my-2';
+	export let editorClassName = '';
+	export let stickyButtonsClassName = 'top-8';
 
 	let _code = '';
 	$: if (code) {
@@ -297,7 +302,7 @@ __builtins__.input = input`);
 </script>
 
 <div>
-	<div class="relative my-2 flex flex-col rounded-lg" dir="ltr">
+	<div class="relative {className} flex flex-col rounded-lg" dir="ltr">
 		{#if lang === 'mermaid'}
 			{#if mermaidHtml}
 				<SvgPanZoom
@@ -314,13 +319,13 @@ __builtins__.input = input`);
 			</div>
 
 			<div
-				class="sticky top-8 mb-1 py-1 pr-2.5 flex items-center justify-end z-10 text-xs text-black dark:text-white"
+				class="sticky {stickyButtonsClassName} mb-1 py-1 pr-2.5 flex items-center justify-end z-10 text-xs text-black dark:text-white"
 			>
 				<div class="flex items-center gap-0.5 translate-y-[1px]">
 					{#if lang.toLowerCase() === 'python' || lang.toLowerCase() === 'py' || (lang === '' && checkPythonCode(code))}
 						{#if executing}
 							<div class="run-code-button bg-none border-none p-1 cursor-not-allowed">Running</div>
-						{:else if allow_execution}
+						{:else if run}
 							<button
 								class="run-code-button bg-none border-none bg-gray-50 hover:bg-gray-100 dark:bg-gray-850 dark:hover:bg-gray-800 transition rounded-md px-1.5 py-0.5"
 								on:click={async () => {
@@ -349,9 +354,11 @@ __builtins__.input = input`);
 			</div>
 
 			<div
-				class="language-{lang} rounded-t-lg -mt-8 {executing || stdout || stderr || result
-					? ''
-					: 'rounded-b-lg'} overflow-hidden"
+				class="language-{lang} rounded-t-lg -mt-8 {editorClassName
+					? editorClassName
+					: executing || stdout || stderr || result
+						? ''
+						: 'rounded-b-lg'} overflow-hidden"
 			>
 				<div class=" pt-7 bg-gray-50 dark:bg-gray-850"></div>
 				<CodeEditor
diff --git a/src/lib/components/chat/Messages/CodeExecutionModal.svelte b/src/lib/components/chat/Messages/CodeExecutionModal.svelte
index 5ef74736d..1be9391d8 100644
--- a/src/lib/components/chat/Messages/CodeExecutionModal.svelte
+++ b/src/lib/components/chat/Messages/CodeExecutionModal.svelte
@@ -2,6 +2,8 @@
 	import { getContext } from 'svelte';
 	import CodeBlock from './CodeBlock.svelte';
 	import Modal from '$lib/components/common/Modal.svelte';
+	import Spinner from '$lib/components/common/Spinner.svelte';
+	import Badge from '$lib/components/common/Badge.svelte';
 	const i18n = getContext('i18n');
 
 	export let show = false;
@@ -11,22 +13,34 @@
 <Modal size="lg" bind:show>
 	<div>
 		<div class="flex justify-between dark:text-gray-300 px-5 pt-4 pb-2">
-			<div class="text-lg font-medium self-center capitalize">
-				{#if codeExecution?.status == 'OK'}
-					&#x2705; <!-- Checkmark -->
-				{:else if codeExecution?.status == 'ERROR'}
-					&#x274C; <!-- X mark -->
-				{:else if codeExecution?.status == 'PENDING'}
-					&#x23F3; <!-- Hourglass -->
-				{:else}
-					&#x2049;&#xFE0F; <!-- Interrobang -->
+			<div class="text-lg font-medium self-center flex flex-col gap-0.5 capitalize">
+				{#if codeExecution?.result}
+					<div>
+						{#if codeExecution.result?.error}
+							<Badge type="error" content="error" />
+						{:else if codeExecution.result?.output}
+							<Badge type="success" content="success" />
+						{:else}
+							<Badge type="warning" content="incomplete" />
+						{/if}
+					</div>
 				{/if}
 
-				{#if codeExecution?.name}
-					{$i18n.t('Code execution')}: {codeExecution?.name}
-				{:else}
-					{$i18n.t('Code execution')}
-				{/if}
+				<div class="flex gap-2 items-center">
+					{#if !codeExecution?.result}
+						<div>
+							<Spinner className="size-4" />
+						</div>
+					{/if}
+
+					<div>
+						{#if codeExecution?.name}
+							{$i18n.t('Code execution')}: {codeExecution?.name}
+						{:else}
+							{$i18n.t('Code execution')}
+						{/if}
+					</div>
+				</div>
 			</div>
 			<button
 				class="self-center"
@@ -48,62 +62,50 @@
 			</button>
 		</div>
 
-		<div class="flex flex-col md:flex-row w-full px-6 pb-5 md:space-x-4">
+		<div class="flex flex-col md:flex-row w-full px-4 pb-5">
 			<div
 				class="flex flex-col w-full dark:text-gray-200 overflow-y-scroll max-h-[22rem] scrollbar-hidden"
 			>
 				<div class="flex flex-col w-full">
-					<div class="text-sm font-medium dark:text-gray-300">
-						{$i18n.t('Code')}
-					</div>
-
 					<CodeBlock
-						id="codeexec-{codeExecution?.id}-code"
-						lang={codeExecution?.language}
-						code={codeExecution?.code}
-						allow_execution={false}
+						id="code-exec-{codeExecution?.id}-code"
+						lang={codeExecution?.language ?? ''}
+						code={codeExecution?.code ?? ''}
+						className=""
+						editorClassName={codeExecution?.result &&
+						(codeExecution?.result?.error || codeExecution?.result?.output)
+							? 'rounded-b-none'
+							: ''}
+						stickyButtonsClassName="top-0"
+						run={false}
 					/>
 				</div>
-				{#if codeExecution?.error}
-					<div class="flex flex-col w-full">
-						<hr class=" dark:border-gray-850 my-3" />
-						<div class="text-sm dark:text-gray-400">
-							{$i18n.t('Error')}
-						</div>
 
-						<CodeBlock
-							id="codeexec-{codeExecution?.id}-error"
-							lang=""
-							code={codeExecution?.error}
-							allow_execution={false}
-						/>
+				{#if codeExecution?.result && (codeExecution?.result?.error || codeExecution?.result?.output)}
+					<div class="dark:bg-[#202123] dark:text-white px-4 py-4 rounded-b-lg flex flex-col gap-3">
+						{#if codeExecution?.result?.error}
+							<div>
+								<div class=" text-gray-500 text-xs mb-1">{$i18n.t('ERROR')}</div>
+								<div class="text-sm">{codeExecution?.result?.error}</div>
+							</div>
+						{/if}
+						{#if codeExecution?.result?.output}
+							<div>
+								<div class=" text-gray-500 text-xs mb-1">{$i18n.t('OUTPUT')}</div>
+								<div class="text-sm">{codeExecution?.result?.output}</div>
+							</div>
+						{/if}
 					</div>
 				{/if}
-				{#if codeExecution?.output}
+				{#if codeExecution?.result?.files && codeExecution?.result?.files.length > 0}
 					<div class="flex flex-col w-full">
-						<hr class=" dark:border-gray-850 my-3" />
-						<div class="text-sm dark:text-gray-400">
-							{$i18n.t('Output')}
-						</div>
-
-						<CodeBlock
-							id="codeexec-{codeExecution?.id}-output"
-							lang=""
-							code={codeExecution?.output}
-							allow_execution={false}
-						/>
-					</div>
-				{/if}
-				{#if codeExecution?.files && codeExecution?.files.length > 0}
-					<div class="flex flex-col w-full">
-						<hr class=" dark:border-gray-850 my-3" />
+						<hr class=" dark:border-gray-850 my-2" />
 						<div class=" text-sm font-medium dark:text-gray-300">
 							{$i18n.t('Files')}
 						</div>
 						<ul class="mt-1 list-disc pl-4 text-xs">
-							{#each codeExecution?.files as file}
+							{#each codeExecution?.result?.files as file}
 								<li>
-									&#x1F4BE; <!-- Floppy disk -->
 									<a href={file.url} target="_blank">{file.name}</a>
 								</li>
 							{/each}
diff --git a/src/lib/components/chat/Messages/CodeExecutions.svelte b/src/lib/components/chat/Messages/CodeExecutions.svelte
index ae3d47894..2c53eaeb3 100644
--- a/src/lib/components/chat/Messages/CodeExecutions.svelte
+++ b/src/lib/components/chat/Messages/CodeExecutions.svelte
@@ -1,38 +1,33 @@
 <script lang="ts">
 	import CodeExecutionModal from './CodeExecutionModal.svelte';
 	import Spinner from '$lib/components/common/Spinner.svelte';
+	import Check from '$lib/components/icons/Check.svelte';
+	import XMark from '$lib/components/icons/XMark.svelte';
+	import EllipsisHorizontal from '$lib/components/icons/EllipsisHorizontal.svelte';
 
 	export let codeExecutions = [];
 
 	let selectedCodeExecution = null;
 	let showCodeExecutionModal = false;
+
+	$: if (codeExecutions) {
+		updateSelectedCodeExecution();
+	}
+
+	const updateSelectedCodeExecution = () => {
+		if (selectedCodeExecution) {
+			selectedCodeExecution = codeExecutions.find(
+				(execution) => execution.id === selectedCodeExecution.id
+			);
+		}
+	};
 </script>
 
 <CodeExecutionModal bind:show={showCodeExecutionModal} codeExecution={selectedCodeExecution} />
 
 {#if codeExecutions.length > 0}
 	<div class="mt-1 mb-2 w-full flex gap-1 items-center flex-wrap">
-		{#each codeExecutions.map((execution) => {
-			let error = null;
-			let output = null;
-			let files = [];
-			let status = 'PENDING';
-
-			if (execution.result) {
-				output = execution.result.output;
-				if (execution.result.error) {
-					status = 'ERROR';
-					error = execution.result.error;
-				} else {
-					status = 'OK';
-				}
-				if (execution.result.files) {
-					files = execution.result.files;
-				}
-			}
-
-			return { id: execution.id, name: execution.name, code: execution.code, language: execution.language || '', status: status, error: error, output: output, files: files };
-		}) as execution (execution.id)}
+		{#each codeExecutions as execution (execution.id)}
 			<div class="flex gap-1 text-xs font-semibold">
 				<button
 					class="flex dark:text-gray-300 py-1 px-1 bg-gray-50 hover:bg-gray-100 dark:bg-gray-850 dark:hover:bg-gray-800 transition rounded-xl max-w-96"
@@ -41,21 +36,23 @@
 						showCodeExecutionModal = true;
 					}}
 				>
-					<div class="bg-white dark:bg-gray-700 rounded-full size-4">
-						{#if execution.status == 'OK'}
-							&#x2705; <!-- Checkmark -->
-						{:else if execution.status == 'ERROR'}
-							&#x274C; <!-- X mark -->
-						{:else if execution.status == 'PENDING'}
-							<Spinner className="size-4" />
+					<div
+						class="bg-white dark:bg-gray-700 rounded-full size-4 flex items-center justify-center"
+					>
+						{#if execution?.result}
+							{#if execution.result?.error}
+								<XMark />
+							{:else if execution.result?.output}
+								<Check strokeWidth="3" className="size-3" />
+							{:else}
+								<EllipsisHorizontal />
+							{/if}
 						{:else}
-							&#x2049;&#xFE0F; <!-- Interrobang -->
+							<Spinner className="size-4" />
 						{/if}
 					</div>
 					<div
-						class="flex-1 mx-2 line-clamp-1 code-execution-name {execution.status == 'PENDING'
-							? 'pulse'
-							: ''}"
+						class="flex-1 mx-2 line-clamp-1 code-execution-name {execution?.result ? '' : 'pulse'}"
 					>
 						{execution.name}
 					</div>