jj: improve initialization process

Cargo.lock

@@ -8724,7 +8724,9 @@ dependencies = [
  "anyhow",
  "gpui",
  "jj-lib",
+ "project",
  "workspace-hack",
+ "worktree",
 ]
 
 [[package]]

crates/jj/Cargo.toml

@@ -15,4 +15,6 @@ path = "src/jj.rs"
 anyhow.workspace = true
 gpui.workspace = true
 jj-lib.workspace = true
+project.workspace = true
 workspace-hack.workspace = true
+worktree.workspace = true

crates/jj/src/jj_store.rs

@@ -1,7 +1,9 @@
-use std::path::Path;
+use std::collections::HashMap;
 use std::sync::Arc;
 
-use gpui::{App, Entity, Global, prelude::*};
+use gpui::{App, Context, Entity, EventEmitter, Global, Subscription, prelude::*};
+use project::worktree_store::{WorktreeStore, WorktreeStoreEvent};
+use worktree::{Worktree, WorktreeId};
 
 use crate::{JujutsuRepository, RealJujutsuRepository};
 
@@ -11,18 +13,22 @@ struct GlobalJujutsuStore(Entity<JujutsuStore>);
 impl Global for GlobalJujutsuStore {}
 
 pub struct JujutsuStore {
-    repository: Arc<dyn JujutsuRepository>,
+    active_repository: Option<WorktreeId>,
+    repositories: HashMap<WorktreeId, Arc<dyn JujutsuRepository>>,
+    _subscriptions: Vec<Subscription>,
 }
 
+pub enum JujutsuStoreEvent {
+    ActiveRepositoryChanged(Option<WorktreeId>),
+    RepositoryAdded(WorktreeId),
+    RepositoryRemoved(WorktreeId),
+}
+
+impl EventEmitter<JujutsuStoreEvent> for JujutsuStore {}
+
 impl JujutsuStore {
-    pub fn init_global(cx: &mut App) {
-        let Some(repository) = RealJujutsuRepository::new(Path::new(".")).ok() else {
-            return;
-        };
-
-        let repository = Arc::new(repository);
-        let jj_store = cx.new(|cx| JujutsuStore::new(repository, cx));
-
+    pub fn init_global(cx: &mut App, worktree_store: Entity<WorktreeStore>) {
+        let jj_store = cx.new(|cx| JujutsuStore::new(worktree_store, cx));
         cx.set_global(GlobalJujutsuStore(jj_store));
     }
 
@@ -31,11 +37,84 @@ impl JujutsuStore {
             .map(|global| global.0.clone())
     }
 
-    pub fn new(repository: Arc<dyn JujutsuRepository>, _cx: &mut Context<Self>) -> Self {
-        Self { repository }
+    pub fn new(worktree_store: Entity<WorktreeStore>, cx: &mut Context<Self>) -> Self {
+        let _subscriptions = vec![cx.subscribe(&worktree_store, Self::on_worktree_store_event)];
+
+        let mut store = JujutsuStore {
+            active_repository: None,
+            repositories: HashMap::default(),
+            _subscriptions,
+        };
+
+        let existing_worktrees: Vec<_> = worktree_store.read(cx).worktrees().collect();
+
+        for worktree in existing_worktrees {
+            store.scan_worktree_for_jj_repo(&worktree, cx);
+        }
+
+        store
     }
 
-    pub fn repository(&self) -> &Arc<dyn JujutsuRepository> {
-        &self.repository
+    fn on_worktree_store_event(
+        &mut self,
+        _worktree_store: Entity<WorktreeStore>,
+        event: &WorktreeStoreEvent,
+        cx: &mut Context<Self>,
+    ) {
+        match event {
+            WorktreeStoreEvent::WorktreeAdded(worktree) => {
+                self.scan_worktree_for_jj_repo(worktree, cx);
+            }
+            WorktreeStoreEvent::WorktreeRemoved(_, worktree_id) => {
+                if self.repositories.remove(worktree_id).is_some() {
+                    cx.emit(JujutsuStoreEvent::RepositoryRemoved(*worktree_id));
+
+                    if self.active_repository == Some(*worktree_id) {
+                        self.active_repository = None;
+                        cx.emit(JujutsuStoreEvent::ActiveRepositoryChanged(None));
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+
+    fn scan_worktree_for_jj_repo(&mut self, worktree: &Entity<Worktree>, cx: &mut Context<Self>) {
+        let worktree = worktree.read(cx);
+        let worktree_id = worktree.id();
+        let root_path = worktree.abs_path();
+
+        match RealJujutsuRepository::new(&root_path) {
+            Ok(repository) => {
+                let repository = Arc::new(repository);
+
+                self.repositories.insert(worktree_id, repository);
+                cx.emit(JujutsuStoreEvent::RepositoryAdded(worktree_id));
+
+                if self.active_repository.is_none() {
+                    self.active_repository = Some(worktree_id);
+                    cx.emit(JujutsuStoreEvent::ActiveRepositoryChanged(Some(
+                        worktree_id,
+                    )));
+                }
+            }
+            Err(_) => {}
+        }
+    }
+
+    pub fn repository_for_worktree(
+        &self,
+        worktree_id: WorktreeId,
+    ) -> Option<&Arc<dyn JujutsuRepository>> {
+        self.repositories.get(&worktree_id)
+    }
+
+    pub fn active_repository(&self) -> Option<&Arc<dyn JujutsuRepository>> {
+        self.active_repository
+            .and_then(|id| self.repositories.get(&id))
+    }
+
+    pub fn repository(&self) -> Option<&Arc<dyn JujutsuRepository>> {
+        self.active_repository()
     }
 }

crates/jj_ui/src/bookmark_picker.rs

@@ -81,7 +81,11 @@ impl BookmarkPickerDelegate {
         jj_store: Entity<JujutsuStore>,
         cx: &mut Context<BookmarkPicker>,
     ) -> Self {
-        let bookmarks = jj_store.read(cx).repository().list_bookmarks();
+        let bookmarks = jj_store
+            .read(cx)
+            .repository()
+            .map(|repo| repo.list_bookmarks())
+            .unwrap_or_default();
 
         Self {
             picker,

crates/jj_ui/src/jj_ui.rs

@@ -7,9 +7,9 @@ use jj::JujutsuStore;
 use workspace::Workspace;
 
 pub fn init(cx: &mut App) {
-    JujutsuStore::init_global(cx);
-
-    cx.observe_new(|workspace: &mut Workspace, _window, _cx| {
+    cx.observe_new(|workspace: &mut Workspace, _window, cx| {
+        let worktree_store = workspace.project().read(cx).worktree_store();
+        JujutsuStore::init_global(cx, worktree_store);
         bookmark_picker::register(workspace);
     })
     .detach();