1use std::{
2 sync::{Arc, Mutex},
3 time::Duration,
4};
5
6use anyhow::{Context as _, Result};
7use context_server::{ContextServerCommand, ContextServerId};
8use editor::{Editor, EditorElement, EditorStyle};
9use gpui::{
10 Animation, AnimationExt as _, AsyncWindowContext, DismissEvent, Entity, EventEmitter,
11 FocusHandle, Focusable, Task, TextStyle, TextStyleRefinement, Transformation, UnderlineStyle,
12 WeakEntity, percentage, prelude::*,
13};
14use language::{Language, LanguageRegistry};
15use markdown::{Markdown, MarkdownElement, MarkdownStyle};
16use notifications::status_toast::{StatusToast, ToastIcon};
17use project::{
18 context_server_store::{
19 ContextServerStatus, ContextServerStore, registry::ContextServerDescriptorRegistry,
20 },
21 project_settings::{ContextServerSettings, ProjectSettings},
22 worktree_store::WorktreeStore,
23};
24use settings::{Settings as _, update_settings_file};
25use theme::ThemeSettings;
26use ui::{KeyBinding, Modal, ModalFooter, ModalHeader, Section, Tooltip, prelude::*};
27use util::ResultExt as _;
28use workspace::{ModalView, Workspace};
29
30use crate::AddContextServer;
31
32enum ConfigurationTarget {
33 New,
34 Existing {
35 id: ContextServerId,
36 command: ContextServerCommand,
37 },
38 Extension {
39 id: ContextServerId,
40 repository_url: Option<SharedString>,
41 installation: Option<extension::ContextServerConfiguration>,
42 },
43}
44
45enum ConfigurationSource {
46 New {
47 editor: Entity<Editor>,
48 },
49 Existing {
50 editor: Entity<Editor>,
51 },
52 Extension {
53 id: ContextServerId,
54 editor: Option<Entity<Editor>>,
55 repository_url: Option<SharedString>,
56 installation_instructions: Option<Entity<markdown::Markdown>>,
57 settings_validator: Option<jsonschema::Validator>,
58 },
59}
60
61impl ConfigurationSource {
62 fn has_configuration_options(&self) -> bool {
63 !matches!(self, ConfigurationSource::Extension { editor: None, .. })
64 }
65
66 fn is_new(&self) -> bool {
67 matches!(self, ConfigurationSource::New { .. })
68 }
69
70 fn from_target(
71 target: ConfigurationTarget,
72 language_registry: Arc<LanguageRegistry>,
73 jsonc_language: Option<Arc<Language>>,
74 window: &mut Window,
75 cx: &mut App,
76 ) -> Self {
77 fn create_editor(
78 json: String,
79 jsonc_language: Option<Arc<Language>>,
80 window: &mut Window,
81 cx: &mut App,
82 ) -> Entity<Editor> {
83 cx.new(|cx| {
84 let mut editor = Editor::auto_height(4, 16, window, cx);
85 editor.set_text(json, window, cx);
86 editor.set_show_gutter(false, cx);
87 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
88 if let Some(buffer) = editor.buffer().read(cx).as_singleton() {
89 buffer.update(cx, |buffer, cx| buffer.set_language(jsonc_language, cx))
90 }
91 editor
92 })
93 }
94
95 match target {
96 ConfigurationTarget::New => ConfigurationSource::New {
97 editor: create_editor(context_server_input(None), jsonc_language, window, cx),
98 },
99 ConfigurationTarget::Existing { id, command } => ConfigurationSource::Existing {
100 editor: create_editor(
101 context_server_input(Some((id, command))),
102 jsonc_language,
103 window,
104 cx,
105 ),
106 },
107 ConfigurationTarget::Extension {
108 id,
109 repository_url,
110 installation,
111 } => {
112 let settings_validator = installation.as_ref().and_then(|installation| {
113 jsonschema::validator_for(&installation.settings_schema)
114 .context("Failed to load JSON schema for context server settings")
115 .log_err()
116 });
117 let installation_instructions = installation.as_ref().map(|installation| {
118 cx.new(|cx| {
119 Markdown::new(
120 installation.installation_instructions.clone().into(),
121 Some(language_registry.clone()),
122 None,
123 cx,
124 )
125 })
126 });
127 ConfigurationSource::Extension {
128 id,
129 repository_url,
130 installation_instructions,
131 settings_validator,
132 editor: installation.map(|installation| {
133 create_editor(installation.default_settings, jsonc_language, window, cx)
134 }),
135 }
136 }
137 }
138 }
139
140 fn output(&self, cx: &mut App) -> Result<(ContextServerId, ContextServerSettings)> {
141 match self {
142 ConfigurationSource::New { editor } | ConfigurationSource::Existing { editor } => {
143 parse_input(&editor.read(cx).text(cx)).map(|(id, command)| {
144 (
145 id,
146 ContextServerSettings::Custom {
147 enabled: true,
148 command,
149 },
150 )
151 })
152 }
153 ConfigurationSource::Extension {
154 id,
155 editor,
156 settings_validator,
157 ..
158 } => {
159 let text = editor
160 .as_ref()
161 .context("No output available")?
162 .read(cx)
163 .text(cx);
164 let settings = serde_json_lenient::from_str::<serde_json::Value>(&text)?;
165 if let Some(settings_validator) = settings_validator {
166 if let Err(error) = settings_validator.validate(&settings) {
167 return Err(anyhow::anyhow!(error.to_string()));
168 }
169 }
170 Ok((
171 id.clone(),
172 ContextServerSettings::Extension {
173 enabled: true,
174 settings,
175 },
176 ))
177 }
178 }
179 }
180}
181
182fn context_server_input(existing: Option<(ContextServerId, ContextServerCommand)>) -> String {
183 let (name, command, args, env) = match existing {
184 Some((id, cmd)) => {
185 let args = serde_json::to_string(&cmd.args).unwrap();
186 let env = serde_json::to_string(&cmd.env.unwrap_or_default()).unwrap();
187 (id.0.to_string(), cmd.path, args, env)
188 }
189 None => (
190 "some-mcp-server".to_string(),
191 "".to_string(),
192 "[]".to_string(),
193 "{}".to_string(),
194 ),
195 };
196
197 format!(
198 r#"{{
199 /// The name of your MCP server
200 "{name}": {{
201 /// The command which runs the MCP server
202 "command": "{command}",
203 /// The arguments to pass to the MCP server
204 "args": {args},
205 /// The environment variables to set
206 "env": {env}
207 }}
208}}"#
209 )
210}
211
212fn resolve_context_server_extension(
213 id: ContextServerId,
214 worktree_store: Entity<WorktreeStore>,
215 cx: &mut App,
216) -> Task<Option<ConfigurationTarget>> {
217 let registry = ContextServerDescriptorRegistry::default_global(cx).read(cx);
218
219 let Some(descriptor) = registry.context_server_descriptor(&id.0) else {
220 return Task::ready(None);
221 };
222
223 let extension = crate::agent_configuration::resolve_extension_for_context_server(&id, cx);
224 cx.spawn(async move |cx| {
225 let installation = descriptor
226 .configuration(worktree_store, cx)
227 .await
228 .context("Failed to resolve context server configuration")
229 .log_err()
230 .flatten();
231
232 Some(ConfigurationTarget::Extension {
233 id,
234 repository_url: extension
235 .and_then(|(_, manifest)| manifest.repository.clone().map(SharedString::from)),
236 installation,
237 })
238 })
239}
240
241enum State {
242 Idle,
243 Waiting,
244 Error(SharedString),
245}
246
247pub struct ConfigureContextServerModal {
248 context_server_store: Entity<ContextServerStore>,
249 workspace: WeakEntity<Workspace>,
250 source: ConfigurationSource,
251 state: State,
252}
253
254impl ConfigureContextServerModal {
255 pub fn register(
256 workspace: &mut Workspace,
257 language_registry: Arc<LanguageRegistry>,
258 _window: Option<&mut Window>,
259 _cx: &mut Context<Workspace>,
260 ) {
261 workspace.register_action({
262 let language_registry = language_registry.clone();
263 move |_workspace, _: &AddContextServer, window, cx| {
264 let workspace_handle = cx.weak_entity();
265 let language_registry = language_registry.clone();
266 window
267 .spawn(cx, async move |cx| {
268 Self::show_modal(
269 ConfigurationTarget::New,
270 language_registry,
271 workspace_handle,
272 cx,
273 )
274 .await
275 })
276 .detach_and_log_err(cx);
277 }
278 });
279 }
280
281 pub fn show_modal_for_existing_server(
282 server_id: ContextServerId,
283 language_registry: Arc<LanguageRegistry>,
284 workspace: WeakEntity<Workspace>,
285 window: &mut Window,
286 cx: &mut App,
287 ) -> Task<Result<()>> {
288 let Some(settings) = ProjectSettings::get_global(cx)
289 .context_servers
290 .get(&server_id.0)
291 .cloned()
292 .or_else(|| {
293 ContextServerDescriptorRegistry::default_global(cx)
294 .read(cx)
295 .context_server_descriptor(&server_id.0)
296 .map(|_| ContextServerSettings::default_extension())
297 })
298 else {
299 return Task::ready(Err(anyhow::anyhow!("Context server not found")));
300 };
301
302 window.spawn(cx, async move |cx| {
303 let target = match settings {
304 ContextServerSettings::Custom {
305 enabled: _,
306 command,
307 } => Some(ConfigurationTarget::Existing {
308 id: server_id,
309 command,
310 }),
311 ContextServerSettings::Extension { .. } => {
312 match workspace
313 .update(cx, |workspace, cx| {
314 resolve_context_server_extension(
315 server_id,
316 workspace.project().read(cx).worktree_store(),
317 cx,
318 )
319 })
320 .ok()
321 {
322 Some(task) => task.await,
323 None => None,
324 }
325 }
326 };
327
328 match target {
329 Some(target) => Self::show_modal(target, language_registry, workspace, cx).await,
330 None => Err(anyhow::anyhow!("Failed to resolve context server")),
331 }
332 })
333 }
334
335 fn show_modal(
336 target: ConfigurationTarget,
337 language_registry: Arc<LanguageRegistry>,
338 workspace: WeakEntity<Workspace>,
339 cx: &mut AsyncWindowContext,
340 ) -> Task<Result<()>> {
341 cx.spawn(async move |cx| {
342 let jsonc_language = language_registry.language_for_name("jsonc").await.ok();
343 workspace.update_in(cx, |workspace, window, cx| {
344 let workspace_handle = cx.weak_entity();
345 let context_server_store = workspace.project().read(cx).context_server_store();
346 workspace.toggle_modal(window, cx, |window, cx| Self {
347 context_server_store,
348 workspace: workspace_handle,
349 state: State::Idle,
350 source: ConfigurationSource::from_target(
351 target,
352 language_registry,
353 jsonc_language,
354 window,
355 cx,
356 ),
357 })
358 })
359 })
360 }
361
362 fn set_error(&mut self, err: impl Into<SharedString>, cx: &mut Context<Self>) {
363 self.state = State::Error(err.into());
364 cx.notify();
365 }
366
367 fn confirm(&mut self, _: &menu::Confirm, cx: &mut Context<Self>) {
368 self.state = State::Idle;
369 let Some(workspace) = self.workspace.upgrade() else {
370 return;
371 };
372
373 let (id, settings) = match self.source.output(cx) {
374 Ok(val) => val,
375 Err(error) => {
376 self.set_error(error.to_string(), cx);
377 return;
378 }
379 };
380
381 self.state = State::Waiting;
382
383 let existing_server = self.context_server_store.read(cx).get_running_server(&id);
384 if existing_server.is_some() {
385 self.context_server_store.update(cx, |store, cx| {
386 store.stop_server(&id, cx).log_err();
387 });
388 }
389
390 let wait_for_context_server_task =
391 wait_for_context_server(&self.context_server_store, id.clone(), cx);
392 cx.spawn({
393 let id = id.clone();
394 async move |this, cx| {
395 let result = wait_for_context_server_task.await;
396 this.update(cx, |this, cx| match result {
397 Ok(_) => {
398 this.state = State::Idle;
399 this.show_configured_context_server_toast(id, cx);
400 cx.emit(DismissEvent);
401 }
402 Err(err) => {
403 this.set_error(err, cx);
404 }
405 })
406 }
407 })
408 .detach();
409
410 let settings_changed =
411 ProjectSettings::get_global(cx).context_servers.get(&id.0) != Some(&settings);
412
413 if settings_changed {
414 // When we write the settings to the file, the context server will be restarted.
415 workspace.update(cx, |workspace, cx| {
416 let fs = workspace.app_state().fs.clone();
417 update_settings_file::<ProjectSettings>(fs.clone(), cx, |project_settings, _| {
418 project_settings.context_servers.insert(id.0, settings);
419 });
420 });
421 } else if let Some(existing_server) = existing_server {
422 self.context_server_store
423 .update(cx, |store, cx| store.start_server(existing_server, cx));
424 }
425 }
426
427 fn cancel(&mut self, _: &menu::Cancel, cx: &mut Context<Self>) {
428 cx.emit(DismissEvent);
429 }
430
431 fn show_configured_context_server_toast(&self, id: ContextServerId, cx: &mut App) {
432 self.workspace
433 .update(cx, {
434 |workspace, cx| {
435 let status_toast = StatusToast::new(
436 format!("{} configured successfully.", id.0),
437 cx,
438 |this, _cx| {
439 this.icon(ToastIcon::new(IconName::Hammer).color(Color::Muted))
440 .action("Dismiss", |_, _| {})
441 },
442 );
443
444 workspace.toggle_status_toast(status_toast, cx);
445 }
446 })
447 .log_err();
448 }
449}
450
451fn parse_input(text: &str) -> Result<(ContextServerId, ContextServerCommand)> {
452 let value: serde_json::Value = serde_json_lenient::from_str(text)?;
453 let object = value.as_object().context("Expected object")?;
454 anyhow::ensure!(object.len() == 1, "Expected exactly one key-value pair");
455 let (context_server_name, value) = object.into_iter().next().unwrap();
456 let command: ContextServerCommand = serde_json::from_value(value.clone())?;
457 Ok((ContextServerId(context_server_name.clone().into()), command))
458}
459
460impl ModalView for ConfigureContextServerModal {}
461
462impl Focusable for ConfigureContextServerModal {
463 fn focus_handle(&self, cx: &App) -> FocusHandle {
464 match &self.source {
465 ConfigurationSource::New { editor } => editor.focus_handle(cx),
466 ConfigurationSource::Existing { editor, .. } => editor.focus_handle(cx),
467 ConfigurationSource::Extension { editor, .. } => editor
468 .as_ref()
469 .map(|editor| editor.focus_handle(cx))
470 .unwrap_or_else(|| cx.focus_handle()),
471 }
472 }
473}
474
475impl EventEmitter<DismissEvent> for ConfigureContextServerModal {}
476
477impl ConfigureContextServerModal {
478 fn render_modal_header(&self) -> ModalHeader {
479 let text: SharedString = match &self.source {
480 ConfigurationSource::New { .. } => "Add MCP Server".into(),
481 ConfigurationSource::Existing { .. } => "Configure MCP Server".into(),
482 ConfigurationSource::Extension { id, .. } => format!("Configure {}", id.0).into(),
483 };
484 ModalHeader::new().headline(text)
485 }
486
487 fn render_modal_description(&self, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
488 const MODAL_DESCRIPTION: &'static str = "Visit the MCP server configuration docs to find all necessary arguments and environment variables.";
489
490 if let ConfigurationSource::Extension {
491 installation_instructions: Some(installation_instructions),
492 ..
493 } = &self.source
494 {
495 div()
496 .pb_2()
497 .text_sm()
498 .child(MarkdownElement::new(
499 installation_instructions.clone(),
500 default_markdown_style(window, cx),
501 ))
502 .into_any_element()
503 } else {
504 Label::new(MODAL_DESCRIPTION)
505 .color(Color::Muted)
506 .into_any_element()
507 }
508 }
509
510 fn render_modal_content(&self, cx: &App) -> AnyElement {
511 let editor = match &self.source {
512 ConfigurationSource::New { editor } => editor,
513 ConfigurationSource::Existing { editor } => editor,
514 ConfigurationSource::Extension { editor, .. } => {
515 let Some(editor) = editor else {
516 return div().into_any_element();
517 };
518 editor
519 }
520 };
521
522 div()
523 .p_2()
524 .rounded_md()
525 .border_1()
526 .border_color(cx.theme().colors().border_variant)
527 .bg(cx.theme().colors().editor_background)
528 .child({
529 let settings = ThemeSettings::get_global(cx);
530 let text_style = TextStyle {
531 color: cx.theme().colors().text,
532 font_family: settings.buffer_font.family.clone(),
533 font_fallbacks: settings.buffer_font.fallbacks.clone(),
534 font_size: settings.buffer_font_size(cx).into(),
535 font_weight: settings.buffer_font.weight,
536 line_height: relative(settings.buffer_line_height.value()),
537 ..Default::default()
538 };
539 EditorElement::new(
540 editor,
541 EditorStyle {
542 background: cx.theme().colors().editor_background,
543 local_player: cx.theme().players().local(),
544 text: text_style,
545 syntax: cx.theme().syntax().clone(),
546 ..Default::default()
547 },
548 )
549 })
550 .into_any_element()
551 }
552
553 fn render_modal_footer(&self, window: &mut Window, cx: &mut Context<Self>) -> ModalFooter {
554 let focus_handle = self.focus_handle(cx);
555 let is_connecting = matches!(self.state, State::Waiting);
556
557 ModalFooter::new()
558 .start_slot::<Button>(
559 if let ConfigurationSource::Extension {
560 repository_url: Some(repository_url),
561 ..
562 } = &self.source
563 {
564 Some(
565 Button::new("open-repository", "Open Repository")
566 .icon(IconName::ArrowUpRight)
567 .icon_color(Color::Muted)
568 .icon_size(IconSize::XSmall)
569 .tooltip({
570 let repository_url = repository_url.clone();
571 move |window, cx| {
572 Tooltip::with_meta(
573 "Open Repository",
574 None,
575 repository_url.clone(),
576 window,
577 cx,
578 )
579 }
580 })
581 .on_click({
582 let repository_url = repository_url.clone();
583 move |_, _, cx| cx.open_url(&repository_url)
584 }),
585 )
586 } else {
587 None
588 },
589 )
590 .end_slot(
591 h_flex()
592 .gap_2()
593 .child(
594 Button::new(
595 "cancel",
596 if self.source.has_configuration_options() {
597 "Cancel"
598 } else {
599 "Dismiss"
600 },
601 )
602 .key_binding(
603 KeyBinding::for_action_in(&menu::Cancel, &focus_handle, window, cx)
604 .map(|kb| kb.size(rems_from_px(12.))),
605 )
606 .on_click(
607 cx.listener(|this, _event, _window, cx| this.cancel(&menu::Cancel, cx)),
608 ),
609 )
610 .children(self.source.has_configuration_options().then(|| {
611 Button::new(
612 "add-server",
613 if self.source.is_new() {
614 "Add Server"
615 } else {
616 "Configure Server"
617 },
618 )
619 .disabled(is_connecting)
620 .key_binding(
621 KeyBinding::for_action_in(&menu::Confirm, &focus_handle, window, cx)
622 .map(|kb| kb.size(rems_from_px(12.))),
623 )
624 .on_click(
625 cx.listener(|this, _event, _window, cx| {
626 this.confirm(&menu::Confirm, cx)
627 }),
628 )
629 })),
630 )
631 }
632
633 fn render_waiting_for_context_server() -> Div {
634 h_flex()
635 .gap_2()
636 .child(
637 Icon::new(IconName::ArrowCircle)
638 .size(IconSize::XSmall)
639 .color(Color::Info)
640 .with_animation(
641 "arrow-circle",
642 Animation::new(Duration::from_secs(2)).repeat(),
643 |icon, delta| icon.transform(Transformation::rotate(percentage(delta))),
644 )
645 .into_any_element(),
646 )
647 .child(
648 Label::new("Waiting for Context Server")
649 .size(LabelSize::Small)
650 .color(Color::Muted),
651 )
652 }
653
654 fn render_modal_error(error: SharedString) -> Div {
655 h_flex()
656 .gap_2()
657 .child(
658 Icon::new(IconName::Warning)
659 .size(IconSize::XSmall)
660 .color(Color::Warning),
661 )
662 .child(
663 div()
664 .w_full()
665 .child(Label::new(error).size(LabelSize::Small).color(Color::Muted)),
666 )
667 }
668}
669
670impl Render for ConfigureContextServerModal {
671 fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
672 div()
673 .elevation_3(cx)
674 .w(rems(34.))
675 .key_context("ConfigureContextServerModal")
676 .on_action(
677 cx.listener(|this, _: &menu::Cancel, _window, cx| this.cancel(&menu::Cancel, cx)),
678 )
679 .on_action(
680 cx.listener(|this, _: &menu::Confirm, _window, cx| {
681 this.confirm(&menu::Confirm, cx)
682 }),
683 )
684 .capture_any_mouse_down(cx.listener(|this, _, window, cx| {
685 this.focus_handle(cx).focus(window);
686 }))
687 .child(
688 Modal::new("configure-context-server", None)
689 .header(self.render_modal_header())
690 .section(
691 Section::new()
692 .child(self.render_modal_description(window, cx))
693 .child(self.render_modal_content(cx))
694 .child(match &self.state {
695 State::Idle => div(),
696 State::Waiting => Self::render_waiting_for_context_server(),
697 State::Error(error) => Self::render_modal_error(error.clone()),
698 }),
699 )
700 .footer(self.render_modal_footer(window, cx)),
701 )
702 }
703}
704
705fn wait_for_context_server(
706 context_server_store: &Entity<ContextServerStore>,
707 context_server_id: ContextServerId,
708 cx: &mut App,
709) -> Task<Result<(), Arc<str>>> {
710 let (tx, rx) = futures::channel::oneshot::channel();
711 let tx = Arc::new(Mutex::new(Some(tx)));
712
713 let subscription = cx.subscribe(context_server_store, move |_, event, _cx| match event {
714 project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
715 match status {
716 ContextServerStatus::Running => {
717 if server_id == &context_server_id {
718 if let Some(tx) = tx.lock().unwrap().take() {
719 let _ = tx.send(Ok(()));
720 }
721 }
722 }
723 ContextServerStatus::Stopped => {
724 if server_id == &context_server_id {
725 if let Some(tx) = tx.lock().unwrap().take() {
726 let _ = tx.send(Err("Context server stopped running".into()));
727 }
728 }
729 }
730 ContextServerStatus::Error(error) => {
731 if server_id == &context_server_id {
732 if let Some(tx) = tx.lock().unwrap().take() {
733 let _ = tx.send(Err(error.clone()));
734 }
735 }
736 }
737 _ => {}
738 }
739 }
740 });
741
742 cx.spawn(async move |_cx| {
743 let result = rx
744 .await
745 .map_err(|_| Arc::from("Context server store was dropped"))?;
746 drop(subscription);
747 result
748 })
749}
750
751pub(crate) fn default_markdown_style(window: &Window, cx: &App) -> MarkdownStyle {
752 let theme_settings = ThemeSettings::get_global(cx);
753 let colors = cx.theme().colors();
754 let mut text_style = window.text_style();
755 text_style.refine(&TextStyleRefinement {
756 font_family: Some(theme_settings.ui_font.family.clone()),
757 font_fallbacks: theme_settings.ui_font.fallbacks.clone(),
758 font_features: Some(theme_settings.ui_font.features.clone()),
759 font_size: Some(TextSize::XSmall.rems(cx).into()),
760 color: Some(colors.text_muted),
761 ..Default::default()
762 });
763
764 MarkdownStyle {
765 base_text_style: text_style.clone(),
766 selection_background_color: colors.element_selection_background,
767 link: TextStyleRefinement {
768 background_color: Some(colors.editor_foreground.opacity(0.025)),
769 underline: Some(UnderlineStyle {
770 color: Some(colors.text_accent.opacity(0.5)),
771 thickness: px(1.),
772 ..Default::default()
773 }),
774 ..Default::default()
775 },
776 ..Default::default()
777 }
778}