From e4181c36ea657ce64fc80f93cb79cef143eef0cb Mon Sep 17 00:00:00 2001 From: rohitkumardanda Date: Wed, 11 Nov 2026 16:42:35 +0000 Subject: [PATCH] fix(agents): roll a local task back to where it started when HEAD moved (#2774) A task working in the project's own checkout compared and rolled back its changes against the current HEAD. Once HEAD moved back or to another branch during the task (git reset --soft, an undone commit, a checkout), the user's own work showed up as the task's changes and Rollback overwrote it; in a repository without commits, Changes and Rollback failed. Local tasks now use the commit they started from: files that are as they were at the start are not listed, and Rollback and the diff base use the start commit, or HEAD when the user committed on top of it so those commits are kept. In a repository without commits new files are listed and rollback deletes them. Worktree tasks are unchanged. Co-authored-by: rohitdanda <10144914+rohitdanda@users.noreply.github.com> --- crates/ide-agent/src/lib.rs | 35 ++++--- crates/ide-agent/src/tests.rs | 152 +++++++++++++++++++++++++++++++ crates/ide-agent/src/worktree.rs | 132 +++++++++++++++++++++++++-- 3 files changed, 295 insertions(+), 24 deletions(-) diff --git a/crates/ide-agent/src/lib.rs b/crates/ide-agent/src/lib.rs index 007eedf6a..fd59668a8 100644 --- a/crates/ide-agent/src/lib.rs +++ b/crates/ide-agent/src/lib.rs @@ -1338,15 +1338,19 @@ impl AgentManager { /// Files the task changed since it started. pub async fn changes(&self, id: &str) -> IdeResult> { let task = self.task(id)?; - let (cwd, base) = task.cwd_and_base(); + let at = task.base(); let before = task.state.lock().before.clone(); - blocking(move || Ok(worktree::without_untouched(worktree::changes(&cwd, &base)?, &before))).await + blocking(move || at.changes(&before)).await } /// Content of `path` (in the task's directory) when the task started. pub async fn base_text(&self, id: &str, path: PathBuf) -> IdeResult> { - let (cwd, base) = self.task(id)?.cwd_and_base(); - blocking(move || worktree::file_at(&cwd, &path, &base)).await + let at = self.task(id)?.base(); + blocking(move || match at.rev()? { + Some(rev) => worktree::file_at(at.cwd(), &path, &rev), + None => Ok(None), + }) + .await } /// Rollback: puts the task's changed files among `rel_paths` (all with `None`) back as they @@ -1356,16 +1360,17 @@ impl AgentManager { if task.state.lock().busy { return Err(IdeError::invalid("the agent is working: stop it before rolling back its changes")); } - let (cwd, base) = task.cwd_and_base(); + let at = task.base(); let (before, saved) = { let s = task.state.lock(); (s.before.clone(), s.saved.as_ref().map(|name| self.shared.paths.store_dir.join(name))) }; let restamped = blocking(move || { - let changes = worktree::without_untouched(worktree::changes(&cwd, &base)?, &before); + let changes = at.changes(&before)?; + let rev = at.rev()?; let mut restamped = Vec::new(); for change in changes.iter().filter(|c| rel_paths.as_ref().is_none_or(|p| p.contains(&c.rel_path))) { - worktree::restore(&cwd, &base, change, &before, saved.as_deref())?; + worktree::restore(at.cwd(), rev.as_deref(), change, &before, saved.as_deref())?; // Back to the user's own edits: not a change of the task any more. if before.contains_key(&change.rel_path) { restamped.push((change.rel_path.clone(), worktree::stamp(&change.path.0))); @@ -1469,15 +1474,15 @@ impl Task { }) } - fn cwd_and_base(&self) -> (PathBuf, String) { + fn base(&self) -> worktree::TaskBase { let s = self.state.lock(); - // In the project's own checkout the user commits too: compare against the current HEAD, - // as the Commit tool window does, so commits made since the task started drop out. - let base = match s.task.isolation { - Isolation::Local => None, - Isolation::Worktree => s.task.base_commit.clone(), - }; - (s.task.cwd.0.clone(), base.unwrap_or_else(|| "HEAD".into())) + let cwd = s.task.cwd.0.clone(); + match s.task.isolation { + Isolation::Local => worktree::TaskBase::Local { cwd, start: s.task.base_commit.clone() }, + Isolation::Worktree => { + worktree::TaskBase::Worktree { cwd, base: s.task.base_commit.clone().unwrap_or_else(|| "HEAD".into()) } + } + } } fn emit_task(&self) { diff --git a/crates/ide-agent/src/tests.rs b/crates/ide-agent/src/tests.rs index 4404097a9..cee413e5c 100644 --- a/crates/ide-agent/src/tests.rs +++ b/crates/ide-agent/src/tests.rs @@ -2869,6 +2869,158 @@ async fn rolling_back_restores_the_start_of_the_task_and_keeps_the_users_own_edi assert!(!dir.join("store").join(format!("{}.before", task.id)).exists()); } +/// A git repository in a temporary folder, for local tasks whose checkout moves under them. +struct Checkout { + _tmp: tempfile::TempDir, + dir: PathBuf, +} + +impl Checkout { + fn new() -> Self { + let tmp = tempfile::tempdir().unwrap(); + let dir = dunce::canonicalize(tmp.path()).unwrap(); + let checkout = Self { _tmp: tmp, dir }; + checkout.git(&["init", "-q", "-b", "main"]); + checkout.git(&["config", "core.autocrlf", "false"]); + checkout.write(".gitignore", "store/\nworktrees/\nagents/\nmodels.json\n"); + checkout + } + + fn git(&self, args: &[&str]) { + let ok = std::process::Command::new("git") + .args(["-c", "user.name=t", "-c", "user.email=t@t", "-c", "commit.gpgsign=false"]) + .args(args) + .current_dir(&self.dir) + .status() + .unwrap() + .success(); + assert!(ok, "{args:?}"); + } + + fn commit(&self, message: &str) { + self.git(&["add", "."]); + self.git(&["commit", "-q", "-m", message]); + } + + fn write(&self, rel: &str, text: &str) { + std::fs::write(self.dir.join(rel), text).unwrap(); + } + + fn read(&self, rel: &str) -> String { + std::fs::read_to_string(self.dir.join(rel)).unwrap() + } + + /// A local task, stopped so it can be rolled back. + async fn task(&self) -> (AgentManager, AgentTask) { + let mgr = manager(&self.dir, Arc::new(FakeAgent::default())); + let task = mgr + .create(NewAgentTask { + project_root: None, + prompt: text("Hello"), + provider: "claude".into(), + config: Vec::new(), + isolation: Isolation::Local, + base: None, + }) + .await + .unwrap(); + wait_for(&mgr, &task.id, AgentStatus::NeedsInput).await; + mgr.cancel(&task.id, None).unwrap(); + wait_for(&mgr, &task.id, AgentStatus::Ready).await; + (mgr, task) + } +} + +async fn changed(mgr: &AgentManager, id: &str) -> Vec { + mgr.changes(id).await.unwrap().into_iter().map(|c| c.rel_path).collect() +} + +#[tokio::test(flavor = "multi_thread")] +async fn work_uncommitted_by_a_reset_during_a_local_task_is_not_rolled_back() { + let repo = Checkout::new(); + repo.write("a.txt", "base\n"); + repo.write("mine.txt", "base\n"); + repo.commit("init"); + repo.write("mine.txt", "my work\n"); + repo.commit("my work"); + let (mgr, task) = repo.task().await; + repo.write("a.txt", "agent\n"); + // The user takes the last commit back into the working tree (`git reset --soft`, an + // undone commit): its changes are the user's, not the agent's. + repo.git(&["reset", "-q", "--soft", "HEAD~1"]); + assert_eq!(changed(&mgr, &task.id).await, ["a.txt"]); + + let base = mgr.base_text(&task.id, repo.dir.join("a.txt")).await.unwrap(); + assert_eq!(base.as_deref(), Some("base\n")); + mgr.rollback_changes(&task.id, None).await.unwrap(); + assert_eq!(repo.read("a.txt"), "base\n"); + assert_eq!(repo.read("mine.txt"), "my work\n", "the user's work survives the rollback"); + assert!(changed(&mgr, &task.id).await.is_empty()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn a_branch_switched_to_during_a_local_task_is_left_alone() { + let repo = Checkout::new(); + repo.write("a.txt", "base\n"); + repo.write("other.txt", "base\n"); + repo.commit("init"); + repo.git(&["checkout", "-q", "-b", "feature"]); + repo.write("other.txt", "feature\n"); + repo.write("only-feature.txt", "feature\n"); + repo.commit("feature"); + repo.git(&["checkout", "-q", "main"]); + repo.write("only-main.txt", "main\n"); + repo.commit("main"); + let (mgr, task) = repo.task().await; + repo.write("a.txt", "agent\n"); + // The agent's uncommitted edit comes along to a branch that doesn't contain the task's start. + repo.git(&["checkout", "-q", "feature"]); + assert_eq!(changed(&mgr, &task.id).await, ["a.txt"]); + + mgr.rollback_changes(&task.id, None).await.unwrap(); + assert_eq!(repo.read("a.txt"), "base\n"); + assert_eq!(repo.read("other.txt"), "feature\n"); + assert_eq!(repo.read("only-feature.txt"), "feature\n"); + assert!(!repo.dir.join("only-main.txt").exists()); + assert!(changed(&mgr, &task.id).await.is_empty()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn commits_made_during_a_local_task_are_kept_by_its_rollback() { + let repo = Checkout::new(); + repo.write("a.txt", "base\n"); + repo.commit("init"); + let (mgr, task) = repo.task().await; + repo.write("a.txt", "agent\n"); + repo.write("b.txt", "agent\n"); + // The user commits what the agent did to `a.txt` so far: it drops out of the task's changes. + repo.git(&["add", "a.txt"]); + repo.git(&["commit", "-q", "-m", "keep"]); + assert_eq!(changed(&mgr, &task.id).await, ["b.txt"]); + repo.write("a.txt", "agent\nagain\n"); + assert_eq!(changed(&mgr, &task.id).await, ["a.txt", "b.txt"]); + + mgr.rollback_changes(&task.id, None).await.unwrap(); + assert_eq!(repo.read("a.txt"), "agent\n", "back to the commit, not to before it"); + assert!(!repo.dir.join("b.txt").exists()); + assert!(changed(&mgr, &task.id).await.is_empty()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn a_local_task_in_a_repository_without_commits_lists_and_rolls_back_its_changes() { + let repo = Checkout::new(); + repo.write("mine.txt", "mine\n"); + let (mgr, task) = repo.task().await; + repo.write("new.txt", "agent\n"); + assert_eq!(changed(&mgr, &task.id).await, ["new.txt"]); + assert_eq!(mgr.base_text(&task.id, repo.dir.join("new.txt")).await.unwrap(), None); + + mgr.rollback_changes(&task.id, None).await.unwrap(); + assert!(!repo.dir.join("new.txt").exists()); + assert_eq!(repo.read("mine.txt"), "mine\n"); + assert!(changed(&mgr, &task.id).await.is_empty()); +} + #[tokio::test(flavor = "multi_thread")] async fn security_policy_auto_approves_whitelist_and_rejects_blacklist() { use ide_core::settings::AgentSecuritySettings; diff --git a/crates/ide-agent/src/worktree.rs b/crates/ide-agent/src/worktree.rs index c72fcefad..0e127e132 100644 --- a/crates/ide-agent/src/worktree.rs +++ b/crates/ide-agent/src/worktree.rs @@ -117,14 +117,92 @@ pub(crate) fn changes(cwd: &Path, base: &str) -> IdeResult> { let tracked: std::collections::HashSet = out.iter().map(|c| c.rel_path.clone()).collect(); for s in repo.status()? { if s.unstaged == Some(ChangeKind::Untracked) && !tracked.contains(&s.rel_path) { - let added = std::fs::read_to_string(&*s.path).map_or(0, |t| t.lines().count() as u32); - out.push(AgentChange { path: s.path, rel_path: s.rel_path, kind: ChangeKind::Added, added, removed: 0 }); + out.push(added(s.path, s.rel_path)); } } out.sort_by(|a, b| a.rel_path.cmp(&b.rel_path)); Ok(out) } +fn added(path: IpcPath, rel_path: String) -> AgentChange { + let added = std::fs::read_to_string(&*path).map_or(0, |t| t.lines().count() as u32); + AgentChange { path, rel_path, kind: ChangeKind::Added, added, removed: 0 } +} + +/// Where a task works and what its changes are measured against. +pub(crate) enum TaskBase { + /// A worktree of its own: everything since the commit it was created from, the agent's + /// own commits included. + Worktree { cwd: PathBuf, base: String }, + /// The project's checkout, where the user works and commits too. `start` is the commit it + /// started from (`None` in a repository without commits, or an imported session). + Local { cwd: PathBuf, start: Option }, +} + +impl TaskBase { + pub(crate) fn cwd(&self) -> &Path { + match self { + Self::Worktree { cwd, .. } | Self::Local { cwd, .. } => cwd, + } + } + + /// The files the task changed, less those changed before it started and untouched since. + pub(crate) fn changes(&self, before: &BTreeMap) -> IdeResult> { + let (cwd, start) = match self { + Self::Worktree { cwd, base } => return Ok(without_untouched(changes(cwd, base)?, before)), + Self::Local { cwd, start } => (cwd, start.as_deref()), + }; + let Some(repo) = discover(cwd)? else { return Ok(Vec::new()) }; + // What is uncommitted, as the Commit tool window lists it: commits made since the task + // started drop out. + let head = repo.branch_info()?.head_commit; + let listed = match &head { + Some(head) => changes(cwd, head)?, + None => { + let mut all: Vec = repo + .status()? + .into_iter() + .filter(|s| s.path.is_file()) + .map(|s| added(s.path, s.rel_path)) + .collect(); + all.sort_by(|a, b| a.rel_path.cmp(&b.rel_path)); + all + } + }; + let listed = without_untouched(listed, before); + let Some(start) = start.filter(|s| head.as_deref() != Some(*s)) else { return Ok(listed) }; + // HEAD moved (a reset, a checkout, a pull): a file that is as it was when the task + // started is the user's doing, not the task's. + Ok(listed + .into_iter() + .filter(|c| { + before.contains_key(&c.rel_path) + || !repo.file_bytes_at(&c.path.0, Some(start)).is_ok_and(|at| at == std::fs::read(&c.path.0).ok()) + }) + .collect()) + } + + /// The revision the task's changed files go back to on rollback, and are compared with: + /// the commit the task started from, or the current `HEAD` once the user committed on top + /// of it, so those commits stay. `None` while the repository has no commits. + pub(crate) fn rev(&self) -> IdeResult> { + let (cwd, start) = match self { + Self::Worktree { base, .. } => return Ok(Some(base.clone())), + Self::Local { cwd, start } => (cwd, start.as_deref()), + }; + let Some(repo) = discover(cwd)? else { return Ok(None) }; + let head = repo.branch_info()?.head_commit; + // A start commit git no longer knows (garbage-collected) can't be gone back to. + let start = start.filter(|s| repo.merge_base(s, s).ok().flatten().is_some()); + Ok(match (start, head) { + (Some(start), Some(head)) if start != head && repo.merge_base(start, &head)?.as_deref() != Some(start) => { + Some(start.to_owned()) + } + (start, head) => head.or_else(|| start.map(str::to_owned)), + }) + } +} + /// Copies the task's `changes` (against `base`) from the worktree `cwd` into the /// local checkout `root`, uncommitted: the user reviews and commits them with the /// VCS tools. Files the checkout changed itself are left alone and reported, so @@ -220,11 +298,11 @@ pub(crate) fn save_before(cwd: &Path, before: &BTreeMap, dir: &Pa } /// Rolls `change` back: to the copy saved when the task started when the file had changes -/// then (`before`, `saved` holding the copies), otherwise to `base` (removed when `base` -/// lacks it). +/// then (`before`, `saved` holding the copies), otherwise to `rev` (removed when `rev` lacks +/// it, or when there is no `rev`: the repository had no commits). pub(crate) fn restore( cwd: &Path, - base: &str, + rev: Option<&str>, change: &AgentChange, before: &BTreeMap, saved: Option<&Path>, @@ -242,10 +320,13 @@ pub(crate) fn restore( })?; Some(std::fs::read(copy)?) } - None => { - let repo = Repo::discover(cwd)?.ok_or_else(|| IdeError::invalid("not in a git repository"))?; - repo.file_bytes_at(path, Some(base))? - } + None => match rev { + Some(rev) => { + let repo = Repo::discover(cwd)?.ok_or_else(|| IdeError::invalid("not in a git repository"))?; + repo.file_bytes_at(path, Some(rev))? + } + None => None, + }, }; match content { Some(bytes) => { @@ -490,6 +571,39 @@ mod tests { assert_eq!(again.accepted, ["a.txt", "gone.txt", "new.txt"]); } + #[test] + fn a_local_task_goes_back_to_its_start_unless_the_user_committed_on_top_of_it() { + let tmp = tempfile::tempdir().unwrap(); + let repo = tmp.path().join("repo"); + std::fs::create_dir(&repo).unwrap(); + git(&repo, &["init", "-q", "-b", "main"]); + let local = |start: Option<&str>| TaskBase::Local { cwd: repo.clone(), start: start.map(str::to_owned) }; + let head = || head_of(&repo).unwrap().1; + let commit = |name: &str| { + std::fs::write(repo.join(name), "x\n").unwrap(); + git(&repo, &["add", "."]); + git(&repo, &["commit", "-q", "-m", name]); + head_of(&repo).unwrap().1 + }; + + assert_eq!(local(None).rev().unwrap(), None, "no commits yet"); + let first = commit("a.txt"); + assert_eq!(local(None).rev().unwrap(), Some(first.clone()), "an imported session: HEAD"); + let start = commit("b.txt"); + assert_eq!(local(Some(&start)).rev().unwrap(), Some(start.clone()), "HEAD where it started"); + let ahead = commit("c.txt"); + assert_eq!(local(Some(&start)).rev().unwrap(), Some(ahead), "committed on top: HEAD"); + git(&repo, &["reset", "-q", "--soft", &first]); + assert_eq!(local(Some(&start)).rev().unwrap(), Some(start.clone()), "moved back: the start"); + git(&repo, &["checkout", "-q", "-b", "side"]); + commit("d.txt"); + assert_eq!(local(Some(&start)).rev().unwrap(), Some(start.clone()), "another branch: the start"); + let gone = "0123456789abcdef0123456789abcdef01234567"; + assert_eq!(local(Some(gone)).rev().unwrap(), Some(head()), "an unknown start: HEAD"); + let worktree = TaskBase::Worktree { cwd: repo.clone(), base: start.clone() }; + assert_eq!(worktree.rev().unwrap(), Some(start)); + } + #[test] fn accept_leaves_files_the_checkout_changed_itself() { let (_tmp, repo, created) = accept_repo();