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