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