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 <benjamin.j.brandt@gmail.com>

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;

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";