@marimo-team/frontend 0.21.2-dev57 → 0.21.2-dev59

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/cells.ts

@@ -29,6 +29,7 @@ import type { CellConfig } from "../network/types";
 import { isRtcEnabled } from "../rtc/state";
 import { createDeepEqualAtom, store } from "../state/jotai";
 import { prepareCellForExecution, transitionCell } from "./cell";
+import { documentTransactionMiddleware } from "./document-changes";
 import { CellId, SCRATCH_CELL_ID, SETUP_CELL_ID } from "./ids";
 import { type CellLog, getCellLogsForMessage } from "./logs";
 import {
@@ -174,6 +175,7 @@ export interface CreateNewCellAction {
  */
 const {
   reducer,
+  addMiddleware,
   createActions,
   useActions,
   valueAtom: notebookAtom,
@@ -756,13 +758,28 @@ const {
   },
   handleCellMessage: (state, message: CellMessage) => {
     const cellId = message.cell_id;
-    const nextState = updateCellRuntimeState({
+    let nextState = updateCellRuntimeState({
       state,
       cellId,
       cellReducer: (cell) => {
         return transitionCell(cell, message);
       },
     });
+    // When a cell is queued for execution, snapshot the current code
+    // as lastCodeRun. This clears staleness for cells executed by the
+    // kernel (e.g. via code_mode). If the user edits during execution,
+    // code !== lastCodeRun keeps the cell stale.
+    if (message.status === "queued") {
+      nextState = updateCellData({
+        state: nextState,
+        cellId,
+        cellReducer: (cell) => ({
+          ...cell,
+          lastCodeRun: cell.code.trim(),
+          edited: false,
+        }),
+      });
+    }
     return {
       ...nextState,
       cellLogs: [...nextState.cellLogs, ...getCellLogsForMessage(message)],
@@ -1405,6 +1422,12 @@ const {
   },
 });
 
+// We apply the middleware here (rather than inline in createReducerAndAtoms)
+// so that the document transaction middleware can import CellActions and
+// strictly type the dispatched actions without creating a circular dependency.
+// @ts-expect-error - TODO: We should have better types for the middleware that are strict
+addMiddleware(documentTransactionMiddleware);
+
 function isCellCodeHidden(state: NotebookState, cellId: CellId): boolean {
   return (
     Boolean(state.cellData[cellId].config.hide_code) &&
@@ -1789,8 +1812,10 @@ export function createTracebackInfoAtom(
  * Use this hook to dispatch cell actions. This hook will not cause a re-render
  * when cells change.
  */
-export function useCellActions(): CellActions {
-  return useActions();
+export function useCellActions(
+  options: { skipMiddleware?: boolean } = {},
+): CellActions {
+  return useActions(options);
 }
 
 /**