fix: wasm loading styles and hide save (#3396)

frontend/src/components/editor/header/filename-input.tsx

@@ -152,6 +152,7 @@ export const FilenameInput = ({
         onFocus={onFocus}
         onBlur={onBlur}
         shouldFilter={false}
+        id="filename-input"
         className="bg-transparent group filename-input"
       >
         <CommandList>

frontend/src/core/wasm/__tests__/state.test.ts

@@ -0,0 +1,91 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { expect, describe, it } from "vitest";
+import { store } from "@/core/state/jotai";
+import { hasAnyOutputAtom } from "../state";
+import { notebookAtom } from "@/core/cells/cells";
+import type { NotebookState } from "@/core/cells/cells";
+import { MultiColumn, CollapsibleTree } from "@/utils/id-tree";
+import type { OutputMessage } from "@/core/kernel/messages";
+import type { CellId } from "@/core/cells/ids";
+import { createRef } from "react";
+import { createCellRuntimeState } from "@/core/cells/types";
+import { createCell } from "@/core/cells/types";
+
+describe("hasAnyOutputAtom", () => {
+  const createNotebookState = (
+    outputs: Array<OutputMessage | null>,
+  ): NotebookState => ({
+    cellIds: new MultiColumn([
+      CollapsibleTree.from(outputs.map((_, i) => `${i}` as CellId)),
+    ]),
+    cellData: Object.fromEntries(
+      outputs.map((_, i) => [
+        `${i}` as CellId,
+        createCell({ id: `${i}` as CellId }),
+      ]),
+    ),
+    cellRuntime: Object.fromEntries(
+      outputs.map((output, i) => [
+        `${i}` as CellId,
+        createCellRuntimeState({
+          output,
+          outline: { items: [] },
+        }),
+      ]),
+    ),
+    cellHandles: Object.fromEntries(
+      outputs.map((_, i) => [`${i}` as CellId, createRef()]),
+    ),
+    cellLogs: [],
+    scrollKey: null,
+    history: [],
+  });
+
+  it("should return false when there are no outputs", () => {
+    store.set(notebookAtom, createNotebookState([null, null]));
+    expect(store.get(hasAnyOutputAtom)).toBe(false);
+  });
+
+  it("should return false when all outputs are empty", () => {
+    store.set(
+      notebookAtom,
+      createNotebookState([
+        { channel: "output", mimetype: "text/plain", data: "", timestamp: 0 },
+        { channel: "output", mimetype: "text/plain", data: "", timestamp: 0 },
+      ]),
+    );
+    expect(store.get(hasAnyOutputAtom)).toBe(false);
+  });
+
+  it("should return true when there is at least one non-empty output", () => {
+    store.set(
+      notebookAtom,
+      createNotebookState([
+        { channel: "output", mimetype: "text/plain", data: "", timestamp: 0 },
+        {
+          channel: "output",
+          mimetype: "text/plain",
+          data: "hello",
+          timestamp: 0,
+        },
+      ]),
+    );
+    expect(store.get(hasAnyOutputAtom)).toBe(true);
+  });
+
+  it("should handle various output types", () => {
+    store.set(
+      notebookAtom,
+      createNotebookState([
+        {
+          channel: "output",
+          mimetype: "application/json",
+          data: { foo: "bar" },
+          timestamp: 0,
+        },
+        { channel: "output", mimetype: "text/plain", data: "", timestamp: 0 },
+      ]),
+    );
+    expect(store.get(hasAnyOutputAtom)).toBe(true);
+  });
+});

frontend/src/core/wasm/state.ts

@@ -1,12 +1,13 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import { atom } from "jotai";
 import { notebookAtom } from "../cells/cells";
+import { isOutputEmpty } from "../cells/outputs";
 
 export const wasmInitializationAtom = atom<string>("Initializing...");
 
 export const hasAnyOutputAtom = atom<boolean>((get) => {
   const notebook = get(notebookAtom);
-  return Object.values(notebook.cellRuntime).some((runtime) =>
-    Boolean(runtime.output),
+  return Object.values(notebook.cellRuntime).some(
+    (runtime) => !isOutputEmpty(runtime.output),
   );
 });

marimo/_server/templates/templates.py

@@ -273,6 +273,19 @@ def wasm_notebook_template(
     """
     body = body.replace("</head>", f"{warning_script}</head>")
 
+    # Hide save button in WASM mode
+    wasm_styles = """
+    <style>
+        #save-button {
+            display: none !important;
+        }
+        #filename-input {
+            display: none !important;
+        }
+    </style>
+    """
+    body = body.replace("</head>", f"{wasm_styles}</head>")
+
     # If has custom css, inline the css and add to the head
     if app_config.css_file:
         css_contents = read_css_file(app_config.css_file, filename=filename)

tests/_server/templates/test_templates.py

@@ -464,5 +464,7 @@ def test_wasm_notebook_template_custom_head(self) -> None:
             assert head in result
             assert '<marimo-wasm hidden="">' in result
             assert '<marimo-code hidden="" data-show-code="false">' in result
+            assert "#save-button" in result
+            assert "#filename-input" in result
         finally:
             os.remove(head_file)