From 98126e298175ed78a5797bacc06d0704a30b1785 Mon Sep 17 00:00:00 2001
From: Darren Schroeder <343840+fdncred@users.noreply.github.com>
Date: Wed, 20 Jul 2022 15:03:29 -0500
Subject: [PATCH] add more shell integration ansi escapes in support of vscode
 (#6087)

* add more shell integration ansi escapes in support of vscode

* clippy
---
 crates/nu-cli/src/prompt_update.rs | 18 ++++++++++++++++
 crates/nu-cli/src/repl.rs          | 33 ++++++++++++++++++------------
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/crates/nu-cli/src/prompt_update.rs b/crates/nu-cli/src/prompt_update.rs
index 9cc9ec36fa..6877f57320 100644
--- a/crates/nu-cli/src/prompt_update.rs
+++ b/crates/nu-cli/src/prompt_update.rs
@@ -15,6 +15,10 @@ pub(crate) const PROMPT_INDICATOR: &str = "PROMPT_INDICATOR";
 pub(crate) const PROMPT_INDICATOR_VI_INSERT: &str = "PROMPT_INDICATOR_VI_INSERT";
 pub(crate) const PROMPT_INDICATOR_VI_NORMAL: &str = "PROMPT_INDICATOR_VI_NORMAL";
 pub(crate) const PROMPT_MULTILINE_INDICATOR: &str = "PROMPT_MULTILINE_INDICATOR";
+// According to Daniel Imms @Tyriar, we need to do these this way:
+// <133 A><prompt><133 B><command><133 C><command output>
+const PRE_PROMPT_MARKER: &str = "\x1b]133;A\x1b\\";
+const POST_PROMPT_MARKER: &str = "\x1b]133;B\x1b\\";
 
 fn get_prompt_string(
     prompt: &str,
@@ -98,6 +102,20 @@ pub(crate) fn update_prompt<'prompt>(
         is_perf_true,
     );
 
+    // Now that we have the prompt string lets ansify it.
+    // <133 A><prompt><133 B><command><133 C><command output>
+    let left_prompt_string = if config.shell_integration {
+        match left_prompt_string {
+            Some(prompt_string) => Some(format!(
+                "{}{}{}",
+                PRE_PROMPT_MARKER, prompt_string, POST_PROMPT_MARKER
+            )),
+            None => left_prompt_string,
+        }
+    } else {
+        left_prompt_string
+    };
+
     let right_prompt_string = get_prompt_string(
         PROMPT_COMMAND_RIGHT,
         config,
diff --git a/crates/nu-cli/src/repl.rs b/crates/nu-cli/src/repl.rs
index a702afd9fc..c7abc78204 100644
--- a/crates/nu-cli/src/repl.rs
+++ b/crates/nu-cli/src/repl.rs
@@ -22,9 +22,14 @@ use std::io::{self, Write};
 use std::{sync::atomic::Ordering, time::Instant};
 use sysinfo::SystemExt;
 
-const PRE_PROMPT_MARKER: &str = "\x1b]133;A\x1b\\";
+// According to Daniel Imms @Tyriar, we need to do these this way:
+// <133 A><prompt><133 B><command><133 C><command output>
+// These first two have been moved to prompt_update to get as close as possible to the prompt.
+// const PRE_PROMPT_MARKER: &str = "\x1b]133;A\x1b\\";
+// const POST_PROMPT_MARKER: &str = "\x1b]133;B\x1b\\";
 const PRE_EXECUTE_MARKER: &str = "\x1b]133;C\x1b\\";
-const CMD_FINISHED_MARKER: &str = "\x1b]133;D\x1b\\";
+// This one is in get_command_finished_marker() now so we can capture the exit codes properly.
+// const CMD_FINISHED_MARKER: &str = "\x1b]133;D;{}\x1b\\";
 const RESET_APPLICATION_MODE: &str = "\x1b[?1l";
 
 pub fn evaluate_repl(
@@ -262,12 +267,6 @@ pub fn evaluate_repl(
         }
 
         let config = engine_state.get_config();
-
-        let shell_integration = config.shell_integration;
-        if shell_integration {
-            run_ansi_sequence(PRE_PROMPT_MARKER)?;
-        }
-
         let prompt =
             prompt_update::update_prompt(config, engine_state, stack, &mut nu_prompt, is_perf_true);
 
@@ -283,6 +282,7 @@ pub fn evaluate_repl(
         }
 
         let input = line_editor.read_line(prompt);
+        let shell_integration = config.shell_integration;
 
         match input {
             Ok(Signal::Success(s)) => {
@@ -425,20 +425,19 @@ pub fn evaluate_repl(
                 }
 
                 if shell_integration {
-                    // FIXME: use variant with exit code, if apropriate
-                    run_ansi_sequence(CMD_FINISHED_MARKER)?;
+                    run_ansi_sequence(&get_command_finished_marker(stack, engine_state))?;
                 }
             }
             Ok(Signal::CtrlC) => {
                 // `Reedline` clears the line content. New prompt is shown
                 if shell_integration {
-                    run_ansi_sequence(CMD_FINISHED_MARKER)?;
+                    run_ansi_sequence(&get_command_finished_marker(stack, engine_state))?;
                 }
             }
             Ok(Signal::CtrlD) => {
                 // When exiting clear to a new line
                 if shell_integration {
-                    run_ansi_sequence(CMD_FINISHED_MARKER)?;
+                    run_ansi_sequence(&get_command_finished_marker(stack, engine_state))?;
                 }
                 println!();
                 break;
@@ -449,7 +448,7 @@ pub fn evaluate_repl(
                     println!("Error: {:?}", err);
                 }
                 if shell_integration {
-                    run_ansi_sequence(CMD_FINISHED_MARKER)?;
+                    run_ansi_sequence(&get_command_finished_marker(stack, engine_state))?;
                 }
             }
         }
@@ -458,6 +457,14 @@ pub fn evaluate_repl(
     Ok(())
 }
 
+pub fn get_command_finished_marker(stack: &Stack, engine_state: &EngineState) -> String {
+    let exit_code = stack
+        .get_env_var(engine_state, "LAST_EXIT_CODE")
+        .and_then(|e| e.as_i64().ok());
+
+    format!("\x1b]133;D;{}\x1b\\", exit_code.unwrap_or(0))
+}
+
 pub fn eval_env_change_hook(
     env_change_hook: Option<Value>,
     engine_state: &mut EngineState,