tenant.openagents/omega

1use std::{collections::BTreeMap, ops::Range, path::Path, sync::Arc};
2
3use anyhow::Result;
4use futures::{StreamExt, TryFutureExt, TryStreamExt, stream::FuturesUnordered};
5use gpui::{App, AppContext, Context, Entity, Subscription, Task};
6use itertools::Itertools;
7use language::{Buffer, BufferEvent};
8use std::collections::HashMap;
9use text::{BufferSnapshot, Point};
10
11use crate::{ProjectPath, buffer_store::BufferStore, worktree_store::WorktreeStore};
12
13#[derive(Clone, Debug)]
14pub struct Bookmark {
15    pub anchor: text::Anchor,
16    pub label: String,
17}
18
19#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
20pub struct SerializedBookmark {
21    pub row: u32,
22    pub label: String,
23}
24
25#[derive(Debug)]
26pub struct BufferBookmarks {
27    buffer: Entity<Buffer>,
28    bookmarks: Vec<Bookmark>,
29    _subscription: Subscription,
30}
31
32impl BufferBookmarks {
33    pub fn new(buffer: Entity<Buffer>, cx: &mut Context<BookmarkStore>) -> Self {
34        let subscription = cx.subscribe(
35            &buffer,
36            |bookmark_store, buffer, event: &BufferEvent, cx| match event {
37                BufferEvent::FileHandleChanged => {
38                    bookmark_store.handle_file_changed(buffer, cx);
39                }
40                _ => {}
41            },
42        );
43
44        Self {
45            buffer,
46            bookmarks: Vec::new(),
47            _subscription: subscription,
48        }
49    }
50
51    pub fn buffer(&self) -> &Entity<Buffer> {
52        &self.buffer
53    }
54
55    pub fn bookmarks(&self) -> &[Bookmark] {
56        &self.bookmarks
57    }
58}
59
60#[derive(Debug)]
61pub enum BookmarkEntry {
62    Loaded(BufferBookmarks),
63    Unloaded(Vec<SerializedBookmark>),
64}
65
66impl BookmarkEntry {
67    pub fn is_empty(&self) -> bool {
68        match self {
69            BookmarkEntry::Loaded(buffer_bookmarks) => buffer_bookmarks.bookmarks.is_empty(),
70            BookmarkEntry::Unloaded(rows) => rows.is_empty(),
71        }
72    }
73
74    fn loaded(&self) -> Option<&BufferBookmarks> {
75        match self {
76            BookmarkEntry::Loaded(buffer_bookmarks) => Some(buffer_bookmarks),
77            BookmarkEntry::Unloaded(_) => None,
78        }
79    }
80}
81
82pub struct BookmarkStore {
83    buffer_store: Entity<BufferStore>,
84    worktree_store: Entity<WorktreeStore>,
85    bookmarks: BTreeMap<Arc<Path>, BookmarkEntry>,
86}
87
88impl BookmarkStore {
89    pub fn new(worktree_store: Entity<WorktreeStore>, buffer_store: Entity<BufferStore>) -> Self {
90        Self {
91            buffer_store,
92            worktree_store,
93            bookmarks: BTreeMap::new(),
94        }
95    }
96
97    pub fn load_serialized_bookmarks(
98        &mut self,
99        bookmark_rows: BTreeMap<Arc<Path>, Vec<SerializedBookmark>>,
100        cx: &mut Context<Self>,
101    ) -> Task<Result<()>> {
102        self.bookmarks.clear();
103
104        for (path, rows) in bookmark_rows {
105            if rows.is_empty() {
106                continue;
107            }
108
109            let count = rows.len();
110            log::debug!("Stored {count} unloaded bookmark(s) at {}", path.display());
111
112            self.bookmarks.insert(path, BookmarkEntry::Unloaded(rows));
113        }
114
115        cx.notify();
116        Task::ready(Ok(()))
117    }
118
119    fn resolve_anchors_if_needed(
120        &mut self,
121        abs_path: &Arc<Path>,
122        buffer: &Entity<Buffer>,
123        cx: &mut Context<Self>,
124    ) {
125        let Some(BookmarkEntry::Unloaded(bookmarks)) = self.bookmarks.get(abs_path) else {
126            return;
127        };
128
129        let snapshot = buffer.read(cx).snapshot();
130        let max_point = snapshot.max_point();
131
132        let bookmarks: Vec<Bookmark> = bookmarks
133            .iter()
134            .filter_map(|bookmark| {
135                let point = Point::new(bookmark.row, 0);
136
137                if point > max_point {
138                    log::warn!(
139                        "Skipping out-of-range bookmark: {} row {} (file has {} rows)",
140                        abs_path.display(),
141                        bookmark.row,
142                        max_point.row
143                    );
144                    return None;
145                }
146
147                let anchor = snapshot.anchor_after(point);
148                Some(Bookmark {
149                    anchor,
150                    label: bookmark.label.clone(),
151                })
152            })
153            .collect();
154
155        if bookmarks.is_empty() {
156            self.bookmarks.remove(abs_path);
157        } else {
158            let mut buffer_bookmarks = BufferBookmarks::new(buffer.clone(), cx);
159            buffer_bookmarks.bookmarks = bookmarks;
160            self.bookmarks
161                .insert(abs_path.clone(), BookmarkEntry::Loaded(buffer_bookmarks));
162        }
163    }
164
165    pub fn abs_path_from_buffer(buffer: &Entity<Buffer>, cx: &App) -> Option<Arc<Path>> {
166        worktree::File::from_dyn(buffer.read(cx).file())
167            .map(|file| file.worktree.read(cx).absolutize(&file.path))
168            .map(Arc::<Path>::from)
169    }
170
171    /// Toggle a bookmark at the given anchor in the buffer.
172    /// If a bookmark already exists on the same row, it will be removed.
173    /// Otherwise, a new bookmark will be added with the given label.
174    pub fn toggle_bookmark(
175        &mut self,
176        buffer: Entity<Buffer>,
177        anchor: text::Anchor,
178        label: String,
179        cx: &mut Context<Self>,
180    ) {
181        let Some(abs_path) = Self::abs_path_from_buffer(&buffer, cx) else {
182            return;
183        };
184
185        self.resolve_anchors_if_needed(&abs_path, &buffer, cx);
186
187        let entry = self
188            .bookmarks
189            .entry(abs_path.clone())
190            .or_insert_with(|| BookmarkEntry::Loaded(BufferBookmarks::new(buffer.clone(), cx)));
191
192        let BookmarkEntry::Loaded(buffer_bookmarks) = entry else {
193            unreachable!("resolve_if_needed should have converted to Loaded");
194        };
195
196        let snapshot = buffer.read(cx).text_snapshot();
197
198        let existing_index = buffer_bookmarks.bookmarks.iter().position(|existing| {
199            existing.anchor.summary::<Point>(&snapshot).row
200                == anchor.summary::<Point>(&snapshot).row
201        });
202
203        if let Some(index) = existing_index {
204            buffer_bookmarks.bookmarks.remove(index);
205            if buffer_bookmarks.bookmarks.is_empty() {
206                self.bookmarks.remove(&abs_path);
207            }
208        } else {
209            buffer_bookmarks.bookmarks.push(Bookmark { anchor, label });
210        }
211
212        cx.notify();
213    }
214
215    pub fn find_bookmark(
216        &mut self,
217        buffer: &Entity<Buffer>,
218        anchor: text::Anchor,
219        cx: &mut Context<Self>,
220    ) -> Option<&Bookmark> {
221        let Some(abs_path) = Self::abs_path_from_buffer(buffer, cx) else {
222            return None;
223        };
224
225        self.resolve_anchors_if_needed(&abs_path, buffer, cx);
226
227        let entry = self
228            .bookmarks
229            .entry(abs_path.clone())
230            .or_insert_with(|| BookmarkEntry::Loaded(BufferBookmarks::new(buffer.clone(), cx)));
231
232        let BookmarkEntry::Loaded(buffer_bookmarks) = entry else {
233            unreachable!("resolve_if_needed should have converted to Loaded");
234        };
235
236        let snapshot = buffer.read(cx).text_snapshot();
237
238        buffer_bookmarks.bookmarks.iter().find(|existing| {
239            existing.anchor.summary::<Point>(&snapshot).row
240                == anchor.summary::<Point>(&snapshot).row
241        })
242    }
243
244    pub fn edit_bookmark(
245        &mut self,
246        buffer: &Entity<Buffer>,
247        anchor: text::Anchor,
248        label: String,
249        cx: &mut Context<Self>,
250    ) {
251        let Some(abs_path) = Self::abs_path_from_buffer(buffer, cx) else {
252            return;
253        };
254
255        self.resolve_anchors_if_needed(&abs_path, buffer, cx);
256
257        let Some(BookmarkEntry::Loaded(buffer_bookmarks)) = self.bookmarks.get_mut(&abs_path)
258        else {
259            return;
260        };
261
262        let snapshot = buffer.read(cx).text_snapshot();
263        let row = anchor.summary::<Point>(&snapshot).row;
264
265        if let Some(bookmark) = buffer_bookmarks
266            .bookmarks
267            .iter_mut()
268            .find(|existing| existing.anchor.summary::<Point>(&snapshot).row == row)
269        {
270            bookmark.label = label;
271            cx.notify();
272        }
273    }
274
275    /// Returns the bookmarks for a given buffer within an optional range.
276    /// Only returns bookmarks that have been resolved to anchors (loaded).
277    /// Unloaded bookmarks for the given buffer will be resolved first.
278    pub fn bookmarks_for_buffer(
279        &mut self,
280        buffer: Entity<Buffer>,
281        range: Range<text::Anchor>,
282        buffer_snapshot: &BufferSnapshot,
283        cx: &mut Context<Self>,
284    ) -> Vec<Bookmark> {
285        let Some(abs_path) = Self::abs_path_from_buffer(&buffer, cx) else {
286            return Vec::new();
287        };
288
289        self.resolve_anchors_if_needed(&abs_path, &buffer, cx);
290
291        let Some(BookmarkEntry::Loaded(file_bookmarks)) = self.bookmarks.get(&abs_path) else {
292            return Vec::new();
293        };
294
295        file_bookmarks
296            .bookmarks
297            .iter()
298            .filter_map({
299                move |bookmark| {
300                    if !buffer_snapshot.can_resolve(&bookmark.anchor) {
301                        return None;
302                    }
303
304                    if bookmark.anchor.cmp(&range.start, buffer_snapshot).is_lt()
305                        || bookmark.anchor.cmp(&range.end, buffer_snapshot).is_gt()
306                    {
307                        return None;
308                    }
309
310                    Some(bookmark.clone())
311                }
312            })
313            .collect()
314    }
315
316    fn handle_file_changed(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
317        let entity_id = buffer.entity_id();
318
319        if buffer
320            .read(cx)
321            .file()
322            .is_none_or(|f| f.disk_state().is_deleted())
323        {
324            self.bookmarks.retain(|_, entry| match entry {
325                BookmarkEntry::Loaded(buffer_bookmarks) => {
326                    buffer_bookmarks.buffer.entity_id() != entity_id
327                }
328                BookmarkEntry::Unloaded(_) => true,
329            });
330            cx.notify();
331            return;
332        }
333
334        if let Some(new_abs_path) = Self::abs_path_from_buffer(&buffer, cx) {
335            if self.bookmarks.contains_key(&new_abs_path) {
336                return;
337            }
338
339            if let Some(old_path) = self
340                .bookmarks
341                .iter()
342                .find(|(_, entry)| match entry {
343                    BookmarkEntry::Loaded(buffer_bookmarks) => {
344                        buffer_bookmarks.buffer.entity_id() == entity_id
345                    }
346                    BookmarkEntry::Unloaded(_) => false,
347                })
348                .map(|(path, _)| path)
349                .cloned()
350            {
351                let Some(entry) = self.bookmarks.remove(&old_path) else {
352                    log::error!(
353                        "Couldn't get bookmarks from old path during buffer rename handling"
354                    );
355                    return;
356                };
357                self.bookmarks.insert(new_abs_path, entry);
358                cx.notify();
359            }
360        }
361    }
362
363    pub fn all_serialized_bookmarks(
364        &self,
365        cx: &App,
366    ) -> BTreeMap<Arc<Path>, Vec<SerializedBookmark>> {
367        self.bookmarks
368            .iter()
369            .filter_map(|(path, entry)| {
370                let mut rows = match entry {
371                    BookmarkEntry::Unloaded(rows) => rows.clone(),
372                    BookmarkEntry::Loaded(buffer_bookmarks) => {
373                        let snapshot = buffer_bookmarks.buffer.read(cx).snapshot();
374                        buffer_bookmarks
375                            .bookmarks
376                            .iter()
377                            .filter_map(|bookmark| {
378                                if !snapshot.can_resolve(&bookmark.anchor) {
379                                    return None;
380                                }
381                                let row =
382                                    snapshot.summary_for_anchor::<Point>(&bookmark.anchor).row;
383                                Some(SerializedBookmark {
384                                    row,
385                                    label: bookmark.label.clone(),
386                                })
387                            })
388                            .collect()
389                    }
390                };
391
392                rows.sort_unstable_by_key(|a| a.row);
393                rows.dedup_by_key(|a| a.row);
394
395                if rows.is_empty() {
396                    None
397                } else {
398                    Some((path.clone(), rows))
399                }
400            })
401            .collect()
402    }
403
404    pub async fn all_bookmark_locations(
405        this: Entity<BookmarkStore>,
406        cx: &mut (impl AppContext + Clone),
407    ) -> Result<HashMap<Entity<Buffer>, Vec<Range<Point>>>> {
408        Self::resolve_all(&this, cx).await?;
409
410        cx.read_entity(&this, |this, cx| {
411            let mut locations: HashMap<_, Vec<_>> = HashMap::new();
412            for bookmarks in this.bookmarks.values().filter_map(BookmarkEntry::loaded) {
413                let snapshot = cx.read_entity(bookmarks.buffer(), |b, _| b.snapshot());
414                let ranges: Vec<Range<Point>> = bookmarks
415                    .bookmarks()
416                    .iter()
417                    .map(|bookmark| {
418                        let row = snapshot.summary_for_anchor::<Point>(&bookmark.anchor).row;
419                        Point::row_range(row..row)
420                    })
421                    .collect();
422
423                locations
424                    .entry(bookmarks.buffer().clone())
425                    .or_default()
426                    .extend(ranges);
427            }
428
429            Ok(locations)
430        })
431    }
432
433    /// Opens buffers for all unloaded bookmark entries and resolves them to anchors. This is used to show all bookmarks in a large multi-buffer.
434    async fn resolve_all(this: &Entity<Self>, cx: &mut (impl AppContext + Clone)) -> Result<()> {
435        let unloaded_paths: Vec<Arc<Path>> = cx.read_entity(&this, |this, _| {
436            this.bookmarks
437                .iter()
438                .filter_map(|(path, entry)| match entry {
439                    BookmarkEntry::Unloaded(_) => Some(path.clone()),
440                    BookmarkEntry::Loaded(_) => None,
441                })
442                .collect_vec()
443        });
444
445        if unloaded_paths.is_empty() {
446            return Ok(());
447        }
448
449        let worktree_store = cx.read_entity(&this, |this, _| this.worktree_store.clone());
450        let buffer_store = cx.read_entity(&this, |this, _| this.buffer_store.clone());
451
452        let open_tasks: FuturesUnordered<_> = unloaded_paths
453            .iter()
454            .map(|path| {
455                open_path(path, &worktree_store, &buffer_store, cx.clone())
456                    .map_err(move |e| (path, e))
457                    .map_ok(move |b| (path, b))
458            })
459            .collect();
460
461        let opened: Vec<_> = open_tasks
462            .inspect_err(|(path, error)| {
463                log::warn!(
464                    "Could not open buffer for bookmarked path {}: {error}",
465                    path.display()
466                )
467            })
468            .filter_map(|res| async move { res.ok() })
469            .collect()
470            .await;
471
472        cx.update_entity(&this, |this, cx| {
473            for (path, buffer) in opened {
474                this.resolve_anchors_if_needed(&path, &buffer, cx);
475            }
476            cx.notify();
477        });
478
479        Ok(())
480    }
481
482    pub fn clear_bookmarks(&mut self, cx: &mut Context<Self>) {
483        self.bookmarks.clear();
484        cx.notify();
485    }
486}
487
488async fn open_path(
489    path: &Path,
490    worktree_store: &Entity<WorktreeStore>,
491    buffer_store: &Entity<BufferStore>,
492    mut cx: impl AppContext,
493) -> Result<Entity<Buffer>> {
494    let (worktree, worktree_path) = cx
495        .update_entity(&worktree_store, |worktree_store, cx| {
496            worktree_store.find_or_create_worktree(path, false, cx)
497        })
498        .await?;
499
500    let project_path = ProjectPath {
501        worktree_id: cx.read_entity(&worktree, |worktree, _| worktree.id()),
502        path: worktree_path,
503    };
504
505    let buffer = cx
506        .update_entity(&buffer_store, |buffer_store, cx| {
507            buffer_store.open_buffer(project_path, cx)
508        })
509        .await?;
510
511    Ok(buffer)
512}
513