1use crate::{
2 example::{Example, ExampleBuffer, ExampleState},
3 headless::EpAppState,
4 paths::{REPOS_DIR, WORKTREES_DIR},
5 progress::{InfoStyle, Progress, Step, StepProgress},
6};
7use anyhow::{Context as _, Result};
8use collections::HashMap;
9use edit_prediction::EditPredictionStore;
10use edit_prediction::udiff::OpenedBuffers;
11use futures::{
12 AsyncWriteExt as _,
13 lock::{Mutex, OwnedMutexGuard},
14};
15use gpui::{AsyncApp, Entity};
16use language::{Anchor, Buffer, LanguageNotFound, ToOffset, ToPoint};
17use project::buffer_store::BufferStoreEvent;
18use project::{Project, ProjectPath};
19use std::{
20 cell::RefCell,
21 fs,
22 path::{Path, PathBuf},
23 sync::Arc,
24};
25use util::{paths::PathStyle, rel_path::RelPath};
26use zeta_prompt::CURSOR_MARKER;
27
28pub async fn run_load_project(
29 example: &mut Example,
30 app_state: Arc<EpAppState>,
31 mut cx: AsyncApp,
32) -> Result<()> {
33 if example.state.is_some() {
34 return Ok(());
35 }
36
37 let progress = Progress::global().start(Step::LoadProject, &example.spec.name);
38
39 let project = setup_project(example, &app_state, &progress, &mut cx).await?;
40
41 let _open_buffers = apply_edit_history(example, &project, &mut cx).await?;
42
43 let (buffer, cursor_position) = cursor_position(example, &project, &mut cx).await?;
44 let (example_buffer, language_name) = buffer.read_with(&cx, |buffer, _cx| {
45 let cursor_point = cursor_position.to_point(&buffer);
46 let language_name = buffer
47 .language()
48 .map(|l| l.name().to_string())
49 .unwrap_or_else(|| "Unknown".to_string());
50 (
51 ExampleBuffer {
52 content: buffer.text(),
53 cursor_row: cursor_point.row,
54 cursor_column: cursor_point.column,
55 cursor_offset: cursor_position.to_offset(&buffer),
56 },
57 language_name,
58 )
59 })?;
60
61 progress.set_info(language_name, InfoStyle::Normal);
62
63 example.buffer = Some(example_buffer);
64 example.state = Some(ExampleState {
65 buffer,
66 project,
67 cursor_position,
68 _open_buffers,
69 });
70 Ok(())
71}
72
73async fn cursor_position(
74 example: &Example,
75 project: &Entity<Project>,
76 cx: &mut AsyncApp,
77) -> Result<(Entity<Buffer>, Anchor)> {
78 let language_registry = project.read_with(cx, |project, _| project.languages().clone())?;
79 let result = language_registry
80 .load_language_for_file_path(&example.spec.cursor_path)
81 .await;
82
83 if let Err(error) = result
84 && !error.is::<LanguageNotFound>()
85 {
86 return Err(error);
87 }
88
89 let worktree = project.read_with(cx, |project, cx| {
90 project
91 .visible_worktrees(cx)
92 .next()
93 .context("No visible worktrees")
94 })??;
95
96 let cursor_path = RelPath::new(&example.spec.cursor_path, PathStyle::Posix)
97 .context("Failed to create RelPath")?
98 .into_arc();
99 let cursor_buffer = project
100 .update(cx, |project, cx| {
101 project.open_buffer(
102 ProjectPath {
103 worktree_id: worktree.read(cx).id(),
104 path: cursor_path,
105 },
106 cx,
107 )
108 })?
109 .await?;
110 let cursor_offset_within_excerpt = example
111 .spec
112 .cursor_position
113 .find(CURSOR_MARKER)
114 .context("missing cursor marker")?;
115 let mut cursor_excerpt = example.spec.cursor_position.clone();
116 cursor_excerpt.replace_range(
117 cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
118 "",
119 );
120 let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
121 let text = buffer.text();
122
123 let mut matches = text.match_indices(&cursor_excerpt);
124 let (excerpt_offset, _) = matches.next().with_context(|| {
125 format!(
126 "\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n.Example: {}\nCursor excerpt did not exist in buffer.",
127 example.spec.name
128 )
129 })?;
130 anyhow::ensure!(
131 matches.next().is_none(),
132 "More than one cursor position match found for {}",
133 &example.spec.name
134 );
135 Ok(excerpt_offset)
136 })??;
137
138 let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
139 let cursor_anchor =
140 cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
141
142 Ok((cursor_buffer, cursor_anchor))
143}
144
145async fn setup_project(
146 example: &mut Example,
147 app_state: &Arc<EpAppState>,
148 step_progress: &StepProgress,
149 cx: &mut AsyncApp,
150) -> Result<Entity<Project>> {
151 let ep_store = cx
152 .update(|cx| EditPredictionStore::try_global(cx))?
153 .context("Store should be initialized at init")?;
154
155 let worktree_path = setup_worktree(example, step_progress).await?;
156
157 if let Some(project) = app_state.project_cache.get(&example.spec.repository_url) {
158 ep_store.update(cx, |ep_store, _| {
159 ep_store.clear_history_for_project(&project);
160 })?;
161 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone())?;
162 let buffers = buffer_store.read_with(cx, |buffer_store, _| {
163 buffer_store.buffers().collect::<Vec<_>>()
164 })?;
165 for buffer in buffers {
166 buffer
167 .update(cx, |buffer, cx| buffer.reload(cx))?
168 .await
169 .ok();
170 }
171 return Ok(project);
172 }
173
174 let project = cx.update(|cx| {
175 Project::local(
176 app_state.client.clone(),
177 app_state.node_runtime.clone(),
178 app_state.user_store.clone(),
179 app_state.languages.clone(),
180 app_state.fs.clone(),
181 None,
182 false,
183 cx,
184 )
185 })?;
186
187 project
188 .update(cx, |project, cx| {
189 project.disable_worktree_scanner(cx);
190 project.create_worktree(&worktree_path, true, cx)
191 })?
192 .await?;
193
194 app_state
195 .project_cache
196 .insert(example.spec.repository_url.clone(), project.clone());
197
198 let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone())?;
199 cx.subscribe(&buffer_store, {
200 let project = project.clone();
201 move |_, event, cx| match event {
202 BufferStoreEvent::BufferAdded(buffer) => {
203 ep_store.update(cx, |store, cx| store.register_buffer(&buffer, &project, cx));
204 }
205 _ => {}
206 }
207 })?
208 .detach();
209
210 Ok(project)
211}
212
213async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Result<PathBuf> {
214 let (repo_owner, repo_name) = example.repo_name().context("failed to get repo name")?;
215 let repo_dir = REPOS_DIR.join(repo_owner.as_ref()).join(repo_name.as_ref());
216 let worktree_path = WORKTREES_DIR
217 .join(repo_owner.as_ref())
218 .join(repo_name.as_ref());
219 let repo_lock = lock_repo(&repo_dir).await;
220
221 if !repo_dir.is_dir() {
222 step_progress.set_substatus(format!("cloning {}", repo_name));
223 fs::create_dir_all(&repo_dir)?;
224 run_git(&repo_dir, &["init"]).await?;
225 run_git(
226 &repo_dir,
227 &["remote", "add", "origin", &example.spec.repository_url],
228 )
229 .await?;
230 }
231
232 // Resolve the example to a revision, fetching it if needed.
233 let revision = run_git(
234 &repo_dir,
235 &[
236 "rev-parse",
237 &format!("{}^{{commit}}", example.spec.revision),
238 ],
239 )
240 .await;
241 let revision = if let Ok(revision) = revision {
242 revision
243 } else {
244 step_progress.set_substatus("fetching");
245 if run_git(
246 &repo_dir,
247 &["fetch", "--depth", "1", "origin", &example.spec.revision],
248 )
249 .await
250 .is_err()
251 {
252 run_git(&repo_dir, &["fetch", "origin"]).await?;
253 }
254 let revision = run_git(&repo_dir, &["rev-parse", "FETCH_HEAD"]).await?;
255 revision
256 };
257
258 // Create the worktree for this example if needed.
259 step_progress.set_substatus("preparing worktree");
260 if worktree_path.is_dir() {
261 run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
262 run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
263 run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
264 } else {
265 let worktree_path_string = worktree_path.to_string_lossy();
266 run_git(
267 &repo_dir,
268 &["branch", "-f", &example.spec.name, revision.as_str()],
269 )
270 .await?;
271 run_git(
272 &repo_dir,
273 &[
274 "worktree",
275 "add",
276 "-f",
277 &worktree_path_string,
278 &example.spec.name,
279 ],
280 )
281 .await?;
282 }
283 drop(repo_lock);
284
285 // Apply the uncommitted diff for this example.
286 if !example.spec.uncommitted_diff.is_empty() {
287 step_progress.set_substatus("applying diff");
288 let mut apply_process = smol::process::Command::new("git")
289 .current_dir(&worktree_path)
290 .args(&["apply", "-"])
291 .stdin(std::process::Stdio::piped())
292 .spawn()?;
293
294 let mut stdin = apply_process.stdin.take().context("Failed to get stdin")?;
295 stdin
296 .write_all(example.spec.uncommitted_diff.as_bytes())
297 .await?;
298 stdin.close().await?;
299 drop(stdin);
300
301 let apply_result = apply_process.output().await?;
302 anyhow::ensure!(
303 apply_result.status.success(),
304 "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
305 apply_result.status,
306 String::from_utf8_lossy(&apply_result.stderr),
307 String::from_utf8_lossy(&apply_result.stdout),
308 );
309 }
310
311 step_progress.clear_substatus();
312 Ok(worktree_path)
313}
314
315async fn apply_edit_history(
316 example: &Example,
317 project: &Entity<Project>,
318 cx: &mut AsyncApp,
319) -> Result<OpenedBuffers> {
320 edit_prediction::udiff::apply_diff(&example.spec.edit_history, project, cx).await
321}
322
323thread_local! {
324 static REPO_LOCKS: RefCell<HashMap<PathBuf, Arc<Mutex<()>>>> = RefCell::new(HashMap::default());
325}
326
327#[must_use]
328pub async fn lock_repo(path: impl AsRef<Path>) -> OwnedMutexGuard<()> {
329 REPO_LOCKS
330 .with(|cell| {
331 cell.borrow_mut()
332 .entry(path.as_ref().to_path_buf())
333 .or_default()
334 .clone()
335 })
336 .lock_owned()
337 .await
338}
339
340async fn run_git(repo_path: &Path, args: &[&str]) -> Result<String> {
341 let output = smol::process::Command::new("git")
342 .current_dir(repo_path)
343 .args(args)
344 .output()
345 .await?;
346
347 anyhow::ensure!(
348 output.status.success(),
349 "`git {}` within `{}` failed with status: {}\nstderr:\n{}\nstdout:\n{}",
350 args.join(" "),
351 repo_path.display(),
352 output.status,
353 String::from_utf8_lossy(&output.stderr),
354 String::from_utf8_lossy(&output.stdout),
355 );
356 Ok(String::from_utf8(output.stdout)?.trim().to_string())
357}