save and restore both, api_messages and display_messages

Author: alexylon

src/history.rs

@@ -10,6 +10,21 @@ const SESSIONS_DIR: &str = "sessions";
 const INDEX_FILE: &str = "index.json";
 const MAX_PREVIEW_LENGTH: usize = 120;
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum DisplayMessage {
+    UserMessage {
+        content: String,
+    },
+    AssistantMessage {
+        content: String,
+    },
+    ToolExecution {
+        tool_name: String,
+        tool_input: serde_json::Value,
+        tool_output: String,
+    },
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct SessionMetadata {
     pub id: String,
@@ -22,7 +37,12 @@ pub struct SessionMetadata {
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Session {
     pub id: String,
-    pub messages: Vec<Message>,
+    /// Messages in API format (for continuing the conversation with Claude)
+    #[serde(default, rename = "messages", alias = "api_messages")]
+    pub api_messages: Vec<Message>,
+    /// Messages in display format (for reconstructing the original UI)
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub display_messages: Vec<DisplayMessage>,
     pub system_prompt: String,
     pub created_at: u64,
     pub updated_at: u64,
@@ -118,6 +138,7 @@ impl HistoryManager {
         &self,
         session_id: &str,
         messages: &[Message],
+        display_messages: &[DisplayMessage],
         system_prompt: &str,
     ) -> Result<()> {
         let now = SystemTime::now()
@@ -129,7 +150,8 @@ impl HistoryManager {
 
         let session = Session {
             id: session_id.to_string(),
-            messages: messages.to_vec(),
+            api_messages: messages.to_vec(),
+            display_messages: display_messages.to_vec(),
             system_prompt: system_prompt.to_string(),
             created_at: if session_path.exists() {
                 let existing: Session = serde_json::from_str(&fs::read_to_string(&session_path)?)?;
@@ -158,13 +180,13 @@ impl HistoryManager {
             }
         };
 
-        let preview = Self::extract_preview(&session.messages);
+        let preview = Self::extract_preview(&session.api_messages);
         let metadata = SessionMetadata {
             id: session.id.clone(),
             preview,
             created_at: session.created_at,
             updated_at: session.updated_at,
-            message_count: session.messages.len(),
+            message_count: session.api_messages.len(),
         };
 
         if let Some(pos) = index.sessions.iter().position(|s| s.id == session.id) {
@@ -256,12 +278,12 @@ mod tests {
         let system_prompt = "Test system prompt";
 
         manager
-            .save_session(&session_id, &messages, system_prompt)
+            .save_session(&session_id, &messages, &[], system_prompt)
             .unwrap();
 
         let loaded = manager.load_session(&session_id).unwrap();
         assert_eq!(loaded.id, session_id);
-        assert_eq!(loaded.messages.len(), 1);
+        assert_eq!(loaded.api_messages.len(), 1);
         assert_eq!(loaded.system_prompt, system_prompt);
     }
 
@@ -272,14 +294,14 @@ mod tests {
 
         let session_id1 = HistoryManager::generate_session_id();
         manager
-            .save_session(&session_id1, &[Message::user("First session")], "System")
+            .save_session(&session_id1, &[Message::user("First session")], &[], "System")
             .unwrap();
 
         std::thread::sleep(std::time::Duration::from_secs(1));
 
         let session_id2 = HistoryManager::generate_session_id();
         manager
-            .save_session(&session_id2, &[Message::user("Second session")], "System")
+            .save_session(&session_id2, &[Message::user("Second session")], &[], "System")
             .unwrap();
 
         let sessions = manager.list_sessions().unwrap();

src/repl.rs

@@ -25,6 +25,7 @@ pub struct Repl {
     max_tokens: u32,
     recursion_depth: u32,
     session_id: String,
+    display_messages: Vec<crate::history::DisplayMessage>,
 }
 
 impl Repl {
@@ -62,6 +63,7 @@ impl Repl {
             max_tokens,
             recursion_depth: 0,
             session_id,
+            display_messages: Vec::new(),
         })
     }
 
@@ -93,6 +95,7 @@ impl Repl {
                         }
                         "clear" => {
                             self.conversation.clear();
+                            self.display_messages.clear();
                             self.session_id = HistoryManager::generate_session_id();
                             println!("{}", "Conversation history cleared.".bright_yellow());
                             continue;
@@ -150,6 +153,11 @@ impl Repl {
 
     fn process_message(&mut self, user_input: &str) -> Result<()> {
         self.conversation.add_user_message(user_input.to_string());
+        
+        // Track display message
+        self.display_messages.push(crate::history::DisplayMessage::UserMessage {
+            content: user_input.to_string(),
+        });
 
         let request = CreateMessageRequest {
             model: self.model.clone(),
@@ -230,6 +238,12 @@ impl Repl {
                 println!("{}", highlighted);
             }
             println!();
+            
+            // Track assistant display message
+            let combined_text = text_output.join("\n");
+            self.display_messages.push(crate::history::DisplayMessage::AssistantMessage {
+                content: combined_text,
+            });
         }
 
         // Store the full assistant response with content blocks
@@ -291,6 +305,13 @@ impl Repl {
 
                         println!("{}", output.dimmed());
                         println!();
+                        
+                        // Track tool execution in display_messages
+                        self.display_messages.push(crate::history::DisplayMessage::ToolExecution {
+                            tool_name: tool_name.clone(),
+                            tool_input: tool_input.clone(),
+                            tool_output: output.clone(),
+                        });
 
                         // Collect tool result instead of adding immediately
                         tool_results.push(crate::api::MessageContentBlock::ToolResult {
@@ -315,6 +336,13 @@ impl Repl {
                         let error_msg = format!("Tool execution failed: {}", e);
                         eprintln!("{} {}", "Error:".bright_red().bold(), error_msg);
                         println!();
+                        
+                        // Track tool execution error in display_messages
+                        self.display_messages.push(crate::history::DisplayMessage::ToolExecution {
+                            tool_name: tool_name.clone(),
+                            tool_input: tool_input.clone(),
+                            tool_output: error_msg.clone(),
+                        });
 
                         // Collect error as tool result
                         tool_results.push(crate::api::MessageContentBlock::ToolResult {
@@ -543,6 +571,7 @@ impl Repl {
         self.history_manager.save_session(
             &self.session_id,
             self.conversation.messages(),
+            &self.display_messages,
             self.conversation.system_prompt(),
         )?;
 
@@ -577,15 +606,74 @@ impl Repl {
 
         self.session_id = session.id.clone();
         self.conversation.clear();
-        self.conversation.restore_messages(session.messages.clone());
+        self.conversation.restore_messages(session.api_messages.clone());
+        self.display_messages = session.display_messages.clone();
 
         println!(
             "{} {} ({} messages)",
             "Loaded session:".bright_green(),
             session.id,
-            session.messages.len()
+            session.api_messages.len()
         );
+        println!();
+        
+        // Display the original conversation
+        self.display_session(&session);
 
         Ok(())
     }
+    
+    fn display_session(&self, session: &crate::history::Session) {
+        if session.display_messages.is_empty() {
+            println!("{}", "Note: This is a legacy session without display history. Only API messages are available.".dimmed());
+            println!();
+            return;
+        }
+        
+        println!("{}", "═".repeat(80).bright_cyan());
+        println!("{}", "Previous Conversation:".bright_cyan().bold());
+        println!("{}", "═".repeat(80).bright_cyan());
+        println!();
+        
+        for display_msg in &session.display_messages {
+            match display_msg {
+                crate::history::DisplayMessage::UserMessage { content } => {
+                    println!("{} {}", ">>>".bright_green(), content);
+                    println!();
+                }
+                crate::history::DisplayMessage::AssistantMessage { content } => {
+                    println!("{}", "Assistant:".bright_blue().bold());
+                    let highlighted = self.highlighter.highlight_text(content);
+                    println!("{}", highlighted);
+                    println!();
+                }
+                crate::history::DisplayMessage::ToolExecution { tool_name, tool_input: _, tool_output } => {
+                    if tool_name == "execute_bash" {
+                        if let Ok(input_val) = serde_json::from_value::<serde_json::Value>(
+                            serde_json::to_value(&tool_output).unwrap_or_default()
+                        ) {
+                            if let Some(command) = input_val.get("command").and_then(|v| v.as_str()) {
+                                println!(
+                                    "{} {}",
+                                    "Executing:".bright_green().bold(),
+                                    command.bright_cyan()
+                                );
+                            }
+                        }
+                    } else {
+                        println!(
+                            "{} {}",
+                            "Using tool:".bright_yellow().bold(),
+                            tool_name.bright_yellow()
+                        );
+                    }
+                    println!("{}", tool_output.dimmed());
+                    println!();
+                }
+            }
+        }
+        
+        println!("{}", "═".repeat(80).bright_cyan());
+        println!();
+    }
 }