From 91c1058d4c175586d22f9c7822663961fe3d62ae Mon Sep 17 00:00:00 2001
From: tylerslaton <mtslaton1@gmail.com>
Date: Mon, 1 Jul 2024 15:43:09 -0400
Subject: [PATCH] fix: bubble errors up from GPTScript correctly

Signed-off-by: tylerslaton <mtslaton1@gmail.com>
---
 components/script.tsx                |  4 +-
 components/script/chatBar.tsx        |  7 ++-
 components/script/chatBar/upload.tsx |  4 +-
 components/script/messages.tsx       | 33 ++++++++------
 components/script/useChatSocket.tsx  | 23 +++++++++-
 package.json                         |  1 +
 server.mjs                           | 66 ++++++++++++++--------------
 7 files changed, 85 insertions(+), 53 deletions(-)

diff --git a/components/script.tsx b/components/script.tsx
index 0c0c8f6f..ca45122a 100644
--- a/components/script.tsx
+++ b/components/script.tsx
@@ -28,7 +28,7 @@ const Script: React.FC<ScriptProps> = ({ file, className, messagesHeight = 'h-fu
 	const [hasParams, setHasParams] = useState(false);
     const [isEmpty, setIsEmpty] = useState(false);
     const { 
-        socket, connected, running, messages, setMessages, restart, interrupt, generating
+        socket, connected, running, messages, setMessages, restart, interrupt, generating, error
     } = useChatSocket(isEmpty);
 
 	useEffect(() => {
@@ -117,7 +117,7 @@ const Script: React.FC<ScriptProps> = ({ file, className, messagesHeight = 'h-fu
 						handleInputChange={handleInputChange}
 						/>
 					) : (
-						<Messages messages={messages} />
+						<Messages restart={restartScript} messages={messages} />
 					)}
 				</div>
 
diff --git a/components/script/chatBar.tsx b/components/script/chatBar.tsx
index bd83d9c2..f9590061 100644
--- a/components/script/chatBar.tsx
+++ b/components/script/chatBar.tsx
@@ -13,6 +13,7 @@ import { GoIssueReopened, GoSquare, GoSquareFill } from "react-icons/go";
 
 const ChatBar = ({
     generating,
+    disabled = false,
     onBack,
     noChat,
     onInterrupt,
@@ -21,6 +22,7 @@ const ChatBar = ({
     onRestart,
 }: {
     generating: boolean;
+    disabled?: boolean;
     onBack: () => void;
     onInterrupt: () => void;
     onMessageSent: (message: string) => void;
@@ -67,8 +69,9 @@ const ChatBar = ({
                     onPress={onRestart}
                 />
             </Tooltip>
-            <Upload onRestart={onRestart}/>
+            <Upload disabled={disabled} onRestart={onRestart}/>
             <Textarea
+                isDisabled={disabled}
                 id="chatInput"
                 autoComplete="off"
                 placeholder="Ask the chat bot something..."
@@ -91,11 +94,13 @@ const ChatBar = ({
                     startContent={<GoSquareFill className="mr-[1px] text-xl" />}
                     isIconOnly
                     radius="full"
+                    isDisabled={disabled}
                     className="ml-2 my-auto text-lg"
                     onPress={onInterrupt}
                 /> : 
                 <Button
                     startContent={<IoMdSend />}
+                    isDisabled={disabled}
                     isIconOnly
                     radius="full"
                     className="ml-2 my-auto text-lg"
diff --git a/components/script/chatBar/upload.tsx b/components/script/chatBar/upload.tsx
index c729f92c..c41a80a5 100644
--- a/components/script/chatBar/upload.tsx
+++ b/components/script/chatBar/upload.tsx
@@ -18,9 +18,10 @@ import {
 
 interface UploadModalProps {
     onRestart: () => void;
+    disabled: boolean;
 }
 
-const UploadModal = ({onRestart}: UploadModalProps) => {
+const UploadModal = ({onRestart, disabled}: UploadModalProps) => {
     const [isOpen, setIsOpen] = useState(false);
     const [selectedFile, setSelectedFile] = useState<File | null>(null);
     const [files, setFiles] = useState<Dirent[]>([]);
@@ -60,6 +61,7 @@ const UploadModal = ({onRestart}: UploadModalProps) => {
         <>
             <Tooltip content="View and manage your workspace" closeDelay={0.5} placement="top">
                 <Button
+                    isDisabled={disabled}
                     startContent={<GoFileDirectory />}
                     isIconOnly
                     radius="full"
diff --git a/components/script/messages.tsx b/components/script/messages.tsx
index 83bc758e..7a97266f 100644
--- a/components/script/messages.tsx
+++ b/components/script/messages.tsx
@@ -29,7 +29,7 @@ const abbreviate = (name: string) => {
 	return firstLetters.slice(0, 2).join('').toUpperCase();
 }
 
-const Message = ({ message, noAvatar }: { message: Message ,noAvatar?: boolean }) => {
+const Message = ({ message, noAvatar, restart }: { message: Message ,noAvatar?: boolean, restart: () => void }) => {
     switch (message.type) {
         case MessageType.User:
             return (
@@ -45,17 +45,15 @@ const Message = ({ message, noAvatar }: { message: Message ,noAvatar?: boolean }
                     <div className="flex gap-2 w-full">
                         { !noAvatar &&
                             <Tooltip 
-                                content={`Sent from ${message.name || "Main"}`}
+                                content={`Sent from ${message.name || "System"}`}
                                 placement="bottom"
                                 closeDelay={0.5}
                             >
                                 <Avatar
                                     showFallback
-                                    name={abbreviate(message.name || 'Main')} 
-                                    icon={!message.name && <GoSquirrel className="text-xl" />}
+                                    name={abbreviate(message.name || "System")} 
                                     className="w-[40px] cursor-default"
-                                    classNames={{base: "bg-white p-6 text-sm border dark:border-none dark:bg-zinc-900"}}
-                                    color={message.error ? "danger": "default"} 
+                                    classNames={{base: `bg-white p-6 text-sm border dark:border-none dark:bg-zinc-900 ${message.error && "border-danger dark:border-danger"}`}}
                                 />
                             </Tooltip>
                         }
@@ -75,14 +73,21 @@ const Message = ({ message, noAvatar }: { message: Message ,noAvatar?: boolean }
                             {message.error && (
                                 <>
                                     <p className="text-danger text-base pl-4 pb-6">{message.error}</p>
-                                    <Button 
-                                        startContent={<GoIssueReopened className="text-lg"/>}
+                                    <Tooltip 
+                                        content="If you are no longer able to chat, click here to restart the script."
+                                        closeDelay={0.5}
+                                        placement="bottom"
                                         color="danger"
-                                        className="ml-4 mb-6"
-                                        onPress={() => window.location.reload()}
                                     >
-                                        Restart Script
-                                    </Button>
+                                        <Button
+                                            startContent={<GoIssueReopened className="text-lg"/>}
+                                            color="danger"
+                                            className="ml-4 mb-6"
+                                            onPress={restart}
+                                        >
+                                            Restart Script
+                                        </Button>
+                                    </Tooltip>
                                 </>
                             )}
                         </div>
@@ -110,10 +115,10 @@ const Message = ({ message, noAvatar }: { message: Message ,noAvatar?: boolean }
     }
 };
 
-const Messages = ({ messages, noAvatar }: { messages: Message[], noAvatar?: boolean }) => (
+const Messages = ({ messages, noAvatar, restart }: { messages: Message[], noAvatar?: boolean, restart: () => void}) => (
 	<div>
 		{messages.map((message, index) => 
-			<Message key={index} message={message} noAvatar={noAvatar} />
+			<Message key={index} restart={restart} message={message} noAvatar={noAvatar} />
 		)}
 	</div>
 );
diff --git a/components/script/useChatSocket.tsx b/components/script/useChatSocket.tsx
index 54bed74a..604fbf92 100644
--- a/components/script/useChatSocket.tsx
+++ b/components/script/useChatSocket.tsx
@@ -13,6 +13,7 @@ const useChatSocket = (isEmpty?: boolean) => {
     const [messages, setMessages] = useState<Message[]>([]);
     const [generating, setGenerating] = useState(false);
     const [running, setRunning] = useState(false);
+    const [error, setError] = useState<string | null>(null);
 
     // Refs
     const socketRef = useRef<Socket | null>(null);
@@ -26,6 +27,8 @@ const useChatSocket = (isEmpty?: boolean) => {
     useEffect(() => { socketRef.current = socket }, [socket]);
 
     const handleError = useCallback((error: string) => {
+        setGenerating(false);
+        setError(error);
         setMessages((prevMessages) => {
             const updatedMessages = [...prevMessages];
             if (latestBotMessageIndex.current !== -1) {
@@ -33,7 +36,13 @@ const useChatSocket = (isEmpty?: boolean) => {
                 updatedMessages[latestBotMessageIndex.current].error = `${error}`
             } else {
                 // If there are no previous messages, create a new error message
-                updatedMessages.push({ type: MessageType.Bot, message: "An error occured before the first message.", error: error });
+                updatedMessages.push(
+                    {
+                        type: MessageType.Bot, 
+                        message: "The script encountered an error. You can either restart the script or try to continue chatting.", 
+                        error 
+                    }
+                );
             }
             return updatedMessages;
         });
@@ -239,6 +248,7 @@ const useChatSocket = (isEmpty?: boolean) => {
     const restart = useCallback(() => {
         trustedRef.current = {};
         setRunning(false);
+        setError(null)
         setMessages([]);
         trustedRepoPrefixesRef.current = ["github.com/gptscript-ai/context"];
         loadSocket();
@@ -267,7 +277,16 @@ const useChatSocket = (isEmpty?: boolean) => {
         }
     }, [running, messages, isEmpty]);
 
-    return { socket, setSocket, connected, setConnected, messages, setMessages, restart, interrupt, generating, running};
+    return { 
+        error,
+        socket, setSocket, 
+        connected, setConnected, 
+        messages, setMessages, 
+        restart, 
+        interrupt, 
+        generating, 
+        running,
+    };
 };
 
 export default useChatSocket;
\ No newline at end of file
diff --git a/package.json b/package.json
index f83ce94e..31238afe 100644
--- a/package.json
+++ b/package.json
@@ -4,6 +4,7 @@
     "private": true,
     "scripts": {
         "dev": "node server.mjs",
+        "debug" : "node --inspect server.mjs",
         "build": "next build",
         "start": "next start",
         "lint": "next lint"
diff --git a/server.mjs b/server.mjs
index cf8cc8ea..f150154b 100644
--- a/server.mjs
+++ b/server.mjs
@@ -28,8 +28,12 @@ app.prepare().then(() => {
                 await runningScript.close();
                 runningScript = null;
             }
-            await dismount(socket);
-            await mount(file, tool, args, workspace, socket, gptscript);
+            try {
+                dismount(socket);
+                await mount(file, tool, args, workspace, socket, gptscript);
+            } catch (e) {
+                socket.emit("error", e);
+            }
 		});
 	});
 
@@ -54,52 +58,48 @@ const mount = async (file, tool, args, workspace, socket, gptscript) => {
 
 	if (tool) opts.subTool = tool;
 
+    // Start the script
+    runningScript = await gptscript.run(file, opts)
     socket.emit("running");
-	runningScript = await gptscript.run(file, opts);
-	runningScript.on(RunEventType.Event, data => socket.emit('progress', {frame: data, state: runningScript.calls}));
 
-	// Handle prompt events
+    // Handle initial runningScript events
+	runningScript.on(RunEventType.Event, data => socket.emit('progress', {frame: data, state: runningScript.calls}));
 	runningScript.on(RunEventType.Prompt, async (data) => socket.emit("promptRequest", {frame: data, state: runningScript.calls}));
-	socket.on("promptResponse", async (data) => await gptscript.promptResponse(data));
-
-	// Handle confirm events
 	runningScript.on(RunEventType.CallConfirm, (data) => socket.emit("confirmRequest", {frame: data, state: runningScript.calls}));
+	socket.on("promptResponse", async (data) => await gptscript.promptResponse(data));
 	socket.on("confirmResponse", async (data) => await gptscript.confirm(data));
-
     socket.on("interrupt", async() => { if (runningScript) runningScript.close() });
+    socket.on('disconnect', () => { if (runningScript) runningScript.close(); runningScript = null; });
 
-	try {
-		socket.on('disconnect', () => {
-			if (runningScript) runningScript.close();
-			runningScript = null;
-		});
+    // Wait for the run to finish and emit any errors that occur. Specifically look for the "Run has been aborted" error
+    // as this is a marker of an interrupt.
+    runningScript.text().catch((e) => e && e != "Run has been aborted" && socket.emit("error", e));
 
-		socket.on('userMessage', async (message) => {
-			// Remove any previous promptResponse or confirmResponse listeners
-			socket.removeAllListeners("promptResponse");
-			socket.removeAllListeners("confirmResponse");
+    // If the user sends a message, we continue and setup the next chat's event listeners
+    socket.on('userMessage', async (message) => {
+        // Remove any previous promptResponse or confirmResponse listeners
+        socket.removeAllListeners("promptResponse");
+        socket.removeAllListeners("confirmResponse");
 
-			// Start the next chat
-			runningScript = runningScript.nextChat(message);
-			runningScript.on(RunEventType.Event, data => socket.emit('progress', {frame: data, state: runningScript.calls}));
+        // Start the next chat
+        runningScript = runningScript.nextChat(message);
 
-			// Handle prompt events
-			runningScript.on(RunEventType.Prompt, async (data) => socket.emit("promptRequest", {frame: data, state: runningScript.calls}));
-			socket.on("promptResponse", async (data) => await gptscript.promptResponse(data));
+        runningScript.on(RunEventType.Event, data => socket.emit('progress', {frame: data, state: runningScript.calls}));
+        runningScript.on(RunEventType.Prompt, async (data) => socket.emit("promptRequest", {frame: data, state: runningScript.calls}));
+        runningScript.on(RunEventType.CallConfirm, (data) => socket.emit("confirmRequest", {frame: data, state: runningScript.calls}));
+        
+        socket.on("promptResponse", async (data) => await gptscript.promptResponse(data));
+        socket.on("confirmResponse", async (data) => await gptscript.confirm(data));
+
+        // Wait for the run to finish and emit any errors that occur
+        runningScript.text().catch((e) => e && e != "Run has been aborted" && socket.emit("error", e));
+    });
 
-			// Handle confirm events
-			runningScript.on(RunEventType.CallConfirm, (data) => socket.emit("confirmRequest", {frame: data, state: runningScript.calls}));
-			socket.on("confirmResponse", async (data) => await gptscript.confirm(data));
-		});
-	} catch (e) {
-		socket.emit('error', e);
-		console.error(e);
-	}
 }
 
 // Only one script is allowed to run at a time in this system. This function is to dismount and
 // previously mounted listeners.
-const dismount = async (socket) => {
+const dismount = (socket) => {
     socket.removeAllListeners("promptResponse");
     socket.removeAllListeners("confirmResponse");
     socket.removeAllListeners("userMessage");