Update export flow for JS outputs

lib/livebook/live_markdown/export.ex

@@ -7,9 +7,9 @@ defmodule Livebook.LiveMarkdown.Export do
     include_outputs? = Keyword.get(opts, :include_outputs, notebook.persist_outputs)
     include_stamp? = Keyword.get(opts, :include_stamp, true)
 
-    js_ref_with_data = if include_outputs?, do: collect_js_output_data(notebook), else: %{}
+    js_ref_with_export = if include_outputs?, do: collect_js_output_export(notebook), else: %{}
 
-    ctx = %{include_outputs?: include_outputs?, js_ref_with_data: js_ref_with_data}
+    ctx = %{include_outputs?: include_outputs?, js_ref_with_export: js_ref_with_export}
 
     iodata = render_notebook(notebook, ctx)
 
@@ -24,26 +24,30 @@ defmodule Livebook.LiveMarkdown.Export do
     {source, footer_warnings}
   end
 
-  defp collect_js_output_data(notebook) do
-    for section <- notebook.sections,
-        %{outputs: outputs} <- section.cells,
-        {_idx, %{type: :js, js_view: %{ref: ref, pid: pid}, export: %{}}} <- outputs do
-      Task.async(fn ->
-        {ref, get_js_output_data(pid, ref)}
-      end)
-    end
+  defp collect_js_output_export(notebook) do
+    for(
+      section <- notebook.sections,
+      %{outputs: outputs} <- section.cells,
+      {_idx, %{type: :js, js_view: js_view, export: export}} <- outputs,
+      export == true or is_map(export),
+      do: {js_view.ref, js_view.pid, export},
+      uniq: true
+    )
+    |> Enum.map(fn {ref, pid, export} ->
+      Task.async(fn -> {ref, get_js_output_export(pid, ref, export)} end)
+    end)
     |> Task.await_many(:infinity)
     |> Map.new()
   end
 
-  defp get_js_output_data(pid, ref) do
-    send(pid, {:connect, self(), %{origin: self(), ref: ref}})
+  defp get_js_output_export(pid, ref, true) do
+    send(pid, {:export, self(), %{ref: ref}})
 
     monitor_ref = Process.monitor(pid)
 
     data =
       receive do
-        {:connect_reply, data, %{ref: ^ref}} -> data
+        {:export_reply, export_result, %{ref: ^ref}} -> export_result
         {:DOWN, ^monitor_ref, :process, _pid, _reason} -> nil
       end
 
@@ -52,6 +56,26 @@ defmodule Livebook.LiveMarkdown.Export do
     data
   end
 
+  # Handle old flow for backward compatibility with Kino <= 0.10.0
+  defp get_js_output_export(pid, ref, %{info_string: info_string, key: key}) do
+    send(pid, {:connect, self(), %{origin: to_string(self()), ref: ref}})
+
+    monitor_ref = Process.monitor(pid)
+
+    data =
+      receive do
+        {:connect_reply, data, %{ref: ^ref}} -> data
+        {:DOWN, ^monitor_ref, :process, _pid, _reason} -> nil
+      end
+
+    Process.demonitor(monitor_ref, [:flush])
+
+    if data do
+      payload = if key && is_map(data), do: data[key], else: data
+      {info_string, payload}
+    end
+  end
+
   defp render_notebook(notebook, ctx) do
     %{setup_section: %{cells: [setup_cell]}} = notebook
 
@@ -221,23 +245,15 @@ defmodule Livebook.LiveMarkdown.Export do
     |> prepend_metadata(%{output: true})
   end
 
-  defp render_output(
-         %{type: :js, export: %{info_string: info_string, key: key}, js_view: %{ref: ref}},
-         ctx
-       )
-       when is_binary(info_string) do
-    data = ctx.js_ref_with_data[ref]
-    payload = if key && is_map(data), do: data[key], else: data
-
-    case encode_js_data(payload) do
-      {:ok, binary} ->
-        delimiter = MarkdownHelpers.code_block_delimiter(binary)
+  defp render_output(%{type: :js, js_view: %{ref: ref}}, ctx) do
+    with {info_string, payload} <- ctx.js_ref_with_export[ref],
+         {:ok, binary} <- encode_js_data(payload) do
+      delimiter = MarkdownHelpers.code_block_delimiter(binary)
 
-        [delimiter, info_string, "\n", binary, "\n", delimiter]
-        |> prepend_metadata(%{output: true})
-
-      _ ->
-        :ignored
+      [delimiter, info_string, "\n", binary, "\n", delimiter]
+      |> prepend_metadata(%{output: true})
+    else
+      _ -> :ignored
     end
   end
 

lib/livebook/runtime.ex

@@ -142,25 +142,24 @@ defprotocol Livebook.Runtime do
 
   ## Export
 
-  The `:export` map describes how the output should be persisted.
-  The output data is put in a Markdown fenced code block.
+  The `:export` specifies whether the given output supports persistence.
+  When enabled, the JS view server should handle the following message:
 
-    * `:info_string` - used as the info string for the Markdown
-      code block
+      {:export, pid(), info :: %{ref: ref()}}
 
-    * `:key` - in case the data is a map and only a specific part
-      should be exported
+  And reply with:
 
+      {:export_reply, export_result, info :: %{ref: ref()}}
+
+  Where `export_result` is a tuple `{info_string, payload}`.
+  `info_string` is used as the info string for the Markdown code block,
+  while `payload` is its content. Payload can be either a string,
+  otherwise it is serialized into JSON.
   """
   @type js_output() :: %{
           type: :js,
           js_view: js_view(),
-          export:
-            nil
-            | %{
-                info_string: String.t(),
-                key: nil | term()
-              }
+          export: boolean()
         }
 
   @typedoc """

lib/livebook/session.ex

@@ -2896,6 +2896,13 @@ defmodule Livebook.Session do
     %{type: :image, content: content, mime_type: mime_type}
   end
 
+  # Rewrite older output format for backward compatibility with Kino <= 0.5.2
+  defp normalize_runtime_output({:js, %{ref: ref, pid: pid, assets: assets, export: export}}) do
+    normalize_runtime_output(
+      {:js, %{js_view: %{ref: ref, pid: pid, assets: assets}, export: export}}
+    )
+  end
+
   defp normalize_runtime_output({:js, info}) do
     %{type: :js, js_view: info.js_view, export: info.export}
   end

lib/livebook/session/data.ex

@@ -1226,19 +1226,6 @@ defmodule Livebook.Session.Data do
     |> update_cell_eval_info!(cell.id, &%{&1 | status: :ready})
   end
 
-  # Rewrite older output format for backward compatibility with Kino <= 0.5.2
-  defp add_cell_output(
-         data_actions,
-         cell,
-         {:js, %{ref: ref, pid: pid, assets: assets, export: export}}
-       ) do
-    add_cell_output(
-      data_actions,
-      cell,
-      {:js, %{js_view: %{ref: ref, pid: pid, assets: assets}, export: export}}
-    )
-  end
-
   defp add_cell_output({data, _} = data_actions, cell, output) do
     {[indexed_output], _counter} = Notebook.index_outputs([output], 0)
 

test/livebook/live_markdown/export_test.exs

@@ -758,7 +758,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
       assert expected_document == document
     end
 
-    test "does not include js output with no export info" do
+    test "does not include js output with export disabled" do
       notebook = %{
         Notebook.new()
         | name: "My Notebook",
@@ -772,15 +772,60 @@ defmodule Livebook.LiveMarkdown.ExportTest do
                     | source: ":ok",
                       outputs: [
                         {0,
-                         {:js,
-                          %{
-                            js_view: %{
-                              ref: "1",
-                              pid: spawn_widget_with_data("1", "data"),
-                              assets: %{archive_path: "", hash: "abcd", js_path: "main.js"}
-                            },
-                            export: nil
-                          }}}
+                         %{
+                           type: :js,
+                           js_view: %{
+                             ref: "1",
+                             pid: spawn_widget_with_data("1", "data"),
+                             assets: %{archive_path: "", hash: "abcd", js_path: "main.js"}
+                           },
+                           export: false
+                         }}
+                      ]
+                  }
+                ]
+            }
+          ]
+      }
+
+      expected_document = """
+      # My Notebook
+
+      ## Section 1
+
+      ```elixir
+      :ok
+      ```
+      """
+
+      {document, []} = Export.notebook_to_livemd(notebook, include_outputs: true)
+
+      assert expected_document == document
+    end
+
+    test "includes js output with export enabled" do
+      notebook = %{
+        Notebook.new()
+        | name: "My Notebook",
+          sections: [
+            %{
+              Notebook.Section.new()
+              | name: "Section 1",
+                cells: [
+                  %{
+                    Notebook.Cell.new(:code)
+                    | source: ":ok",
+                      outputs: [
+                        {0,
+                         %{
+                           type: :js,
+                           js_view: %{
+                             ref: "1",
+                             pid: spawn_widget_with_export("1", {"mermaid", "graph TD;\nA-->B;"}),
+                             assets: %{archive_path: "", hash: "abcd", js_path: "main.js"}
+                           },
+                           export: true
+                         }}
                       ]
                   }
                 ]
@@ -796,14 +841,21 @@ defmodule Livebook.LiveMarkdown.ExportTest do
       ```elixir
       :ok
       ```
+
+      <!-- livebook:{"output":true} -->
+
+      ```mermaid
+      graph TD;
+      A-->B;
+      ```
       """
 
       {document, []} = Export.notebook_to_livemd(notebook, include_outputs: true)
 
       assert expected_document == document
     end
 
-    test "includes js output if export info is set" do
+    test "includes js output with legacy export info" do
       notebook = %{
         Notebook.new()
         | name: "My Notebook",
@@ -1439,6 +1491,15 @@ defmodule Livebook.LiveMarkdown.ExportTest do
     metadata
   end
 
+  defp spawn_widget_with_export(ref, export_result) do
+    spawn(fn ->
+      receive do
+        {:export, pid, %{ref: ^ref}} ->
+          send(pid, {:export_reply, export_result, %{ref: ref}})
+      end
+    end)
+  end
+
   defp spawn_widget_with_data(ref, data) do
     spawn(fn ->
       receive do