feat(gui): background terminals and more (#2303)

* fixes & allow for background terminals to stay running

* status indicators etc

4 files changed, 177 insertions, 92 deletions

diff --git a/ui/src/components/runbooks/List.tsx b/ui/src/components/runbooks/List.tsx
index be6e84f5..72c1b3b3 100644
--- a/ui/src/components/runbooks/List.tsx
+++ b/ui/src/components/runbooks/List.tsx
@@ -14,6 +14,7 @@ import {
   DropdownTrigger,
   DropdownMenu,
   DropdownItem,
+  Badge,
 } from "@nextui-org/react";
 
 import { EllipsisVerticalIcon } from "lucide-react";
@@ -22,32 +23,39 @@ import { DateTime } from "luxon";
 
 import { NotebookPenIcon } from "lucide-react";
 import Runbook from "@/state/runbooks/runbook";
-import { useStore } from "@/state/store";
+import { AtuinState, useStore } from "@/state/store";
 
 const NoteSidebar = () => {
-  const runbooks = useStore((state) => state.runbooks);
-  const refreshRunbooks = useStore((state) => state.refreshRunbooks);
+  const runbooks = useStore((state: AtuinState) => state.runbooks);
+  const refreshRunbooks = useStore(
+    (state: AtuinState) => state.refreshRunbooks,
+  );
 
-  const currentRunbook = useStore((state) => state.currentRunbook);
-  const setCurrentRunbook = useStore((state) => state.setCurrentRunbook);
+  const currentRunbook = useStore((state: AtuinState) => state.currentRunbook);
+  const setCurrentRunbook = useStore(
+    (state: AtuinState) => state.setCurrentRunbook,
+  );
+  const runbookInfo = useStore((state: AtuinState) => state.runbookInfo);
 
   useEffect(() => {
     refreshRunbooks();
   }, []);
 
   return (
-    <div className="min-w-48 h-screen flex flex-col border-r-1">
+    <div className="w-48 flex flex-col border-r-1">
       <div className="overflow-y-auto flex-grow">
         <Listbox
           hideSelectedIcon
-          items={runbooks}
+          items={runbooks.map((runbook) => {
+            return [runbook, runbookInfo[runbook.id]];
+          })}
           variant="flat"
           aria-label="Runbook list"
           selectionMode="single"
           selectedKeys={[currentRunbook]}
           itemClasses={{ base: "data-[selected=true]:bg-gray-200" }}
           topContent={
-            <ButtonGroup>
+            <ButtonGroup className="z-20">
               <Tooltip showArrow content="New Runbook" closeDelay={50}>
                 <Button
                   isIconOnly
@@ -66,7 +74,7 @@ const NoteSidebar = () => {
             </ButtonGroup>
           }
         >
-          {(runbook) => (
+          {([runbook, info]) => (
             <ListboxItem
               key={runbook.id}
               onPress={() => {
@@ -75,14 +83,26 @@ const NoteSidebar = () => {
               textValue={runbook.name || "Untitled"}
               endContent={
                 <Dropdown>
-                  <DropdownTrigger className="bg-transparent">
-                    <Button isIconOnly>
-                      <EllipsisVerticalIcon
-                        size="16px"
-                        className="bg-transparent"
-                      />
-                    </Button>
-                  </DropdownTrigger>
+                  <Badge
+                    content={info?.ptys}
+                    color="primary"
+                    style={
+                      info && info?.ptys > 0
+                        ? {}
+                        : {
+                            display: "none",
+                          }
+                    }
+                  >
+                    <DropdownTrigger className="bg-transparent">
+                      <Button isIconOnly>
+                        <EllipsisVerticalIcon
+                          size="16px"
+                          className="bg-transparent"
+                        />
+                      </Button>
+                    </DropdownTrigger>
+                  </Badge>
                   <DropdownMenu aria-label="Dynamic Actions">
                     <DropdownItem
                       key={"delete"}
diff --git a/ui/src/components/runbooks/editor/Editor.tsx b/ui/src/components/runbooks/editor/Editor.tsx
index b70436de..98a6a282 100644
--- a/ui/src/components/runbooks/editor/Editor.tsx
+++ b/ui/src/components/runbooks/editor/Editor.tsx
@@ -4,6 +4,8 @@ import "@blocknote/core/fonts/inter.css";
 import "@blocknote/mantine/style.css";
 import "./index.css";
 
+import { Spinner } from "@nextui-org/react";
+
 import {
   BlockNoteSchema,
   BlockNoteEditor,
@@ -102,7 +104,7 @@ export default function Editor() {
 
   const debouncedOnChange = useDebounceCallback(onChange, 1000);
 
-  const fetchName = (): String => {
+  const fetchName = (): string => {
     // Infer the title from the first text block
 
     let blocks = editor.document;
@@ -119,7 +121,11 @@ export default function Editor() {
   };
 
   if (editor === undefined) {
-    return "Loading content...";
+    return (
+      <div className="flex w-full h-full flex-col justify-center items-center">
+        <Spinner />
+      </div>
+    );
   }
 
   // Renders the editor instance.
diff --git a/ui/src/components/runbooks/editor/blocks/RunBlock/index.tsx b/ui/src/components/runbooks/editor/blocks/RunBlock/index.tsx
index c7386806..15653611 100644
--- a/ui/src/components/runbooks/editor/blocks/RunBlock/index.tsx
+++ b/ui/src/components/runbooks/editor/blocks/RunBlock/index.tsx
@@ -1,7 +1,9 @@
+import React from "react";
 import { createReactBlockSpec } from "@blocknote/react";
 import "./index.css";
 
 import CodeMirror from "@uiw/react-codemirror";
+import { keymap } from "@codemirror/view";
 import { langs } from "@uiw/codemirror-extensions-langs";
 
 import { Play, Square } from "lucide-react";
@@ -12,58 +14,81 @@ import { invoke } from "@tauri-apps/api/core";
 import Terminal from "./terminal.tsx";
 
 import "@xterm/xterm/css/xterm.css";
+import { AtuinState, useStore } from "@/state/store.ts";
 
 interface RunBlockProps {
   onChange: (val: string) => void;
-  onPlay?: () => void;
-  onStop?: () => void;
+  onRun?: (pty: string) => void;
+  onStop?: (pty: string) => void;
   id: string;
   code: string;
   type: string;
+  pty: string;
   isEditable: boolean;
 }
 
 const RunBlock = ({
   onChange,
-  onPlay,
   id,
   code,
   isEditable,
+  onRun,
+  onStop,
+  pty,
 }: RunBlockProps) => {
-  console.log(code);
-  const [isRunning, setIsRunning] = useState(false);
-  const [showTerminal, setShowTerminal] = useState(false);
   const [value, setValue] = useState<String>(code);
+  const cleanupPtyTerm = useStore((store: AtuinState) => store.cleanupPtyTerm);
+  const terminals = useStore((store: AtuinState) => store.terminals);
 
-  const [pty, setPty] = useState<string | null>(null);
+  const [currentRunbook, incRunbookPty, decRunbookPty] = useStore(
+    (store: AtuinState) => [
+      store.currentRunbook,
+      store.incRunbookPty,
+      store.decRunbookPty,
+    ],
+  );
+
+  const isRunning = pty !== null;
 
-  const handleToggle = async (event: any) => {
-    event.stopPropagation();
+  const handleToggle = async (event: any | null) => {
+    if (event) event.stopPropagation();
 
     // If there's no code, don't do anything
     if (!value) return;
 
-    setIsRunning(!isRunning);
-    setShowTerminal(!isRunning);
-
     if (isRunning) {
-      // send sigkill
-      console.log("sending sigkill");
       await invoke("pty_kill", { pid: pty });
+
+      terminals[pty].terminal.dispose();
+      cleanupPtyTerm(pty);
+
+      if (onStop) onStop(pty);
+      decRunbookPty(currentRunbook);
     }
 
     if (!isRunning) {
-      if (onPlay) onPlay();
-
       let pty = await invoke<string>("pty_open");
-      setPty(pty);
-      console.log(pty);
+      if (onRun) onRun(pty);
+
+      incRunbookPty(currentRunbook);
 
       let val = !value.endsWith("\n") ? value + "\r\n" : value;
       await invoke("pty_write", { pid: pty, data: val });
     }
   };
 
+  const handleCmdEnter = (view) => {
+    handleToggle(null);
+    return true;
+  };
+
+  const customKeymap = keymap.of([
+    {
+      key: "Mod-Enter",
+      run: handleCmdEnter,
+    },
+  ]);
+
   return (
     <div className="w-full !max-w-full !outline-none overflow-none">
       <div className="flex flex-row items-start">
@@ -96,12 +121,12 @@ const RunBlock = ({
               setValue(val);
               onChange(val);
             }}
-            extensions={[...extensions(), langs.shell()]}
+            extensions={[customKeymap, ...extensions(), langs.shell()]}
             basicSetup={false}
           />
           <div
             className={`overflow-hidden transition-all duration-300 ease-in-out min-w-0 ${
-              showTerminal ? "block" : "hidden"
+              isRunning ? "block" : "hidden"
             }`}
           >
             {pty && <Terminal pty={pty} />}
@@ -120,6 +145,7 @@ export default createReactBlockSpec(
         default: "bash",
       },
       code: { default: "" },
+      pty: { default: null },
     },
     content: "none",
   },
@@ -130,7 +156,18 @@ export default createReactBlockSpec(
         editor.updateBlock(block, {
           props: { ...block.props, code: val },
         });
-        console.log(block.props);
+      };
+
+      const onRun = (pty: string) => {
+        editor.updateBlock(block, {
+          props: { ...block.props, pty: pty },
+        });
+      };
+
+      const onStop = (pty: string) => {
+        editor.updateBlock(block, {
+          props: { ...block.props, pty: null },
+        });
       };
 
       return (
@@ -139,7 +176,10 @@ export default createReactBlockSpec(
           id={block?.id}
           code={block.props.code}
           type={block.props.type}
+          pty={block.props.pty}
           isEditable={editor.isEditable}
+          onRun={onRun}
+          onStop={onStop}
         />
       );
     },
diff --git a/ui/src/components/runbooks/editor/blocks/RunBlock/terminal.tsx b/ui/src/components/runbooks/editor/blocks/RunBlock/terminal.tsx
index fa203fc9..cb490887 100644
--- a/ui/src/components/runbooks/editor/blocks/RunBlock/terminal.tsx
+++ b/ui/src/components/runbooks/editor/blocks/RunBlock/terminal.tsx
@@ -1,79 +1,98 @@
-/*
-export const openTerm = (pty: string, id: string) => {
-  const term = new Terminal({
-    fontSize: 12,
-    fontFamily: "Courier New",
-  });
-
-  let element = document.getElementById(id);
-  term.open(element);
-
-  //term.onResize(onResize(pty));
+import { useState, useEffect, useRef } from "react";
+import { listen } from "@tauri-apps/api/event";
+import "@xterm/xterm/css/xterm.css";
+import { useStore } from "@/state/store";
 
-  //const fitAddon = new FitAddon();
-  //term.loadAddon(fitAddon);
-  //term.loadAddon(new WebglAddon());
+const usePersistentTerminal = (pty: string) => {
+  const newPtyTerm = useStore((store) => store.newPtyTerm);
+  const terminals = useStore((store) => store.terminals);
+  const [isReady, setIsReady] = useState(false);
 
-  /*
-  const onSize = (e) => {
-    e.stopPropagation();
-    fitAddon.fit();
-  };
-  fitAddon.fit();
+  useEffect(() => {
+    if (!terminals.hasOwnProperty(pty)) {
+      // create a new terminal and store it in the store.
+      // this means we can resume the same instance even across mount/dismount
+      newPtyTerm(pty);
+    }
 
-  window.addEventListener("resize", onSize, false);
-  */
+    setIsReady(true);
 
-import { useEffect, useRef } from "react";
-import { Terminal } from "@xterm/xterm";
-import { listen } from "@tauri-apps/api/event";
-import { FitAddon } from "@xterm/addon-fit";
-import { WebglAddon } from "@xterm/addon-webgl";
-import "@xterm/xterm/css/xterm.css";
-import { invoke } from "@tauri-apps/api/core";
+    return () => {
+      // We don't dispose of the terminal when the component unmounts
+    };
+  }, [pty, terminals, newPtyTerm]);
 
-const onResize = (pty: string) => async (size: any) => {
-  await invoke("pty_resize", {
-    pid: pty,
-    cols: size.cols,
-    rows: size.rows,
-  });
+  return { terminalData: terminals[pty], isReady };
 };
 
 const TerminalComponent = ({ pty }: any) => {
   const terminalRef = useRef(null);
+  const { terminalData, isReady } = usePersistentTerminal(pty);
+  const [isAttached, setIsAttached] = useState(false);
+  const cleanupListenerRef = useRef<(() => void) | null>(null);
 
   useEffect(() => {
+    // no pty? no terminal
     if (pty == null) return;
-    if (!terminalRef.current) return;
-
-    const terminal = new Terminal();
-    const fitAddon = new FitAddon();
 
-    terminal.open(terminalRef.current);
-    terminal.loadAddon(new WebglAddon());
-    terminal.loadAddon(fitAddon);
-    terminal.onResize(onResize(pty));
+    // the terminal may still be being created so hold off
+    if (!isReady) return;
 
-    fitAddon.fit();
     const windowResize = () => {
-      fitAddon.fit();
+      if (!terminalData || !terminalData.fitAddon) return;
+
+      terminalData.fitAddon.fit();
     };
 
+    // terminal object needs attaching to a ref to a div
+    if (!isAttached && terminalData && terminalData.terminal) {
+      // If it's never been attached, attach it
+      if (!terminalData.terminal.element && terminalRef.current) {
+        terminalData.terminal.open(terminalRef.current);
+
+        // it might have been previously attached, but need moving elsewhere
+      } else if (terminalData && terminalRef.current) {
+        // @ts-ignore
+        terminalRef.current.appendChild(terminalData.terminal.element);
+      }
+
+      terminalData.fitAddon.fit();
+      setIsAttached(true);
+
+      window.addEventListener("resize", windowResize);
+    }
+
     listen(`pty-${pty}`, (event: any) => {
-      terminal.write(event.payload);
-    }).then(() => {
-      console.log("Listening for pty events");
+      terminalData.terminal.write(event.payload);
+    }).then((ul) => {
+      cleanupListenerRef.current = ul;
     });
 
-    window.addEventListener("resize", windowResize);
-
     // Customize further as needed
     return () => {
-      terminal.dispose();
+      if (
+        terminalData &&
+        terminalData.terminal &&
+        terminalData.terminal.element
+      ) {
+        // Instead of removing, we just detach
+        if (terminalData.terminal.element.parentElement) {
+          terminalData.terminal.element.parentElement.removeChild(
+            terminalData.terminal.element,
+          );
+        }
+        setIsAttached(false);
+      }
+
+      if (cleanupListenerRef.current) {
+        cleanupListenerRef.current();
+      }
+
       window.removeEventListener("resize", windowResize);
     };
-  }, [pty]);
+  }, [terminalData, isReady]);
+
+  if (!isReady) return null;
 
   return (
     <div className="!max-w-full min-w-0 overflow-hidden" ref={terminalRef} />