From 3f5c688a10cbf56ab0bb68af0eb416d591458ea2 Mon Sep 17 00:00:00 2001
From: Jarrod Sibbison <jsibbison@squareup.com>
Date: Tue, 17 Dec 2024 10:58:58 +1100
Subject: [PATCH] Restore command history when resuming session

---
 crates/goose-cli/src/commands/session.rs | 38 ++++++++++++++++--------
 crates/goose-cli/src/prompt.rs           |  3 ++
 crates/goose-cli/src/prompt/rustyline.rs | 13 ++++++++
 crates/goose-cli/src/session.rs          |  8 ++++-
 4 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/crates/goose-cli/src/commands/session.rs b/crates/goose-cli/src/commands/session.rs
index ea8c8fcb4733..e2d7659ad386 100644
--- a/crates/goose-cli/src/commands/session.rs
+++ b/crates/goose-cli/src/commands/session.rs
@@ -57,18 +57,11 @@ pub fn build_session<'a>(
         Err(_) => Box::new(RustylinePrompt::new()),
     };
 
-    println!(
-        "{} {} {} {} {}",
-        style("starting session |").dim(),
-        style("provider:").dim(),
-        style(loaded_profile.provider).cyan().dim(),
-        style("model:").dim(),
-        style(loaded_profile.model).cyan().dim(),
-    );
-    println!(
-        "    {} {}",
-        style("logging to").dim(),
-        style(session_file.display()).dim().cyan(),
+    display_session_info(
+        resume,
+        loaded_profile.provider,
+        loaded_profile.model,
+        session_file.as_path(),
     );
     Box::new(Session::new(agent, prompt, session_file))
 }
@@ -150,6 +143,27 @@ fn load_profile(profile_name: Option<String>) -> Box<Profile> {
     loaded_profile
 }
 
+fn display_session_info(resume: bool, provider: String, model: String, session_file: &Path) {
+    let start_session_msg = if resume {
+        "resuming session |"
+    } else {
+        "starting session |"
+    };
+    println!(
+        "{} {} {} {} {}",
+        style(start_session_msg).dim(),
+        style("provider:").dim(),
+        style(provider).cyan().dim(),
+        style("model:").dim(),
+        style(model).cyan().dim(),
+    );
+    println!(
+        "    {} {}",
+        style("logging to").dim(),
+        style(session_file.display()).dim().cyan(),
+    );
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/crates/goose-cli/src/prompt.rs b/crates/goose-cli/src/prompt.rs
index e4bb2d60233c..2ba736c95042 100644
--- a/crates/goose-cli/src/prompt.rs
+++ b/crates/goose-cli/src/prompt.rs
@@ -12,6 +12,9 @@ pub trait Prompt {
     fn show_busy(&mut self);
     fn hide_busy(&self);
     fn close(&self);
+    /// Load the user's message history into the prompt for command history navigation. First message is the oldest message.
+    /// When history is supported by the prompt.
+    fn load_user_message_history(&mut self, _messages: Vec<Message>) {}
     fn goose_ready(&self) {
         println!("\n");
         println!("Goose is running! Enter your instructions, or try asking what goose can do.");
diff --git a/crates/goose-cli/src/prompt/rustyline.rs b/crates/goose-cli/src/prompt/rustyline.rs
index c80bcac06427..d6dd97791451 100644
--- a/crates/goose-cli/src/prompt/rustyline.rs
+++ b/crates/goose-cli/src/prompt/rustyline.rs
@@ -9,6 +9,7 @@ use super::{
 use anyhow::Result;
 use cliclack::spinner;
 use goose::message::Message;
+use mcp_core::Role;
 use rustyline::DefaultEditor;
 
 const PROMPT: &str = "\x1b[1m\x1b[38;5;30m( O)> \x1b[0m";
@@ -131,6 +132,18 @@ impl Prompt for RustylinePrompt {
         }
     }
 
+    fn load_user_message_history(&mut self, messages: Vec<Message>) {
+        for message in messages.into_iter().filter(|m| m.role == Role::User) {
+            for content in message.content {
+                if let Some(text) = content.as_text() {
+                    if let Err(e) = self.editor.add_history_entry(text) {
+                        eprintln!("Failed to add to history: {}", e);
+                    }
+                }
+            }
+        }
+    }
+
     fn close(&self) {
         // No cleanup required
     }
diff --git a/crates/goose-cli/src/session.rs b/crates/goose-cli/src/session.rs
index cc1ce4b96899..0d0270d3eafc 100644
--- a/crates/goose-cli/src/session.rs
+++ b/crates/goose-cli/src/session.rs
@@ -102,7 +102,11 @@ pub struct Session<'a> {
 }
 
 impl<'a> Session<'a> {
-    pub fn new(agent: Box<dyn Agent>, prompt: Box<dyn Prompt + 'a>, session_file: PathBuf) -> Self {
+    pub fn new(
+        agent: Box<dyn Agent>,
+        mut prompt: Box<dyn Prompt + 'a>,
+        session_file: PathBuf,
+    ) -> Self {
         let messages = match readable_session_file(&session_file) {
             Ok(file) => deserialize_messages(file).unwrap_or_else(|e| {
                 eprintln!(
@@ -117,6 +121,8 @@ impl<'a> Session<'a> {
             }
         };
 
+        prompt.load_user_message_history(messages.clone());
+
         Session {
             agent,
             prompt,