fix(resolve): preserve submodule worktrees during cleanup

Author: dkattan

Cargo.lock

@@ -21,7 +21,12 @@ use gitbutler_repo::{RepositoryExt as _, SignaturePurpose, signature};
 use gitbutler_stack::VirtualBranchesHandle;
 use gitbutler_workspace::{
     branch_trees::{WorkspaceState, update_uncommitted_changes_with_tree},
-    submodules::has_submodules_configured,
+    submodules::{
+        configured_submodule_paths,
+        has_submodules_configured,
+        is_submodule_related_path,
+        remove_untracked_excluding_submodule_paths,
+    },
 };
 use serde::Serialize;
 
@@ -152,62 +157,6 @@ fn get_uncommited_changes(ctx: &Context) -> Result<git2::Oid> {
     Ok(uncommited_changes)
 }
 
-fn sync_configured_submodules(repo: &git2::Repository) {
-    let Ok(mut submodules) = repo.submodules() else {
-        return;
-    };
-
-    for submodule in &mut submodules {
-        let _ = submodule.update(true, None);
-    }
-}
-
-fn configured_submodule_paths(repo: &git2::Repository) -> Vec<String> {
-    let mut paths = HashSet::new();
-
-    if let Ok(submodules) = repo.submodules() {
-        for submodule in submodules {
-            paths.insert(submodule.path().to_string_lossy().to_string());
-        }
-    }
-
-    // When only .git/modules entries exist, derive logical submodule paths from that layout.
-    let modules_root = repo.path().join("modules");
-    if modules_root.exists() {
-        let mut stack = vec![modules_root.clone()];
-        while let Some(dir) = stack.pop() {
-            let Ok(entries) = std::fs::read_dir(&dir) else {
-                continue;
-            };
-
-            for entry in entries.flatten() {
-                let path = entry.path();
-                if !path.is_dir() {
-                    continue;
-                }
-
-                if path.join("HEAD").is_file()
-                    && let Ok(relative) = path.strip_prefix(&modules_root)
-                {
-                    paths.insert(relative.to_string_lossy().to_string());
-                }
-
-                stack.push(path);
-            }
-        }
-    }
-
-    let mut output = paths.into_iter().collect::<Vec<_>>();
-    output.sort();
-    output
-}
-
-fn is_submodule_related_path(path: &str, submodule_paths: &[String]) -> bool {
-    submodule_paths
-        .iter()
-    .any(|sm| path == sm || path.strip_prefix(sm).is_some_and(|rest| rest.starts_with('/')))
-}
-
 fn collect_checkout_paths(
     repo: &git2::Repository,
     baseline: &git2::Tree,
@@ -262,15 +211,14 @@ fn checkout_edit_branch(ctx: &Context, commit: git2::Commit) -> Result<()> {
         repo.set_head(EDIT_BRANCH_REF)?;
         let mut checkout_head = CheckoutBuilder::new();
         checkout_head.force();
-        if !has_submodules {
-            checkout_head.remove_untracked(true);
-        } else {
+        if has_submodules {
             for path in &checkout_head_paths {
                 checkout_head.path(path);
             }
         }
         if !has_submodules || !checkout_head_paths.is_empty() {
             repo.checkout_head(Some(&mut checkout_head))?;
+            remove_untracked_excluding_submodule_paths(repo)?;
         }
 
         // Checkout the commit as unstaged changes
@@ -284,21 +232,15 @@ fn checkout_edit_branch(ctx: &Context, commit: git2::Commit) -> Result<()> {
         )?;
         let mut checkout_index = CheckoutBuilder::new();
         checkout_index.force().conflict_style_diff3(true);
-        if !has_submodules {
-            checkout_index.remove_untracked(true);
-        } else {
+        if has_submodules {
             for path in &checkout_index_paths {
                 checkout_index.path(path);
             }
         }
 
         if !has_submodules || !checkout_index_paths.is_empty() {
             repo.checkout_index(Some(&mut index), Some(&mut checkout_index))?;
-        }
-
-        // Keep configured submodule worktrees populated after moving into edit mode.
-        if has_submodules {
-            sync_configured_submodules(repo);
+            remove_untracked_excluding_submodule_paths(repo)?;
         }
 
         Ok(())
@@ -333,15 +275,25 @@ pub(crate) fn abort_and_return_to_workspace(
     force: bool,
     perm: &mut RepoExclusive,
 ) -> Result<()> {
-    if !force && !changes_from_initial(ctx, perm.read_permission())?.is_empty() {
+    let repo = &*ctx.git2_repo.get()?;
+    let changes = changes_from_initial(ctx, perm.read_permission())?;
+    if !force && !changes.is_empty() {
+        let submodule_paths = configured_submodule_paths(repo);
+        let has_submodule_or_gitlink_changes = changes.iter().any(|change| {
+            is_submodule_related_path(&change.path.to_str_lossy(), &submodule_paths)
+        });
+
+        if has_submodule_or_gitlink_changes {
+            bail!(
+                "The working tree differs from the original commit, including submodule or gitlink changes. A forced abort is necessary and will revert changes made during edit mode.\nIf you are seeing this message, please report it as a bug. The UI should have prevented this line getting hit."
+            );
+        }
+
         bail!(
             "The working tree differs from the original commit. A forced abort is necessary.\nIf you are seeing this message, please report it as a bug. The UI should have prevented this line getting hit."
         );
     }
 
-    let repo = &*ctx.git2_repo.get()?;
-    let has_submodules = has_submodules_configured(repo);
-
     // Checkout gitbutler workspace branch
     repo.set_head(WORKSPACE_BRANCH_REF)
         .context("Failed to set head reference")?;
@@ -351,10 +303,13 @@ pub(crate) fn abort_and_return_to_workspace(
 
     let mut checkout_tree = CheckoutBuilder::new();
     checkout_tree.force();
-    if !has_submodules {
-        checkout_tree.remove_untracked(true);
-    }
     repo.checkout_tree(uncommited_changes.as_object(), Some(&mut checkout_tree))?;
+    remove_untracked_excluding_submodule_paths(repo)?;
+
+    // Keep index in sync with HEAD after forceful checkout + cleanup.
+    let mut index = repo.index()?;
+    index.read_tree(&repo.head()?.peel_to_tree()?)?;
+    index.write()?;
 
     Ok(())
 }

crates/gitbutler-edit-mode/src/lib.rs

@@ -6,12 +6,27 @@ use gitbutler_edit_mode::commands::{
 };
 use gitbutler_operating_modes::{EDIT_BRANCH_REF, WORKSPACE_BRANCH_REF};
 use gitbutler_stack::VirtualBranchesHandle;
+use std::process::Command;
 use tempfile::TempDir;
 
 fn command_ctx(folder: &str) -> Result<(Context, TempDir)> {
     gitbutler_testsupport::writable::fixture("edit_mode.sh", folder)
 }
 
+fn run_git(cwd: &std::path::Path, args: &[&str]) -> Result<String> {
+    let output = Command::new("git").args(args).current_dir(cwd).output()?;
+    if !output.status.success() {
+        anyhow::bail!(
+            "git command failed in {}: git {}\nstderr: {}",
+            cwd.display(),
+            args.join(" "),
+            String::from_utf8_lossy(&output.stderr)
+        );
+    }
+
+    Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
+}
+
 // Fixture:
 // * xxx (HEAD -> gitbutler/workspace) GitButler Workspace Commit
 // * xxx foobar
@@ -86,6 +101,8 @@ fn abort_requires_force_when_changes_were_made() -> Result<()> {
     drop(repo);
 
     std::fs::write(worktree_dir.join("file"), "edited during edit mode\n")?;
+    let untracked_path = worktree_dir.join("new-untracked-during-edit-mode.txt");
+    std::fs::write(&untracked_path, "temporary file\n")?;
 
     let result = abort_and_return_to_workspace(&mut ctx, false);
     assert!(result.is_err());
@@ -101,6 +118,10 @@ fn abort_requires_force_when_changes_were_made() -> Result<()> {
         ctx.git2_repo.get()?.head()?.name(),
         Some(WORKSPACE_BRANCH_REF)
     );
+    assert!(
+        !untracked_path.exists(),
+        "forced abort should clean untracked files in non-submodule repos"
+    );
 
     Ok(())
 }
@@ -136,3 +157,110 @@ fn save_and_return_to_workspace_preserves_submodule_worktree() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn abort_preserves_preexisting_dirty_and_diverged_submodule_state() -> Result<()> {
+    let (mut ctx, _tempdir) = command_ctx("save_and_return_to_workspace_preserves_submodule_worktree")?;
+    let (foobar, worktree_dir) = {
+        let repo = ctx.git2_repo.get()?;
+        let foobar = repo.head()?.peel_to_commit()?.parent(0)?.id();
+        (foobar, repo.path().parent().unwrap().to_path_buf())
+    };
+    let submodule_dir = worktree_dir.join("submodules/test-module");
+
+    run_git(
+        &submodule_dir,
+        &["config", "user.name", "Submodule Author"],
+    )?;
+    run_git(
+        &submodule_dir,
+        &["config", "user.email", "submodule@example.com"],
+    )?;
+    std::fs::write(submodule_dir.join("diverged.txt"), "diverged commit\n")?;
+    run_git(&submodule_dir, &["add", "diverged.txt"])?;
+    run_git(&submodule_dir, &["commit", "-m", "local diverged commit"])?;
+
+    std::fs::write(submodule_dir.join("dirty.txt"), "dirty worktree change\n")?;
+
+    let baseline_submodule_head = run_git(&submodule_dir, &["rev-parse", "HEAD"])?;
+    let baseline_submodule_status = run_git(&submodule_dir, &["status", "--porcelain"])?;
+    let baseline_superproject_submodule_status =
+        run_git(&worktree_dir, &["status", "--porcelain", "submodules/test-module"])?;
+
+    let vb_state = VirtualBranchesHandle::new(ctx.project_data_dir());
+    let stacks = vb_state.list_stacks_in_workspace()?;
+    let stack = stacks.first().unwrap();
+
+    enter_edit_mode(&mut ctx, foobar, stack.id)?;
+    abort_and_return_to_workspace(&mut ctx, true)?;
+
+    let final_submodule_head = run_git(&submodule_dir, &["rev-parse", "HEAD"])?;
+    let final_submodule_status = run_git(&submodule_dir, &["status", "--porcelain"])?;
+    let final_superproject_submodule_status =
+        run_git(&worktree_dir, &["status", "--porcelain", "submodules/test-module"])?;
+
+    assert_eq!(
+        final_submodule_head, baseline_submodule_head,
+        "abort should preserve pre-existing submodule commit divergence"
+    );
+    assert_eq!(
+        final_submodule_status, baseline_submodule_status,
+        "abort should preserve pre-existing dirty submodule working tree state"
+    );
+    assert_eq!(
+        final_superproject_submodule_status, baseline_superproject_submodule_status,
+        "abort should restore the same superproject-visible submodule state"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn abort_requires_force_warns_about_submodule_and_gitlink_reversion() -> Result<()> {
+    let (mut ctx, _tempdir) = command_ctx("save_and_return_to_workspace_preserves_submodule_worktree")?;
+    let (foobar, worktree_dir) = {
+        let repo = ctx.git2_repo.get()?;
+        let foobar = repo.head()?.peel_to_commit()?.parent(0)?.id();
+        (foobar, repo.path().parent().unwrap().to_path_buf())
+    };
+    let submodule_dir = worktree_dir.join("submodules/test-module");
+
+    let vb_state = VirtualBranchesHandle::new(ctx.project_data_dir());
+    let stacks = vb_state.list_stacks_in_workspace()?;
+    let stack = stacks.first().unwrap();
+
+    enter_edit_mode(&mut ctx, foobar, stack.id)?;
+
+    run_git(
+        &submodule_dir,
+        &["config", "user.name", "Submodule Author"],
+    )?;
+    run_git(
+        &submodule_dir,
+        &["config", "user.email", "submodule@example.com"],
+    )?;
+    std::fs::write(
+        submodule_dir.join("during-edit-mode-gitlink-change.txt"),
+        "changed during edit mode\n",
+    )?;
+    run_git(&submodule_dir, &["add", "during-edit-mode-gitlink-change.txt"])?;
+    run_git(&submodule_dir, &["commit", "-m", "gitlink change during edit mode"])?;
+
+    // Stage the submodule entry so superproject tree diff includes the gitlink change.
+    run_git(&worktree_dir, &["add", "submodules/test-module"])?;
+
+    let result = abort_and_return_to_workspace(&mut ctx, false);
+    assert!(result.is_err());
+    let err = result.err().unwrap().to_string();
+    assert!(
+        err.contains("including submodule or gitlink changes"),
+        "expected submodule/gitlink warning in abort error, got: {err}"
+    );
+    assert!(
+        err.contains("will revert changes made during edit mode"),
+        "expected explicit reversion warning in abort error, got: {err}"
+    );
+
+    Ok(())
+}
+

crates/gitbutler-edit-mode/tests/edit_mode.rs

@@ -23,5 +23,8 @@ anyhow.workspace = true
 gix.workspace = true
 git2.workspace = true
 
+[dev-dependencies]
+tempfile.workspace = true
+
 [lints]
 workspace = true

crates/gitbutler-workspace/Cargo.toml

@@ -8,7 +8,11 @@ use gitbutler_cherry_pick::RepositoryExt;
 use gitbutler_repo::RepositoryExt as _;
 use gitbutler_stack::VirtualBranchesHandle;
 
-use crate::{submodules::has_submodules_configured, workspace_base, workspace_base_from_heads};
+use crate::{
+    submodules::remove_untracked_excluding_submodule_paths,
+    workspace_base,
+    workspace_base_from_heads,
+};
 
 /// A snapshot of the workspace at a point in time.
 #[derive(Debug)]
@@ -97,7 +101,6 @@ pub fn update_uncommitted_changes_with_tree(
     _perm: &mut RepoExclusive,
 ) -> Result<()> {
     let repo = &*ctx.git2_repo.get()?;
-    let has_submodules = has_submodules_configured(repo);
 
     if let Some(worktree_id) = old_uncommitted_changes {
         let mut new_uncommitted_changes =
@@ -113,13 +116,9 @@ pub fn update_uncommitted_changes_with_tree(
 
         let mut checkout = git2::build::CheckoutBuilder::new();
         checkout.force().conflict_style_diff3(true);
-        // Submodule worktrees can be deleted by aggressive untracked cleanup.
-        // Keep that cleanup only for repos without configured submodules.
-        if !has_submodules {
-            checkout.remove_untracked(true);
-        }
 
         repo.checkout_index(Some(&mut new_uncommitted_changes), Some(&mut checkout))?;
+        remove_untracked_excluding_submodule_paths(repo)?;
     } else {
         let old_tree_id = merge_workspace(repo, old)?.to_gix();
         let new_tree_id = merge_workspace(repo, new)?.to_gix();