From 4d13f01c89ef4fbba9ece2ee3fe6a61b1d7f32e3 Mon Sep 17 00:00:00 2001 From: Rio Fujita Date: Mon, 18 May 2026 04:38:37 +0900 Subject: [PATCH] project: Context server updates on worktree changes (#51244) Before you mark this PR as ready for review, make sure that you have: - [x] Added a solid test coverage and/or screenshots from doing manual testing - [x] Done a self-review taking into account security and performance aspects - [x] Aligned any UI changes with the UI checklist Release Notes: - Fixed context server availability updates when a new worktree is added to or removed from a project. --------- Co-authored-by: Ben Brandt --- crates/project/src/context_server_store.rs | 12 ++- .../tests/integration/context_server_store.rs | 87 +++++++++++++++++++ 2 files changed, 98 insertions(+), 1 deletion(-) diff --git a/crates/project/src/context_server_store.rs b/crates/project/src/context_server_store.rs index de2e1e3ceff..6e231453e41 100644 --- a/crates/project/src/context_server_store.rs +++ b/crates/project/src/context_server_store.rs @@ -28,7 +28,7 @@ use util::{ResultExt as _, rel_path::RelPath}; use crate::{ DisableAiSettings, Project, project_settings::{ContextServerSettings, ProjectSettings}, - worktree_store::WorktreeStore, + worktree_store::{WorktreeStore, WorktreeStoreEvent}, }; /// Maximum timeout for context server requests @@ -454,6 +454,16 @@ impl ContextServerStore { this.available_context_servers_changed(cx); } })); + subscriptions.push(cx.subscribe(&worktree_store, |this, _store, event, cx| { + if matches!( + event, + WorktreeStoreEvent::WorktreeAdded(_) + | WorktreeStoreEvent::WorktreeRemoved(_, _) + ) && !DisableAiSettings::get_global(cx).disable_ai + { + this.available_context_servers_changed(cx); + } + })); } let ai_disabled = DisableAiSettings::get_global(cx).disable_ai; diff --git a/crates/project/tests/integration/context_server_store.rs b/crates/project/tests/integration/context_server_store.rs index 5b68e11bb95..f9dbce84a17 100644 --- a/crates/project/tests/integration/context_server_store.rs +++ b/crates/project/tests/integration/context_server_store.rs @@ -664,6 +664,93 @@ async fn test_context_server_respects_disable_ai(cx: &mut TestAppContext) { }); } +#[gpui::test] +async fn test_context_server_refreshed_when_worktree_added(cx: &mut TestAppContext) { + const SERVER_1_ID: &str = "mcp-1"; + + let server_1_id = ContextServerId(SERVER_1_ID.into()); + + let (fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await; + fs.insert_tree(path!("/second"), json!({"other.rs": ""})) + .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())), + )) + })); + }); + + set_context_server_configuration( + vec![( + server_1_id.0.clone(), + settings::ContextServerSettingsContent::Stdio { + enabled: true, + remote: false, + command: ContextServerCommand { + path: "somebinary".into(), + args: vec!["arg".to_string()], + env: None, + timeout: None, + }, + }, + )], + cx, + ); + + { + let _server_events = assert_server_events( + &store, + vec![ + (server_1_id.clone(), ContextServerStatus::Starting), + (server_1_id.clone(), ContextServerStatus::Running), + ], + cx, + ); + cx.run_until_parked(); + } + + // Witness that adding a worktree triggers the store to refresh available + // servers (via `cx.notify` after `maintain_servers`). Without the + // `WorktreeStoreEvent::WorktreeAdded` subscription in `ContextServerStore`, + // this counter would remain zero. + let notify_count = Rc::new(RefCell::new(0usize)); + let _notify_subscription = cx.update(|cx| { + let count = notify_count.clone(); + cx.observe(&store, move |_, _| { + *count.borrow_mut() += 1; + }) + }); + + { + let _server_events = assert_server_events(&store, vec![], cx); + let _ = project.update(cx, |project, cx| { + project.find_or_create_worktree(path!("/second"), true, cx) + }); + cx.run_until_parked(); + } + + cx.update(|cx| { + assert!( + *notify_count.borrow() > 0, + "Adding a worktree should trigger the context server store to refresh" + ); + assert!( + store.read(cx).server_ids().contains(&server_1_id), + "Configured server list should still include the server after a worktree is added" + ); + assert_eq!( + store.read(cx).status_for_server(&server_1_id), + Some(ContextServerStatus::Running), + "Server should still be running after a worktree is added" + ); + }); +} + #[gpui::test] async fn test_server_ids_includes_disabled_servers(cx: &mut TestAppContext) { const ENABLED_SERVER_ID: &str = "enabled-server";