diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index 5881d109e..fe6c277f1 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -4,7 +4,7 @@
 	import hljs from 'highlight.js';
 	import 'highlight.js/styles/github-dark.min.css';
 	import { loadPyodide } from 'pyodide';
-	import { tick } from 'svelte';
+	import { onMount, tick } from 'svelte';
 	import PyodideWorker from '$lib/workers/pyodide.worker?worker';
 
 	export let id = '';
@@ -12,6 +12,7 @@
 	export let lang = '';
 	export let code = '';
 
+	let highlightedCode = null;
 	let executing = false;
 
 	let stdout = null;
@@ -202,60 +203,60 @@ __builtins__.input = input`);
 		};
 	};
 
-	$: highlightedCode = code ? hljs.highlightAuto(code, hljs.getLanguage(lang)?.aliases).value : '';
+	$: if (code) {
+		highlightedCode = hljs.highlightAuto(code, hljs.getLanguage(lang)?.aliases).value || code;
+	}
 </script>
 
-{#if code}
-	<div class="mb-4" dir="ltr">
-		<div
-			class="flex justify-between bg-[#202123] text-white text-xs px-4 pt-1 pb-0.5 rounded-t-lg overflow-x-auto"
-		>
-			<div class="p-1">{@html lang}</div>
+<div class="mb-4" dir="ltr">
+	<div
+		class="flex justify-between bg-[#202123] text-white text-xs px-4 pt-1 pb-0.5 rounded-t-lg overflow-x-auto"
+	>
+		<div class="p-1">{@html lang}</div>
 
-			<div class="flex items-center">
-				{#if lang === 'python' || (lang === '' && checkPythonCode(code))}
-					{#if executing}
-						<div class="copy-code-button bg-none border-none p-1 cursor-not-allowed">Running</div>
-					{:else}
-						<button
-							class="copy-code-button bg-none border-none p-1"
-							on:click={() => {
-								executePython(code);
-							}}>Run</button
-						>
-					{/if}
+		<div class="flex items-center">
+			{#if lang === 'python' || (lang === '' && checkPythonCode(code))}
+				{#if executing}
+					<div class="copy-code-button bg-none border-none p-1 cursor-not-allowed">Running</div>
+				{:else}
+					<button
+						class="copy-code-button bg-none border-none p-1"
+						on:click={() => {
+							executePython(code);
+						}}>Run</button
+					>
 				{/if}
-				<button class="copy-code-button bg-none border-none p-1" on:click={copyCode}
-					>{copied ? 'Copied' : 'Copy Code'}</button
-				>
-			</div>
+			{/if}
+			<button class="copy-code-button bg-none border-none p-1" on:click={copyCode}
+				>{copied ? 'Copied' : 'Copy Code'}</button
+			>
 		</div>
-
-		<pre
-			class=" hljs p-4 px-5 overflow-x-auto"
-			style="border-top-left-radius: 0px; border-top-right-radius: 0px; {(executing ||
-				stdout ||
-				stderr ||
-				result) &&
-				'border-bottom-left-radius: 0px; border-bottom-right-radius: 0px;'}"><code
-				class="language-{lang} rounded-t-none whitespace-pre">{@html highlightedCode || code}</code
-			></pre>
-
-		<div
-			id="plt-canvas-{id}"
-			class="bg-[#202123] text-white max-w-full overflow-x-auto scrollbar-hidden"
-		/>
-
-		{#if executing}
-			<div class="bg-[#202123] text-white px-4 py-4 rounded-b-lg">
-				<div class=" text-gray-500 text-xs mb-1">STDOUT/STDERR</div>
-				<div class="text-sm">Running...</div>
-			</div>
-		{:else if stdout || stderr || result}
-			<div class="bg-[#202123] text-white px-4 py-4 rounded-b-lg">
-				<div class=" text-gray-500 text-xs mb-1">STDOUT/STDERR</div>
-				<div class="text-sm">{stdout || stderr || result}</div>
-			</div>
-		{/if}
 	</div>
-{/if}
+
+	<pre
+		class=" hljs p-4 px-5 overflow-x-auto"
+		style="border-top-left-radius: 0px; border-top-right-radius: 0px; {(executing ||
+			stdout ||
+			stderr ||
+			result) &&
+			'border-bottom-left-radius: 0px; border-bottom-right-radius: 0px;'}"><code
+			class="language-{lang} rounded-t-none whitespace-pre">{@html highlightedCode || code}</code
+		></pre>
+
+	<div
+		id="plt-canvas-{id}"
+		class="bg-[#202123] text-white max-w-full overflow-x-auto scrollbar-hidden"
+	/>
+
+	{#if executing}
+		<div class="bg-[#202123] text-white px-4 py-4 rounded-b-lg">
+			<div class=" text-gray-500 text-xs mb-1">STDOUT/STDERR</div>
+			<div class="text-sm">Running...</div>
+		</div>
+	{:else if stdout || stderr || result}
+		<div class="bg-[#202123] text-white px-4 py-4 rounded-b-lg">
+			<div class=" text-gray-500 text-xs mb-1">STDOUT/STDERR</div>
+			<div class="text-sm">{stdout || stderr || result}</div>
+		</div>
+	{/if}
+</div>
diff --git a/src/lib/components/chat/Messages/ResponseMessage.svelte b/src/lib/components/chat/Messages/ResponseMessage.svelte
index 44bcea038..f1693a94f 100644
--- a/src/lib/components/chat/Messages/ResponseMessage.svelte
+++ b/src/lib/components/chat/Messages/ResponseMessage.svelte
@@ -440,8 +440,8 @@
 									{#if token.type === 'code'}
 										<CodeBlock
 											id={`${message.id}-${tokenIdx}`}
-											lang={token.lang}
-											code={revertSanitizedResponseContent(token.text)}
+											lang={token?.lang ?? ''}
+											code={revertSanitizedResponseContent(token?.text ?? '')}
 										/>
 									{:else}
 										{@html marked.parse(token.raw, {