Implementation of a persistent shell

package.json

@@ -25,7 +25,7 @@
   },
   "devDependencies": {
     "@commander-js/extra-typings": "^11.0.0",
-    "@types/node": "^20.5.9",
+    "@types/node": "^20.6.0",
     "typescript": "^5.2.2"
   }
 }

src/commands/do.ts

@@ -6,7 +6,7 @@ import {
   getRelevantExecutables,
   shellCommandFromToolCall,
 } from "../llm";
-import { runCommand } from "../runCommand";
+import { PersistentShell } from "../runCommand";
 import {
   logAssistant,
   logCommand,
@@ -32,6 +32,7 @@ export const doCommand = new Command("do")
   )
   .option("--path", "Get a list of executables available on PATH")
   .action(async (instruction, opts) => {
+    const shell = new PersistentShell();
     const messages: ChatMessage[] = [];
 
     if (!instruction) {
@@ -59,7 +60,7 @@ export const doCommand = new Command("do")
       console.log("Getting environment information...");
       const GET_PATH_EXECUTABLES_COMMAND =
         "echo $PATH | tr ':' '\n' | xargs -I {} ls {} 2>/dev/null | sort -u | tr '\n' ','";
-      const pathExecutables = await runCommand({
+      const pathExecutables = await shell.executeCommand({
         command: GET_PATH_EXECUTABLES_COMMAND,
         args: [],
         explanation: "Retrieving all executables in your PATH",
@@ -143,7 +144,7 @@ export const doCommand = new Command("do")
         const shellCommand = shellCommandFromToolCall(toolCall);
         logCommand(`${shellCommand.command} ${shellCommand.args.join(" ")}`);
         logExplanation(shellCommand.explanation);
-        const shellCommandOutput = await runCommand(shellCommand);
+        const shellCommandOutput = await shell.executeCommand(shellCommand);
         logCommandOutput(shellCommandOutput);
         const toolResponse =
           chatMessageFromShellCommandOutput(shellCommandOutput);

src/ctags.ts

@@ -1,8 +1,9 @@
 import { ShellCommandOutput } from "./llm";
-import { runCommand } from "./runCommand";
+import { PersistentShell } from "./runCommand";
 
 const runCtagCommand = (): Promise<ShellCommandOutput> => {
-  return runCommand({
+  const shell = new PersistentShell();
+  return shell.executeCommand({
     command: "ctags --fields=+S --output-format=json -R .",
     args: [],
     explanation: "Generate ctags for the current project",

src/runCommand.ts

@@ -1,29 +1,97 @@
-import { spawn } from "child_process";
+import { ChildProcessWithoutNullStreams, spawn } from "node:child_process";
 import { ShellCommand, ShellCommandOutput } from "./llm";
+import { randomString } from "./utils/rand";
 
-export async function runCommand(
-  cmd: ShellCommand
-): Promise<ShellCommandOutput> {
-  return new Promise((resolve, reject) => {
-    const command = spawn(`${cmd.command} ${cmd.args.join(" ")}`, {
-      shell: true,
+export class PersistentShell {
+  process: ChildProcessWithoutNullStreams;
+  stdout: string;
+  stderr: string;
+  listeners: Map<string, () => void>;
+
+  constructor(shell?: string) {
+    this.process = this.createShellProcess(shell);
+    this.stdout = "";
+    this.stderr = "";
+    this.listeners = new Map();
+
+    this.process.stdout.on("data", (data) => {
+      data = data.toString();
+      this.stdout += data;
+      this.listeners.forEach((listener) => listener());
+    });
+    this.process.stderr.on("data", (data) => {
+      data = data.toString();
+      this.stderr += data;
+      this.listeners.forEach((listener) => listener());
     });
 
-    const commandOutput: ShellCommandOutput = {
-      stdout: "",
-      stderr: "",
-      exitCode: null,
-    };
+    this.process.on("exit", (exitCode) => {
+      console.log(`shell exit code ${exitCode}`);
+    });
 
-    command.stdout.on("data", (output) => {
-      commandOutput.stdout += output.toString();
+    this.process.on("error", (err) => {
+      console.error(`shell error: ${err}`);
     });
-    command.stderr.on("data", (output) => {
-      commandOutput.stderr += output.toString();
+    this.process.on("message", (msg) => {
+      console.error(`shell message: ${msg}`);
     });
-    command.on("close", (code) => {
-      commandOutput.exitCode = code;
-      resolve(commandOutput);
+    this.process.on("disconnect", () => {
+      console.log("shell disconnected");
+    });
+  }
+
+  createShellProcess(shell?: string) {
+    let shellFile = "/bin/sh";
+    if (process.platform === "win32") {
+      shellFile = process.env.comspec || "cmd.exe";
+    } else if (process.platform === "android") {
+      shellFile = "/system/bin/sh";
+    }
+
+    shellFile = shell || shellFile;
+
+    return spawn(shellFile, {
+      shell: false,
+    });
+  }
+
+  executeCommand(cmd: ShellCommand): Promise<ShellCommandOutput> {
+    const uniqueCommandId = `end_command_${randomString(16)}`;
+
+    return new Promise((resolve, reject) => {
+      // This listener runs on every chunk received to either stdout or stderr
+      const onChunk = () => {
+        // Check the entire stdout for the unique end command id
+        if (this.stdout.includes(uniqueCommandId)) {
+          const splitStdout = this.stdout.split("\n");
+          while (splitStdout.pop() !== uniqueCommandId) {}
+          const exitCode = parseInt(splitStdout.pop() || "0");
+          const commandStdout = splitStdout.join("\n");
+
+          // Remove the end command id from the stdout
+          this.stdout = this.stdout.replace(uniqueCommandId, "");
+
+          // Remove the listener for this command
+          if (!this.listeners.delete(uniqueCommandId)) {
+            throw new Error(
+              `removing listener for command ${uniqueCommandId} failed`
+            );
+          }
+
+          resolve({
+            stdout: commandStdout,
+            stderr: this.stderr,
+            exitCode,
+          });
+        }
+      };
+      this.listeners.set(uniqueCommandId, onChunk);
+
+      // Write the main command, and the exit code retrieval, and the unique end command id
+      // to stdin on the shell.
+      this.process.stdin.write(`${cmd.command} ${cmd.args.join(" ")}\n`);
+      this.process.stdin.write(`echo $?\n`);
+      this.process.stdin.write(`echo ${uniqueCommandId}\n`);
     });
-  });
+  }
 }

src/utils/rand.ts

@@ -0,0 +1,5 @@
+import crypto from "crypto";
+
+export function randomString(length: number) {
+  return crypto.randomBytes(length).toString("hex").slice(0, length);
+}