Add support for mix run flags to the Mix runtime (#1108)

* Add support for mix run flags to the Mix runtime

* Support mix run flags in runtime config string

* Add flag validation

* Update lib/livebook_web/live/session_live/mix_standalone_live.ex

Co-authored-by: José Valim <jose.valim@dashbit.co>

Co-authored-by: José Valim <jose.valim@dashbit.co>

README.md

@@ -153,7 +153,7 @@ The following environment variables configure Livebook:
 
   * LIVEBOOK_DEFAULT_RUNTIME - sets the runtime type that is used by default
     when none is started explicitly for the given notebook. Must be either
-    "standalone" (Elixir standalone), "mix[:PATH]" (Mix standalone),
+    "standalone" (Elixir standalone), "mix[:PATH][:FLAGS]" (Mix standalone),
     "attached:NODE:COOKIE" (Attached node) or "embedded" (Embedded).
     Defaults to "standalone".
 

lib/livebook/config.ex

@@ -244,10 +244,16 @@ defmodule Livebook.Config do
             )
         end
 
-      "mix:" <> path ->
+      "mix:" <> config ->
+        {path, flags} = parse_mix_config!(config)
+
         case mix_path(path) do
           {:ok, path} ->
-            Livebook.Runtime.MixStandalone.new(path)
+            if Livebook.Utils.valid_cli_flags?(flags) do
+              Livebook.Runtime.MixStandalone.new(path, flags)
+            else
+              abort!(~s{"#{flags}" is not a valid flag sequence})
+            end
 
           :error ->
             abort!(~s{"#{path}" does not point to a Mix project})
@@ -264,6 +270,13 @@ defmodule Livebook.Config do
     end
   end
 
+  defp parse_mix_config!(config) do
+    case String.split(config, ":", parts: 2) do
+      [path] -> {path, ""}
+      [path, flags] -> {path, flags}
+    end
+  end
+
   defp mix_path(path) do
     path = Path.expand(path)
     mixfile = Path.join(path, "mix.exs")

lib/livebook/runtime/mix_standalone.ex

@@ -1,5 +1,5 @@
 defmodule Livebook.Runtime.MixStandalone do
-  defstruct [:node, :server_pid, :project_path]
+  defstruct [:node, :server_pid, :project_path, :flags]
 
   # A runtime backed by a standalone Elixir node managed by Livebook.
   #
@@ -13,6 +13,7 @@ defmodule Livebook.Runtime.MixStandalone do
 
   @type t :: %__MODULE__{
           project_path: String.t(),
+          flags: String.t(),
           node: node() | nil,
           server_pid: pid() | nil
         }
@@ -20,9 +21,9 @@ defmodule Livebook.Runtime.MixStandalone do
   @doc """
   Returns a new runtime instance.
   """
-  @spec new(String.t()) :: t()
-  def new(project_path) do
-    %__MODULE__{project_path: project_path}
+  @spec new(String.t(), String.t()) :: t()
+  def new(project_path, flags \\ "") do
+    %__MODULE__{project_path: project_path, flags: flags}
   end
 
   @doc """
@@ -46,6 +47,7 @@ defmodule Livebook.Runtime.MixStandalone do
   @spec connect_async(t(), Emitter.t()) :: :ok
   def connect_async(runtime, emitter) do
     %{project_path: project_path} = runtime
+    flags = OptionParser.split(runtime.flags)
     output_emitter = Emitter.mapper(emitter, fn output -> {:output, output} end)
 
     spawn_link(fn ->
@@ -59,7 +61,8 @@ defmodule Livebook.Runtime.MixStandalone do
              :ok <- run_mix_task("deps.get", project_path, output_emitter),
              :ok <- run_mix_task("compile", project_path, output_emitter),
              eval = child_node_eval_string(),
-             port = start_elixir_mix_node(elixir_path, child_node, eval, argv, project_path),
+             port =
+               start_elixir_mix_node(elixir_path, child_node, flags, eval, argv, project_path),
              {:ok, server_pid} <- parent_init_sequence(child_node, port, emitter: output_emitter) do
           runtime = %{runtime | node: child_node, server_pid: server_pid}
           Emitter.emit(emitter, {:ok, runtime})
@@ -115,7 +118,7 @@ defmodule Livebook.Runtime.MixStandalone do
     end
   end
 
-  defp start_elixir_mix_node(elixir_path, node_name, eval, argv, project_path) do
+  defp start_elixir_mix_node(elixir_path, node_name, flags, eval, argv, project_path) do
     # Here we create a port to start the system process in a non-blocking way.
     Port.open({:spawn_executable, elixir_path}, [
       :binary,
@@ -127,7 +130,8 @@ defmodule Livebook.Runtime.MixStandalone do
       cd: project_path,
       args:
         elixir_flags(node_name) ++
-          ["-S", "mix", "run", "--eval", eval, "--" | Enum.map(argv, &to_string/1)]
+          ["-S", "mix", "run", "--eval", eval | flags] ++
+          ["--" | Enum.map(argv, &to_string/1)]
     ])
   end
 end
@@ -138,13 +142,11 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.MixStandalone do
   def describe(runtime) do
     [
       {"Type", "Mix standalone"},
-      {"Project", runtime.project_path}
-    ] ++
-      if connected?(runtime) do
-        [{"Node name", Atom.to_string(runtime.node)}]
-      else
-        []
-      end
+      {"Project", runtime.project_path},
+      runtime.flags != "" && {"Flags", runtime.flags},
+      connected?(runtime) && {"Node name", Atom.to_string(runtime.node)}
+    ]
+    |> Enum.filter(&is_tuple/1)
   end
 
   def connect(runtime) do
@@ -166,7 +168,7 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.MixStandalone do
   end
 
   def duplicate(runtime) do
-    Livebook.Runtime.MixStandalone.new(runtime.project_path)
+    Livebook.Runtime.MixStandalone.new(runtime.project_path, runtime.flags)
   end
 
   def evaluate_code(runtime, code, locator, base_locator, opts \\ []) do

lib/livebook/runtime/standalone_init.ex

@@ -110,7 +110,7 @@ defmodule Livebook.Runtime.StandaloneInit do
           loop.(loop)
 
         {:DOWN, ^port_ref, :port, _object, _reason} ->
-          {:error, "Elixir process terminated unexpectedly"}
+          {:error, "Elixir terminated unexpectedly, please check the terminal for errors"}
       after
         # Use a longer timeout to account for longer child node startup,
         # as may happen when starting with Mix.

lib/livebook/utils.ex

@@ -193,6 +193,30 @@ defmodule Livebook.Utils do
   @spec valid_hex_color?(String.t()) :: boolean()
   def valid_hex_color?(hex_color), do: hex_color =~ ~r/^#[0-9a-fA-F]{6}$/
 
+  @doc ~S"""
+  Validates if the given string forms valid CLI flags.
+
+  ## Examples
+
+      iex> Livebook.Utils.valid_cli_flags?("")
+      true
+
+      iex> Livebook.Utils.valid_cli_flags?("--arg1 value --arg2 'value'")
+      true
+
+      iex> Livebook.Utils.valid_cli_flags?("--arg1 \"")
+      false
+  """
+  @spec valid_cli_flags?(String.t()) :: boolean()
+  def valid_cli_flags?(flags) do
+    try do
+      OptionParser.split(flags)
+      true
+    rescue
+      _ -> false
+    end
+  end
+
   @doc """
   Changes the first letter in the given string to upper case.
 

lib/livebook_cli/server.ex

@@ -45,7 +45,7 @@ defmodule LivebookCLI.Server do
                            explicitly for the given notebook, defaults to standalone
                            Supported options:
                              * standalone - Elixir standalone
-                             * mix[:PATH] - Mix standalone
+                             * mix[:PATH][:FLAGS] - Mix standalone
                              * attached:NODE:COOKIE - Attached
                              * embedded - Embedded
       --home               The home path for the Livebook instance

lib/livebook_web/live/session_live/mix_standalone_live.ex

@@ -20,8 +20,9 @@ defmodule LivebookWeb.SessionLive.MixStandaloneLive do
        session: session,
        status: :initial,
        current_runtime: current_runtime,
-       file: initial_file(current_runtime),
+       data: initial_data(current_runtime),
        outputs: [],
+       outputs_version: 0,
        emitter: nil
      ), temporary_assigns: [outputs: []]}
   end
@@ -41,27 +42,40 @@ defmodule LivebookWeb.SessionLive.MixStandaloneLive do
         inside a Mix project because the dependencies of the project
         itself have been installed instead.
       </p>
-      <%= if @status == :initial do %>
+      <%= if @status != :initializing do %>
         <div class="h-full h-52">
           <.live_component module={LivebookWeb.FileSelectComponent}
               id="mix-project-dir"
-              file={@file}
+              file={@data.file}
               extnames={[]}
               running_files={[]}
-              submit_event={if(disabled?(@file.path), do: nil, else: :init)}
+              submit_event={if(data_valid?(@data), do: :init, else: nil)}
               file_system_select_disabled={true} />
         </div>
-        <button class="button-base button-blue" phx-click="init" disabled={disabled?(@file.path)}>
-          <%= if(matching_runtime?(@current_runtime, @file.path), do: "Reconnect", else: "Connect") %>
-        </button>
+        <form phx-change="validate" phx-submit="init">
+          <div>
+            <div class="input-label"><code>mix run</code> command-line flags</div>
+            <input class="input"
+              type="text"
+              name="flags"
+              value={@data.flags}
+              spellcheck="false"
+              autocomplete="off" />
+          </div>
+          <button class="mt-5 button-base button-blue"
+            type="submit"
+            disabled={not data_valid?(@data)}>
+            <%= if(matching_runtime?(@current_runtime, @data), do: "Reconnect", else: "Connect") %>
+          </button>
+        </form>
       <% end %>
       <%= if @status != :initial do %>
         <div class="markdown">
           <pre><code class="max-h-40 overflow-y-auto tiny-scrollbar"
-            id="mix-standalone-init-output"
+            id={"mix-standalone-init-output-#{@outputs_version}"}
             phx-update="append"
             phx-hook="ScrollOnUpdate"
-            ><%= for {output, i} <- @outputs do %><span id={"output-#{i}"}><%= ansi_string_to_html(output) %></span><% end %></code></pre>
+            ><%= for {output, i} <- @outputs do %><span id={"mix-standalone-init-output-#{@outputs_version}-#{i}"}><%= ansi_string_to_html(output) %></span><% end %></code></pre>
         </div>
       <% end %>
     </div>
@@ -73,9 +87,13 @@ defmodule LivebookWeb.SessionLive.MixStandaloneLive do
     handle_init(socket)
   end
 
+  def handle_event("validate", %{"flags" => flags}, socket) do
+    {:noreply, update(socket, :data, &%{&1 | flags: flags})}
+  end
+
   @impl true
   def handle_info({:set_file, file, _info}, socket) do
-    {:noreply, assign(socket, :file, file)}
+    {:noreply, update(socket, :data, &%{&1 | file: file})}
   end
 
   def handle_info(:init, socket) do
@@ -104,31 +122,40 @@ defmodule LivebookWeb.SessionLive.MixStandaloneLive do
 
   defp handle_init(socket) do
     emitter = Utils.Emitter.new(self())
-    runtime = Runtime.MixStandalone.new(socket.assigns.file.path)
+    runtime = Runtime.MixStandalone.new(socket.assigns.data.file.path, socket.assigns.data.flags)
     Runtime.MixStandalone.connect_async(runtime, emitter)
-    {:noreply, assign(socket, status: :initializing, emitter: emitter)}
+
+    {:noreply,
+     socket
+     |> assign(status: :initializing, emitter: emitter, outputs: [])
+     |> update(:outputs_version, &(&1 + 1))}
   end
 
   defp add_output(socket, output) do
     assign(socket, outputs: socket.assigns.outputs ++ [{output, Utils.random_id()}])
   end
 
-  defp initial_file(%Runtime.MixStandalone{} = current_runtime) do
-    FileSystem.File.local(current_runtime.project_path)
+  defp initial_data(%Runtime.MixStandalone{project_path: project_path, flags: flags}) do
+    file =
+      project_path
+      |> FileSystem.Utils.ensure_dir_path()
+      |> FileSystem.File.local()
+
+    %{file: file, flags: flags}
   end
 
-  defp initial_file(_runtime) do
-    Livebook.Config.local_filesystem_home()
+  defp initial_data(_runtime) do
+    %{file: Livebook.Config.local_filesystem_home(), flags: ""}
   end
 
-  defp matching_runtime?(%Runtime.MixStandalone{} = runtime, path) do
-    Path.expand(runtime.project_path) == Path.expand(path)
+  defp matching_runtime?(%Runtime.MixStandalone{} = runtime, data) do
+    Path.expand(runtime.project_path) == Path.expand(data.file.path)
   end
 
   defp matching_runtime?(_runtime, _path), do: false
 
-  defp disabled?(path) do
-    not mix_project_root?(path)
+  defp data_valid?(data) do
+    mix_project_root?(data.file.path) and Livebook.Utils.valid_cli_flags?(data.flags)
   end
 
   defp mix_project_root?(path) do