From a6d87cbf4b9bfce876e666e886a9b17a9ff7a623 Mon Sep 17 00:00:00 2001
From: romainPellerin <591564+romainPellerin@users.noreply.github.com>
Date: Wed, 30 Sep 2026 08:46:07 +0200
Subject: [PATCH] fix(vm): recover overlay journal before generation cleanup

Signed-off-by: romainPellerin <591564+romainPellerin@users.noreply.github.com>
---
 crates/openshell-driver-vm/README.md     |  5 +++
 crates/openshell-driver-vm/src/driver.rs | 40 +++++++++++++++++++++++
 e2e/rust/tests/vm_overlay.rs             | 41 ++++++++++++++++++++++++
 3 files changed, 86 insertions(+)

diff --git a/crates/openshell-driver-vm/README.md b/crates/openshell-driver-vm/README.md
index b2103a6957..5259d4ef8d 100644
--- a/crates/openshell-driver-vm/README.md
+++ b/crates/openshell-driver-vm/README.md
@@ -244,6 +244,11 @@ marked sandboxes without launching compute. Start removes the marker and uses
 the normal persisted restore path with the existing overlay. Delete removes the
 entire sandbox state directory, including a stop marker and overlay.
 
+After stopping the VM, the driver recovers the overlay's ext4 journal before
+removing the previous generation's guest authentication files. If recovery
+fails, cleanup returns an error and retains the host generation markers so a
+later attempt can recover the disk before changing it.
+
 The host control writes and syncs a terminal tombstone when the canonical main
 process exits, before it reports completion and while it retains the boundary
 for exec and forwarding. Driver startup reports that sandbox as terminal
diff --git a/crates/openshell-driver-vm/src/driver.rs b/crates/openshell-driver-vm/src/driver.rs
index 58ad21892e..98ce42dcf2 100644
--- a/crates/openshell-driver-vm/src/driver.rs
+++ b/crates/openshell-driver-vm/src/driver.rs
@@ -6178,6 +6178,7 @@ async fn remove_runtime_generation_material(state_dir: &Path) -> Result<(), Stri
         let overlay = sandbox_runtime_disk_paths(state_dir).overlay_disk;
         let generation_for_cleanup = generation.clone();
         tokio::task::spawn_blocking(move || {
+            recover_rootfs_image(&overlay)?;
             let tls = guest_boundary_tls_paths(&generation_for_cleanup);
             for guest_path in [
                 PathBuf::from(guest_boundary_config_path(&generation_for_cleanup)),
@@ -8888,6 +8889,45 @@ mod tests {
         )
     }
 
+    #[tokio::test]
+    async fn generation_cleanup_preserves_markers_when_overlay_recovery_fails() {
+        let temp = tempfile::tempdir().unwrap();
+        let state_dir = temp.path();
+        tokio::fs::write(
+            state_dir.join(HOST_BOUNDARY_GENERATION_FILE),
+            b"generation-1\n",
+        )
+        .await
+        .unwrap();
+        tokio::fs::write(state_dir.join(HOST_AUTH_BUNDLE_FILE), b"auth")
+            .await
+            .unwrap();
+        tokio::fs::write(state_dir.join(HOST_RUNTIME_DESCRIPTOR_FILE), b"descriptor")
+            .await
+            .unwrap();
+        let overlay = sandbox_runtime_disk_paths(state_dir).overlay_disk;
+        tokio::fs::write(&overlay, b"invalid ext4 image")
+            .await
+            .unwrap();
+
+        assert!(remove_runtime_generation_material(state_dir).await.is_err());
+
+        for marker in [
+            HOST_BOUNDARY_GENERATION_FILE,
+            HOST_AUTH_BUNDLE_FILE,
+            HOST_RUNTIME_DESCRIPTOR_FILE,
+        ] {
+            assert!(
+                state_dir.join(marker).is_file(),
+                "lost {marker} before recovery"
+            );
+        }
+        assert_eq!(
+            tokio::fs::read(overlay).await.unwrap(),
+            b"invalid ext4 image"
+        );
+    }
+
     #[tokio::test]
     async fn explicit_start_clears_stop_and_terminal_markers() {
         let temp = tempfile::tempdir().unwrap();
diff --git a/e2e/rust/tests/vm_overlay.rs b/e2e/rust/tests/vm_overlay.rs
index 8d1797ef27..578f841f1f 100644
--- a/e2e/rust/tests/vm_overlay.rs
+++ b/e2e/rust/tests/vm_overlay.rs
@@ -53,5 +53,46 @@ async fn vm_overlay() {
         output.status.code(),
     );
 
+    for cycle in 0..3 {
+        let output = openshell_cmd()
+            .args(["sandbox", "exec", "--name", &sandbox.name, "--no-tty", "--"])
+            .args(["sh", "-c"])
+            .arg(format!(
+                "set -eu; mkdir -p /sandbox/journal-{cycle}; \
+                 for i in $(seq 1 128); do printf '%s\\n' \"$i\" > /sandbox/journal-{cycle}/$i; done; sync"
+            ))
+            .output()
+            .await
+            .expect("write overlay before stop");
+        assert!(output.status.success(), "overlay writes failed: {output:?}");
+
+        for action in ["stop", "start"] {
+            let output = openshell_cmd()
+                .args(["sandbox", action, &sandbox.name])
+                .output()
+                .await
+                .expect("run sandbox lifecycle command");
+            assert!(
+                output.status.success(),
+                "{action} failed on cycle {cycle}: {output:?}"
+            );
+        }
+
+        let output = openshell_cmd()
+            .args(["sandbox", "exec", "--name", &sandbox.name, "--no-tty", "--"])
+            .args(["sh", "-c"])
+            .arg(format!(
+                "set -eu; test \"$(cat /sandbox/overlay-check)\" = overlay-write; \
+                 for i in $(seq 1 128); do test \"$(cat /sandbox/journal-{cycle}/$i)\" = \"$i\"; done"
+            ))
+            .output()
+            .await
+            .expect("read overlay after restart");
+        assert!(
+            output.status.success(),
+            "overlay contents changed on cycle {cycle}: {output:?}"
+        );
+    }
+
     sandbox.cleanup().await;
 }