Client: ask before updating, wait for running work; reopen this computer's AI session

Mastertech4.0/src/app_state.rs

@@ -165,6 +165,13 @@ pub struct MastertechContext {
     pub copied_items_rx: Receiver<String>,
     pub github_releases: Vec<GithubRelease>,
     pub github_releases_channel: (Sender<Vec<GithubRelease>>, Receiver<Vec<GithubRelease>>),
+    /// Version of the release being downloaded.
+    pub announced_release: Option<String>,
+    pub pending_update: Option<crate::utilities::update_policy::PendingUpdate>,
+    /// The menu's Update was clicked, so the next finished download installs without a prompt.
+    pub update_requested: bool,
+    pub update_mode: crate::utilities::update_policy::UpdateMode,
+    pub next_update_check: Option<std::time::Instant>,
     pub seb_channel: (Sender<Vec<CarboniteResponse>>, Receiver<Vec<CarboniteResponse>>),
     pub get_settings: bool,
     pub client_friendly_name: String,
@@ -388,6 +395,11 @@ impl MasterTechApp {
             progress: (0.0, 0.0),
             special_part_order: Default::default(),
             github_releases: Default::default(),
+            announced_release: None,
+            pending_update: None,
+            update_requested: false,
+            update_mode: crate::utilities::update_policy::load_mode(),
+            next_update_check: None,
             toolbox: FileSystem::new(),
 
             // Data table shit

Mastertech4.0/src/data/receive_github.rs

@@ -1,34 +1,38 @@
-use egui::ViewportCommand;
+use std::time::{Duration, Instant};
+
+use displays::modals::update_toast::{self, UpdateChoice};
+use displays::{ToastMessage, get_toast_sender};
+use egui::{UserAttentionType, ViewportCommand};
 use log::{debug, error, info};
 use semver::Version;
 use tokio::spawn;
-use displays::{get_toast_sender, ToastMessage};
 
+use crate::utilities::update_policy::{Activity, PendingUpdate, SNOOZE, UpdateStep, next_step};
 use crate::{app_state::MasterTechApp, tabs::github::self_updater::run};
 
 impl MasterTechApp {
     pub fn receive_github(&mut self, ctx: &eframe::egui::Context) {
         // Track download progress for toast notifications
         while let Ok(res) = self.context.bytes_rx.try_recv() {
             ctx.request_repaint();
-            
+
             // Calculate previous progress percentage before updating
             let prev_pct = if self.context.progress.1 > 0.0 {
                 (self.context.progress.0 / self.context.progress.1 * 100.0) as u32
             } else {
                 0
             };
-            
+
             self.context.progress.1 = res.1 as f32;
             self.context.progress.0 = res.0 as f32;
-            
+
             // Calculate current progress percentage
             let current_pct = if self.context.progress.1 > 0.0 {
                 (self.context.progress.0 / self.context.progress.1 * 100.0) as u32
             } else {
                 0
             };
-            
+
             // Show progress toast when crossing milestones (25%, 50%, 75%)
             for milestone in [25u32, 50, 75] {
                 if prev_pct < milestone && current_pct >= milestone {
@@ -38,51 +42,30 @@ impl MasterTechApp {
                     ));
                 }
             }
-            
+
             if res.1 > 0 && res.0 == res.1 {
                 self.context.progress = (0.0, 0.0);
-                #[cfg(target_os = "windows")]
-                {
-                    use crate::utilities::safe_swap;
-                    let toast_tx = get_toast_sender();
-                    let applied = safe_swap::staged_update_path()
-                        .and_then(|staged| safe_swap::apply_staged_update(&staged, res.1));
-                    match applied {
-                        Ok(exe) => match safe_swap::relaunch(&exe, &[]) {
-                            Ok(()) => {
-                                let _ = toast_tx.try_send(ToastMessage::Success(
-                                    "Update installed! Restarting...".to_string(),
-                                ));
-                                ctx.send_viewport_cmd(ViewportCommand::Close);
-                            }
-                            Err(e) => {
-                                log::error!("update installed but relaunch failed: {e:?}");
-                                let _ = toast_tx.try_send(ToastMessage::Warning(
-                                    "Update installed — restart MasterTech to finish.".to_string(),
-                                ));
-                            }
-                        },
-                        Err(e) => {
-                            log::error!("self-update failed: {e:?}");
-                            if let Ok(staged) = safe_swap::staged_update_path() {
-                                let _ = std::fs::remove_file(staged);
-                            }
-                            let _ = toast_tx.try_send(ToastMessage::Error(format!(
-                                "Update failed — still running the current version. {e}"
-                            )));
-                        }
-                    }
-                }
+                let version = self
+                    .context
+                    .announced_release
+                    .clone()
+                    .unwrap_or_else(|| "update".to_string());
+                info!("update {version} downloaded ({} bytes)", res.1);
+                let accepted = std::mem::take(&mut self.context.update_requested);
+                self.context.pending_update = Some(PendingUpdate::new(res.1, version, accepted));
+                self.context.next_update_check = None;
             }
         }
 
+        self.drive_pending_update(ctx);
+
         if let Ok(releases) = self.context.github_releases_channel.1.try_recv() {
             debug!("Releases: {releases:?}");
             let os = std::env::consts::OS;
             let current_version =
                 Version::parse(env!("CARGO_PKG_VERSION")).expect("Invalid version format");
 
-            for release in releases.iter() {
+            for release in releases.iter().filter(|r| !r.draft && !r.prerelease) {
                 let Ok(github_release_version) =
                     Version::parse(release.tag_name.trim_start_matches('v'))
                 else {
@@ -104,6 +87,7 @@ impl MasterTechApp {
 
                 let client = self.context.client.clone();
                 info!("Found a new release! {:?}", &github_release_version);
+                self.context.announced_release = Some(github_release_version.to_string());
 
                 let toast_tx = get_toast_sender();
                 let _ = toast_tx.try_send(ToastMessage::Info(format!(
@@ -122,4 +106,96 @@ impl MasterTechApp {
             self.context.github_releases = releases;
         }
     }
-}
+
+    /// Installs, offers or holds a downloaded update, checking at most every two seconds.
+    fn drive_pending_update(&mut self, ctx: &eframe::egui::Context) {
+        if self.context.pending_update.is_none() {
+            return;
+        }
+        let now = Instant::now();
+        if self.context.next_update_check.is_some_and(|at| now < at) {
+            return;
+        }
+        self.context.next_update_check = Some(now + Duration::from_secs(2));
+
+        let reason = Activity::current().reason();
+        let Some(pending) = self.context.pending_update.as_mut() else {
+            return;
+        };
+        match update_toast::take_choice() {
+            Some(UpdateChoice::Now) => pending.accepted = true,
+            Some(UpdateChoice::Later) => pending.snoozed_until = Some(now + SNOOZE),
+            None => {}
+        }
+
+        match next_step(self.context.update_mode, reason.is_some(), pending.accepted) {
+            UpdateStep::Wait => {
+                if pending.waiting_on != reason {
+                    let why = reason.as_deref().unwrap_or_default();
+                    info!("update {} waits: {why}", pending.version);
+                    if pending.accepted && pending.waiting_on.is_none() {
+                        let _ = get_toast_sender().try_send(ToastMessage::Info(format!(
+                            "MasterTech will update once nothing is running ({why})."
+                        )));
+                    }
+                    pending.waiting_on = reason;
+                }
+            }
+            UpdateStep::Prompt => {
+                pending.waiting_on = None;
+                if pending.snoozed_until.is_some_and(|until| now < until) {
+                    return;
+                }
+                if update_toast::post(&mut self.context.shared_ctx.toasts, &pending.version) {
+                    ctx.send_viewport_cmd(ViewportCommand::RequestUserAttention(
+                        UserAttentionType::Informational,
+                    ));
+                }
+            }
+            UpdateStep::Install => {
+                update_toast::retire();
+                let size = pending.size;
+                self.context.pending_update = None;
+                install_update(ctx, size);
+            }
+        }
+    }
+}
+
+/// Swaps in the staged executable and relaunches MasterTech.
+fn install_update(ctx: &eframe::egui::Context, size: u64) {
+    #[cfg(target_os = "windows")]
+    {
+        use crate::utilities::safe_swap;
+        let toast_tx = get_toast_sender();
+        let applied = safe_swap::staged_update_path()
+            .and_then(|staged| safe_swap::apply_staged_update(&staged, size));
+        match applied {
+            Ok(exe) => match safe_swap::relaunch(&exe, &[]) {
+                Ok(()) => {
+                    let _ = toast_tx.try_send(ToastMessage::Success(
+                        "Update installed! Restarting...".to_string(),
+                    ));
+                    ctx.send_viewport_cmd(ViewportCommand::Close);
+                }
+                Err(e) => {
+                    log::error!("update installed but relaunch failed: {e:?}");
+                    let _ = toast_tx.try_send(ToastMessage::Warning(
+                        "Update installed — restart MasterTech to finish.".to_string(),
+                    ));
+                }
+            },
+            Err(e) => {
+                log::error!("self-update failed: {e:?}");
+                if let Ok(staged) = safe_swap::staged_update_path() {
+                    let _ = std::fs::remove_file(staged);
+                }
+                let _ = toast_tx.try_send(ToastMessage::Error(format!(
+                    "Update failed — still running the current version. {e}"
+                )));
+            }
+        }
+    }
+    #[cfg(not(target_os = "windows"))]
+    let _ = (ctx, size);
+}

Mastertech4.0/src/first_run.rs

@@ -243,6 +243,27 @@ impl MasterTechApp {
                         log::error!("Error getting github releases: {e:?}");
                     }
                 });
+
+                crate::utilities::update_policy::spawn_agent_poll();
+                // Reopens this machine's AI session in the Ai tab.
+                spawn(async move {
+                    let Ok(cs) = tokio::task::spawn_blocking(|| {
+                        crate::filesystem::get_client_hash().connection_string
+                    })
+                    .await
+                    else {
+                        return;
+                    };
+                    match database::schema::agent_thread::AgentThread::reopen_for_connection(&cs)
+                        .await
+                    {
+                        Ok(Some(thread)) => {
+                            displays::ui_data::agent_session_notify::request_open(thread.id, true)
+                        }
+                        Ok(None) => {}
+                        Err(e) => log::debug!("first_run -> session reopen lookup failed: {e}"),
+                    }
+                });
             } 
         }
 

Mastertech4.0/src/pages/menu_bar.rs

@@ -2,6 +2,7 @@ use eframe::egui::{Button, Color32, ComboBox, FontId, Frame, Layout, RichText, S
 use database::{schema::{utilities::{get_completed_tasks_for_store, get_store_users, get_tasks_for_store}, Store}, db};
 use egui::{PopupCloseBehavior, UiKind, containers::menu::{MenuButton, MenuConfig}, style::StyleModifier};
 use crate::tabs::github::self_updater::run;
+use crate::utilities::update_policy::{self, UpdateMode};
 use displays::{app_state::{default_tree, AppState, MainPages}, pages::view_menu, plugins::push_widget_anchor, tabs::{TabContext, WorkMode}, ui_tools::{store_picker::presta_store_options, theme}, TaskUiActions};
 use crate::app_state::MasterTechApp;
 use std::collections::BTreeSet;
@@ -169,12 +170,36 @@ impl MasterTechApp {
                                 .ui(ui)
                                 .clicked()
                                 {
-                                    let client = self.context.client.clone();
-                                    let tx = self.context.bytes_tx.clone();
-
-                                    spawn(async move {
-                                        let _ = run(client, tx.clone()).await;
-                                    });
+                                    if let Some(pending) = self.context.pending_update.as_mut() {
+                                        pending.accepted = true;
+                                    } else {
+                                        self.context.update_requested = true;
+                                        let client = self.context.client.clone();
+                                        let tx = self.context.bytes_tx.clone();
+
+                                        spawn(async move {
+                                            let _ = run(client, tx.clone()).await;
+                                        });
+                                    }
+                                }
+
+                                let mut auto = self.context.update_mode == UpdateMode::Auto;
+                                if ui
+                                    .checkbox(&mut auto, "Install updates automatically")
+                                    .on_hover_text(
+                                        "Off: a notification asks before updating. Either way, updates wait \
+                                         while remote work or an AI session runs on this computer.",
+                                    )
+                                    .changed()
+                                {
+                                    let mode = if auto { UpdateMode::Auto } else { UpdateMode::Prompt };
+                                    self.context.update_mode = mode;
+                                    if let Err(e) = update_policy::save_mode(mode) {
+                                        error!("menu_bar -> saving the update mode failed: {e}");
+                                        let _ = displays::get_toast_sender().try_send(displays::ToastMessage::Error(
+                                            format!("Couldn't save the update setting: {e}"),
+                                        ));
+                                    }
                                 }
 
                                 if Button::new(

Mastertech4.0/src/terminal_mode/websockets/mod.rs

@@ -2972,6 +2972,7 @@ if (Test-Path $path) {{
                 // (CPS key fetch, antivirus installs) execute in the background.
                 let tx = self.command_tx.clone();
                 tokio::spawn(async move {
+                let _script_run = crate::utilities::update_policy::ScriptRunGuard::start();
 
                 let task_epoch = Instant::now();
                 // Millis from task_epoch when the current script started; AtomicU64 keeps the future Send.

Mastertech4.0/src/utilities/mod.rs

@@ -9,6 +9,7 @@ pub mod scripts;
 
 pub mod network;
 pub mod no_window;
+pub mod update_policy;
 
 pub use crypto::{load_encrypted_user_data, save_encrypted_user_data};