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