1 files changed, 13 insertions, 17 deletions

diff --git a/src/config/session.rs b/src/config/session.rs
index eccffa9..c36c9db 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -9,7 +9,7 @@ use inquire::{validator::Validation, Confirm, Text};
 use serde::{Deserialize, Serialize};
 use serde_json::json;
 use std::collections::HashMap;
-use std::fs::{self, read_to_string};
+use std::fs::{read_to_string, write};
 use std::path::Path;
 
 #[derive(Debug, Clone, Default, Deserialize, Serialize)]
@@ -86,10 +86,6 @@ impl Session {
         Ok(session)
     }
 
-    pub fn is_temp(&self) -> bool {
-        self.name == TEMP_SESSION_NAME
-    }
-
     pub fn is_empty(&self) -> bool {
         self.messages.is_empty() && self.compressed_messages.is_empty()
     }
@@ -218,12 +214,7 @@ impl Session {
             }
         }
 
-        if lines.last() == Some(&String::new()) {
-            lines.pop();
-        }
-
-        let output = lines.join("\n");
-        Ok(output)
+        Ok(lines.join("\n"))
     }
 
     pub fn tokens_usage(&self) -> (usize, f32) {
@@ -302,6 +293,7 @@ impl Session {
     pub fn exit(&mut self, session_dir: &Path, is_repl: bool) -> Result<()> {
         let save_session = self.save_session();
         if self.dirty && save_session != Some(false) {
+            let mut session_name = self.name().to_string();
             if save_session.is_none() {
                 if !is_repl {
                     return Ok(());
@@ -310,8 +302,8 @@ impl Session {
                 if !ans {
                     return Ok(());
                 }
-                if self.is_temp() {
-                    self.name = Text::new("Session name:")
+                while session_name == TEMP_SESSION_NAME {
+                    session_name = Text::new("Session name:")
                         .with_validator(|input: &str| {
                             if input.trim().is_empty() {
                                 Ok(Validation::Invalid("This field is required".into()))
@@ -322,20 +314,20 @@ impl Session {
                         .prompt()?;
                 }
             }
-            let session_path = session_dir.join(format!("{}.yaml", self.name()));
-            self.save(&session_path, is_repl)?;
+            let session_path = session_dir.join(format!("{session_name}.yaml"));
+            self.save(&session_name, &session_path, is_repl)?;
         }
         Ok(())
     }
 
-    pub fn save(&mut self, session_path: &Path, is_repl: bool) -> Result<()> {
+    pub fn save(&mut self, session_name: &str, session_path: &Path, is_repl: bool) -> Result<()> {
         ensure_parent_exists(session_path)?;
 
         self.path = Some(session_path.display().to_string());
 
         let content = serde_yaml::to_string(&self)
             .with_context(|| format!("Failed to serde session {}", self.name))?;
-        fs::write(session_path, content).with_context(|| {
+        write(session_path, content).with_context(|| {
             format!(
                 "Failed to write session {} to {}",
                 self.name,
@@ -347,6 +339,10 @@ impl Session {
             println!("✨ Saved session to '{}'", session_path.display());
         }
 
+        if self.name() != session_name {
+            self.name = session_name.to_string()
+        }
+
         self.dirty = false;
 
         Ok(())