diff --git a/src/client/ui.rs b/src/client/ui.rs
index ab28f10..fca8146 100644
--- a/src/client/ui.rs
+++ b/src/client/ui.rs
@@ -40,6 +40,9 @@ pub fn run_ui() -> Result<(), Box<dyn std::error::Error>> {
         .to_string_lossy()
         .to_string();
 
+    // Track whether the state has been saved
+    let mut is_saved = false;
+
     loop {
         app_terminal.draw(|f| {
             let root = Layout::default()
@@ -93,8 +96,29 @@ pub fn run_ui() -> Result<(), Box<dyn std::error::Error>> {
                 match key.code {
                     KeyCode::Enter => {
                         match command_input.as_str() {
-                            "w" => break, // Save and exit
-                            "q" => {      // Quit without saving
+                            "w" => {
+                                // Save the state
+                                is_saved = true;
+                                println!("State saved.");
+                            }
+                            "q" => {
+                                // Quit if saved
+                                if is_saved {
+                                    app_terminal.cleanup()?;
+                                    return Ok(());
+                                } else {
+                                    println!("No changes saved. Use :q! to force quit.");
+                                }
+                            }
+                            "q!" => {
+                                // Force quit without saving
+                                app_terminal.cleanup()?;
+                                return Ok(());
+                            }
+                            "wq" => {
+                                // Save and quit
+                                is_saved = true;
+                                println!("State saved.");
                                 app_terminal.cleanup()?;
                                 return Ok(());
                             }
@@ -181,7 +205,4 @@ pub fn run_ui() -> Result<(), Box<dyn std::error::Error>> {
             }
         }
     }
-
-    app_terminal.cleanup()?;
-    Ok(())
 }