@marimo-team/islands 0.25.3-dev197 → 0.25.3-dev198

package/src/components/slides/minimap.tsx

@@ -2,10 +2,9 @@
 
 import { useDeleteCellCallback } from "@/components/editor/cell/useDeleteCell";
 import { outputIsStale } from "@/core/cells/cell";
-import { useCellIds } from "@/core/cells/cells";
 import { useCellActions } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
-import type { CellColumnId } from "@/utils/id-tree";
+import { type CellColumnId, type Layout, useLayout } from "@/core/cells/layout";
 import { useEffect, useRef, useState } from "react";
 import type { ICellRendererProps } from "../editor/renderers/types";
 import type {
@@ -50,7 +49,6 @@ import { SLIDE_TYPE_OPTIONS_BY_VALUE } from "./slide-form";
 
 type Props = ICellRendererProps<SlidesLayout>;
 type SlideCell = Props["cells"][number];
-type CellIdsState = ReturnType<typeof useCellIds>;
 type DropPosition = "before" | "after";
 
 interface ProjectedDropTarget {
@@ -229,7 +227,7 @@ export const SlidesMinimap = ({
   slideTypes,
   onSlideClick,
 }: SlidesMinimapProps) => {
-  const cellIds = useCellIds();
+  const layout = useLayout();
   const { moveCellToIndex, createCell } = useCellActions();
   const deleteCell = useDeleteCellCallback();
   const containerRef = useRef<HTMLDivElement>(null);
@@ -297,7 +295,7 @@ export const SlidesMinimap = ({
     try {
       if (activeId && dropTarget) {
         const resolvedTarget = resolveDropTarget({
-          cellIds,
+          layout,
           activeId,
           target: dropTarget,
         });
@@ -751,24 +749,27 @@ function projectDropTarget(
 }
 
 function resolveDropTarget({
-  cellIds,
+  layout,
   activeId,
   target,
 }: {
-  cellIds: CellIdsState;
+  layout: Layout;
   activeId: CellId;
   target: ProjectedDropTarget;
 }): ResolvedDropTarget | null {
   if (activeId === target.overId) {
     return null;
   }
-  if (cellIds.colLength !== 1) {
+  if (layout.columns.length !== 1) {
     Logger.warn("Multi-column mode is not supported");
     return null;
   }
 
-  const column = cellIds.findWithId(target.overId);
-  const overIndex = column.indexOfOrThrow(target.overId);
+  const [column] = layout.columns;
+  const overIndex = column.ids.indexOf(target.overId);
+  if (overIndex === -1) {
+    return null;
+  }
 
   return {
     cellId: activeId,

package/src/components/sort/SortableCellsProvider.tsx

@@ -22,10 +22,8 @@ import React, { useCallback, useMemo, useState } from "react";
 import type { CellId } from "@/core/cells/ids";
 import { useAppConfig } from "@/core/config/config";
 import { Arrays } from "@/utils/arrays";
-import type { CellColumnId, MultiColumn } from "@/utils/id-tree";
 import { invariant } from "@/utils/invariant";
-import { getNotebook } from "../../core/cells/cells";
-import { findColumn, layoutAtom } from "../../core/cells/layout";
+import { findColumn, isColumnId, layoutAtom } from "../../core/cells/layout";
 import { store } from "../../core/state/jotai";
 import { useCellActions } from "../../core/cells/actions";
 import { useEvent } from "../../hooks/useEvent";
@@ -50,9 +48,6 @@ const SortableCellsProviderInternal = ({
     useCellActions();
 
   const [activeId, setActiveId] = useState<UniqueIdentifier | null>(null);
-  const [clonedItems, setClonedItems] = useState<MultiColumn<CellId> | null>(
-    null,
-  );
 
   const config = useAppConfig();
   const modifiers = useMemo(() => {
@@ -76,20 +71,10 @@ const SortableCellsProviderInternal = ({
 
   const handleDragStart = useEvent((event: DragStartEvent) => {
     setActiveId(event.active.id);
-    const notebook = getNotebook();
-    setClonedItems(notebook.cellIds);
   });
 
   const handleDragCancel = useEvent(() => {
-    // TODO: restore cloned items
-    if (clonedItems) {
-      // Reset items to their original state in case items have been
-      // Dragged across containers
-      // setItems(clonedItems);
-    }
-
     setActiveId(null);
-    setClonedItems(null);
   });
 
   /**
@@ -236,9 +221,5 @@ const SortableCellsProviderInternal = ({
 export const SortableCellsProvider = React.memo(SortableCellsProviderInternal);
 
 function isCellId(id: UniqueIdentifier): id is CellId {
-  return typeof id === "string" && !id.startsWith("tree_");
-}
-
-function isColumnId(id: UniqueIdentifier): id is CellColumnId {
-  return typeof id === "string" && id.startsWith("tree_");
+  return typeof id === "string" && !isColumnId(id);
 }

package/src/components/tracing/tracing.tsx

@@ -209,7 +209,7 @@ const TraceBlockBody: React.FC<{
       const elapsedTime = cellRun.elapsedTime ?? 0;
       return {
         cell: cellRun.cellId,
-        cellNum: cellIds.inOrderIds.indexOf(cellRun.cellId),
+        cellNum: cellIds.indexOf(cellRun.cellId),
         startTimestamp: formatChartTime(cellRun.startTime),
         endTimestamp: formatChartTime(cellRun.startTime + elapsedTime),
         elapsedTime: formatElapsedTime(elapsedTime * 1000),

package/src/core/ai/__tests__/staged-cells.test.ts

@@ -43,7 +43,7 @@ vi.mock("../../cells/actions", () => ({
 vi.mock("../../cells/cells", async () => {
   const { atom } = await import("jotai");
   return {
-    notebookAtom: atom({ cellData: {}, cellIds: { inOrderIds: [] } }),
+    notebookAtom: atom({ cellData: {}, cellIds: [] }),
     cellHandleAtom: vi.fn(() => ({
       read: vi.fn(() => mockCellHandle),
     })),

package/src/core/ai/context/providers/__tests__/cell-output.test.ts

@@ -34,9 +34,7 @@ describe("CellOutputContextProvider", () => {
   beforeEach(() => {
     // Create a basic mock notebook state
     mockNotebook = {
-      cellIds: {
-        inOrderIds: ["cell1", "cell2", "cell3"],
-      },
+      cellIds: ["cell1", "cell2", "cell3"],
       cellData: {
         cell1: {
           id: "cell1",
@@ -207,7 +205,7 @@ describe("Cell output utility functions", () => {
 
     beforeEach(() => {
       const mockNotebook = {
-        cellIds: { inOrderIds: [] },
+        cellIds: [],
         cellData: {},
         cellRuntime: {},
       } as unknown as NotebookState;
@@ -226,7 +224,7 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test"] },
+          cellIds: ["test"],
           cellData: {
             test: { id: "test", name: "", code: "" },
           },
@@ -261,7 +259,7 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test"] },
+          cellIds: ["test"],
           cellData: {
             test: { id: "test", name: "", code: "" },
           },
@@ -295,7 +293,7 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test"] },
+          cellIds: ["test"],
           cellData: {
             test: { id: "test", name: "", code: "" },
           },
@@ -338,7 +336,7 @@ describe("Cell output utility functions", () => {
       } as unknown as Document;
 
       const mockNotebook = {
-        cellIds: { inOrderIds: [] },
+        cellIds: [],
         cellData: {},
         cellRuntime: {},
       } as unknown as NotebookState;
@@ -369,7 +367,7 @@ describe("Cell output utility functions", () => {
       ).mockReturnValue(mockDiv);
 
       mockStore = createMockStore({
-        cellIds: { inOrderIds: ["test"] },
+        cellIds: ["test"],
         cellData: {
           test: { id: "test", name: "", code: "" },
         },
@@ -395,9 +393,7 @@ describe("Cell output utility functions", () => {
 
     beforeEach(() => {
       mockNotebook = {
-        cellIds: {
-          inOrderIds: ["cell1", "cell2"],
-        },
+        cellIds: ["cell1", "cell2"],
         cellData: {
           cell1: {
             id: "cell1",

package/src/core/ai/context/providers/cell-output.ts

@@ -93,7 +93,7 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
     const notebook = this.store.get(notebookAtom);
     const items: CellOutputContextItem[] = [];
 
-    for (const cellId of notebook.cellIds.inOrderIds) {
+    for (const cellId of notebook.cellIds) {
       const cellContextData = getCellContextData(cellId, notebook);
       const cellOutput = cellContextData.cellOutput;
 
@@ -358,7 +358,7 @@ export function getCellContextData(
   const cellRuntime = notebook.cellRuntime[cellId];
 
   const cellData = notebook.cellData[cellId];
-  const cellIndex = notebook.cellIds.inOrderIds.indexOf(cellId);
+  const cellIndex = notebook.cellIds.indexOf(cellId);
   const cellName = displayCellName(cellData.name, cellIndex);
 
   let consoleOutputs: BaseOutput[] | undefined;

package/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts

@@ -4,12 +4,18 @@ import { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import { getDefaultStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { mockDocumentModel } from "@/__mocks__/document";
 import { MockNotebook } from "@/__mocks__/notebook";
 import { cellId } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
-import { MultiColumn } from "@/utils/id-tree";
+import {
+  cell,
+  CONFIG,
+  document as documentOf,
+  withCells,
+} from "@/core/notebook/document/__tests__/fixtures";
 import { cellConfigExtension } from "../../../codemirror/config/extension";
 import { adaptiveLanguageConfiguration } from "../../../codemirror/language/extension";
 import { stagedAICellsAtom } from "../../staged-cells";
@@ -61,7 +67,6 @@ describe("EditNotebookTool", () => {
   let tool: EditNotebookTool;
   let cellId1: CellId;
   let cellId2: CellId;
-  let cellId3: CellId;
   let toolContext: {
     addStagedCell: ReturnType<typeof vi.fn>;
     createCell: ReturnType<typeof vi.fn>;
@@ -86,7 +91,6 @@ describe("EditNotebookTool", () => {
 
     cellId1 = cellId("cell-1");
     cellId2 = cellId("cell-2");
-    cellId3 = cellId("cell-3");
 
     // Reset mocks
     vi.clearAllMocks();
@@ -443,27 +447,37 @@ describe("EditNotebookTool", () => {
           [cellId2]: { code: "x = 2" },
         },
       });
-      // Create multi-column layout
-      notebook.cellIds = MultiColumn.from([[cellId1], [cellId2]]);
       store.set(notebookAtom, notebook);
-
-      const newCode = "y = 2";
-      const result = await tool.handler(
-        {
-          edit: {
-            type: "add_cell",
-            position: { type: "column_end", columnIndex: 1 },
-            code: newCode,
-          },
-        },
-        toolContext as never,
+      // The layout reads columns from the document's cells.
+      const { stop } = mockDocumentModel(
+        store,
+        withCells(documentOf(), {
+          c1: null,
+          c2: null,
+          [cellId1]: cell(cellId1, "x = 1"),
+          [cellId2]: cell(cellId2, "x = 2", {
+            config: { ...CONFIG, column: 1 },
+          }),
+        }),
       );
 
-      expect(result.status).toBe("success");
+      try {
+        const result = await tool.handler(
+          {
+            edit: {
+              type: "add_cell",
+              position: { type: "column_end", columnIndex: 1 },
+              code: "y = 2",
+            },
+          },
+          toolContext as never,
+        );
 
-      // Check that a new cell was staged
-      const stagedCells = store.get(stagedAICellsAtom);
-      expect(stagedCells.size).toBe(1);
+        expect(result.status).toBe("success");
+        expect(store.get(stagedAICellsAtom).size).toBe(1);
+      } finally {
+        stop();
+      }
     });
 
     it("should throw error when cell ID doesn't exist for position", async () => {
@@ -698,52 +712,6 @@ describe("EditNotebookTool", () => {
     });
   });
 
-  describe("validation", () => {
-    it("should validate cell exists in multi-column notebook", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "x = 1" },
-          [cellId2]: { code: "x = 2" },
-          [cellId3]: { code: "x = 3" },
-        },
-      });
-      // Create multi-column layout
-      notebook.cellIds = MultiColumn.from([[cellId1, cellId2], [cellId3]]);
-      store.set(notebookAtom, notebook);
-
-      // Should not throw for cells in different columns
-      const editorView = createMockEditorView("x = 1");
-      notebook.cellHandles[cellId1] = { current: { editorView } } as never;
-      notebook.cellHandles[cellId3] = { current: { editorView } } as never;
-
-      await expect(
-        tool.handler(
-          {
-            edit: {
-              type: "update_cell",
-              position: { cellId: cellId1 },
-              code: "y = 1",
-            },
-          },
-          toolContext as never,
-        ),
-      ).resolves.toBeDefined();
-
-      await expect(
-        tool.handler(
-          {
-            edit: {
-              type: "update_cell",
-              position: { cellId: cellId3 },
-              code: "y = 3",
-            },
-          },
-          toolContext as never,
-        ),
-      ).resolves.toBeDefined();
-    });
-  });
-
   describe("return value", () => {
     it("should return success status with next steps", async () => {
       const notebook = MockNotebook.notebookState({

package/src/core/ai/tools/edit-notebook-tool.ts

@@ -7,7 +7,11 @@ import { type NotebookState, notebookAtom } from "@/core/cells/cells";
 import type { CreateCellOptions } from "@/core/cells/actions";
 import { CellId } from "@/core/cells/ids";
 import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
-import type { CellColumnId } from "@/utils/id-tree";
+import {
+  type CellColumnId,
+  type Layout,
+  layoutAtom,
+} from "@/core/cells/layout";
 import { stagedAICellsAtom } from "../staged-cells";
 import {
   type AiTool,
@@ -201,7 +205,10 @@ export class EditNotebookTool implements AiTool<
             }
 
             placement = {
-              endOf: this.getColumnId(position.columnIndex, notebook),
+              endOf: this.getColumnId(
+                position.columnIndex,
+                store.get(layoutAtom),
+              ),
             };
             break;
 
@@ -259,8 +266,7 @@ export class EditNotebookTool implements AiTool<
   };
 
   private validateCellIdExists(cellId: CellId, notebook: NotebookState) {
-    const cellIds = notebook.cellIds;
-    if (!cellIds.getColumns().some((column) => column.idSet.has(cellId))) {
+    if (!notebook.cellIds.includes(cellId)) {
       throw new ToolExecutionError("Cell not found", {
         code: "CELL_NOT_FOUND",
         isRetryable: false,
@@ -269,12 +275,8 @@ export class EditNotebookTool implements AiTool<
     }
   }
 
-  private getColumnId(
-    columnIndex: number,
-    notebook: NotebookState,
-  ): CellColumnId {
-    const cellIds = notebook.cellIds;
-    const columns = cellIds.getColumns();
+  private getColumnId(columnIndex: number, layout: Layout): CellColumnId {
+    const { columns } = layout;
 
     if (columnIndex < 0 || columnIndex >= columns.length) {
       throw new ToolExecutionError("Column index is out of range", {

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

@@ -5,7 +5,6 @@ import { afterEach, describe, expect, it, vi } from "vitest";
 import { cellId } from "@/__tests__/branded";
 import { createStore } from "jotai";
 import { cell, CONFIG } from "@/core/notebook/document/__tests__/fixtures";
-import type { CellIndex } from "@/utils/id-tree";
 import {
   addSetupCellAtom,
   collapseAllCellsAtom,
@@ -17,9 +16,11 @@ import {
   focusCellAtom,
   moveToNextCellAtom,
   showCellIfHiddenAtom,
+  splitCellAtom,
   undoDeleteCellAtom,
+  undoSplitCellAtom,
 } from "../actions";
-import { moveCellsRelativeToAtom } from "../arrange";
+import { moveCellAtom, moveCellsRelativeToAtom } from "../arrange";
 import {
   cellDataAtom,
   cellHandleAtom,
@@ -95,7 +96,7 @@ describe("createCellAtom", () => {
   });
 
   it("names the column of a cell placed first in a later column", () => {
-    const { store, notebook, cellIn } = setup({
+    const { store, cellIn } = setup({
       c2: cell("c2", "y = 2", { config: { ...CONFIG, column: 1 } }),
     });
 
@@ -103,12 +104,14 @@ describe("createCellAtom", () => {
 
     expect(cellIn(id as string)?.config.column).toBe(1);
     expect(cellIn(c2)?.config.column).toBeNull();
-    expect(notebook().cellIds.atOrThrow(1).inOrderIds).toEqual([id, c2]);
-    expect(notebook().cellIds.atOrThrow(0).inOrderIds).toEqual([c1]);
+    expect(store.get(layoutAtom).columns.map((c) => c.ids)).toEqual([
+      [c1],
+      [id, c2],
+    ]);
   });
 
   it("places a cell in the column of its anchor, whatever config it carries", () => {
-    const { store, order, notebook, cellIn } = setup({
+    const { store, order, cellIn } = setup({
       c2: cell("c2", "y = 2", { config: { ...CONFIG, column: 1 } }),
     });
 
@@ -124,21 +127,23 @@ describe("createCellAtom", () => {
       hide_code: true,
       column: null,
     });
-    expect(notebook().cellIds.atOrThrow(1).inOrderIds).toEqual([c2, id]);
+    expect(store.get(layoutAtom).columns[1].ids).toEqual([c2, id]);
   });
 
-  it("names the column when appending to an empty column", () => {
-    const { store, notebook, cellIn } = setup();
-    store.set(notebookAtom, (state) => ({
-      ...state,
-      cellIds: state.cellIds.addColumn(state.cellIds.atOrThrow(0).id),
-    }));
-    const empty = notebook().cellIds.atOrThrow(1).id;
+  it("fills an emptied column, which keeps its id", () => {
+    const { store, columns, columnIds, cellIn } = setup({
+      c3: cell("c3", "z = 3", { config: { ...CONFIG, column: 1 } }),
+    });
+    const [first, second] = columnIds();
+    // Moving its only cell out leaves the column on the page, to add to.
+    store.set(moveCellAtom, { cellId: c3, direction: "left" });
+    expect(columns()).toEqual([[c3, c1, c2], []]);
 
-    const id = store.set(createCellAtom, { endOf: empty });
+    const id = store.set(createCellAtom, { endOf: second });
 
     expect(cellIn(id as string)?.config.column).toBe(1);
-    expect(notebook().cellIds.atOrThrow(1).inOrderIds).toEqual([id]);
+    expect(columns()).toEqual([[c3, c1, c2], [id]]);
+    expect(columnIds()).toEqual([first, second]);
   });
 
   it("creates nothing when a cell already has the code", () => {
@@ -161,6 +166,40 @@ describe("createCellAtom", () => {
   });
 });
 
+describe("splitCellAtom", () => {
+  it("splits a cell at the cursor and undoes it", () => {
+    const { store, order, notebook, cellIn, mountEditor } = setup({
+      c1: cell("c1", "line1\nline2"),
+    });
+    const view = mountEditor(c1);
+    view.dispatch({ selection: { anchor: 5 } });
+
+    const id = store.set(splitCellAtom, { cellId: c1 }) as CellId;
+
+    expect(order()).toEqual([c1, id, c2]);
+    expect(cellIn(c1)?.code).toBe("line1");
+    expect(cellIn(id)?.code).toBe("line2");
+    expect(view.state.doc.toString()).toBe("line1");
+    expect(notebook().scrollKey).toBe(id);
+
+    store.set(undoSplitCellAtom, {
+      cellId: c1,
+      newCellId: id,
+      snapshot: "line1\nline2",
+    });
+
+    expect(order()).toEqual([c1, c2]);
+    expect(cellIn(c1)?.code).toBe("line1\nline2");
+    expect(view.state.doc.toString()).toBe("line1\nline2");
+  });
+
+  it("does nothing for a cell without an editor", () => {
+    const { store, order } = setup();
+    expect(store.set(splitCellAtom, { cellId: c1 })).toBeNull();
+    expect(order()).toEqual([c1, c2]);
+  });
+});
+
 describe("moveToNextCellAtom", () => {
   it("focuses the neighbouring cell within the column", () => {
     const { store, order } = setup();
@@ -218,14 +257,14 @@ describe("addSetupCellAtom", () => {
   });
 
   it("places the setup cell in an empty first column", () => {
-    const { store, notebook, cellIn } = setup({
+    const { store, columns, notebook, cellIn } = setup({
       c1: cell(c1, "x = 1", { config: { ...CONFIG, column: 1 } }),
     });
-    expect(notebook().cellIds.atOrThrow(0).inOrderIds).toEqual([]);
+    expect(columns()[0]).toEqual([]);
 
     store.set(addSetupCellAtom, { code: "import marimo as mo" });
 
-    expect(notebook().cellIds.atOrThrow(0).inOrderIds).toEqual([SETUP_CELL_ID]);
+    expect(columns()[0]).toEqual([SETUP_CELL_ID]);
     expect(cellIn(SETUP_CELL_ID)?.config.column).toBe(0);
     expect(notebook().scrollKey).toBe(SETUP_CELL_ID);
   });
@@ -297,16 +336,16 @@ describe("deleteCellsAtom", () => {
   });
 
   it("names the column of a cell that followed a deleted column start", () => {
-    const { store, notebook, cellIn } = setup({
+    const { store, cellIn } = setup({
       c2: cell("c2", "y = 2", { config: { ...CONFIG, column: 1 } }),
       c3: cell("c3", "z = 3"),
     });
-    expect(notebook().cellIds.atOrThrow(1).inOrderIds).toEqual([c2, c3]);
+    expect(store.get(layoutAtom).columns[1].ids).toEqual([c2, c3]);
 
     store.set(deleteCellsAtom, { cellIds: [c2] });
 
     expect(cellIn(c3)?.config.column).toBe(1);
-    expect(notebook().cellIds.atOrThrow(1).inOrderIds).toEqual([c3]);
+    expect(store.get(layoutAtom).columns[1].ids).toEqual([c3]);
   });
 
   it("forgets a deleted cell's editor-only state and its atoms", () => {
@@ -386,14 +425,14 @@ describe("undoDeleteCellAtom", () => {
     const { store, order, notebook } = setup({
       c3: cell("c3", "z = 3"),
     });
-    const column = notebook().cellIds.atOrThrow(0);
+    const column = store.get(layoutAtom).columns[0];
     store.set(moveCellsRelativeToAtom, {
       cellIds: [c1, c2],
       targetCellId: c3,
       position: "after",
       previousPlacements: [
-        { columnId: column.id, index: 0 as CellIndex },
-        { columnId: column.id, index: 1 as CellIndex },
+        { columnId: column.id, index: 0 },
+        { columnId: column.id, index: 1 },
       ],
     });
     expect(order()).toEqual([c3, c1, c2]);