From f976e0cb479f5e69453fd60d4b81d7835953441e Mon Sep 17 00:00:00 2001
From: reuben olinsky <reubeno@users.noreply.github.com>
Date: Fri, 25 Oct 2024 08:42:46 -0700
Subject: [PATCH] fix: handle PS2 prompts that require prompt-expansion (#239)

---
 brush-core/src/shell.rs                    | 36 +++++++++++++---------
 brush-interactive/src/interactive_shell.rs |  2 +-
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/brush-core/src/shell.rs b/brush-core/src/shell.rs
index d1a4b7d0..c102f207 100644
--- a/brush-core/src/shell.rs
+++ b/brush-core/src/shell.rs
@@ -685,11 +685,31 @@ impl Shell {
 
     /// Composes the shell's prompt, applying all appropriate expansions.
     pub async fn compose_prompt(&mut self) -> Result<String, error::Error> {
+        self.prompt_from_var_or_default("PS1", self.default_prompt())
+            .await
+    }
+
+    /// Compose's the shell's alternate-side prompt, applying all appropriate expansions.
+    #[allow(clippy::unused_async)]
+    pub async fn compose_alt_side_prompt(&mut self) -> Result<String, error::Error> {
+        Ok(String::new())
+    }
+
+    /// Composes the shell's continuation prompt.
+    pub async fn compose_continuation_prompt(&mut self) -> Result<String, error::Error> {
+        self.prompt_from_var_or_default("PS2", "> ").await
+    }
+
+    async fn prompt_from_var_or_default(
+        &mut self,
+        var_name: &str,
+        default: &str,
+    ) -> Result<String, error::Error> {
         // Retrieve the spec.
-        let ps1 = self.parameter_or_default("PS1", self.default_prompt());
+        let prompt_spec = self.parameter_or_default(var_name, default);
 
         // Expand it.
-        let formatted_prompt = prompt::expand_prompt(self, ps1.as_ref())?;
+        let formatted_prompt = prompt::expand_prompt(self, prompt_spec.as_ref())?;
 
         // NOTE: We're having difficulty with xterm escape sequences going through rustyline;
         // so we strip them here.
@@ -703,18 +723,6 @@ impl Shell {
         Ok(formatted_prompt)
     }
 
-    /// Compose's the shell's alternate-side prompt, applying all appropriate expansions.
-    #[allow(clippy::unused_async)]
-    pub async fn compose_alt_side_prompt(&mut self) -> Result<String, error::Error> {
-        Ok(String::new())
-    }
-
-    /// Retrieve's the shell's continuation prompt.
-    pub fn continuation_prompt(&self) -> Result<String, error::Error> {
-        // TODO: Evaluate expansion?
-        Ok(self.parameter_or_default("PS2", "> "))
-    }
-
     /// Returns the exit status of the last command executed in this shell.
     pub fn last_result(&self) -> u8 {
         self.last_exit_status
diff --git a/brush-interactive/src/interactive_shell.rs b/brush-interactive/src/interactive_shell.rs
index 20e1e192..3e24390d 100644
--- a/brush-interactive/src/interactive_shell.rs
+++ b/brush-interactive/src/interactive_shell.rs
@@ -126,7 +126,7 @@ pub trait InteractiveShell {
             let prompt = InteractivePrompt {
                 prompt: shell_mut.as_mut().compose_prompt().await?,
                 alt_side_prompt: shell_mut.as_mut().compose_alt_side_prompt().await?,
-                continuation_prompt: shell_mut.as_mut().continuation_prompt()?,
+                continuation_prompt: shell_mut.as_mut().compose_continuation_prompt().await?,
             };
 
             drop(shell_mut);