2 files changed, 29 insertions, 0 deletions
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 08509be..1b3d802 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -283,6 +283,24 @@ impl Config {
         Ok(())
     }
 
+    pub fn get_state(&self) -> State {
+        if let Some(session) = &self.session {
+            if session.is_empty() {
+                if session.role.is_some() {
+                    State::EmptySessionWithRole
+                } else {
+                    State::EmptySession
+                }
+            } else {
+                State::Session
+            }
+        } else if self.role.is_some() {
+            State::Role
+        } else {
+            State::Normal
+        }
+    }
+
     pub fn get_temperature(&self) -> Option<f64> {
         self.temperature
     }
@@ -806,6 +824,15 @@ impl Keybindings {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
+pub enum State {
+    Normal,
+    Role,
+    EmptySession,
+    EmptySessionWithRole,
+    Session,
+}
+
 fn create_config_file(config_path: &Path) -> Result<()> {
     let ans = Confirm::new("No config file, create a new one?")
         .with_default(true)
diff --git a/src/config/session.rs b/src/config/session.rs
index 644255f..e135c58 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -20,6 +20,7 @@ pub struct Session {
     model_id: String,
     temperature: Option<f64>,
     messages: Vec<Message>,
+    #[serde(default)]
     data_urls: HashMap<String, String>,
     #[serde(skip)]
     pub name: String,
@@ -248,6 +249,7 @@ impl Session {
             role: MessageRole::Assistant,
             content: MessageContent::Text(output.to_string()),
         });
+        self.role = None;
         self.dirty = true;
         Ok(())
     }