Improve context server lifecycle management (#20622)

This optimizes and fixes bugs in our logic for maintaining a set of
running context servers, based on the combination of the user's
`context_servers` settings and their installed extensions.

Release Notes:

- N/A

---------

Co-authored-by: Marshall <[email protected]>
Co-authored-by: Marshall Bowers <[email protected]>

crates/assistant/src/context_store.rs

@@ -8,9 +8,8 @@ use anyhow::{anyhow, Context as _, Result};
 use client::{proto, telemetry::Telemetry, Client, TypedEnvelope};
 use clock::ReplicaId;
 use collections::HashMap;
-use command_palette_hooks::CommandPaletteFilter;
-use context_servers::manager::{ContextServerManager, ContextServerSettings};
-use context_servers::{ContextServerFactoryRegistry, CONTEXT_SERVERS_NAMESPACE};
+use context_servers::manager::ContextServerManager;
+use context_servers::ContextServerFactoryRegistry;
 use fs::Fs;
 use futures::StreamExt;
 use fuzzy::StringMatchCandidate;
@@ -22,7 +21,6 @@ use paths::contexts_dir;
 use project::Project;
 use regex::Regex;
 use rpc::AnyProtoClient;
-use settings::{Settings as _, SettingsStore};
 use std::{
     cmp::Reverse,
     ffi::OsStr,
@@ -111,7 +109,11 @@ impl ContextStore {
             let (mut events, _) = fs.watch(contexts_dir(), CONTEXT_WATCH_DURATION).await;
 
             let this = cx.new_model(|cx: &mut ModelContext<Self>| {
-                let context_server_manager = cx.new_model(|_cx| ContextServerManager::new());
+                let context_server_factory_registry =
+                    ContextServerFactoryRegistry::default_global(cx);
+                let context_server_manager = cx.new_model(|cx| {
+                    ContextServerManager::new(context_server_factory_registry, project.clone(), cx)
+                });
                 let mut this = Self {
                     contexts: Vec::new(),
                     contexts_metadata: Vec::new(),
@@ -148,91 +150,16 @@ impl ContextStore {
                 this.handle_project_changed(project.clone(), cx);
                 this.synchronize_contexts(cx);
                 this.register_context_server_handlers(cx);
-
-                if project.read(cx).is_local() {
-                    // TODO: At the time when we construct the `ContextStore` we may not have yet initialized the extensions.
-                    // In order to register the context servers when the extension is loaded, we're periodically looping to
-                    // see if there are context servers to register.
-                    //
-                    // I tried doing this in a subscription on the `ExtensionStore`, but it never seemed to fire.
-                    //
-                    // We should find a more elegant way to do this.
-                    let context_server_factory_registry =
-                        ContextServerFactoryRegistry::default_global(cx);
-                    cx.spawn(|context_store, mut cx| async move {
-                        loop {
-                            let mut servers_to_register = Vec::new();
-                            for (_id, factory) in
-                                context_server_factory_registry.context_server_factories()
-                            {
-                                if let Some(server) = factory(project.clone(), &cx).await.log_err()
-                                {
-                                    servers_to_register.push(server);
-                                }
-                            }
-
-                            let Some(_) = context_store
-                                .update(&mut cx, |this, cx| {
-                                    this.context_server_manager.update(cx, |this, cx| {
-                                        for server in servers_to_register {
-                                            this.add_server(server, cx).detach_and_log_err(cx);
-                                        }
-                                    })
-                                })
-                                .log_err()
-                            else {
-                                break;
-                            };
-
-                            smol::Timer::after(Duration::from_millis(100)).await;
-                        }
-
-                        anyhow::Ok(())
-                    })
-                    .detach_and_log_err(cx);
-                }
-
                 this
             })?;
             this.update(&mut cx, |this, cx| this.reload(cx))?
                 .await
                 .log_err();
 
-            this.update(&mut cx, |this, cx| {
-                this.watch_context_server_settings(cx);
-            })
-            .log_err();
-
             Ok(this)
         })
     }
 
-    fn watch_context_server_settings(&self, cx: &mut ModelContext<Self>) {
-        cx.observe_global::<SettingsStore>(move |this, cx| {
-            this.context_server_manager.update(cx, |manager, cx| {
-                let location = this.project.read(cx).worktrees(cx).next().map(|worktree| {
-                    settings::SettingsLocation {
-                        worktree_id: worktree.read(cx).id(),
-                        path: Path::new(""),
-                    }
-                });
-                let settings = ContextServerSettings::get(location, cx);
-
-                manager.maintain_servers(settings, cx);
-
-                let has_any_context_servers = !manager.servers().is_empty();
-                CommandPaletteFilter::update_global(cx, |filter, _cx| {
-                    if has_any_context_servers {
-                        filter.show_namespace(CONTEXT_SERVERS_NAMESPACE);
-                    } else {
-                        filter.hide_namespace(CONTEXT_SERVERS_NAMESPACE);
-                    }
-                });
-            })
-        })
-        .detach();
-    }
-
     async fn handle_advertise_contexts(
         this: Model<Self>,
         envelope: TypedEnvelope<proto::AdvertiseContexts>,

crates/assistant/src/slash_command/context_server_command.rs

@@ -27,7 +27,7 @@ pub struct ContextServerSlashCommand {
 impl ContextServerSlashCommand {
     pub fn new(
         server_manager: Model<ContextServerManager>,
-        server: &Arc<dyn ContextServer>,
+        server: &Arc<ContextServer>,
         prompt: Prompt,
     ) -> Self {
         Self {

crates/collab/Cargo.toml

@@ -78,6 +78,7 @@ uuid.workspace = true
 
 [dev-dependencies]
 assistant = { workspace = true, features = ["test-support"] }
+context_servers.workspace = true
 async-trait.workspace = true
 audio.workspace = true
 call = { workspace = true, features = ["test-support"] }

crates/collab/src/tests/integration_tests.rs

@@ -6486,6 +6486,8 @@ async fn test_context_collaboration_with_reconnect(
         assert_eq!(project.collaborators().len(), 1);
     });
 
+    cx_a.update(context_servers::init);
+    cx_b.update(context_servers::init);
     let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
     let context_store_a = cx_a
         .update(|cx| {

crates/command_palette_hooks/src/command_palette_hooks.rs

@@ -39,11 +39,13 @@ impl CommandPaletteFilter {
     }
 
     /// Updates the global [`CommandPaletteFilter`] using the given closure.
-    pub fn update_global<F, R>(cx: &mut AppContext, update: F) -> R
+    pub fn update_global<F>(cx: &mut AppContext, update: F)
     where
-        F: FnOnce(&mut Self, &mut AppContext) -> R,
+        F: FnOnce(&mut Self, &mut AppContext),
     {
-        cx.update_global(|this: &mut GlobalCommandPaletteFilter, cx| update(&mut this.0, cx))
+        if cx.has_global::<GlobalCommandPaletteFilter>() {
+            cx.update_global(|this: &mut GlobalCommandPaletteFilter, cx| update(&mut this.0, cx))
+        }
     }
 
     /// Returns whether the given [`Action`] is hidden by the filter.

crates/context_servers/Cargo.toml

@@ -13,7 +13,6 @@ path = "src/context_servers.rs"
 
 [dependencies]
 anyhow.workspace = true
-async-trait.workspace = true
 collections.workspace = true
 command_palette_hooks.workspace = true
 futures.workspace = true

crates/context_servers/src/context_servers.rs

@@ -8,7 +8,6 @@ use command_palette_hooks::CommandPaletteFilter;
 use gpui::{actions, AppContext};
 use settings::Settings;
 
-pub use crate::manager::ContextServer;
 use crate::manager::ContextServerSettings;
 pub use crate::registry::ContextServerFactoryRegistry;