diff --git a/app/package.json b/app/package.json
index fd27f521..31d2a3ec 100644
--- a/app/package.json
+++ b/app/package.json
@@ -1,6 +1,6 @@
 {
   "name": "app",
-  "version": "1.57.7",
+  "version": "1.58.0",
   "main": "module/module.js",
   "license": "MIT",
   "scripts": {
diff --git a/app/src/components/AiToolbar.tsx b/app/src/components/AiToolbar.tsx
index 3b1451c6..742ac0e8 100644
--- a/app/src/components/AiToolbar.tsx
+++ b/app/src/components/AiToolbar.tsx
@@ -1,5 +1,5 @@
 import { Button2, IconButton2, Textarea } from "../ui/Shared";
-import { CaretDown, CaretUp, MagicWand } from "phosphor-react";
+import { CaretDown, CaretUp, MagicWand, Stop } from "phosphor-react";
 import cx from "classnames";
 import { t, Trans } from "@lingui/macro";
 import { createExamples } from "../pages/createExamples";
@@ -13,6 +13,9 @@ import {
   usePromptStore,
   useRunAiWithStore,
 } from "../lib/usePromptStore";
+import { getDefaultText } from "../lib/getDefaultText";
+import { useMemo } from "react";
+import { useDoc } from "../lib/useDoc";
 
 function getModeDescription(mode: Mode): string {
   const prompts = createExamples();
@@ -41,7 +44,7 @@ export function AiToolbar() {
   const isOpen = usePromptStore((state) => state.isOpen);
   const currentMode = usePromptStore((state) => state.mode);
   const isRunning = usePromptStore((state) => state.isRunning);
-  const runAiWithStore = useRunAiWithStore();
+  const { runAi, cancelAi } = useRunAiWithStore();
   const diff = usePromptStore((state) => state.diff);
 
   const toggleOpen = () => setIsOpen(!isOpen);
@@ -57,6 +60,11 @@ export function AiToolbar() {
 
   const currentText = usePromptStore((state) => state.currentText);
 
+  const text = useDoc((state) => state.text);
+  const defaultText = useMemo(() => {
+    return getDefaultText();
+  }, []);
+  const isTextEditable = Boolean(text) && text !== defaultText;
   const showAcceptDiffButton = diff && !isRunning;
 
   return (
@@ -72,21 +80,28 @@ export function AiToolbar() {
             })}
           />
           {!showAcceptDiffButton ? (
-            (["prompt", "convert", "edit"] as Mode[]).map((mode) => (
-              <Button2
-                key={mode}
-                color={mode === currentMode && isOpen ? "purple" : "default"}
-                size="xs"
-                onClick={() => handleModeChange(mode)}
-                className={cx({
-                  "dark:hover:bg-neutral-700": mode !== currentMode,
-                  "dark:bg-purple-700 dark:text-purple-100":
-                    mode === currentMode && isOpen,
-                })}
-              >
-                {getModeTitle(mode)}
-              </Button2>
-            ))
+            (["prompt", "convert", "edit"] as Mode[]).map((mode) => {
+              // If the mode is edit and the text is not editable, don't show the button
+              if (mode === "edit" && !isTextEditable) {
+                return null;
+              }
+
+              return (
+                <Button2
+                  key={mode}
+                  color={mode === currentMode && isOpen ? "purple" : "default"}
+                  size="xs"
+                  onClick={() => handleModeChange(mode)}
+                  className={cx("disabled:opacity-50", {
+                    "dark:hover:bg-neutral-700": mode !== currentMode,
+                    "dark:bg-purple-700 dark:text-purple-100":
+                      mode === currentMode && isOpen,
+                  })}
+                >
+                  {getModeTitle(mode)}
+                </Button2>
+              );
+            })
           ) : (
             <span className="text-sm text-purple-600 dark:text-white">
               <Trans>Keep changes?</Trans>
@@ -94,19 +109,33 @@ export function AiToolbar() {
           )}
         </div>
         {!showAcceptDiffButton ? (
-          <IconButton2
-            onClick={toggleOpen}
-            color="default"
-            size="xs"
-            className="flex items-center justify-center dark:bg-neutral-800 dark:text-neutral-400"
-            isLoading={isRunning}
-          >
-            {!isOpen ? (
-              <CaretDown size={16} weight="bold" />
-            ) : (
-              <CaretUp size={16} weight="bold" />
-            )}
-          </IconButton2>
+          <>
+            <div className="relative">
+              <IconButton2
+                onClick={toggleOpen}
+                color="default"
+                size="xs"
+                className="flex items-center justify-center dark:bg-neutral-800 dark:text-neutral-400"
+                isLoading={isRunning}
+              >
+                {!isOpen ? (
+                  <CaretDown size={16} weight="bold" />
+                ) : (
+                  <CaretUp size={16} weight="bold" />
+                )}
+              </IconButton2>
+              {isRunning ? (
+                <IconButton2
+                  color="red"
+                  size="xs"
+                  className="flex items-center justify-center !absolute top-0 left-0 opacity-0 hover:opacity-100"
+                  onClick={cancelAi}
+                >
+                  <Stop size={16} weight="bold" />
+                </IconButton2>
+              ) : null}
+            </div>
+          </>
         ) : (
           <div className="flex space-x-2">
             <Button2 color="green" size="xs" onClick={acceptDiff}>
@@ -135,7 +164,7 @@ export function AiToolbar() {
             onKeyDown={(e) => {
               if (e.key === "Enter" && !e.shiftKey) {
                 e.preventDefault();
-                runAiWithStore();
+                runAi();
               }
             }}
           />
@@ -145,7 +174,7 @@ export function AiToolbar() {
               size="xs"
               className="dark:bg-purple-700 dark:hover:bg-purple-600 dark:text-purple-100"
               disabled={isRunning}
-              onClick={runAiWithStore}
+              onClick={runAi}
               data-session-activity={`Run AI: ${currentMode}`}
             >
               {!isRunning ? t`Submit` : "..."}
diff --git a/app/src/components/ConvertOnPasteOverlay.tsx b/app/src/components/ConvertOnPasteOverlay.tsx
index b02ad270..7b5e9f3c 100644
--- a/app/src/components/ConvertOnPasteOverlay.tsx
+++ b/app/src/components/ConvertOnPasteOverlay.tsx
@@ -50,7 +50,7 @@ export function ConvertOnPasteOverlay() {
  * This is positioned at the bottom across the whole screen
  */
 function Overlay() {
-  const runAiWithStore = useRunAiWithStore();
+  const { runAi: runAiWithStore } = useRunAiWithStore();
   const isRunning = usePromptStore((s) => s.isRunning);
   const pasted = useEditorStore((s) => s.userPasted);
   return (
diff --git a/app/src/components/LoadTemplateDialog.tsx b/app/src/components/LoadTemplateDialog.tsx
index f22f5f44..06e9fc28 100644
--- a/app/src/components/LoadTemplateDialog.tsx
+++ b/app/src/components/LoadTemplateDialog.tsx
@@ -61,7 +61,6 @@ export function LoadTemplateDialog() {
         <EditorActionTextButton
           icon={PiShapesDuotone}
           data-session-activity="Load Template: Open Dialog"
-          iconClassName="fill-zinc-500"
         >
           <Trans>Examples</Trans>
         </EditorActionTextButton>
diff --git a/app/src/components/TextEditor.tsx b/app/src/components/TextEditor.tsx
index 2707cc2d..72c6f2ae 100644
--- a/app/src/components/TextEditor.tsx
+++ b/app/src/components/TextEditor.tsx
@@ -55,15 +55,22 @@ export function TextEditor({ extendOptions = {}, ...props }: TextEditorProps) {
         options={{
           ...editorOptions,
           ...extendOptions,
-          renderSideBySide: false,
+          renderSideBySide: true,
         }}
         loading={<Loading />}
+        theme={mode === "light" ? "vs" : "vs-dark"}
+        onMount={(_, monaco) => {
+          monaco.editor.setTheme(mode === "light" ? "vs" : "vs-dark");
+        }}
         wrapperProps={{
           "data-testid": "Editor",
-          className: classNames("bg-white dark:bg-neutral-900", {
-            "overflow-hidden": isDragging,
-            "cursor-wait pointer-events-none opacity-50": convertIsRunning,
-          }),
+          className: classNames(
+            "bg-white dark:bg-neutral-900 dark:text-white",
+            {
+              "overflow-hidden": isDragging,
+              "cursor-wait pointer-events-none opacity-50": convertIsRunning,
+            }
+          ),
         }}
       />
     );
diff --git a/app/src/lib/runAi.ts b/app/src/lib/runAi.ts
index 97125b19..8505deb5 100644
--- a/app/src/lib/runAi.ts
+++ b/app/src/lib/runAi.ts
@@ -39,13 +39,15 @@ function writeToEditor(text: string) {
  * Runs an AI endpoint and streams the response back into the editor
  */
 export async function runAi({
-  prompt,
   endpoint,
+  prompt,
   sid,
+  signal,
 }: {
+  endpoint: string;
   prompt: string;
-  endpoint: "prompt" | "convert" | "edit";
   sid?: string;
+  signal?: AbortSignal;
 }) {
   let accumulated = "";
 
@@ -65,6 +67,7 @@ export async function runAi({
         Authorization: sid ? `Bearer ${sid}` : "",
       },
       body: JSON.stringify({ prompt, document: useDoc.getState().text }),
+      signal,
     })
       .then((response) => {
         if (response.ok && response.body) {
@@ -81,8 +84,8 @@ export async function runAi({
           if (endpoint === "edit") {
             // No setup...
           } else {
-            editor.pushUndoStop(); // Make sure you can undo the changes
-            writeToEditor(""); // Clear the editor content
+            editor.pushUndoStop();
+            writeToEditor("");
           }
 
           const processText = ({
@@ -91,42 +94,42 @@ export async function runAi({
           }: ReadableStreamReadResult<Uint8Array>): Promise<void> => {
             if (done) {
               setLastResult(accumulated);
+              resolve(accumulated);
               return Promise.resolve();
             }
 
             const text = decoder.decode(value, { stream: true });
             accumulated += text;
 
-            // If we are editing, we want to set the diff
             if (endpoint === "edit") {
               setDiff(accumulated);
             } else {
-              // If we are not editing, we want to write the text to the editor
               writeToEditor(accumulated);
             }
 
-            return reader.read().then(processText);
+            return reader.read().then(processText).catch(reject);
           };
 
-          reader
-            .read()
-            .then(processText)
-            .finally(() => {
-              editor.pushUndoStop();
-              resolve(accumulated);
-            });
+          reader.read().then(processText).catch(reject);
         } else {
           if (response.status === 429) {
             reject(new Error(RATE_LIMIT_EXCEEDED));
+          } else {
+            reject(
+              new Error(
+                t`Sorry, there was an error converting the text to a flowchart. Try again later.`
+              )
+            );
           }
-          reject(
-            new Error(
-              t`Sorry, there was an error converting the text to a flowchart. Try again later.`
-            )
-          );
         }
       })
-      .catch(reject);
+      .catch((error) => {
+        if (error.name === "AbortError") {
+          reject(new Error("Operation canceled"));
+        } else {
+          reject(error);
+        }
+      });
   });
 }
 
diff --git a/app/src/lib/usePromptStore.ts b/app/src/lib/usePromptStore.ts
index 945e6d2d..bdf6b947 100644
--- a/app/src/lib/usePromptStore.ts
+++ b/app/src/lib/usePromptStore.ts
@@ -1,7 +1,6 @@
 import { create } from "zustand";
 import { RATE_LIMIT_EXCEEDED, runAi } from "./runAi";
-import { isError } from "./helpers";
-import { useCallback, useContext } from "react";
+import { useCallback, useContext, useState } from "react";
 import { useHasProAccess } from "./hooks";
 import { showPaywall } from "./usePaywallModalStore";
 import { t } from "@lingui/macro";
@@ -90,11 +89,14 @@ export function useRunAiWithStore() {
   const hasProAccess = useHasProAccess();
   const customer = useContext(AppContext).customer;
   const sid = customer?.subscription?.id;
+  const [abortController, setAbortController] =
+    useState<AbortController | null>(null);
 
   const handleError = useCallback(
     (error: Error) => {
-      if (!hasProAccess && error.message === RATE_LIMIT_EXCEEDED) {
-        // Show paywall
+      if (error.name === "AbortError") {
+        setError(t`Operation canceled`);
+      } else if (!hasProAccess && error.message === RATE_LIMIT_EXCEEDED) {
         showPaywall({
           title: t`Get Unlimited AI Requests`,
           content: t`You've used all your free AI conversions. Upgrade to Pro for unlimited AI use, custom themes, private sharing, and more. Keep creating amazing flowcharts effortlessly!`,
@@ -113,25 +115,28 @@ export function useRunAiWithStore() {
     [hasProAccess]
   );
 
-  return useCallback(() => {
+  const runAiCallback = useCallback(() => {
     const store = usePromptStore.getState();
     if (store.isRunning) return;
 
-    // close the toolbar
     setIsOpen(false);
     startConvert();
 
-    // If we're creating, we need to unfreeze the editor
     if (store.mode === "convert" || store.mode === "prompt") {
       unfreezeDoc();
     }
 
-    runAi({ endpoint: store.mode, prompt: store.currentText, sid })
-      .catch((err) => {
-        if (isError(err)) handleError(err);
-      })
+    const newAbortController = new AbortController();
+    setAbortController(newAbortController);
+
+    runAi({
+      endpoint: store.mode,
+      prompt: store.currentText,
+      sid,
+      signal: newAbortController.signal,
+    })
+      .catch(handleError)
       .then((result) => {
-        // Just in case there is an error, run repair text on the result
         if (result) {
           const text = repairText(result);
           if (text) {
@@ -142,6 +147,17 @@ export function useRunAiWithStore() {
       .finally(() => {
         stopConvert();
         useEditorStore.setState({ userPasted: "" });
+        setAbortController(null);
       });
   }, [handleError, sid]);
+
+  const cancelAi = useCallback(() => {
+    if (abortController) {
+      abortController.abort();
+      stopConvert();
+      setAbortController(null);
+    }
+  }, [abortController]);
+
+  return { runAi: runAiCallback, cancelAi };
 }