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