Fix more errors found when retrieving context for a huge example batch (#47039)

Max Brunsfeld created

Release Notes:

- N/A

Change summary

crates/edit_prediction_cli/src/load_project.rs | 30 ++++++++++++++++---
crates/edit_prediction_cli/src/main.rs         | 14 +++++++++
2 files changed, 39 insertions(+), 5 deletions(-)

Detailed changes

crates/edit_prediction_cli/src/load_project.rs 🔗

@@ -259,7 +259,16 @@ async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Resu
         fs::remove_file(&index_lock).ok();
     }
 
-    if !repo_dir.is_dir() {
+    let mut git_repo_exists = false;
+    if repo_dir.is_dir() {
+        if git::run_git(&repo_dir, &["status"]).await.is_ok() {
+            git_repo_exists = true;
+        } else {
+            fs::remove_dir_all(&repo_dir).ok();
+        }
+    }
+
+    if !git_repo_exists {
         step_progress.set_substatus(format!("cloning {}", repo_name.name));
         fs::create_dir_all(&repo_dir)?;
         git::run_git(&repo_dir, &["init"]).await?;
@@ -276,11 +285,22 @@ async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Resu
 
     // Create the worktree for this example if needed.
     step_progress.set_substatus("preparing worktree");
+
+    let mut worktree_exists = false;
     if worktree_path.is_dir() {
-        git::run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
-        git::run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
-        git::run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
-    } else {
+        if git::run_git(&worktree_path, &["clean", "--force", "-d"])
+            .await
+            .is_ok()
+        {
+            git::run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
+            git::run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
+            worktree_exists = true;
+        } else {
+            fs::remove_dir_all(&worktree_path).ok();
+        }
+    }
+
+    if !worktree_exists {
         let worktree_path_string = worktree_path.to_string_lossy();
         let branch_name = example.spec.filename();
         git::run_git(

crates/edit_prediction_cli/src/main.rs 🔗

@@ -650,6 +650,20 @@ fn main() {
                                 }
                             }
 
+                            if let Some(state) =
+                                repo_examples.first().and_then(|e| e.state.as_ref())
+                            {
+                                let mut cx = cx.clone();
+                                if let Some(ep_store) =
+                                    cx.update(|cx| EditPredictionStore::try_global(cx))
+                                {
+                                    let project = state.project.clone();
+                                    ep_store.update(&mut cx, |store, _| {
+                                        store.remove_project(&project);
+                                    });
+                                }
+                            }
+
                             app_state
                                 .project_cache
                                 .remove(&repo_examples.first().unwrap().spec.repository_url);