mirror of
https://github.com/zed-industries/zed.git
synced 2026-05-28 17:54:14 +00:00
0d832bc6d5
In the interactive MCP OAuth flow, the MCP client registers itself with the authorization in one of three ways: - Client ID Metadata Document aka CIMD (recommended default). This is already implemented: https://zed.dev/oauth/client-metadata.json. - Dynamic Client Registration (DCR). This is the traditional method. Also already implemented in Zed. - Pre-registration: the client is registered out of band, typically in the IdP or SaaS provider's UI. You get a client id and maybe a client secret, that have to be provided by the MCP client when it wants to exchange an access token. This is what this pull request is about. This PR has two main parts: - Allow users to configure a client id and optional client secret for an MCP server in their configuration, under a new `oauth` key, and take it into account - Make the MCP server state and the configuration modal aware of the intermediate states (client secret missing) and error cases stemming from client pre-registration. The client secret can be stored either in the system keychain or in plain text in the MCP server configuration. The UI tries to steer user towards the more secure option: the keychain. <img width="715" height="201" alt="Screenshot 2026-04-10 at 16 48 06" src="https://github.com/user-attachments/assets/5e64103e-6746-4ef0-8bd9-533d492b6912" /> <img width="884" height="544" alt="Screenshot 2026-04-10 at 16 47 07" src="https://github.com/user-attachments/assets/0e35bb3c-cbc4-4e8c-a713-66323597b2e2" /> <img width="785" height="558" alt="Screenshot 2026-04-10 at 16 47 23" src="https://github.com/user-attachments/assets/03339187-1508-461a-87ae-a7c2647df9a5" /> Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content is consistent with the [UI/UX checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Closes https://github.com/issues/assigned?issue=zed-industries%7Czed%7C52198 **Note for the reviewer: I know how busy the AI team is at the moment so please treat this as low priority, we don't have signal that this is a highly desired feature. It's a rather large PR, so I'm happy to pair review / walk through it.** Release Notes: - Added support for OAuth client pre-registration (client id, client secret) to the built-in MCP client.
1169 lines
36 KiB
Rust
1169 lines
36 KiB
Rust
use anyhow::Result;
|
|
use context_server::test::create_fake_transport;
|
|
use context_server::{ContextServer, ContextServerId};
|
|
use gpui::{AppContext, AsyncApp, Entity, Subscription, Task, TestAppContext, UpdateGlobal as _};
|
|
use http_client::{FakeHttpClient, Response};
|
|
use project::context_server_store::registry::ContextServerDescriptorRegistry;
|
|
use project::context_server_store::*;
|
|
use project::project_settings::ContextServerSettings;
|
|
use project::worktree_store::WorktreeStore;
|
|
use project::{
|
|
DisableAiSettings, FakeFs, Project, context_server_store::registry::ContextServerDescriptor,
|
|
project_settings::ProjectSettings,
|
|
};
|
|
use serde_json::json;
|
|
use settings::settings_content::SaturatingBool;
|
|
use settings::{ContextServerCommand, Settings, SettingsStore};
|
|
use std::sync::Arc;
|
|
use std::{cell::RefCell, path::PathBuf, rc::Rc};
|
|
use util::path;
|
|
|
|
#[gpui::test]
|
|
async fn test_context_server_status(cx: &mut TestAppContext) {
|
|
const SERVER_1_ID: &str = "mcp-1";
|
|
const SERVER_2_ID: &str = "mcp-2";
|
|
|
|
let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
|
|
|
|
let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
|
|
let store = cx.new(|cx| {
|
|
ContextServerStore::test(
|
|
registry.clone(),
|
|
project.read(cx).worktree_store(),
|
|
Some(project.downgrade()),
|
|
cx,
|
|
)
|
|
});
|
|
|
|
let server_1_id = ContextServerId(SERVER_1_ID.into());
|
|
let server_2_id = ContextServerId(SERVER_2_ID.into());
|
|
|
|
let server_1 = Arc::new(ContextServer::new(
|
|
server_1_id.clone(),
|
|
Arc::new(create_fake_transport(SERVER_1_ID, cx.executor())),
|
|
));
|
|
let server_2 = Arc::new(ContextServer::new(
|
|
server_2_id.clone(),
|
|
Arc::new(create_fake_transport(SERVER_2_ID, cx.executor())),
|
|
));
|
|
|
|
store.update(cx, |store, cx| store.test_start_server(server_1, cx));
|
|
|
|
cx.run_until_parked();
|
|
|
|
cx.update(|cx| {
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_1_id),
|
|
Some(ContextServerStatus::Running)
|
|
);
|
|
assert_eq!(store.read(cx).status_for_server(&server_2_id), None);
|
|
});
|
|
|
|
store.update(cx, |store, cx| {
|
|
store.test_start_server(server_2.clone(), cx)
|
|
});
|
|
|
|
cx.run_until_parked();
|
|
|
|
cx.update(|cx| {
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_1_id),
|
|
Some(ContextServerStatus::Running)
|
|
);
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_2_id),
|
|
Some(ContextServerStatus::Running)
|
|
);
|
|
});
|
|
|
|
store
|
|
.update(cx, |store, cx| store.stop_server(&server_2_id, cx))
|
|
.unwrap();
|
|
|
|
cx.update(|cx| {
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_1_id),
|
|
Some(ContextServerStatus::Running)
|
|
);
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_2_id),
|
|
Some(ContextServerStatus::Stopped)
|
|
);
|
|
});
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_context_server_status_events(cx: &mut TestAppContext) {
|
|
const SERVER_1_ID: &str = "mcp-1";
|
|
const SERVER_2_ID: &str = "mcp-2";
|
|
|
|
let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
|
|
|
|
let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
|
|
let store = cx.new(|cx| {
|
|
ContextServerStore::test(
|
|
registry.clone(),
|
|
project.read(cx).worktree_store(),
|
|
Some(project.downgrade()),
|
|
cx,
|
|
)
|
|
});
|
|
|
|
let server_1_id = ContextServerId(SERVER_1_ID.into());
|
|
let server_2_id = ContextServerId(SERVER_2_ID.into());
|
|
|
|
let server_1 = Arc::new(ContextServer::new(
|
|
server_1_id.clone(),
|
|
Arc::new(create_fake_transport(SERVER_1_ID, cx.executor())),
|
|
));
|
|
let server_2 = Arc::new(ContextServer::new(
|
|
server_2_id.clone(),
|
|
Arc::new(create_fake_transport(SERVER_2_ID, cx.executor())),
|
|
));
|
|
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![
|
|
(server_1_id.clone(), ContextServerStatus::Starting),
|
|
(server_1_id, ContextServerStatus::Running),
|
|
(server_2_id.clone(), ContextServerStatus::Starting),
|
|
(server_2_id.clone(), ContextServerStatus::Running),
|
|
(server_2_id.clone(), ContextServerStatus::Stopped),
|
|
],
|
|
cx,
|
|
);
|
|
|
|
store.update(cx, |store, cx| store.test_start_server(server_1, cx));
|
|
|
|
cx.run_until_parked();
|
|
|
|
store.update(cx, |store, cx| {
|
|
store.test_start_server(server_2.clone(), cx)
|
|
});
|
|
|
|
cx.run_until_parked();
|
|
|
|
store
|
|
.update(cx, |store, cx| store.stop_server(&server_2_id, cx))
|
|
.unwrap();
|
|
}
|
|
|
|
#[gpui::test(iterations = 25)]
|
|
async fn test_context_server_concurrent_starts(cx: &mut TestAppContext) {
|
|
const SERVER_1_ID: &str = "mcp-1";
|
|
|
|
let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
|
|
|
|
let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
|
|
let store = cx.new(|cx| {
|
|
ContextServerStore::test(
|
|
registry.clone(),
|
|
project.read(cx).worktree_store(),
|
|
Some(project.downgrade()),
|
|
cx,
|
|
)
|
|
});
|
|
|
|
let server_id = ContextServerId(SERVER_1_ID.into());
|
|
|
|
let server_with_same_id_1 = Arc::new(ContextServer::new(
|
|
server_id.clone(),
|
|
Arc::new(create_fake_transport(SERVER_1_ID, cx.executor())),
|
|
));
|
|
let server_with_same_id_2 = Arc::new(ContextServer::new(
|
|
server_id.clone(),
|
|
Arc::new(create_fake_transport(SERVER_1_ID, cx.executor())),
|
|
));
|
|
|
|
// If we start another server with the same id, we should report that we stopped the previous one
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![
|
|
(server_id.clone(), ContextServerStatus::Starting),
|
|
(server_id.clone(), ContextServerStatus::Stopped),
|
|
(server_id.clone(), ContextServerStatus::Starting),
|
|
(server_id.clone(), ContextServerStatus::Running),
|
|
],
|
|
cx,
|
|
);
|
|
|
|
store.update(cx, |store, cx| {
|
|
store.test_start_server(server_with_same_id_1.clone(), cx)
|
|
});
|
|
store.update(cx, |store, cx| {
|
|
store.test_start_server(server_with_same_id_2.clone(), cx)
|
|
});
|
|
|
|
cx.run_until_parked();
|
|
|
|
cx.update(|cx| {
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_id),
|
|
Some(ContextServerStatus::Running)
|
|
);
|
|
});
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_context_server_maintain_servers_loop(cx: &mut TestAppContext) {
|
|
const SERVER_1_ID: &str = "mcp-1";
|
|
const SERVER_2_ID: &str = "mcp-2";
|
|
|
|
let server_1_id = ContextServerId(SERVER_1_ID.into());
|
|
let server_2_id = ContextServerId(SERVER_2_ID.into());
|
|
|
|
let fake_descriptor_1 = Arc::new(FakeContextServerDescriptor::new(SERVER_1_ID));
|
|
|
|
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, cx| {
|
|
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())),
|
|
))
|
|
}));
|
|
store.registry().update(cx, |registry, cx| {
|
|
registry.register_context_server_descriptor(SERVER_1_ID.into(), fake_descriptor_1, cx);
|
|
});
|
|
});
|
|
|
|
set_context_server_configuration(
|
|
vec![(
|
|
server_1_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Extension {
|
|
enabled: true,
|
|
remote: false,
|
|
settings: json!({
|
|
"somevalue": true
|
|
}),
|
|
},
|
|
)],
|
|
cx,
|
|
);
|
|
|
|
// Ensure that mcp-1 starts up
|
|
{
|
|
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();
|
|
}
|
|
|
|
// Ensure that mcp-1 is restarted when the configuration was changed
|
|
{
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![
|
|
(server_1_id.clone(), ContextServerStatus::Stopped),
|
|
(server_1_id.clone(), ContextServerStatus::Starting),
|
|
(server_1_id.clone(), ContextServerStatus::Running),
|
|
],
|
|
cx,
|
|
);
|
|
set_context_server_configuration(
|
|
vec![(
|
|
server_1_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Extension {
|
|
enabled: true,
|
|
remote: false,
|
|
settings: json!({
|
|
"somevalue": false
|
|
}),
|
|
},
|
|
)],
|
|
cx,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
}
|
|
|
|
// Ensure that mcp-1 is not restarted when the configuration was not changed
|
|
{
|
|
let _server_events = assert_server_events(&store, vec![], cx);
|
|
set_context_server_configuration(
|
|
vec![(
|
|
server_1_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Extension {
|
|
enabled: true,
|
|
remote: false,
|
|
settings: json!({
|
|
"somevalue": false
|
|
}),
|
|
},
|
|
)],
|
|
cx,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
}
|
|
|
|
// Ensure that mcp-2 is started once it is added to the settings
|
|
{
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![
|
|
(server_2_id.clone(), ContextServerStatus::Starting),
|
|
(server_2_id.clone(), ContextServerStatus::Running),
|
|
],
|
|
cx,
|
|
);
|
|
set_context_server_configuration(
|
|
vec![
|
|
(
|
|
server_1_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Extension {
|
|
enabled: true,
|
|
remote: false,
|
|
settings: json!({
|
|
"somevalue": false
|
|
}),
|
|
},
|
|
),
|
|
(
|
|
server_2_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,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
}
|
|
|
|
// Ensure that mcp-2 is restarted once the args have changed
|
|
{
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![
|
|
(server_2_id.clone(), ContextServerStatus::Stopped),
|
|
(server_2_id.clone(), ContextServerStatus::Starting),
|
|
(server_2_id.clone(), ContextServerStatus::Running),
|
|
],
|
|
cx,
|
|
);
|
|
set_context_server_configuration(
|
|
vec![
|
|
(
|
|
server_1_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Extension {
|
|
enabled: true,
|
|
remote: false,
|
|
settings: json!({
|
|
"somevalue": false
|
|
}),
|
|
},
|
|
),
|
|
(
|
|
server_2_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Stdio {
|
|
enabled: true,
|
|
remote: false,
|
|
command: ContextServerCommand {
|
|
path: "somebinary".into(),
|
|
args: vec!["anotherArg".to_string()],
|
|
env: None,
|
|
timeout: None,
|
|
},
|
|
},
|
|
),
|
|
],
|
|
cx,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
}
|
|
|
|
// Ensure that mcp-2 is removed once it is removed from the settings
|
|
{
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![(server_2_id.clone(), ContextServerStatus::Stopped)],
|
|
cx,
|
|
);
|
|
set_context_server_configuration(
|
|
vec![(
|
|
server_1_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Extension {
|
|
enabled: true,
|
|
remote: false,
|
|
settings: json!({
|
|
"somevalue": false
|
|
}),
|
|
},
|
|
)],
|
|
cx,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
|
|
cx.update(|cx| {
|
|
assert_eq!(store.read(cx).status_for_server(&server_2_id), None);
|
|
});
|
|
}
|
|
|
|
// Ensure that nothing happens if the settings do not change
|
|
{
|
|
let _server_events = assert_server_events(&store, vec![], cx);
|
|
set_context_server_configuration(
|
|
vec![(
|
|
server_1_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Extension {
|
|
enabled: true,
|
|
remote: false,
|
|
settings: json!({
|
|
"somevalue": false
|
|
}),
|
|
},
|
|
)],
|
|
cx,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
|
|
cx.update(|cx| {
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_1_id),
|
|
Some(ContextServerStatus::Running)
|
|
);
|
|
assert_eq!(store.read(cx).status_for_server(&server_2_id), None);
|
|
});
|
|
}
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_context_server_enabled_disabled(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;
|
|
|
|
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,
|
|
);
|
|
|
|
// Ensure that mcp-1 starts up
|
|
{
|
|
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();
|
|
}
|
|
|
|
// Ensure that mcp-1 is stopped once it is disabled.
|
|
{
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![(server_1_id.clone(), ContextServerStatus::Stopped)],
|
|
cx,
|
|
);
|
|
set_context_server_configuration(
|
|
vec![(
|
|
server_1_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Stdio {
|
|
enabled: false,
|
|
remote: false,
|
|
command: ContextServerCommand {
|
|
path: "somebinary".into(),
|
|
args: vec!["arg".to_string()],
|
|
env: None,
|
|
timeout: None,
|
|
},
|
|
},
|
|
)],
|
|
cx,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
}
|
|
|
|
// Ensure that mcp-1 is started once it is enabled again.
|
|
{
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![
|
|
(server_1_id.clone(), ContextServerStatus::Starting),
|
|
(server_1_id.clone(), ContextServerStatus::Running),
|
|
],
|
|
cx,
|
|
);
|
|
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()],
|
|
timeout: None,
|
|
env: None,
|
|
},
|
|
},
|
|
)],
|
|
cx,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
}
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_context_server_respects_disable_ai(cx: &mut TestAppContext) {
|
|
const SERVER_1_ID: &str = "mcp-1";
|
|
|
|
let server_1_id = ContextServerId(SERVER_1_ID.into());
|
|
|
|
// Set up SettingsStore with disable_ai: true in user settings BEFORE creating project
|
|
cx.update(|cx| {
|
|
let settings_store = SettingsStore::test(cx);
|
|
cx.set_global(settings_store);
|
|
DisableAiSettings::register(cx);
|
|
// Set disable_ai via user settings (not override_global) so it persists through recompute_values
|
|
SettingsStore::update_global(cx, |store, cx| {
|
|
store.update_user_settings(cx, |content| {
|
|
content.project.disable_ai = Some(SaturatingBool(true));
|
|
});
|
|
});
|
|
});
|
|
|
|
// Now create the project (ContextServerStore will see disable_ai = true)
|
|
let fs = FakeFs::new(cx.executor());
|
|
fs.insert_tree(path!("/test"), json!({"code.rs": ""})).await;
|
|
let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).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,
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
|
|
// Verify that no server started because AI is disabled
|
|
cx.update(|cx| {
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_1_id),
|
|
None,
|
|
"Server should not start when disable_ai is true"
|
|
);
|
|
});
|
|
|
|
// Enable AI and verify server starts
|
|
{
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![
|
|
(server_1_id.clone(), ContextServerStatus::Starting),
|
|
(server_1_id.clone(), ContextServerStatus::Running),
|
|
],
|
|
cx,
|
|
);
|
|
cx.update(|cx| {
|
|
SettingsStore::update_global(cx, |store, cx| {
|
|
store.update_user_settings(cx, |content| {
|
|
content.project.disable_ai = Some(SaturatingBool(false));
|
|
});
|
|
});
|
|
});
|
|
cx.run_until_parked();
|
|
}
|
|
|
|
// Disable AI again and verify server stops
|
|
{
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![(server_1_id.clone(), ContextServerStatus::Stopped)],
|
|
cx,
|
|
);
|
|
cx.update(|cx| {
|
|
SettingsStore::update_global(cx, |store, cx| {
|
|
store.update_user_settings(cx, |content| {
|
|
content.project.disable_ai = Some(SaturatingBool(true));
|
|
});
|
|
});
|
|
});
|
|
cx.run_until_parked();
|
|
}
|
|
|
|
// Verify server is stopped
|
|
cx.update(|cx| {
|
|
assert_eq!(
|
|
store.read(cx).status_for_server(&server_1_id),
|
|
Some(ContextServerStatus::Stopped),
|
|
"Server should be stopped when disable_ai is true"
|
|
);
|
|
});
|
|
}
|
|
|
|
#[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";
|
|
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,
|
|
) {
|
|
cx.update(|cx| {
|
|
SettingsStore::update_global(cx, |store, cx| {
|
|
store.update_user_settings(cx, |content| {
|
|
content.project.context_servers.clear();
|
|
for (id, config) in context_servers {
|
|
content.project.context_servers.insert(id, config);
|
|
}
|
|
});
|
|
})
|
|
});
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_remote_context_server(cx: &mut TestAppContext) {
|
|
const SERVER_ID: &str = "remote-server";
|
|
let server_id = ContextServerId(SERVER_ID.into());
|
|
let server_url = "http://example.com/api";
|
|
|
|
let client = FakeHttpClient::create(|_| async move {
|
|
use http_client::AsyncBody;
|
|
|
|
let response = Response::builder()
|
|
.status(200)
|
|
.header("Content-Type", "application/json")
|
|
.body(AsyncBody::from(
|
|
serde_json::to_string(&json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 0,
|
|
"result": {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {},
|
|
"serverInfo": {
|
|
"name": "test-server",
|
|
"version": "1.0.0"
|
|
}
|
|
}
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap();
|
|
Ok(response)
|
|
});
|
|
cx.update(|cx| cx.set_http_client(client));
|
|
|
|
let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await;
|
|
|
|
let store = project.read_with(cx, |project, _| project.context_server_store());
|
|
|
|
set_context_server_configuration(
|
|
vec![(
|
|
server_id.0.clone(),
|
|
settings::ContextServerSettingsContent::Http {
|
|
enabled: true,
|
|
url: server_url.to_string(),
|
|
headers: Default::default(),
|
|
timeout: None,
|
|
oauth: None,
|
|
},
|
|
)],
|
|
cx,
|
|
);
|
|
|
|
let _server_events = assert_server_events(
|
|
&store,
|
|
vec![
|
|
(server_id.clone(), ContextServerStatus::Starting),
|
|
(server_id.clone(), ContextServerStatus::Running),
|
|
],
|
|
cx,
|
|
);
|
|
cx.run_until_parked();
|
|
}
|
|
|
|
struct ServerEvents {
|
|
received_event_count: Rc<RefCell<usize>>,
|
|
expected_event_count: usize,
|
|
_subscription: Subscription,
|
|
}
|
|
|
|
impl Drop for ServerEvents {
|
|
fn drop(&mut self) {
|
|
let actual_event_count = *self.received_event_count.borrow();
|
|
assert_eq!(
|
|
actual_event_count, self.expected_event_count,
|
|
"
|
|
Expected to receive {} context server store events, but received {} events",
|
|
self.expected_event_count, actual_event_count
|
|
);
|
|
}
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_context_server_global_timeout(cx: &mut TestAppContext) {
|
|
cx.update(|cx| {
|
|
let settings_store = SettingsStore::test(cx);
|
|
cx.set_global(settings_store);
|
|
SettingsStore::update_global(cx, |store, cx| {
|
|
store
|
|
.set_user_settings(r#"{"context_server_timeout": 90}"#, cx)
|
|
.expect("Failed to set test user settings");
|
|
});
|
|
});
|
|
|
|
let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
|
|
|
|
let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
|
|
let store = cx.new(|cx| {
|
|
ContextServerStore::test(
|
|
registry.clone(),
|
|
project.read(cx).worktree_store(),
|
|
Some(project.downgrade()),
|
|
cx,
|
|
)
|
|
});
|
|
|
|
let mut async_cx = cx.to_async();
|
|
let result = ContextServerStore::create_context_server(
|
|
store.downgrade(),
|
|
ContextServerId("test-server".into()),
|
|
Arc::new(ContextServerConfiguration::Http {
|
|
url: url::Url::parse("http://localhost:8080").expect("Failed to parse test URL"),
|
|
headers: Default::default(),
|
|
timeout: None,
|
|
oauth: None,
|
|
}),
|
|
&mut async_cx,
|
|
)
|
|
.await;
|
|
|
|
assert!(
|
|
result.is_ok(),
|
|
"Server should be created successfully with global timeout"
|
|
);
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_context_server_per_server_timeout_override(cx: &mut TestAppContext) {
|
|
const SERVER_ID: &str = "test-server";
|
|
|
|
cx.update(|cx| {
|
|
let settings_store = SettingsStore::test(cx);
|
|
cx.set_global(settings_store);
|
|
SettingsStore::update_global(cx, |store, cx| {
|
|
store
|
|
.set_user_settings(r#"{"context_server_timeout": 60}"#, cx)
|
|
.expect("Failed to set test user settings");
|
|
});
|
|
});
|
|
|
|
let (_fs, project) = setup_context_server_test(
|
|
cx,
|
|
json!({"code.rs": ""}),
|
|
vec![(
|
|
SERVER_ID.into(),
|
|
ContextServerSettings::Http {
|
|
enabled: true,
|
|
url: "http://localhost:8080".to_string(),
|
|
headers: Default::default(),
|
|
timeout: Some(120),
|
|
oauth: None,
|
|
},
|
|
)],
|
|
)
|
|
.await;
|
|
|
|
let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
|
|
let store = cx.new(|cx| {
|
|
ContextServerStore::test(
|
|
registry.clone(),
|
|
project.read(cx).worktree_store(),
|
|
Some(project.downgrade()),
|
|
cx,
|
|
)
|
|
});
|
|
|
|
let mut async_cx = cx.to_async();
|
|
let result = ContextServerStore::create_context_server(
|
|
store.downgrade(),
|
|
ContextServerId("test-server".into()),
|
|
Arc::new(ContextServerConfiguration::Http {
|
|
url: url::Url::parse("http://localhost:8080").expect("Failed to parse test URL"),
|
|
headers: Default::default(),
|
|
timeout: Some(120),
|
|
oauth: None,
|
|
}),
|
|
&mut async_cx,
|
|
)
|
|
.await;
|
|
|
|
assert!(
|
|
result.is_ok(),
|
|
"Server should be created successfully with per-server timeout override"
|
|
);
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_context_server_stdio_timeout(cx: &mut TestAppContext) {
|
|
let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
|
|
|
|
let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
|
|
let store = cx.new(|cx| {
|
|
ContextServerStore::test(
|
|
registry.clone(),
|
|
project.read(cx).worktree_store(),
|
|
Some(project.downgrade()),
|
|
cx,
|
|
)
|
|
});
|
|
|
|
let mut async_cx = cx.to_async();
|
|
let result = ContextServerStore::create_context_server(
|
|
store.downgrade(),
|
|
ContextServerId("stdio-server".into()),
|
|
Arc::new(ContextServerConfiguration::Custom {
|
|
command: ContextServerCommand {
|
|
path: "/usr/bin/node".into(),
|
|
args: vec!["server.js".into()],
|
|
env: None,
|
|
timeout: Some(180000),
|
|
},
|
|
remote: false,
|
|
}),
|
|
&mut async_cx,
|
|
)
|
|
.await;
|
|
|
|
assert!(
|
|
result.is_ok(),
|
|
"Stdio server should be created successfully with timeout"
|
|
);
|
|
}
|
|
|
|
fn assert_server_events(
|
|
store: &Entity<ContextServerStore>,
|
|
expected_events: Vec<(ContextServerId, ContextServerStatus)>,
|
|
cx: &mut TestAppContext,
|
|
) -> ServerEvents {
|
|
cx.update(|cx| {
|
|
let mut ix = 0;
|
|
let received_event_count = Rc::new(RefCell::new(0));
|
|
let expected_event_count = expected_events.len();
|
|
let subscription = cx.subscribe(store, {
|
|
let received_event_count = received_event_count.clone();
|
|
move |_, event, _| {
|
|
let ServerStatusChangedEvent {
|
|
server_id: actual_server_id,
|
|
status: actual_status,
|
|
} = 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;
|
|
}
|
|
});
|
|
ServerEvents {
|
|
expected_event_count,
|
|
received_event_count,
|
|
_subscription: subscription,
|
|
}
|
|
})
|
|
}
|
|
|
|
async fn setup_context_server_test(
|
|
cx: &mut TestAppContext,
|
|
files: serde_json::Value,
|
|
context_server_configurations: Vec<(Arc<str>, ContextServerSettings)>,
|
|
) -> (Arc<FakeFs>, Entity<Project>) {
|
|
cx.update(|cx| {
|
|
let settings_store = SettingsStore::test(cx);
|
|
cx.set_global(settings_store);
|
|
let mut settings = ProjectSettings::get_global(cx).clone();
|
|
for (id, config) in context_server_configurations {
|
|
settings.context_servers.insert(id, config);
|
|
}
|
|
ProjectSettings::override_global(settings, cx);
|
|
});
|
|
|
|
let fs = FakeFs::new(cx.executor());
|
|
fs.insert_tree(path!("/test"), files).await;
|
|
let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
|
|
|
|
(fs, project)
|
|
}
|
|
|
|
struct FakeContextServerDescriptor {
|
|
path: PathBuf,
|
|
}
|
|
|
|
impl FakeContextServerDescriptor {
|
|
fn new(path: impl Into<PathBuf>) -> Self {
|
|
Self { path: path.into() }
|
|
}
|
|
}
|
|
|
|
impl ContextServerDescriptor for FakeContextServerDescriptor {
|
|
fn command(
|
|
&self,
|
|
_worktree_store: Entity<WorktreeStore>,
|
|
_cx: &AsyncApp,
|
|
) -> Task<Result<ContextServerCommand>> {
|
|
Task::ready(Ok(ContextServerCommand {
|
|
path: self.path.clone(),
|
|
args: vec!["arg1".to_string(), "arg2".to_string()],
|
|
env: None,
|
|
timeout: None,
|
|
}))
|
|
}
|
|
|
|
fn configuration(
|
|
&self,
|
|
_worktree_store: Entity<WorktreeStore>,
|
|
_cx: &AsyncApp,
|
|
) -> Task<Result<Option<::extension::ContextServerConfiguration>>> {
|
|
Task::ready(Ok(None))
|
|
}
|
|
}
|