1use crate::{
2 example::{Example, ExampleState},
3 git,
4 headless::EpAppState,
5 progress::{ExampleProgress, InfoStyle, Step, StepProgress},
6};
7use anyhow::{Context as _, Result};
8use edit_prediction::{
9 EditPredictionStore,
10 cursor_excerpt::{compute_cursor_excerpt, compute_syntax_ranges},
11 udiff::{OpenedBuffers, refresh_worktree_entries, strip_diff_path_prefix},
12};
13use futures::AsyncWriteExt as _;
14use gpui::{AsyncApp, Entity};
15use language::{Anchor, Buffer, LanguageNotFound, ToOffset};
16use project::{Project, ProjectPath, buffer_store::BufferStoreEvent};
17use std::{fs, path::PathBuf, sync::Arc};
18use zeta_prompt::ZetaPromptInput;
19
20pub async fn run_load_project(
21 example: &mut Example,
22 app_state: Arc<EpAppState>,
23 example_progress: &ExampleProgress,
24 mut cx: AsyncApp,
25) -> Result<()> {
26 if example.state.is_some() {
27 return Ok(());
28 }
29
30 let progress = example_progress.start(Step::LoadProject);
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 let ep_store = cx
38 .update(|cx| EditPredictionStore::try_global(cx))
39 .context("EditPredictionStore not initialized")?;
40
41 let recent_paths: Vec<ProjectPath> = open_buffers
42 .buffers()
43 .filter_map(|buffer| {
44 buffer.read_with(&cx, |buffer, cx| {
45 buffer
46 .file()
47 .map(|file| ProjectPath::from_file(file.as_ref(), cx))
48 })
49 })
50 .collect();
51
52 ep_store.update(&mut cx, |store, cx| {
53 store.set_recent_paths_for_project(&project, recent_paths, cx);
54 });
55
56 progress.set_substatus("resolving cursor");
57 let (buffer, cursor_position) =
58 cursor_position(example, &project, &open_buffers, &mut cx).await?;
59 buffer
60 .read_with(&cx, |buffer, _| buffer.parsing_idle())
61 .await;
62
63 let events: Vec<Arc<zeta_prompt::Event>> = ep_store.update(&mut cx, |store, cx| {
64 store
65 .edit_history_for_project(&project, cx)
66 .into_iter()
67 .map(|e| e.event)
68 .collect()
69 });
70
71 let existing_related_files = example
72 .prompt_inputs
73 .take()
74 .and_then(|inputs| inputs.related_files);
75
76 let (prompt_inputs, language_name) = buffer.read_with(&cx, |buffer, _cx| {
77 let snapshot = buffer.snapshot();
78 let cursor_offset = cursor_position.to_offset(&snapshot);
79 let language_name = buffer
80 .language()
81 .map(|l| l.name().to_string())
82 .unwrap_or_else(|| "Unknown".to_string());
83
84 let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) =
85 compute_cursor_excerpt(&snapshot, cursor_offset);
86
87 let cursor_excerpt: Arc<str> = buffer
88 .text_for_range(excerpt_offset_range.clone())
89 .collect::<String>()
90 .into();
91 let syntax_ranges = compute_syntax_ranges(&snapshot, cursor_offset, &excerpt_offset_range);
92 let excerpt_ranges = zeta_prompt::compute_legacy_excerpt_ranges(
93 &cursor_excerpt,
94 cursor_offset_in_excerpt,
95 &syntax_ranges,
96 );
97
98 (
99 ZetaPromptInput {
100 cursor_path: example.spec.cursor_path.clone(),
101 cursor_excerpt,
102 cursor_offset_in_excerpt,
103 excerpt_start_row: Some(excerpt_point_range.start.row),
104 events,
105 related_files: existing_related_files,
106 active_buffer_diagnostics: vec![],
107 excerpt_ranges,
108 syntax_ranges: Some(syntax_ranges),
109 in_open_source_repo: false,
110 can_collect_data: false,
111 experiment: None,
112 repo_url: None,
113 },
114 language_name,
115 )
116 });
117
118 progress.set_info(language_name, InfoStyle::Normal);
119
120 example.prompt_inputs = Some(prompt_inputs);
121 example.state = Some(ExampleState {
122 buffer,
123 project,
124 cursor_position,
125 _open_buffers: open_buffers,
126 });
127 Ok(())
128}
129
130async fn cursor_position(
131 example: &Example,
132 project: &Entity<Project>,
133 open_buffers: &OpenedBuffers,
134 cx: &mut AsyncApp,
135) -> Result<(Entity<Buffer>, Anchor)> {
136 let language_registry = project.read_with(cx, |project, _| project.languages().clone());
137 let result = language_registry
138 .load_language_for_file_path(&example.spec.cursor_path)
139 .await;
140
141 if let Err(error) = result
142 && !error.is::<LanguageNotFound>()
143 {
144 return Err(error);
145 }
146
147 let cursor_path_str = example.spec.cursor_path.to_string_lossy();
148 // Also try cursor path with first component stripped - old examples may have
149 // paths like "zed/crates/foo.rs" instead of "crates/foo.rs".
150 let cursor_path_without_prefix: PathBuf =
151 example.spec.cursor_path.components().skip(1).collect();
152 let cursor_path_without_prefix_str = cursor_path_without_prefix.to_string_lossy();
153
154 // We try open_buffers first because the file might be new and not saved to disk
155 let cursor_buffer = if let Some(buffer) = open_buffers.get(cursor_path_str.as_ref()) {
156 buffer.clone()
157 } else if let Some(buffer) = open_buffers.get(cursor_path_without_prefix_str.as_ref()) {
158 buffer.clone()
159 } else {
160 // Since the worktree scanner is disabled, manually refresh entries for the cursor path.
161 if let Some(worktree) = project.read_with(cx, |project, cx| project.worktrees(cx).next()) {
162 refresh_worktree_entries(&worktree, [&*example.spec.cursor_path], cx).await?;
163 }
164
165 let cursor_path = project
166 .read_with(cx, |project, cx| {
167 project
168 .find_project_path(&example.spec.cursor_path, cx)
169 .or_else(|| project.find_project_path(&cursor_path_without_prefix, cx))
170 })
171 .with_context(|| {
172 format!(
173 "failed to find cursor path {}",
174 example.spec.cursor_path.display()
175 )
176 })?;
177
178 project
179 .update(cx, |project, cx| project.open_buffer(cursor_path, cx))
180 .await?
181 };
182
183 let (cursor_excerpt, cursor_offset_within_excerpt) = example.spec.cursor_excerpt()?;
184
185 let excerpt_offset = cursor_buffer.read_with(&*cx, |buffer, _cx| {
186 let text = buffer.text();
187
188 let mut matches = text.match_indices(&cursor_excerpt);
189 let (excerpt_offset, _) = matches.next().with_context(|| {
190 format!("Cursor excerpt did not exist in buffer:\n\n{cursor_excerpt}\n",)
191 })?;
192 anyhow::ensure!(
193 matches.next().is_none(),
194 "More than one cursor position match found",
195 );
196 Ok(excerpt_offset)
197 })?;
198
199 let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
200 let cursor_anchor =
201 cursor_buffer.read_with(&*cx, |buffer, _| buffer.anchor_after(cursor_offset));
202
203 Ok((cursor_buffer, cursor_anchor))
204}
205
206async fn setup_project(
207 example: &mut Example,
208 app_state: &Arc<EpAppState>,
209 step_progress: &StepProgress,
210 cx: &mut AsyncApp,
211) -> Result<Entity<Project>> {
212 let ep_store = cx
213 .update(|cx| EditPredictionStore::try_global(cx))
214 .context("Store should be initialized at init")?;
215
216 let worktree_path = setup_worktree(example, step_progress).await?;
217
218 let project = cx.update(|cx| {
219 Project::local(
220 app_state.client.clone(),
221 app_state.node_runtime.clone(),
222 app_state.user_store.clone(),
223 app_state.languages.clone(),
224 app_state.fs.clone(),
225 None,
226 project::LocalProjectFlags {
227 init_worktree_trust: false,
228 watch_global_configs: false,
229 },
230 cx,
231 )
232 });
233
234 project
235 .update(cx, |project, cx| {
236 project.disable_worktree_scanner(cx);
237 project.create_worktree(&worktree_path, true, cx)
238 })
239 .await?;
240
241 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone());
242 cx.subscribe(&buffer_store, {
243 let project = project.downgrade();
244 let ep_store = ep_store.downgrade();
245 move |_, event, cx| match event {
246 BufferStoreEvent::BufferAdded(buffer) => {
247 let Some(project) = project.upgrade() else {
248 return;
249 };
250 ep_store
251 .update(cx, |store, cx| store.register_buffer(&buffer, &project, cx))
252 .ok();
253 }
254 _ => {}
255 }
256 })
257 .detach();
258
259 Ok(project)
260}
261
262async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Result<PathBuf> {
263 let repo_name = example.repo_name().context("failed to get repo name")?;
264 let repo_dir = git::repo_path_for_url(&example.spec.repository_url)?;
265 let worktree_path = repo_name.worktree_path();
266 let repo_lock = git::lock_repo(&repo_dir).await;
267
268 // Clean up any stale git lock files from previous crashed runs.
269 // Safe-ish since we have our own lock.
270 // WARNING: Can corrupt worktrees if multiple processes of the CLI are running.
271 let worktree_git_dir = repo_dir
272 .join(".git/worktrees")
273 .join(repo_name.name.as_ref());
274 for lock_file in &["index.lock", "HEAD.lock", "config.lock"] {
275 let worktree_lock_path = worktree_git_dir.join(lock_file);
276 let repo_lock_path = repo_dir.join(".git").join(lock_file);
277 if worktree_lock_path.exists() {
278 fs::remove_file(&worktree_lock_path).ok();
279 }
280 if repo_lock_path.exists() {
281 fs::remove_file(&repo_lock_path).ok();
282 }
283 }
284
285 let mut git_repo_exists = false;
286 if repo_dir.is_dir() {
287 if git::run_git(&repo_dir, &["remote", "get-url", "origin"])
288 .await
289 .map_or(false, |origin| origin.trim() == example.spec.repository_url)
290 {
291 git_repo_exists = true;
292 } else {
293 fs::remove_dir_all(&repo_dir).ok();
294 }
295 }
296
297 if !git_repo_exists {
298 step_progress.set_substatus(format!("cloning {}", repo_name.name));
299 fs::create_dir_all(&repo_dir)?;
300 git::run_git(&repo_dir, &["init"]).await?;
301 git::run_git(
302 &repo_dir,
303 &["remote", "add", "origin", &example.spec.repository_url],
304 )
305 .await?;
306 }
307
308 // Resolve the example to a revision, fetching it if needed.
309 step_progress.set_substatus("fetching");
310 let revision = git::fetch_if_needed(&repo_dir, &example.spec.revision).await?;
311
312 // Clean up any stale worktree registrations from previous crashed runs.
313 git::run_git(&repo_dir, &["worktree", "prune"]).await.ok();
314
315 // Create the worktree for this example if needed.
316 step_progress.set_substatus("preparing worktree");
317
318 // Check if worktree exists and is valid (not just a directory from a crashed run).
319 let worktree_valid = worktree_path.is_dir()
320 && git::run_git(&worktree_path, &["rev-parse", "--git-dir"])
321 .await
322 .is_ok();
323
324 if worktree_valid {
325 git::run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
326 git::run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
327 git::run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
328 } else {
329 let worktree_path_string = worktree_path.to_string_lossy();
330
331 // Clean up invalid worktree directory and registration if they exist.
332 if worktree_path.exists() {
333 fs::remove_dir_all(&worktree_path).ok();
334 }
335 git::run_git(
336 &repo_dir,
337 &["worktree", "remove", "--force", &worktree_path_string],
338 )
339 .await
340 .ok();
341
342 let branch_name = example.spec.filename();
343 git::run_git(
344 &repo_dir,
345 &["branch", "-f", &branch_name, revision.as_str()],
346 )
347 .await?;
348 git::run_git(
349 &repo_dir,
350 &["worktree", "add", "-f", &worktree_path_string, &branch_name],
351 )
352 .await?;
353 }
354 drop(repo_lock);
355
356 if !example.spec.uncommitted_diff.is_empty() {
357 step_progress.set_substatus("applying diff");
358
359 // old examples had full paths in the uncommitted diff.
360 let uncommitted_diff =
361 strip_diff_path_prefix(&example.spec.uncommitted_diff, &repo_name.name);
362
363 let mut apply_process = smol::process::Command::new("git")
364 .current_dir(&worktree_path)
365 .args(&["apply", "-"])
366 .stdin(std::process::Stdio::piped())
367 .spawn()?;
368
369 let mut stdin = apply_process.stdin.take().context("Failed to get stdin")?;
370 stdin.write_all(uncommitted_diff.as_bytes()).await?;
371 stdin.close().await?;
372 drop(stdin);
373
374 let apply_result = apply_process.output().await?;
375 anyhow::ensure!(
376 apply_result.status.success(),
377 "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
378 apply_result.status,
379 String::from_utf8_lossy(&apply_result.stderr),
380 String::from_utf8_lossy(&apply_result.stdout),
381 );
382 }
383
384 step_progress.clear_substatus();
385 Ok(worktree_path)
386}
387
388async fn apply_edit_history(
389 example: &Example,
390 project: &Entity<Project>,
391 cx: &mut AsyncApp,
392) -> Result<OpenedBuffers> {
393 edit_prediction::udiff::apply_diff(&example.spec.edit_history, project, cx).await
394}