From 98ea659af6e5be39e3360093aa6e66c1ea801c1a Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Tue, 18 Feb 2025 17:44:43 -0500 Subject: [PATCH] assistant2: Fix thread history only working in one Zed window (#25119) This PR fixes an issue where the thread history would only work in one Zed window at a time. The backing LMDB database can only be opened once per Zed instance. However, the `ThreadStore` has one instance per Zed window. To fix this, we need to create the `heed` environment once and store it as a global, and then reference the same environment across all of the `ThreadStore`s. Release Notes: - N/A --- crates/assistant2/src/assistant.rs | 1 + crates/assistant2/src/assistant_panel.rs | 3 ++ crates/assistant2/src/thread_store.rs | 60 ++++++++++++++++-------- 3 files changed, 44 insertions(+), 20 deletions(-) diff --git a/crates/assistant2/src/assistant.rs b/crates/assistant2/src/assistant.rs index 0d1303ea90..30127a59d5 100644 --- a/crates/assistant2/src/assistant.rs +++ b/crates/assistant2/src/assistant.rs @@ -66,6 +66,7 @@ pub fn init( cx: &mut App, ) { AssistantSettings::register(cx); + thread_store::init(cx); assistant_panel::init(cx); inline_assistant::init( diff --git a/crates/assistant2/src/assistant_panel.rs b/crates/assistant2/src/assistant_panel.rs index e12100907f..f20436146e 100644 --- a/crates/assistant2/src/assistant_panel.rs +++ b/crates/assistant2/src/assistant_panel.rs @@ -323,6 +323,9 @@ impl AssistantPanel { } fn open_history(&mut self, window: &mut Window, cx: &mut Context) { + self.thread_store + .update(cx, |thread_store, cx| thread_store.reload(cx)) + .detach_and_log_err(cx); self.active_view = ActiveView::History; self.history.focus_handle(cx).focus(window); cx.notify(); diff --git a/crates/assistant2/src/thread_store.rs b/crates/assistant2/src/thread_store.rs index 3e7ee82da9..58caad37b6 100644 --- a/crates/assistant2/src/thread_store.rs +++ b/crates/assistant2/src/thread_store.rs @@ -9,7 +9,9 @@ use context_server::manager::ContextServerManager; use context_server::{ContextServerFactoryRegistry, ContextServerTool}; use futures::future::{self, BoxFuture, Shared}; use futures::FutureExt as _; -use gpui::{prelude::*, App, BackgroundExecutor, Context, Entity, SharedString, Task}; +use gpui::{ + prelude::*, App, BackgroundExecutor, Context, Entity, Global, ReadGlobal, SharedString, Task, +}; use heed::types::SerdeBincode; use heed::Database; use language_model::Role; @@ -19,6 +21,10 @@ use util::ResultExt as _; use crate::thread::{MessageId, Thread, ThreadId}; +pub fn init(cx: &mut App) { + ThreadsDatabase::init(cx); +} + pub struct ThreadStore { #[allow(unused)] project: Entity, @@ -26,7 +32,6 @@ pub struct ThreadStore { context_server_manager: Entity, context_server_tool_ids: HashMap, Vec>, threads: Vec, - database_future: Shared, Arc>>>, } impl ThreadStore { @@ -41,24 +46,12 @@ impl ThreadStore { ContextServerManager::new(context_server_factory_registry, project.clone(), cx) }); - let executor = cx.background_executor().clone(); - let database_future = executor - .spawn({ - let executor = executor.clone(); - let database_path = paths::support_dir().join("threads/threads-db.0.mdb"); - async move { ThreadsDatabase::new(database_path, executor) } - }) - .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new))) - .boxed() - .shared(); - let this = Self { project, tools, context_server_manager, context_server_tool_ids: HashMap::default(), threads: Vec::new(), - database_future, }; this.register_context_server_handlers(cx); this.reload(cx).detach_and_log_err(cx); @@ -94,7 +87,7 @@ impl ThreadStore { cx: &mut Context, ) -> Task>> { let id = id.clone(); - let database_future = self.database_future.clone(); + let database_future = ThreadsDatabase::global_future(cx); cx.spawn(|this, mut cx| async move { let database = database_future.await.map_err(|err| anyhow!(err))?; let thread = database @@ -127,7 +120,7 @@ impl ThreadStore { (id, thread) }); - let database_future = self.database_future.clone(); + let database_future = ThreadsDatabase::global_future(cx); cx.spawn(|this, mut cx| async move { let database = database_future.await.map_err(|err| anyhow!(err))?; database.save_thread(metadata, thread).await?; @@ -138,7 +131,7 @@ impl ThreadStore { pub fn delete_thread(&mut self, id: &ThreadId, cx: &mut Context) -> Task> { let id = id.clone(); - let database_future = self.database_future.clone(); + let database_future = ThreadsDatabase::global_future(cx); cx.spawn(|this, mut cx| async move { let database = database_future.await.map_err(|err| anyhow!(err))?; database.delete_thread(id.clone()).await?; @@ -149,8 +142,8 @@ impl ThreadStore { }) } - fn reload(&self, cx: &mut Context) -> Task> { - let database_future = self.database_future.clone(); + pub fn reload(&self, cx: &mut Context) -> Task> { + let database_future = ThreadsDatabase::global_future(cx); cx.spawn(|this, mut cx| async move { let threads = database_future .await @@ -253,13 +246,40 @@ pub struct SavedMessage { pub text: String, } -struct ThreadsDatabase { +struct GlobalThreadsDatabase( + Shared, Arc>>>, +); + +impl Global for GlobalThreadsDatabase {} + +pub(crate) struct ThreadsDatabase { executor: BackgroundExecutor, env: heed::Env, threads: Database, SerdeBincode>, } impl ThreadsDatabase { + fn global_future( + cx: &mut App, + ) -> Shared, Arc>>> { + GlobalThreadsDatabase::global(cx).0.clone() + } + + fn init(cx: &mut App) { + let executor = cx.background_executor().clone(); + let database_future = executor + .spawn({ + let executor = executor.clone(); + let database_path = paths::support_dir().join("threads/threads-db.0.mdb"); + async move { ThreadsDatabase::new(database_path, executor) } + }) + .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new))) + .boxed() + .shared(); + + cx.set_global(GlobalThreadsDatabase(database_future)); + } + pub fn new(path: PathBuf, executor: BackgroundExecutor) -> Result { std::fs::create_dir_all(&path)?;