configure_context_server_modal.rs

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