fix(staging): never drop the staging set aside for a rewrite

CLAUDE.md

@@ -81,7 +81,10 @@ because nothing was autostashed — unless `reset_mixed_to` or `reset_hard_to` i
 set, which marks a caller that moved HEAD before the rebase existed (`commit`,
 `absorb`) and so has its own work to take back. A caller that unstages before
 its rebase without moving HEAD keeps a `staging::StagedAside` guard instead
-(Spec 014).
+(Spec 014). `StagedAside::handed_over()` may only name an owner that is durable
+or has already run: called before a `rebase_abort_then_cleanup` closure it
+drops the patch on exactly the path that skips that closure, so it goes inside
+one.
 
 Every rebase autostashes, and that replay reaches the working tree only: a
 staged modification comes back unstaged on a rebase that completed just as it

specs/006-commit.md

@@ -72,6 +72,10 @@ Create the commit at HEAD, then relocate it to the target via one Weave operatio
 
 If relocation conflicts, pause and save `.git/loom/state.json`. The new commit remains recoverable in the working tree through mixed reset; original staged changes MUST be restored on abort.
 
+Staged files set aside so they cannot join this commit MUST come back on every
+failure path, including the ones after the commit is created but before
+`state.json` exists, where no rollback can return them.
+
 - `loom continue` resolves/completes relocation.
 - `loom abort` cancels and restores original history and staged state.
 

specs/007-fold.md

@@ -216,6 +216,10 @@ These operations pause on rebase conflict and save the listed `LoomState.context
 
 `loom continue` dispatches `after_continue`, removes `_loom-track`, restores pre-existing staged changes from `LoomState.rollback`, and prints the operation's success message. `loom abort` restores original history, staged state, and working-tree state (Spec 014). The post-continue unapplied-patch exception is defined above.
 
+Staged files set aside so they cannot join the fold MUST come back on every
+failure path, including the ones reached before `state.json` exists, where no
+rollback can return them.
+
 Multiple moves (to a branch or next to a commit), all `-c` moves, all `-p` forms, and CommitFile move failures save no resumable state and auto-rollback as specified in their sections.
 
 ## General invariants and prerequisites

src/commit.rs

@@ -9,7 +9,7 @@ use crate::core::changeid;
 use crate::core::graph;
 use crate::core::msg;
 use crate::core::repo;
-use crate::core::staging;
+use crate::core::staging::{self, StagedAside};
 use crate::core::transaction::{self, LoomState, Rollback};
 use crate::core::weave::{self, RebaseOutcome, Weave};
 use crate::git;
@@ -94,17 +94,13 @@ pub fn run(
 
     // Stage files, saving aside any pre-existing staged files not in the
     // target list so they don't accidentally end up in this commit.
-    let saved_staged = if patch {
+    let staged_aside = if patch {
         resolve_staging_patch(&repo, &workdir, &files, theme)?
     } else {
         resolve_staging(&repo, &workdir, &files)?
     };
 
-    // Restore the saved staged work if the index turns out to be empty.
-    if let Err(e) = repo::verify_has_staged_changes(&repo) {
-        git::restore_staged_patch(&workdir, &saved_staged);
-        return Err(e);
-    }
+    repo::verify_has_staged_changes(&repo)?;
 
     let git_opts: Vec<&str> = git_args.iter().map(String::as_str).collect();
     let do_commit = || -> Result<()> {
@@ -124,7 +120,9 @@ pub fn run(
     // "origin/main").
     if loose {
         let result = do_commit();
-        git::restore_staged_patch(&workdir, &saved_staged);
+        // Put back either way: a loose commit writes no state file, so there
+        // is no later owner for the patch and no rollback to hold it.
+        staged_aside.restore();
         result?;
         let new_head = repo::head_oid(&repo)?;
         msg::success(&format!(
@@ -139,25 +137,15 @@ pub fn run(
     // Resolve branch target (may create a new branch at merge-base).
     // Returns whether the branch was newly created — only newly-created
     // branches are deleted on rollback (not pre-existing empty ones).
-    // A cancelled prompt here must put the saved-aside files back in the index.
     let (branch_name, branch_is_new) =
-        match resolve_branch_target(&repo, &info, &workdir, branch.as_deref()) {
-            Ok(resolved) => resolved,
-            Err(e) => {
-                git::restore_staged_patch(&workdir, &saved_staged);
-                return Err(e);
-            }
-        };
+        resolve_branch_target(&repo, &info, &workdir, branch.as_deref())?;
 
     // Empty branches (pointing at merge-base) need a branch section and
     // merge entry created in the Weave before moving the commit there.
     let branch_is_empty =
         is_branch_at_merge_base(&repo, &branch_name, info.upstream.merge_base_oid)?;
 
-    if let Err(e) = do_commit() {
-        git::restore_staged_patch(&workdir, &saved_staged);
-        return Err(e);
-    }
+    do_commit()?;
 
     let head_oid = repo::head_oid(&repo)?;
 
@@ -184,7 +172,7 @@ pub fn run(
     }
     let ctx = CommitContext {
         branch_name: branch_name.clone(),
-        saved_staged: Some(saved_staged.clone()),
+        saved_staged: Some(staged_aside.patch().to_string()),
     };
     let state = LoomState {
         command: "commit".to_string(),
@@ -200,6 +188,9 @@ pub fn run(
         protect: vec![head_oid.to_string()],
     };
     transaction::save(&git_dir, &state)?;
+    // The state file owns the patch from here: `post_commit` puts it back on
+    // success, `Rollback` on abort.
+    let saved_staged = staged_aside.release();
 
     let base = graph.base_oid.to_string();
 
@@ -256,49 +247,46 @@ fn post_commit(workdir: &Path, branch_name: &str, saved_staged: &str) -> Result<
 ///
 /// With specific files, other staged files are saved aside and unstaged first,
 /// so they neither show in the picker nor leak into this commit. Returns that
-/// saved patch for restoration after the commit; a cancelled picker restores
-/// it and errors.
-fn resolve_staging_patch(
+/// saved patch for restoration after the commit; a picker that is cancelled or
+/// fails puts it back and errors.
+fn resolve_staging_patch<'a>(
     repo: &Repository,
-    workdir: &std::path::Path,
+    workdir: &'a Path,
     files: &[String],
     theme: &graph::Theme,
-) -> Result<String> {
+) -> Result<StagedAside<'a>> {
     // Save aside other staged files when specific files are targeted.
     let filter = staging::filter_paths(repo, files)?;
     let saved_staged = match &filter {
         Some(paths) => {
             let path_refs: Vec<&str> = paths.iter().map(|s| s.as_str()).collect();
             staging::save_and_unstage_other_staged(repo, workdir, &path_refs)?
         }
-        None => String::new(),
+        None => StagedAside::none(workdir),
     };
 
-    let confirmed = staging::run_hunk_picker(repo, workdir, filter.as_deref(), theme)?;
-    if confirmed.is_none() {
-        git::restore_staged_patch(workdir, &saved_staged);
+    if staging::run_hunk_picker(repo, workdir, filter.as_deref(), theme)?.is_none() {
         return Err(msg::cancelled());
     }
-
     Ok(saved_staged)
 }
 
 /// Resolve staging from the file arguments: an empty list uses the index
 /// as-is, `zz` stages everything, and named files are staged after any other
 /// pre-existing staged file is saved aside and unstaged so it cannot leak into
 /// this commit. Returns that saved patch for later restoration.
-fn resolve_staging(
+fn resolve_staging<'a>(
     repo: &Repository,
-    workdir: &std::path::Path,
+    workdir: &'a Path,
     files: &[String],
-) -> Result<String> {
+) -> Result<StagedAside<'a>> {
     if files.is_empty() {
-        return Ok(String::new());
+        return Ok(StagedAside::none(workdir));
     }
 
     if files.iter().any(|f| f == "zz") {
         git::stage_all(workdir)?;
-        return Ok(String::new());
+        return Ok(StagedAside::none(workdir));
     }
 
     let resolved_paths = resolve_file_args(repo, files)?;

src/commit_test.rs

@@ -177,6 +177,66 @@ fn commit_to_non_woven_branch_fails() {
     assert!(result.unwrap_err().to_string().contains("not woven"));
 }
 
+/// The worst window: the commit is made, but the state file that would carry
+/// the set-aside work is not written yet, so nothing else can return it —
+/// `loom abort` has no rollback to read. A file where `.git/loom` must be a
+/// directory is what makes `transaction::save` fail here — which also blocks
+/// the last-resort patch dump the restore falls back on, so the one asserted
+/// below is the real one, not the fallback.
+#[test]
+fn a_commit_failing_after_it_is_created_puts_back_the_staging() {
+    let test_repo = setup_with_woven_branch();
+    test_repo.write_file("kept.txt", "staged before the commit");
+    test_repo.stage_files(&["kept.txt"]);
+    test_repo.write_file("file.txt", "content");
+    std::fs::write(test_repo.repo.path().join("loom"), "not a directory").unwrap();
+    let before = test_repo.head_oid();
+
+    let result = test_repo.in_dir(|| {
+        run(
+            Some("feature-a".to_string()),
+            Some("Message".to_string()),
+            vec!["file.txt".to_string()],
+        )
+    });
+
+    assert!(result.is_err(), "the state file cannot be written");
+    // Without this the test still passes if the failure ever moves earlier,
+    // leaving the window it is named after unguarded.
+    assert_ne!(
+        test_repo.head_oid(),
+        before,
+        "the commit must already exist"
+    );
+    let status = test_repo.status_porcelain();
+    assert!(status.contains("A  kept.txt"), "{status}");
+}
+
+/// A failure between staging and the commit puts the set-aside work back too:
+/// which files were staged is state git keeps no second copy of.
+#[test]
+fn a_failed_commit_puts_back_the_staging_it_set_aside() {
+    let test_repo = TestRepo::new_with_remote();
+    test_repo.commit("A1", "a1.txt");
+    test_repo.create_branch_tracking("not-woven", "origin/main");
+
+    test_repo.write_file("kept.txt", "staged before the commit");
+    test_repo.stage_files(&["kept.txt"]);
+    test_repo.write_file("file.txt", "content");
+
+    let result = test_repo.in_dir(|| {
+        run(
+            Some("not-woven".to_string()),
+            Some("Message".to_string()),
+            vec!["file.txt".to_string()],
+        )
+    });
+
+    assert!(result.is_err());
+    let status = test_repo.status_porcelain();
+    assert!(status.contains("A  kept.txt"), "{status}");
+}
+
 #[test]
 fn commit_to_new_branch_creates_and_weaves() {
     let test_repo = TestRepo::new_with_remote();

src/core/staging.rs

@@ -1,5 +1,6 @@
 use anyhow::{Result, bail};
 use git2::Repository;
+use std::cell::Cell;
 use std::path::Path;
 
 use crate::core::diff::{self, parse_hunk_start};
@@ -499,43 +500,106 @@ pub(crate) fn collect_commit_hunks(
     Ok(entries)
 }
 
-/// Save and unstage all currently staged changes, returning a patch to restore them later.
+/// Staged work an operation set aside, put back into the index when this
+/// drops — every `?` in between included, which is what keeps it (Spec 006).
 ///
-/// Returns an empty string if nothing is staged. Callers must call
-/// `git::restore_staged_patch` with the returned patch when the operation
-/// completes (or is rolled back), so pre-existing staged work is never lost.
-pub(crate) fn save_and_unstage_staged(repo: &Repository, workdir: &Path) -> Result<String> {
+/// Hand it on where someone else owns the restore: a state file, or a worktree
+/// snapshot taken before the unstaging. That owner has to be one this guard
+/// cannot outlive uncollected — see [`StagedAside::handed_over`].
+#[must_use = "dropping the guard right away puts the work straight back, undoing the unstage"]
+pub(crate) struct StagedAside<'a> {
+    workdir: &'a Path,
+    patch: String,
+    armed: Cell<bool>,
+}
+
+impl<'a> StagedAside<'a> {
+    fn new(workdir: &'a Path, patch: String) -> Self {
+        StagedAside {
+            workdir,
+            patch,
+            armed: Cell::new(true),
+        }
+    }
+
+    /// A guard over nothing: the operation set no staged work aside.
+    pub(crate) fn none(workdir: &'a Path) -> Self {
+        StagedAside::new(workdir, String::new())
+    }
+
+    pub(crate) fn patch(&self) -> &str {
+        &self.patch
+    }
+
+    /// Put it back now, rather than wherever this would have dropped. Does
+    /// nothing once [`StagedAside::handed_over`] has run.
+    pub(crate) fn restore(self) {}
+
+    /// Take the patch back: the caller restores it from here.
+    #[must_use = "the set-aside patch is lost unless a new owner keeps it"]
+    pub(crate) fn release(mut self) -> String {
+        self.armed.set(false);
+        std::mem::take(&mut self.patch)
+    }
+
+    /// Something else holds the patch now, so this must not put it back a
+    /// second time.
+    ///
+    /// Data safety: that owner must be durable (a state file `loom abort`
+    /// reads) or already have run. A rollback closure is neither until it
+    /// runs, and [`git::rebase_abort_then_cleanup`] skips its closure when the
+    /// abort fails — so call this from inside such a closure, never before it.
+    /// Disarming early there drops the patch on the one path that cannot get
+    /// it back; leaving the guard armed parks it instead.
+    pub(crate) fn handed_over(&self) {
+        self.armed.set(false);
+    }
+}
+
+impl Drop for StagedAside<'_> {
+    fn drop(&mut self) {
+        if self.armed.get() {
+            git::restore_loom_unstaged(self.workdir, &self.patch);
+        }
+    }
+}
+
+/// Save and unstage all currently staged changes, so an amend or rebase below
+/// leaves them out. Restored when the returned guard drops.
+pub(crate) fn save_and_unstage_staged<'a>(
+    repo: &Repository,
+    workdir: &'a Path,
+) -> Result<StagedAside<'a>> {
     let staged = repo::get_staged_files(repo)?;
     if staged.is_empty() {
-        return Ok(String::new());
+        return Ok(StagedAside::new(workdir, String::new()));
     }
     let refs: Vec<&str> = staged.iter().map(|s| s.as_str()).collect();
     let patch = git::diff_cached_files(workdir, &refs)?;
     git::unstage_files(workdir, &refs)?;
-    Ok(patch)
+    Ok(StagedAside::new(workdir, patch))
 }
 
 /// Save the staged diff for files that are staged but NOT in `target_files`,
-/// then unstage them so they don't leak into the upcoming commit.
-///
-/// Returns the patch as a string (may be empty if nothing to save).
-pub(crate) fn save_and_unstage_other_staged(
+/// then unstage them so they don't leak into the upcoming commit. Restored
+/// when the returned guard drops.
+pub(crate) fn save_and_unstage_other_staged<'a>(
     repo: &Repository,
-    workdir: &Path,
+    workdir: &'a Path,
     target_files: &[&str],
-) -> Result<String> {
+) -> Result<StagedAside<'a>> {
     let staged = repo::get_staged_files(repo)?;
     let other: Vec<&str> = staged
         .iter()
         .filter(|f| !target_files.contains(&f.as_str()))
         .map(|s| s.as_str())
         .collect();
     if other.is_empty() {
-        return Ok(String::new());
+        return Ok(StagedAside::new(workdir, String::new()));
     }
     let patch = git::diff_cached_files(workdir, &other)?;
     git::unstage_files(workdir, &other)?;
-    Ok(patch)
+    Ok(StagedAside::new(workdir, patch))
 }
 
 /// The single entry a submodule contributes to a picker: one object id, with no