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