tenant.openagents/omega

1use anyhow::Result;
2use buffer_diff::BufferDiff;
3use collections::HashSet;
4use futures::StreamExt;
5use git::{
6    repository::RepoPath,
7    status::{DiffTreeType, FileStatus, StatusCode, TrackedStatus, TreeDiff, TreeDiffStatus},
8};
9use gpui::{
10    App, AsyncApp, AsyncWindowContext, Context, Entity, EventEmitter, SharedString, Subscription,
11    Task, WeakEntity, Window,
12};
13
14use language::Buffer;
15use text::BufferId;
16use util::ResultExt;
17use ztracing::instrument;
18
19use crate::{
20    ConflictSet, Project,
21    git_store::{GitStoreEvent, Repository, RepositoryEvent},
22};
23
24#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
25pub enum DiffBase {
26    Head,
27    Index,
28    Staged,
29    Merge { base_ref: SharedString },
30}
31
32impl DiffBase {
33    pub fn is_merge_base(&self) -> bool {
34        matches!(self, DiffBase::Merge { .. })
35    }
36}
37
38pub struct DiffBufferList {
39    diff_base: DiffBase,
40    repo: Option<Entity<Repository>>,
41    project: Entity<Project>,
42    base_commit: Option<SharedString>,
43    head_commit: Option<SharedString>,
44    tree_diff: Option<TreeDiff>,
45    tree_diff_update_needed: bool,
46    tree_diff_base_task: Option<Task<()>>,
47    _subscription: Subscription,
48    update_needed: postage::watch::Sender<()>,
49    _task: Task<()>,
50}
51
52pub enum BranchDiffEvent {
53    FileListChanged,
54    DiffBaseChanged,
55}
56
57impl EventEmitter<BranchDiffEvent> for DiffBufferList {}
58
59impl DiffBufferList {
60    pub fn new(
61        source: DiffBase,
62        project: Entity<Project>,
63        window: &mut Window,
64        cx: &mut Context<Self>,
65    ) -> Self {
66        let git_store = project.read(cx).git_store().clone();
67        let repo = git_store.read(cx).active_repository();
68        let git_store_subscription = cx.subscribe_in(
69            &git_store,
70            window,
71            move |this, _git_store, event, _window, cx| {
72                let should_update = match event {
73                    GitStoreEvent::ActiveRepositoryChanged(new_repo_id) => {
74                        this.repo.is_none() && new_repo_id.is_some()
75                    }
76                    GitStoreEvent::RepositoryUpdated(
77                        event_repo_id,
78                        RepositoryEvent::StatusesChanged | RepositoryEvent::HeadChanged,
79                        _,
80                    ) => this
81                        .repo
82                        .as_ref()
83                        .is_some_and(|r| r.read(cx).snapshot().id == *event_repo_id),
84                    _ => false,
85                };
86
87                if should_update {
88                    cx.emit(BranchDiffEvent::FileListChanged);
89                    *this.update_needed.borrow_mut() = ();
90                }
91            },
92        );
93
94        let (send, recv) = postage::watch::channel::<()>();
95        let worker = window.spawn(cx, {
96            let this = cx.weak_entity();
97            async |cx| Self::handle_status_updates(this, recv, cx).await
98        });
99
100        Self {
101            diff_base: source,
102            repo,
103            project,
104            tree_diff: None,
105            tree_diff_update_needed: false,
106            tree_diff_base_task: None,
107            base_commit: None,
108            head_commit: None,
109            _subscription: git_store_subscription,
110            _task: worker,
111            update_needed: send,
112        }
113    }
114
115    pub fn diff_base(&self) -> &DiffBase {
116        &self.diff_base
117    }
118
119    pub fn set_repo(&mut self, repo: Option<Entity<Repository>>, cx: &mut Context<Self>) {
120        let same_repo = match (self.repo.as_ref(), repo.as_ref()) {
121            (Some(current), Some(new)) => current.read(cx).id == new.read(cx).id,
122            (None, None) => true,
123            _ => false,
124        };
125        if same_repo {
126            return;
127        }
128
129        self.repo = repo;
130        self.tree_diff = None;
131        self.tree_diff_update_needed = self.diff_base.is_merge_base();
132        self.tree_diff_base_task = None;
133        self.base_commit = None;
134        self.head_commit = None;
135        cx.emit(BranchDiffEvent::FileListChanged);
136        *self.update_needed.borrow_mut() = ();
137    }
138
139    pub fn set_diff_base(&mut self, diff_base: DiffBase, cx: &mut Context<Self>) {
140        if self.diff_base == diff_base {
141            return;
142        }
143
144        self.tree_diff_update_needed = diff_base.is_merge_base();
145        self.tree_diff = None;
146        self.tree_diff_base_task = None;
147        self.diff_base = diff_base;
148        self.base_commit = None;
149        self.head_commit = None;
150
151        cx.emit(BranchDiffEvent::DiffBaseChanged);
152        if self.tree_diff_update_needed {
153            *self.update_needed.borrow_mut() = ();
154        }
155    }
156
157    pub async fn handle_status_updates(
158        this: WeakEntity<Self>,
159        mut recv: postage::watch::Receiver<()>,
160        cx: &mut AsyncWindowContext,
161    ) {
162        this.update(cx, |this, cx| this.spawn_reload_tree_diff(cx))
163            .log_err();
164        while recv.next().await.is_some() {
165            let Ok(()) = this.update(cx, |this, cx| {
166                let mut needs_update = this.tree_diff_update_needed;
167                this.tree_diff_update_needed = false;
168
169                if this.repo.is_none() {
170                    let active_repo = this
171                        .project
172                        .read(cx)
173                        .git_store()
174                        .read(cx)
175                        .active_repository();
176                    if active_repo.is_some() {
177                        this.repo = active_repo;
178                        needs_update = true;
179                    }
180                } else if let Some(repo) = this.repo.as_ref() {
181                    repo.update(cx, |repo, _| {
182                        if let Some(branch) = &repo.branch
183                            && let DiffBase::Merge { base_ref } = &this.diff_base
184                            && let Some(commit) = branch.most_recent_commit.as_ref()
185                            && &branch.ref_name == base_ref
186                            && this.base_commit.as_ref() != Some(&commit.sha)
187                        {
188                            this.base_commit = Some(commit.sha.clone());
189                            needs_update = true;
190                        }
191
192                        if repo.head_commit.as_ref().map(|c| &c.sha) != this.head_commit.as_ref() {
193                            this.head_commit = repo.head_commit.as_ref().map(|c| c.sha.clone());
194                            needs_update = true;
195                        }
196                    })
197                }
198
199                if needs_update {
200                    this.spawn_reload_tree_diff(cx);
201                }
202            }) else {
203                return;
204            };
205        }
206    }
207
208    pub fn status_for_buffer_id(&self, buffer_id: BufferId, cx: &App) -> Option<FileStatus> {
209        let (repo, path) = self
210            .project
211            .read(cx)
212            .git_store()
213            .read(cx)
214            .repository_and_path_for_buffer_id(buffer_id, cx)?;
215        if self.repo() == Some(&repo) {
216            return self.merge_statuses(
217                repo.read(cx)
218                    .status_for_path(&path)
219                    .map(|status| status.status),
220                self.tree_diff
221                    .as_ref()
222                    .and_then(|diff| diff.entries.get(&path)),
223            );
224        }
225        None
226    }
227
228    pub fn merge_statuses(
229        &self,
230        diff_from_head: Option<FileStatus>,
231        diff_from_merge_base: Option<&TreeDiffStatus>,
232    ) -> Option<FileStatus> {
233        match (diff_from_head, diff_from_merge_base) {
234            (None, None) => None,
235            (Some(diff_from_head), None) => Some(diff_from_head),
236            (Some(diff_from_head @ FileStatus::Unmerged(_)), _) => Some(diff_from_head),
237
238            // file does not exist in HEAD
239            // but *does* exist in work-tree
240            // and *does* exist in merge-base
241            (
242                Some(FileStatus::Untracked)
243                | Some(FileStatus::Tracked(TrackedStatus {
244                    index_status: StatusCode::Added,
245                    worktree_status: _,
246                })),
247                Some(_),
248            ) => Some(FileStatus::Tracked(TrackedStatus {
249                index_status: StatusCode::Modified,
250                worktree_status: StatusCode::Modified,
251            })),
252
253            // file exists in HEAD
254            // but *does not* exist in work-tree
255            (Some(diff_from_head), Some(diff_from_merge_base)) if diff_from_head.is_deleted() => {
256                match diff_from_merge_base {
257                    TreeDiffStatus::Added => None, // unchanged, didn't exist in merge base or worktree
258                    _ => Some(diff_from_head),
259                }
260            }
261
262            // file exists in HEAD
263            // and *does* exist in work-tree
264            (Some(FileStatus::Tracked(_)), Some(tree_status)) => {
265                Some(FileStatus::Tracked(TrackedStatus {
266                    index_status: match tree_status {
267                        TreeDiffStatus::Added { .. } => StatusCode::Added,
268                        _ => StatusCode::Modified,
269                    },
270                    worktree_status: match tree_status {
271                        TreeDiffStatus::Added => StatusCode::Added,
272                        _ => StatusCode::Modified,
273                    },
274                }))
275            }
276
277            (_, Some(diff_from_merge_base)) => {
278                Some(diff_status_to_file_status(diff_from_merge_base))
279            }
280        }
281    }
282
283    fn spawn_reload_tree_diff(&mut self, cx: &mut Context<Self>) {
284        if !self.diff_base.is_merge_base() {
285            return;
286        }
287
288        let task = cx.spawn(async move |this, cx| {
289            Self::reload_tree_diff(this, cx).await.log_err();
290        });
291
292        self.tree_diff_base_task = Some(task);
293        cx.notify();
294    }
295
296    pub fn is_tree_base_loading(&self) -> bool {
297        self.tree_diff_base_task
298            .as_ref()
299            .is_some_and(|task| !task.is_ready())
300    }
301
302    pub async fn reload_tree_diff(this: WeakEntity<Self>, cx: &mut AsyncApp) -> Result<()> {
303        let task = this.update(cx, |this, cx| {
304            let DiffBase::Merge { base_ref } = this.diff_base.clone() else {
305                return None;
306            };
307            let Some(repo) = this.repo.as_ref() else {
308                this.tree_diff.take();
309                return None;
310            };
311            repo.update(cx, |repo, cx| {
312                Some(repo.diff_tree(
313                    DiffTreeType::MergeBase {
314                        base: base_ref,
315                        head: "HEAD".into(),
316                    },
317                    cx,
318                ))
319            })
320        })?;
321        let Some(task) = task else { return Ok(()) };
322
323        let diff = task.await??;
324        this.update(cx, |this, cx| {
325            this.tree_diff = Some(diff);
326            cx.emit(BranchDiffEvent::FileListChanged);
327            cx.notify();
328        })
329    }
330
331    pub fn repo(&self) -> Option<&Entity<Repository>> {
332        self.repo.as_ref()
333    }
334
335    #[instrument(skip_all)]
336    pub fn load_buffers(&mut self, cx: &mut Context<Self>) -> Vec<DiffBuffer> {
337        let mut output = Vec::default();
338        let Some(repo) = self.repo.clone() else {
339            return output;
340        };
341        if self.diff_base.is_merge_base() && self.tree_diff.is_none() {
342            return output;
343        }
344
345        self.project.update(cx, |_project, cx| {
346            let mut seen = HashSet::default();
347
348            for item in repo.read(cx).cached_status() {
349                seen.insert(item.repo_path.clone());
350                let branch_diff = self
351                    .tree_diff
352                    .as_ref()
353                    .and_then(|t| t.entries.get(&item.repo_path))
354                    .cloned();
355                let Some(status) = (match self.diff_base {
356                    DiffBase::Head | DiffBase::Merge { .. } => {
357                        self.merge_statuses(Some(item.status), branch_diff.as_ref())
358                    }
359                    DiffBase::Index => item.status.staging().has_unstaged().then_some(item.status),
360                    DiffBase::Staged => item.status.staging().has_staged().then_some(item.status),
361                }) else {
362                    continue;
363                };
364                if !status.has_changes() {
365                    continue;
366                }
367
368                let Some(project_path) =
369                    repo.read(cx).repo_path_to_project_path(&item.repo_path, cx)
370                else {
371                    continue;
372                };
373                let task = Self::load_buffer(
374                    self.diff_base.clone(),
375                    branch_diff,
376                    project_path,
377                    repo.clone(),
378                    cx,
379                );
380
381                output.push(DiffBuffer {
382                    repo_path: item.repo_path.clone(),
383                    load: task,
384                    file_status: item.status,
385                });
386            }
387            let Some(tree_diff) = self.tree_diff.as_ref() else {
388                return;
389            };
390
391            for (path, branch_diff) in tree_diff.entries.iter() {
392                if seen.contains(&path) {
393                    continue;
394                }
395
396                let Some(project_path) = repo.read(cx).repo_path_to_project_path(&path, cx) else {
397                    continue;
398                };
399                let task = Self::load_buffer(
400                    self.diff_base.clone(),
401                    Some(branch_diff.clone()),
402                    project_path,
403                    repo.clone(),
404                    cx,
405                );
406
407                let file_status = diff_status_to_file_status(branch_diff);
408
409                output.push(DiffBuffer {
410                    repo_path: path.clone(),
411                    load: task,
412                    file_status,
413                });
414            }
415        });
416        output
417    }
418
419    #[instrument(skip_all)]
420    fn load_buffer(
421        diff_base: DiffBase,
422        branch_diff: Option<git::status::TreeDiffStatus>,
423        project_path: crate::ProjectPath,
424        repo: Entity<Repository>,
425        cx: &Context<'_, Project>,
426    ) -> Task<Result<LoadedDiffBuffer>> {
427        let task = cx.spawn(async move |project, cx| {
428            let buffer = project
429                .update(cx, |project, cx| project.open_buffer(project_path, cx))?
430                .await?;
431
432            let main_buffer = buffer.clone();
433            let load_conflict_set = diff_base != DiffBase::Staged;
434            let (display_buffer, changes) = match diff_base {
435                DiffBase::Head => {
436                    let diff = project
437                        .update(cx, |project, cx| {
438                            project.open_uncommitted_diff(buffer.clone(), cx)
439                        })?
440                        .await?;
441                    (buffer, diff)
442                }
443                DiffBase::Index => {
444                    let diff = project
445                        .update(cx, |project, cx| {
446                            project.open_unstaged_diff(buffer.clone(), cx)
447                        })?
448                        .await?;
449                    (buffer, diff)
450                }
451                DiffBase::Staged => {
452                    let (diff, index_buffer) = project
453                        .update(cx, |project, cx| {
454                            project.open_staged_diff(buffer.clone(), cx)
455                        })?
456                        .await?;
457                    (index_buffer, diff)
458                }
459                DiffBase::Merge { .. } => {
460                    let diff = if let Some(entry) = branch_diff {
461                        let oid = match entry {
462                            git::status::TreeDiffStatus::Added { .. } => None,
463                            git::status::TreeDiffStatus::Modified { old, .. }
464                            | git::status::TreeDiffStatus::Deleted { old } => Some(old),
465                        };
466                        project
467                            .update(cx, |project, cx| {
468                                project.git_store().update(cx, |git_store, cx| {
469                                    git_store.open_diff_since(oid, buffer.clone(), repo, cx)
470                                })
471                            })?
472                            .await?
473                    } else {
474                        project
475                            .update(cx, |project, cx| {
476                                project.open_uncommitted_diff(buffer.clone(), cx)
477                            })?
478                            .await?
479                    };
480                    (buffer, diff)
481                }
482            };
483            let conflict_set = if load_conflict_set {
484                Some(
485                    project
486                        .update(cx, |project, cx| {
487                            project.git_store().update(cx, |git_store, cx| {
488                                git_store.open_conflict_set(main_buffer.clone(), cx)
489                            })
490                        })?
491                        .await,
492                )
493            } else {
494                None
495            };
496            Ok(LoadedDiffBuffer {
497                display_buffer,
498                main_buffer,
499                diff: changes,
500                conflict_set,
501            })
502        });
503        task
504    }
505}
506
507fn diff_status_to_file_status(branch_diff: &git::status::TreeDiffStatus) -> FileStatus {
508    let file_status = match branch_diff {
509        git::status::TreeDiffStatus::Added { .. } => FileStatus::Tracked(TrackedStatus {
510            index_status: StatusCode::Added,
511            worktree_status: StatusCode::Added,
512        }),
513        git::status::TreeDiffStatus::Modified { .. } => FileStatus::Tracked(TrackedStatus {
514            index_status: StatusCode::Modified,
515            worktree_status: StatusCode::Modified,
516        }),
517        git::status::TreeDiffStatus::Deleted { .. } => FileStatus::Tracked(TrackedStatus {
518            index_status: StatusCode::Deleted,
519            worktree_status: StatusCode::Deleted,
520        }),
521    };
522    file_status
523}
524
525#[derive(Debug)]
526pub struct LoadedDiffBuffer {
527    pub display_buffer: Entity<Buffer>,
528    pub main_buffer: Entity<Buffer>,
529    pub diff: Entity<BufferDiff>,
530    pub conflict_set: Option<Entity<ConflictSet>>,
531}
532
533#[derive(Debug)]
534pub struct DiffBuffer {
535    pub repo_path: RepoPath,
536    pub file_status: FileStatus,
537    pub load: Task<Result<LoadedDiffBuffer>>,
538}
539