@marimo-team/islands 0.21.2-dev6 → 0.21.2-dev61

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

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { beforeEach, describe, expect, it } from "vitest";
+import { cellId } from "@/__tests__/branded";
 import type { CellMessage } from "@/core/kernel/messages";
 import { invariant } from "@/utils/invariant";
 import {
@@ -22,13 +23,13 @@ describe("RunsState Reducer", () => {
   let state: RunsState;
 
   const runId = "run1" as RunId;
-  const cellId = "cell1";
+  const testCellId = cellId("cell1");
   const timestamp = Date.now();
   const code = "print('Hello World')";
 
   const cellNotification: CellMessage = {
     run_id: runId,
-    cell_id: cellId,
+    cell_id: testCellId,
     timestamp,
     status: "queued",
   };
@@ -61,9 +62,9 @@ describe("RunsState Reducer", () => {
       runStartTime: timestamp,
       cellRuns: new Map([
         [
-          cellId,
+          testCellId,
           {
-            cellId,
+            cellId: testCellId,
             code: code.slice(0, MAX_CODE_LENGTH),
             elapsedTime: 0,
             startTime: timestamp,
@@ -98,7 +99,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId2,
-          cell_id: "cell2",
+          cell_id: cellId("cell2"),
           timestamp,
           status: "queued",
         },
@@ -124,7 +125,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp: timestamp + 1000,
           status: "running",
         },
@@ -145,7 +146,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp: runStartTimestamp + 5000,
           status: "success",
         },
@@ -171,8 +172,8 @@ describe("RunsState Reducer", () => {
         type: "addCellNotification",
         payload: {
           cellNotification: {
-            run_id: `run${i}`,
-            cell_id: `cell${i}`,
+            run_id: `run${i}` as RunId,
+            cell_id: cellId(`cell${i}`),
             timestamp: timestamp,
             status: "queued",
           },
@@ -198,7 +199,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp,
           status: "queued",
         },
@@ -220,7 +221,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp: errorTimestamp,
           status: "running",
           output: {
@@ -248,7 +249,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp: errorTimestamp,
           status: "running",
           output: {
@@ -273,7 +274,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp,
           output: {
             channel: "marimo-error",
@@ -288,7 +289,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp: timestamp + 2000,
           status: "running", // shouldn't happen
         },
@@ -312,7 +313,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId2,
-          cell_id: "cell2",
+          cell_id: cellId("cell2"),
           timestamp: timestamp + 1000,
           status: "queued",
         },
@@ -325,7 +326,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId3,
-          cell_id: "cell3",
+          cell_id: cellId("cell3"),
           timestamp: timestamp + 2000,
           status: "queued",
         },
@@ -345,7 +346,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: "cell2",
+          cell_id: cellId("cell2"),
           timestamp: timestamp + 1000,
           status: "queued",
         },
@@ -367,7 +368,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp: timestamp + 1000,
           status: "running",
         },
@@ -386,7 +387,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp,
           status: "queued",
         },
@@ -405,7 +406,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: cellId,
+          cell_id: testCellId,
           timestamp,
           status: "queued",
         },
@@ -419,7 +420,7 @@ describe("RunsState Reducer", () => {
       payload: {
         cellNotification: {
           run_id: runId,
-          cell_id: "cell2",
+          cell_id: cellId("cell2"),
           timestamp: timestamp + 1000,
           status: "queued",
         },

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

@@ -1,7 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { type CellId, HTMLCellId } from "@/core/cells/ids";
+import { cellId } from "@/__tests__/branded";
+import { HTMLCellId } from "@/core/cells/ids";
 import { Logger } from "@/utils/Logger";
 
 // Mock the getCellEditorView function
@@ -20,12 +21,12 @@ vi.mock("@/core/codemirror/extensions", () => ({
 const { scrollCellIntoView } = await import("@/core/cells/scrollCellIntoView");
 
 describe("scrollCellIntoView", () => {
-  const cellId = "test-cell-id" as CellId;
+  const cid = cellId("test-cell-id");
   let cellElement: HTMLElement;
 
   beforeEach(() => {
     cellElement = document.createElement("div");
-    cellElement.id = HTMLCellId.create(cellId);
+    cellElement.id = HTMLCellId.create(cid);
     cellElement.scrollIntoView = vi.fn();
     document.body.append(cellElement);
 
@@ -42,7 +43,7 @@ describe("scrollCellIntoView", () => {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     mockGetCellEditorView.mockReturnValue(mockEditor as any);
 
-    scrollCellIntoView(cellId);
+    scrollCellIntoView(cid);
 
     expect(mockScrollActiveLineIntoView).toHaveBeenCalledWith(mockEditor, {
       behavior: "instant",
@@ -51,7 +52,7 @@ describe("scrollCellIntoView", () => {
   });
 
   it("should scroll cell element when editor is not focused", () => {
-    scrollCellIntoView(cellId);
+    scrollCellIntoView(cid);
 
     expect(mockScrollActiveLineIntoView).not.toHaveBeenCalled();
     expect(cellElement.scrollIntoView).toHaveBeenCalledWith({
@@ -64,10 +65,10 @@ describe("scrollCellIntoView", () => {
     const warnSpy = vi.spyOn(Logger, "warn").mockImplementation(vi.fn());
     cellElement.remove();
 
-    scrollCellIntoView(cellId);
+    scrollCellIntoView(cid);
 
     expect(warnSpy).toHaveBeenCalledWith(
-      `[CellFocusManager] scrollCellIntoView: element not found: ${cellId}`,
+      `[CellFocusManager] scrollCellIntoView: element not found: ${cid}`,
     );
     warnSpy.mockRestore();
   });

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

@@ -4,11 +4,12 @@ import type * as api from "@marimo-team/marimo-api";
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { Mocks } from "@/__mocks__/common";
+import { cellId } from "@/__tests__/branded";
 import { parseOutline } from "@/core/dom/outline";
 import { MultiColumn, visibleForTesting } from "@/utils/id-tree";
 import { invariant } from "@/utils/invariant";
 import { Logger } from "@/utils/Logger";
-import { type CellId, SETUP_CELL_ID } from "../ids";
+import { SETUP_CELL_ID } from "../ids";
 import { notebookStateFromSession } from "../session";
 
 // Mock dependencies
@@ -22,7 +23,7 @@ type SessionCell = api.Session["NotebookSessionV1"]["cells"][0];
 type NotebookCell = api.Notebook["NotebookV1"]["cells"][0];
 
 // Test constants
-const CELL_1 = "cell-1" as CellId;
+const CELL_1 = cellId("cell-1");
 
 describe("notebookStateFromSession", () => {
   beforeEach(() => {
@@ -671,9 +672,9 @@ describe("notebookStateFromSession", () => {
       );
       // Should have correct code and output for each cell
       for (const code of ["a", "b", "c", "d", "e", "f"]) {
-        const cellId = `cell-${code}` as CellId;
-        expect(result.cellData[cellId].code).toBe(code);
-        expect(result.cellRuntime[cellId].output).toEqual({
+        const cid = cellId(`cell-${code}`);
+        expect(result.cellData[cid].code).toBe(code);
+        expect(result.cellRuntime[cid].output).toEqual({
           channel: "output",
           data: `${code.toUpperCase()}!`,
           mimetype: "text/plain",
@@ -751,26 +752,26 @@ describe("notebookStateFromSession", () => {
       );
 
       // Should have correct code for each cell
-      expect(result.cellData["cell-a" as CellId].code).toBe("a");
-      expect(result.cellData["cell-c" as CellId].code).toBe("c");
-      expect(result.cellData["cell-z" as CellId].code).toBe("z");
-      expect(result.cellData["cell-e" as CellId].code).toBe("e");
-      expect(result.cellData["cell-g" as CellId].code).toBe("g");
+      expect(result.cellData[cellId("cell-a")].code).toBe("a");
+      expect(result.cellData[cellId("cell-c")].code).toBe("c");
+      expect(result.cellData[cellId("cell-z")].code).toBe("z");
+      expect(result.cellData[cellId("cell-e")].code).toBe("e");
+      expect(result.cellData[cellId("cell-g")].code).toBe("g");
 
       // Should have session outputs for matching cells (a, c, e)
-      expect(result.cellRuntime["cell-a" as CellId].output).toEqual({
+      expect(result.cellRuntime[cellId("cell-a")].output).toEqual({
         channel: "output",
         data: "A!",
         mimetype: "text/plain",
         timestamp: 0,
       });
-      expect(result.cellRuntime["cell-c" as CellId].output).toEqual({
+      expect(result.cellRuntime[cellId("cell-c")].output).toEqual({
         channel: "output",
         data: "C!",
         mimetype: "text/plain",
         timestamp: 0,
       });
-      expect(result.cellRuntime["cell-e" as CellId].output).toEqual({
+      expect(result.cellRuntime[cellId("cell-e")].output).toEqual({
         channel: "output",
         data: "E!",
         mimetype: "text/plain",
@@ -778,8 +779,8 @@ describe("notebookStateFromSession", () => {
       });
 
       // Should have no output for new cells (z, g) - they get stub session cells
-      expect(result.cellRuntime["cell-z" as CellId].output).toBeNull();
-      expect(result.cellRuntime["cell-g" as CellId].output).toBeNull();
+      expect(result.cellRuntime[cellId("cell-z")].output).toBeNull();
+      expect(result.cellRuntime[cellId("cell-g")].output).toBeNull();
 
       // Should log warning about different cells
       expect(Logger.warn).toHaveBeenCalledWith(
@@ -875,18 +876,18 @@ describe("notebookStateFromSession", () => {
       );
 
       // Should have correct code from notebook
-      expect(result.cellData["cell-1" as CellId].code).toBe("1 / 0");
-      expect(result.cellData["cell-2" as CellId].code).toBe('mo.md("Hello")');
-      expect(result.cellData["cell-3" as CellId].code).toBe(
+      expect(result.cellData[cellId("cell-1")].code).toBe("1 / 0");
+      expect(result.cellData[cellId("cell-2")].code).toBe('mo.md("Hello")');
+      expect(result.cellData[cellId("cell-3")].code).toBe(
         "x = mo.ui.slider(0, 10)",
       );
 
       // cell-1: No matching session cell (hash is null), gets stub session cell
-      expect(result.cellRuntime["cell-1" as CellId].output).toBeNull();
-      expect(result.cellRuntime["cell-1" as CellId].consoleOutputs).toEqual([]);
+      expect(result.cellRuntime[cellId("cell-1")].output).toBeNull();
+      expect(result.cellRuntime[cellId("cell-1")].consoleOutputs).toEqual([]);
 
       // cell-2: Matches session cell-1 by hash (moMd), gets its output
-      expect(result.cellRuntime["cell-2" as CellId].output).toEqual({
+      expect(result.cellRuntime[cellId("cell-2")].output).toEqual({
         channel: "output",
         data: "Welcome to marimo!",
         mimetype: "text/markdown",
@@ -894,7 +895,7 @@ describe("notebookStateFromSession", () => {
       });
 
       // cell-3: Matches session cell-2 by hash (slider), gets its output
-      expect(result.cellRuntime["cell-3" as CellId].output).toEqual({
+      expect(result.cellRuntime[cellId("cell-3")].output).toEqual({
         channel: "output",
         data: "",
         mimetype: "text/plain",

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,
@@ -755,14 +757,29 @@ const {
     });
   },
   handleCellMessage: (state, message: CellMessage) => {
-    const cellId = message.cell_id as CellId;
-    const nextState = updateCellRuntimeState({
+    const cellId = message.cell_id;
+    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);
 }
 
 /**