3 files changed, 66 insertions, 2 deletions

diff --git a/crates/atuin-client/src/plugin.rs b/crates/atuin-client/src/plugin.rs
index 21a2bcef..6f351bf1 100644
--- a/crates/atuin-client/src/plugin.rs
+++ b/crates/atuin-client/src/plugin.rs
@@ -42,7 +42,7 @@ impl OfficialPluginRegistry {
                 "Update atuin to the latest version",
                 "The 'atuin update' command is provided by the atuin-update plugin.\n\
                  It is only installed if you used the install script\n  \
-                 If you used a package manager (brew, apt, etc), please continue to use it for updates"
+                 If you used a package manager (brew, apt, etc), please continue to use it for updates",
             ),
         );
     }
@@ -68,6 +68,59 @@ impl Default for OfficialPluginRegistry {
     }
 }
 
+pub struct PluginContext {
+    #[cfg(windows)]
+    _update_on_windows: Option<UpdateOnWindowsContext>,
+}
+
+impl PluginContext {
+    pub fn new(_subcommand: &str) -> Self {
+        PluginContext {
+            #[cfg(windows)]
+            _update_on_windows: (_subcommand == "update").then(UpdateOnWindowsContext::new),
+        }
+    }
+}
+
+impl Drop for PluginContext {
+    fn drop(&mut self) {}
+}
+
+#[cfg(windows)]
+struct UpdateOnWindowsContext {
+    initial_exe: Option<std::path::PathBuf>,
+}
+
+#[cfg(windows)]
+impl UpdateOnWindowsContext {
+    const OLD_FILE_NAME: &'static str = "atuin.old";
+
+    pub fn new() -> Self {
+        // Windows doesn't let you overwrite a running exe, but it lets you rename it,
+        // so make some room for atuin-update to install the new version.
+        let initial_exe = std::env::current_exe().ok().and_then(|exe| {
+            std::fs::rename(&exe, exe.with_file_name(Self::OLD_FILE_NAME)).ok()?;
+            Some(exe)
+        });
+
+        Self { initial_exe }
+    }
+}
+
+#[cfg(windows)]
+impl Drop for UpdateOnWindowsContext {
+    fn drop(&mut self) {
+        if let Some(exe) = &self.initial_exe
+            && !exe.exists()
+        {
+            // The update failed, roll back the current exe to its initial name.
+            std::fs::rename(exe.with_file_name(Self::OLD_FILE_NAME), exe).unwrap_or_else(|e| {
+                eprintln!("Failed to roll back the update, you may need to reinstall Atuin: {e}");
+            });
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/crates/atuin/src/command/client/daemon.rs b/crates/atuin/src/command/client/daemon.rs
index cd769373..483829e9 100644
--- a/crates/atuin/src/command/client/daemon.rs
+++ b/crates/atuin/src/command/client/daemon.rs
@@ -331,6 +331,7 @@ async fn wait_until_ready(settings: &Settings, timeout: Duration) -> Result<Hist
     }
 }
 
+#[allow(clippy::unnecessary_wraps)]
 fn ensure_autostart_supported(settings: &Settings) -> Result<()> {
     #[cfg(unix)]
     if settings.daemon.systemd_socket {
diff --git a/crates/atuin/src/command/external.rs b/crates/atuin/src/command/external.rs
index 657aea56..5d875e9d 100644
--- a/crates/atuin/src/command/external.rs
+++ b/crates/atuin/src/command/external.rs
@@ -3,7 +3,7 @@ use std::process::Command;
 use std::{io, process};
 
 #[cfg(feature = "client")]
-use atuin_client::plugin::OfficialPluginRegistry;
+use atuin_client::plugin::{OfficialPluginRegistry, PluginContext};
 use clap::CommandFactory;
 use clap::builder::{StyledStr, Styles};
 use eyre::Result;
@@ -16,6 +16,9 @@ pub fn run(args: &[String]) -> Result<()> {
     let mut cmd = Command::new(&bin);
     cmd.args(&args[1..]);
 
+    #[cfg(feature = "client")]
+    let context = PluginContext::new(subcommand);
+
     let spawn_result = match cmd.spawn() {
         Ok(child) => Ok(child),
         Err(e) => match e.kind() {
@@ -33,11 +36,18 @@ pub fn run(args: &[String]) -> Result<()> {
             if status.success() {
                 Ok(())
             } else {
+                #[cfg(feature = "client")]
+                drop(context);
+
                 process::exit(status.code().unwrap_or(1));
             }
         }
         Err(e) => {
             eprintln!("{}", e.ansi());
+
+            #[cfg(feature = "client")]
+            drop(context);
+
             process::exit(1);
         }
     }