fix: serialize complete upgrade operations

src/cli/upgrade.rs

@@ -1620,6 +1620,7 @@ impl Cli {
         options: UpgradeOptions,
     ) -> Result<UpgradeEnvSummary, String> {
         let _transaction_lock = lock_upgrade_transaction(name, &self.env, &self.cwd)?;
+        let _operation_lock = self.environment_service().lock_operation(name)?;
         self.upgrade_env_locked(name, target, options)
     }
 
@@ -1694,7 +1695,7 @@ impl Cli {
                     ),
                 });
             }
-            let mut transaction = self.begin_upgrade_transaction(
+            let mut transaction = self.begin_upgrade_transaction_locked(
                 env_name,
                 UpgradeTransactionPlan {
                     source: UpgradeHistoryBinding {
@@ -1710,6 +1711,8 @@ impl Cli {
                 },
                 &[current.name.clone(), target_runtime_name.clone()],
                 options.rollback_enabled,
+                "pre-upgrade",
+                None,
             )?;
             let prepared = match self.prepare_isolated_upgrade_target(env_name, target, resolved) {
                 Ok(prepared) => prepared,
@@ -1756,7 +1759,7 @@ impl Cli {
             };
             let publish_result = if binding_changed {
                 self.environment_service()
-                    .set_runtime(env_name, prepared.name.as_str())
+                    .set_runtime_locked(env_name, prepared.name.as_str())
                     .map(|_| ())
             } else {
                 self.runtime_service().refresh_supervisor_if_present()
@@ -1774,8 +1777,12 @@ impl Cli {
                     format!("failed to publish upgraded runtime: {error}"),
                 );
             }
-            let service_result =
-                self.reconcile_upgraded_service(env_name, service.as_ref(), binding_changed, true);
+            let service_result = self.reconcile_upgraded_service_locked(
+                env_name,
+                service.as_ref(),
+                binding_changed,
+                true,
+            );
             let (service_action, service_note) = match service_result {
                 Ok(result) => result,
                 Err(error) => {
@@ -1920,7 +1927,7 @@ impl Cli {
                     ),
                 });
             }
-            let mut transaction = self.begin_upgrade_transaction(
+            let mut transaction = self.begin_upgrade_transaction_locked(
                 env_name,
                 UpgradeTransactionPlan {
                     source: UpgradeHistoryBinding {
@@ -1936,6 +1943,8 @@ impl Cli {
                 },
                 &[current.name.clone(), target_runtime_name.clone()],
                 options.rollback_enabled,
+                "pre-upgrade",
+                None,
             )?;
             let prepared = match self.prepare_isolated_upgrade_target(env_name, &target, resolved) {
                 Ok(prepared) => prepared,
@@ -1999,7 +2008,7 @@ impl Cli {
                 );
             }
             let service_result =
-                self.reconcile_upgraded_service(env_name, service.as_ref(), false, changed);
+                self.reconcile_upgraded_service_locked(env_name, service.as_ref(), false, changed);
             let (service_action, service_note) = match service_result {
                 Ok(result) => result,
                 Err(error) => {
@@ -2089,7 +2098,7 @@ impl Cli {
                 note: Some("dry run: no runtime, env, service, or snapshot changed".to_string()),
             });
         }
-        let mut transaction = self.begin_upgrade_transaction(
+        let mut transaction = self.begin_upgrade_transaction_locked(
             env_name,
             UpgradeTransactionPlan {
                 source: UpgradeHistoryBinding {
@@ -2105,6 +2114,8 @@ impl Cli {
             },
             std::slice::from_ref(&current.name),
             options.rollback_enabled,
+            "pre-upgrade",
+            None,
         )?;
         let updated = match self.with_progress(format!("Updating runtime {}", current.name), || {
             self.with_isolated_runtime_mutation(env_name, &current.name, || {
@@ -2162,7 +2173,7 @@ impl Cli {
             );
         }
         let service_result =
-            self.reconcile_upgraded_service(env_name, service.as_ref(), false, true);
+            self.reconcile_upgraded_service_locked(env_name, service.as_ref(), false, true);
         let (service_action, service_note) = match service_result {
             Ok(result) => result,
             Err(error) => {
@@ -2301,7 +2312,7 @@ impl Cli {
             });
         }
 
-        let mut transaction = self.begin_upgrade_transaction(
+        let mut transaction = self.begin_upgrade_transaction_locked(
             env_name,
             UpgradeTransactionPlan {
                 source: UpgradeHistoryBinding {
@@ -2317,6 +2328,8 @@ impl Cli {
             },
             std::slice::from_ref(&target_runtime_name),
             options.rollback_enabled,
+            "pre-upgrade",
+            None,
         )?;
         let prepared = match self.prepare_isolated_upgrade_target(env_name, target, resolved) {
             Ok(prepared) => prepared,
@@ -2362,7 +2375,7 @@ impl Cli {
         };
         if let Err(error) = self
             .environment_service()
-            .set_runtime(env_name, prepared.name.as_str())
+            .set_runtime_locked(env_name, prepared.name.as_str())
         {
             return self.rollback_failed_upgrade(
                 env_name,
@@ -2377,7 +2390,7 @@ impl Cli {
             );
         }
         let service_result =
-            self.reconcile_upgraded_service(env_name, service.as_ref(), true, true);
+            self.reconcile_upgraded_service_locked(env_name, service.as_ref(), true, true);
         let (service_action, service_note) = match service_result {
             Ok(result) => result,
             Err(error) => {
@@ -2613,7 +2626,7 @@ impl Cli {
         operation()
     }
 
-    fn reconcile_upgraded_service(
+    fn reconcile_upgraded_service_locked(
         &self,
         env_name: &str,
         service: Option<&ServiceSummary>,
@@ -2633,7 +2646,7 @@ impl Cli {
         if service.running {
             let restart = self
                 .with_progress(format!("Restarting service for {env_name}"), || {
-                    self.service_service().restart(env_name)
+                    self.service_service().restart_locked(env_name)
                 })?;
             let note = join_optional_warnings(
                 join_warnings(&restart.warnings),
@@ -2644,7 +2657,7 @@ impl Cli {
 
         if binding_changed || runtime_changed {
             let start = self.with_progress(format!("Starting service for {env_name}"), || {
-                self.service_service().start(env_name)
+                self.service_service().start_locked(env_name)
             })?;
             let note = join_optional_warnings(
                 join_warnings(&start.warnings),
@@ -2877,24 +2890,6 @@ impl Cli {
             .map_err(|error| format!("{name} failed: {error}"))
     }
 
-    fn begin_upgrade_transaction(
-        &self,
-        env_name: &str,
-        plan: UpgradeTransactionPlan,
-        runtime_names: &[String],
-        rollback_enabled: bool,
-    ) -> Result<UpgradeTransaction, String> {
-        let _operation_lock = self.environment_service().lock_operation(env_name)?;
-        self.begin_upgrade_transaction_locked(
-            env_name,
-            plan,
-            runtime_names,
-            rollback_enabled,
-            "pre-upgrade",
-            None,
-        )
-    }
-
     #[allow(clippy::too_many_arguments)]
     fn begin_upgrade_transaction_locked(
         &self,
@@ -3241,7 +3236,7 @@ impl Cli {
             return Ok(summary);
         }
 
-        let rollback_result = self.rollback_upgrade(env_name, &transaction);
+        let rollback_result = self.rollback_upgrade_locked(env_name, &transaction);
         let snapshot_id = transaction.snapshot_id.clone();
         let mut summary = match rollback_result {
             Ok(()) => UpgradeEnvSummary {
@@ -3287,32 +3282,13 @@ impl Cli {
         Ok(summary)
     }
 
-    fn rollback_upgrade(
+    fn rollback_upgrade_locked(
         &self,
         env_name: &str,
         transaction: &UpgradeTransaction,
     ) -> Result<(), String> {
         // Restore runtime bytes and metadata before the snapshot republishes supervisor
         // state; otherwise rollback can briefly advertise the failed runtime revision.
-        for runtime_backup in &transaction.runtime_backups {
-            self.restore_runtime_backup(runtime_backup)?;
-        }
-        self.environment_service()
-            .restore_snapshot(RestoreEnvSnapshotOptions {
-                env_name: env_name.to_string(),
-                snapshot_id: transaction.snapshot_id.clone(),
-            })?;
-        for runtime_name in &transaction.created_runtime_names {
-            self.remove_runtime_created_during_upgrade(runtime_name)?;
-        }
-        Ok(())
-    }
-
-    fn rollback_upgrade_locked(
-        &self,
-        env_name: &str,
-        transaction: &UpgradeTransaction,
-    ) -> Result<(), String> {
         for runtime_backup in &transaction.runtime_backups {
             self.restore_runtime_backup(runtime_backup)?;
         }

src/env/binding.rs

@@ -22,6 +22,14 @@ impl<'a> EnvironmentService<'a> {
 
     pub fn set_runtime(&self, name: &str, runtime_name: &str) -> Result<EnvMeta, String> {
         let _lock = self.lock_operation(name)?;
+        self.set_runtime_locked(name, runtime_name)
+    }
+
+    pub(crate) fn set_runtime_locked(
+        &self,
+        name: &str,
+        runtime_name: &str,
+    ) -> Result<EnvMeta, String> {
         let mut meta = get_environment(name, self.env, self.cwd)?;
         if runtime_name.eq_ignore_ascii_case("none") {
             meta.default_runtime = None;

tests/upgrade_command_tests.rs

@@ -3,7 +3,7 @@ mod support;
 use std::collections::BTreeMap;
 use std::fs;
 use std::path::{Path, PathBuf};
-use std::process::Command;
+use std::process::{Command, Stdio};
 use std::thread::{self, sleep};
 use std::time::Duration;
 
@@ -136,6 +136,14 @@ case "$1" in
     ;;
   update)
     if [ "$2" = "finalize" ]; then
+      if [ -n "${{OCM_TEST_UPDATE_FINALIZE_STARTED:-}}" ]; then
+        : > "$OCM_TEST_UPDATE_FINALIZE_STARTED"
+      fi
+      if [ -n "${{OCM_TEST_UPDATE_FINALIZE_RELEASE:-}}" ]; then
+        while [ ! -e "$OCM_TEST_UPDATE_FINALIZE_RELEASE" ]; do
+          sleep 0.05
+        done
+      fi
       if [ "${{OCM_TEST_FAIL_UPDATE_FINALIZE:-}}" = "1" ]; then
         echo "forced update finalize failure" >&2
         exit 23
@@ -2517,6 +2525,110 @@ fn upgrade_can_switch_env_to_an_installed_runtime() {
     );
 }
 
+#[test]
+fn upgrade_holds_the_environment_operation_lock_until_completion() {
+    let root = TestDir::new("upgrade-operation-lock");
+    let cwd = root.child("workspace");
+    fs::create_dir_all(&cwd).unwrap();
+
+    let old_runtime = root.child("old-openclaw");
+    let new_runtime = root.child("new-openclaw");
+    write_executable_script(&old_runtime, &recording_openclaw_script("old-openclaw"));
+    write_executable_script(&new_runtime, &recording_openclaw_script("new-openclaw"));
+
+    let mut env = ocm_env(&root);
+    for (name, runtime) in [("old-local", &old_runtime), ("new-local", &new_runtime)] {
+        let add = run_ocm(
+            &cwd,
+            &env,
+            &[
+                "runtime",
+                "add",
+                name,
+                "--path",
+                &runtime.display().to_string(),
+            ],
+        );
+        assert!(add.status.success(), "{}", stderr(&add));
+    }
+
+    let create = run_ocm(
+        &cwd,
+        &env,
+        &["env", "create", "demo", "--runtime", "old-local"],
+    );
+    assert!(create.status.success(), "{}", stderr(&create));
+
+    let finalize_started = root.child("upgrade-finalize-started");
+    let finalize_release = root.child("upgrade-finalize-release");
+    env.insert(
+        "OCM_TEST_UPDATE_FINALIZE_STARTED".to_string(),
+        path_string(&finalize_started),
+    );
+    env.insert(
+        "OCM_TEST_UPDATE_FINALIZE_RELEASE".to_string(),
+        path_string(&finalize_release),
+    );
+
+    let mut upgrade_command = Command::new(env!("CARGO_BIN_EXE_ocm"));
+    upgrade_command
+        .current_dir(&cwd)
+        .args(["upgrade", "demo", "--runtime", "new-local"])
+        .env_clear()
+        .envs(&env)
+        .stdin(Stdio::null())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped());
+    let upgrade = upgrade_command.spawn().unwrap();
+
+    for _ in 0..200 {
+        if finalize_started.exists() {
+            break;
+        }
+        sleep(Duration::from_millis(25));
+    }
+    if !finalize_started.exists() {
+        fs::write(&finalize_release, "").unwrap();
+        let output = upgrade.wait_with_output().unwrap();
+        panic!(
+            "upgrade did not reach finalization: {}",
+            String::from_utf8_lossy(&output.stderr)
+        );
+    }
+
+    let mut binding_command = Command::new(env!("CARGO_BIN_EXE_ocm"));
+    binding_command
+        .current_dir(&cwd)
+        .args(["env", "set-runtime", "demo", "old-local"])
+        .env_clear()
+        .envs(&env)
+        .stdin(Stdio::null())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped());
+    let mut binding = binding_command.spawn().unwrap();
+    sleep(Duration::from_millis(250));
+
+    if let Some(status) = binding.try_wait().unwrap() {
+        fs::write(&finalize_release, "").unwrap();
+        let upgrade = upgrade.wait_with_output().unwrap();
+        panic!(
+            "binding mutation completed with {status} while upgrade held the environment operation lock; upgrade stderr: {}",
+            String::from_utf8_lossy(&upgrade.stderr)
+        );
+    }
+
+    fs::write(&finalize_release, "").unwrap();
+    let upgrade = upgrade.wait_with_output().unwrap();
+    assert!(upgrade.status.success(), "{}", stderr(&upgrade));
+    let binding = binding.wait_with_output().unwrap();
+    assert!(binding.status.success(), "{}", stderr(&binding));
+
+    let show = run_ocm(&cwd, &env, &["env", "show", "demo", "--json"]);
+    assert!(show.status.success(), "{}", stderr(&show));
+    let env_json: Value = serde_json::from_str(&stdout(&show)).unwrap();
+    assert_eq!(env_json["defaultRuntime"], "old-local");
+}
+
 #[test]
 fn upgrade_reuses_the_bound_named_runtime_without_retaining_recovery_bytes() {
     let root = TestDir::new("upgrade-reuse-bound-runtime");