tenant.openagents/omega

1use crate::{EditPredictionStore, StoredEvent, zeta};
2
3use anyhow::{Context as _, Result};
4use buffer_diff::BufferDiffSnapshot;
5use collections::HashMap;
6use gpui::{Context, Entity, Task};
7use language::{Buffer, BufferSnapshot, ToPoint as _};
8use project::{Project, ProjectPath, WorktreeId};
9use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc};
10use text::{OffsetRangeExt as _, Point};
11use util::rel_path::RelPath;
12
13const MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024;
14
15pub type UncommittedDiffSnapshot = Vec<(Arc<Path>, BufferSnapshot, BufferDiffSnapshot)>;
16pub type UncommittedDiffResult = std::result::Result<UncommittedDiffSnapshot, Arc<anyhow::Error>>;
17
18pub use zeta_prompt::udiff::CURSOR_POSITION_MARKER;
19
20pub(crate) struct CapturedPredictionContext {
21    pub(crate) repository_url: Option<String>,
22    pub(crate) revision: Option<String>,
23    pub(crate) uncommitted_diff: Option<String>,
24    pub(crate) buffer_diagnostics: Vec<zeta_prompt::ActiveBufferDiagnostic>,
25    pub(crate) editable_context: Vec<zeta_prompt::RelatedFile>,
26}
27
28pub(crate) fn capture_prediction_context(
29    project: Entity<Project>,
30    buffer: Entity<Buffer>,
31    cursor_anchor: language::Anchor,
32    stored_events: Vec<StoredEvent>,
33    repository_url: Option<String>,
34    revision: Option<String>,
35    editable_context_task: Task<Result<Vec<zeta_prompt::RelatedFile>>>,
36    cx: &mut Context<EditPredictionStore>,
37) -> Option<Task<Result<CapturedPredictionContext>>> {
38    let snapshot = buffer.read(cx).snapshot();
39    let worktree_id = snapshot.file()?.worktree_id(cx);
40    let uncommitted_diff_task =
41        uncommitted_diffs_for_events(project, worktree_id, stored_events, cx);
42
43    Some(cx.spawn(async move |_this, cx| {
44        let uncommitted_diff_snapshot = match uncommitted_diff_task.await {
45            Ok(snapshot) => Some(snapshot),
46            Err(error) => {
47                log::debug!("failed to capture uncommitted diff: {error:?}");
48                None
49            }
50        };
51
52        let uncommitted_diff = if let Some(uncommitted_diff_snapshot) = uncommitted_diff_snapshot {
53            let estimated_uncommitted_diff_size = uncommitted_diff_snapshot
54                .iter()
55                .map(|(_, buffer_snapshot, diff_snapshot)| {
56                    diff_snapshot
57                        .hunks(buffer_snapshot)
58                        .map(|hunk| {
59                            hunk.diff_base_byte_range.len()
60                                + hunk.range.to_offset(buffer_snapshot).len()
61                        })
62                        .sum::<usize>()
63                })
64                .sum::<usize>();
65
66            if estimated_uncommitted_diff_size <= MAX_UNCOMMITTED_DIFF_SIZE {
67                let uncommitted_diff = cx
68                    .background_executor()
69                    .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) })
70                    .await;
71                (uncommitted_diff.len() <= MAX_UNCOMMITTED_DIFF_SIZE).then_some(uncommitted_diff)
72            } else {
73                None
74            }
75        } else {
76            None
77        };
78
79        let buffer_diagnostics = zeta::active_buffer_diagnostics(
80            &snapshot,
81            Point::new(0, 0)..snapshot.max_point(),
82            cursor_anchor.to_point(&snapshot).row,
83            100,
84        );
85        let editable_context = match editable_context_task.await {
86            Ok(editable_context) => editable_context,
87            Err(error) => {
88                log::debug!("failed to capture editable context: {error:?}");
89                Vec::new()
90            }
91        };
92
93        Ok(CapturedPredictionContext {
94            repository_url,
95            revision,
96            uncommitted_diff,
97            buffer_diagnostics,
98            editable_context,
99        })
100    }))
101}
102
103pub fn uncommitted_diffs_for_events(
104    project: Entity<Project>,
105    worktree_id: WorktreeId,
106    events: Vec<StoredEvent>,
107    cx: &Context<'_, EditPredictionStore>,
108) -> Task<UncommittedDiffResult> {
109    let git_store = project.read_with(cx, |project, _| project.git_store().clone());
110
111    cx.spawn(async move |_store, cx| {
112        let (worktree_root_name, worktree_abs_path, path_style) = project
113            .read_with(cx, |project, cx| {
114                let worktree = project.worktree_for_id(worktree_id, cx)?;
115                let worktree = worktree.read(cx);
116                let path_style = worktree.path_style();
117                let root_name = RelPath::new(Path::new(worktree.root_name_str()), path_style)
118                    .ok()?
119                    .into_owned();
120                Some((root_name, worktree.abs_path(), path_style))
121            })
122            .context("failed to find worktree for uncommitted diff capture")
123            .map_err(Arc::new)?;
124
125        let events_with_paths = events
126            .into_iter()
127            .filter_map(|stored_event| {
128                let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref();
129                let path = if let Ok(path) = RelPath::new(path, path_style) {
130                    path.strip_prefix(&worktree_root_name).ok()?.into_arc()
131                } else {
132                    let path = path.strip_prefix(worktree_abs_path.as_ref()).ok()?;
133                    RelPath::new(path, path_style).ok()?.into_arc()
134                };
135                let project_path = ProjectPath { worktree_id, path };
136                let relative_path: Arc<Path> = project_path.path.as_std_path().into();
137                Some((stored_event, project_path, relative_path))
138            })
139            .collect::<Vec<_>>();
140
141        let mut snapshots_by_path: HashMap<Arc<Path>, (BufferSnapshot, BufferDiffSnapshot)> =
142            HashMap::default();
143        for (stored_event, project_path, relative_path) in events_with_paths.iter().rev() {
144            if snapshots_by_path.contains_key(relative_path) {
145                continue;
146            }
147
148            let buffer = project
149                .update(cx, |project, cx| {
150                    project.open_buffer(project_path.clone(), cx)
151                })
152                .await
153                .context("failed to open buffer for uncommitted diff capture")
154                .map_err(Arc::new)?;
155            let buffer_id = buffer.read_with(cx, |buffer, _| buffer.remote_id());
156            let file_context = stored_event.file_context.clone();
157            let cached_diff = file_context
158                .as_ref()
159                .and_then(|file_context| {
160                    file_context
161                        .read_with(cx, |file_context, _| file_context.uncommitted_diff.clone())
162                })
163                // The cached diff is keyed by path, but its hunk anchors are pinned to a
164                // specific buffer. If that buffer was closed and reopened, `open_buffer`
165                // hands back a buffer with a new `BufferId`; reusing the stale diff against
166                // it would mix anchors from different buffers and panic. Drop the cache in
167                // that case so the diff is recomputed for the current buffer.
168                .filter(|diff| diff.read_with(cx, |diff, _| diff.buffer_id) == buffer_id);
169            let diff = match cached_diff {
170                Some(diff) => diff,
171                None => {
172                    let diff = git_store
173                        .update(cx, |git_store, cx| {
174                            git_store.open_uncommitted_diff(buffer.clone(), cx)
175                        })
176                        .await
177                        .context("failed to open uncommitted diff for capture")
178                        .map_err(Arc::new)?;
179                    if let Some(file_context) = file_context {
180                        file_context.update(cx, |file_context, _| {
181                            file_context.uncommitted_diff = Some(diff.clone());
182                        });
183                    }
184                    diff
185                }
186            };
187
188            let buffer_snapshot = buffer.update(cx, |buffer, _| buffer.snapshot());
189            let diff_snapshot = diff.update(cx, |diff, cx| diff.snapshot(cx));
190            snapshots_by_path.insert(relative_path.clone(), (buffer_snapshot, diff_snapshot));
191        }
192
193        let uncommitted_diff_snapshots = snapshots_by_path
194            .into_iter()
195            .map(|(relative_path, (snapshot, diff_snapshot))| {
196                (relative_path, snapshot, diff_snapshot)
197            })
198            .collect();
199
200        Ok(uncommitted_diff_snapshots)
201    })
202}
203
204pub fn compute_cursor_excerpt(
205    snapshot: &language::BufferSnapshot,
206    cursor_anchor: language::Anchor,
207) -> (String, usize, Range<Point>) {
208    use text::ToOffset as _;
209    use text::ToPoint as _;
210
211    let cursor_offset = cursor_anchor.to_offset(snapshot);
212    let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) =
213        crate::cursor_excerpt::compute_cursor_excerpt(snapshot, cursor_offset);
214    let syntax_ranges = crate::cursor_excerpt::compute_syntax_ranges(
215        snapshot,
216        cursor_offset,
217        &excerpt_offset_range,
218    );
219    let excerpt_text: String = snapshot.text_for_range(excerpt_point_range).collect();
220    let (_, context_range) = zeta_prompt::compute_editable_and_context_ranges(
221        &excerpt_text,
222        cursor_offset_in_excerpt,
223        &syntax_ranges,
224        100,
225        50,
226    );
227    let context_text = excerpt_text[context_range.clone()].to_string();
228    let cursor_in_context = cursor_offset_in_excerpt.saturating_sub(context_range.start);
229    let context_buffer_start =
230        (excerpt_offset_range.start + context_range.start).to_point(snapshot);
231    let context_buffer_end = (excerpt_offset_range.start + context_range.end).to_point(snapshot);
232    (
233        context_text,
234        cursor_in_context,
235        context_buffer_start..context_buffer_end,
236    )
237}
238
239pub(crate) fn compute_uncommitted_diff(snapshot: UncommittedDiffSnapshot) -> String {
240    let mut uncommitted_diff = String::new();
241    let mut snapshots_by_path = snapshot;
242    snapshots_by_path.sort_by(|(left_path, _, _), (right_path, _, _)| left_path.cmp(right_path));
243    for (relative_path, buffer_snapshot, diff_snapshot) in snapshots_by_path {
244        let base_snapshot = diff_snapshot.base_text();
245        let is_existing_file = diff_snapshot.base_text_exists();
246
247        let new_path_str = relative_path.to_string_lossy();
248        let old_path_str = if is_existing_file {
249            new_path_str.as_ref()
250        } else {
251            "/dev/null"
252        };
253        writeln!(
254            uncommitted_diff,
255            "--- {}{old_path_str}",
256            if is_existing_file { "a/" } else { "" }
257        )
258        .ok();
259        writeln!(uncommitted_diff, "+++ b/{new_path_str}").ok();
260
261        if !is_existing_file {
262            let new_text = buffer_snapshot.text();
263            writeln!(
264                uncommitted_diff,
265                "@@ -0,0 +1,{} @@",
266                new_text.lines().count()
267            )
268            .ok();
269            for line in new_text.lines() {
270                writeln!(uncommitted_diff, "+{line}").ok();
271            }
272            continue;
273        }
274
275        let mut ranges: Vec<(Range<u32>, Range<u32>)> = Vec::new();
276        for hunk in (&diff_snapshot).hunks(&buffer_snapshot) {
277            let old_start = base_snapshot
278                .offset_to_point(hunk.diff_base_byte_range.start)
279                .row;
280            let old_end =
281                exclusive_end_row(base_snapshot.offset_to_point(hunk.diff_base_byte_range.end));
282            let new_start = hunk.range.start.row;
283            let new_end = exclusive_end_row(hunk.range.end);
284            let old_range = old_start.saturating_sub(3)..old_end + 3;
285            let new_range = new_start.saturating_sub(3)..new_end + 3;
286
287            if let Some((last_old_range, last_new_range)) = ranges.last_mut()
288                && (old_range.start <= last_old_range.end || new_range.start <= last_new_range.end)
289            {
290                last_old_range.end = last_old_range.end.max(old_range.end);
291                last_new_range.end = last_new_range.end.max(new_range.end);
292                continue;
293            }
294            ranges.push((old_range, new_range));
295        }
296
297        for (old_range, new_range) in ranges {
298            uncommitted_diff.push_str(&language::unified_diff_with_offsets(
299                &base_snapshot
300                    .text_for_range(
301                        Point::new(old_range.start, 0)
302                            ..row_start_or_max(base_snapshot, old_range.end),
303                    )
304                    .collect::<String>(),
305                &buffer_snapshot
306                    .text_for_range(
307                        Point::new(new_range.start, 0)
308                            ..row_start_or_max(&buffer_snapshot, new_range.end),
309                    )
310                    .collect::<String>(),
311                old_range.start,
312                new_range.start,
313            ));
314        }
315        if !uncommitted_diff.ends_with('\n') {
316            uncommitted_diff.push('\n');
317        }
318    }
319    uncommitted_diff
320}
321
322fn row_start_or_max(snapshot: &language::BufferSnapshot, row: u32) -> Point {
323    if row >= snapshot.max_point().row {
324        snapshot.max_point()
325    } else {
326        Point::new(row, 0)
327    }
328}
329
330fn exclusive_end_row(point: Point) -> u32 {
331    if point.column == 0 {
332        point.row
333    } else {
334        point.row + 1
335    }
336}
337
338pub fn format_cursor_excerpt(
339    excerpt: &str,
340    cursor_offset: usize,
341    line_comment_prefix: &str,
342) -> String {
343    let cursor_line_start = excerpt[..cursor_offset]
344        .rfind('\n')
345        .map(|pos| pos + 1)
346        .unwrap_or(0);
347    let cursor_line_end = excerpt[cursor_line_start..]
348        .find('\n')
349        .map(|pos| cursor_line_start + pos + 1)
350        .unwrap_or(excerpt.len());
351    let cursor_line = &excerpt[cursor_line_start..cursor_line_end];
352    let cursor_line_indent = &cursor_line[..cursor_line.len() - cursor_line.trim_start().len()];
353    let cursor_column = cursor_offset - cursor_line_start;
354
355    let mut marker_line = String::new();
356    if cursor_column < line_comment_prefix.len() {
357        for _ in 0..cursor_column {
358            marker_line.push(' ');
359        }
360        marker_line.push_str(line_comment_prefix);
361        write!(marker_line, " <{}", CURSOR_POSITION_MARKER).unwrap();
362    } else {
363        if cursor_column >= cursor_line_indent.len() + line_comment_prefix.len() {
364            marker_line.push_str(cursor_line_indent);
365        }
366        marker_line.push_str(line_comment_prefix);
367        while marker_line.len() < cursor_column {
368            marker_line.push(' ');
369        }
370        write!(marker_line, "^{}", CURSOR_POSITION_MARKER).unwrap();
371    }
372
373    let mut result = String::with_capacity(excerpt.len() + marker_line.len() + 2);
374    result.push_str(&excerpt[..cursor_line_end]);
375    if !result.ends_with('\n') {
376        result.push('\n');
377    }
378    result.push_str(&marker_line);
379    if cursor_line_end < excerpt.len() {
380        result.push('\n');
381        result.push_str(&excerpt[cursor_line_end..]);
382    }
383    result
384}
385