From a35b975d84e842e871e589bca59b2eb892c4bbe8 Mon Sep 17 00:00:00 2001
From: Marc Schreiber <info@schrieveslaach.de>
Date: Sun, 17 Apr 2022 05:03:02 +0200
Subject: [PATCH] Shell Integration (#5162)

This commit renders ANSI chars in order to provide shell integrations
such Kitty's opening feature that captures the output of the last
command in a pager such as less.

Fixes #5138
---
 crates/nu-cli/src/prompt.rs        | 53 ++++++++++++++++++++++--------
 crates/nu-cli/src/prompt_update.rs |  4 +++
 crates/nu-cli/src/repl.rs          |  7 ++++
 3 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/crates/nu-cli/src/prompt.rs b/crates/nu-cli/src/prompt.rs
index fb843dcea8..361baec682 100644
--- a/crates/nu-cli/src/prompt.rs
+++ b/crates/nu-cli/src/prompt.rs
@@ -16,6 +16,7 @@ pub struct NushellPrompt {
     default_vi_insert_prompt_indicator: Option<String>,
     default_vi_normal_prompt_indicator: Option<String>,
     default_multiline_indicator: Option<String>,
+    shell_integration: bool,
 }
 
 impl Default for NushellPrompt {
@@ -33,6 +34,7 @@ impl NushellPrompt {
             default_vi_insert_prompt_indicator: None,
             default_vi_normal_prompt_indicator: None,
             default_multiline_indicator: None,
+            shell_integration: false,
         }
     }
 
@@ -82,20 +84,34 @@ impl NushellPrompt {
     fn default_wrapped_custom_string(&self, str: String) -> String {
         format!("({})", str)
     }
+
+    pub(crate) fn enable_shell_integration(&mut self) {
+        self.shell_integration = true
+    }
 }
 
 impl Prompt for NushellPrompt {
     fn render_prompt_left(&self) -> Cow<str> {
-        if let Some(prompt_string) = &self.left_prompt_string {
-            prompt_string.replace('\n', "\r\n").into()
+        // Just before starting to draw the PS1 prompt send the escape code (see
+        // https://sw.kovidgoyal.net/kitty/shell-integration/#notes-for-shell-developers)
+        let mut prompt = if self.shell_integration {
+            String::from("\x1b]133;A\x1b\\")
         } else {
-            let default = DefaultPrompt::new();
-            default
-                .render_prompt_left()
-                .to_string()
-                .replace('\n', "\r\n")
-                .into()
-        }
+            String::new()
+        };
+
+        prompt.push_str(&match &self.left_prompt_string {
+            Some(prompt_string) => prompt_string.replace('\n', "\r\n"),
+            None => {
+                let default = DefaultPrompt::new();
+                default
+                    .render_prompt_left()
+                    .to_string()
+                    .replace('\n', "\r\n")
+            }
+        });
+
+        prompt.into()
     }
 
     fn render_prompt_right(&self) -> Cow<str> {
@@ -136,10 +152,21 @@ impl Prompt for NushellPrompt {
     }
 
     fn render_prompt_multiline_indicator(&self) -> Cow<str> {
-        match &self.default_multiline_indicator {
-            Some(indicator) => indicator.as_str().into(),
-            None => "::: ".into(),
-        }
+        // Just before starting to draw the PS1 prompt send the escape code (see
+        // https://sw.kovidgoyal.net/kitty/shell-integration/#notes-for-shell-developers)
+        let mut prompt = if self.shell_integration {
+            String::from("\x1b]133;A;k=s\x1b\\")
+        } else {
+            String::new()
+        };
+
+        prompt.push_str(
+            self.default_multiline_indicator
+                .as_ref()
+                .unwrap_or(&String::from("::: ")),
+        );
+
+        prompt.into()
     }
 
     fn render_prompt_history_search_indicator(
diff --git a/crates/nu-cli/src/prompt_update.rs b/crates/nu-cli/src/prompt_update.rs
index 9cc9ec36fa..ecaafdf543 100644
--- a/crates/nu-cli/src/prompt_update.rs
+++ b/crates/nu-cli/src/prompt_update.rs
@@ -147,6 +147,10 @@ pub(crate) fn update_prompt<'prompt>(
         (prompt_vi_insert_string, prompt_vi_normal_string),
     );
 
+    if config.use_ansi_coloring {
+        nu_prompt.enable_shell_integration();
+    }
+
     let ret_val = nu_prompt as &dyn Prompt;
     if is_perf_true {
         info!("update_prompt {}:{}:{}", file!(), line!(), column!());
diff --git a/crates/nu-cli/src/repl.rs b/crates/nu-cli/src/repl.rs
index 97e8d55397..7f190e6f78 100644
--- a/crates/nu-cli/src/repl.rs
+++ b/crates/nu-cli/src/repl.rs
@@ -231,6 +231,13 @@ pub fn evaluate_repl(
         }
 
         let input = line_editor.read_line(prompt);
+
+        if config.use_ansi_coloring {
+            // Just before running a command/program, send the escape code (see
+            // https://sw.kovidgoyal.net/kitty/shell-integration/#notes-for-shell-developers)
+            print!("\x1b]133;C\x1b\\");
+        }
+
         match input {
             Ok(Signal::Success(s)) => {
                 let start_time = Instant::now();