1use crate::{
2 example::{Example, ExamplePromptInputs, ExampleState},
3 git,
4 headless::EpAppState,
5 progress::{ExampleProgress, InfoStyle, Step, StepProgress},
6};
7use anyhow::{Context as _, Result};
8use edit_prediction::{
9 EditPredictionStore,
10 udiff::{OpenedBuffers, refresh_worktree_entries, strip_diff_path_prefix},
11};
12use futures::AsyncWriteExt as _;
13use gpui::{AsyncApp, Entity};
14use language::{Anchor, Buffer, LanguageNotFound, ToOffset, ToPoint};
15use project::{Project, ProjectPath, buffer_store::BufferStoreEvent};
16use std::{fs, path::PathBuf, sync::Arc};
17
18pub async fn run_load_project(
19 example: &mut Example,
20 app_state: Arc<EpAppState>,
21 example_progress: &ExampleProgress,
22 mut cx: AsyncApp,
23) -> Result<()> {
24 if example.state.is_some() {
25 return Ok(());
26 }
27
28 let progress = example_progress.start(Step::LoadProject);
29
30 let project = setup_project(example, &app_state, &progress, &mut cx).await?;
31
32 progress.set_substatus("applying edit history");
33 let open_buffers = apply_edit_history(example, &project, &mut cx).await?;
34
35 let ep_store = cx
36 .update(|cx| EditPredictionStore::try_global(cx))
37 .context("EditPredictionStore not initialized")?;
38
39 let recent_paths: Vec<ProjectPath> = open_buffers
40 .buffers()
41 .filter_map(|buffer| {
42 buffer.read_with(&cx, |buffer, cx| {
43 buffer
44 .file()
45 .map(|file| ProjectPath::from_file(file.as_ref(), cx))
46 })
47 })
48 .collect();
49
50 ep_store.update(&mut cx, |store, cx| {
51 store.set_recent_paths_for_project(&project, recent_paths, cx);
52 });
53
54 progress.set_substatus("resolving cursor");
55 let (buffer, cursor_position) =
56 cursor_position(example, &project, &open_buffers, &mut cx).await?;
57 buffer
58 .read_with(&cx, |buffer, _| buffer.parsing_idle())
59 .await;
60
61 let edit_history = ep_store.update(&mut cx, |store, cx| {
62 store
63 .edit_history_for_project(&project, cx)
64 .into_iter()
65 .map(|e| e.event)
66 .collect()
67 });
68
69 let (prompt_inputs, language_name) = buffer.read_with(&cx, |buffer, _cx| {
70 let cursor_point = cursor_position.to_point(&buffer);
71 let language_name = buffer
72 .language()
73 .map(|l| l.name().to_string())
74 .unwrap_or_else(|| "Unknown".to_string());
75 (
76 ExamplePromptInputs {
77 content: buffer.text(),
78 cursor_row: cursor_point.row,
79 cursor_column: cursor_point.column,
80 cursor_offset: cursor_position.to_offset(&buffer),
81 edit_history,
82 related_files: example
83 .prompt_inputs
84 .take()
85 .map(|inputs| inputs.related_files)
86 .unwrap_or_default(),
87 },
88 language_name,
89 )
90 });
91
92 progress.set_info(language_name, InfoStyle::Normal);
93
94 example.prompt_inputs = Some(prompt_inputs);
95 example.state = Some(ExampleState {
96 buffer,
97 project,
98 cursor_position,
99 _open_buffers: open_buffers,
100 });
101 Ok(())
102}
103
104async fn cursor_position(
105 example: &Example,
106 project: &Entity<Project>,
107 open_buffers: &OpenedBuffers,
108 cx: &mut AsyncApp,
109) -> Result<(Entity<Buffer>, Anchor)> {
110 let language_registry = project.read_with(cx, |project, _| project.languages().clone());
111 let result = language_registry
112 .load_language_for_file_path(&example.spec.cursor_path)
113 .await;
114
115 if let Err(error) = result
116 && !error.is::<LanguageNotFound>()
117 {
118 return Err(error);
119 }
120
121 let cursor_path_str = example.spec.cursor_path.to_string_lossy();
122 // Also try cursor path with first component stripped - old examples may have
123 // paths like "zed/crates/foo.rs" instead of "crates/foo.rs".
124 let cursor_path_without_prefix: PathBuf =
125 example.spec.cursor_path.components().skip(1).collect();
126 let cursor_path_without_prefix_str = cursor_path_without_prefix.to_string_lossy();
127
128 // We try open_buffers first because the file might be new and not saved to disk
129 let cursor_buffer = if let Some(buffer) = open_buffers.get(cursor_path_str.as_ref()) {
130 buffer.clone()
131 } else if let Some(buffer) = open_buffers.get(cursor_path_without_prefix_str.as_ref()) {
132 buffer.clone()
133 } else {
134 // Since the worktree scanner is disabled, manually refresh entries for the cursor path.
135 if let Some(worktree) = project.read_with(cx, |project, cx| project.worktrees(cx).next()) {
136 refresh_worktree_entries(&worktree, [&*example.spec.cursor_path], cx).await?;
137 }
138
139 let cursor_path = project
140 .read_with(cx, |project, cx| {
141 project
142 .find_project_path(&example.spec.cursor_path, cx)
143 .or_else(|| project.find_project_path(&cursor_path_without_prefix, cx))
144 })
145 .with_context(|| {
146 format!(
147 "failed to find cursor path {}",
148 example.spec.cursor_path.display()
149 )
150 })?;
151
152 project
153 .update(cx, |project, cx| project.open_buffer(cursor_path, cx))
154 .await?
155 };
156
157 let (cursor_excerpt, cursor_offset_within_excerpt) = example.spec.cursor_excerpt()?;
158
159 let excerpt_offset = cursor_buffer.read_with(&*cx, |buffer, _cx| {
160 let text = buffer.text();
161
162 let mut matches = text.match_indices(&cursor_excerpt);
163 let (excerpt_offset, _) = matches.next().with_context(|| {
164 format!("Cursor excerpt did not exist in buffer:\n\n{cursor_excerpt}\n",)
165 })?;
166 anyhow::ensure!(
167 matches.next().is_none(),
168 "More than one cursor position match found",
169 );
170 Ok(excerpt_offset)
171 })?;
172
173 let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
174 let cursor_anchor =
175 cursor_buffer.read_with(&*cx, |buffer, _| buffer.anchor_after(cursor_offset));
176
177 Ok((cursor_buffer, cursor_anchor))
178}
179
180async fn setup_project(
181 example: &mut Example,
182 app_state: &Arc<EpAppState>,
183 step_progress: &StepProgress,
184 cx: &mut AsyncApp,
185) -> Result<Entity<Project>> {
186 let ep_store = cx
187 .update(|cx| EditPredictionStore::try_global(cx))
188 .context("Store should be initialized at init")?;
189
190 let worktree_path = setup_worktree(example, step_progress).await?;
191
192 if let Some(project) = app_state.project_cache.get(&example.spec.repository_url) {
193 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone());
194 let buffers = buffer_store.read_with(cx, |buffer_store, _| {
195 buffer_store.buffers().collect::<Vec<_>>()
196 });
197 for buffer in buffers {
198 buffer.update(cx, |buffer, cx| buffer.reload(cx)).await.ok();
199 }
200 ep_store.update(cx, |ep_store, _| {
201 ep_store.clear_history_for_project(&project);
202 });
203 return Ok(project);
204 }
205
206 let project = cx.update(|cx| {
207 Project::local(
208 app_state.client.clone(),
209 app_state.node_runtime.clone(),
210 app_state.user_store.clone(),
211 app_state.languages.clone(),
212 app_state.fs.clone(),
213 None,
214 false,
215 cx,
216 )
217 });
218
219 project
220 .update(cx, |project, cx| {
221 project.disable_worktree_scanner(cx);
222 project.create_worktree(&worktree_path, true, cx)
223 })
224 .await?;
225
226 app_state
227 .project_cache
228 .insert(example.spec.repository_url.clone(), project.clone());
229
230 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone());
231 cx.subscribe(&buffer_store, {
232 let project = project.clone();
233 move |_, event, cx| match event {
234 BufferStoreEvent::BufferAdded(buffer) => {
235 ep_store.update(cx, |store, cx| store.register_buffer(&buffer, &project, cx));
236 }
237 _ => {}
238 }
239 })
240 .detach();
241
242 Ok(project)
243}
244
245async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Result<PathBuf> {
246 let repo_name = example.repo_name().context("failed to get repo name")?;
247 let repo_dir = git::repo_path_for_url(&example.spec.repository_url)?;
248 let worktree_path = repo_name.worktree_path();
249 let repo_lock = git::lock_repo(&repo_dir).await;
250
251 // Clean up any stale git lock files from previous crashed runs.
252 // Safe-ish since we have our own lock.
253 // WARNING: Can corrupt worktrees if multiple processes of the CLI are running.
254 let worktree_git_dir = repo_dir
255 .join(".git/worktrees")
256 .join(repo_name.name.as_ref());
257 for lock_file in &["index.lock", "HEAD.lock", "config.lock"] {
258 let worktree_lock_path = worktree_git_dir.join(lock_file);
259 let repo_lock_path = repo_dir.join(".git").join(lock_file);
260 if worktree_lock_path.exists() {
261 fs::remove_file(&worktree_lock_path).ok();
262 }
263 if repo_lock_path.exists() {
264 fs::remove_file(&repo_lock_path).ok();
265 }
266 }
267
268 let mut git_repo_exists = false;
269 if repo_dir.is_dir() {
270 if git::run_git(&repo_dir, &["remote", "get-url", "origin"])
271 .await
272 .map_or(false, |origin| origin.trim() == example.spec.repository_url)
273 {
274 git_repo_exists = true;
275 } else {
276 fs::remove_dir_all(&repo_dir).ok();
277 }
278 }
279
280 if !git_repo_exists {
281 step_progress.set_substatus(format!("cloning {}", repo_name.name));
282 fs::create_dir_all(&repo_dir)?;
283 git::run_git(&repo_dir, &["init"]).await?;
284 git::run_git(
285 &repo_dir,
286 &["remote", "add", "origin", &example.spec.repository_url],
287 )
288 .await?;
289 }
290
291 // Resolve the example to a revision, fetching it if needed.
292 step_progress.set_substatus("fetching");
293 let revision = git::fetch_if_needed(&repo_dir, &example.spec.revision).await?;
294
295 // Clean up any stale worktree registrations from previous crashed runs.
296 git::run_git(&repo_dir, &["worktree", "prune"]).await.ok();
297
298 // Create the worktree for this example if needed.
299 step_progress.set_substatus("preparing worktree");
300
301 // Check if worktree exists and is valid (not just a directory from a crashed run).
302 let worktree_valid = worktree_path.is_dir()
303 && git::run_git(&worktree_path, &["rev-parse", "--git-dir"])
304 .await
305 .is_ok();
306
307 if worktree_valid {
308 git::run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
309 git::run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
310 git::run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
311 } else {
312 let worktree_path_string = worktree_path.to_string_lossy();
313
314 // Clean up invalid worktree directory and registration if they exist.
315 if worktree_path.exists() {
316 fs::remove_dir_all(&worktree_path).ok();
317 }
318 git::run_git(
319 &repo_dir,
320 &["worktree", "remove", "--force", &worktree_path_string],
321 )
322 .await
323 .ok();
324
325 let branch_name = example.spec.filename();
326 git::run_git(
327 &repo_dir,
328 &["branch", "-f", &branch_name, revision.as_str()],
329 )
330 .await?;
331 git::run_git(
332 &repo_dir,
333 &["worktree", "add", "-f", &worktree_path_string, &branch_name],
334 )
335 .await?;
336 }
337 drop(repo_lock);
338
339 if !example.spec.uncommitted_diff.is_empty() {
340 step_progress.set_substatus("applying diff");
341
342 // old examples had full paths in the uncommitted diff.
343 let uncommitted_diff =
344 strip_diff_path_prefix(&example.spec.uncommitted_diff, &repo_name.name);
345
346 let mut apply_process = smol::process::Command::new("git")
347 .current_dir(&worktree_path)
348 .args(&["apply", "-"])
349 .stdin(std::process::Stdio::piped())
350 .spawn()?;
351
352 let mut stdin = apply_process.stdin.take().context("Failed to get stdin")?;
353 stdin.write_all(uncommitted_diff.as_bytes()).await?;
354 stdin.close().await?;
355 drop(stdin);
356
357 let apply_result = apply_process.output().await?;
358 anyhow::ensure!(
359 apply_result.status.success(),
360 "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
361 apply_result.status,
362 String::from_utf8_lossy(&apply_result.stderr),
363 String::from_utf8_lossy(&apply_result.stdout),
364 );
365 }
366
367 step_progress.clear_substatus();
368 Ok(worktree_path)
369}
370
371async fn apply_edit_history(
372 example: &Example,
373 project: &Entity<Project>,
374 cx: &mut AsyncApp,
375) -> Result<OpenedBuffers> {
376 edit_prediction::udiff::apply_diff(&example.spec.edit_history, project, cx).await
377}