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