1 files changed, 17 insertions, 4 deletions

diff --git a/cmd/deploy/host/update.go b/cmd/deploy/host/update.go
index 6f1b43b..aa47ed8 100644
--- a/cmd/deploy/host/update.go
+++ b/cmd/deploy/host/update.go
@@ -40,11 +40,11 @@ func runUpdate(cmd *cobra.Command, args []string) error {
         yes, _ := cmd.Flags().GetBool("yes")
         if !yes {
                 fmt.Printf("This will run apt update && apt upgrade on %s\n", host)
-                fmt.Print("Continue? [y/N]: ")
+                fmt.Print("Continue? [Y/n]: ")
                 reader := bufio.NewReader(os.Stdin)
                 response, _ := reader.ReadString('\n')
                 response = strings.TrimSpace(response)
-                if response != "y" && response != "Y" {
+                if response == "n" || response == "N" {
                         fmt.Println("Aborted.")
                         return nil
                 }
@@ -70,8 +70,21 @@ func runUpdate(cmd *cobra.Command, args []string) error {
 
         fmt.Println()
         if output, err := client.Run("[ -f /var/run/reboot-required ] && echo 'yes' || echo 'no'"); err == nil {
-                if output == "yes\n" {
-                        fmt.Println("Note: A reboot is required to complete the update.")
+                if strings.TrimSpace(output) == "yes" {
+                        fmt.Print("A reboot is required to complete the update. Reboot now? [Y/n]: ")
+                        reader := bufio.NewReader(os.Stdin)
+                        response, _ := reader.ReadString('\n')
+                        response = strings.TrimSpace(response)
+                        if response == "" || response == "y" || response == "Y" {
+                                fmt.Println("Rebooting...")
+                                if _, err := client.RunSudo("reboot"); err != nil {
+                                        // reboot command often returns an error as connection drops
+                                        // this is expected behavior
+                                }
+                                fmt.Println("Reboot initiated. The host will be back online shortly.")
+                                return nil
+                        }
+                        fmt.Println("Skipping reboot. Run 'sudo reboot' manually when ready.")
                 }
         }