context_server_store.rs

   1use anyhow::Result;
   2use context_server::test::create_fake_transport;
   3use context_server::{ContextServer, ContextServerId};
   4use gpui::{AppContext, AsyncApp, Entity, Subscription, Task, TestAppContext, UpdateGlobal as _};
   5use http_client::{FakeHttpClient, Response};
   6use project::context_server_store::registry::ContextServerDescriptorRegistry;
   7use project::context_server_store::*;
   8use project::project_settings::ContextServerSettings;
   9use project::worktree_store::WorktreeStore;
  10use project::{
  11    DisableAiSettings, FakeFs, Project, context_server_store::registry::ContextServerDescriptor,
  12    project_settings::ProjectSettings,
  13};
  14use serde_json::json;
  15use settings::settings_content::SaturatingBool;
  16use settings::{ContextServerCommand, Settings, SettingsStore};
  17use std::sync::Arc;
  18use std::{cell::RefCell, path::PathBuf, rc::Rc};
  19use util::path;
  20
  21#[gpui::test]
  22async fn test_context_server_status(cx: &mut TestAppContext) {
  23    const SERVER_1_ID: &str = "mcp-1";
  24    const SERVER_2_ID: &str = "mcp-2";
  25
  26    let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
  27
  28    let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
  29    let store = cx.new(|cx| {
  30        ContextServerStore::test(
  31            registry.clone(),
  32            project.read(cx).worktree_store(),
  33            Some(project.downgrade()),
  34            cx,
  35        )
  36    });
  37
  38    let server_1_id = ContextServerId(SERVER_1_ID.into());
  39    let server_2_id = ContextServerId(SERVER_2_ID.into());
  40
  41    let server_1 = Arc::new(ContextServer::new(
  42        server_1_id.clone(),
  43        Arc::new(create_fake_transport(SERVER_1_ID, cx.executor())),
  44    ));
  45    let server_2 = Arc::new(ContextServer::new(
  46        server_2_id.clone(),
  47        Arc::new(create_fake_transport(SERVER_2_ID, cx.executor())),
  48    ));
  49
  50    store.update(cx, |store, cx| store.test_start_server(server_1, cx));
  51
  52    cx.run_until_parked();
  53
  54    cx.update(|cx| {
  55        assert_eq!(
  56            store.read(cx).status_for_server(&server_1_id),
  57            Some(ContextServerStatus::Running)
  58        );
  59        assert_eq!(store.read(cx).status_for_server(&server_2_id), None);
  60    });
  61
  62    store.update(cx, |store, cx| {
  63        store.test_start_server(server_2.clone(), cx)
  64    });
  65
  66    cx.run_until_parked();
  67
  68    cx.update(|cx| {
  69        assert_eq!(
  70            store.read(cx).status_for_server(&server_1_id),
  71            Some(ContextServerStatus::Running)
  72        );
  73        assert_eq!(
  74            store.read(cx).status_for_server(&server_2_id),
  75            Some(ContextServerStatus::Running)
  76        );
  77    });
  78
  79    store
  80        .update(cx, |store, cx| store.stop_server(&server_2_id, cx))
  81        .unwrap();
  82
  83    cx.update(|cx| {
  84        assert_eq!(
  85            store.read(cx).status_for_server(&server_1_id),
  86            Some(ContextServerStatus::Running)
  87        );
  88        assert_eq!(
  89            store.read(cx).status_for_server(&server_2_id),
  90            Some(ContextServerStatus::Stopped)
  91        );
  92    });
  93}
  94
  95#[gpui::test]
  96async fn test_context_server_status_events(cx: &mut TestAppContext) {
  97    const SERVER_1_ID: &str = "mcp-1";
  98    const SERVER_2_ID: &str = "mcp-2";
  99
 100    let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
 101
 102    let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
 103    let store = cx.new(|cx| {
 104        ContextServerStore::test(
 105            registry.clone(),
 106            project.read(cx).worktree_store(),
 107            Some(project.downgrade()),
 108            cx,
 109        )
 110    });
 111
 112    let server_1_id = ContextServerId(SERVER_1_ID.into());
 113    let server_2_id = ContextServerId(SERVER_2_ID.into());
 114
 115    let server_1 = Arc::new(ContextServer::new(
 116        server_1_id.clone(),
 117        Arc::new(create_fake_transport(SERVER_1_ID, cx.executor())),
 118    ));
 119    let server_2 = Arc::new(ContextServer::new(
 120        server_2_id.clone(),
 121        Arc::new(create_fake_transport(SERVER_2_ID, cx.executor())),
 122    ));
 123
 124    let _server_events = assert_server_events(
 125        &store,
 126        vec![
 127            (server_1_id.clone(), ContextServerStatus::Starting),
 128            (server_1_id, ContextServerStatus::Running),
 129            (server_2_id.clone(), ContextServerStatus::Starting),
 130            (server_2_id.clone(), ContextServerStatus::Running),
 131            (server_2_id.clone(), ContextServerStatus::Stopped),
 132        ],
 133        cx,
 134    );
 135
 136    store.update(cx, |store, cx| store.test_start_server(server_1, cx));
 137
 138    cx.run_until_parked();
 139
 140    store.update(cx, |store, cx| {
 141        store.test_start_server(server_2.clone(), cx)
 142    });
 143
 144    cx.run_until_parked();
 145
 146    store
 147        .update(cx, |store, cx| store.stop_server(&server_2_id, cx))
 148        .unwrap();
 149}
 150
 151#[gpui::test(iterations = 25)]
 152async fn test_context_server_concurrent_starts(cx: &mut TestAppContext) {
 153    const SERVER_1_ID: &str = "mcp-1";
 154
 155    let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
 156
 157    let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
 158    let store = cx.new(|cx| {
 159        ContextServerStore::test(
 160            registry.clone(),
 161            project.read(cx).worktree_store(),
 162            Some(project.downgrade()),
 163            cx,
 164        )
 165    });
 166
 167    let server_id = ContextServerId(SERVER_1_ID.into());
 168
 169    let server_with_same_id_1 = Arc::new(ContextServer::new(
 170        server_id.clone(),
 171        Arc::new(create_fake_transport(SERVER_1_ID, cx.executor())),
 172    ));
 173    let server_with_same_id_2 = Arc::new(ContextServer::new(
 174        server_id.clone(),
 175        Arc::new(create_fake_transport(SERVER_1_ID, cx.executor())),
 176    ));
 177
 178    // If we start another server with the same id, we should report that we stopped the previous one
 179    let _server_events = assert_server_events(
 180        &store,
 181        vec![
 182            (server_id.clone(), ContextServerStatus::Starting),
 183            (server_id.clone(), ContextServerStatus::Stopped),
 184            (server_id.clone(), ContextServerStatus::Starting),
 185            (server_id.clone(), ContextServerStatus::Running),
 186        ],
 187        cx,
 188    );
 189
 190    store.update(cx, |store, cx| {
 191        store.test_start_server(server_with_same_id_1.clone(), cx)
 192    });
 193    store.update(cx, |store, cx| {
 194        store.test_start_server(server_with_same_id_2.clone(), cx)
 195    });
 196
 197    cx.run_until_parked();
 198
 199    cx.update(|cx| {
 200        assert_eq!(
 201            store.read(cx).status_for_server(&server_id),
 202            Some(ContextServerStatus::Running)
 203        );
 204    });
 205}
 206
 207#[gpui::test]
 208async fn test_context_server_maintain_servers_loop(cx: &mut TestAppContext) {
 209    const SERVER_1_ID: &str = "mcp-1";
 210    const SERVER_2_ID: &str = "mcp-2";
 211
 212    let server_1_id = ContextServerId(SERVER_1_ID.into());
 213    let server_2_id = ContextServerId(SERVER_2_ID.into());
 214
 215    let fake_descriptor_1 = Arc::new(FakeContextServerDescriptor::new(SERVER_1_ID));
 216
 217    let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
 218
 219    let executor = cx.executor();
 220    let store = project.read_with(cx, |project, _| project.context_server_store());
 221    store.update(cx, |store, cx| {
 222        store.set_context_server_factory(Box::new(move |id, _| {
 223            Arc::new(ContextServer::new(
 224                id.clone(),
 225                Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
 226            ))
 227        }));
 228        store.registry().update(cx, |registry, cx| {
 229            registry.register_context_server_descriptor(SERVER_1_ID.into(), fake_descriptor_1, cx);
 230        });
 231    });
 232
 233    set_context_server_configuration(
 234        vec![(
 235            server_1_id.0.clone(),
 236            settings::ContextServerSettingsContent::Extension {
 237                enabled: true,
 238                remote: false,
 239                settings: json!({
 240                    "somevalue": true
 241                }),
 242            },
 243        )],
 244        cx,
 245    );
 246
 247    // Ensure that mcp-1 starts up
 248    {
 249        let _server_events = assert_server_events(
 250            &store,
 251            vec![
 252                (server_1_id.clone(), ContextServerStatus::Starting),
 253                (server_1_id.clone(), ContextServerStatus::Running),
 254            ],
 255            cx,
 256        );
 257        cx.run_until_parked();
 258    }
 259
 260    // Ensure that mcp-1 is restarted when the configuration was changed
 261    {
 262        let _server_events = assert_server_events(
 263            &store,
 264            vec![
 265                (server_1_id.clone(), ContextServerStatus::Stopped),
 266                (server_1_id.clone(), ContextServerStatus::Starting),
 267                (server_1_id.clone(), ContextServerStatus::Running),
 268            ],
 269            cx,
 270        );
 271        set_context_server_configuration(
 272            vec![(
 273                server_1_id.0.clone(),
 274                settings::ContextServerSettingsContent::Extension {
 275                    enabled: true,
 276                    remote: false,
 277                    settings: json!({
 278                        "somevalue": false
 279                    }),
 280                },
 281            )],
 282            cx,
 283        );
 284
 285        cx.run_until_parked();
 286    }
 287
 288    // Ensure that mcp-1 is not restarted when the configuration was not changed
 289    {
 290        let _server_events = assert_server_events(&store, vec![], cx);
 291        set_context_server_configuration(
 292            vec![(
 293                server_1_id.0.clone(),
 294                settings::ContextServerSettingsContent::Extension {
 295                    enabled: true,
 296                    remote: false,
 297                    settings: json!({
 298                        "somevalue": false
 299                    }),
 300                },
 301            )],
 302            cx,
 303        );
 304
 305        cx.run_until_parked();
 306    }
 307
 308    // Ensure that mcp-2 is started once it is added to the settings
 309    {
 310        let _server_events = assert_server_events(
 311            &store,
 312            vec![
 313                (server_2_id.clone(), ContextServerStatus::Starting),
 314                (server_2_id.clone(), ContextServerStatus::Running),
 315            ],
 316            cx,
 317        );
 318        set_context_server_configuration(
 319            vec![
 320                (
 321                    server_1_id.0.clone(),
 322                    settings::ContextServerSettingsContent::Extension {
 323                        enabled: true,
 324                        remote: false,
 325                        settings: json!({
 326                            "somevalue": false
 327                        }),
 328                    },
 329                ),
 330                (
 331                    server_2_id.0.clone(),
 332                    settings::ContextServerSettingsContent::Stdio {
 333                        enabled: true,
 334                        remote: false,
 335                        command: ContextServerCommand {
 336                            path: "somebinary".into(),
 337                            args: vec!["arg".to_string()],
 338                            env: None,
 339                            timeout: None,
 340                        },
 341                    },
 342                ),
 343            ],
 344            cx,
 345        );
 346
 347        cx.run_until_parked();
 348    }
 349
 350    // Ensure that mcp-2 is restarted once the args have changed
 351    {
 352        let _server_events = assert_server_events(
 353            &store,
 354            vec![
 355                (server_2_id.clone(), ContextServerStatus::Stopped),
 356                (server_2_id.clone(), ContextServerStatus::Starting),
 357                (server_2_id.clone(), ContextServerStatus::Running),
 358            ],
 359            cx,
 360        );
 361        set_context_server_configuration(
 362            vec![
 363                (
 364                    server_1_id.0.clone(),
 365                    settings::ContextServerSettingsContent::Extension {
 366                        enabled: true,
 367                        remote: false,
 368                        settings: json!({
 369                            "somevalue": false
 370                        }),
 371                    },
 372                ),
 373                (
 374                    server_2_id.0.clone(),
 375                    settings::ContextServerSettingsContent::Stdio {
 376                        enabled: true,
 377                        remote: false,
 378                        command: ContextServerCommand {
 379                            path: "somebinary".into(),
 380                            args: vec!["anotherArg".to_string()],
 381                            env: None,
 382                            timeout: None,
 383                        },
 384                    },
 385                ),
 386            ],
 387            cx,
 388        );
 389
 390        cx.run_until_parked();
 391    }
 392
 393    // Ensure that mcp-2 is removed once it is removed from the settings
 394    {
 395        let _server_events = assert_server_events(
 396            &store,
 397            vec![(server_2_id.clone(), ContextServerStatus::Stopped)],
 398            cx,
 399        );
 400        set_context_server_configuration(
 401            vec![(
 402                server_1_id.0.clone(),
 403                settings::ContextServerSettingsContent::Extension {
 404                    enabled: true,
 405                    remote: false,
 406                    settings: json!({
 407                        "somevalue": false
 408                    }),
 409                },
 410            )],
 411            cx,
 412        );
 413
 414        cx.run_until_parked();
 415
 416        cx.update(|cx| {
 417            assert_eq!(store.read(cx).status_for_server(&server_2_id), None);
 418        });
 419    }
 420
 421    // Ensure that nothing happens if the settings do not change
 422    {
 423        let _server_events = assert_server_events(&store, vec![], cx);
 424        set_context_server_configuration(
 425            vec![(
 426                server_1_id.0.clone(),
 427                settings::ContextServerSettingsContent::Extension {
 428                    enabled: true,
 429                    remote: false,
 430                    settings: json!({
 431                        "somevalue": false
 432                    }),
 433                },
 434            )],
 435            cx,
 436        );
 437
 438        cx.run_until_parked();
 439
 440        cx.update(|cx| {
 441            assert_eq!(
 442                store.read(cx).status_for_server(&server_1_id),
 443                Some(ContextServerStatus::Running)
 444            );
 445            assert_eq!(store.read(cx).status_for_server(&server_2_id), None);
 446        });
 447    }
 448}
 449
 450#[gpui::test]
 451async fn test_context_server_enabled_disabled(cx: &mut TestAppContext) {
 452    const SERVER_1_ID: &str = "mcp-1";
 453
 454    let server_1_id = ContextServerId(SERVER_1_ID.into());
 455
 456    let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
 457
 458    let executor = cx.executor();
 459    let store = project.read_with(cx, |project, _| project.context_server_store());
 460    store.update(cx, |store, _| {
 461        store.set_context_server_factory(Box::new(move |id, _| {
 462            Arc::new(ContextServer::new(
 463                id.clone(),
 464                Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
 465            ))
 466        }));
 467    });
 468
 469    set_context_server_configuration(
 470        vec![(
 471            server_1_id.0.clone(),
 472            settings::ContextServerSettingsContent::Stdio {
 473                enabled: true,
 474                remote: false,
 475                command: ContextServerCommand {
 476                    path: "somebinary".into(),
 477                    args: vec!["arg".to_string()],
 478                    env: None,
 479                    timeout: None,
 480                },
 481            },
 482        )],
 483        cx,
 484    );
 485
 486    // Ensure that mcp-1 starts up
 487    {
 488        let _server_events = assert_server_events(
 489            &store,
 490            vec![
 491                (server_1_id.clone(), ContextServerStatus::Starting),
 492                (server_1_id.clone(), ContextServerStatus::Running),
 493            ],
 494            cx,
 495        );
 496        cx.run_until_parked();
 497    }
 498
 499    // Ensure that mcp-1 is stopped once it is disabled.
 500    {
 501        let _server_events = assert_server_events(
 502            &store,
 503            vec![(server_1_id.clone(), ContextServerStatus::Stopped)],
 504            cx,
 505        );
 506        set_context_server_configuration(
 507            vec![(
 508                server_1_id.0.clone(),
 509                settings::ContextServerSettingsContent::Stdio {
 510                    enabled: false,
 511                    remote: false,
 512                    command: ContextServerCommand {
 513                        path: "somebinary".into(),
 514                        args: vec!["arg".to_string()],
 515                        env: None,
 516                        timeout: None,
 517                    },
 518                },
 519            )],
 520            cx,
 521        );
 522
 523        cx.run_until_parked();
 524    }
 525
 526    // Ensure that mcp-1 is started once it is enabled again.
 527    {
 528        let _server_events = assert_server_events(
 529            &store,
 530            vec![
 531                (server_1_id.clone(), ContextServerStatus::Starting),
 532                (server_1_id.clone(), ContextServerStatus::Running),
 533            ],
 534            cx,
 535        );
 536        set_context_server_configuration(
 537            vec![(
 538                server_1_id.0.clone(),
 539                settings::ContextServerSettingsContent::Stdio {
 540                    enabled: true,
 541                    remote: false,
 542                    command: ContextServerCommand {
 543                        path: "somebinary".into(),
 544                        args: vec!["arg".to_string()],
 545                        timeout: None,
 546                        env: None,
 547                    },
 548                },
 549            )],
 550            cx,
 551        );
 552
 553        cx.run_until_parked();
 554    }
 555}
 556
 557#[gpui::test]
 558async fn test_context_server_respects_disable_ai(cx: &mut TestAppContext) {
 559    const SERVER_1_ID: &str = "mcp-1";
 560
 561    let server_1_id = ContextServerId(SERVER_1_ID.into());
 562
 563    // Set up SettingsStore with disable_ai: true in user settings BEFORE creating project
 564    cx.update(|cx| {
 565        let settings_store = SettingsStore::test(cx);
 566        cx.set_global(settings_store);
 567        DisableAiSettings::register(cx);
 568        // Set disable_ai via user settings (not override_global) so it persists through recompute_values
 569        SettingsStore::update_global(cx, |store, cx| {
 570            store.update_user_settings(cx, |content| {
 571                content.project.disable_ai = Some(SaturatingBool(true));
 572            });
 573        });
 574    });
 575
 576    // Now create the project (ContextServerStore will see disable_ai = true)
 577    let fs = FakeFs::new(cx.executor());
 578    fs.insert_tree(path!("/test"), json!({"code.rs": ""})).await;
 579    let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
 580
 581    let executor = cx.executor();
 582    let store = project.read_with(cx, |project, _| project.context_server_store());
 583    store.update(cx, |store, _| {
 584        store.set_context_server_factory(Box::new(move |id, _| {
 585            Arc::new(ContextServer::new(
 586                id.clone(),
 587                Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
 588            ))
 589        }));
 590    });
 591
 592    set_context_server_configuration(
 593        vec![(
 594            server_1_id.0.clone(),
 595            settings::ContextServerSettingsContent::Stdio {
 596                enabled: true,
 597                remote: false,
 598                command: ContextServerCommand {
 599                    path: "somebinary".into(),
 600                    args: vec!["arg".to_string()],
 601                    env: None,
 602                    timeout: None,
 603                },
 604            },
 605        )],
 606        cx,
 607    );
 608
 609    cx.run_until_parked();
 610
 611    // Verify that no server started because AI is disabled
 612    cx.update(|cx| {
 613        assert_eq!(
 614            store.read(cx).status_for_server(&server_1_id),
 615            None,
 616            "Server should not start when disable_ai is true"
 617        );
 618    });
 619
 620    // Enable AI and verify server starts
 621    {
 622        let _server_events = assert_server_events(
 623            &store,
 624            vec![
 625                (server_1_id.clone(), ContextServerStatus::Starting),
 626                (server_1_id.clone(), ContextServerStatus::Running),
 627            ],
 628            cx,
 629        );
 630        cx.update(|cx| {
 631            SettingsStore::update_global(cx, |store, cx| {
 632                store.update_user_settings(cx, |content| {
 633                    content.project.disable_ai = Some(SaturatingBool(false));
 634                });
 635            });
 636        });
 637        cx.run_until_parked();
 638    }
 639
 640    // Disable AI again and verify server stops
 641    {
 642        let _server_events = assert_server_events(
 643            &store,
 644            vec![(server_1_id.clone(), ContextServerStatus::Stopped)],
 645            cx,
 646        );
 647        cx.update(|cx| {
 648            SettingsStore::update_global(cx, |store, cx| {
 649                store.update_user_settings(cx, |content| {
 650                    content.project.disable_ai = Some(SaturatingBool(true));
 651                });
 652            });
 653        });
 654        cx.run_until_parked();
 655    }
 656
 657    // Verify server is stopped
 658    cx.update(|cx| {
 659        assert_eq!(
 660            store.read(cx).status_for_server(&server_1_id),
 661            Some(ContextServerStatus::Stopped),
 662            "Server should be stopped when disable_ai is true"
 663        );
 664    });
 665}
 666
 667#[gpui::test]
 668async fn test_server_ids_includes_disabled_servers(cx: &mut TestAppContext) {
 669    const ENABLED_SERVER_ID: &str = "enabled-server";
 670    const DISABLED_SERVER_ID: &str = "disabled-server";
 671
 672    let enabled_server_id = ContextServerId(ENABLED_SERVER_ID.into());
 673    let disabled_server_id = ContextServerId(DISABLED_SERVER_ID.into());
 674
 675    let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
 676
 677    let executor = cx.executor();
 678    let store = project.read_with(cx, |project, _| project.context_server_store());
 679    store.update(cx, |store, _| {
 680        store.set_context_server_factory(Box::new(move |id, _| {
 681            Arc::new(ContextServer::new(
 682                id.clone(),
 683                Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
 684            ))
 685        }));
 686    });
 687
 688    // Configure one enabled and one disabled server
 689    set_context_server_configuration(
 690        vec![
 691            (
 692                enabled_server_id.0.clone(),
 693                settings::ContextServerSettingsContent::Stdio {
 694                    enabled: true,
 695                    remote: false,
 696                    command: ContextServerCommand {
 697                        path: "somebinary".into(),
 698                        args: vec![],
 699                        env: None,
 700                        timeout: None,
 701                    },
 702                },
 703            ),
 704            (
 705                disabled_server_id.0.clone(),
 706                settings::ContextServerSettingsContent::Stdio {
 707                    enabled: false,
 708                    remote: false,
 709                    command: ContextServerCommand {
 710                        path: "somebinary".into(),
 711                        args: vec![],
 712                        env: None,
 713                        timeout: None,
 714                    },
 715                },
 716            ),
 717        ],
 718        cx,
 719    );
 720
 721    cx.run_until_parked();
 722
 723    // Verify that server_ids includes both enabled and disabled servers
 724    cx.update(|cx| {
 725        let server_ids = store.read(cx).server_ids().to_vec();
 726        assert!(
 727            server_ids.contains(&enabled_server_id),
 728            "server_ids should include enabled server"
 729        );
 730        assert!(
 731            server_ids.contains(&disabled_server_id),
 732            "server_ids should include disabled server"
 733        );
 734    });
 735
 736    // Verify that the enabled server is running and the disabled server is not
 737    cx.read(|cx| {
 738        assert_eq!(
 739            store.read(cx).status_for_server(&enabled_server_id),
 740            Some(ContextServerStatus::Running),
 741            "enabled server should be running"
 742        );
 743        // Disabled server should not be in the servers map (status returns None)
 744        // but should still be in server_ids
 745        assert_eq!(
 746            store.read(cx).status_for_server(&disabled_server_id),
 747            None,
 748            "disabled server should not have a status (not in servers map)"
 749        );
 750    });
 751}
 752
 753fn set_context_server_configuration(
 754    context_servers: Vec<(Arc<str>, settings::ContextServerSettingsContent)>,
 755    cx: &mut TestAppContext,
 756) {
 757    cx.update(|cx| {
 758        SettingsStore::update_global(cx, |store, cx| {
 759            store.update_user_settings(cx, |content| {
 760                content.project.context_servers.clear();
 761                for (id, config) in context_servers {
 762                    content.project.context_servers.insert(id, config);
 763                }
 764            });
 765        })
 766    });
 767}
 768
 769#[gpui::test]
 770async fn test_remote_context_server(cx: &mut TestAppContext) {
 771    const SERVER_ID: &str = "remote-server";
 772    let server_id = ContextServerId(SERVER_ID.into());
 773    let server_url = "http://example.com/api";
 774
 775    let client = FakeHttpClient::create(|_| async move {
 776        use http_client::AsyncBody;
 777
 778        let response = Response::builder()
 779            .status(200)
 780            .header("Content-Type", "application/json")
 781            .body(AsyncBody::from(
 782                serde_json::to_string(&json!({
 783                    "jsonrpc": "2.0",
 784                    "id": 0,
 785                    "result": {
 786                        "protocolVersion": "2024-11-05",
 787                        "capabilities": {},
 788                        "serverInfo": {
 789                            "name": "test-server",
 790                            "version": "1.0.0"
 791                        }
 792                    }
 793                }))
 794                .unwrap(),
 795            ))
 796            .unwrap();
 797        Ok(response)
 798    });
 799    cx.update(|cx| cx.set_http_client(client));
 800
 801    let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await;
 802
 803    let store = project.read_with(cx, |project, _| project.context_server_store());
 804
 805    set_context_server_configuration(
 806        vec![(
 807            server_id.0.clone(),
 808            settings::ContextServerSettingsContent::Http {
 809                enabled: true,
 810                url: server_url.to_string(),
 811                headers: Default::default(),
 812                timeout: None,
 813            },
 814        )],
 815        cx,
 816    );
 817
 818    let _server_events = assert_server_events(
 819        &store,
 820        vec![
 821            (server_id.clone(), ContextServerStatus::Starting),
 822            (server_id.clone(), ContextServerStatus::Running),
 823        ],
 824        cx,
 825    );
 826    cx.run_until_parked();
 827}
 828
 829struct ServerEvents {
 830    received_event_count: Rc<RefCell<usize>>,
 831    expected_event_count: usize,
 832    _subscription: Subscription,
 833}
 834
 835impl Drop for ServerEvents {
 836    fn drop(&mut self) {
 837        let actual_event_count = *self.received_event_count.borrow();
 838        assert_eq!(
 839            actual_event_count, self.expected_event_count,
 840            "
 841               Expected to receive {} context server store events, but received {} events",
 842            self.expected_event_count, actual_event_count
 843        );
 844    }
 845}
 846
 847#[gpui::test]
 848async fn test_context_server_global_timeout(cx: &mut TestAppContext) {
 849    cx.update(|cx| {
 850        let settings_store = SettingsStore::test(cx);
 851        cx.set_global(settings_store);
 852        SettingsStore::update_global(cx, |store, cx| {
 853            store
 854                .set_user_settings(r#"{"context_server_timeout": 90}"#, cx)
 855                .expect("Failed to set test user settings");
 856        });
 857    });
 858
 859    let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
 860
 861    let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
 862    let store = cx.new(|cx| {
 863        ContextServerStore::test(
 864            registry.clone(),
 865            project.read(cx).worktree_store(),
 866            Some(project.downgrade()),
 867            cx,
 868        )
 869    });
 870
 871    let mut async_cx = cx.to_async();
 872    let result = ContextServerStore::create_context_server(
 873        store.downgrade(),
 874        ContextServerId("test-server".into()),
 875        Arc::new(ContextServerConfiguration::Http {
 876            url: url::Url::parse("http://localhost:8080").expect("Failed to parse test URL"),
 877            headers: Default::default(),
 878            timeout: None,
 879        }),
 880        &mut async_cx,
 881    )
 882    .await;
 883
 884    assert!(
 885        result.is_ok(),
 886        "Server should be created successfully with global timeout"
 887    );
 888}
 889
 890#[gpui::test]
 891async fn test_context_server_per_server_timeout_override(cx: &mut TestAppContext) {
 892    const SERVER_ID: &str = "test-server";
 893
 894    cx.update(|cx| {
 895        let settings_store = SettingsStore::test(cx);
 896        cx.set_global(settings_store);
 897        SettingsStore::update_global(cx, |store, cx| {
 898            store
 899                .set_user_settings(r#"{"context_server_timeout": 60}"#, cx)
 900                .expect("Failed to set test user settings");
 901        });
 902    });
 903
 904    let (_fs, project) = setup_context_server_test(
 905        cx,
 906        json!({"code.rs": ""}),
 907        vec![(
 908            SERVER_ID.into(),
 909            ContextServerSettings::Http {
 910                enabled: true,
 911                url: "http://localhost:8080".to_string(),
 912                headers: Default::default(),
 913                timeout: Some(120),
 914            },
 915        )],
 916    )
 917    .await;
 918
 919    let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
 920    let store = cx.new(|cx| {
 921        ContextServerStore::test(
 922            registry.clone(),
 923            project.read(cx).worktree_store(),
 924            Some(project.downgrade()),
 925            cx,
 926        )
 927    });
 928
 929    let mut async_cx = cx.to_async();
 930    let result = ContextServerStore::create_context_server(
 931        store.downgrade(),
 932        ContextServerId("test-server".into()),
 933        Arc::new(ContextServerConfiguration::Http {
 934            url: url::Url::parse("http://localhost:8080").expect("Failed to parse test URL"),
 935            headers: Default::default(),
 936            timeout: Some(120),
 937        }),
 938        &mut async_cx,
 939    )
 940    .await;
 941
 942    assert!(
 943        result.is_ok(),
 944        "Server should be created successfully with per-server timeout override"
 945    );
 946}
 947
 948#[gpui::test]
 949async fn test_context_server_stdio_timeout(cx: &mut TestAppContext) {
 950    let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
 951
 952    let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
 953    let store = cx.new(|cx| {
 954        ContextServerStore::test(
 955            registry.clone(),
 956            project.read(cx).worktree_store(),
 957            Some(project.downgrade()),
 958            cx,
 959        )
 960    });
 961
 962    let mut async_cx = cx.to_async();
 963    let result = ContextServerStore::create_context_server(
 964        store.downgrade(),
 965        ContextServerId("stdio-server".into()),
 966        Arc::new(ContextServerConfiguration::Custom {
 967            command: ContextServerCommand {
 968                path: "/usr/bin/node".into(),
 969                args: vec!["server.js".into()],
 970                env: None,
 971                timeout: Some(180000),
 972            },
 973            remote: false,
 974        }),
 975        &mut async_cx,
 976    )
 977    .await;
 978
 979    assert!(
 980        result.is_ok(),
 981        "Stdio server should be created successfully with timeout"
 982    );
 983}
 984
 985fn assert_server_events(
 986    store: &Entity<ContextServerStore>,
 987    expected_events: Vec<(ContextServerId, ContextServerStatus)>,
 988    cx: &mut TestAppContext,
 989) -> ServerEvents {
 990    cx.update(|cx| {
 991        let mut ix = 0;
 992        let received_event_count = Rc::new(RefCell::new(0));
 993        let expected_event_count = expected_events.len();
 994        let subscription = cx.subscribe(store, {
 995            let received_event_count = received_event_count.clone();
 996            move |_, event, _| {
 997                let ServerStatusChangedEvent {
 998                    server_id: actual_server_id,
 999                    status: actual_status,
1000                } = event;
1001                let (expected_server_id, expected_status) = &expected_events[ix];
1002
1003                assert_eq!(
1004                    actual_server_id, expected_server_id,
1005                    "Expected different server id at index {}",
1006                    ix
1007                );
1008                assert_eq!(
1009                    actual_status, expected_status,
1010                    "Expected different status at index {}",
1011                    ix
1012                );
1013                ix += 1;
1014                *received_event_count.borrow_mut() += 1;
1015            }
1016        });
1017        ServerEvents {
1018            expected_event_count,
1019            received_event_count,
1020            _subscription: subscription,
1021        }
1022    })
1023}
1024
1025async fn setup_context_server_test(
1026    cx: &mut TestAppContext,
1027    files: serde_json::Value,
1028    context_server_configurations: Vec<(Arc<str>, ContextServerSettings)>,
1029) -> (Arc<FakeFs>, Entity<Project>) {
1030    cx.update(|cx| {
1031        let settings_store = SettingsStore::test(cx);
1032        cx.set_global(settings_store);
1033        let mut settings = ProjectSettings::get_global(cx).clone();
1034        for (id, config) in context_server_configurations {
1035            settings.context_servers.insert(id, config);
1036        }
1037        ProjectSettings::override_global(settings, cx);
1038    });
1039
1040    let fs = FakeFs::new(cx.executor());
1041    fs.insert_tree(path!("/test"), files).await;
1042    let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
1043
1044    (fs, project)
1045}
1046
1047struct FakeContextServerDescriptor {
1048    path: PathBuf,
1049}
1050
1051impl FakeContextServerDescriptor {
1052    fn new(path: impl Into<PathBuf>) -> Self {
1053        Self { path: path.into() }
1054    }
1055}
1056
1057impl ContextServerDescriptor for FakeContextServerDescriptor {
1058    fn command(
1059        &self,
1060        _worktree_store: Entity<WorktreeStore>,
1061        _cx: &AsyncApp,
1062    ) -> Task<Result<ContextServerCommand>> {
1063        Task::ready(Ok(ContextServerCommand {
1064            path: self.path.clone(),
1065            args: vec!["arg1".to_string(), "arg2".to_string()],
1066            env: None,
1067            timeout: None,
1068        }))
1069    }
1070
1071    fn configuration(
1072        &self,
1073        _worktree_store: Entity<WorktreeStore>,
1074        _cx: &AsyncApp,
1075    ) -> Task<Result<Option<::extension::ContextServerConfiguration>>> {
1076        Task::ready(Ok(None))
1077    }
1078}