Apoorva Vermaproject: Restart context servers when their working directory changes (#61365)ac44e9f3 months agoHistory
1932 lines · 61.9 KB · RustRawDownload
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.    {

Lines 1–500 of 1932.