Ability to stop AI Process (#750)

tone-row · Oct 9, 2024 · b6be165 · b6be165
1 parent 87fb751
commit b6be165
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 50 deletions.
diff --git 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";
@@ -44,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);
@@ -109,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}>
@@ -150,7 +164,7 @@ export function AiToolbar() {
             onKeyDown={(e) => {
               if (e.key === "Enter" && !e.shiftKey) {
                 e.preventDefault();
-                runAiWithStore();
+                runAi();
               }
             }}
           />
@@ -160,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
@@ -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/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
@@ -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 };
 }