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