diff --git a/src/lib/components/DesktopEditor.svelte b/src/lib/components/DesktopEditor.svelte
index 47612015..d73b20d3 100644
--- a/src/lib/components/DesktopEditor.svelte
+++ b/src/lib/components/DesktopEditor.svelte
@@ -170,7 +170,14 @@
       const newText = editorMode === 'code' ? code : mermaid;
       if (newText !== currentText) {
         editor.setScrollTop(0);
-        editor.setValue(newText);
+        editor.pushUndoStop();
+        editor.executeEdits('updateCode', [
+          {
+            range: model.getFullModelRange(),
+            text: newText
+          }
+        ]);
+        editor.pushUndoStop();
         currentText = newText;
         renderAIPromptGutterGlyphIcon();
       }