Skip to repository content1933 lines · 63.4 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T05:40:33.634Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
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_context_server_refreshed_when_worktree_added(cx: &mut TestAppContext) {
669 const SERVER_1_ID: &str = "mcp-1";
670
671 let server_1_id = ContextServerId(SERVER_1_ID.into());
672
673 let (fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
674 fs.insert_tree(path!("/second"), json!({"other.rs": ""}))
675 .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 set_context_server_configuration(
689 vec![(
690 server_1_id.0.clone(),
691 settings::ContextServerSettingsContent::Stdio {
692 enabled: true,
693 remote: false,
694 command: ContextServerCommand {
695 path: "somebinary".into(),
696 args: vec!["arg".to_string()],
697 env: None,
698 timeout: None,
699 },
700 },
701 )],
702 cx,
703 );
704
705 {
706 let _server_events = assert_server_events(
707 &store,
708 vec![
709 (server_1_id.clone(), ContextServerStatus::Starting),
710 (server_1_id.clone(), ContextServerStatus::Running),
711 ],
712 cx,
713 );
714 cx.run_until_parked();
715 }
716
717 // Witness that adding a worktree triggers the store to refresh available
718 // servers (via `cx.notify` after `maintain_servers`). Without the
719 // `WorktreeStoreEvent::WorktreeAdded` subscription in `ContextServerStore`,
720 // this counter would remain zero.
721 let notify_count = Rc::new(RefCell::new(0usize));
722 let _notify_subscription = cx.update(|cx| {
723 let count = notify_count.clone();
724 cx.observe(&store, move |_, _| {
725 *count.borrow_mut() += 1;
726 })
727 });
728
729 {
730 let _server_events = assert_server_events(&store, vec![], cx);
731 let _ = project.update(cx, |project, cx| {
732 project.find_or_create_worktree(path!("/second"), true, cx)
733 });
734 cx.run_until_parked();
735 }
736
737 cx.update(|cx| {
738 assert!(
739 *notify_count.borrow() > 0,
740 "Adding a worktree should trigger the context server store to refresh"
741 );
742 assert!(
743 store.read(cx).server_ids().contains(&server_1_id),
744 "Configured server list should still include the server after a worktree is added"
745 );
746 assert_eq!(
747 store.read(cx).status_for_server(&server_1_id),
748 Some(ContextServerStatus::Running),
749 "Server should still be running after a worktree is added"
750 );
751 });
752}
753
754#[gpui::test]
755async fn test_stdio_server_restarts_when_project_root_becomes_available(cx: &mut TestAppContext) {
756 const SERVER_ID: &str = "mcp-1";
757 let server_id = ContextServerId(SERVER_ID.into());
758
759 let fs = FakeFs::new(cx.executor());
760 fs.insert_tree(
761 path!("/"),
762 json!({"lonely.rs": "", "project": {"code.rs": ""}}),
763 )
764 .await;
765
766 cx.update(|cx| {
767 let settings_store = SettingsStore::test(cx);
768 cx.set_global(settings_store);
769 });
770
771 // Open a single-file worktree, whose `root_dir()` is `None`, so the server is
772 // spawned with no working directory even though it is configured.
773 let project = Project::test(fs.clone(), [path!("/lonely.rs").as_ref()], cx).await;
774
775 let executor = cx.executor();
776 let store = project.read_with(cx, |project, _| project.context_server_store());
777 store.update(cx, |store, _| {
778 store.set_context_server_factory(Box::new(move |id, _| {
779 Arc::new(ContextServer::new(
780 id.clone(),
781 Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
782 ))
783 }));
784 });
785
786 // Configure the server globally so it starts against the file worktree.
787 {
788 let _server_events = assert_server_events(
789 &store,
790 vec![
791 (server_id.clone(), ContextServerStatus::Starting),
792 (server_id.clone(), ContextServerStatus::Running),
793 ],
794 cx,
795 );
796 set_context_server_configuration(
797 vec![(
798 server_id.0.clone(),
799 settings::ContextServerSettingsContent::Stdio {
800 enabled: true,
801 remote: false,
802 command: ContextServerCommand {
803 path: "somebinary".into(),
804 args: vec!["arg".to_string()],
805 env: None,
806 timeout: None,
807 },
808 },
809 )],
810 cx,
811 );
812 cx.run_until_parked();
813 }
814
815 // Adding a directory worktree makes the project root resolvable. Since the
816 // server was started with working directory `None`, it must be restarted so
817 // it picks up the new root — otherwise it keeps running under Zed's own cwd.
818 {
819 let _server_events = assert_server_events(
820 &store,
821 vec![
822 (server_id.clone(), ContextServerStatus::Stopped),
823 (server_id.clone(), ContextServerStatus::Starting),
824 (server_id.clone(), ContextServerStatus::Running),
825 ],
826 cx,
827 );
828 project
829 .update(cx, |project, cx| {
830 project.find_or_create_worktree(path!("/project"), true, cx)
831 })
832 .await
833 .expect("Failed to add worktree");
834 cx.run_until_parked();
835 }
836
837 cx.update(|cx| {
838 assert_eq!(
839 store.read(cx).status_for_server(&server_id),
840 Some(ContextServerStatus::Running),
841 "Server should be running again after restarting with the project root"
842 );
843 });
844}
845
846#[gpui::test]
847async fn test_server_ids_includes_disabled_servers(cx: &mut TestAppContext) {
848 const ENABLED_SERVER_ID: &str = "enabled-server";
849 const DISABLED_SERVER_ID: &str = "disabled-server";
850
851 let enabled_server_id = ContextServerId(ENABLED_SERVER_ID.into());
852 let disabled_server_id = ContextServerId(DISABLED_SERVER_ID.into());
853
854 let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
855
856 let executor = cx.executor();
857 let store = project.read_with(cx, |project, _| project.context_server_store());
858 store.update(cx, |store, _| {
859 store.set_context_server_factory(Box::new(move |id, _| {
860 Arc::new(ContextServer::new(
861 id.clone(),
862 Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
863 ))
864 }));
865 });
866
867 // Configure one enabled and one disabled server
868 set_context_server_configuration(
869 vec![
870 (
871 enabled_server_id.0.clone(),
872 settings::ContextServerSettingsContent::Stdio {
873 enabled: true,
874 remote: false,
875 command: ContextServerCommand {
876 path: "somebinary".into(),
877 args: vec![],
878 env: None,
879 timeout: None,
880 },
881 },
882 ),
883 (
884 disabled_server_id.0.clone(),
885 settings::ContextServerSettingsContent::Stdio {
886 enabled: false,
887 remote: false,
888 command: ContextServerCommand {
889 path: "somebinary".into(),
890 args: vec![],
891 env: None,
892 timeout: None,
893 },
894 },
895 ),
896 ],
897 cx,
898 );
899
900 cx.run_until_parked();
901
902 // Verify that server_ids includes both enabled and disabled servers
903 cx.update(|cx| {
904 let server_ids = store.read(cx).server_ids().to_vec();
905 assert!(
906 server_ids.contains(&enabled_server_id),
907 "server_ids should include enabled server"
908 );
909 assert!(
910 server_ids.contains(&disabled_server_id),
911 "server_ids should include disabled server"
912 );
913 });
914
915 // Verify that the enabled server is running and the disabled server is not
916 cx.read(|cx| {
917 assert_eq!(
918 store.read(cx).status_for_server(&enabled_server_id),
919 Some(ContextServerStatus::Running),
920 "enabled server should be running"
921 );
922 // Disabled server should not be in the servers map (status returns None)
923 // but should still be in server_ids
924 assert_eq!(
925 store.read(cx).status_for_server(&disabled_server_id),
926 None,
927 "disabled server should not have a status (not in servers map)"
928 );
929 });
930}
931
932fn set_context_server_configuration(
933 context_servers: Vec<(Arc<str>, settings::ContextServerSettingsContent)>,
934 cx: &mut TestAppContext,
935) {
936 cx.update(|cx| {
937 SettingsStore::update_global(cx, |store, cx| {
938 store.update_user_settings(cx, |content| {
939 content.project.context_servers.clear();
940 for (id, config) in context_servers {
941 content.project.context_servers.insert(id, config);
942 }
943 });
944 })
945 });
946}
947
948#[gpui::test]
949async fn test_remote_context_server(cx: &mut TestAppContext) {
950 const SERVER_ID: &str = "remote-server";
951 let server_id = ContextServerId(SERVER_ID.into());
952 let server_url = "http://example.com/api";
953
954 let client = FakeHttpClient::create(|_| async move {
955 use http_client::AsyncBody;
956
957 let response = Response::builder()
958 .status(200)
959 .header("Content-Type", "application/json")
960 .body(AsyncBody::from(
961 serde_json::to_string(&json!({
962 "jsonrpc": "2.0",
963 "id": 0,
964 "result": {
965 "protocolVersion": "2024-11-05",
966 "capabilities": {},
967 "serverInfo": {
968 "name": "test-server",
969 "version": "1.0.0"
970 }
971 }
972 }))
973 .unwrap(),
974 ))
975 .unwrap();
976 Ok(response)
977 });
978 cx.update(|cx| cx.set_http_client(client));
979
980 let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await;
981
982 let store = project.read_with(cx, |project, _| project.context_server_store());
983
984 set_context_server_configuration(
985 vec![(
986 server_id.0.clone(),
987 settings::ContextServerSettingsContent::Http {
988 enabled: true,
989 url: server_url.to_string(),
990 headers: Default::default(),
991 timeout: None,
992 oauth: None,
993 },
994 )],
995 cx,
996 );
997
998 let _server_events = assert_server_events(
999 &store,
1000 vec![
1001 (server_id.clone(), ContextServerStatus::Starting),
1002 (server_id.clone(), ContextServerStatus::Running),
1003 ],
1004 cx,
1005 );
1006 cx.run_until_parked();
1007}
1008
1009// A server may accept `initialize` unauthenticated (returns 200) yet only send
1010// `WWW-Authenticate` on a later request such as `tools/list` / `tools/call`.
1011// That post-initialize 401 must initiate the OAuth flow instead of surfacing as
1012// an opaque request failure while the server stays "Running".
1013#[gpui::test]
1014async fn test_http_server_authenticates_on_post_init_401(cx: &mut TestAppContext) {
1015 use context_server::transport::TransportError;
1016
1017 const SERVER_ID: &str = "auth-server";
1018 let server_id = ContextServerId(SERVER_ID.into());
1019
1020 set_fake_mcp_http_client(cx, |message| {
1021 if message.contains("\"method\":\"initialize\"") {
1022 Ok(initialize_response())
1023 } else if message.contains("notifications/initialized") {
1024 Ok(notification_accepted_response())
1025 } else {
1026 Ok(unauthorized_response())
1027 }
1028 });
1029
1030 let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await;
1031 let store = project.read_with(cx, |project, _| project.context_server_store());
1032
1033 set_http_context_server_configuration(&server_id, cx);
1034
1035 {
1036 let _server_events = assert_server_events(
1037 &store,
1038 vec![
1039 (server_id.clone(), ContextServerStatus::Starting),
1040 (server_id.clone(), ContextServerStatus::Running),
1041 ],
1042 cx,
1043 );
1044 cx.run_until_parked();
1045 }
1046
1047 let client = store.read_with(cx, |store, _| {
1048 store
1049 .get_running_server(&server_id)
1050 .expect("server should be running")
1051 .client()
1052 .expect("running server should have a client")
1053 });
1054
1055 {
1056 let _server_events = assert_server_events(
1057 &store,
1058 vec![
1059 (server_id.clone(), ContextServerStatus::Starting),
1060 (server_id.clone(), ContextServerStatus::AuthRequired),
1061 ],
1062 cx,
1063 );
1064
1065 // The 401 tears down the client, and the request that carried it fails
1066 // with the typed error.
1067 let error = client
1068 .request::<context_server::types::requests::ListTools>(())
1069 .await
1070 .expect_err("request challenged with a 401 should fail");
1071 assert!(
1072 matches!(
1073 error.downcast_ref::<TransportError>(),
1074 Some(TransportError::AuthRequired { .. })
1075 ),
1076 "expected an AuthRequired error, got: {error}"
1077 );
1078
1079 cx.run_until_parked();
1080 }
1081
1082 cx.update(|cx| {
1083 assert_eq!(
1084 store.read(cx).status_for_server(&server_id),
1085 Some(ContextServerStatus::AuthRequired),
1086 "server should require authentication after a post-initialize 401"
1087 );
1088 });
1089}
1090
1091// The first 401 may even arrive on a notification (`notifications/initialized`
1092// here): the send fails with no request in flight to carry a typed error back,
1093// and the client is dead by the time anything notices. Watching the transport
1094// shutdown must still move the server into `AuthRequired`.
1095#[gpui::test]
1096async fn test_http_server_authenticates_on_notification_401(cx: &mut TestAppContext) {
1097 const SERVER_ID: &str = "auth-server";
1098 let server_id = ContextServerId(SERVER_ID.into());
1099
1100 set_fake_mcp_http_client(cx, |message| {
1101 if message.contains("\"method\":\"initialize\"") {
1102 Ok(initialize_response())
1103 } else {
1104 Ok(unauthorized_response())
1105 }
1106 });
1107
1108 let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await;
1109 let store = project.read_with(cx, |project, _| project.context_server_store());
1110
1111 set_http_context_server_configuration(&server_id, cx);
1112
1113 {
1114 let _server_events = assert_server_events(
1115 &store,
1116 vec![
1117 (server_id.clone(), ContextServerStatus::Starting),
1118 (server_id.clone(), ContextServerStatus::Running),
1119 (server_id.clone(), ContextServerStatus::Starting),
1120 (server_id.clone(), ContextServerStatus::AuthRequired),
1121 ],
1122 cx,
1123 );
1124 cx.run_until_parked();
1125 }
1126
1127 cx.update(|cx| {
1128 assert_eq!(
1129 store.read(cx).status_for_server(&server_id),
1130 Some(ContextServerStatus::AuthRequired),
1131 "server should require authentication after a 401 on a notification"
1132 );
1133 });
1134}
1135
1136// A transport failure that is not an authentication challenge must not touch
1137// the server's state: no spurious auth flow, and (as before the transport
1138// watch existed) the server stays `Running`.
1139#[gpui::test]
1140async fn test_http_server_ignores_non_auth_transport_failure(cx: &mut TestAppContext) {
1141 const SERVER_ID: &str = "flaky-server";
1142 let server_id = ContextServerId(SERVER_ID.into());
1143
1144 set_fake_mcp_http_client(cx, |message| {
1145 if message.contains("\"method\":\"initialize\"") {
1146 Ok(initialize_response())
1147 } else if message.contains("notifications/initialized") {
1148 Ok(notification_accepted_response())
1149 } else {
1150 Err(anyhow::anyhow!("connection reset"))
1151 }
1152 });
1153
1154 let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await;
1155 let store = project.read_with(cx, |project, _| project.context_server_store());
1156
1157 set_http_context_server_configuration(&server_id, cx);
1158
1159 {
1160 let _server_events = assert_server_events(
1161 &store,
1162 vec![
1163 (server_id.clone(), ContextServerStatus::Starting),
1164 (server_id.clone(), ContextServerStatus::Running),
1165 ],
1166 cx,
1167 );
1168 cx.run_until_parked();
1169
1170 let client = store.read_with(cx, |store, _| {
1171 store
1172 .get_running_server(&server_id)
1173 .expect("server should be running")
1174 .client()
1175 .expect("running server should have a client")
1176 });
1177
1178 client
1179 .request::<context_server::types::requests::ListTools>(())
1180 .await
1181 .expect_err("request should fail when the transport errors");
1182
1183 cx.run_until_parked();
1184 // Dropping the events guard asserts no further status change happened.
1185 }
1186
1187 cx.update(|cx| {
1188 assert_eq!(
1189 store.read(cx).status_for_server(&server_id),
1190 Some(ContextServerStatus::Running),
1191 "a non-auth transport failure should not change the server state"
1192 );
1193 });
1194}
1195
1196// A server may also require authentication on `initialize` itself. The
1197// challenge is read from the transport slot rather than the returned error, so
1198// the 401 is recognized even if another error (e.g. the request timeout) wins
1199// the race to become the reported startup failure.
1200#[gpui::test]
1201async fn test_http_server_authenticates_on_initialize_401(cx: &mut TestAppContext) {
1202 const SERVER_ID: &str = "auth-server";
1203 let server_id = ContextServerId(SERVER_ID.into());
1204
1205 set_fake_mcp_http_client(cx, |_message| Ok(unauthorized_response()));
1206
1207 let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await;
1208 let store = project.read_with(cx, |project, _| project.context_server_store());
1209
1210 set_http_context_server_configuration(&server_id, cx);
1211
1212 {
1213 let _server_events = assert_server_events(
1214 &store,
1215 vec![
1216 (server_id.clone(), ContextServerStatus::Starting),
1217 (server_id.clone(), ContextServerStatus::AuthRequired),
1218 ],
1219 cx,
1220 );
1221 cx.run_until_parked();
1222 }
1223
1224 cx.update(|cx| {
1225 assert_eq!(
1226 store.read(cx).status_for_server(&server_id),
1227 Some(ContextServerStatus::AuthRequired),
1228 "server should require authentication after a 401 on initialize"
1229 );
1230 });
1231}
1232
1233// Restarting a server reuses its transport (e.g. via the MCP settings page),
1234// so a challenge recorded by a previous client generation must not leak into
1235// the next one: after a successful restart, a non-auth transport failure must
1236// not trip a spurious auth flow on the stale challenge.
1237#[gpui::test]
1238async fn test_http_server_restart_clears_stale_auth_challenge(cx: &mut TestAppContext) {
1239 use std::sync::atomic::{AtomicBool, Ordering};
1240
1241 const SERVER_ID: &str = "auth-server";
1242 let server_id = ContextServerId(SERVER_ID.into());
1243
1244 let restarted = Arc::new(AtomicBool::new(false));
1245 set_fake_mcp_http_client(cx, {
1246 let restarted = restarted.clone();
1247 move |message| {
1248 if message.contains("\"method\":\"initialize\"") {
1249 Ok(initialize_response())
1250 } else if message.contains("notifications/initialized") {
1251 Ok(notification_accepted_response())
1252 } else if restarted.load(Ordering::SeqCst) {
1253 Err(anyhow::anyhow!("connection reset"))
1254 } else {
1255 Ok(unauthorized_response())
1256 }
1257 }
1258 });
1259
1260 let (_fs, project) = setup_context_server_test(cx, json!({ "code.rs": "" }), vec![]).await;
1261 let store = project.read_with(cx, |project, _| project.context_server_store());
1262
1263 set_http_context_server_configuration(&server_id, cx);
1264 cx.run_until_parked();
1265
1266 // A post-initialize 401 records a challenge on the transport and moves the
1267 // server into AuthRequired.
1268 let client = store.read_with(cx, |store, _| {
1269 store
1270 .get_running_server(&server_id)
1271 .expect("server should be running")
1272 .client()
1273 .expect("running server should have a client")
1274 });
1275 client
1276 .request::<context_server::types::requests::ListTools>(())
1277 .await
1278 .expect_err("request challenged with a 401 should fail");
1279 // Drop our handle so the dead client fully goes away, as it does in
1280 // production once the store has stopped it: a lingering client would
1281 // compete with its successor for the reused transport's response channel.
1282 drop(client);
1283 cx.run_until_parked();
1284 cx.update(|cx| {
1285 assert_eq!(
1286 store.read(cx).status_for_server(&server_id),
1287 Some(ContextServerStatus::AuthRequired),
1288 );
1289 });
1290
1291 // The user restarts the same server instance instead of authenticating,
1292 // and the server no longer challenges.
1293 restarted.store(true, Ordering::SeqCst);
1294 store.update(cx, |store, cx| {
1295 let server = store.get_server(&server_id).expect("server should exist");
1296 store.start_server(server, cx);
1297 });
1298 cx.run_until_parked();
1299 cx.update(|cx| {
1300 assert_eq!(
1301 store.read(cx).status_for_server(&server_id),
1302 Some(ContextServerStatus::Running),
1303 );
1304 });
1305
1306 let client = store.read_with(cx, |store, _| {
1307 store
1308 .get_running_server(&server_id)
1309 .expect("server should be running")
1310 .client()
1311 .expect("running server should have a client")
1312 });
1313 client
1314 .request::<context_server::types::requests::ListTools>(())
1315 .await
1316 .expect_err("request should fail when the transport errors");
1317 cx.run_until_parked();
1318
1319 cx.update(|cx| {
1320 assert_eq!(
1321 store.read(cx).status_for_server(&server_id),
1322 Some(ContextServerStatus::Running),
1323 "a stale challenge from a previous client generation must not trigger auth"
1324 );
1325 });
1326}
1327
1328fn set_http_context_server_configuration(server_id: &ContextServerId, cx: &mut TestAppContext) {
1329 set_context_server_configuration(
1330 vec![(
1331 server_id.0.clone(),
1332 settings::ContextServerSettingsContent::Http {
1333 enabled: true,
1334 url: "https://mcp.example.com/mcp".to_string(),
1335 headers: Default::default(),
1336 timeout: None,
1337 oauth: None,
1338 },
1339 )],
1340 cx,
1341 );
1342}
1343
1344/// A fake HTTP client that serves the OAuth discovery documents (CIMD-capable,
1345/// so no dynamic client registration is needed) and routes MCP endpoint POSTs
1346/// to `respond_to_mcp_message` by the JSON-RPC message in the request body.
1347fn set_fake_mcp_http_client(
1348 cx: &mut TestAppContext,
1349 respond_to_mcp_message: impl Fn(&str) -> Result<Response<http_client::AsyncBody>>
1350 + Send
1351 + Sync
1352 + 'static,
1353) {
1354 let respond_to_mcp_message = Arc::new(respond_to_mcp_message);
1355 let client = FakeHttpClient::create(move |request| {
1356 let respond_to_mcp_message = respond_to_mcp_message.clone();
1357 async move {
1358 let uri = request.uri().to_string();
1359 let discovery_document = if uri.contains("oauth-protected-resource") {
1360 Some(json!({
1361 "resource": "https://mcp.example.com",
1362 "authorization_servers": ["https://auth.example.com"],
1363 "scopes_supported": ["mcp:read"]
1364 }))
1365 } else if uri.contains("oauth-authorization-server") {
1366 Some(json!({
1367 "issuer": "https://auth.example.com",
1368 "authorization_endpoint": "https://auth.example.com/authorize",
1369 "token_endpoint": "https://auth.example.com/token",
1370 "code_challenge_methods_supported": ["S256"],
1371 "client_id_metadata_document_supported": true
1372 }))
1373 } else {
1374 None
1375 };
1376 if let Some(document) = discovery_document {
1377 return Ok(json_response(document));
1378 }
1379
1380 let mut body = request.into_body();
1381 let mut message = String::new();
1382 futures::AsyncReadExt::read_to_string(&mut body, &mut message).await?;
1383 respond_to_mcp_message(&message)
1384 }
1385 });
1386 cx.update(|cx| cx.set_http_client(client));
1387}
1388
1389fn json_response(body: serde_json::Value) -> Response<http_client::AsyncBody> {
1390 Response::builder()
1391 .status(200)
1392 .header("Content-Type", "application/json")
1393 .body(http_client::AsyncBody::from(body.to_string()))
1394 .unwrap()
1395}
1396
1397fn initialize_response() -> Response<http_client::AsyncBody> {
1398 json_response(json!({
1399 "jsonrpc": "2.0",
1400 "id": 0,
1401 "result": {
1402 "protocolVersion": "2024-11-05",
1403 "capabilities": {},
1404 "serverInfo": { "name": "test-server", "version": "1.0.0" }
1405 }
1406 }))
1407}
1408
1409fn notification_accepted_response() -> Response<http_client::AsyncBody> {
1410 Response::builder()
1411 .status(202)
1412 .body(http_client::AsyncBody::empty())
1413 .unwrap()
1414}
1415
1416fn unauthorized_response() -> Response<http_client::AsyncBody> {
1417 Response::builder()
1418 .status(401)
1419 .header("WWW-Authenticate", "Bearer")
1420 .body(http_client::AsyncBody::empty())
1421 .unwrap()
1422}
1423
1424struct ServerEvents {
1425 received_event_count: Rc<RefCell<usize>>,
1426 expected_event_count: usize,
1427 _subscription: Subscription,
1428}
1429
1430impl Drop for ServerEvents {
1431 fn drop(&mut self) {
1432 let actual_event_count = *self.received_event_count.borrow();
1433 assert_eq!(
1434 actual_event_count, self.expected_event_count,
1435 "
1436 Expected to receive {} context server store events, but received {} events",
1437 self.expected_event_count, actual_event_count
1438 );
1439 }
1440}
1441
1442#[gpui::test]
1443async fn test_context_server_global_timeout(cx: &mut TestAppContext) {
1444 cx.update(|cx| {
1445 let settings_store = SettingsStore::test(cx);
1446 cx.set_global(settings_store);
1447 SettingsStore::update_global(cx, |store, cx| {
1448 store
1449 .set_user_settings(r#"{"context_server_timeout": 90}"#, cx)
1450 .expect("Failed to set test user settings");
1451 });
1452 });
1453
1454 let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
1455
1456 let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
1457 let store = cx.new(|cx| {
1458 ContextServerStore::test(
1459 registry.clone(),
1460 project.read(cx).worktree_store(),
1461 Some(project.downgrade()),
1462 cx,
1463 )
1464 });
1465
1466 let mut async_cx = cx.to_async();
1467 let result = ContextServerStore::create_context_server(
1468 store.downgrade(),
1469 ContextServerId("test-server".into()),
1470 Arc::new(ContextServerConfiguration::Http {
1471 url: url::Url::parse("http://localhost:8080").expect("Failed to parse test URL"),
1472 headers: Default::default(),
1473 timeout: None,
1474 oauth: None,
1475 }),
1476 &mut async_cx,
1477 )
1478 .await;
1479
1480 assert!(
1481 result.is_ok(),
1482 "Server should be created successfully with global timeout"
1483 );
1484}
1485
1486#[gpui::test]
1487async fn test_context_server_per_server_timeout_override(cx: &mut TestAppContext) {
1488 const SERVER_ID: &str = "test-server";
1489
1490 cx.update(|cx| {
1491 let settings_store = SettingsStore::test(cx);
1492 cx.set_global(settings_store);
1493 SettingsStore::update_global(cx, |store, cx| {
1494 store
1495 .set_user_settings(r#"{"context_server_timeout": 60}"#, cx)
1496 .expect("Failed to set test user settings");
1497 });
1498 });
1499
1500 let (_fs, project) = setup_context_server_test(
1501 cx,
1502 json!({"code.rs": ""}),
1503 vec![(
1504 SERVER_ID.into(),
1505 ContextServerSettings::Http {
1506 enabled: true,
1507 url: "http://localhost:8080".to_string(),
1508 headers: Default::default(),
1509 timeout: Some(120),
1510 oauth: None,
1511 },
1512 )],
1513 )
1514 .await;
1515
1516 let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
1517 let store = cx.new(|cx| {
1518 ContextServerStore::test(
1519 registry.clone(),
1520 project.read(cx).worktree_store(),
1521 Some(project.downgrade()),
1522 cx,
1523 )
1524 });
1525
1526 let mut async_cx = cx.to_async();
1527 let result = ContextServerStore::create_context_server(
1528 store.downgrade(),
1529 ContextServerId("test-server".into()),
1530 Arc::new(ContextServerConfiguration::Http {
1531 url: url::Url::parse("http://localhost:8080").expect("Failed to parse test URL"),
1532 headers: Default::default(),
1533 timeout: Some(120),
1534 oauth: None,
1535 }),
1536 &mut async_cx,
1537 )
1538 .await;
1539
1540 assert!(
1541 result.is_ok(),
1542 "Server should be created successfully with per-server timeout override"
1543 );
1544}
1545
1546#[gpui::test]
1547async fn test_context_server_stdio_timeout(cx: &mut TestAppContext) {
1548 let (_fs, project) = setup_context_server_test(cx, json!({"code.rs": ""}), vec![]).await;
1549
1550 let registry = cx.new(|_| ContextServerDescriptorRegistry::new());
1551 let store = cx.new(|cx| {
1552 ContextServerStore::test(
1553 registry.clone(),
1554 project.read(cx).worktree_store(),
1555 Some(project.downgrade()),
1556 cx,
1557 )
1558 });
1559
1560 let mut async_cx = cx.to_async();
1561 let result = ContextServerStore::create_context_server(
1562 store.downgrade(),
1563 ContextServerId("stdio-server".into()),
1564 Arc::new(ContextServerConfiguration::Custom {
1565 command: ContextServerCommand {
1566 path: "/usr/bin/node".into(),
1567 args: vec!["server.js".into()],
1568 env: None,
1569 timeout: Some(180000),
1570 },
1571 remote: false,
1572 }),
1573 &mut async_cx,
1574 )
1575 .await;
1576
1577 assert!(
1578 result.is_ok(),
1579 "Stdio server should be created successfully with timeout"
1580 );
1581}
1582
1583#[gpui::test]
1584async fn test_multi_worktree_context_server_settings(cx: &mut TestAppContext) {
1585 const SERVER_A: &str = "server-from-project-a";
1586 const SERVER_B: &str = "server-from-project-b";
1587
1588 let server_a_id = ContextServerId(SERVER_A.into());
1589 let server_b_id = ContextServerId(SERVER_B.into());
1590
1591 let fs = FakeFs::new(cx.executor());
1592 fs.insert_tree(
1593 path!("/project_a"),
1594 json!({
1595 ".zed": {
1596 "settings.json": serde_json::to_string(&json!({
1597 "context_servers": {
1598 "server-from-project-a": {
1599 "command": "server-a-binary",
1600 "args": []
1601 }
1602 }
1603 })).unwrap()
1604 },
1605 "code.rs": ""
1606 }),
1607 )
1608 .await;
1609 fs.insert_tree(
1610 path!("/project_b"),
1611 json!({
1612 ".zed": {
1613 "settings.json": serde_json::to_string(&json!({
1614 "context_servers": {
1615 "server-from-project-b": {
1616 "command": "server-b-binary",
1617 "args": []
1618 }
1619 }
1620 })).unwrap()
1621 },
1622 "code.rs": ""
1623 }),
1624 )
1625 .await;
1626
1627 cx.update(|cx| {
1628 let settings_store = SettingsStore::test(cx);
1629 cx.set_global(settings_store);
1630 });
1631
1632 // Create project with only project_a initially
1633 let project = Project::test(fs.clone(), [path!("/project_a").as_ref()], cx).await;
1634
1635 let executor = cx.executor();
1636 let store = project.read_with(cx, |project, _| project.context_server_store());
1637 store.update(cx, |store, _| {
1638 store.set_context_server_factory(Box::new(move |id, _| {
1639 Arc::new(ContextServer::new(
1640 id.clone(),
1641 Arc::new(create_fake_transport(id.0.to_string(), executor.clone())),
1642 ))
1643 }));
1644 });
1645
1646 cx.run_until_parked();
1647
1648 // Only server-a should be configured
1649 cx.update(|cx| {
1650 let configured = store.read(cx).configured_server_ids();
1651 assert!(
1652 configured.contains(&server_a_id),
1653 "server-a should be configured from project_a"
1654 );
1655 assert!(
1656 !configured.contains(&server_b_id),
1657 "server-b should not be configured yet"
1658 );
1659 });
1660
1661 // Add project_b as a second worktree
1662 project
1663 .update(cx, |project, cx| {
1664 project.find_or_create_worktree(path!("/project_b"), true, cx)
1665 })
1666 .await
1667 .expect("Failed to add second worktree");
1668
1669 cx.run_until_parked();
1670
1671 // Both servers should now be configured
1672 cx.update(|cx| {
1673 let configured = store.read(cx).configured_server_ids();
1674 assert!(
1675 configured.contains(&server_a_id),
1676 "server-a should still be configured from project_a"
1677 );
1678 assert!(
1679 configured.contains(&server_b_id),
1680 "server-b should now be configured from project_b"
1681 );
1682 });
1683}
1684
1685#[gpui::test]
1686async fn test_multi_worktree_duplicate_server_first_wins(cx: &mut TestAppContext) {
1687 const SHARED_SERVER: &str = "shared-server";
1688
1689 let fs = FakeFs::new(cx.executor());
1690 fs.insert_tree(
1691 path!("/project_a"),
1692 json!({
1693 ".zed": {
1694 "settings.json": serde_json::to_string(&json!({
1695 "context_servers": {
1696 "shared-server": {
1697 "command": "binary-from-a",
1698 "args": ["arg-a"]
1699 }
1700 }
1701 })).unwrap()
1702 },
1703 "code.rs": ""
1704 }),
1705 )
1706 .await;
1707 fs.insert_tree(
1708 path!("/project_b"),
1709 json!({
1710 ".zed": {
1711 "settings.json": serde_json::to_string(&json!({
1712 "context_servers": {
1713 "shared-server": {
1714 "command": "binary-from-b",
1715 "args": ["arg-b"]
1716 }
1717 }
1718 })).unwrap()
1719 },
1720 "code.rs": ""
1721 }),
1722 )
1723 .await;
1724
1725 cx.update(|cx| {
1726 let settings_store = SettingsStore::test(cx);
1727 cx.set_global(settings_store);
1728 });
1729
1730 // Create project with both worktrees
1731 let project = Project::test(
1732 fs.clone(),
1733 [path!("/project_a").as_ref(), path!("/project_b").as_ref()],
1734 cx,
1735 )
1736 .await;
1737
1738 cx.run_until_parked();
1739
1740 let store = project.read_with(cx, |project, _| project.context_server_store());
1741
1742 // The server should appear exactly once
1743 cx.update(|cx| {
1744 let configured = store.read(cx).configured_server_ids();
1745 let count = configured
1746 .iter()
1747 .filter(|id| id.0.as_ref() == SHARED_SERVER)
1748 .count();
1749 assert_eq!(count, 1, "duplicate server ID should appear exactly once");
1750 });
1751}
1752
1753#[gpui::test]
1754async fn test_is_server_enabled(cx: &mut TestAppContext) {
1755 // We'll be setting up 4 different servers in order to test the following
1756 // scenarios:
1757 //
1758 // 1. Explicit Settings, Enabled
1759 // 2. Explicit Settings, Disabled
1760 // 3. No Settings, Registry Descriptor, Enabled
1761 // 4. No Settings, No Descriptor, Disabled
1762 const SERVER_1_ID: &str = "mcp-1";
1763 const SERVER_2_ID: &str = "mcp-2";
1764 const SERVER_3_ID: &str = "mcp-3";
1765 const SERVER_4_ID: &str = "mcp-4";
1766
1767 let (_fs, project) = setup_context_server_test(
1768 cx,
1769 json!({"code.rs": ""}),
1770 vec![
1771 (
1772 SERVER_1_ID.into(),
1773 ContextServerSettings::Extension {
1774 enabled: true,
1775 remote: false,
1776 settings: json!({}),
1777 },
1778 ),
1779 (
1780 SERVER_2_ID.into(),
1781 ContextServerSettings::Extension {
1782 enabled: false,
1783 remote: false,
1784 settings: json!({}),
1785 },
1786 ),
1787 ],
1788 )
1789 .await;
1790
1791 let registry = cx.new(|cx| {
1792 let mut registry = ContextServerDescriptorRegistry::new();
1793 let descriptor = Arc::new(FakeContextServerDescriptor::new(SERVER_3_ID));
1794 registry.register_context_server_descriptor(SERVER_3_ID.into(), descriptor, cx);
1795
1796 registry
1797 });
1798
1799 let store = cx.new(|cx| {
1800 ContextServerStore::test(
1801 registry.clone(),
1802 project.read(cx).worktree_store(),
1803 Some(project.downgrade()),
1804 cx,
1805 )
1806 });
1807
1808 // Sanity check before proceeding, confirm server 1 and 2 have settings,
1809 // server 3 is present in the registry while server 4 does not meet any of
1810 // the conditions.
1811 cx.update(|cx| {
1812 let settings = ProjectSettings::get_global(cx);
1813
1814 assert!(settings.context_servers.contains_key(SERVER_1_ID));
1815 assert!(settings.context_servers.contains_key(SERVER_2_ID));
1816 assert!(!settings.context_servers.contains_key(SERVER_3_ID));
1817 assert!(!settings.context_servers.contains_key(SERVER_4_ID));
1818 });
1819
1820 registry.update(cx, |registry, _cx| {
1821 let descriptors = registry.context_server_descriptors();
1822 assert_eq!(descriptors.len(), 1);
1823 assert_eq!(descriptors[0].0.as_ref(), SERVER_3_ID);
1824 });
1825
1826 let server_1_id = ContextServerId(SERVER_1_ID.into());
1827 let server_2_id = ContextServerId(SERVER_2_ID.into());
1828 let server_3_id = ContextServerId(SERVER_3_ID.into());
1829 let server_4_id = ContextServerId(SERVER_4_ID.into());
1830
1831 store.read_with(cx, |store, cx| {
1832 assert!(store.is_server_enabled(&server_1_id, cx));
1833 assert!(!store.is_server_enabled(&server_2_id, cx));
1834 assert!(store.is_server_enabled(&server_3_id, cx));
1835 assert!(!store.is_server_enabled(&server_4_id, cx));
1836 })
1837}
1838
1839fn assert_server_events(
1840 store: &Entity<ContextServerStore>,
1841 expected_events: Vec<(ContextServerId, ContextServerStatus)>,
1842 cx: &mut TestAppContext,
1843) -> ServerEvents {
1844 cx.update(|cx| {
1845 let mut ix = 0;
1846 let received_event_count = Rc::new(RefCell::new(0));
1847 let expected_event_count = expected_events.len();
1848 let subscription = cx.subscribe(store, {
1849 let received_event_count = received_event_count.clone();
1850 move |_, event, _| {
1851 let ServerStatusChangedEvent {
1852 server_id: actual_server_id,
1853 status: actual_status,
1854 } = event;
1855 let (expected_server_id, expected_status) = &expected_events[ix];
1856
1857 assert_eq!(
1858 actual_server_id, expected_server_id,
1859 "Expected different server id at index {}",
1860 ix
1861 );
1862 assert_eq!(
1863 actual_status, expected_status,
1864 "Expected different status at index {}",
1865 ix
1866 );
1867 ix += 1;
1868 *received_event_count.borrow_mut() += 1;
1869 }
1870 });
1871 ServerEvents {
1872 expected_event_count,
1873 received_event_count,
1874 _subscription: subscription,
1875 }
1876 })
1877}
1878
1879async fn setup_context_server_test(
1880 cx: &mut TestAppContext,
1881 files: serde_json::Value,
1882 context_server_configurations: Vec<(Arc<str>, ContextServerSettings)>,
1883) -> (Arc<FakeFs>, Entity<Project>) {
1884 cx.update(|cx| {
1885 let settings_store = SettingsStore::test(cx);
1886 cx.set_global(settings_store);
1887 let mut settings = ProjectSettings::get_global(cx).clone();
1888 for (id, config) in context_server_configurations {
1889 settings.context_servers.insert(id, config);
1890 }
1891 ProjectSettings::override_global(settings, cx);
1892 });
1893
1894 let fs = FakeFs::new(cx.executor());
1895 fs.insert_tree(path!("/test"), files).await;
1896 let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
1897
1898 (fs, project)
1899}
1900
1901struct FakeContextServerDescriptor {
1902 path: PathBuf,
1903}
1904
1905impl FakeContextServerDescriptor {
1906 fn new(path: impl Into<PathBuf>) -> Self {
1907 Self { path: path.into() }
1908 }
1909}
1910
1911impl ContextServerDescriptor for FakeContextServerDescriptor {
1912 fn command(
1913 &self,
1914 _worktree_store: Entity<WorktreeStore>,
1915 _cx: &AsyncApp,
1916 ) -> Task<Result<ContextServerCommand>> {
1917 Task::ready(Ok(ContextServerCommand {
1918 path: self.path.clone(),
1919 args: vec!["arg1".to_string(), "arg2".to_string()],
1920 env: None,
1921 timeout: None,
1922 }))
1923 }
1924
1925 fn configuration(
1926 &self,
1927 _worktree_store: Entity<WorktreeStore>,
1928 _cx: &AsyncApp,
1929 ) -> Task<Result<Option<::extension::ContextServerConfiguration>>> {
1930 Task::ready(Ok(None))
1931 }
1932}
1933