@marimo-team/islands 0.25.3-dev192 → 0.25.3-dev196

package/src/core/meta/__tests__/code-visibility.test.tsx

@@ -1,17 +1,21 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { renderHook } from "@testing-library/react";
-import { createStore, Provider } from "jotai";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { atom, createStore, Provider } from "jotai";
 import type { ReactNode } from "react";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { partial } from "@/__mocks__/partial";
 import { showCodeInRunModeAtom } from "@/core/meta/state";
 import { type AppMode, kioskModeAtom, viewStateAtom } from "@/core/mode";
+import type { Document, NotebookReplica } from "@/core/notebook/replica";
 import {
   remoteAdapter,
   runtimeAdapterAtom,
   staticAdapter,
   type RuntimeAdapter,
 } from "@/core/runtime/adapter";
+import { activeNotebookReplicaAtom } from "@/core/session/active-notebook";
 import { useNotebookCodeAvailable } from "../code-visibility";
 
 interface StoreOpts {
@@ -19,6 +23,39 @@ interface StoreOpts {
   kiosk?: boolean;
   showInRunMode?: boolean;
   runtime?: RuntimeAdapter;
+  /** Each cell's source as the host sent it; null is withheld. */
+  codes?: ReadonlyArray<string | null>;
+}
+
+function document(codes: ReadonlyArray<string | null>): Document {
+  const cells = codes.map((code, index): Protocol.Cell => ({
+    id: `cell-${index}` as Protocol.CellId,
+    notebook_id: "nb",
+    code,
+    name: "_",
+    config: {
+      column: null,
+      disabled: false,
+      hide_code: false,
+      expand_output: false,
+    },
+    etag: `etag-${index}`,
+    code_hash: `hash-${index}`,
+    version: 1,
+  }));
+  return {
+    revision: {
+      notebook_id: "nb",
+      order: cells.map((cell) => cell.id),
+      app_config: {},
+      etag: "etag",
+      saved_etag: "etag",
+      order_etag: "order",
+      app_config_etag: "config",
+      version: 1,
+    },
+    cells: new Map(cells.map((cell) => [cell.id, cell])),
+  };
 }
 
 function makeStore({
@@ -26,23 +63,32 @@ function makeStore({
   kiosk = false,
   showInRunMode = true,
   runtime = remoteAdapter,
+  codes = ["x = 1", ""],
 }: StoreOpts = {}) {
   const store = createStore();
   store.set(viewStateAtom, { mode, cellAnchor: null });
   store.set(kioskModeAtom, kiosk);
   store.set(showCodeInRunModeAtom, showInRunMode);
   store.set(runtimeAdapterAtom, runtime);
+  store.set(
+    activeNotebookReplicaAtom,
+    partial<NotebookReplica>(
+      { documentAtom: atom(document(codes)) },
+      "replica",
+    ),
+  );
   return store;
 }
 
-function wrap(store: ReturnType<typeof createStore>) {
-  return ({ children }: { children: ReactNode }) => (
+function render(store: ReturnType<typeof createStore>) {
+  const wrapper = ({ children }: { children: ReactNode }) => (
     <Provider store={store}>{children}</Provider>
   );
+  return renderHook(() => useNotebookCodeAvailable(), { wrapper }).result
+    .current;
 }
 
-const cellsWithCode = [{ code: "x = 1" }, { code: "" }];
-const cellsWithoutCode = [{ code: "" }, { code: "" }];
+const withheld = [null, null];
 
 const originalHref = window.location.href;
 
@@ -60,21 +106,11 @@ describe("useNotebookCodeAvailable", () => {
   });
 
   it("returns true in edit mode regardless of cells", () => {
-    const store = makeStore({ mode: "edit" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(render(makeStore({ mode: "edit", codes: withheld }))).toBe(true);
   });
 
   it("returns true in present mode (only reachable from edit)", () => {
-    const store = makeStore({ mode: "present" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(render(makeStore({ mode: "present", codes: withheld }))).toBe(true);
   });
 
   it("returns true in kiosk mode even when read mode and no code", () => {
@@ -82,88 +118,60 @@ describe("useNotebookCodeAvailable", () => {
       mode: "read",
       kiosk: true,
       showInRunMode: false,
+      codes: withheld,
     });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(render(store)).toBe(true);
   });
 
   it("returns true in read mode when at least one cell has code", () => {
-    const store = makeStore({ mode: "read" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(render(makeStore({ mode: "read" }))).toBe(true);
+  });
+
+  it("counts an empty cell as code the host sent", () => {
+    expect(render(makeStore({ mode: "read", codes: ["", null] }))).toBe(true);
+  });
+
+  it("returns false in read mode when the host withheld every source", () => {
+    expect(render(makeStore({ mode: "read", codes: withheld }))).toBe(false);
   });
 
-  it("returns false in read mode when every cell.code is empty (server stripped)", () => {
+  it("returns false in read mode before a document arrives", () => {
     const store = makeStore({ mode: "read" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(false);
+    store.set(activeNotebookReplicaAtom, null);
+    expect(render(store)).toBe(false);
   });
 
   it("returns false in read mode when host opts out via showAppCode", () => {
-    const store = makeStore({ mode: "read", showInRunMode: false });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithCode),
-      { wrapper: wrap(store) },
+    expect(render(makeStore({ mode: "read", showInRunMode: false }))).toBe(
+      false,
     );
-    expect(result.current).toBe(false);
   });
 
   it("returns false in read mode when ?include-code=false", () => {
     setSearch("?include-code=false");
-    const store = makeStore({ mode: "read" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(false);
+    expect(render(makeStore({ mode: "read" }))).toBe(false);
   });
 
   it("ignores ?include-code=false in kiosk mode", () => {
     setSearch("?include-code=false");
-    const store = makeStore({ mode: "read", kiosk: true });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(
+      render(makeStore({ mode: "read", kiosk: true, codes: withheld })),
+    ).toBe(true);
   });
 
   it("checks embedded code in a static kiosk", () => {
-    const store = makeStore({
+    const options = {
       mode: "read",
       kiosk: true,
       runtime: staticAdapter,
-    });
-
-    const withoutCode = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    const withCode = renderHook(() => useNotebookCodeAvailable(cellsWithCode), {
-      wrapper: wrap(store),
-    });
-
-    expect(withoutCode.result.current).toBe(false);
-    expect(withCode.result.current).toBe(true);
+    } as const;
+    expect(render(makeStore({ ...options, codes: withheld }))).toBe(false);
+    expect(render(makeStore(options))).toBe(true);
   });
 
   it("returns false for non-edit/non-read modes (home, gallery)", () => {
     for (const mode of ["home", "gallery"] as const) {
-      const store = makeStore({ mode });
-      const { result } = renderHook(
-        () => useNotebookCodeAvailable(cellsWithCode),
-        { wrapper: wrap(store) },
-      );
-      expect(result.current).toBe(false);
+      expect(render(makeStore({ mode }))).toBe(false);
     }
   });
 });

package/src/core/meta/code-visibility.ts

@@ -1,22 +1,31 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtomValue } from "jotai";
+import { atom, useAtomValue } from "jotai";
 import { KnownQueryParams } from "@/core/constants";
 import { showCodeInRunModeAtom } from "@/core/meta/state";
 import { kioskModeAtom, viewStateAtom } from "@/core/mode";
 import { runtimeAdapterAtom } from "@/core/runtime/adapter";
+import { confirmedDocumentAtom } from "@/core/session/active-notebook";
 import { logNever } from "@/utils/assertNever";
 
+/** Whether any cell's source reached this client. A host withholds source as null. */
+export const codeAvailableAtom = atom((get) => {
+  const document = get(confirmedDocumentAtom);
+  return (
+    document !== null &&
+    [...document.cells.values()].some((cell) => cell.code !== null)
+  );
+});
+
 /**
- * Whether the notebook source code reached the frontend and can be rendered.
+ * Whether the notebook source code can be rendered.
  *
- * In `marimo run` the server omits cell sources unless `--include-code` is
- * set (every `cell.code` arrives as `""`). Use this to gate any "show code"
- * / "copy code" / "download .py" affordance.
+ * In `marimo run` the server withholds cell sources unless `--include-code`
+ * is set. Use this to gate any "show code" / "copy code" / "download .py"
+ * affordance.
  */
-export function useNotebookCodeAvailable(
-  cells: ReadonlyArray<{ code: string }>,
-): boolean {
+export function useNotebookCodeAvailable(): boolean {
+  const available = useAtomValue(codeAvailableAtom);
   const kioskMode = useAtomValue(kioskModeAtom);
   const { mode } = useAtomValue(viewStateAtom);
   const showInRunMode = useAtomValue(showCodeInRunModeAtom);
@@ -41,7 +50,7 @@ export function useNotebookCodeAvailable(
       if (params.get(KnownQueryParams.includeCode) === "false") {
         return false;
       }
-      return cells.some((cell) => Boolean(cell.code));
+      return available;
     }
     default:
       logNever(mode);

package/src/core/notebook/__tests__/operations.test.ts

@@ -0,0 +1,51 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { NotebookOutput } from "@marimo-team/client/protocol";
+import { atom, createStore } from "jotai";
+import { describe, expect, it } from "vitest";
+import { partial } from "@/__mocks__/partial";
+import { kioskModeAtom } from "@/core/mode";
+import type { NotebookReplica } from "@/core/notebook/replica";
+import { activeNotebookReplicaAtom } from "@/core/session/active-notebook";
+import { canEditAtom, notebookOperationsAtom } from "../operations";
+
+function connect(
+  store: ReturnType<typeof createStore>,
+  operations: readonly string[] | null,
+) {
+  const notebook =
+    operations === null
+      ? null
+      : ({ id: "nb", operations } as unknown as NotebookOutput);
+  store.set(
+    activeNotebookReplicaAtom,
+    partial<NotebookReplica>({ notebookAtom: atom(notebook) }, "replica"),
+  );
+}
+
+describe("notebook operations", () => {
+  it("advertises nothing until the notebook is ready", () => {
+    const store = createStore();
+    expect(store.get(notebookOperationsAtom)).toEqual([]);
+    expect(store.get(canEditAtom)).toBe(false);
+    connect(store, null);
+    expect(store.get(canEditAtom)).toBe(false);
+  });
+
+  it("allows editing when the host lists notebook.edit", () => {
+    const store = createStore();
+    connect(store, ["notebook.edit", "runtime.run"]);
+    expect(store.get(canEditAtom)).toBe(true);
+    connect(store, ["runtime.run"]);
+    expect(store.get(canEditAtom)).toBe(false);
+  });
+
+  it("withholds editing from a kiosk connection", () => {
+    const store = createStore();
+    connect(store, ["notebook.edit"]);
+    store.set(kioskModeAtom, true);
+    expect(store.get(canEditAtom)).toBe(false);
+    store.set(kioskModeAtom, false);
+    expect(store.get(canEditAtom)).toBe(true);
+  });
+});

package/src/core/notebook/document/__tests__/codemirror.test.ts

@@ -1,6 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { history, undo } from "@codemirror/commands";
+import { MySQL, PostgreSQL } from "@codemirror/lang-sql";
+import { language as activeLanguage } from "@codemirror/language";
 import { EditorSelection, EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import type { CellId } from "@marimo-team/client/protocol";
@@ -65,6 +67,21 @@ function languages(id: CellId) {
       addSetupCell: vi.fn(),
       deleteCells: vi.fn(),
       undoDeleteCell: vi.fn(),
+      focusCell: vi.fn(),
+      focusTopCell: vi.fn(),
+      focusBottomCell: vi.fn(),
+      collapseCell: vi.fn(),
+      expandCell: vi.fn(),
+      collapseAllCells: vi.fn(),
+      expandAllCells: vi.fn(),
+      showCellIfHidden: vi.fn(),
+      moveCell: vi.fn(),
+      moveCellToIndex: vi.fn(),
+      moveCellsRelativeTo: vi.fn(),
+      dropCellOverCell: vi.fn(),
+      dropCellOverColumn: vi.fn(),
+      sendToTop: vi.fn(),
+      sendToBottom: vi.fn(),
       onRun: vi.fn(),
       deleteCell: vi.fn(),
       afterToggleMarkdown: vi.fn(),
@@ -504,14 +521,23 @@ describe("cellCodeBinding", () => {
     });
 
     it("sets the SQL dialect from the host's engine", async () => {
-      const { store, source } = setup({
+      const { store, source, view } = setup({
         adapters: true,
         initial: withCells(document(), {
           c1: cell("c1", 'df = mo.sql("SELECT 1")'),
         }),
       });
-      await settle();
-      const dialect = vi.spyOn(SCHEMA_CACHE, "getDialect");
+      await Promise.resolve();
+      const getDialect = SCHEMA_CACHE.getDialect.bind(SCHEMA_CACHE);
+      const dialects = new Map([
+        ["warehouse", PostgreSQL],
+        ["lake", MySQL],
+      ]);
+      const dialect = vi
+        .spyOn(SCHEMA_CACHE, "getDialect")
+        .mockImplementation(
+          (engine) => dialects.get(engine) ?? getDialect(engine),
+        );
       try {
         store.set(
           source,
@@ -519,8 +545,8 @@ describe("cellCodeBinding", () => {
             c1: cell("c1", '_df = mo.sql(f"""SELECT 1""", engine=warehouse)'),
           }),
         );
-        await settle();
-        expect(dialect).toHaveBeenLastCalledWith("warehouse");
+        await Promise.resolve();
+        expect(view.state.facet(activeLanguage)).toBe(PostgreSQL.language);
 
         store.set(
           source,
@@ -528,8 +554,8 @@ describe("cellCodeBinding", () => {
             c1: cell("c1", '_df = mo.sql(f"""SELECT 1""", engine=lake)'),
           }),
         );
-        await settle();
-        expect(dialect).toHaveBeenLastCalledWith("lake");
+        await Promise.resolve();
+        expect(view.state.facet(activeLanguage)).toBe(MySQL.language);
       } finally {
         dialect.mockRestore();
       }

package/src/core/notebook/index.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-/** React hooks for the current notebook. */
+/** What the editor reads about the current notebook. */
 
 export {
   useNotebookAttachments,
@@ -16,3 +16,4 @@ export {
   useNotebookDocumentEdit,
   useNotebookDocumentSync,
 } from "./document";
+export { canEditAtom, notebookOperationsAtom } from "./operations";

package/src/core/notebook/operations.ts

@@ -0,0 +1,30 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** What the host lets this client do to the notebook. */
+
+import { atom } from "jotai";
+import { kioskModeAtom } from "../mode";
+import { activeNotebookReplicaAtom } from "../session/active-notebook";
+
+const NO_OPERATIONS: readonly string[] = [];
+
+/** Operations the host advertises for this notebook, none until its stream is ready. */
+export const notebookOperationsAtom = atom((get) => {
+  const replica = get(activeNotebookReplicaAtom);
+  const notebook = replica === null ? null : get(replica.notebookAtom);
+  return notebook?.operations ?? NO_OPERATIONS;
+});
+
+/**
+ * Whether this client may edit the document.
+ *
+ * TODO(MO-8028): The host advertises the same operations to every
+ * connection of a notebook, so a kiosk connection, an interactor that may
+ * not edit, narrows it here. Once the host narrows `operations` per
+ * connection by consumer capability, the kiosk term goes.
+ */
+export const canEditAtom = atom(
+  (get) =>
+    !get(kioskModeAtom) &&
+    get(notebookOperationsAtom).includes("notebook.edit"),
+);

package/src/core/saving/__tests__/state.test.ts

@@ -0,0 +1,111 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import { atom, createStore } from "jotai";
+import { describe, expect, it } from "vitest";
+import { partial } from "@/__mocks__/partial";
+import type {
+  DocumentModel,
+  DocumentSyncStatus,
+} from "@/core/notebook/document";
+import type { Document, NotebookReplica } from "@/core/notebook/replica";
+import {
+  activeDocumentModelAtom,
+  activeNotebookReplicaAtom,
+} from "@/core/session/active-notebook";
+import { initialLayoutState, layoutStateAtom } from "../../layout/state";
+import { documentVersionAtom, needsSaveAtom, savedLayoutAtom } from "../state";
+
+function revision(
+  overrides: Partial<Protocol.Document> = {},
+): Protocol.Document {
+  return {
+    notebook_id: "nb",
+    order: [],
+    app_config: {},
+    etag: "v2",
+    saved_etag: "v2",
+    order_etag: "order",
+    app_config_etag: "config",
+    version: 2,
+    ...overrides,
+  };
+}
+
+function connect(
+  store: ReturnType<typeof createStore>,
+  document: Document | null,
+  status: DocumentSyncStatus["status"] = "idle",
+) {
+  store.set(
+    activeNotebookReplicaAtom,
+    partial<NotebookReplica>({ documentAtom: atom(document) }, "replica"),
+  );
+  store.set(
+    activeDocumentModelAtom,
+    partial<DocumentModel>(
+      { syncStatusAtom: atom({ status } as DocumentSyncStatus) },
+      "model",
+    ),
+  );
+}
+
+describe("needsSaveAtom", () => {
+  it("is false without a connected document", () => {
+    const store = createStore();
+    expect(store.get(needsSaveAtom)).toBe(false);
+    expect(store.get(documentVersionAtom)).toBeNull();
+    connect(store, null);
+    expect(store.get(needsSaveAtom)).toBe(false);
+  });
+
+  it("compares the host's etag with what it last wrote", () => {
+    const store = createStore();
+    connect(store, { revision: revision(), cells: new Map() });
+    expect(store.get(needsSaveAtom)).toBe(false);
+    expect(store.get(documentVersionAtom)).toBe(2);
+    connect(store, { revision: revision({ etag: "v3" }), cells: new Map() });
+    expect(store.get(needsSaveAtom)).toBe(true);
+  });
+
+  it("treats a document never written as unsaved", () => {
+    const store = createStore();
+    connect(store, {
+      revision: revision({ saved_etag: null }),
+      cells: new Map(),
+    });
+    expect(store.get(needsSaveAtom)).toBe(true);
+  });
+
+  it.each(["pending", "sending", "confirming", "error"] as const)(
+    "counts edits the host has not confirmed (%s)",
+    (status) => {
+      const store = createStore();
+      connect(store, { revision: revision(), cells: new Map() }, status);
+      expect(store.get(needsSaveAtom)).toBe(true);
+    },
+  );
+
+  it("ignores a model that is not running", () => {
+    const store = createStore();
+    connect(store, { revision: revision(), cells: new Map() }, "unavailable");
+    expect(store.get(needsSaveAtom)).toBe(false);
+  });
+
+  it("compares the layout with the one last saved", () => {
+    const store = createStore();
+    connect(store, { revision: revision(), cells: new Map() });
+    store.set(layoutStateAtom, {
+      selectedLayout: "grid",
+      layoutData: { grid: null },
+    });
+    expect(store.get(needsSaveAtom)).toBe(false);
+    store.set(savedLayoutAtom, initialLayoutState());
+    expect(store.get(needsSaveAtom)).toBe(true);
+    store.set(savedLayoutAtom, {
+      selectedLayout: "grid",
+      layoutData: { grid: null },
+    });
+    expect(store.get(needsSaveAtom)).toBe(false);
+  });
+});