diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index 305aa84ff..d0dcc8973 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -270,6 +270,10 @@ __builtins__.input = input`);
 		render();
 	}
 
+	$: if (lang) {
+		dispatch('code', { lang });
+	}
+
 	onMount(async () => {
 		console.log('codeblock', lang, code);
 		if (document.documentElement.classList.contains('dark')) {
diff --git a/src/lib/components/chat/Messages/ContentRenderer.svelte b/src/lib/components/chat/Messages/ContentRenderer.svelte
index 5c30c4586..ea612034c 100644
--- a/src/lib/components/chat/Messages/ContentRenderer.svelte
+++ b/src/lib/components/chat/Messages/ContentRenderer.svelte
@@ -21,7 +21,7 @@
 			await tick();
 
 			// Check if the event target is within the content container
-			if (!contentContainerElement.contains(event.target)) return;
+			if (!contentContainerElement?.contains(event.target)) return;
 
 			let selection = window.getSelection();
 
@@ -72,6 +72,7 @@
 		}}
 		on:code={(e) => {
 			const { lang } = e.detail;
+			console.log('code', lang);
 			if (lang === 'html') {
 				showArtifacts.set(true);
 				showControls.set(true);
diff --git a/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte b/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
index be86ff4f1..7fc61c327 100644
--- a/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
@@ -33,11 +33,6 @@
 			<MarkdownInlineTokens id={`${id}-${tokenIdx}-h`} tokens={token.tokens} />
 		</svelte:element>
 	{:else if token.type === 'code'}
-		{#if token.lang === 'html'}
-			{dispatch('code', {
-				lang: token.lang
-			})}
-		{/if}
 		{#if token.raw.includes('```')}
 			<CodeBlock
 				id={`${id}-${tokenIdx}`}
@@ -45,6 +40,9 @@
 				lang={token?.lang ?? ''}
 				code={revertSanitizedResponseContent(token?.text ?? '')}
 				{save}
+				on:code={(e) => {
+					dispatch('code', e.detail);
+				}}
 				on:save={(e) => {
 					dispatch('update', {
 						raw: token.raw,
diff --git a/src/lib/components/common/CodeEditor.svelte b/src/lib/components/common/CodeEditor.svelte
index 38028626b..8c79ac303 100644
--- a/src/lib/components/common/CodeEditor.svelte
+++ b/src/lib/components/common/CodeEditor.svelte
@@ -49,10 +49,7 @@
 	let editorLanguage = new Compartment();
 
 	const getLang = async () => {
-		console.log(languages);
-
 		const language = languages.find((l) => l.alias.includes(lang));
-
 		return await language?.load();
 	};