From b7cb6f89f137e20f1d6dbb1f66e40f319c978a45 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Wed, 8 Mar 2023 11:27:51 +0800
Subject: [PATCH] fix: repl set save true not work if not rerun(#34)

---
 src/config.rs       | 36 +++++++++++++++---------------------
 src/main.rs         |  3 +--
 src/repl/handler.rs | 26 ++++++--------------------
 3 files changed, 22 insertions(+), 43 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index dc07c95..8ac99b7 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -103,29 +103,14 @@ impl Config {
         Ok(path)
     }
 
-    pub fn open_message_file(&self) -> Result<Option<File>> {
+    pub fn save_message(&self, input: &str, output: &str) -> Result<()> {
         if !self.save {
-            return Ok(None);
-        }
-        let path = Config::messages_file()?;
-        let file: Option<File> = if self.save {
-            let file = OpenOptions::new()
-                .create(true)
-                .append(true)
-                .open(&path)
-                .with_context(|| format!("Failed to create/append {}", path.display()))?;
-            Some(file)
-        } else {
-            None
-        };
-        Ok(file)
-    }
-
-    pub fn save_message(&self, file: Option<&mut File>, input: &str, output: &str) -> Result<()> {
-        if output.is_empty() || !self.save || file.is_none() {
             return Ok(());
         }
-        let file = file.unwrap();
+        let mut file = self.open_message_file()?;
+        if output.is_empty() || !self.save {
+            return Ok(());
+        }
         let timestamp = now();
         let output = match self.role.as_ref() {
             None => {
@@ -261,7 +246,7 @@ impl Config {
         match key {
             "api_key" => {
                 if unset {
-                    return Ok("Not allowed".into());
+                    return Ok("Error: Not allowed".into());
                 } else {
                     self.api_key = value.to_string();
                 }
@@ -298,6 +283,15 @@ impl Config {
         Ok("".into())
     }
 
+    fn open_message_file(&self) -> Result<File> {
+        let path = Config::messages_file()?;
+        OpenOptions::new()
+            .create(true)
+            .append(true)
+            .open(&path)
+            .with_context(|| format!("Failed to create/append {}", path.display()))
+    }
+
     fn load_roles(&mut self) -> Result<()> {
         let path = Self::roles_file()?;
         if !path.exists() {
diff --git a/src/main.rs b/src/main.rs
index dfce165..7d79b07 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -70,7 +70,6 @@ fn start_directive(
     input: &str,
     no_stream: bool,
 ) -> Result<()> {
-    let mut file = config.borrow().open_message_file()?;
     let prompt = config.borrow().get_prompt();
     let highlight = config.borrow().highlight && stdout().is_terminal();
     let output = if no_stream {
@@ -94,7 +93,7 @@ fn start_directive(
         wg.wait();
         output
     };
-    config.borrow().save_message(file.as_mut(), input, &output)
+    config.borrow().save_message(input, &output)
 }
 
 fn start_interactive(client: ChatGptClient, config: SharedConfig) -> Result<()> {
diff --git a/src/repl/handler.rs b/src/repl/handler.rs
index 384fce0..e971ff8 100644
--- a/src/repl/handler.rs
+++ b/src/repl/handler.rs
@@ -7,7 +7,6 @@ use anyhow::Result;
 use crossbeam::channel::Sender;
 use crossbeam::sync::WaitGroup;
 use std::cell::RefCell;
-use std::fs::File;
 
 use super::abort::SharedAbortSignal;
 
@@ -23,30 +22,21 @@ pub enum ReplCmd {
 pub struct ReplCmdHandler {
     client: ChatGptClient,
     config: SharedConfig,
-    state: RefCell<ReplCmdHandlerState>,
+    reply: RefCell<String>,
     abort: SharedAbortSignal,
 }
 
-pub struct ReplCmdHandlerState {
-    reply: String,
-    save_file: Option<File>,
-}
-
 impl ReplCmdHandler {
     pub fn init(
         client: ChatGptClient,
         config: SharedConfig,
         abort: SharedAbortSignal,
     ) -> Result<Self> {
-        let save_file = config.as_ref().borrow().open_message_file()?;
-        let state = RefCell::new(ReplCmdHandlerState {
-            save_file,
-            reply: String::new(),
-        });
+        let reply = RefCell::new(String::new());
         Ok(Self {
             client,
             config,
-            state,
+            reply,
             abort,
         })
     }
@@ -55,7 +45,7 @@ impl ReplCmdHandler {
         match cmd {
             ReplCmd::Submit(input) => {
                 if input.is_empty() {
-                    self.state.borrow_mut().reply.clear();
+                    self.reply.borrow_mut().clear();
                     return Ok(());
                 }
                 let highlight = self.config.borrow().highlight;
@@ -72,12 +62,8 @@ impl ReplCmdHandler {
                 );
                 wg.wait();
                 let buffer = ret?;
-                self.config.borrow().save_message(
-                    self.state.borrow_mut().save_file.as_mut(),
-                    &input,
-                    &buffer,
-                )?;
-                self.state.borrow_mut().reply = buffer;
+                self.config.borrow().save_message(&input, &buffer)?;
+                *self.reply.borrow_mut() = buffer;
             }
             ReplCmd::SetRole(name) => {
                 let output = self.config.borrow_mut().change_role(&name);