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

* fixes & allow for background terminals to stay running

* status indicators etc

8 files changed, 427 insertions, 201 deletions

diff --git a/ui/backend/src/run/pty.rs b/ui/backend/src/run/pty.rs
index 819dc7d0..2af617dd 100644
--- a/ui/backend/src/run/pty.rs
+++ b/ui/backend/src/run/pty.rs
@@ -86,9 +86,17 @@ pub(crate) async fn pty_kill(
     pid: uuid::Uuid,
     state: tauri::State<'_, AtuinState>,
 ) -> Result<(), String> {
-    let pty = state.pty_sessions.write().await.remove(&pid).unwrap();
-    pty.kill_child().await.map_err(|e|e.to_string())?;
-    println!("RIP {pid:?}");
+    let pty = state.pty_sessions.write().await.remove(&pid);
+
+    match pty {
+        Some(pty)=>{
+
+            pty.kill_child().await.map_err(|e|e.to_string())?;
+            println!("RIP {pid:?}");
+        }
+        None=>{}
+    }
+
 
     Ok(())
 }
diff --git a/ui/src/App.tsx b/ui/src/App.tsx
index 27c57207..5963d31e 100644
--- a/ui/src/App.tsx
+++ b/ui/src/App.tsx
@@ -131,127 +131,136 @@ function App() {
 
   return (
     <div
-      className="flex h-dvh w-screen select-none"
-      style={{ maxWidth: "100vw" }}
+      className="flex w-screen select-none"
+      style={{ maxWidth: "100vw", height: "calc(100dvh - 2rem)" }}
     >
-      <div className="relative flex h-full flex-col !border-r-small border-divider transition-width pb-6 pt-9 items-center">
-        <div className="flex items-center gap-0 px-3 justify-center">
-          <div className="flex h-8 w-8">
-            <img src={icon} alt="icon" className="h-8 w-8" />
+      <div className="flex w-full">
+        <div className="relative flex flex-col !border-r-small border-divider transition-width pb-6 pt-4 items-center">
+          <div className="flex items-center gap-0 px-3 justify-center">
+            <div className="flex h-8 w-8">
+              <img src={icon} alt="icon" className="h-8 w-8" />
+            </div>
           </div>
-        </div>
-
-        <ScrollShadow className="-mr-6 h-full max-h-full py-6 pr-6">
-          <Sidebar
-            defaultSelectedKey="home"
-            isCompact={true}
-            items={navigation}
-          />
-        </ScrollShadow>
 
-        <Spacer y={2} />
+          <ScrollShadow className="-mr-6 h-full max-h-full py-6 pr-6">
+            <Sidebar
+              defaultSelectedKey="home"
+              isCompact={true}
+              items={navigation}
+              className="z-50"
+            />
+          </ScrollShadow>
 
-        <div className="flex items-center gap-3 px-3">
-          <Dropdown showArrow placement="right-start">
-            <DropdownTrigger>
-              <Button disableRipple isIconOnly radius="full" variant="light">
-                <Avatar
-                  isBordered
-                  className="flex-none"
-                  size="sm"
-                  name={user.username || ""}
-                />
-              </Button>
-            </DropdownTrigger>
-            <DropdownMenu aria-label="Custom item styles">
-              <DropdownItem
-                key="profile"
-                isReadOnly
-                className="h-14 opacity-100"
-                textValue="Signed in as"
-              >
-                <User
-                  avatarProps={{
-                    size: "sm",
-                    name: user.username || "Anonymous User",
-                    showFallback: true,
-                    imgProps: {
-                      className: "transition-none",
-                    },
-                  }}
-                  classNames={{
-                    name: "text-default-600",
-                    description: "text-default-500",
-                  }}
-                  description={
-                    user.bio || (user.username && "No bio") || "Sign up now"
-                  }
-                  name={user.username || "Anonymous User"}
-                />
-              </DropdownItem>
+          <Spacer y={2} />
 
-              <DropdownItem
-                key="settings"
-                description="Configure Atuin"
-                startContent={<Icon icon="solar:settings-linear" width={24} />}
-              >
-                Settings
-              </DropdownItem>
+          <div className="flex items-center gap-3 px-3">
+            <Dropdown showArrow placement="right-start">
+              <DropdownTrigger>
+                <Button disableRipple isIconOnly radius="full" variant="light">
+                  <Avatar
+                    isBordered
+                    className="flex-none"
+                    size="sm"
+                    name={user.username || ""}
+                  />
+                </Button>
+              </DropdownTrigger>
+              <DropdownMenu aria-label="Custom item styles">
+                <DropdownItem
+                  key="profile"
+                  isReadOnly
+                  className="h-14 opacity-100"
+                  textValue="Signed in as"
+                >
+                  <User
+                    avatarProps={{
+                      size: "sm",
+                      name: user.username || "Anonymous User",
+                      showFallback: true,
+                      imgProps: {
+                        className: "transition-none",
+                      },
+                    }}
+                    classNames={{
+                      name: "text-default-600",
+                      description: "text-default-500",
+                    }}
+                    description={
+                      user.bio || (user.username && "No bio") || "Sign up now"
+                    }
+                    name={user.username || "Anonymous User"}
+                  />
+                </DropdownItem>
 
-              <DropdownSection aria-label="Help & Feedback">
                 <DropdownItem
-                  key="help_and_feedback"
-                  description="Get in touch"
-                  onPress={() => open("https://forum.atuin.sh")}
+                  key="settings"
+                  description="Configure Atuin"
                   startContent={
-                    <Icon width={24} icon="solar:question-circle-linear" />
+                    <Icon icon="solar:settings-linear" width={24} />
                   }
                 >
-                  Help & Feedback
+                  Settings
                 </DropdownItem>
 
-                {(user.username && (
+                <DropdownSection aria-label="Help & Feedback">
                   <DropdownItem
-                    key="logout"
+                    key="help_and_feedback"
+                    description="Get in touch"
+                    onPress={() => open("https://forum.atuin.sh")}
                     startContent={
-                      <Icon width={24} icon="solar:logout-broken" />
+                      <Icon width={24} icon="solar:question-circle-linear" />
                     }
-                    onClick={() => {
-                      logout();
-                      refreshUser();
-                    }}
-                  >
-                    Log Out
-                  </DropdownItem>
-                )) || (
-                  <DropdownItem
-                    key="signup"
-                    description="Sync, backup and share your data"
-                    className="bg-emerald-100"
-                    startContent={<KeyRoundIcon size="18px" />}
-                    onPress={onOpen}
                   >
-                    Log in or Register
+                    Help & Feedback
                   </DropdownItem>
-                )}
-              </DropdownSection>
-            </DropdownMenu>
-          </Dropdown>
+
+                  {(user.username && (
+                    <DropdownItem
+                      key="logout"
+                      startContent={
+                        <Icon width={24} icon="solar:logout-broken" />
+                      }
+                      onClick={() => {
+                        logout();
+                        refreshUser();
+                      }}
+                    >
+                      Log Out
+                    </DropdownItem>
+                  )) || (
+                    <DropdownItem
+                      key="signup"
+                      description="Sync, backup and share your data"
+                      className="bg-emerald-100"
+                      startContent={<KeyRoundIcon size="18px" />}
+                      onPress={onOpen}
+                    >
+                      Log in or Register
+                    </DropdownItem>
+                  )}
+                </DropdownSection>
+              </DropdownMenu>
+            </Dropdown>
+          </div>
         </div>
-      </div>
 
-      {renderMain(section)}
+        {renderMain(section)}
 
-      <Toaster />
-      <Modal isOpen={isOpen} onOpenChange={onOpenChange} placement="top-center">
-        <ModalContent className="p-8">
-          {(onClose) => (
-            <>
-              <LoginOrRegister onClose={onClose} />
-            </>
-          )}
-        </ModalContent>
-      </Modal>
+        <Toaster />
+        <Modal
+          isOpen={isOpen}
+          onOpenChange={onOpenChange}
+          placement="top-center"
+        >
+          <ModalContent className="p-8">
+            {(onClose) => (
+              <>
+                <LoginOrRegister onClose={onClose} />
+              </>
+            )}
+          </ModalContent>
+        </Modal>
+      </div>
     </div>
   );
 }
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} />
diff --git a/ui/src/main.tsx b/ui/src/main.tsx
index 58cfd77b..96d570a3 100644
--- a/ui/src/main.tsx
+++ b/ui/src/main.tsx
@@ -8,8 +8,14 @@ ReactDOM.createRoot(document.getElementById("root") as HTMLElement).render(
   <React.StrictMode>
     <NextUIProvider>
       <main className="text-foreground bg-background">
-        <div data-tauri-drag-region className="w-full min-h-8 absolute z-10" />
-        <App />
+        <div
+          data-tauri-drag-region
+          className="w-full min-h-8 z-10 border-b-1"
+        />
+
+        <div className="z-20 ">
+          <App />
+        </div>
       </main>
     </NextUIProvider>
   </React.StrictMode>,
diff --git a/ui/src/state/store.ts b/ui/src/state/store.ts
index cde2da17..c6d3c152 100644
--- a/ui/src/state/store.ts
+++ b/ui/src/state/store.ts
@@ -19,11 +19,24 @@ import { invoke } from "@tauri-apps/api/core";
 import { sessionToken, settings } from "./client";
 import { getWeekInfo } from "@/lib/utils";
 import Runbook from "./runbooks/runbook";
+import { Terminal } from "@xterm/xterm";
+import { FitAddon } from "@xterm/addon-fit";
+import { WebglAddon } from "@xterm/addon-webgl";
+
+export class TerminalData {
+  terminal: Terminal;
+  fitAddon: FitAddon;
+
+  constructor(terminal: Terminal, fit: FitAddon) {
+    this.terminal = terminal;
+    this.fitAddon = fit;
+  }
+}
 
 // I'll probs want to slice this up at some point, but for now a
 // big blobby lump of state is fine.
 // Totally just hoping that structure will be emergent in the future.
-interface AtuinState {
+export interface AtuinState {
   user: User;
   homeInfo: HomeInfo;
   aliases: Alias[];
@@ -32,7 +45,7 @@ interface AtuinState {
   calendar: any[];
   weekStart: number;
   runbooks: Runbook[];
-  currentRunbook: String | null;
+  currentRunbook: string | null;
 
   refreshHomeInfo: () => void;
   refreshCalendar: () => void;
@@ -44,6 +57,16 @@ interface AtuinState {
   historyNextPage: (query?: string) => void;
 
   setCurrentRunbook: (id: String) => void;
+  setPtyTerm: (pty: string, terminal: any) => void;
+  newPtyTerm: (pty: string, runbook: string) => TerminalData;
+  cleanupPtyTerm: (pty: string) => void;
+
+  terminals: { [pty: string]: TerminalData };
+
+  // Store ephemeral state for runbooks, that is not persisted to the database
+  runbookInfo: { [runbook: string]: { ptys: number } };
+  incRunbookPty: (runbook: string) => void;
+  decRunbookPty: (runbook: string) => void;
 }
 
 let state = (set: any, get: any): AtuinState => ({
@@ -55,6 +78,8 @@ let state = (set: any, get: any): AtuinState => ({
   calendar: [],
   runbooks: [],
   currentRunbook: "",
+  terminals: {},
+  runbookInfo: {},
 
   weekStart: getWeekInfo().firstDay,
 
@@ -158,8 +183,101 @@ let state = (set: any, get: any): AtuinState => ({
   setCurrentRunbook: (id: String) => {
     set({ currentRunbook: id });
   },
+
+  setPtyTerm: (pty: string, terminal: TerminalData) => {
+    set({
+      terminals: { ...get().terminals, [pty]: terminal },
+    });
+  },
+
+  cleanupPtyTerm: (pty: string) => {
+    set((state: AtuinState) => {
+      const terminals = Object.keys(state.terminals).reduce(
+        (terms: { [pty: string]: TerminalData }, key) => {
+          if (key !== pty) {
+            terms[key] = state.terminals[key];
+          }
+          return terms;
+        },
+        {},
+      );
+
+      return { terminals };
+    });
+  },
+
+  newPtyTerm: (pty: string) => {
+    let terminal = new Terminal();
+
+    // TODO: fallback to canvas, also some sort of setting to allow disabling webgl usage
+    // probs fine for now though, it's widely supported. maybe issues on linux.
+    terminal.loadAddon(new WebglAddon());
+
+    let fitAddon = new FitAddon();
+    terminal.loadAddon(fitAddon);
+
+    const onResize = (size: { cols: number; rows: number }) => {
+      invoke("pty_resize", {
+        pid: pty,
+        cols: size.cols,
+        rows: size.rows,
+      });
+    };
+
+    terminal.onResize(onResize);
+
+    let td = new TerminalData(terminal, fitAddon);
+
+    set({
+      terminals: { ...get().terminals, [pty]: td },
+    });
+
+    return td;
+  },
+
+  incRunbookPty: (runbook: string) => {
+    set((state: AtuinState) => {
+      let oldVal = state.runbookInfo[runbook] || { ptys: 0 };
+      let newVal = { ptys: oldVal.ptys + 1 };
+      console.log(newVal);
+
+      return {
+        runbookInfo: {
+          ...state.runbookInfo,
+          [runbook]: newVal,
+        },
+      };
+    });
+  },
+
+  decRunbookPty: (runbook: string) => {
+    set((state: AtuinState) => {
+      let newVal = state.runbookInfo[runbook];
+      if (!newVal) {
+        return;
+      }
+
+      newVal.ptys--;
+
+      return {
+        runbookInfo: {
+          ...state.runbookInfo,
+          [runbook]: newVal,
+        },
+      };
+    });
+  },
 });
 
 export const useStore = create<AtuinState>()(
-  persist(state, { name: "atuin-storage" }),
+  persist(state, {
+    name: "atuin-storage",
+
+    // don't serialize the terminals map
+    // it won't work as JSON. too cyclical
+    partialize: (state) =>
+      Object.fromEntries(
+        Object.entries(state).filter(([key]) => !["terminals"].includes(key)),
+      ),
+  }),
 );