@marimo-team/frontend 0.21.2-dev6 → 0.21.2-dev62

package/src/core/cells/__tests__/document-roundtrip.test.ts

@@ -0,0 +1,376 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/**
+ * Round-trip tests: perform actions on a primary notebook (producing changes
+ * via the middleware), then apply those changes to a replica notebook via
+ * applyTransactionChanges. The two should converge to identical document state.
+ *
+ * This catches drift between what the middleware emits and what
+ * apply-transaction consumes.
+ */
+
+import { python } from "@codemirror/lang-python";
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import {
+  afterAll,
+  afterEach,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  it,
+} from "vitest";
+import { cellId } from "@/__tests__/branded";
+import type { CellHandle } from "@/components/editor/notebook-cell";
+import { adaptiveLanguageConfiguration } from "@/core/codemirror/language/extension";
+import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
+import { MultiColumn } from "@/utils/id-tree";
+import { exportedForTesting, type NotebookState } from "../cells";
+import {
+  applyTransactionChanges,
+  exportedForTesting as middlewareExports,
+} from "../document-changes";
+import { CellId } from "../ids";
+
+const { initialNotebookState, reducer, createActions } = exportedForTesting;
+const { drainChanges } = middlewareExports;
+
+function createEditor(code: string) {
+  const state = EditorState.create({
+    doc: code,
+    extensions: [
+      python(),
+      adaptiveLanguageConfiguration({
+        cellId: cellId("cell1"),
+        completionConfig: {
+          activate_on_typing: true,
+          signature_hint_on_typing: false,
+          copilot: false,
+          codeium_api_key: null,
+        },
+        hotkeys: new OverridingHotkeyProvider({}),
+        placeholderType: "marimo-import",
+        lspConfig: {},
+      }),
+    ],
+  });
+  return new EditorView({ state, parent: document.body });
+}
+
+// --- Primary notebook: performs actions, middleware produces changes ---
+
+let primary: NotebookState;
+
+const primaryActions = createActions((action) => {
+  primary = reducer(primary, action);
+  for (const [cellIdString, handle] of Object.entries(primary.cellHandles)) {
+    // @ts-expect-error - Object.entries doesn't know keys are CellId
+    const cid: CellId = cellIdString;
+    if (!handle.current) {
+      const view = createEditor(primary.cellData[cid].code);
+      const h: CellHandle = { editorView: view, editorViewOrNull: view };
+      primary.cellHandles[cid] = { current: h };
+    }
+  }
+});
+
+// --- Replica notebook: receives changes via applyTransactionChanges ---
+
+let replica: NotebookState;
+
+const replicaActions = createActions((action) => {
+  replica = reducer(replica, action);
+  for (const [cellIdString, handle] of Object.entries(replica.cellHandles)) {
+    // @ts-expect-error - Object.entries doesn't know keys are CellId
+    const cid: CellId = cellIdString;
+    if (!handle.current) {
+      const view = createEditor(replica.cellData[cid].code);
+      const h: CellHandle = { editorView: view, editorViewOrNull: view };
+      replica.cellHandles[cid] = { current: h };
+    }
+  }
+});
+
+let i = 0;
+const originalCreate = CellId.create.bind(CellId);
+
+beforeAll(() => {
+  CellId.create = () => cellId(`${i++}`);
+});
+
+beforeEach(() => {
+  i = 0;
+  primary = initialNotebookState();
+  primary.cellIds = MultiColumn.from([]);
+  drainChanges();
+});
+
+afterEach(() => {
+  middlewareExports.cancelPendingChanges();
+});
+
+afterAll(() => {
+  CellId.create = originalCreate;
+});
+
+/** Set up both notebooks with the same initial cells. */
+function setup(...codes: string[]) {
+  for (const code of codes) {
+    primaryActions.createNewCell({
+      cellId: "__end__",
+      before: false,
+      code,
+      newCellId: CellId.create(),
+    });
+  }
+
+  // Apply the setup changes to the replica so both start identical
+  const setupChanges = drainChanges();
+  replica = initialNotebookState();
+  replica.cellIds = MultiColumn.from([]);
+  applyTransactionChanges(
+    setupChanges,
+    replicaActions,
+    () => replica.cellIds.inOrderIds,
+  );
+  // Drain any changes the replica's middleware produced
+  drainChanges();
+}
+
+/**
+ * Drain changes from the primary's middleware and apply them to the replica.
+ */
+function sync() {
+  const changes = drainChanges();
+  applyTransactionChanges(
+    changes,
+    replicaActions,
+    () => replica.cellIds.inOrderIds,
+  );
+  // Drain any changes the replica's middleware produced (we don't want those)
+  drainChanges();
+}
+
+/**
+ * Extract the document-relevant state: cell ordering, code, name, config.
+ * This is the "NotebookDocument" equivalent — what the Python side tracks.
+ *
+ * TODO(column-config): config.column is excluded because the column
+ * reducers (addColumnBreakpoint, dropOverNewColumn, moveColumn, etc.)
+ * update cellIds (MultiColumn structure) but don't sync config.column
+ * on affected cells. The middleware correctly emits set-config changes with
+ * the new column index, but the primary's config.column stays stale,
+ * causing a mismatch with the replica. Fix: have the column reducers
+ * update config.column as part of their state transition, then remove
+ * the { column: _, ...config } exclusion here.
+ */
+function documentSnapshot(state: NotebookState) {
+  return state.cellIds.inOrderIds.map((id) => {
+    const { column: _, ...config } = state.cellData[id].config;
+    return {
+      id,
+      code: state.cellData[id].code,
+      name: state.cellData[id].name,
+      config,
+    };
+  });
+}
+
+describe("document round-trip", () => {
+  it("initial setup converges", () => {
+    setup("a", "b", "c");
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("createNewCell at end", () => {
+    setup("a", "b");
+    primaryActions.createNewCell({
+      cellId: "__end__",
+      before: false,
+      code: "c",
+      newCellId: CellId.create(),
+    });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("createNewCell before first cell", () => {
+    setup("a", "b");
+    const [a] = primary.cellIds.inOrderIds;
+    primaryActions.createNewCell({
+      cellId: a,
+      before: true,
+      code: "before-a",
+      newCellId: CellId.create(),
+    });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("createNewCell between cells", () => {
+    setup("a", "b", "c");
+    const [a] = primary.cellIds.inOrderIds;
+    primaryActions.createNewCell({
+      cellId: a,
+      before: false,
+      code: "between",
+      newCellId: CellId.create(),
+    });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("createNewCell with hideCode config", () => {
+    setup("a");
+    primaryActions.createNewCell({
+      cellId: "__end__",
+      before: false,
+      code: "hidden",
+      newCellId: CellId.create(),
+      hideCode: true,
+    });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("deleteCell", () => {
+    setup("a", "b", "c");
+    const [, b] = primary.cellIds.inOrderIds;
+    primaryActions.deleteCell({ cellId: b });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("updateCellCode", () => {
+    setup("a", "b");
+    const [a] = primary.cellIds.inOrderIds;
+    primaryActions.updateCellCode({
+      cellId: a,
+      code: "updated",
+      formattingChange: false,
+    });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("updateCellName", () => {
+    setup("a");
+    const [a] = primary.cellIds.inOrderIds;
+    primaryActions.updateCellName({ cellId: a, name: "my_var" });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("updateCellConfig", () => {
+    setup("a");
+    const [a] = primary.cellIds.inOrderIds;
+    primaryActions.updateCellConfig({
+      cellId: a,
+      config: { hide_code: true, disabled: true },
+    });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("moveCell down", () => {
+    setup("a", "b", "c");
+    const [a] = primary.cellIds.inOrderIds;
+    primaryActions.moveCell({ cellId: a, before: false });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("sendToTop", () => {
+    setup("a", "b", "c");
+    const c = primary.cellIds.inOrderIds[2];
+    primaryActions.sendToTop({ cellId: c });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("sendToBottom", () => {
+    setup("a", "b", "c");
+    const [a] = primary.cellIds.inOrderIds;
+    primaryActions.sendToBottom({ cellId: a });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("dropCellOverCell", () => {
+    setup("a", "b", "c");
+    const [a, , c] = primary.cellIds.inOrderIds;
+    primaryActions.dropCellOverCell({ cellId: c, overCellId: a });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("multiple operations in sequence", () => {
+    setup("a", "b", "c");
+
+    // Add a cell
+    const [a] = primary.cellIds.inOrderIds;
+    primaryActions.createNewCell({
+      cellId: a,
+      before: false,
+      code: "new",
+      newCellId: CellId.create(),
+    });
+    sync();
+
+    // Rename it
+    const newId = primary.cellIds.inOrderIds[1];
+    primaryActions.updateCellName({ cellId: newId, name: "inserted" });
+    sync();
+
+    // Move it to top
+    primaryActions.sendToTop({ cellId: newId });
+    sync();
+
+    // Update code on another cell
+    const last =
+      primary.cellIds.inOrderIds[primary.cellIds.inOrderIds.length - 1];
+    primaryActions.updateCellCode({
+      cellId: last,
+      code: "modified",
+      formattingChange: false,
+    });
+    sync();
+
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("create then delete", () => {
+    setup("a", "b");
+    primaryActions.createNewCell({
+      cellId: "__end__",
+      before: false,
+      code: "temporary",
+      newCellId: CellId.create(),
+    });
+    sync();
+
+    const newId =
+      primary.cellIds.inOrderIds[primary.cellIds.inOrderIds.length - 1];
+    primaryActions.deleteCell({ cellId: newId });
+    sync();
+
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("addColumnBreakpoint", () => {
+    setup("a", "b", "c");
+    const [, b] = primary.cellIds.inOrderIds;
+    primaryActions.addColumnBreakpoint({ cellId: b });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+
+  it("dropOverNewColumn", () => {
+    setup("a", "b", "c");
+    const [, b] = primary.cellIds.inOrderIds;
+    primaryActions.dropOverNewColumn({ cellId: b });
+    sync();
+    expect(documentSnapshot(primary)).toEqual(documentSnapshot(replica));
+  });
+});

package/src/core/cells/__tests__/focus.test.ts

@@ -1,15 +1,16 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+
 import { beforeEach, describe, expect, it } from "vitest";
-import type { CellId } from "@/core/cells/ids";
+import { cellId } from "@/__tests__/branded";
 import type { CellFocusState } from "../focus";
 import { exportedForTesting } from "../focus";
 
 const { initialState, reducer, createActions } = exportedForTesting;
 
 const CellIds = {
-  a: "a" as CellId,
-  b: "b" as CellId,
-  c: "c" as CellId,
+  a: cellId("a"),
+  b: cellId("b"),
+  c: cellId("c"),
 };
 
 describe("cell focus reducer", () => {

package/src/core/cells/__tests__/logs.test.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
 import type { CellMessage } from "../../kernel/messages";
 import { formatLogTimestamp, getCellLogsForMessage } from "../logs";
 
@@ -18,7 +19,7 @@ describe("getCellLogsForMessage", () => {
 
   test("handles text/plain MIME type on stdout", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-1",
+      cell_id: cellId("cell-1"),
       console: [
         {
           mimetype: "text/plain",
@@ -46,7 +47,7 @@ describe("getCellLogsForMessage", () => {
 
   test("handles text/plain MIME type on stderr", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-2",
+      cell_id: cellId("cell-2"),
       console: [
         {
           mimetype: "text/plain",
@@ -74,7 +75,7 @@ describe("getCellLogsForMessage", () => {
 
   test("handles text/html MIME type and strips HTML tags", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-3",
+      cell_id: cellId("cell-3"),
       console: [
         {
           mimetype: "text/html",
@@ -102,7 +103,7 @@ describe("getCellLogsForMessage", () => {
 
   test("handles text/html MIME type on stderr", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-4",
+      cell_id: cellId("cell-4"),
       console: [
         {
           mimetype: "text/html",
@@ -130,7 +131,7 @@ describe("getCellLogsForMessage", () => {
 
   test("handles application/vnd.marimo+traceback MIME type and strips HTML", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-5",
+      cell_id: cellId("cell-5"),
       console: [
         {
           mimetype: "application/vnd.marimo+traceback",
@@ -156,7 +157,7 @@ describe("getCellLogsForMessage", () => {
 
   test("handles multiple console outputs with different MIME types", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-7",
+      cell_id: cellId("cell-7"),
       console: [
         {
           mimetype: "text/plain",
@@ -196,7 +197,7 @@ describe("getCellLogsForMessage", () => {
     vi.spyOn(Date, "now").mockReturnValue(now);
 
     const cellMessage: CellMessage = {
-      cell_id: "cell-8",
+      cell_id: cellId("cell-8"),
       console: [
         {
           mimetype: "text/plain",
@@ -219,7 +220,7 @@ describe("getCellLogsForMessage", () => {
 
   test("ignores unsupported MIME types", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-9",
+      cell_id: cellId("cell-9"),
       console: [
         {
           mimetype: "application/json",
@@ -241,7 +242,7 @@ describe("getCellLogsForMessage", () => {
 
   test("ignores non-logging channels", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-10",
+      cell_id: cellId("cell-10"),
       console: [
         {
           mimetype: "text/plain",
@@ -263,7 +264,7 @@ describe("getCellLogsForMessage", () => {
 
   test("returns empty array when console is null", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-11",
+      cell_id: cellId("cell-11"),
       console: null as unknown as CellMessage["console"],
       output: null,
       status: "idle",
@@ -278,7 +279,7 @@ describe("getCellLogsForMessage", () => {
 
   test("handles complex HTML with nested elements in text/html", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-12",
+      cell_id: cellId("cell-12"),
       console: [
         {
           mimetype: "text/html",
@@ -301,7 +302,7 @@ describe("getCellLogsForMessage", () => {
 
   test("handles marimo-error channel as stderr level", () => {
     const cellMessage: CellMessage = {
-      cell_id: "cell-13",
+      cell_id: cellId("cell-13"),
       console: [
         {
           mimetype: "text/plain",

package/src/core/cells/__tests__/pending-delete-service.test.tsx

@@ -4,11 +4,11 @@ import { act, renderHook, waitFor } from "@testing-library/react";
 import { createStore, Provider } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { MockNotebook } from "@/__mocks__/notebook";
+import { variableName } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import { CellId } from "@/core/cells/ids";
 import { createCellRuntimeState } from "@/core/cells/types";
 import { variablesAtom } from "@/core/variables/state";
-import type { VariableName } from "@/core/variables/types";
 import type { Milliseconds } from "@/utils/time";
 import {
   usePendingDelete,
@@ -138,8 +138,8 @@ describe("pending-delete-service", () => {
 
     store.set(notebookAtom, notebook);
     store.set(variablesAtom, {
-      ["x" as VariableName]: {
-        name: "x" as VariableName,
+      [variableName("x")]: {
+        name: variableName("x"),
         declaredBy: [cell1Id],
         usedBy: [cell2Id],
       },