Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
This commit is contained in:
Marshall Bowers 2025-03-19 11:37:48 -04:00 • committed by GitHub
parent 06ffdc6791
commit 410a942d57
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 98 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

87
crates/assistant2/src/assistant_configuration.rs
View file

@ -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 {

2
crates/assistant_context_editor/src/context_store.rs
View file

@ -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);

33
crates/context_server/src/manager.rs
View file

@ -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())