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

agent: Fix MCP server handler subscription race condition (#32133)

Browse source 
Closes #32132

Release Notes:

- Fixed MCP server handler subscription race condition causing tools to
not load.

---------

Co-authored-by: Bennet Bo Fenner <bennet@zed.dev>
Co-authored-by: Bennet Bo Fenner <bennetbo@gmx.de>
This commit is contained in:
Jonathan LEI 2025-06-06 21:32:06 +08:00 • committed by GitHub
parent 380d8c5662
commit 6ea4d2b30d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 113 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

73
crates/agent/src/thread_store.rs
View file

@ -514,11 +514,14 @@ impl ThreadStore {
}
fn register_context_server_handlers(&self, cx: &mut Context<Self>) {
cx.subscribe(
&self.project.read(cx).context_server_store(),
Self::handle_context_server_event,
)
let context_server_store = self.project.read(cx).context_server_store();
cx.subscribe(&context_server_store, Self::handle_context_server_event)
.detach();
// Check for any servers that were already running before the handler was registered
for server in context_server_store.read(cx).running_servers() {
self.load_context_server_tools(server.id(), context_server_store.clone(), cx);
}
}
fn handle_context_server_event(
@ -533,14 +536,31 @@ impl ThreadStore {
match status {
ContextServerStatus::Starting => {}
ContextServerStatus::Running => {
if let Some(server) =
context_server_store.read(cx).get_running_server(server_id)
{
let context_server_manager = context_server_store.clone();
cx.spawn({
let server = server.clone();
let server_id = server_id.clone();
async move |this, cx| {
self.load_context_server_tools(server_id.clone(), context_server_store, cx);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
if let Some(tool_ids) = self.context_server_tool_ids.remove(server_id) {
tool_working_set.update(cx, |tool_working_set, _| {
tool_working_set.remove(&tool_ids);
});
}
}
}
}
}
}
fn load_context_server_tools(
&self,
server_id: ContextServerId,
context_server_store: Entity<ContextServerStore>,
cx: &mut Context<Self>,
) {
let Some(server) = context_server_store.read(cx).get_running_server(&server_id) else {
return;
};
let tool_working_set = self.tools.clone();
cx.spawn(async move |this, cx| {
let Some(protocol) = server.client() else {
return;
};
@ -553,17 +573,12 @@ impl ThreadStore {
.tools
.into_iter()
.map(|tool| {
log::info!(
"registering context server tool: {:?}",
tool.name
);
tool_working_set.insert(Arc::new(
ContextServerTool::new(
context_server_manager.clone(),
log::info!("registering context server tool: {:?}", tool.name);
tool_working_set.insert(Arc::new(ContextServerTool::new(
context_server_store.clone(),
server.id(),
tool,
),
))
)))
})
.collect::<Vec<_>>()
})
@ -571,29 +586,15 @@ impl ThreadStore {
if let Some(tool_ids) = tool_ids {
this.update(cx, |this, _| {
this.context_server_tool_ids
.insert(server_id, tool_ids);
this.context_server_tool_ids.insert(server_id, tool_ids);
})
.log_err();
}
}
}
}
})
.detach();
}
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
if let Some(tool_ids) = self.context_server_tool_ids.remove(server_id) {
tool_working_set.update(cx, |tool_working_set, _| {
tool_working_set.remove(&tool_ids);
});
}
}
}
}
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]

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

@ -809,33 +809,56 @@ impl ContextStore {
}
fn register_context_server_handlers(&self, cx: &mut Context<Self>) {
cx.subscribe(
&self.project.read(cx).context_server_store(),
Self::handle_context_server_event,
)
let context_server_store = self.project.read(cx).context_server_store();
cx.subscribe(&context_server_store, Self::handle_context_server_event)
.detach();
// Check for any servers that were already running before the handler was registered
for server in context_server_store.read(cx).running_servers() {
self.load_context_server_slash_commands(server.id(), context_server_store.clone(), cx);
}
}
fn handle_context_server_event(
&mut self,
context_server_manager: Entity<ContextServerStore>,
context_server_store: Entity<ContextServerStore>,
event: &project::context_server_store::Event,
cx: &mut Context<Self>,
) {
let slash_command_working_set = self.slash_commands.clone();
match event {
project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
match status {
ContextServerStatus::Running => {
if let Some(server) = context_server_manager
.read(cx)
.get_running_server(server_id)
self.load_context_server_slash_commands(
server_id.clone(),
context_server_store.clone(),
cx,
);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
if let Some(slash_command_ids) =
self.context_server_slash_command_ids.remove(server_id)
{
let context_server_manager = context_server_manager.clone();
cx.spawn({
let server = server.clone();
let server_id = server_id.clone();
async move |this, cx| {
self.slash_commands.remove(&slash_command_ids);
}
}
_ => {}
}
}
}
}
fn load_context_server_slash_commands(
&self,
server_id: ContextServerId,
context_server_store: Entity<ContextServerStore>,
cx: &mut Context<Self>,
) {
let Some(server) = context_server_store.read(cx).get_running_server(&server_id) else {
return;
};
let slash_command_working_set = self.slash_commands.clone();
cx.spawn(async move |this, cx| {
let Some(protocol) = server.client() else {
return;
};
@ -846,13 +869,10 @@ impl ContextStore {
.into_iter()
.filter(assistant_slash_commands::acceptable_prompt)
.map(|prompt| {
log::info!(
"registering context server command: {:?}",
prompt.name
);
log::info!("registering context server command: {:?}", prompt.name);
slash_command_working_set.insert(Arc::new(
assistant_slash_commands::ContextServerSlashCommand::new(
context_server_manager.clone(),
context_server_store.clone(),
server.id(),
prompt,
),
@ -860,28 +880,14 @@ impl ContextStore {
})
.collect::<Vec<_>>();
this.update( cx, |this, _cx| {
this.update(cx, |this, _cx| {
this.context_server_slash_command_ids
.insert(server_id.clone(), slash_command_ids);
})
.log_err();
}
}
}
})
.detach();
}
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
if let Some(slash_command_ids) =
self.context_server_slash_command_ids.remove(server_id)
{
slash_command_working_set.remove(&slash_command_ids);
}
}
_ => {}
}
}
}
}
}