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

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

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| {
@@ -205,6 +206,9 @@ impl AssistantConfiguration {
                                     .border_b_1()
                                     .border_color(cx.theme().colors().border)
                             })
+                            .child(
+                                h_flex()
+                                    .gap_2()
                                     .child(
                                         Disclosure::new("tool-list-disclosure", are_tools_expanded)
                                             .on_click(cx.listener({
@@ -225,7 +229,48 @@ impl AssistantConfiguration {
                                         Color::Error
                                     }))
                                     .child(Label::new(context_server.id()))
-                            .child(Label::new(format!("{tool_count} tools")).color(Color::Muted)),
+                                    .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())