vndangkhoa/zed
1
0
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2026-06-01 03:14:56 +07:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

agent: Fix disabled MCP servers disappearing from UI after restart (#47758)

Browse source 
Closes #42888

Release Notes:

- Fixed disabled MCP servers disappearing from UI after restart

---------

Co-authored-by: MrSubidubi <finn@zed.dev>
This commit is contained in:
Oliver Azevedo Barnes 2026-02-08 23:56:20 +00:00 committed by GitHub
parent 3b31224448
commit a96d777988
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 216 additions and 129 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
crates/agent/src/agent.rs
View file

@ -667,7 +667,7 @@ impl NativeAgent {
fn handle_context_server_store_updated(
&mut self,
_store: Entity<project::context_server_store::ContextServerStore>,
_event: &project::context_server_store::Event,
_event: &project::context_server_store::ServerStatusChangedEvent,
cx: &mut Context<Self>,
) {
self.update_available_commands(cx);

36
crates/agent/src/tools/context_server_registry.rs
View file

@ -247,31 +247,29 @@ impl ContextServerRegistry {
fn handle_context_server_store_event(
&mut self,
_: Entity<ContextServerStore>,
event: &project::context_server_store::Event,
event: &project::context_server_store::ServerStatusChangedEvent,
cx: &mut Context<Self>,
) {
match event {
project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
match status {
ContextServerStatus::Starting => {}
ContextServerStatus::Running => {
self.reload_tools_for_server(server_id.clone(), cx);
self.reload_prompts_for_server(server_id.clone(), cx);
let project::context_server_store::ServerStatusChangedEvent { server_id, status } = event;
match status {
ContextServerStatus::Starting => {}
ContextServerStatus::Running => {
self.reload_tools_for_server(server_id.clone(), cx);
self.reload_prompts_for_server(server_id.clone(), cx);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
if let Some(registered_server) = self.registered_servers.remove(server_id) {
if !registered_server.tools.is_empty() {
cx.emit(ContextServerRegistryEvent::ToolsChanged);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
if let Some(registered_server) = self.registered_servers.remove(server_id) {
if !registered_server.tools.is_empty() {
cx.emit(ContextServerRegistryEvent::ToolsChanged);
}
if !registered_server.prompts.is_empty() {
cx.emit(ContextServerRegistryEvent::PromptsChanged);
}
}
cx.notify();
if !registered_server.prompts.is_empty() {
cx.emit(ContextServerRegistryEvent::PromptsChanged);
}
}
cx.notify();
}
}
};
}
}

52
crates/agent_ui/src/agent_configuration.rs
View file

@ -518,25 +518,7 @@ impl AgentConfiguration {
window: &mut Window,
cx: &mut Context<Self>,
) -> impl IntoElement {
let mut context_server_ids = self
.context_server_store
.read(cx)
.server_ids(cx)
.into_iter()
.collect::<Vec<_>>();
// Sort context servers: ones without mcp-server- prefix first, then prefixed ones
context_server_ids.sort_by(|a, b| {
const MCP_PREFIX: &str = "mcp-server-";
match (a.0.strip_prefix(MCP_PREFIX), b.0.strip_prefix(MCP_PREFIX)) {
// If one has mcp-server- prefix and other doesn't, non-mcp comes first
(Some(_), None) => std::cmp::Ordering::Greater,
(None, Some(_)) => std::cmp::Ordering::Less,
// If both have same prefix status, sort by appropriate key
(Some(a), Some(b)) => a.cmp(b),
(None, None) => a.0.cmp(&b.0),
}
});
let context_server_ids = self.context_server_store.read(cx).server_ids();
let add_server_popover = PopoverMenu::new("add-server-popover")
.trigger(
@ -594,7 +576,7 @@ impl AgentConfiguration {
.pr_5()
.w_full()
.gap_1()
.map(|mut parent| {
.map(|parent| {
if context_server_ids.is_empty() {
parent.child(
h_flex()
@ -611,23 +593,17 @@ impl AgentConfiguration {
),
)
} else {
for (index, context_server_id) in
context_server_ids.into_iter().enumerate()
{
if index > 0 {
parent = parent.child(
Divider::horizontal()
.color(DividerColor::BorderFaded)
.into_any_element(),
);
}
parent = parent.child(self.render_context_server(
context_server_id,
window,
cx,
));
}
parent
parent.children(itertools::intersperse_with(
context_server_ids.iter().cloned().map(|context_server_id| {
self.render_context_server(context_server_id, window, cx)
.into_any_element()
}),
|| {
Divider::horizontal()
.color(DividerColor::BorderFaded)
.into_any_element()
},
))
}
}),
)
@ -637,7 +613,7 @@ impl AgentConfiguration {
&self,
context_server_id: ContextServerId,
window: &mut Window,
cx: &mut Context<Self>,
cx: &Context<Self>,
) -> impl use<> + IntoElement {
let server_status = self
.context_server_store

48
crates/agent_ui/src/agent_configuration/configure_context_server_modal.rs
View file

@ -880,32 +880,32 @@ fn wait_for_context_server(
let (tx, rx) = futures::channel::oneshot::channel();
let tx = Arc::new(Mutex::new(Some(tx)));
let subscription = cx.subscribe(context_server_store, move |_, event, _cx| match event {
project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
match status {
ContextServerStatus::Running => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
let _ = tx.send(Ok(()));
}
let subscription = cx.subscribe(context_server_store, move |_, event, _cx| {
let project::context_server_store::ServerStatusChangedEvent { server_id, status } = event;
match status {
ContextServerStatus::Running => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
let _ = tx.send(Ok(()));
}
ContextServerStatus::Stopped => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
let _ = tx.send(Err("Context server stopped running".into()));
}
}
ContextServerStatus::Error(error) => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
let _ = tx.send(Err(error.clone()));
}
}
_ => {}
}
ContextServerStatus::Stopped => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
let _ = tx.send(Err("Context server stopped running".into()));
}
}
ContextServerStatus::Error(error) => {
if server_id == &context_server_id
&& let Some(tx) = tx.lock().unwrap().take()
{
let _ = tx.send(Err(error.clone()));
}
}
_ => {}
}
});

36
crates/assistant_text_thread/src/text_thread_store.rs
View file

@ -888,29 +888,27 @@ impl TextThreadStore {
fn handle_context_server_event(
&mut self,
context_server_store: Entity<ContextServerStore>,
event: &project::context_server_store::Event,
event: &project::context_server_store::ServerStatusChangedEvent,
cx: &mut Context<Self>,
) {
match event {
project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
match status {
ContextServerStatus::Running => {
self.load_context_server_slash_commands(
server_id.clone(),
context_server_store,
cx,
);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
if let Some(slash_command_ids) =
self.context_server_slash_command_ids.remove(server_id)
{
self.slash_commands.remove(&slash_command_ids);
}
}
_ => {}
let project::context_server_store::ServerStatusChangedEvent { server_id, status } = event;
match status {
ContextServerStatus::Running => {
self.load_context_server_slash_commands(
server_id.clone(),
context_server_store,
cx,
);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
if let Some(slash_command_ids) =
self.context_server_slash_command_ids.remove(server_id)
{
self.slash_commands.remove(&slash_command_ids);
}
}
_ => {}
}
}

52
crates/project/src/context_server_store.rs
View file

@ -10,6 +10,7 @@ use collections::{HashMap, HashSet};
use context_server::{ContextServer, ContextServerCommand, ContextServerId};
use futures::{FutureExt as _, future::join_all};
use gpui::{App, AsyncApp, Context, Entity, EventEmitter, Subscription, Task, WeakEntity, actions};
use itertools::Itertools;
use registry::ContextServerDescriptorRegistry;
use remote::RemoteClient;
use rpc::{AnyProtoClient, TypedEnvelope, proto};
@ -203,6 +204,7 @@ pub struct ContextServerStore {
state: ContextServerStoreState,
context_server_settings: HashMap<Arc<str>, ContextServerSettings>,
servers: HashMap<ContextServerId, ContextServerState>,
server_ids: Vec<ContextServerId>,
worktree_store: Entity<WorktreeStore>,
project: Option<WeakEntity<Project>>,
registry: Entity<ContextServerDescriptorRegistry>,
@ -212,14 +214,12 @@ pub struct ContextServerStore {
_subscriptions: Vec<Subscription>,
}
pub enum Event {
ServerStatusChanged {
server_id: ContextServerId,
status: ContextServerStatus,
},
pub struct ServerStatusChangedEvent {
pub server_id: ContextServerId,
pub status: ContextServerStatus,
}
impl EventEmitter<Event> for ContextServerStore {}
impl EventEmitter<ServerStatusChangedEvent> for ContextServerStore {}
impl ContextServerStore {
pub fn local(
@ -394,6 +394,7 @@ impl ContextServerStore {
registry,
needs_server_update: false,
servers: HashMap::default(),
server_ids: Default::default(),
update_servers_task: None,
context_server_factory,
};
@ -426,8 +427,16 @@ impl ContextServerStore {
self.servers.get(id).map(|state| state.configuration())
}
pub fn server_ids(&self, cx: &App) -> HashSet<ContextServerId> {
self.servers
/// Returns a sorted slice of available unique context server IDs. Within the
/// slice, context servers which have `mcp-server-` as a prefix in their ID will
/// appear after servers that do not have this prefix in their ID.
pub fn server_ids(&self) -> &[ContextServerId] {
self.server_ids.as_slice()
}
fn populate_server_ids(&mut self, cx: &App) {
self.server_ids = self
.servers
.keys()
.cloned()
.chain(
@ -437,7 +446,27 @@ impl ContextServerStore {
.into_iter()
.map(|(id, _)| ContextServerId(id)),
)
.collect()
.chain(
self.context_server_settings
.keys()
.map(|id| ContextServerId(id.clone())),
)
.unique()
.sorted_unstable_by(
// Sort context servers: ones without mcp-server- prefix first, then prefixed ones
|a, b| {
const MCP_PREFIX: &str = "mcp-server-";
match (a.0.strip_prefix(MCP_PREFIX), b.0.strip_prefix(MCP_PREFIX)) {
// If one has mcp-server- prefix and other doesn't, non-mcp comes first
(Some(_), None) => std::cmp::Ordering::Greater,
(None, Some(_)) => std::cmp::Ordering::Less,
// If both have same prefix status, sort by appropriate key
(Some(a), Some(b)) => a.cmp(b),
(None, None) => a.0.cmp(&b.0),
}
},
)
.collect();
}
pub fn running_servers(&self) -> Vec<Arc<ContextServer>> {
@ -591,7 +620,7 @@ impl ContextServerStore {
.remove(id)
.context("Context server not found")?;
drop(state);
cx.emit(Event::ServerStatusChanged {
cx.emit(ServerStatusChangedEvent {
server_id: id.clone(),
status: ContextServerStatus::Stopped,
});
@ -808,7 +837,7 @@ impl ContextServerStore {
) {
let status = ContextServerStatus::from_state(&state);
self.servers.insert(id.clone(), state);
cx.emit(Event::ServerStatusChanged {
cx.emit(ServerStatusChangedEvent {
server_id: id,
status,
});
@ -825,6 +854,7 @@ impl ContextServerStore {
}
this.update(cx, |this, cx| {
this.populate_server_ids(cx);
this.update_servers_task.take();
if this.needs_server_update {
this.available_context_servers_changed(cx);

119
crates/project/tests/integration/context_server_store.rs
View file

@ -553,6 +553,92 @@ async fn test_context_server_enabled_disabled(cx: &mut TestAppContext) {
}
}
#[gpui::test]
async fn test_server_ids_includes_disabled_servers(cx: &mut TestAppContext) {
const ENABLED_SERVER_ID: &str = "enabled-server";
const DISABLED_SERVER_ID: &str = "disabled-server";
let enabled_server_id = ContextServerId(ENABLED_SERVER_ID.into());
let disabled_server_id = ContextServerId(DISABLED_SERVER_ID.into());
let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
let executor = cx.executor();
let store = project.read_with(cx, |project, _| project.context_server_store());
store.update(cx, |store, _| {
store.set_context_server_factory(Box::new(move |id, _| {
Arc::new(ContextServer::new(
id.clone(),
Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
))
}));
});
// Configure one enabled and one disabled server
set_context_server_configuration(
vec![
(
enabled_server_id.0.clone(),
settings::ContextServerSettingsContent::Stdio {
enabled: true,
remote: false,
command: ContextServerCommand {
path: "somebinary".into(),
args: vec![],
env: None,
timeout: None,
},
},
),
(
disabled_server_id.0.clone(),
settings::ContextServerSettingsContent::Stdio {
enabled: false,
remote: false,
command: ContextServerCommand {
path: "somebinary".into(),
args: vec![],
env: None,
timeout: None,
},
},
),
],
cx,
);
cx.run_until_parked();
// Verify that server_ids includes both enabled and disabled servers
cx.update(|cx| {
let server_ids = store.read(cx).server_ids().to_vec();
assert!(
server_ids.contains(&enabled_server_id),
"server_ids should include enabled server"
);
assert!(
server_ids.contains(&disabled_server_id),
"server_ids should include disabled server"
);
});
// Verify that the enabled server is running and the disabled server is not
cx.read(|cx| {
assert_eq!(
store.read(cx).status_for_server(&enabled_server_id),
Some(ContextServerStatus::Running),
"enabled server should be running"
);
// Disabled server should not be in the servers map (status returns None)
// but should still be in server_ids
assert_eq!(
store.read(cx).status_for_server(&disabled_server_id),
None,
"disabled server should not have a status (not in servers map)"
);
});
}
fn set_context_server_configuration(
context_servers: Vec<(Arc<str>, settings::ContextServerSettingsContent)>,
cx: &mut TestAppContext,
@ -796,26 +882,25 @@ fn assert_server_events(
let expected_event_count = expected_events.len();
let subscription = cx.subscribe(store, {
let received_event_count = received_event_count.clone();
move |_, event, _| match event {
Event::ServerStatusChanged {
move |_, event, _| {
let ServerStatusChangedEvent {
server_id: actual_server_id,
status: actual_status,
} => {
let (expected_server_id, expected_status) = &expected_events[ix];
} = event;
let (expected_server_id, expected_status) = &expected_events[ix];
assert_eq!(
actual_server_id, expected_server_id,
"Expected different server id at index {}",
ix
);
assert_eq!(
actual_status, expected_status,
"Expected different status at index {}",
ix
);
ix += 1;
*received_event_count.borrow_mut() += 1;
}
assert_eq!(
actual_server_id, expected_server_id,
"Expected different server id at index {}",
ix
);
assert_eq!(
actual_status, expected_status,
"Expected different status at index {}",
ix
);
ix += 1;
*received_event_count.borrow_mut() += 1;
}
});
ServerEvents {