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)).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        let wait_for_context_server_task =
383            wait_for_context_server(&self.context_server_store, id.clone(), cx);
384        cx.spawn({
385            let id = id.clone();
386            async move |this, cx| {
387                let result = wait_for_context_server_task.await;
388                this.update(cx, |this, cx| match result {
389                    Ok(_) => {
390                        this.state = State::Idle;
391                        this.show_configured_context_server_toast(id, cx);
392                        cx.emit(DismissEvent);
393                    }
394                    Err(err) => {
395                        this.set_error(err, cx);
396                    }
397                })
398            }
399        })
400        .detach();
401
402        // When we write the settings to the file, the context server will be restarted.
403        workspace.update(cx, |workspace, cx| {
404            let fs = workspace.app_state().fs.clone();
405            update_settings_file::<ProjectSettings>(fs.clone(), cx, |project_settings, _| {
406                project_settings.context_servers.insert(id.0, settings);
407            });
408        });
409    }
410
411    fn cancel(&mut self, _: &menu::Cancel, cx: &mut Context<Self>) {
412        cx.emit(DismissEvent);
413    }
414
415    fn show_configured_context_server_toast(&self, id: ContextServerId, cx: &mut App) {
416        self.workspace
417            .update(cx, {
418                |workspace, cx| {
419                    let status_toast = StatusToast::new(
420                        format!("{} configured successfully.", id.0),
421                        cx,
422                        |this, _cx| {
423                            this.icon(ToastIcon::new(IconName::Hammer).color(Color::Muted))
424                                .action("Dismiss", |_, _| {})
425                        },
426                    );
427
428                    workspace.toggle_status_toast(status_toast, cx);
429                }
430            })
431            .log_err();
432    }
433}
434
435fn parse_input(text: &str) -> Result<(ContextServerId, ContextServerCommand)> {
436    let value: serde_json::Value = serde_json_lenient::from_str(text)?;
437    let object = value.as_object().context("Expected object")?;
438    anyhow::ensure!(object.len() == 1, "Expected exactly one key-value pair");
439    let (context_server_name, value) = object.into_iter().next().unwrap();
440    let command: ContextServerCommand = serde_json::from_value(value.clone())?;
441    Ok((ContextServerId(context_server_name.clone().into()), command))
442}
443
444impl ModalView for ConfigureContextServerModal {}
445
446impl Focusable for ConfigureContextServerModal {
447    fn focus_handle(&self, cx: &App) -> FocusHandle {
448        match &self.source {
449            ConfigurationSource::New { editor } => editor.focus_handle(cx),
450            ConfigurationSource::Existing { editor, .. } => editor.focus_handle(cx),
451            ConfigurationSource::Extension { editor, .. } => editor
452                .as_ref()
453                .map(|editor| editor.focus_handle(cx))
454                .unwrap_or_else(|| cx.focus_handle()),
455        }
456    }
457}
458
459impl EventEmitter<DismissEvent> for ConfigureContextServerModal {}
460
461impl ConfigureContextServerModal {
462    fn render_modal_header(&self) -> ModalHeader {
463        let text: SharedString = match &self.source {
464            ConfigurationSource::New { .. } => "Add MCP Server".into(),
465            ConfigurationSource::Existing { .. } => "Configure MCP Server".into(),
466            ConfigurationSource::Extension { id, .. } => format!("Configure {}", id.0).into(),
467        };
468        ModalHeader::new().headline(text)
469    }
470
471    fn render_modal_description(&self, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
472        const MODAL_DESCRIPTION: &'static str = "Visit the MCP server configuration docs to find all necessary arguments and environment variables.";
473
474        if let ConfigurationSource::Extension {
475            installation_instructions: Some(installation_instructions),
476            ..
477        } = &self.source
478        {
479            div()
480                .pb_2()
481                .text_sm()
482                .child(MarkdownElement::new(
483                    installation_instructions.clone(),
484                    default_markdown_style(window, cx),
485                ))
486                .into_any_element()
487        } else {
488            Label::new(MODAL_DESCRIPTION)
489                .color(Color::Muted)
490                .into_any_element()
491        }
492    }
493
494    fn render_modal_content(&self, cx: &App) -> AnyElement {
495        let editor = match &self.source {
496            ConfigurationSource::New { editor } => editor,
497            ConfigurationSource::Existing { editor } => editor,
498            ConfigurationSource::Extension { editor, .. } => {
499                let Some(editor) = editor else {
500                    return div().into_any_element();
501                };
502                editor
503            }
504        };
505
506        div()
507            .p_2()
508            .rounded_md()
509            .border_1()
510            .border_color(cx.theme().colors().border_variant)
511            .bg(cx.theme().colors().editor_background)
512            .child({
513                let settings = ThemeSettings::get_global(cx);
514                let text_style = TextStyle {
515                    color: cx.theme().colors().text,
516                    font_family: settings.buffer_font.family.clone(),
517                    font_fallbacks: settings.buffer_font.fallbacks.clone(),
518                    font_size: settings.buffer_font_size(cx).into(),
519                    font_weight: settings.buffer_font.weight,
520                    line_height: relative(settings.buffer_line_height.value()),
521                    ..Default::default()
522                };
523                EditorElement::new(
524                    editor,
525                    EditorStyle {
526                        background: cx.theme().colors().editor_background,
527                        local_player: cx.theme().players().local(),
528                        text: text_style,
529                        syntax: cx.theme().syntax().clone(),
530                        ..Default::default()
531                    },
532                )
533            })
534            .into_any_element()
535    }
536
537    fn render_modal_footer(&self, window: &mut Window, cx: &mut Context<Self>) -> ModalFooter {
538        let focus_handle = self.focus_handle(cx);
539        let is_connecting = matches!(self.state, State::Waiting);
540
541        ModalFooter::new()
542            .start_slot::<Button>(
543                if let ConfigurationSource::Extension {
544                    repository_url: Some(repository_url),
545                    ..
546                } = &self.source
547                {
548                    Some(
549                        Button::new("open-repository", "Open Repository")
550                            .icon(IconName::ArrowUpRight)
551                            .icon_color(Color::Muted)
552                            .icon_size(IconSize::XSmall)
553                            .tooltip({
554                                let repository_url = repository_url.clone();
555                                move |window, cx| {
556                                    Tooltip::with_meta(
557                                        "Open Repository",
558                                        None,
559                                        repository_url.clone(),
560                                        window,
561                                        cx,
562                                    )
563                                }
564                            })
565                            .on_click({
566                                let repository_url = repository_url.clone();
567                                move |_, _, cx| cx.open_url(&repository_url)
568                            }),
569                    )
570                } else {
571                    None
572                },
573            )
574            .end_slot(
575                h_flex()
576                    .gap_2()
577                    .child(
578                        Button::new(
579                            "cancel",
580                            if self.source.has_configuration_options() {
581                                "Cancel"
582                            } else {
583                                "Dismiss"
584                            },
585                        )
586                        .key_binding(
587                            KeyBinding::for_action_in(&menu::Cancel, &focus_handle, window, cx)
588                                .map(|kb| kb.size(rems_from_px(12.))),
589                        )
590                        .on_click(
591                            cx.listener(|this, _event, _window, cx| this.cancel(&menu::Cancel, cx)),
592                        ),
593                    )
594                    .children(self.source.has_configuration_options().then(|| {
595                        Button::new(
596                            "add-server",
597                            if self.source.is_new() {
598                                "Add Server"
599                            } else {
600                                "Configure Server"
601                            },
602                        )
603                        .disabled(is_connecting)
604                        .key_binding(
605                            KeyBinding::for_action_in(&menu::Confirm, &focus_handle, window, cx)
606                                .map(|kb| kb.size(rems_from_px(12.))),
607                        )
608                        .on_click(
609                            cx.listener(|this, _event, _window, cx| {
610                                this.confirm(&menu::Confirm, cx)
611                            }),
612                        )
613                    })),
614            )
615    }
616
617    fn render_waiting_for_context_server() -> Div {
618        h_flex()
619            .gap_2()
620            .child(
621                Icon::new(IconName::ArrowCircle)
622                    .size(IconSize::XSmall)
623                    .color(Color::Info)
624                    .with_animation(
625                        "arrow-circle",
626                        Animation::new(Duration::from_secs(2)).repeat(),
627                        |icon, delta| icon.transform(Transformation::rotate(percentage(delta))),
628                    )
629                    .into_any_element(),
630            )
631            .child(
632                Label::new("Waiting for Context Server")
633                    .size(LabelSize::Small)
634                    .color(Color::Muted),
635            )
636    }
637
638    fn render_modal_error(error: SharedString) -> Div {
639        h_flex()
640            .gap_2()
641            .child(
642                Icon::new(IconName::Warning)
643                    .size(IconSize::XSmall)
644                    .color(Color::Warning),
645            )
646            .child(
647                div()
648                    .w_full()
649                    .child(Label::new(error).size(LabelSize::Small).color(Color::Muted)),
650            )
651    }
652}
653
654impl Render for ConfigureContextServerModal {
655    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
656        div()
657            .elevation_3(cx)
658            .w(rems(34.))
659            .key_context("ConfigureContextServerModal")
660            .on_action(
661                cx.listener(|this, _: &menu::Cancel, _window, cx| this.cancel(&menu::Cancel, cx)),
662            )
663            .on_action(
664                cx.listener(|this, _: &menu::Confirm, _window, cx| {
665                    this.confirm(&menu::Confirm, cx)
666                }),
667            )
668            .capture_any_mouse_down(cx.listener(|this, _, window, cx| {
669                this.focus_handle(cx).focus(window);
670            }))
671            .child(
672                Modal::new("configure-context-server", None)
673                    .header(self.render_modal_header())
674                    .section(
675                        Section::new()
676                            .child(self.render_modal_description(window, cx))
677                            .child(self.render_modal_content(cx))
678                            .child(match &self.state {
679                                State::Idle => div(),
680                                State::Waiting => Self::render_waiting_for_context_server(),
681                                State::Error(error) => Self::render_modal_error(error.clone()),
682                            }),
683                    )
684                    .footer(self.render_modal_footer(window, cx)),
685            )
686    }
687}
688
689fn wait_for_context_server(
690    context_server_store: &Entity<ContextServerStore>,
691    context_server_id: ContextServerId,
692    cx: &mut App,
693) -> Task<Result<(), Arc<str>>> {
694    let (tx, rx) = futures::channel::oneshot::channel();
695    let tx = Arc::new(Mutex::new(Some(tx)));
696
697    let subscription = cx.subscribe(context_server_store, move |_, event, _cx| match event {
698        project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
699            match status {
700                ContextServerStatus::Running => {
701                    if server_id == &context_server_id {
702                        if let Some(tx) = tx.lock().unwrap().take() {
703                            let _ = tx.send(Ok(()));
704                        }
705                    }
706                }
707                ContextServerStatus::Stopped => {
708                    if server_id == &context_server_id {
709                        if let Some(tx) = tx.lock().unwrap().take() {
710                            let _ = tx.send(Err("Context server stopped running".into()));
711                        }
712                    }
713                }
714                ContextServerStatus::Error(error) => {
715                    if server_id == &context_server_id {
716                        if let Some(tx) = tx.lock().unwrap().take() {
717                            let _ = tx.send(Err(error.clone()));
718                        }
719                    }
720                }
721                _ => {}
722            }
723        }
724    });
725
726    cx.spawn(async move |_cx| {
727        let result = rx.await.unwrap();
728        drop(subscription);
729        result
730    })
731}
732
733pub(crate) fn default_markdown_style(window: &Window, cx: &App) -> MarkdownStyle {
734    let theme_settings = ThemeSettings::get_global(cx);
735    let colors = cx.theme().colors();
736    let mut text_style = window.text_style();
737    text_style.refine(&TextStyleRefinement {
738        font_family: Some(theme_settings.ui_font.family.clone()),
739        font_fallbacks: theme_settings.ui_font.fallbacks.clone(),
740        font_features: Some(theme_settings.ui_font.features.clone()),
741        font_size: Some(TextSize::XSmall.rems(cx).into()),
742        color: Some(colors.text_muted),
743        ..Default::default()
744    });
745
746    MarkdownStyle {
747        base_text_style: text_style.clone(),
748        selection_background_color: colors.element_selection_background,
749        link: TextStyleRefinement {
750            background_color: Some(colors.editor_foreground.opacity(0.025)),
751            underline: Some(UnderlineStyle {
752                color: Some(colors.text_accent.opacity(0.5)),
753                thickness: px(1.),
754                ..Default::default()
755            }),
756            ..Default::default()
757        },
758        ..Default::default()
759    }
760}