Register and allow to install ocamlmerlin-mlx

src/command_api.ml

@@ -68,6 +68,7 @@ module Internal = struct
   let select_sandbox = unit_handle "select-sandbox"
   let install_ocaml_lsp_server = unit_handle "install-ocaml-lsp-server"
   let upgrade_ocaml_lsp_server = unit_handle "update-ocaml-lsp-server"
+  let install_ocamlmerlin_mlx = unit_handle "install-ocamlmerlin-mlx"
   let restart_language_server = unit_handle "server.restart"
   let select_sandbox_and_open_terminal = unit_handle "open-terminal-select"
   let open_terminal = unit_handle "open-terminal"

src/extension_commands.ml

@@ -90,6 +90,27 @@ let _upgrade_ocaml_lsp_server =
   command Command_api.Internal.upgrade_ocaml_lsp_server callback
 ;;
 
+let _install_ocamlmerlin_mlx =
+  let callback (instance : Extension_instance.t) () =
+    let open Promise.Syntax in
+    let (_ : unit Promise.t) =
+      let sandbox = Extension_instance.sandbox instance in
+      let* ocamlmerlin_mlx_present =
+        Extension_instance.check_ocamlmerlin_mlx_available sandbox
+      in
+      match ocamlmerlin_mlx_present with
+      | Ok () ->
+        show_message `Info "ocamlmerlin-mlx is already installed." |> Promise.return
+      | Error _ ->
+        let* () = Extension_instance.install_ocamlmerlin_mlx sandbox in
+        show_message `Info "Installation of ocamlmerlin-mlx completed successfully.";
+        Extension_instance.start_language_server instance
+    in
+    ()
+  in
+  command Command_api.Internal.install_ocamlmerlin_mlx callback
+;;
+
 let _install_dune_lsp_server =
   let callback (instance : Extension_instance.t) () =
     let open Promise.Syntax in

src/extension_instance.ml

@@ -16,6 +16,7 @@ type t =
   ; mutable standard_hover : bool option
   ; mutable dune_diagnostics : bool option
   ; mutable syntax_documentation : bool option
+  ; mutable prompted_for_ocamlmerlin_mlx : bool
   }
 
 let sandbox t = t.sandbox
@@ -158,6 +159,18 @@ let check_ocaml_lsp_available (sandbox : Sandbox.t) =
             current sandbox.")
 ;;
 
+let check_ocamlmerlin_mlx_available (sandbox : Sandbox.t) =
+  let ocamlmerlin_mlx_version sandbox =
+    Sandbox.get_command sandbox "ocamlmerlin-mlx" [ "--version" ] `Tool
+  in
+  let cwd = Sandbox.workspace_root () in
+  Cmd.output ?cwd (ocamlmerlin_mlx_version sandbox)
+  |> Promise.Result.fold
+       ~ok:(fun (_ : string) -> ())
+       ~error:(fun (_ : string) ->
+         "\"ocamlmerlin-mlx\" is not installed in the current sandbox.")
+;;
+
 module Language_server_init : sig
   val start_language_server : t -> unit Promise.t
 end = struct
@@ -319,6 +332,14 @@ let upgrade_ocaml_lsp_server sandbox =
   ()
 ;;
 
+let install_ocamlmerlin_mlx sandbox =
+  let open Promise.Syntax in
+  let* () = Sandbox.install_packages sandbox [ "ocamlmerlin-mlx" ] in
+  let* () = Command_api.(execute Internal.refresh_switches) () in
+  let+ () = Command_api.(execute Internal.refresh_sandbox) () in
+  ()
+;;
+
 module Sandbox_info : sig
   val make : Sandbox.t -> StatusBarItem.t
   val update : StatusBarItem.t -> new_sandbox:Sandbox.t -> unit
@@ -364,6 +385,7 @@ let make () =
   ; standard_hover = None
   ; dune_diagnostics = None
   ; syntax_documentation = None
+  ; prompted_for_ocamlmerlin_mlx = false
   }
 ;;
 
@@ -483,6 +505,48 @@ let open_terminal sandbox =
 
 let ast_editor_state t = t.ast_editor_state
 
+let suggest_or_install_ocamlmerlin_mlx t =
+  let open Promise.Syntax in
+  let install_mlx_text = "Install ocamlmerlin-mlx" in
+  let select_different_sandbox = "Select a different Sandbox" in
+  let* selection =
+    Window.showInformationMessage
+      ~message:
+        "MLX support requires \"ocamlmerlin-mlx\". Without it, the language server may crash \
+         when opening .mlx files."
+      ~choices:[ install_mlx_text, `Install_mlx; select_different_sandbox, `Select_sandbox ]
+      ()
+  in
+  match selection with
+  | Some `Install_mlx ->
+    let+ () = Command_api.(execute Internal.install_ocamlmerlin_mlx) () in
+    ()
+  | Some `Select_sandbox ->
+    let+ () = Command_api.(execute Internal.select_sandbox) () in
+    ()
+  | _ -> Promise.return ()
+;;
+
+let check_mlx_file_opened t (document : TextDocument.t) =
+  let file_name = TextDocument.fileName document in
+  if String.is_suffix file_name ~suffix:".mlx" && not t.prompted_for_ocamlmerlin_mlx
+  then (
+    t.prompted_for_ocamlmerlin_mlx <- true;
+    let (_ : unit Promise.t) =
+      let open Promise.Syntax in
+      let* ocamlmerlin_mlx_present = check_ocamlmerlin_mlx_available t.sandbox in
+      match ocamlmerlin_mlx_present with
+      | Ok () -> Promise.return ()
+      | Error _ -> suggest_or_install_ocamlmerlin_mlx t
+    in
+    ())
+;;
+
+let register_mlx_check t =
+  let listener document = check_mlx_file_opened t document in
+  Workspace.onDidOpenTextDocument ~listener ()
+;;
+
 let disposable t =
   Disposable.make ~dispose:(fun () ->
     StatusBarItem.dispose t.sandbox_info;

src/extension_instance.mli

@@ -8,6 +8,7 @@ val set_sandbox : t -> Sandbox.t -> unit
 val language_client : t -> LanguageClient.t option
 val ocaml_lsp : t -> Ocaml_lsp.t option
 val check_ocaml_lsp_available : Sandbox.t -> (unit, string) result Promise.t
+val check_ocamlmerlin_mlx_available : Sandbox.t -> (unit, string) result Promise.t
 
 val start_documentation_server
   :  t
@@ -20,6 +21,7 @@ val ocaml_version_exn : t -> Ocaml_version.t
 val start_language_server : t -> unit Promise.t
 val install_ocaml_lsp_server : Sandbox.t -> unit Promise.t
 val upgrade_ocaml_lsp_server : Sandbox.t -> unit Promise.t
+val install_ocamlmerlin_mlx : Sandbox.t -> unit Promise.t
 val suggest_to_run_dune_pkg_lock : unit -> unit
 
 val set_configuration
@@ -34,6 +36,7 @@ val set_configuration
 
 val open_terminal : Sandbox.t -> unit
 val disposable : t -> Disposable.t
+val register_mlx_check : t -> Disposable.t
 val repl : t -> Terminal_sandbox.t option
 val set_repl : t -> Terminal.t -> unit
 val close_repl : t -> unit

src/vscode_ocaml_platform.ml

@@ -41,6 +41,7 @@ let activate (extension : ExtensionContext.t) =
     extension
     ~disposable:(Extension_instance.disposable instance);
   ExtensionContext.subscribe extension ~disposable:(notify_configuration_changes instance);
+  ExtensionContext.subscribe extension ~disposable:(Extension_instance.register_mlx_check instance);
   Dune_formatter.register extension instance;
   Dune_task_provider.register extension instance;
   Treeview_switches.register extension instance;