1use std::{
2 path::{Path, PathBuf},
3 sync::{
4 Arc,
5 atomic::{AtomicU64, AtomicUsize},
6 },
7};
8
9use anyhow::{Context as _, Result, anyhow, bail};
10use collections::HashMap;
11use fs::{Fs, copy_recursive};
12use futures::{FutureExt, future::Shared};
13use gpui::{
14 App, AppContext as _, AsyncApp, Context, Entity, EntityId, EventEmitter, Global, Task,
15 WeakEntity,
16};
17use itertools::Either;
18use rpc::{
19 AnyProtoClient, ErrorExt, TypedEnvelope,
20 proto::{self, REMOTE_SERVER_PROJECT_ID},
21};
22use text::ReplicaId;
23use util::{
24 ResultExt,
25 paths::{PathStyle, RemotePathBuf, SanitizedPath},
26 rel_path::RelPath,
27};
28use worktree::{
29 CreatedEntry, Entry, ProjectEntryId, UpdatedEntriesSet, UpdatedGitRepositoriesSet, Worktree,
30 WorktreeId,
31};
32
33use crate::{ProjectPath, trusted_worktrees::TrustedWorktrees};
34
35enum WorktreeStoreState {
36 Local {
37 fs: Arc<dyn Fs>,
38 },
39 Remote {
40 upstream_client: AnyProtoClient,
41 upstream_project_id: u64,
42 path_style: PathStyle,
43 },
44}
45
46#[derive(Clone)]
47pub struct WorktreeIdCounter(Arc<AtomicU64>);
48
49impl Default for WorktreeIdCounter {
50 fn default() -> Self {
51 Self(Arc::new(AtomicU64::new(1)))
52 }
53}
54
55impl WorktreeIdCounter {
56 pub fn get(cx: &mut App) -> Self {
57 cx.default_global::<Self>().clone()
58 }
59
60 fn next(&self) -> u64 {
61 self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
62 }
63}
64
65impl Global for WorktreeIdCounter {}
66
67pub struct WorktreeStore {
68 next_entry_id: Arc<AtomicUsize>,
69 next_worktree_id: WorktreeIdCounter,
70 downstream_client: Option<(AnyProtoClient, u64)>,
71 retain_worktrees: bool,
72 worktrees: Vec<WorktreeHandle>,
73 worktrees_reordered: bool,
74 scanning_enabled: bool,
75 #[allow(clippy::type_complexity)]
76 loading_worktrees:
77 HashMap<Arc<SanitizedPath>, Shared<Task<Result<Entity<Worktree>, Arc<anyhow::Error>>>>>,
78 state: WorktreeStoreState,
79}
80
81#[derive(Debug)]
82pub enum WorktreeStoreEvent {
83 WorktreeAdded(Entity<Worktree>),
84 WorktreeRemoved(EntityId, WorktreeId),
85 WorktreeReleased(EntityId, WorktreeId),
86 WorktreeOrderChanged,
87 WorktreeUpdateSent(Entity<Worktree>),
88 WorktreeUpdatedEntries(WorktreeId, UpdatedEntriesSet),
89 WorktreeUpdatedGitRepositories(WorktreeId, UpdatedGitRepositoriesSet),
90 WorktreeDeletedEntry(WorktreeId, ProjectEntryId),
91}
92
93impl EventEmitter<WorktreeStoreEvent> for WorktreeStore {}
94
95impl WorktreeStore {
96 pub fn init(client: &AnyProtoClient) {
97 client.add_entity_request_handler(Self::handle_create_project_entry);
98 client.add_entity_request_handler(Self::handle_copy_project_entry);
99 client.add_entity_request_handler(Self::handle_delete_project_entry);
100 client.add_entity_request_handler(Self::handle_expand_project_entry);
101 client.add_entity_request_handler(Self::handle_expand_all_for_project_entry);
102 }
103
104 pub fn init_remote(client: &AnyProtoClient) {
105 client.add_entity_request_handler(Self::handle_allocate_worktree_id);
106 }
107
108 pub fn local(
109 retain_worktrees: bool,
110 fs: Arc<dyn Fs>,
111 next_worktree_id: WorktreeIdCounter,
112 ) -> Self {
113 Self {
114 next_entry_id: Default::default(),
115 next_worktree_id,
116 loading_worktrees: Default::default(),
117 downstream_client: None,
118 worktrees: Vec::new(),
119 worktrees_reordered: false,
120 scanning_enabled: true,
121 retain_worktrees,
122 state: WorktreeStoreState::Local { fs },
123 }
124 }
125
126 pub fn remote(
127 retain_worktrees: bool,
128 upstream_client: AnyProtoClient,
129 upstream_project_id: u64,
130 path_style: PathStyle,
131 next_worktree_id: WorktreeIdCounter,
132 ) -> Self {
133 Self {
134 next_entry_id: Default::default(),
135 next_worktree_id,
136 loading_worktrees: Default::default(),
137 downstream_client: None,
138 worktrees: Vec::new(),
139 worktrees_reordered: false,
140 scanning_enabled: true,
141 retain_worktrees,
142 state: WorktreeStoreState::Remote {
143 upstream_client,
144 upstream_project_id,
145 path_style,
146 },
147 }
148 }
149
150 pub fn next_worktree_id(&self) -> impl Future<Output = Result<WorktreeId>> + use<> {
151 let strategy = match (&self.state, &self.downstream_client) {
152 // we are a remote server, the client is in charge of assigning worktree ids
153 (WorktreeStoreState::Local { .. }, Some((client, REMOTE_SERVER_PROJECT_ID))) => {
154 Either::Left(client.clone())
155 }
156 // we are just a local zed project, we can assign ids
157 (WorktreeStoreState::Local { .. }, _) => Either::Right(self.next_worktree_id.next()),
158 // we are connected to a remote server, we are in charge of assigning worktree ids
159 (WorktreeStoreState::Remote { .. }, _) => Either::Right(self.next_worktree_id.next()),
160 };
161 async move {
162 match strategy {
163 Either::Left(client) => Ok(client
164 .request(proto::AllocateWorktreeId {
165 project_id: REMOTE_SERVER_PROJECT_ID,
166 })
167 .await?
168 .worktree_id),
169 Either::Right(id) => Ok(id),
170 }
171 .map(WorktreeId::from_proto)
172 }
173 }
174
175 pub fn disable_scanner(&mut self) {
176 self.scanning_enabled = false;
177 }
178
179 /// Iterates through all worktrees, including ones that don't appear in the project panel
180 pub fn worktrees(&self) -> impl '_ + DoubleEndedIterator<Item = Entity<Worktree>> {
181 self.worktrees
182 .iter()
183 .filter_map(move |worktree| worktree.upgrade())
184 }
185
186 /// Iterates through all user-visible worktrees, the ones that appear in the project panel.
187 pub fn visible_worktrees<'a>(
188 &'a self,
189 cx: &'a App,
190 ) -> impl 'a + DoubleEndedIterator<Item = Entity<Worktree>> {
191 self.worktrees()
192 .filter(|worktree| worktree.read(cx).is_visible())
193 }
194
195 /// Iterates through all user-visible worktrees (directories and files that appear in the project panel) and other, invisible single files that could appear e.g. due to drag and drop.
196 pub fn visible_worktrees_and_single_files<'a>(
197 &'a self,
198 cx: &'a App,
199 ) -> impl 'a + DoubleEndedIterator<Item = Entity<Worktree>> {
200 self.worktrees()
201 .filter(|worktree| worktree.read(cx).is_visible() || worktree.read(cx).is_single_file())
202 }
203
204 pub fn worktree_for_id(&self, id: WorktreeId, cx: &App) -> Option<Entity<Worktree>> {
205 self.worktrees()
206 .find(|worktree| worktree.read(cx).id() == id)
207 }
208
209 pub fn worktree_for_entry(
210 &self,
211 entry_id: ProjectEntryId,
212 cx: &App,
213 ) -> Option<Entity<Worktree>> {
214 self.worktrees()
215 .find(|worktree| worktree.read(cx).contains_entry(entry_id))
216 }
217
218 pub fn find_worktree(
219 &self,
220 abs_path: impl AsRef<Path>,
221 cx: &App,
222 ) -> Option<(Entity<Worktree>, Arc<RelPath>)> {
223 let abs_path = SanitizedPath::new(abs_path.as_ref());
224 for tree in self.worktrees() {
225 let path_style = tree.read(cx).path_style();
226 if let Some(relative_path) =
227 path_style.strip_prefix(abs_path.as_ref(), tree.read(cx).abs_path().as_ref())
228 {
229 return Some((tree.clone(), relative_path.into_arc()));
230 }
231 }
232 None
233 }
234
235 pub fn project_path_for_absolute_path(&self, abs_path: &Path, cx: &App) -> Option<ProjectPath> {
236 self.find_worktree(abs_path, cx)
237 .map(|(worktree, relative_path)| ProjectPath {
238 worktree_id: worktree.read(cx).id(),
239 path: relative_path,
240 })
241 }
242
243 pub fn absolutize(&self, project_path: &ProjectPath, cx: &App) -> Option<PathBuf> {
244 let worktree = self.worktree_for_id(project_path.worktree_id, cx)?;
245 Some(worktree.read(cx).absolutize(&project_path.path))
246 }
247
248 pub fn path_style(&self) -> PathStyle {
249 match &self.state {
250 WorktreeStoreState::Local { .. } => PathStyle::local(),
251 WorktreeStoreState::Remote { path_style, .. } => *path_style,
252 }
253 }
254
255 pub fn find_or_create_worktree(
256 &mut self,
257 abs_path: impl AsRef<Path>,
258 visible: bool,
259 cx: &mut Context<Self>,
260 ) -> Task<Result<(Entity<Worktree>, Arc<RelPath>)>> {
261 let abs_path = abs_path.as_ref();
262 if let Some((tree, relative_path)) = self.find_worktree(abs_path, cx) {
263 Task::ready(Ok((tree, relative_path)))
264 } else {
265 let worktree = self.create_worktree(abs_path, visible, cx);
266 cx.background_spawn(async move { Ok((worktree.await?, RelPath::empty().into())) })
267 }
268 }
269
270 pub fn entry_for_id<'a>(&'a self, entry_id: ProjectEntryId, cx: &'a App) -> Option<&'a Entry> {
271 self.worktrees()
272 .find_map(|worktree| worktree.read(cx).entry_for_id(entry_id))
273 }
274
275 pub fn worktree_and_entry_for_id<'a>(
276 &'a self,
277 entry_id: ProjectEntryId,
278 cx: &'a App,
279 ) -> Option<(Entity<Worktree>, &'a Entry)> {
280 self.worktrees().find_map(|worktree| {
281 worktree
282 .read(cx)
283 .entry_for_id(entry_id)
284 .map(|e| (worktree.clone(), e))
285 })
286 }
287
288 pub fn entry_for_path<'a>(&'a self, path: &ProjectPath, cx: &'a App) -> Option<&'a Entry> {
289 self.worktree_for_id(path.worktree_id, cx)?
290 .read(cx)
291 .entry_for_path(&path.path)
292 }
293
294 pub fn copy_entry(
295 &mut self,
296 entry_id: ProjectEntryId,
297 new_project_path: ProjectPath,
298 cx: &mut Context<Self>,
299 ) -> Task<Result<Option<Entry>>> {
300 let Some(old_worktree) = self.worktree_for_entry(entry_id, cx) else {
301 return Task::ready(Err(anyhow!("no such worktree")));
302 };
303 let Some(old_entry) = old_worktree.read(cx).entry_for_id(entry_id) else {
304 return Task::ready(Err(anyhow!("no such entry")));
305 };
306 let Some(new_worktree) = self.worktree_for_id(new_project_path.worktree_id, cx) else {
307 return Task::ready(Err(anyhow!("no such worktree")));
308 };
309
310 match &self.state {
311 WorktreeStoreState::Local { fs } => {
312 let old_abs_path = old_worktree.read(cx).absolutize(&old_entry.path);
313 let new_abs_path = new_worktree.read(cx).absolutize(&new_project_path.path);
314 let fs = fs.clone();
315 let copy = cx.background_spawn(async move {
316 copy_recursive(
317 fs.as_ref(),
318 &old_abs_path,
319 &new_abs_path,
320 Default::default(),
321 )
322 .await
323 });
324
325 cx.spawn(async move |_, cx| {
326 copy.await?;
327 new_worktree
328 .update(cx, |this, cx| {
329 this.as_local_mut().unwrap().refresh_entry(
330 new_project_path.path,
331 None,
332 cx,
333 )
334 })
335 .await
336 })
337 }
338 WorktreeStoreState::Remote {
339 upstream_client,
340 upstream_project_id,
341 ..
342 } => {
343 let response = upstream_client.request(proto::CopyProjectEntry {
344 project_id: *upstream_project_id,
345 entry_id: entry_id.to_proto(),
346 new_path: new_project_path.path.to_proto(),
347 new_worktree_id: new_project_path.worktree_id.to_proto(),
348 });
349 cx.spawn(async move |_, cx| {
350 let response = response.await?;
351 match response.entry {
352 Some(entry) => new_worktree
353 .update(cx, |worktree, cx| {
354 worktree.as_remote_mut().unwrap().insert_entry(
355 entry,
356 response.worktree_scan_id as usize,
357 cx,
358 )
359 })
360 .await
361 .map(Some),
362 None => Ok(None),
363 }
364 })
365 }
366 }
367 }
368
369 pub fn rename_entry(
370 &mut self,
371 entry_id: ProjectEntryId,
372 new_project_path: ProjectPath,
373 cx: &mut Context<Self>,
374 ) -> Task<Result<CreatedEntry>> {
375 let Some(old_worktree) = self.worktree_for_entry(entry_id, cx) else {
376 return Task::ready(Err(anyhow!("no such worktree")));
377 };
378 let Some(old_entry) = old_worktree.read(cx).entry_for_id(entry_id).cloned() else {
379 return Task::ready(Err(anyhow!("no such entry")));
380 };
381 let Some(new_worktree) = self.worktree_for_id(new_project_path.worktree_id, cx) else {
382 return Task::ready(Err(anyhow!("no such worktree")));
383 };
384
385 match &self.state {
386 WorktreeStoreState::Local { fs } => {
387 let abs_old_path = old_worktree.read(cx).absolutize(&old_entry.path);
388 let new_worktree_ref = new_worktree.read(cx);
389 let is_root_entry = new_worktree_ref
390 .root_entry()
391 .is_some_and(|e| e.id == entry_id);
392 let abs_new_path = if is_root_entry {
393 let abs_path = new_worktree_ref.abs_path();
394 let Some(root_parent_path) = abs_path.parent() else {
395 return Task::ready(Err(anyhow!("no parent for path {:?}", abs_path)));
396 };
397 root_parent_path.join(new_project_path.path.as_std_path())
398 } else {
399 new_worktree_ref.absolutize(&new_project_path.path)
400 };
401
402 let fs = fs.clone();
403 let case_sensitive = new_worktree
404 .read(cx)
405 .as_local()
406 .unwrap()
407 .fs_is_case_sensitive();
408
409 let do_rename =
410 async move |fs: &dyn Fs, old_path: &Path, new_path: &Path, overwrite| {
411 fs.rename(
412 &old_path,
413 &new_path,
414 fs::RenameOptions {
415 overwrite,
416 ..fs::RenameOptions::default()
417 },
418 )
419 .await
420 .with_context(|| format!("renaming {old_path:?} into {new_path:?}"))
421 };
422
423 let rename = cx.background_spawn({
424 let abs_new_path = abs_new_path.clone();
425 async move {
426 // If we're on a case-insensitive FS and we're doing a case-only rename (i.e. `foobar` to `FOOBAR`)
427 // we want to overwrite, because otherwise we run into a file-already-exists error.
428 let overwrite = !case_sensitive
429 && abs_old_path != abs_new_path
430 && abs_old_path.to_str().map(|p| p.to_lowercase())
431 == abs_new_path.to_str().map(|p| p.to_lowercase());
432
433 // The directory we're renaming into might not exist yet
434 if let Err(e) =
435 do_rename(fs.as_ref(), &abs_old_path, &abs_new_path, overwrite).await
436 {
437 if let Some(err) = e.downcast_ref::<std::io::Error>()
438 && err.kind() == std::io::ErrorKind::NotFound
439 {
440 if let Some(parent) = abs_new_path.parent() {
441 fs.create_dir(parent).await.with_context(|| {
442 format!("creating parent directory {parent:?}")
443 })?;
444 return do_rename(
445 fs.as_ref(),
446 &abs_old_path,
447 &abs_new_path,
448 overwrite,
449 )
450 .await;
451 }
452 }
453 return Err(e);
454 }
455 Ok(())
456 }
457 });
458
459 cx.spawn(async move |_, cx| {
460 rename.await?;
461 Ok(new_worktree
462 .update(cx, |this, cx| {
463 let local = this.as_local_mut().unwrap();
464 if is_root_entry {
465 // We eagerly update `abs_path` and refresh this worktree.
466 // Otherwise, the FS watcher would do it on the `RootUpdated` event,
467 // but with a noticeable delay, so we handle it proactively.
468 local.update_abs_path_and_refresh(
469 SanitizedPath::new_arc(&abs_new_path),
470 cx,
471 );
472 Task::ready(Ok(this.root_entry().cloned()))
473 } else {
474 // First refresh the parent directory (in case it was newly created)
475 if let Some(parent) = new_project_path.path.parent() {
476 let _ = local.refresh_entries_for_paths(vec![parent.into()]);
477 }
478 // Then refresh the new path
479 local.refresh_entry(
480 new_project_path.path.clone(),
481 Some(old_entry.path),
482 cx,
483 )
484 }
485 })
486 .await?
487 .map(CreatedEntry::Included)
488 .unwrap_or_else(|| CreatedEntry::Excluded {
489 abs_path: abs_new_path,
490 }))
491 })
492 }
493 WorktreeStoreState::Remote {
494 upstream_client,
495 upstream_project_id,
496 ..
497 } => {
498 let response = upstream_client.request(proto::RenameProjectEntry {
499 project_id: *upstream_project_id,
500 entry_id: entry_id.to_proto(),
501 new_path: new_project_path.path.to_proto(),
502 new_worktree_id: new_project_path.worktree_id.to_proto(),
503 });
504 cx.spawn(async move |_, cx| {
505 let response = response.await?;
506 match response.entry {
507 Some(entry) => new_worktree
508 .update(cx, |worktree, cx| {
509 worktree.as_remote_mut().unwrap().insert_entry(
510 entry,
511 response.worktree_scan_id as usize,
512 cx,
513 )
514 })
515 .await
516 .map(CreatedEntry::Included),
517 None => {
518 let abs_path = new_worktree.read_with(cx, |worktree, _| {
519 worktree.absolutize(&new_project_path.path)
520 });
521 Ok(CreatedEntry::Excluded { abs_path })
522 }
523 }
524 })
525 }
526 }
527 }
528 pub fn create_worktree(
529 &mut self,
530 abs_path: impl AsRef<Path>,
531 visible: bool,
532 cx: &mut Context<Self>,
533 ) -> Task<Result<Entity<Worktree>>> {
534 let abs_path: Arc<SanitizedPath> = SanitizedPath::new_arc(&abs_path);
535 let is_via_collab = matches!(&self.state, WorktreeStoreState::Remote { upstream_client, .. } if upstream_client.is_via_collab());
536 if !self.loading_worktrees.contains_key(&abs_path) {
537 let task = match &self.state {
538 WorktreeStoreState::Remote {
539 upstream_client,
540 path_style,
541 ..
542 } => {
543 if upstream_client.is_via_collab() {
544 Task::ready(Err(Arc::new(anyhow!("cannot create worktrees via collab"))))
545 } else {
546 let abs_path = RemotePathBuf::new(abs_path.to_string(), *path_style);
547 self.create_remote_worktree(upstream_client.clone(), abs_path, visible, cx)
548 }
549 }
550 WorktreeStoreState::Local { fs } => {
551 self.create_local_worktree(fs.clone(), abs_path.clone(), visible, cx)
552 }
553 };
554
555 self.loading_worktrees
556 .insert(abs_path.clone(), task.shared());
557 }
558 let task = self.loading_worktrees.get(&abs_path).unwrap().clone();
559 cx.spawn(async move |this, cx| {
560 let result = task.await;
561 this.update(cx, |this, _| this.loading_worktrees.remove(&abs_path))
562 .ok();
563 match result {
564 Ok(worktree) => {
565 if !is_via_collab {
566 if let Some((trusted_worktrees, worktree_store)) = this
567 .update(cx, |_, cx| {
568 TrustedWorktrees::try_get_global(cx).zip(Some(cx.entity()))
569 })
570 .ok()
571 .flatten()
572 {
573 trusted_worktrees.update(cx, |trusted_worktrees, cx| {
574 trusted_worktrees.can_trust(
575 &worktree_store,
576 worktree.read(cx).id(),
577 cx,
578 );
579 });
580 }
581 }
582 Ok(worktree)
583 }
584 Err(err) => Err((*err).cloned()),
585 }
586 })
587 }
588
589 fn create_remote_worktree(
590 &mut self,
591 client: AnyProtoClient,
592 abs_path: RemotePathBuf,
593 visible: bool,
594 cx: &mut Context<Self>,
595 ) -> Task<Result<Entity<Worktree>, Arc<anyhow::Error>>> {
596 let path_style = abs_path.path_style();
597 let mut abs_path = abs_path.to_string();
598 // If we start with `/~` that means the ssh path was something like `ssh://user@host/~/home-dir-folder/`
599 // in which case want to strip the leading the `/`.
600 // On the host-side, the `~` will get expanded.
601 // That's what git does too: https://github.com/libgit2/libgit2/issues/3345#issuecomment-127050850
602 if abs_path.starts_with("/~") {
603 abs_path = abs_path[1..].to_string();
604 }
605 if abs_path.is_empty() {
606 abs_path = "~/".to_string();
607 }
608
609 cx.spawn(async move |this, cx| {
610 let this = this.upgrade().context("Dropped worktree store")?;
611
612 let path = RemotePathBuf::new(abs_path, path_style);
613 let response = client
614 .request(proto::AddWorktree {
615 project_id: REMOTE_SERVER_PROJECT_ID,
616 path: path.to_proto(),
617 visible,
618 })
619 .await?;
620
621 if let Some(existing_worktree) = this.read_with(cx, |this, cx| {
622 this.worktree_for_id(WorktreeId::from_proto(response.worktree_id), cx)
623 }) {
624 return Ok(existing_worktree);
625 }
626
627 let root_path_buf = PathBuf::from(response.canonicalized_path.clone());
628 let root_name = root_path_buf
629 .file_name()
630 .map(|n| n.to_string_lossy().into_owned())
631 .unwrap_or(root_path_buf.to_string_lossy().into_owned());
632
633 let worktree = cx.update(|cx| {
634 Worktree::remote(
635 REMOTE_SERVER_PROJECT_ID,
636 ReplicaId::REMOTE_SERVER,
637 proto::WorktreeMetadata {
638 id: response.worktree_id,
639 root_name,
640 visible,
641 abs_path: response.canonicalized_path,
642 },
643 client,
644 path_style,
645 cx,
646 )
647 });
648
649 this.update(cx, |this, cx| {
650 this.add(&worktree, cx);
651 });
652 Ok(worktree)
653 })
654 }
655
656 fn create_local_worktree(
657 &mut self,
658 fs: Arc<dyn Fs>,
659 abs_path: Arc<SanitizedPath>,
660 visible: bool,
661 cx: &mut Context<Self>,
662 ) -> Task<Result<Entity<Worktree>, Arc<anyhow::Error>>> {
663 let next_entry_id = self.next_entry_id.clone();
664 let scanning_enabled = self.scanning_enabled;
665
666 let next_worktree_id = self.next_worktree_id();
667
668 cx.spawn(async move |this, cx| {
669 let worktree_id = next_worktree_id.await?;
670 let worktree = Worktree::local(
671 SanitizedPath::cast_arc(abs_path.clone()),
672 visible,
673 fs,
674 next_entry_id,
675 scanning_enabled,
676 worktree_id,
677 cx,
678 )
679 .await?;
680
681 this.update(cx, |this, cx| this.add(&worktree, cx))?;
682
683 if visible {
684 cx.update(|cx| {
685 cx.add_recent_document(abs_path.as_path());
686 });
687 }
688
689 Ok(worktree)
690 })
691 }
692
693 pub fn add(&mut self, worktree: &Entity<Worktree>, cx: &mut Context<Self>) {
694 let worktree_id = worktree.read(cx).id();
695 debug_assert!(self.worktrees().all(|w| w.read(cx).id() != worktree_id));
696
697 let push_strong_handle = self.retain_worktrees || worktree.read(cx).is_visible();
698 let handle = if push_strong_handle {
699 WorktreeHandle::Strong(worktree.clone())
700 } else {
701 WorktreeHandle::Weak(worktree.downgrade())
702 };
703 if self.worktrees_reordered {
704 self.worktrees.push(handle);
705 } else {
706 let i = match self
707 .worktrees
708 .binary_search_by_key(&Some(worktree.read(cx).abs_path()), |other| {
709 other.upgrade().map(|worktree| worktree.read(cx).abs_path())
710 }) {
711 Ok(i) | Err(i) => i,
712 };
713 self.worktrees.insert(i, handle);
714 }
715
716 cx.emit(WorktreeStoreEvent::WorktreeAdded(worktree.clone()));
717 self.send_project_updates(cx);
718
719 let handle_id = worktree.entity_id();
720 cx.subscribe(worktree, |_, worktree, event, cx| {
721 let worktree_id = worktree.read(cx).id();
722 match event {
723 worktree::Event::UpdatedEntries(changes) => {
724 cx.emit(WorktreeStoreEvent::WorktreeUpdatedEntries(
725 worktree_id,
726 changes.clone(),
727 ));
728 }
729 worktree::Event::UpdatedGitRepositories(set) => {
730 cx.emit(WorktreeStoreEvent::WorktreeUpdatedGitRepositories(
731 worktree_id,
732 set.clone(),
733 ));
734 }
735 worktree::Event::DeletedEntry(id) => {
736 cx.emit(WorktreeStoreEvent::WorktreeDeletedEntry(worktree_id, *id))
737 }
738 }
739 })
740 .detach();
741 cx.observe_release(worktree, move |this, worktree, cx| {
742 cx.emit(WorktreeStoreEvent::WorktreeReleased(
743 handle_id,
744 worktree.id(),
745 ));
746 cx.emit(WorktreeStoreEvent::WorktreeRemoved(
747 handle_id,
748 worktree.id(),
749 ));
750 this.send_project_updates(cx);
751 })
752 .detach();
753 }
754
755 pub fn remove_worktree(&mut self, id_to_remove: WorktreeId, cx: &mut Context<Self>) {
756 self.worktrees.retain(|worktree| {
757 if let Some(worktree) = worktree.upgrade() {
758 if worktree.read(cx).id() == id_to_remove {
759 cx.emit(WorktreeStoreEvent::WorktreeRemoved(
760 worktree.entity_id(),
761 id_to_remove,
762 ));
763 false
764 } else {
765 true
766 }
767 } else {
768 false
769 }
770 });
771 self.send_project_updates(cx);
772 }
773
774 pub fn set_worktrees_reordered(&mut self, worktrees_reordered: bool) {
775 self.worktrees_reordered = worktrees_reordered;
776 }
777
778 fn upstream_client(&self) -> Option<(AnyProtoClient, u64)> {
779 match &self.state {
780 WorktreeStoreState::Remote {
781 upstream_client,
782 upstream_project_id,
783 ..
784 } => Some((upstream_client.clone(), *upstream_project_id)),
785 WorktreeStoreState::Local { .. } => None,
786 }
787 }
788
789 pub fn set_worktrees_from_proto(
790 &mut self,
791 worktrees: Vec<proto::WorktreeMetadata>,
792 replica_id: ReplicaId,
793 cx: &mut Context<Self>,
794 ) -> Result<()> {
795 let mut old_worktrees_by_id = self
796 .worktrees
797 .drain(..)
798 .filter_map(|worktree| {
799 let worktree = worktree.upgrade()?;
800 Some((worktree.read(cx).id(), worktree))
801 })
802 .collect::<HashMap<_, _>>();
803
804 let (client, project_id) = self.upstream_client().context("invalid project")?;
805
806 for worktree in worktrees {
807 if let Some(old_worktree) =
808 old_worktrees_by_id.remove(&WorktreeId::from_proto(worktree.id))
809 {
810 let push_strong_handle =
811 self.retain_worktrees || old_worktree.read(cx).is_visible();
812 let handle = if push_strong_handle {
813 WorktreeHandle::Strong(old_worktree.clone())
814 } else {
815 WorktreeHandle::Weak(old_worktree.downgrade())
816 };
817 self.worktrees.push(handle);
818 } else {
819 self.add(
820 &Worktree::remote(
821 project_id,
822 replica_id,
823 worktree,
824 client.clone(),
825 self.path_style(),
826 cx,
827 ),
828 cx,
829 );
830 }
831 }
832 self.send_project_updates(cx);
833
834 Ok(())
835 }
836
837 pub fn move_worktree(
838 &mut self,
839 source: WorktreeId,
840 destination: WorktreeId,
841 cx: &mut Context<Self>,
842 ) -> Result<()> {
843 if source == destination {
844 return Ok(());
845 }
846
847 let mut source_index = None;
848 let mut destination_index = None;
849 for (i, worktree) in self.worktrees.iter().enumerate() {
850 if let Some(worktree) = worktree.upgrade() {
851 let worktree_id = worktree.read(cx).id();
852 if worktree_id == source {
853 source_index = Some(i);
854 if destination_index.is_some() {
855 break;
856 }
857 } else if worktree_id == destination {
858 destination_index = Some(i);
859 if source_index.is_some() {
860 break;
861 }
862 }
863 }
864 }
865
866 let source_index =
867 source_index.with_context(|| format!("Missing worktree for id {source}"))?;
868 let destination_index =
869 destination_index.with_context(|| format!("Missing worktree for id {destination}"))?;
870
871 if source_index == destination_index {
872 return Ok(());
873 }
874
875 let worktree_to_move = self.worktrees.remove(source_index);
876 self.worktrees.insert(destination_index, worktree_to_move);
877 self.worktrees_reordered = true;
878 cx.emit(WorktreeStoreEvent::WorktreeOrderChanged);
879 cx.notify();
880 Ok(())
881 }
882
883 pub fn disconnected_from_host(&mut self, cx: &mut App) {
884 for worktree in &self.worktrees {
885 if let Some(worktree) = worktree.upgrade() {
886 worktree.update(cx, |worktree, _| {
887 if let Some(worktree) = worktree.as_remote_mut() {
888 worktree.disconnected_from_host();
889 }
890 });
891 }
892 }
893 }
894
895 pub fn send_project_updates(&mut self, cx: &mut Context<Self>) {
896 let Some((downstream_client, project_id)) = self.downstream_client.clone() else {
897 return;
898 };
899
900 let update = proto::UpdateProject {
901 project_id,
902 worktrees: self.worktree_metadata_protos(cx),
903 };
904
905 // collab has bad concurrency guarantees, so we send requests in serial.
906 let update_project = if downstream_client.is_via_collab() {
907 Some(downstream_client.request(update))
908 } else {
909 downstream_client.send(update).log_err();
910 None
911 };
912 cx.spawn(async move |this, cx| {
913 if let Some(update_project) = update_project {
914 update_project.await?;
915 }
916
917 this.update(cx, |this, cx| {
918 let worktrees = this.worktrees().collect::<Vec<_>>();
919
920 for worktree in worktrees {
921 worktree.update(cx, |worktree, cx| {
922 let client = downstream_client.clone();
923 worktree.observe_updates(project_id, cx, {
924 move |update| {
925 let client = client.clone();
926 async move {
927 if client.is_via_collab() {
928 client
929 .request(update)
930 .map(|result| result.log_err().is_some())
931 .await
932 } else {
933 client.send(update).log_err().is_some()
934 }
935 }
936 }
937 });
938 });
939
940 cx.emit(WorktreeStoreEvent::WorktreeUpdateSent(worktree.clone()))
941 }
942
943 anyhow::Ok(())
944 })
945 })
946 .detach_and_log_err(cx);
947 }
948
949 pub fn worktree_metadata_protos(&self, cx: &App) -> Vec<proto::WorktreeMetadata> {
950 self.worktrees()
951 .map(|worktree| {
952 let worktree = worktree.read(cx);
953 proto::WorktreeMetadata {
954 id: worktree.id().to_proto(),
955 root_name: worktree.root_name_str().to_owned(),
956 visible: worktree.is_visible(),
957 abs_path: worktree.abs_path().to_string_lossy().into_owned(),
958 }
959 })
960 .collect()
961 }
962
963 pub fn shared(
964 &mut self,
965 remote_id: u64,
966 downstream_client: AnyProtoClient,
967 cx: &mut Context<Self>,
968 ) {
969 self.retain_worktrees = true;
970 self.downstream_client = Some((downstream_client, remote_id));
971
972 // When shared, retain all worktrees
973 for worktree_handle in self.worktrees.iter_mut() {
974 match worktree_handle {
975 WorktreeHandle::Strong(_) => {}
976 WorktreeHandle::Weak(worktree) => {
977 if let Some(worktree) = worktree.upgrade() {
978 *worktree_handle = WorktreeHandle::Strong(worktree);
979 }
980 }
981 }
982 }
983 // Only send project updates if we share in a collaborative mode.
984 // Otherwise we are the remote server which is currently constructing
985 // worktree store before the client actually has set up its message
986 // handlers.
987 if remote_id != REMOTE_SERVER_PROJECT_ID {
988 self.send_project_updates(cx);
989 }
990 }
991
992 pub fn unshared(&mut self, cx: &mut Context<Self>) {
993 self.retain_worktrees = false;
994 self.downstream_client.take();
995
996 // When not shared, only retain the visible worktrees
997 for worktree_handle in self.worktrees.iter_mut() {
998 if let WorktreeHandle::Strong(worktree) = worktree_handle {
999 let is_visible = worktree.update(cx, |worktree, _| {
1000 worktree.stop_observing_updates();
1001 worktree.is_visible()
1002 });
1003 if !is_visible {
1004 *worktree_handle = WorktreeHandle::Weak(worktree.downgrade());
1005 }
1006 }
1007 }
1008 }
1009
1010 pub async fn handle_create_project_entry(
1011 this: Entity<Self>,
1012 envelope: TypedEnvelope<proto::CreateProjectEntry>,
1013 mut cx: AsyncApp,
1014 ) -> Result<proto::ProjectEntryResponse> {
1015 let worktree = this.update(&mut cx, |this, cx| {
1016 let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id);
1017 this.worktree_for_id(worktree_id, cx)
1018 .context("worktree not found")
1019 })?;
1020 Worktree::handle_create_entry(worktree, envelope.payload, cx).await
1021 }
1022
1023 pub async fn handle_copy_project_entry(
1024 this: Entity<Self>,
1025 envelope: TypedEnvelope<proto::CopyProjectEntry>,
1026 mut cx: AsyncApp,
1027 ) -> Result<proto::ProjectEntryResponse> {
1028 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
1029 let new_worktree_id = WorktreeId::from_proto(envelope.payload.new_worktree_id);
1030 let new_project_path = (
1031 new_worktree_id,
1032 RelPath::from_proto(&envelope.payload.new_path)?,
1033 );
1034 let (scan_id, entry) = this.update(&mut cx, |this, cx| {
1035 let Some((_, project_id)) = this.downstream_client else {
1036 bail!("no downstream client")
1037 };
1038 let Some(entry) = this.entry_for_id(entry_id, cx) else {
1039 bail!("no such entry");
1040 };
1041 if entry.is_private && project_id != REMOTE_SERVER_PROJECT_ID {
1042 bail!("entry is private")
1043 }
1044
1045 let new_worktree = this
1046 .worktree_for_id(new_worktree_id, cx)
1047 .context("no such worktree")?;
1048 let scan_id = new_worktree.read(cx).scan_id();
1049 anyhow::Ok((
1050 scan_id,
1051 this.copy_entry(entry_id, new_project_path.into(), cx),
1052 ))
1053 })?;
1054 let entry = entry.await?;
1055 Ok(proto::ProjectEntryResponse {
1056 entry: entry.as_ref().map(|entry| entry.into()),
1057 worktree_scan_id: scan_id as u64,
1058 })
1059 }
1060
1061 pub async fn handle_delete_project_entry(
1062 this: Entity<Self>,
1063 envelope: TypedEnvelope<proto::DeleteProjectEntry>,
1064 mut cx: AsyncApp,
1065 ) -> Result<proto::ProjectEntryResponse> {
1066 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
1067 let worktree = this.update(&mut cx, |this, cx| {
1068 let Some((_, project_id)) = this.downstream_client else {
1069 bail!("no downstream client")
1070 };
1071 let Some(entry) = this.entry_for_id(entry_id, cx) else {
1072 bail!("no entry")
1073 };
1074 if entry.is_private && project_id != REMOTE_SERVER_PROJECT_ID {
1075 bail!("entry is private")
1076 }
1077 this.worktree_for_entry(entry_id, cx)
1078 .context("worktree not found")
1079 })?;
1080 Worktree::handle_delete_entry(worktree, envelope.payload, cx).await
1081 }
1082
1083 pub async fn handle_rename_project_entry(
1084 this: Entity<Self>,
1085 request: proto::RenameProjectEntry,
1086 mut cx: AsyncApp,
1087 ) -> Result<proto::ProjectEntryResponse> {
1088 let entry_id = ProjectEntryId::from_proto(request.entry_id);
1089 let new_worktree_id = WorktreeId::from_proto(request.new_worktree_id);
1090 let rel_path = RelPath::from_proto(&request.new_path)
1091 .with_context(|| format!("received invalid relative path {:?}", &request.new_path))?;
1092
1093 let (scan_id, task) = this.update(&mut cx, |this, cx| {
1094 let worktree = this
1095 .worktree_for_entry(entry_id, cx)
1096 .context("no such worktree")?;
1097
1098 let Some((_, project_id)) = this.downstream_client else {
1099 bail!("no downstream client")
1100 };
1101 let entry = worktree
1102 .read(cx)
1103 .entry_for_id(entry_id)
1104 .ok_or_else(|| anyhow!("missing entry"))?;
1105 if entry.is_private && project_id != REMOTE_SERVER_PROJECT_ID {
1106 bail!("entry is private")
1107 }
1108
1109 let scan_id = worktree.read(cx).scan_id();
1110 anyhow::Ok((
1111 scan_id,
1112 this.rename_entry(entry_id, (new_worktree_id, rel_path).into(), cx),
1113 ))
1114 })?;
1115 Ok(proto::ProjectEntryResponse {
1116 entry: match &task.await? {
1117 CreatedEntry::Included(entry) => Some(entry.into()),
1118 CreatedEntry::Excluded { .. } => None,
1119 },
1120 worktree_scan_id: scan_id as u64,
1121 })
1122 }
1123
1124 pub async fn handle_expand_project_entry(
1125 this: Entity<Self>,
1126 envelope: TypedEnvelope<proto::ExpandProjectEntry>,
1127 mut cx: AsyncApp,
1128 ) -> Result<proto::ExpandProjectEntryResponse> {
1129 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
1130 let worktree = this
1131 .update(&mut cx, |this, cx| this.worktree_for_entry(entry_id, cx))
1132 .context("invalid request")?;
1133 Worktree::handle_expand_entry(worktree, envelope.payload, cx).await
1134 }
1135
1136 pub async fn handle_expand_all_for_project_entry(
1137 this: Entity<Self>,
1138 envelope: TypedEnvelope<proto::ExpandAllForProjectEntry>,
1139 mut cx: AsyncApp,
1140 ) -> Result<proto::ExpandAllForProjectEntryResponse> {
1141 let entry_id = ProjectEntryId::from_proto(envelope.payload.entry_id);
1142 let worktree = this
1143 .update(&mut cx, |this, cx| this.worktree_for_entry(entry_id, cx))
1144 .context("invalid request")?;
1145 Worktree::handle_expand_all_for_entry(worktree, envelope.payload, cx).await
1146 }
1147
1148 pub async fn handle_allocate_worktree_id(
1149 _this: Entity<Self>,
1150 _envelope: TypedEnvelope<proto::AllocateWorktreeId>,
1151 cx: AsyncApp,
1152 ) -> Result<proto::AllocateWorktreeIdResponse> {
1153 let worktree_id = cx.update(|cx| WorktreeIdCounter::get(cx).next());
1154 Ok(proto::AllocateWorktreeIdResponse { worktree_id })
1155 }
1156
1157 pub fn fs(&self) -> Option<Arc<dyn Fs>> {
1158 match &self.state {
1159 WorktreeStoreState::Local { fs } => Some(fs.clone()),
1160 WorktreeStoreState::Remote { .. } => None,
1161 }
1162 }
1163}
1164
1165#[derive(Clone, Debug)]
1166enum WorktreeHandle {
1167 Strong(Entity<Worktree>),
1168 Weak(WeakEntity<Worktree>),
1169}
1170
1171impl WorktreeHandle {
1172 fn upgrade(&self) -> Option<Entity<Worktree>> {
1173 match self {
1174 WorktreeHandle::Strong(handle) => Some(handle.clone()),
1175 WorktreeHandle::Weak(handle) => handle.upgrade(),
1176 }
1177 }
1178}