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