oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

context_store: Refactor state management (#29910)

Browse Source 
Because we instantiated `ContextServerManager` both in `agent` and
`assistant-context-editor`, and these two entities track the running MCP
servers separately, we were effectively running every MCP server twice.

This PR moves the `ContextServerManager` into the project crate (now
called `ContextServerStore`). The store can be accessed via a project
instance. This ensures that we only instantiate one `ContextServerStore`
per project.

Also, this PR adds a bunch of tests to ensure that the
`ContextServerStore` behaves correctly (Previously there were none).

Closes #28714
Closes #29530

Release Notes:

- N/A
This commit is contained in:
Bennet Bo Fenner
2025-05-05 21:36:12 +02:00
committed by GitHub
parent 8199664a5a
commit 9cb5ffac25
43 changed files with 1568 additions and 1047 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/project/src/context_server_store.rs
+1129
View File
File diff suppressed because it is too large Load Diff
crates/project/src/context_server_store/extension.rs
+121
View File
@@ -0,0 +1,121 @@
use std::sync::Arc;
use anyhow::Result;
use context_server::ContextServerCommand;
use extension::{
ContextServerConfiguration, Extension, ExtensionContextServerProxy, ExtensionHostProxy,
ProjectDelegate,
};
use gpui::{App, AsyncApp, Entity, Task};
use crate::worktree_store::WorktreeStore;
use super::registry::{self, ContextServerDescriptorRegistry};
pub fn init(cx: &mut App) {
let proxy = ExtensionHostProxy::default_global(cx);
proxy.register_context_server_proxy(ContextServerDescriptorRegistryProxy {
context_server_factory_registry: ContextServerDescriptorRegistry::default_global(cx),
});
}
struct ExtensionProject {
worktree_ids: Vec<u64>,
}
impl ProjectDelegate for ExtensionProject {
fn worktree_ids(&self) -> Vec<u64> {
self.worktree_ids.clone()
}
}
struct ContextServerDescriptor {
id: Arc<str>,
extension: Arc<dyn Extension>,
}
fn extension_project(
worktree_store: Entity<WorktreeStore>,
cx: &mut AsyncApp,
) -> Result<Arc<ExtensionProject>> {
worktree_store.update(cx, |worktree_store, cx| {
Arc::new(ExtensionProject {
worktree_ids: worktree_store
.visible_worktrees(cx)
.map(|worktree| worktree.read(cx).id().to_proto())
.collect(),
})
})
}
impl registry::ContextServerDescriptor for ContextServerDescriptor {
fn command(
&self,
worktree_store: Entity<WorktreeStore>,
cx: &AsyncApp,
) -> Task<Result<ContextServerCommand>> {
let id = self.id.clone();
let extension = self.extension.clone();
cx.spawn(async move |cx| {
let extension_project = extension_project(worktree_store, cx)?;
let mut command = extension
.context_server_command(id.clone(), extension_project.clone())
.await?;
command.command = extension
.path_from_extension(command.command.as_ref())
.to_string_lossy()
.to_string();
log::info!("loaded command for context server {id}: {command:?}");
Ok(ContextServerCommand {
path: command.command,
args: command.args,
env: Some(command.env.into_iter().collect()),
})
})
}
fn configuration(
&self,
worktree_store: Entity<WorktreeStore>,
cx: &AsyncApp,
) -> Task<Result<Option<ContextServerConfiguration>>> {
let id = self.id.clone();
let extension = self.extension.clone();
cx.spawn(async move |cx| {
let extension_project = extension_project(worktree_store, cx)?;
let configuration = extension
.context_server_configuration(id.clone(), extension_project)
.await?;
log::debug!("loaded configuration for context server {id}: {configuration:?}");
Ok(configuration)
})
}
}
struct ContextServerDescriptorRegistryProxy {
context_server_factory_registry: Entity<ContextServerDescriptorRegistry>,
}
impl ExtensionContextServerProxy for ContextServerDescriptorRegistryProxy {
fn register_context_server(&self, extension: Arc<dyn Extension>, id: Arc<str>, cx: &mut App) {
self.context_server_factory_registry
.update(cx, |registry, _| {
registry.register_context_server_descriptor(
id.clone(),
Arc::new(ContextServerDescriptor { id, extension })
as Arc<dyn registry::ContextServerDescriptor>,
)
});
}
fn unregister_context_server(&self, server_id: Arc<str>, cx: &mut App) {
self.context_server_factory_registry
.update(cx, |registry, _| {
registry.unregister_context_server_descriptor_by_id(&server_id)
});
}
}
crates/project/src/context_server_store/registry.rs
+77
View File
@@ -0,0 +1,77 @@
use std::sync::Arc;
use anyhow::Result;
use collections::HashMap;
use context_server::ContextServerCommand;
use extension::ContextServerConfiguration;
use gpui::{App, AppContext as _, AsyncApp, Entity, Global, Task};
use crate::worktree_store::WorktreeStore;
pub trait ContextServerDescriptor {
fn command(
&self,
worktree_store: Entity<WorktreeStore>,
cx: &AsyncApp,
) -> Task<Result<ContextServerCommand>>;
fn configuration(
&self,
worktree_store: Entity<WorktreeStore>,
cx: &AsyncApp,
) -> Task<Result<Option<ContextServerConfiguration>>>;
}
struct GlobalContextServerDescriptorRegistry(Entity<ContextServerDescriptorRegistry>);
impl Global for GlobalContextServerDescriptorRegistry {}
#[derive(Default)]
pub struct ContextServerDescriptorRegistry {
context_servers: HashMap<Arc<str>, Arc<dyn ContextServerDescriptor>>,
}
impl ContextServerDescriptorRegistry {
/// Returns the global [`ContextServerDescriptorRegistry`].
///
/// Inserts a default [`ContextServerDescriptorRegistry`] if one does not yet exist.
pub fn default_global(cx: &mut App) -> Entity<Self> {
if !cx.has_global::<GlobalContextServerDescriptorRegistry>() {
let registry = cx.new(|_| Self::new());
cx.set_global(GlobalContextServerDescriptorRegistry(registry));
}
cx.global::<GlobalContextServerDescriptorRegistry>()
.0
.clone()
}
pub fn new() -> Self {
Self {
context_servers: HashMap::default(),
}
}
pub fn context_server_descriptors(&self) -> Vec<(Arc<str>, Arc<dyn ContextServerDescriptor>)> {
self.context_servers
.iter()
.map(|(id, factory)| (id.clone(), factory.clone()))
.collect()
}
pub fn context_server_descriptor(&self, id: &str) -> Option<Arc<dyn ContextServerDescriptor>> {
self.context_servers.get(id).cloned()
}
/// Registers the provided [`ContextServerDescriptor`].
pub fn register_context_server_descriptor(
&mut self,
id: Arc<str>,
descriptor: Arc<dyn ContextServerDescriptor>,
) {
self.context_servers.insert(id, descriptor);
}
/// Unregisters the [`ContextServerDescriptor`] for the server with the given ID.
pub fn unregister_context_server_descriptor_by_id(&mut self, server_id: &str) {
self.context_servers.remove(server_id);
}
}
crates/project/src/git_store/git_traversal.rs
+2 -3
View File
@@ -244,9 +244,8 @@ mod tests {
use git::status::{FileStatus, StatusCode, TrackedSummary, UnmergedStatus, UnmergedStatusCode};
use gpui::TestAppContext;
use serde_json::json;
use settings::{Settings as _, SettingsStore};
use settings::SettingsStore;
use util::path;
use worktree::WorktreeSettings;
const CONFLICT: FileStatus = FileStatus::Unmerged(UnmergedStatus {
first_head: UnmergedStatusCode::Updated,
@@ -682,7 +681,7 @@ mod tests {
cx.update(|cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
WorktreeSettings::register(cx);
Project::init_settings(cx);
});
}
crates/project/src/project.rs
+19
View File
@@ -1,6 +1,7 @@
pub mod buffer_store;
mod color_extractor;
pub mod connection_manager;
pub mod context_server_store;
pub mod debounced_delay;
pub mod debugger;
pub mod git_store;
@@ -23,6 +24,7 @@ mod project_tests;
mod direnv;
mod environment;
use buffer_diff::BufferDiff;
use context_server_store::ContextServerStore;
pub use environment::{EnvironmentErrorMessage, ProjectEnvironmentEvent};
use git_store::{Repository, RepositoryId};
pub mod search_history;
@@ -182,6 +184,7 @@ pub struct Project {
client_subscriptions: Vec<client::Subscription>,
worktree_store: Entity<WorktreeStore>,
buffer_store: Entity<BufferStore>,
context_server_store: Entity<ContextServerStore>,
image_store: Entity<ImageStore>,
lsp_store: Entity<LspStore>,
_subscriptions: Vec<gpui::Subscription>,
@@ -845,6 +848,7 @@ impl Project {
ToolchainStore::init(&client);
DapStore::init(&client, cx);
BreakpointStore::init(&client);
context_server_store::init(cx);
}
pub fn local(
@@ -865,6 +869,9 @@ impl Project {
cx.subscribe(&worktree_store, Self::on_worktree_store_event)
.detach();
let context_server_store =
cx.new(|cx| ContextServerStore::new(worktree_store.clone(), cx));
let environment = cx.new(|_| ProjectEnvironment::new(env));
let toolchain_store = cx.new(|cx| {
ToolchainStore::local(
@@ -965,6 +972,7 @@ impl Project {
buffer_store,
image_store,
lsp_store,
context_server_store,
join_project_response_message_id: 0,
client_state: ProjectClientState::Local,
git_store,
@@ -1025,6 +1033,9 @@ impl Project {
cx.subscribe(&worktree_store, Self::on_worktree_store_event)
.detach();
let context_server_store =
cx.new(|cx| ContextServerStore::new(worktree_store.clone(), cx));
let buffer_store = cx.new(|cx| {
BufferStore::remote(
worktree_store.clone(),
@@ -1109,6 +1120,7 @@ impl Project {
buffer_store,
image_store,
lsp_store,
context_server_store,
breakpoint_store,
dap_store,
join_project_response_message_id: 0,
@@ -1267,6 +1279,8 @@ impl Project {
let image_store = cx.new(|cx| {
ImageStore::remote(worktree_store.clone(), client.clone().into(), remote_id, cx)
})?;
let context_server_store =
cx.new(|cx| ContextServerStore::new(worktree_store.clone(), cx))?;
let environment = cx.new(|_| ProjectEnvironment::new(None))?;
@@ -1360,6 +1374,7 @@ impl Project {
image_store,
worktree_store: worktree_store.clone(),
lsp_store: lsp_store.clone(),
context_server_store,
active_entry: None,
collaborators: Default::default(),
join_project_response_message_id: response.message_id,
@@ -1590,6 +1605,10 @@ impl Project {
self.worktree_store.clone()
}
pub fn context_server_store(&self) -> Entity<ContextServerStore> {
self.context_server_store.clone()
}
pub fn buffer_for_id(&self, remote_id: BufferId, cx: &App) -> Option<Entity<Buffer>> {
self.buffer_store.read(cx).get(remote_id)
}
crates/project/src/project_settings.rs
+52
View File
@@ -1,5 +1,6 @@
use anyhow::Context as _;
use collections::HashMap;
use context_server::ContextServerCommand;
use dap::adapters::DebugAdapterName;
use fs::Fs;
use futures::StreamExt as _;
@@ -51,6 +52,10 @@ pub struct ProjectSettings {
#[serde(default)]
pub dap: HashMap<DebugAdapterName, DapSettings>,
/// Settings for context servers used for AI-related features.
#[serde(default)]
pub context_servers: HashMap<Arc<str>, ContextServerConfiguration>,
/// Configuration for Diagnostics-related features.
#[serde(default)]
pub diagnostics: DiagnosticsSettings,
@@ -78,6 +83,19 @@ pub struct DapSettings {
pub binary: Option<String>,
}
#[derive(Deserialize, Serialize, Clone, PartialEq, Eq, JsonSchema, Debug, Default)]
pub struct ContextServerConfiguration {
/// The command to run this context server.
///
/// This will override the command set by an extension.
pub command: Option<ContextServerCommand>,
/// The settings for this context server.
///
/// Consult the documentation for the context server to see what settings
/// are supported.
pub settings: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct NodeBinarySettings {
/// The path to the Node binary.
@@ -376,6 +394,40 @@ impl Settings for ProjectSettings {
}
}
#[derive(Deserialize)]
struct VsCodeContextServerCommand {
command: String,
args: Option<Vec<String>>,
env: Option<HashMap<String, String>>,
// note: we don't support envFile and type
}
impl From<VsCodeContextServerCommand> for ContextServerCommand {
fn from(cmd: VsCodeContextServerCommand) -> Self {
Self {
path: cmd.command,
args: cmd.args.unwrap_or_default(),
env: cmd.env,
}
}
}
if let Some(mcp) = vscode.read_value("mcp").and_then(|v| v.as_object()) {
current
.context_servers
.extend(mcp.iter().filter_map(|(k, v)| {
Some((
k.clone().into(),
ContextServerConfiguration {
command: Some(
serde_json::from_value::<VsCodeContextServerCommand>(v.clone())
.ok()?
.into(),
),
settings: None,
},
))
}));
}
// TODO: translate lsp settings for rust-analyzer and other popular ones to old.lsp
}
}