diff --git a/server/src/rpc/store.rs b/server/src/rpc/store.rs index cd4fff53ec87a4239e75287ba0dfe5a3be4f4a3c..1d77e6b5effdee42cc800ab11467aa33e049ad4f 100644 --- a/server/src/rpc/store.rs +++ b/server/src/rpc/store.rs @@ -305,16 +305,20 @@ impl Store { } let connection_ids = worktree.connection_ids(); - - if let Some(_) = worktree.share.take() { - for connection_id in &connection_ids { - if let Some(connection) = self.connections.get_mut(connection_id) { + let collaborator_ids = worktree.collaborator_user_ids.clone(); + if let Some(share) = worktree.share.take() { + for connection_id in share.guest_connection_ids.into_keys() { + if let Some(connection) = self.connections.get_mut(&connection_id) { connection.worktrees.remove(&worktree_id); } } + + #[cfg(test)] + self.check_invariants(); + Ok(UnsharedWorktree { connection_ids, - collaborator_ids: worktree.collaborator_user_ids.clone(), + collaborator_ids, }) } else { Err(anyhow!("worktree is not shared"))? @@ -352,9 +356,13 @@ impl Store { } share.active_replica_ids.insert(replica_id); share.guest_connection_ids.insert(connection_id, replica_id); + + #[cfg(test)] + self.check_invariants(); + Ok(JoinedWorktree { replica_id, - worktree, + worktree: &self.worktrees[&worktree_id], }) } @@ -367,9 +375,19 @@ impl Store { let share = worktree.share.as_mut()?; let replica_id = share.guest_connection_ids.remove(&connection_id)?; share.active_replica_ids.remove(&replica_id); + + let connection = self.connections.get_mut(&connection_id)?; + connection.worktrees.remove(&worktree_id); + + let connection_ids = worktree.connection_ids(); + let collaborator_ids = worktree.collaborator_user_ids.clone(); + + #[cfg(test)] + self.check_invariants(); + Some(LeftWorktree { - connection_ids: worktree.connection_ids(), - collaborator_ids: worktree.collaborator_user_ids.clone(), + connection_ids, + collaborator_ids, }) }