toolchains: Fix persistence by not relying on unstable worktree id (#45357)

Piotr Osiewicz and Dino created

Closes #42268
We've migrated user selections when a given workspace has a single
worktree (as then we could determine what the target worktree is).

Release Notes:

- python: Fixed selected virtual environments not being
persisted/deserialized correctly within long-running Zed sessions (where
multiple different projects might've been opened). This is a breaking
change for users of multi-worktree projects - your selected toolchain
for those projects will be reset.

Co-authored-by: Dino <dino@zed.dev>

Change summary

crates/language/src/toolchain.rs                    |  7 
crates/project/src/project.rs                       |  7 +
crates/project/src/toolchain_store.rs               | 26 ++++
crates/project/src/x.py                             |  1 
crates/toolchain_selector/src/active_toolchain.rs   |  9 +
crates/toolchain_selector/src/toolchain_selector.rs | 43 ++++++-
crates/workspace/src/persistence.rs                 | 81 +++++++++++---
crates/workspace/src/workspace.rs                   | 31 +++++
8 files changed, 165 insertions(+), 40 deletions(-)

Detailed changes

crates/language/src/toolchain.rs 🔗

@@ -4,7 +4,10 @@
 //! which is a set of tools used to interact with the projects written in said language.
 //! For example, a Python project can have an associated virtual environment; a Rust project can have a toolchain override.
 
-use std::{path::PathBuf, sync::Arc};
+use std::{
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 
 use async_trait::async_trait;
 use collections::HashMap;
@@ -36,7 +39,7 @@ pub struct Toolchain {
 /// - Only in the subproject they're currently in.
 #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
 pub enum ToolchainScope {
-    Subproject(WorktreeId, Arc<RelPath>),
+    Subproject(Arc<Path>, Arc<RelPath>),
     Project,
     /// Available in all projects on this box. It wouldn't make sense to show suggestions across machines.
     Global,

crates/project/src/project.rs 🔗

@@ -1330,7 +1330,12 @@ impl Project {
             cx.subscribe(&buffer_store, Self::on_buffer_store_event)
                 .detach();
             let toolchain_store = cx.new(|cx| {
-                ToolchainStore::remote(REMOTE_SERVER_PROJECT_ID, remote.read(cx).proto_client(), cx)
+                ToolchainStore::remote(
+                    REMOTE_SERVER_PROJECT_ID,
+                    worktree_store.clone(),
+                    remote.read(cx).proto_client(),
+                    cx,
+                )
             });
             let task_store = cx.new(|cx| {
                 TaskStore::remote(

crates/project/src/toolchain_store.rs 🔗

@@ -32,6 +32,7 @@ use crate::{
 pub struct ToolchainStore {
     mode: ToolchainStoreInner,
     user_toolchains: BTreeMap<ToolchainScope, IndexSet<Toolchain>>,
+    worktree_store: Entity<WorktreeStore>,
     _sub: Subscription,
 }
 
@@ -66,7 +67,7 @@ impl ToolchainStore {
     ) -> Self {
         let entity = cx.new(|_| LocalToolchainStore {
             languages,
-            worktree_store,
+            worktree_store: worktree_store.clone(),
             project_environment,
             active_toolchains: Default::default(),
             manifest_tree,
@@ -77,12 +78,18 @@ impl ToolchainStore {
         });
         Self {
             mode: ToolchainStoreInner::Local(entity),
+            worktree_store,
             user_toolchains: Default::default(),
             _sub,
         }
     }
 
-    pub(super) fn remote(project_id: u64, client: AnyProtoClient, cx: &mut Context<Self>) -> Self {
+    pub(super) fn remote(
+        project_id: u64,
+        worktree_store: Entity<WorktreeStore>,
+        client: AnyProtoClient,
+        cx: &mut Context<Self>,
+    ) -> Self {
         let entity = cx.new(|_| RemoteToolchainStore { client, project_id });
         let _sub = cx.subscribe(&entity, |_, _, e: &ToolchainStoreEvent, cx| {
             cx.emit(e.clone())
@@ -90,6 +97,7 @@ impl ToolchainStore {
         Self {
             mode: ToolchainStoreInner::Remote(entity),
             user_toolchains: Default::default(),
+            worktree_store,
             _sub,
         }
     }
@@ -165,12 +173,22 @@ impl ToolchainStore {
         language_name: LanguageName,
         cx: &mut Context<Self>,
     ) -> Task<Option<Toolchains>> {
+        let Some(worktree) = self
+            .worktree_store
+            .read(cx)
+            .worktree_for_id(path.worktree_id, cx)
+        else {
+            return Task::ready(None);
+        };
+        let target_root_path = worktree.read_with(cx, |this, _| this.abs_path());
+
         let user_toolchains = self
             .user_toolchains
             .iter()
             .filter(|(scope, _)| {
-                if let ToolchainScope::Subproject(worktree_id, relative_path) = scope {
-                    path.worktree_id == *worktree_id && relative_path.starts_with(&path.path)
+                if let ToolchainScope::Subproject(subproject_root_path, relative_path) = scope {
+                    target_root_path == *subproject_root_path
+                        && relative_path.starts_with(&path.path)
                 } else {
                     true
                 }

crates/toolchain_selector/src/active_toolchain.rs 🔗

@@ -198,10 +198,17 @@ impl ActiveToolchain {
                     .or_else(|| toolchains.toolchains.first())
                     .cloned();
                 if let Some(toolchain) = &default_choice {
+                    let worktree_root_path = project
+                        .read_with(cx, |this, cx| {
+                            this.worktree_for_id(worktree_id, cx)
+                                .map(|worktree| worktree.read(cx).abs_path())
+                        })
+                        .ok()
+                        .flatten()?;
                     workspace::WORKSPACE_DB
                         .set_toolchain(
                             workspace_id,
-                            worktree_id,
+                            worktree_root_path,
                             relative_path.clone(),
                             toolchain.clone(),
                         )

crates/toolchain_selector/src/toolchain_selector.rs 🔗

@@ -1,6 +1,7 @@
 mod active_toolchain;
 
 pub use active_toolchain::ActiveToolchain;
+use anyhow::Context as _;
 use convert_case::Casing as _;
 use editor::Editor;
 use file_finder::OpenPathDelegate;
@@ -62,6 +63,7 @@ struct AddToolchainState {
     language_name: LanguageName,
     root_path: ProjectPath,
     weak: WeakEntity<ToolchainSelector>,
+    worktree_root_path: Arc<Path>,
 }
 
 struct ScopePickerState {
@@ -99,12 +101,17 @@ impl AddToolchainState {
         root_path: ProjectPath,
         window: &mut Window,
         cx: &mut Context<ToolchainSelector>,
-    ) -> Entity<Self> {
+    ) -> anyhow::Result<Entity<Self>> {
         let weak = cx.weak_entity();
-
-        cx.new(|cx| {
+        let worktree_root_path = project
+            .read(cx)
+            .worktree_for_id(root_path.worktree_id, cx)
+            .map(|worktree| worktree.read(cx).abs_path())
+            .context("Could not find worktree")?;
+        Ok(cx.new(|cx| {
             let (lister, rx) = Self::create_path_browser_delegate(project.clone(), cx);
             let picker = cx.new(|cx| Picker::uniform_list(lister, window, cx));
+
             Self {
                 state: AddState::Path {
                     _subscription: cx.subscribe(&picker, |_, _, _: &DismissEvent, cx| {
@@ -118,8 +125,9 @@ impl AddToolchainState {
                 language_name,
                 root_path,
                 weak,
+                worktree_root_path,
             }
-        })
+        }))
     }
 
     fn create_path_browser_delegate(
@@ -237,7 +245,15 @@ impl AddToolchainState {
                 // Suggest a default scope based on the applicability.
                 let scope = if let Some(project_path) = resolved_toolchain_path {
                     if !root_path.path.as_ref().is_empty() && project_path.starts_with(&root_path) {
-                        ToolchainScope::Subproject(root_path.worktree_id, root_path.path)
+                        let worktree_root_path = project
+                            .read_with(cx, |this, cx| {
+                                this.worktree_for_id(root_path.worktree_id, cx)
+                                    .map(|worktree| worktree.read(cx).abs_path())
+                            })
+                            .ok()
+                            .flatten()
+                            .context("Could not find a worktree with a given worktree ID")?;
+                        ToolchainScope::Subproject(worktree_root_path, root_path.path)
                     } else {
                         ToolchainScope::Project
                     }
@@ -400,7 +416,7 @@ impl Render for AddToolchainState {
                         ToolchainScope::Global,
                         ToolchainScope::Project,
                         ToolchainScope::Subproject(
-                            self.root_path.worktree_id,
+                            self.worktree_root_path.clone(),
                             self.root_path.path.clone(),
                         ),
                     ];
@@ -693,7 +709,7 @@ impl ToolchainSelector {
         cx: &mut Context<Self>,
     ) {
         if matches!(self.state, State::Search(_)) {
-            self.state = State::AddToolchain(AddToolchainState::new(
+            let Ok(state) = AddToolchainState::new(
                 self.project.clone(),
                 self.language_name.clone(),
                 ProjectPath {
@@ -702,7 +718,10 @@ impl ToolchainSelector {
                 },
                 window,
                 cx,
-            ));
+            ) else {
+                return;
+            };
+            self.state = State::AddToolchain(state);
             self.state.focus_handle(cx).focus(window, cx);
             cx.notify();
         }
@@ -899,11 +918,17 @@ impl PickerDelegate for ToolchainSelectorDelegate {
             {
                 let workspace = self.workspace.clone();
                 let worktree_id = self.worktree_id;
+                let worktree_abs_path_root = self.worktree_abs_path_root.clone();
                 let path = self.relative_path.clone();
                 let relative_path = self.relative_path.clone();
                 cx.spawn_in(window, async move |_, cx| {
                     workspace::WORKSPACE_DB
-                        .set_toolchain(workspace_id, worktree_id, relative_path, toolchain.clone())
+                        .set_toolchain(
+                            workspace_id,
+                            worktree_abs_path_root,
+                            relative_path,
+                            toolchain.clone(),
+                        )
                         .await
                         .log_err();
                     workspace

crates/workspace/src/persistence.rs 🔗

@@ -24,7 +24,6 @@ use project::{
 };
 
 use language::{LanguageName, Toolchain, ToolchainScope};
-use project::WorktreeId;
 use remote::{
     DockerConnectionOptions, RemoteConnectionOptions, SshConnectionOptions, WslConnectionOptions,
 };
@@ -845,6 +844,44 @@ impl Domain for WorkspaceDb {
                 host_name TEXT
             ) STRICT;
         ),
+        sql!(CREATE TABLE toolchains2 (
+            workspace_id INTEGER,
+            worktree_root_path TEXT NOT NULL,
+            language_name TEXT NOT NULL,
+            name TEXT NOT NULL,
+            path TEXT NOT NULL,
+            raw_json TEXT NOT NULL,
+            relative_worktree_path TEXT NOT NULL,
+            PRIMARY KEY (workspace_id, worktree_root_path, language_name, relative_worktree_path)) STRICT;
+            INSERT OR REPLACE INTO toolchains2
+                // The `instr(paths, '\n') = 0` part allows us to find all
+                // workspaces that have a single worktree, as `\n` is used as a
+                // separator when serializing the workspace paths, so if no `\n` is
+                // found, we know we have a single worktree.
+                SELECT toolchains.workspace_id, paths, language_name, name, path, raw_json, relative_worktree_path FROM toolchains INNER JOIN workspaces ON toolchains.workspace_id = workspaces.workspace_id AND instr(paths, '\n') = 0;
+            DROP TABLE toolchains;
+            ALTER TABLE toolchains2 RENAME TO toolchains;
+        ),
+        sql!(CREATE TABLE user_toolchains2 (
+            remote_connection_id INTEGER,
+            workspace_id INTEGER NOT NULL,
+            worktree_root_path TEXT NOT NULL,
+            relative_worktree_path TEXT NOT NULL,
+            language_name TEXT NOT NULL,
+            name TEXT NOT NULL,
+            path TEXT NOT NULL,
+            raw_json TEXT NOT NULL,
+
+            PRIMARY KEY (workspace_id, worktree_root_path, relative_worktree_path, language_name, name, path, raw_json)) STRICT;
+            INSERT OR REPLACE INTO user_toolchains2
+                // The `instr(paths, '\n') = 0` part allows us to find all
+                // workspaces that have a single worktree, as `\n` is used as a
+                // separator when serializing the workspace paths, so if no `\n` is
+                // found, we know we have a single worktree.
+                SELECT user_toolchains.remote_connection_id, user_toolchains.workspace_id, paths, relative_worktree_path, language_name, name, path, raw_json  FROM user_toolchains INNER JOIN workspaces ON user_toolchains.workspace_id = workspaces.workspace_id AND instr(paths, '\n') = 0;
+            DROP TABLE user_toolchains;
+            ALTER TABLE user_toolchains2 RENAME TO user_toolchains;
+        ),
     ];
 
     // Allow recovering from bad migration that was initially shipped to nightly
@@ -1030,11 +1067,11 @@ impl WorkspaceDb {
         workspace_id: WorkspaceId,
         remote_connection_id: Option<RemoteConnectionId>,
     ) -> BTreeMap<ToolchainScope, IndexSet<Toolchain>> {
-        type RowKind = (WorkspaceId, u64, String, String, String, String, String);
+        type RowKind = (WorkspaceId, String, String, String, String, String, String);
 
         let toolchains: Vec<RowKind> = self
             .select_bound(sql! {
-                SELECT workspace_id, worktree_id, relative_worktree_path,
+                SELECT workspace_id, worktree_root_path, relative_worktree_path,
                 language_name, name, path, raw_json
                 FROM user_toolchains WHERE remote_connection_id IS ?1 AND (
                       workspace_id IN (0, ?2)
@@ -1048,7 +1085,7 @@ impl WorkspaceDb {
 
         for (
             _workspace_id,
-            worktree_id,
+            worktree_root_path,
             relative_worktree_path,
             language_name,
             name,
@@ -1058,22 +1095,24 @@ impl WorkspaceDb {
         {
             // INTEGER's that are primary keys (like workspace ids, remote connection ids and such) start at 1, so we're safe to
             let scope = if _workspace_id == WorkspaceId(0) {
-                debug_assert_eq!(worktree_id, u64::MAX);
+                debug_assert_eq!(worktree_root_path, String::default());
                 debug_assert_eq!(relative_worktree_path, String::default());
                 ToolchainScope::Global
             } else {
                 debug_assert_eq!(workspace_id, _workspace_id);
                 debug_assert_eq!(
-                    worktree_id == u64::MAX,
+                    worktree_root_path == String::default(),
                     relative_worktree_path == String::default()
                 );
 
                 let Some(relative_path) = RelPath::unix(&relative_worktree_path).log_err() else {
                     continue;
                 };
-                if worktree_id != u64::MAX && relative_worktree_path != String::default() {
+                if worktree_root_path != String::default()
+                    && relative_worktree_path != String::default()
+                {
                     ToolchainScope::Subproject(
-                        WorktreeId::from_usize(worktree_id as usize),
+                        Arc::from(worktree_root_path.as_ref()),
                         relative_path.into(),
                     )
                 } else {
@@ -1159,13 +1198,13 @@ impl WorkspaceDb {
 
                 for (scope, toolchains) in workspace.user_toolchains {
                     for toolchain in toolchains {
-                        let query = sql!(INSERT OR REPLACE INTO user_toolchains(remote_connection_id, workspace_id, worktree_id, relative_worktree_path, language_name, name, path, raw_json) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8));
-                        let (workspace_id, worktree_id, relative_worktree_path) = match scope {
-                            ToolchainScope::Subproject(worktree_id, ref path) => (Some(workspace.id), Some(worktree_id), Some(path.as_unix_str().to_owned())),
+                        let query = sql!(INSERT OR REPLACE INTO user_toolchains(remote_connection_id, workspace_id, worktree_root_path, relative_worktree_path, language_name, name, path, raw_json) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8));
+                        let (workspace_id, worktree_root_path, relative_worktree_path) = match scope {
+                            ToolchainScope::Subproject(ref worktree_root_path, ref path) => (Some(workspace.id), Some(worktree_root_path.to_string_lossy().into_owned()), Some(path.as_unix_str().to_owned())),
                             ToolchainScope::Project => (Some(workspace.id), None, None),
                             ToolchainScope::Global => (None, None, None),
                         };
-                        let args = (remote_connection_id, workspace_id.unwrap_or(WorkspaceId(0)), worktree_id.map_or(usize::MAX,|id| id.to_usize()), relative_worktree_path.unwrap_or_default(),
+                        let args = (remote_connection_id, workspace_id.unwrap_or(WorkspaceId(0)), worktree_root_path.unwrap_or_default(), relative_worktree_path.unwrap_or_default(),
                         toolchain.language_name.as_ref().to_owned(), toolchain.name.to_string(), toolchain.path.to_string(), toolchain.as_json.to_string());
                         if let Err(err) = conn.exec_bound(query)?(args) {
                             log::error!("{err}");
@@ -1844,24 +1883,24 @@ impl WorkspaceDb {
     pub(crate) async fn toolchains(
         &self,
         workspace_id: WorkspaceId,
-    ) -> Result<Vec<(Toolchain, WorktreeId, Arc<RelPath>)>> {
+    ) -> Result<Vec<(Toolchain, Arc<Path>, Arc<RelPath>)>> {
         self.write(move |this| {
             let mut select = this
                 .select_bound(sql!(
                     SELECT
-                        name, path, worktree_id, relative_worktree_path, language_name, raw_json
+                        name, path, worktree_root_path, relative_worktree_path, language_name, raw_json
                     FROM toolchains
                     WHERE workspace_id = ?
                 ))
                 .context("select toolchains")?;
 
-            let toolchain: Vec<(String, String, u64, String, String, String)> =
+            let toolchain: Vec<(String, String, String, String, String, String)> =
                 select(workspace_id)?;
 
             Ok(toolchain
                 .into_iter()
                 .filter_map(
-                    |(name, path, worktree_id, relative_worktree_path, language, json)| {
+                    |(name, path, worktree_root_path, relative_worktree_path, language, json)| {
                         Some((
                             Toolchain {
                                 name: name.into(),
@@ -1869,7 +1908,7 @@ impl WorkspaceDb {
                                 language_name: LanguageName::new(&language),
                                 as_json: serde_json::Value::from_str(&json).ok()?,
                             },
-                            WorktreeId::from_proto(worktree_id),
+                           Arc::from(worktree_root_path.as_ref()),
                             RelPath::from_proto(&relative_worktree_path).log_err()?,
                         ))
                     },
@@ -1882,18 +1921,18 @@ impl WorkspaceDb {
     pub async fn set_toolchain(
         &self,
         workspace_id: WorkspaceId,
-        worktree_id: WorktreeId,
+        worktree_root_path: Arc<Path>,
         relative_worktree_path: Arc<RelPath>,
         toolchain: Toolchain,
     ) -> Result<()> {
         log::debug!(
-            "Setting toolchain for workspace, worktree: {worktree_id:?}, relative path: {relative_worktree_path:?}, toolchain: {}",
+            "Setting toolchain for workspace, worktree: {worktree_root_path:?}, relative path: {relative_worktree_path:?}, toolchain: {}",
             toolchain.name
         );
         self.write(move |conn| {
             let mut insert = conn
                 .exec_bound(sql!(
-                    INSERT INTO toolchains(workspace_id, worktree_id, relative_worktree_path, language_name, name, path, raw_json) VALUES (?, ?, ?, ?, ?,  ?, ?)
+                    INSERT INTO toolchains(workspace_id, worktree_root_path, relative_worktree_path, language_name, name, path, raw_json) VALUES (?, ?, ?, ?, ?,  ?, ?)
                     ON CONFLICT DO
                     UPDATE SET
                         name = ?5,
@@ -1904,7 +1943,7 @@ impl WorkspaceDb {
 
             insert((
                 workspace_id,
-                worktree_id.to_usize(),
+                worktree_root_path.to_string_lossy().into_owned(),
                 relative_worktree_path.as_unix_str(),
                 toolchain.language_name.as_ref(),
                 toolchain.name.as_ref(),

crates/workspace/src/workspace.rs 🔗

@@ -1697,8 +1697,22 @@ impl Workspace {
 
             let toolchains = DB.toolchains(workspace_id).await?;
 
-            for (toolchain, worktree_id, path) in toolchains {
+            for (toolchain, worktree_path, path) in toolchains {
                 let toolchain_path = PathBuf::from(toolchain.path.clone().to_string());
+                let Some(worktree_id) = project_handle.read_with(cx, |this, cx| {
+                    this.find_worktree(&worktree_path, cx)
+                        .and_then(|(worktree, rel_path)| {
+                            if rel_path.is_empty() {
+                                Some(worktree.read(cx).id())
+                            } else {
+                                None
+                            }
+                        })
+                })?
+                else {
+                    // We did not find a worktree with a given path, but that's whatever.
+                    continue;
+                };
                 if !app_state.fs.is_file(toolchain_path.as_path()).await {
                     continue;
                 }
@@ -8217,9 +8231,22 @@ async fn open_remote_project_inner(
     cx: &mut AsyncApp,
 ) -> Result<Vec<Option<Box<dyn ItemHandle>>>> {
     let toolchains = DB.toolchains(workspace_id).await?;
-    for (toolchain, worktree_id, path) in toolchains {
+    for (toolchain, worktree_path, path) in toolchains {
         project
             .update(cx, |this, cx| {
+                let Some(worktree_id) =
+                    this.find_worktree(&worktree_path, cx)
+                        .and_then(|(worktree, rel_path)| {
+                            if rel_path.is_empty() {
+                                Some(worktree.read(cx).id())
+                            } else {
+                                None
+                            }
+                        })
+                else {
+                    return Task::ready(None);
+                };
+
                 this.activate_toolchain(ProjectPath { worktree_id, path }, toolchain, cx)
             })?
             .await;