diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index ba6d6fc06..7c64e7446 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -5,14 +5,16 @@
 
 	import { v4 as uuidv4 } from 'uuid';
 
-	import { getContext, getAllContexts, onMount } from 'svelte';
+	import { getContext, getAllContexts, onMount, tick, createEventDispatcher } from 'svelte';
 	import { copyToClipboard } from '$lib/utils';
 
 	import 'highlight.js/styles/github-dark.min.css';
 
 	import PyodideWorker from '$lib/workers/pyodide.worker?worker';
+	import CodeEditor from '$lib/components/common/CodeEditor.svelte';
 
 	const i18n = getContext('i18n');
+	const dispatch = createEventDispatcher();
 
 	export let id = '';
 
@@ -20,6 +22,15 @@
 	export let lang = '';
 	export let code = '';
 
+	let _code = '';
+	$: if (code) {
+		updateCode();
+	}
+
+	const updateCode = () => {
+		_code = code;
+	};
+
 	let _token = null;
 
 	let mermaidHtml = null;
@@ -32,6 +43,18 @@
 	let result = null;
 
 	let copied = false;
+	let saved = false;
+
+	const saveCode = () => {
+		saved = true;
+
+		code = _code;
+		dispatch('save', code);
+
+		setTimeout(() => {
+			saved = false;
+		}, 1000);
+	};
 
 	const copyCode = async () => {
 		copied = true;
@@ -233,22 +256,11 @@ __builtins__.input = input`);
 			(async () => {
 				await drawMermaidDiagram();
 			})();
-		} else {
-			// Function to perform the code highlighting
-			const highlightCode = () => {
-				highlightedCode = hljs.highlightAuto(code, hljs.getLanguage(lang)?.aliases).value || code;
-			};
-
-			// Clear the previous timeout if it exists
-			clearTimeout(debounceTimeout);
-			// Set a new timeout to debounce the code highlighting
-			debounceTimeout = setTimeout(highlightCode, 10);
 		}
 	};
 
 	$: if (token) {
 		if (JSON.stringify(token) !== JSON.stringify(_token)) {
-			console.log('hi');
 			_token = token;
 		}
 	}
@@ -295,28 +307,50 @@ __builtins__.input = input`);
 					{:else}
 						<button
 							class="copy-code-button bg-none border-none p-1"
-							on:click={() => {
+							on:click={async () => {
+								code = _code;
+								await tick();
 								executePython(code);
 							}}>{$i18n.t('Run')}</button
 						>
 					{/if}
 				{/if}
+
+				<button class="copy-code-button bg-none border-none p-1" on:click={saveCode}>
+					{saved ? $i18n.t('Saved') : $i18n.t('Save')}
+				</button>
+
 				<button class="copy-code-button bg-none border-none p-1" on:click={copyCode}
 					>{copied ? $i18n.t('Copied') : $i18n.t('Copy')}</button
 				>
 			</div>
 		</div>
 
-		<pre
+		<div
+			class="language-{lang} rounded-t-none {executing || stdout || stderr || result
+				? ''
+				: 'rounded-b-lg'} overflow-hidden"
+		>
+			<CodeEditor
+				value={code}
+				{id}
+				{lang}
+				on:save={() => {
+					saveCode();
+				}}
+				on:change={(e) => {
+					_code = e.detail.value;
+				}}
+			/>
+		</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"
-				>{#if highlightedCode}{@html highlightedCode}{:else}{code}{/if}</code
-			></pre>
+				'border-bottom-left-radius: 0px; border-bottom-right-radius: 0px;'}"><code></code></pre> -->
 
 		<div
 			id="plt-canvas-{id}"
diff --git a/src/lib/components/chat/Messages/ContentRenderer.svelte b/src/lib/components/chat/Messages/ContentRenderer.svelte
index f6393dbe4..6daf6dfa0 100644
--- a/src/lib/components/chat/Messages/ContentRenderer.svelte
+++ b/src/lib/components/chat/Messages/ContentRenderer.svelte
@@ -56,7 +56,14 @@
 </script>
 
 <div bind:this={contentContainerElement}>
-	<Markdown {id} {content} {model} />
+	<Markdown
+		{id}
+		{content}
+		{model}
+		on:update={(e) => {
+			dispatch('update', e.detail);
+		}}
+	/>
 </div>
 
 {#if floatingButtons}
diff --git a/src/lib/components/chat/Messages/Markdown.svelte b/src/lib/components/chat/Messages/Markdown.svelte
index e5373d39d..cdf0d5f5a 100644
--- a/src/lib/components/chat/Messages/Markdown.svelte
+++ b/src/lib/components/chat/Messages/Markdown.svelte
@@ -7,6 +7,9 @@
 	import markedKatexExtension from '$lib/utils/marked/katex-extension';
 
 	import MarkdownTokens from './Markdown/MarkdownTokens.svelte';
+	import { createEventDispatcher } from 'svelte';
+
+	const dispatch = createEventDispatcher();
 
 	export let id;
 	export let content;
@@ -31,5 +34,11 @@
 </script>
 
 {#key id}
-	<MarkdownTokens {tokens} {id} />
+	<MarkdownTokens
+		{tokens}
+		{id}
+		on:update={(e) => {
+			dispatch('update', e.detail);
+		}}
+	/>
 {/key}
diff --git a/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte b/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
index a1df4b67e..931dea598 100644
--- a/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
@@ -1,16 +1,18 @@
 <script lang="ts">
 	import DOMPurify from 'dompurify';
-	import { onMount } from 'svelte';
+	import { createEventDispatcher, onMount } from 'svelte';
 	import { marked, type Token } from 'marked';
 	import { revertSanitizedResponseContent, unescapeHtml } from '$lib/utils';
 
+	import { WEBUI_BASE_URL } from '$lib/constants';
+
 	import CodeBlock from '$lib/components/chat/Messages/CodeBlock.svelte';
 	import MarkdownInlineTokens from '$lib/components/chat/Messages/Markdown/MarkdownInlineTokens.svelte';
 	import KatexRenderer from './KatexRenderer.svelte';
-	import { WEBUI_BASE_URL } from '$lib/constants';
-	import { stringify } from 'postcss';
 	import Collapsible from '$lib/components/common/Collapsible.svelte';
 
+	const dispatch = createEventDispatcher();
+
 	export let id: string;
 	export let tokens: Token[];
 	export let top = true;
@@ -34,6 +36,12 @@
 			{token}
 			lang={token?.lang ?? ''}
 			code={revertSanitizedResponseContent(token?.text ?? '')}
+			on:save={(e) => {
+				dispatch('update', {
+					oldContent: token.text,
+					newContent: e.detail
+				});
+			}}
 		/>
 	{:else if token.type === 'table'}
 		<div class="scrollbar-hidden relative whitespace-nowrap overflow-x-auto max-w-full">
diff --git a/src/lib/components/chat/Messages/ResponseMessage.svelte b/src/lib/components/chat/Messages/ResponseMessage.svelte
index 51f6c2fdb..dbd247153 100644
--- a/src/lib/components/chat/Messages/ResponseMessage.svelte
+++ b/src/lib/components/chat/Messages/ResponseMessage.svelte
@@ -479,6 +479,15 @@
 										id={message.id}
 										content={message.content}
 										{model}
+										on:update={(e) => {
+											const { oldContent, newContent } = e.detail;
+
+											history.messages[message.id].content = history.messages[
+												message.id
+											].content.replace(oldContent, newContent);
+
+											dispatch('update');
+										}}
 										on:explain={(e) => {
 											dispatch(
 												'submit',
diff --git a/src/lib/components/common/CodeEditor.svelte b/src/lib/components/common/CodeEditor.svelte
index 9453c6ce3..f3cf3621a 100644
--- a/src/lib/components/common/CodeEditor.svelte
+++ b/src/lib/components/common/CodeEditor.svelte
@@ -8,6 +8,8 @@
 
 	import { indentUnit } from '@codemirror/language';
 	import { python } from '@codemirror/lang-python';
+	import { javascript } from '@codemirror/lang-javascript';
+
 	import { oneDark } from '@codemirror/theme-one-dark';
 
 	import { onMount, createEventDispatcher, getContext } from 'svelte';
@@ -19,15 +21,41 @@
 
 	export let boilerplate = '';
 	export let value = '';
+	let _value = '';
+
+	$: if (value) {
+		updateValue();
+	}
+
+	const updateValue = () => {
+		_value = value;
+		if (codeEditor) {
+			codeEditor.dispatch({
+				changes: [{ from: 0, to: codeEditor.state.doc.length, insert: _value }]
+			});
+		}
+	};
+
+	export let id = '';
+	export let lang = '';
 
 	let codeEditor;
 
 	let isDarkMode = false;
 	let editorTheme = new Compartment();
 
+	const getLang = () => {
+		if (lang === 'python') {
+			return python();
+		} else if (lang === 'javascript') {
+			return javascript();
+		}
+		return python();
+	};
+
 	export const formatPythonCodeHandler = async () => {
 		if (codeEditor) {
-			const res = await formatPythonCode(value).catch((error) => {
+			const res = await formatPythonCode(_value).catch((error) => {
 				toast.error(error);
 				return null;
 			});
@@ -49,12 +77,13 @@
 	let extensions = [
 		basicSetup,
 		keymap.of([{ key: 'Tab', run: acceptCompletion }, indentWithTab]),
-		python(),
+		getLang(),
 		indentUnit.of('    '),
 		placeholder('Enter your code here...'),
 		EditorView.updateListener.of((e) => {
 			if (e.docChanged) {
-				value = e.state.doc.toString();
+				_value = e.state.doc.toString();
+				dispatch('change', { value: _value });
 			}
 		}),
 		editorTheme.of([])
@@ -66,16 +95,18 @@
 			value = boilerplate;
 		}
 
+		_value = value;
+
 		// Check if html class has dark mode
 		isDarkMode = document.documentElement.classList.contains('dark');
 
 		// python code editor, highlight python code
 		codeEditor = new EditorView({
 			state: EditorState.create({
-				doc: value,
+				doc: _value,
 				extensions: extensions
 			}),
-			parent: document.getElementById('code-textarea')
+			parent: document.getElementById(`code-textarea-${id}`)
 		});
 
 		if (isDarkMode) {
@@ -133,4 +164,4 @@
 	});
 </script>
 
-<div id="code-textarea" class="h-full w-full" />
+<div id="code-textarea-{id}" class="h-full w-full" />
diff --git a/src/lib/components/workspace/Functions/FunctionEditor.svelte b/src/lib/components/workspace/Functions/FunctionEditor.svelte
index 074d0bbe6..3f85c9f2a 100644
--- a/src/lib/components/workspace/Functions/FunctionEditor.svelte
+++ b/src/lib/components/workspace/Functions/FunctionEditor.svelte
@@ -21,6 +21,15 @@
 		description: ''
 	};
 	export let content = '';
+	let _content = '';
+
+	$: if (content) {
+		updateContent();
+	}
+
+	const updateContent = () => {
+		_content = content;
+	};
 
 	$: if (name && !edit && !clone) {
 		id = name.replace(/\s+/g, '_').toLowerCase();
@@ -336,10 +345,14 @@ class Pipe:
 
 				<div class="mb-2 flex-1 overflow-auto h-0 rounded-lg">
 					<CodeEditor
-						bind:value={content}
 						bind:this={codeEditor}
+						value={content}
 						{boilerplate}
+						on:change={(e) => {
+							_content = e.detail.value;
+						}}
 						on:save={() => {
+							content = _content;
 							if (formElement) {
 								formElement.requestSubmit();
 							}
diff --git a/src/lib/components/workspace/Tools/ToolkitEditor.svelte b/src/lib/components/workspace/Tools/ToolkitEditor.svelte
index 718ddb39d..61413821d 100644
--- a/src/lib/components/workspace/Tools/ToolkitEditor.svelte
+++ b/src/lib/components/workspace/Tools/ToolkitEditor.svelte
@@ -22,6 +22,15 @@
 		description: ''
 	};
 	export let content = '';
+	let _content = '';
+
+	$: if (content) {
+		updateContent();
+	}
+
+	const updateContent = () => {
+		_content = content;
+	};
 
 	$: if (name && !edit && !clone) {
 		id = name.replace(/\s+/g, '_').toLowerCase();
@@ -224,10 +233,14 @@ class Tools:
 
 				<div class="mb-2 flex-1 overflow-auto h-0 rounded-lg">
 					<CodeEditor
-						bind:value={content}
 						bind:this={codeEditor}
+						value={content}
 						{boilerplate}
+						on:change={(e) => {
+							_content = e.detail.value;
+						}}
 						on:save={() => {
+							content = _content;
 							if (formElement) {
 								formElement.requestSubmit();
 							}