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 if let Some(project) = app_state.project_cache.get(&example.spec.repository_url) {
194 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone());
195 let buffers = buffer_store.read_with(cx, |buffer_store, _| {
196 buffer_store.buffers().collect::<Vec<_>>()
197 });
198 for buffer in buffers {
199 buffer.update(cx, |buffer, cx| buffer.reload(cx)).await.ok();
200 }
201 ep_store.update(cx, |ep_store, _| {
202 ep_store.clear_history_for_project(&project);
203 });
204 return Ok(project);
205 }
206
207 let project = cx.update(|cx| {
208 Project::local(
209 app_state.client.clone(),
210 app_state.node_runtime.clone(),
211 app_state.user_store.clone(),
212 app_state.languages.clone(),
213 app_state.fs.clone(),
214 None,
215 project::LocalProjectFlags {
216 init_worktree_trust: false,
217 watch_global_configs: false,
218 },
219 cx,
220 )
221 });
222
223 project
224 .update(cx, |project, cx| {
225 project.disable_worktree_scanner(cx);
226 project.create_worktree(&worktree_path, true, cx)
227 })
228 .await?;
229
230 app_state
231 .project_cache
232 .insert(example.spec.repository_url.clone(), project.clone());
233
234 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone());
235 cx.subscribe(&buffer_store, {
236 let project = project.downgrade();
237 let ep_store = ep_store.downgrade();
238 move |_, event, cx| match event {
239 BufferStoreEvent::BufferAdded(buffer) => {
240 let Some(project) = project.upgrade() else {
241 return;
242 };
243 ep_store
244 .update(cx, |store, cx| store.register_buffer(&buffer, &project, cx))
245 .ok();
246 }
247 _ => {}
248 }
249 })
250 .detach();
251
252 Ok(project)
253}
254
255async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Result<PathBuf> {
256 let repo_name = example.repo_name().context("failed to get repo name")?;
257 let repo_dir = git::repo_path_for_url(&example.spec.repository_url)?;
258 let worktree_path = repo_name.worktree_path();
259 let repo_lock = git::lock_repo(&repo_dir).await;
260
261 // Clean up any stale git lock files from previous crashed runs.
262 // Safe-ish since we have our own lock.
263 // WARNING: Can corrupt worktrees if multiple processes of the CLI are running.
264 let worktree_git_dir = repo_dir
265 .join(".git/worktrees")
266 .join(repo_name.name.as_ref());
267 for lock_file in &["index.lock", "HEAD.lock", "config.lock"] {
268 let worktree_lock_path = worktree_git_dir.join(lock_file);
269 let repo_lock_path = repo_dir.join(".git").join(lock_file);
270 if worktree_lock_path.exists() {
271 fs::remove_file(&worktree_lock_path).ok();
272 }
273 if repo_lock_path.exists() {
274 fs::remove_file(&repo_lock_path).ok();
275 }
276 }
277
278 let mut git_repo_exists = false;
279 if repo_dir.is_dir() {
280 if git::run_git(&repo_dir, &["remote", "get-url", "origin"])
281 .await
282 .map_or(false, |origin| origin.trim() == example.spec.repository_url)
283 {
284 git_repo_exists = true;
285 } else {
286 fs::remove_dir_all(&repo_dir).ok();
287 }
288 }
289
290 if !git_repo_exists {
291 step_progress.set_substatus(format!("cloning {}", repo_name.name));
292 fs::create_dir_all(&repo_dir)?;
293 git::run_git(&repo_dir, &["init"]).await?;
294 git::run_git(
295 &repo_dir,
296 &["remote", "add", "origin", &example.spec.repository_url],
297 )
298 .await?;
299 }
300
301 // Resolve the example to a revision, fetching it if needed.
302 step_progress.set_substatus("fetching");
303 let revision = git::fetch_if_needed(&repo_dir, &example.spec.revision).await?;
304
305 // Clean up any stale worktree registrations from previous crashed runs.
306 git::run_git(&repo_dir, &["worktree", "prune"]).await.ok();
307
308 // Create the worktree for this example if needed.
309 step_progress.set_substatus("preparing worktree");
310
311 // Check if worktree exists and is valid (not just a directory from a crashed run).
312 let worktree_valid = worktree_path.is_dir()
313 && git::run_git(&worktree_path, &["rev-parse", "--git-dir"])
314 .await
315 .is_ok();
316
317 if worktree_valid {
318 git::run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
319 git::run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
320 git::run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
321 } else {
322 let worktree_path_string = worktree_path.to_string_lossy();
323
324 // Clean up invalid worktree directory and registration if they exist.
325 if worktree_path.exists() {
326 fs::remove_dir_all(&worktree_path).ok();
327 }
328 git::run_git(
329 &repo_dir,
330 &["worktree", "remove", "--force", &worktree_path_string],
331 )
332 .await
333 .ok();
334
335 let branch_name = example.spec.filename();
336 git::run_git(
337 &repo_dir,
338 &["branch", "-f", &branch_name, revision.as_str()],
339 )
340 .await?;
341 git::run_git(
342 &repo_dir,
343 &["worktree", "add", "-f", &worktree_path_string, &branch_name],
344 )
345 .await?;
346 }
347 drop(repo_lock);
348
349 if !example.spec.uncommitted_diff.is_empty() {
350 step_progress.set_substatus("applying diff");
351
352 // old examples had full paths in the uncommitted diff.
353 let uncommitted_diff =
354 strip_diff_path_prefix(&example.spec.uncommitted_diff, &repo_name.name);
355
356 let mut apply_process = smol::process::Command::new("git")
357 .current_dir(&worktree_path)
358 .args(&["apply", "-"])
359 .stdin(std::process::Stdio::piped())
360 .spawn()?;
361
362 let mut stdin = apply_process.stdin.take().context("Failed to get stdin")?;
363 stdin.write_all(uncommitted_diff.as_bytes()).await?;
364 stdin.close().await?;
365 drop(stdin);
366
367 let apply_result = apply_process.output().await?;
368 anyhow::ensure!(
369 apply_result.status.success(),
370 "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
371 apply_result.status,
372 String::from_utf8_lossy(&apply_result.stderr),
373 String::from_utf8_lossy(&apply_result.stdout),
374 );
375 }
376
377 step_progress.clear_substatus();
378 Ok(worktree_path)
379}
380
381async fn apply_edit_history(
382 example: &Example,
383 project: &Entity<Project>,
384 cx: &mut AsyncApp,
385) -> Result<OpenedBuffers> {
386 edit_prediction::udiff::apply_diff(&example.spec.edit_history, project, cx).await
387}