Capturing command line prompt

2 files changed, 114 insertions, 13 deletions

diff --git a/extension/src/terminal/terminalEmulator.ts b/extension/src/terminal/terminalEmulator.ts
index ff2060f8..8e49737e 100644
--- a/extension/src/terminal/terminalEmulator.ts
+++ b/extension/src/terminal/terminalEmulator.ts
@@ -3,6 +3,7 @@
 import * as vscode from "vscode";
 import os = require("os");
 import stripAnsi from "strip-ansi";
+import { longestCommonSubsequence } from "../util/lcs";
 
 function loadNativeModule<T>(id: string): T | null {
   try {
@@ -71,13 +72,20 @@ export class CapturedTerminal {
 
   private dataEndsInPrompt(strippedData: string): boolean {
     const lines = strippedData.split("\n");
-    const last_line = lines[lines.length - 1];
+    const lastLine = lines[lines.length - 1];
+
     return (
-      (lines.length > 0 &&
-        (last_line.includes("bash-") || last_line.includes(") $ ")) &&
-        last_line.includes("$")) ||
-      (last_line.includes("]> ") && last_line.includes(") [")) ||
-      (last_line.includes(" (") && last_line.includes(")>"))
+      lines.length > 0 &&
+      (((lastLine.includes("bash-") || lastLine.includes(") $ ")) &&
+        lastLine.includes("$")) ||
+        (lastLine.includes("]> ") && lastLine.includes(") [")) ||
+        (lastLine.includes(" (") && lastLine.includes(")>")) ||
+        (typeof this.commandPromptString !== "undefined" &&
+          (lastLine.includes(this.commandPromptString) ||
+            this.commandPromptString.length -
+              longestCommonSubsequence(lastLine, this.commandPromptString)
+                .length <
+              3)))
     );
   }
 
@@ -96,12 +104,6 @@ export class CapturedTerminal {
   }
 
   async runCommand(command: string): Promise<string> {
-    if (!this.hasRunCommand) {
-      this.hasRunCommand = true;
-      // Let the first bash- prompt appear and let python env be opened
-      // await this.waitForCommandToFinish();
-    }
-
     if (this.commandQueue.length === 0) {
       return new Promise(async (resolve, reject) => {
         this.commandQueue.push([command, resolve]);
@@ -109,8 +111,12 @@ export class CapturedTerminal {
         while (this.commandQueue.length > 0) {
           const [command, resolve] = this.commandQueue.shift()!;
 
+          // Refresh the command prompt string every time in case it changes
+          await this.refreshCommandPromptString();
+
           this.terminal.sendText(command);
-          resolve(await this.waitForCommandToFinish());
+          const output = await this.waitForCommandToFinish();
+          resolve(output);
         }
       });
     } else {
@@ -138,6 +144,24 @@ export class CapturedTerminal {
     }
   }
 
+  private runningClearToGetPrompt: boolean = false;
+  private seenClear: boolean = false;
+  private commandPromptString: string | undefined = undefined;
+  private resolveMeWhenCommandPromptStringFound:
+    | ((_: unknown) => void)
+    | undefined = undefined;
+
+  private async refreshCommandPromptString(): Promise<string | undefined> {
+    // Sends a message that will be received by the terminal to get the command prompt string, see the onData method below in constructor.
+    this.runningClearToGetPrompt = true;
+    this.terminal.sendText("echo");
+    const promise = new Promise((resolve, reject) => {
+      this.resolveMeWhenCommandPromptStringFound = resolve;
+    });
+    await promise;
+    return this.commandPromptString;
+  }
+
   constructor(
     options: { name: string } & Partial<vscode.ExtensionTerminalOptions>,
     onCommandOutput?: (output: string) => void
@@ -165,7 +189,54 @@ export class CapturedTerminal {
     this.writeEmitter = new vscode.EventEmitter<string>();
 
     this.ptyProcess.onData((data: any) => {
+      if (this.runningClearToGetPrompt) {
+        if (
+          stripAnsi(data)
+            .split("\n")
+            .flatMap((line) => line.split("\r"))
+            .find((line) => line.trim() === "echo") !== undefined
+        ) {
+          this.seenClear = true;
+          return;
+        } else if (this.seenClear) {
+          const strippedLines = stripAnsi(data)
+            .split("\r")
+            .filter(
+              (line) =>
+                line.trim().length > 0 &&
+                line.trim() !== "%" &&
+                line.trim() !== "⏎"
+            );
+          const lastLine = strippedLines[strippedLines.length - 1] || "";
+          const lines = lastLine
+            .split("\n")
+            .filter(
+              (line) =>
+                line.trim().length > 0 &&
+                line.trim() !== "%" &&
+                line.trim() !== "⏎"
+            );
+          const commandPromptString = (lines[lines.length - 1] || "").trim();
+          if (
+            commandPromptString.length > 0 &&
+            !commandPromptString.includes("echo")
+          ) {
+            this.runningClearToGetPrompt = false;
+            this.seenClear = false;
+            this.commandPromptString = commandPromptString;
+            console.log(
+              "Found command prompt string: " + this.commandPromptString
+            );
+            if (this.resolveMeWhenCommandPromptStringFound) {
+              this.resolveMeWhenCommandPromptStringFound(undefined);
+            }
+          }
+          return;
+        }
+      }
+
       // Pass data through to terminal
+      data = data.replace("⏎", "");
       this.writeEmitter.fire(data);
 
       this.splitByCommandsListener(data);
diff --git a/extension/src/util/lcs.ts b/extension/src/util/lcs.ts
new file mode 100644
index 00000000..17ea63f9
--- /dev/null
+++ b/extension/src/util/lcs.ts
@@ -0,0 +1,30 @@
+export function longestCommonSubsequence(a: string, b: string) {
+  const lengths: number[][] = [];
+  for (let i = 0; i <= a.length; i++) {
+    lengths[i] = [];
+    for (let j = 0; j <= b.length; j++) {
+      if (i === 0 || j === 0) {
+        lengths[i][j] = 0;
+      } else if (a[i - 1] === b[j - 1]) {
+        lengths[i][j] = lengths[i - 1][j - 1] + 1;
+      } else {
+        lengths[i][j] = Math.max(lengths[i - 1][j], lengths[i][j - 1]);
+      }
+    }
+  }
+  let result = "";
+  let x = a.length;
+  let y = b.length;
+  while (x !== 0 && y !== 0) {
+    if (lengths[x][y] === lengths[x - 1][y]) {
+      x--;
+    } else if (lengths[x][y] === lengths[x][y - 1]) {
+      y--;
+    } else {
+      result = a[x - 1] + result;
+      x--;
+      y--;
+    }
+  }
+  return result;
+}