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();