feat: Improved readability of session message exchange replays

src/repl/mod.rs

@@ -7,7 +7,10 @@ use self::completer::ReplCompleter;
 use self::highlighter::ReplHighlighter;
 use self::prompt::ReplPrompt;
 
-use crate::client::{call_chat_completions, call_chat_completions_streaming, init_client, oauth};
+use crate::client::{
+    Message, MessageRole, call_chat_completions, call_chat_completions_streaming, init_client,
+    oauth,
+};
 use crate::config::{
     AgentVariables, AppConfig, AssertState, Input, LastMessage, RequestContext, StateFlags,
     macro_execute,
@@ -370,6 +373,31 @@ Type ".help" for additional help.
             if !compressed.is_empty() || !active.is_empty() {
                 let app = Arc::clone(&self.ctx.read().app.config);
                 replay::render(app.as_ref(), &compressed, &active)?;
+                let last_msgs: &[Message] = if !active.is_empty() {
+                    &active
+                } else {
+                    &compressed
+                };
+                let assistant_text = last_msgs
+                    .iter()
+                    .rev()
+                    .find(|m| m.role == MessageRole::Assistant)
+                    .and_then(|m| m.content.as_text())
+                    .map(str::to_string);
+                if let Some(output) = assistant_text {
+                    let user_text = last_msgs
+                        .iter()
+                        .rev()
+                        .find(|m| m.role == MessageRole::User)
+                        .and_then(|m| m.content.as_text())
+                        .unwrap_or("")
+                        .to_string();
+                    let ctx = self.ctx.read();
+                    if let Ok(input) = Input::from_str(&ctx, &user_text, None) {
+                        drop(ctx);
+                        self.ctx.write().last_message = Some(LastMessage::new(input, output));
+                    }
+                }
             }
         }
 
@@ -1119,16 +1147,28 @@ pub async fn run_repl_command(
                 }
             },
             ".copy" => {
-                let output = match ctx
+                let output = ctx
                     .last_message
                     .as_ref()
                     .filter(|v| !v.output.is_empty())
                     .map(|v| v.output.clone())
-                {
-                    Some(v) => v,
-                    None => bail!("No chat response to copy"),
-                };
-                set_text(&output).context("Failed to copy the last chat response")?;
+                    .or_else(|| {
+                        ctx.session.as_ref().and_then(|s| {
+                            s.messages()
+                                .iter()
+                                .rev()
+                                .chain(s.compressed_messages().iter().rev())
+                                .find(|m| m.role == MessageRole::Assistant)
+                                .and_then(|m| m.content.as_text())
+                                .map(str::to_string)
+                        })
+                    });
+                match output {
+                    Some(v) if !v.is_empty() => {
+                        set_text(&v).context("Failed to copy the last chat response")?;
+                    }
+                    _ => bail!("No chat response to copy"),
+                }
             }
             ".exit" => match args {
                 Some("role") => {

src/repl/replay.rs

@@ -2,7 +2,7 @@ use anyhow::Result;
 
 use crate::client::{Message, MessageRole};
 use crate::config::{AppConfig, Session};
-use crate::utils::dimmed_text;
+use crate::utils::{dimmed_text, replay_label_text};
 
 pub fn snapshot(session: &Session) -> (Vec<Message>, Vec<Message>) {
     (
@@ -40,13 +40,14 @@ fn render_messages(app: &AppConfig, messages: &[Message]) -> Result<()> {
         match message.role {
             MessageRole::User => {
                 if let Some(text) = message.content.as_text() {
-                    println!("{}", dimmed_text("You:"));
+                    println!("{}", replay_label_text("You:"));
                     println!("{text}");
                     println!();
                 }
             }
             MessageRole::Assistant => {
                 if let Some(text) = message.content.as_text() {
+                    println!("{}", replay_label_text("Assistant:"));
                     app.print_markdown(text)?;
                     println!();
                 }

src/utils/mod.rs

@@ -75,6 +75,7 @@ static TOOL_DIM_COLOR: OnceLock<Color> = OnceLock::new();
 static TOOL_FN_COLOR: OnceLock<Color> = OnceLock::new();
 static TOOL_KEY_COLOR: OnceLock<Color> = OnceLock::new();
 static TOOL_WARN_COLOR: OnceLock<Color> = OnceLock::new();
+static REPLAY_LABEL_COLOR: OnceLock<Color> = OnceLock::new();
 
 pub fn init_tool_colors(theme: &Theme) {
     fn resolve(theme: &Theme, scope_str: &str) -> Option<Color> {
@@ -96,6 +97,16 @@ pub fn init_tool_colors(theme: &Theme) {
     if let Some(c) = resolve(theme, "string") {
         let _ = TOOL_WARN_COLOR.set(c);
     }
+    let replay_color = Scope::new("entity.name")
+        .ok()
+        .and_then(|scope| {
+            let style = Highlighter::new(theme).style_mod_for_stack(&[scope]);
+            style.foreground.or(theme.settings.foreground)
+        })
+        .map(|fg| Color::Rgb(fg.r, fg.g, fg.b));
+    if let Some(c) = replay_color {
+        let _ = REPLAY_LABEL_COLOR.set(c);
+    }
 }
 
 pub fn now() -> String {
@@ -236,6 +247,18 @@ pub fn cyan_bold_text(input: &str) -> String {
         .to_string()
 }
 
+pub fn replay_label_text(input: &str) -> String {
+    if *NO_COLOR {
+        return input.to_string();
+    }
+    let color = REPLAY_LABEL_COLOR.get().copied().unwrap_or(Color::Green);
+    nu_ansi_term::Style::new()
+        .fg(color)
+        .bold()
+        .paint(input)
+        .to_string()
+}
+
 pub fn magenta_text(input: &str) -> String {
     if *NO_COLOR {
         return input.to_string();