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}