tenant.openagents/omega

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::Zeta2PromptInput;
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            Zeta2PromptInput {
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                repo_url: None,
112            },
113            language_name,
114        )
115    });
116
117    progress.set_info(language_name, InfoStyle::Normal);
118
119    example.prompt_inputs = Some(prompt_inputs);
120    example.state = Some(ExampleState {
121        buffer,
122        project,
123        cursor_position,
124        _open_buffers: open_buffers,
125    });
126    Ok(())
127}
128
129async fn cursor_position(
130    example: &Example,
131    project: &Entity<Project>,
132    open_buffers: &OpenedBuffers,
133    cx: &mut AsyncApp,
134) -> Result<(Entity<Buffer>, Anchor)> {
135    let language_registry = project.read_with(cx, |project, _| project.languages().clone());
136    let result = language_registry
137        .load_language_for_file_path(&example.spec.cursor_path)
138        .await;
139
140    if let Err(error) = result
141        && !error.is::<LanguageNotFound>()
142    {
143        return Err(error);
144    }
145
146    let cursor_path_str = example.spec.cursor_path.to_string_lossy();
147    // Also try cursor path with first component stripped - old examples may have
148    // paths like "zed/crates/foo.rs" instead of "crates/foo.rs".
149    let cursor_path_without_prefix: PathBuf =
150        example.spec.cursor_path.components().skip(1).collect();
151    let cursor_path_without_prefix_str = cursor_path_without_prefix.to_string_lossy();
152
153    // We try open_buffers first because the file might be new and not saved to disk
154    let cursor_buffer = if let Some(buffer) = open_buffers.get(cursor_path_str.as_ref()) {
155        buffer.clone()
156    } else if let Some(buffer) = open_buffers.get(cursor_path_without_prefix_str.as_ref()) {
157        buffer.clone()
158    } else {
159        // Since the worktree scanner is disabled, manually refresh entries for the cursor path.
160        if let Some(worktree) = project.read_with(cx, |project, cx| project.worktrees(cx).next()) {
161            refresh_worktree_entries(&worktree, [&*example.spec.cursor_path], cx).await?;
162        }
163
164        let cursor_path = project
165            .read_with(cx, |project, cx| {
166                project
167                    .find_project_path(&example.spec.cursor_path, cx)
168                    .or_else(|| project.find_project_path(&cursor_path_without_prefix, cx))
169            })
170            .with_context(|| {
171                format!(
172                    "failed to find cursor path {}",
173                    example.spec.cursor_path.display()
174                )
175            })?;
176
177        project
178            .update(cx, |project, cx| project.open_buffer(cursor_path, cx))
179            .await?
180    };
181
182    let (cursor_excerpt, cursor_offset_within_excerpt) = example.spec.cursor_excerpt()?;
183
184    let excerpt_offset = cursor_buffer.read_with(&*cx, |buffer, _cx| {
185        let text = buffer.text();
186
187        let mut matches = text.match_indices(&cursor_excerpt);
188        let (excerpt_offset, _) = matches.next().with_context(|| {
189            format!("Cursor excerpt did not exist in buffer:\n\n{cursor_excerpt}\n",)
190        })?;
191        anyhow::ensure!(
192            matches.next().is_none(),
193            "More than one cursor position match found",
194        );
195        Ok(excerpt_offset)
196    })?;
197
198    let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
199    let cursor_anchor =
200        cursor_buffer.read_with(&*cx, |buffer, _| buffer.anchor_after(cursor_offset));
201
202    Ok((cursor_buffer, cursor_anchor))
203}
204
205async fn setup_project(
206    example: &mut Example,
207    app_state: &Arc<EpAppState>,
208    step_progress: &StepProgress,
209    cx: &mut AsyncApp,
210) -> Result<Entity<Project>> {
211    let ep_store = cx
212        .update(|cx| EditPredictionStore::try_global(cx))
213        .context("Store should be initialized at init")?;
214
215    let worktree_path = setup_worktree(example, step_progress).await?;
216
217    let project = cx.update(|cx| {
218        Project::local(
219            app_state.client.clone(),
220            app_state.node_runtime.clone(),
221            app_state.user_store.clone(),
222            app_state.languages.clone(),
223            app_state.fs.clone(),
224            None,
225            project::LocalProjectFlags {
226                init_worktree_trust: false,
227                watch_global_configs: false,
228            },
229            cx,
230        )
231    });
232
233    project
234        .update(cx, |project, cx| {
235            project.disable_worktree_scanner(cx);
236            project.create_worktree(&worktree_path, true, cx)
237        })
238        .await?;
239
240    let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone());
241    cx.subscribe(&buffer_store, {
242        let project = project.downgrade();
243        let ep_store = ep_store.downgrade();
244        move |_, event, cx| match event {
245            BufferStoreEvent::BufferAdded(buffer) => {
246                let Some(project) = project.upgrade() else {
247                    return;
248                };
249                ep_store
250                    .update(cx, |store, cx| store.register_buffer(&buffer, &project, cx))
251                    .ok();
252            }
253            _ => {}
254        }
255    })
256    .detach();
257
258    Ok(project)
259}
260
261async fn setup_worktree(example: &Example, step_progress: &StepProgress) -> Result<PathBuf> {
262    let repo_name = example.repo_name().context("failed to get repo name")?;
263    let repo_dir = git::repo_path_for_url(&example.spec.repository_url)?;
264    let worktree_path = repo_name.worktree_path();
265    let repo_lock = git::lock_repo(&repo_dir).await;
266
267    // Clean up any stale git lock files from previous crashed runs.
268    // Safe-ish since we have our own lock.
269    // WARNING: Can corrupt worktrees if multiple processes of the CLI are running.
270    let worktree_git_dir = repo_dir
271        .join(".git/worktrees")
272        .join(repo_name.name.as_ref());
273    for lock_file in &["index.lock", "HEAD.lock", "config.lock"] {
274        let worktree_lock_path = worktree_git_dir.join(lock_file);
275        let repo_lock_path = repo_dir.join(".git").join(lock_file);
276        if worktree_lock_path.exists() {
277            fs::remove_file(&worktree_lock_path).ok();
278        }
279        if repo_lock_path.exists() {
280            fs::remove_file(&repo_lock_path).ok();
281        }
282    }
283
284    let mut git_repo_exists = false;
285    if repo_dir.is_dir() {
286        if git::run_git(&repo_dir, &["remote", "get-url", "origin"])
287            .await
288            .map_or(false, |origin| origin.trim() == example.spec.repository_url)
289        {
290            git_repo_exists = true;
291        } else {
292            fs::remove_dir_all(&repo_dir).ok();
293        }
294    }
295
296    if !git_repo_exists {
297        step_progress.set_substatus(format!("cloning {}", repo_name.name));
298        fs::create_dir_all(&repo_dir)?;
299        git::run_git(&repo_dir, &["init"]).await?;
300        git::run_git(
301            &repo_dir,
302            &["remote", "add", "origin", &example.spec.repository_url],
303        )
304        .await?;
305    }
306
307    // Resolve the example to a revision, fetching it if needed.
308    step_progress.set_substatus("fetching");
309    let revision = git::fetch_if_needed(&repo_dir, &example.spec.revision).await?;
310
311    // Clean up any stale worktree registrations from previous crashed runs.
312    git::run_git(&repo_dir, &["worktree", "prune"]).await.ok();
313
314    // Create the worktree for this example if needed.
315    step_progress.set_substatus("preparing worktree");
316
317    // Check if worktree exists and is valid (not just a directory from a crashed run).
318    let worktree_valid = worktree_path.is_dir()
319        && git::run_git(&worktree_path, &["rev-parse", "--git-dir"])
320            .await
321            .is_ok();
322
323    if worktree_valid {
324        git::run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
325        git::run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
326        git::run_git(&worktree_path, &["checkout", revision.as_str()]).await?;
327    } else {
328        let worktree_path_string = worktree_path.to_string_lossy();
329
330        // Clean up invalid worktree directory and registration if they exist.
331        if worktree_path.exists() {
332            fs::remove_dir_all(&worktree_path).ok();
333        }
334        git::run_git(
335            &repo_dir,
336            &["worktree", "remove", "--force", &worktree_path_string],
337        )
338        .await
339        .ok();
340
341        let branch_name = example.spec.filename();
342        git::run_git(
343            &repo_dir,
344            &["branch", "-f", &branch_name, revision.as_str()],
345        )
346        .await?;
347        git::run_git(
348            &repo_dir,
349            &["worktree", "add", "-f", &worktree_path_string, &branch_name],
350        )
351        .await?;
352    }
353    drop(repo_lock);
354
355    if !example.spec.uncommitted_diff.is_empty() {
356        step_progress.set_substatus("applying diff");
357
358        // old examples had full paths in the uncommitted diff.
359        let uncommitted_diff =
360            strip_diff_path_prefix(&example.spec.uncommitted_diff, &repo_name.name);
361
362        let mut apply_process = smol::process::Command::new("git")
363            .current_dir(&worktree_path)
364            .args(&["apply", "-"])
365            .stdin(std::process::Stdio::piped())
366            .spawn()?;
367
368        let mut stdin = apply_process.stdin.take().context("Failed to get stdin")?;
369        stdin.write_all(uncommitted_diff.as_bytes()).await?;
370        stdin.close().await?;
371        drop(stdin);
372
373        let apply_result = apply_process.output().await?;
374        anyhow::ensure!(
375            apply_result.status.success(),
376            "Failed to apply uncommitted diff patch with status: {}\nstderr:\n{}\nstdout:\n{}",
377            apply_result.status,
378            String::from_utf8_lossy(&apply_result.stderr),
379            String::from_utf8_lossy(&apply_result.stdout),
380        );
381    }
382
383    step_progress.clear_substatus();
384    Ok(worktree_path)
385}
386
387async fn apply_edit_history(
388    example: &Example,
389    project: &Entity<Project>,
390    cx: &mut AsyncApp,
391) -> Result<OpenedBuffers> {
392    edit_prediction::udiff::apply_diff(&example.spec.edit_history, project, cx).await
393}
394