assistant2: Add ability to start and stop context servers (#27080)

Marshall Bowers created

This PR adds the ability to start and stop context servers from within
the configuration view in the Assistant panel:


https://github.com/user-attachments/assets/93c3a7cb-d799-4286-88ba-c13cc26e959a

Release Notes:

- N/A

Change summary

crates/assistant2/src/assistant_configuration.rs     | 87 ++++++++++---
crates/assistant_context_editor/src/context_store.rs |  2 
crates/context_server/src/manager.rs                 | 33 ++++
3 files changed, 98 insertions(+), 24 deletions(-)

Detailed changes

crates/assistant2/src/assistant_configuration.rs 🔗

@@ -5,7 +5,8 @@ use collections::HashMap;
 use context_server::manager::ContextServerManager;
 use gpui::{Action, AnyView, App, Entity, EventEmitter, FocusHandle, Focusable, Subscription};
 use language_model::{LanguageModelProvider, LanguageModelProviderId, LanguageModelRegistry};
-use ui::{prelude::*, Disclosure, Divider, DividerColor, ElevationIndex, Indicator};
+use ui::{prelude::*, Disclosure, Divider, DividerColor, ElevationIndex, Indicator, Switch};
+use util::ResultExt as _;
 use zed_actions::assistant::DeployPromptLibrary;
 
 pub struct AssistantConfiguration {
@@ -158,7 +159,7 @@ impl AssistantConfiguration {
     }
 
     fn render_context_servers_section(&mut self, cx: &mut Context<Self>) -> impl IntoElement {
-        let context_servers = self.context_server_manager.read(cx).servers().clone();
+        let context_servers = self.context_server_manager.read(cx).all_servers().clone();
         let tools_by_source = self.tools.tools_by_source(cx);
         let empty = Vec::new();
 
@@ -197,7 +198,7 @@ impl AssistantConfiguration {
                     .bg(cx.theme().colors().editor_background)
                     .child(
                         h_flex()
-                            .gap_2()
+                            .justify_between()
                             .px_2()
                             .py_1()
                             .when(are_tools_expanded, |element| {
@@ -206,26 +207,70 @@ impl AssistantConfiguration {
                                     .border_color(cx.theme().colors().border)
                             })
                             .child(
-                                Disclosure::new("tool-list-disclosure", are_tools_expanded)
-                                    .on_click(cx.listener({
-                                        let context_server_id = context_server.id();
-                                        move |this, _event, _window, _cx| {
-                                            let is_open = this
-                                                .expanded_context_server_tools
-                                                .entry(context_server_id.clone())
-                                                .or_insert(false);
+                                h_flex()
+                                    .gap_2()
+                                    .child(
+                                        Disclosure::new("tool-list-disclosure", are_tools_expanded)
+                                            .on_click(cx.listener({
+                                                let context_server_id = context_server.id();
+                                                move |this, _event, _window, _cx| {
+                                                    let is_open = this
+                                                        .expanded_context_server_tools
+                                                        .entry(context_server_id.clone())
+                                                        .or_insert(false);
 
-                                            *is_open = !*is_open;
-                                        }
-                                    })),
+                                                    *is_open = !*is_open;
+                                                }
+                                            })),
+                                    )
+                                    .child(Indicator::dot().color(if is_running {
+                                        Color::Success
+                                    } else {
+                                        Color::Error
+                                    }))
+                                    .child(Label::new(context_server.id()))
+                                    .child(
+                                        Label::new(format!("{tool_count} tools"))
+                                            .color(Color::Muted),
+                                    ),
                             )
-                            .child(Indicator::dot().color(if is_running {
-                                Color::Success
-                            } else {
-                                Color::Error
-                            }))
-                            .child(Label::new(context_server.id()))
-                            .child(Label::new(format!("{tool_count} tools")).color(Color::Muted)),
+                            .child(h_flex().child(
+                                Switch::new("context-server-switch", is_running.into()).on_click({
+                                    let context_server_manager =
+                                        self.context_server_manager.clone();
+                                    let context_server = context_server.clone();
+                                    move |state, _window, cx| match state {
+                                        ToggleState::Unselected | ToggleState::Indeterminate => {
+                                            context_server_manager.update(cx, |this, cx| {
+                                                this.stop_server(context_server.clone(), cx)
+                                                    .log_err();
+                                            });
+                                        }
+                                        ToggleState::Selected => {
+                                            cx.spawn({
+                                                let context_server_manager =
+                                                    context_server_manager.clone();
+                                                let context_server = context_server.clone();
+                                                async move |cx| {
+                                                    if let Some(start_server_task) =
+                                                        context_server_manager
+                                                            .update(cx, |this, cx| {
+                                                                this.start_server(
+                                                                    context_server,
+                                                                    cx,
+                                                                )
+                                                            })
+                                                            .log_err()
+                                                    {
+                                                        start_server_task.await.log_err();
+                                                    }
+                                                }
+                                            })
+                                            .detach();
+                                        }
+                                    }
+                                }),
+                            )),
                     )
                     .map(|parent| {
                         if !are_tools_expanded {

crates/assistant_context_editor/src/context_store.rs 🔗

@@ -818,7 +818,7 @@ impl ContextStore {
         cx.update_entity(
             &self.context_server_manager,
             |context_server_manager, cx| {
-                for server in context_server_manager.servers() {
+                for server in context_server_manager.running_servers() {
                     context_server_manager
                         .restart_server(&server.id(), cx)
                         .detach_and_log_err(cx);

crates/context_server/src/manager.rs 🔗

@@ -155,7 +155,7 @@ impl ContextServerManager {
                 Self::maintain_servers(this.clone(), cx).await?;
 
                 this.update(cx, |this, cx| {
-                    let has_any_context_servers = !this.servers().is_empty();
+                    let has_any_context_servers = !this.running_servers().is_empty();
                     if has_any_context_servers {
                         CommandPaletteFilter::update_global(cx, |filter, _cx| {
                             filter.show_namespace(CONTEXT_SERVERS_NAMESPACE);
@@ -180,6 +180,31 @@ impl ContextServerManager {
             .cloned()
     }
 
+    pub fn start_server(
+        &self,
+        server: Arc<ContextServer>,
+        cx: &mut Context<Self>,
+    ) -> Task<anyhow::Result<()>> {
+        cx.spawn(async move |this, cx| {
+            let id = server.id.clone();
+            server.start(&cx).await?;
+            this.update(cx, |_, cx| cx.emit(Event::ServerStarted { server_id: id }))?;
+            Ok(())
+        })
+    }
+
+    pub fn stop_server(
+        &self,
+        server: Arc<ContextServer>,
+        cx: &mut Context<Self>,
+    ) -> anyhow::Result<()> {
+        server.stop()?;
+        cx.emit(Event::ServerStopped {
+            server_id: server.id(),
+        });
+        Ok(())
+    }
+
     pub fn restart_server(
         &mut self,
         id: &Arc<str>,
@@ -206,7 +231,11 @@ impl ContextServerManager {
         })
     }
 
-    pub fn servers(&self) -> Vec<Arc<ContextServer>> {
+    pub fn all_servers(&self) -> Vec<Arc<ContextServer>> {
+        self.servers.values().cloned().collect()
+    }
+
+    pub fn running_servers(&self) -> Vec<Arc<ContextServer>> {
         self.servers
             .values()
             .filter(|server| server.client().is_some())