@marimo-team/islands 0.25.3-dev26 → 0.25.3-dev261

package/src/components/editor/navigation/focus-utils.ts

@@ -3,6 +3,7 @@
 import type { createStore } from "jotai";
 import { getCellEditorView } from "@/core/cells/cells";
 import { type CellId, HTMLCellId } from "@/core/cells/ids";
+import { scrollCellIntoView } from "@/core/cells/scrollCellIntoView";
 import { Logger } from "@/utils/Logger";
 
 export function focusCellEditor(
@@ -19,10 +20,10 @@ export function focusCellEditor(
   }
 }
 
-export function focusCell(cellId: CellId): void {
+export function focusCell(cellId: CellId, options?: FocusOptions): void {
   const element = document.getElementById(HTMLCellId.create(cellId));
   if (element) {
-    tryFocus(element);
+    tryFocus(element, options);
   } else {
     Logger.warn(`[CellFocusManager] focusCell: element not found: ${cellId}`);
   }
@@ -38,6 +39,24 @@ export function raf2(callback: () => void): void {
   });
 }
 
+/**
+ * Keep the focused cell in view after a layout shift, such as its code being
+ * hidden again.
+ */
+export function keepFocusedCellInView(): void {
+  raf2(() => {
+    const activeElement = document.activeElement;
+    if (!activeElement) {
+      return;
+    }
+    const focusedCell = HTMLCellId.findElementThroughShadowDOMs(activeElement);
+    if (!focusedCell) {
+      return;
+    }
+    scrollCellIntoView(HTMLCellId.parse(focusedCell.id));
+  });
+}
+
 /**
  * Checks if the cell is focused at the top level.
  */
@@ -48,9 +67,9 @@ export function isAnyCellFocused(): boolean {
   );
 }
 
-export function tryFocus(dom: HTMLElement) {
+export function tryFocus(dom: HTMLElement, options?: FocusOptions) {
   try {
-    dom.focus();
+    dom.focus(options);
   } catch {
     Logger.warn("[CellFocusManager] element may not be focusable", dom);
   }

package/src/components/editor/navigation/multi-cell-action-toolbar.tsx

@@ -30,18 +30,15 @@ import {
   DropdownMenuSeparator,
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
-import {
-  getCellEditorView,
-  hasOnlyOneCellAtom,
-  notebookAtom,
-  useCellActions,
-} from "@/core/cells/cells";
+import { getCellEditorView, hasOnlyOneCellAtom } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
+import { columnOf, layoutAtom } from "@/core/cells/layout";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import type { CellId } from "@/core/cells/ids";
 import { usePendingDeleteService } from "@/core/cells/pending-delete-service";
 import { formatEditorViews } from "@/core/codemirror/format";
 import { userConfigAtom } from "@/core/config/config";
 import type { HotkeyAction } from "@/core/hotkeys/hotkeys";
-import { useRequestClient } from "@/core/network/requests";
 import type { CellConfig } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
 import { useEventListener } from "@/hooks/useEventListener";
@@ -112,20 +109,15 @@ const CellStateDropdown: React.FC<{
 };
 
 export function useMultiCellActionButtons(cellIds: CellId[]) {
-  const {
-    updateCellConfig,
-    moveCell,
-    clearCellOutput,
-    sendToTop,
-    sendToBottom,
-  } = useCellActions();
+  const { moveCell, clearCellOutput, sendToTop, sendToBottom } =
+    useCellActions();
+  const edit = useNotebookDocumentEdit();
   const deleteCell = useDeleteManyCellsCallback();
   const hasOnlyOneCell = useAtomValue(hasOnlyOneCellAtom);
   const selectionActions = useCellSelectionActions();
   const runCells = useRunCells();
   const pendingDeleteService = usePendingDeleteService();
   const userConfig = useAtomValue(userConfigAtom);
-  const { saveCellConfig } = useRequestClient();
 
   const selectedCount = cellIds.length;
 
@@ -147,10 +139,8 @@ export function useMultiCellActionButtons(cellIds: CellId[]) {
       /// If moving down, make sure the last cell is not at the bottom of the notebook
       if (direction === "down") {
         const lastCellId = cellIds[cellIds.length - 1];
-        const notebook = store.get(notebookAtom);
-        const isLast =
-          notebook.cellIds.findWithId(lastCellId).last() === lastCellId;
-        if (isLast) {
+        const layout = store.get(layoutAtom);
+        if (columnOf(layout, lastCellId).visible.at(-1) === lastCellId) {
           return;
         }
       }
@@ -158,10 +148,8 @@ export function useMultiCellActionButtons(cellIds: CellId[]) {
       // If moving up, make sure the first cell is not at the top of the notebook
       if (direction === "up") {
         const firstCellId = cellIds[0];
-        const notebook = store.get(notebookAtom);
-        const isFirst =
-          notebook.cellIds.findWithId(firstCellId).first() === firstCellId;
-        if (isFirst) {
+        const layout = store.get(layoutAtom);
+        if (columnOf(layout, firstCellId).visible[0] === firstCellId) {
           return;
         }
       }
@@ -205,16 +193,12 @@ export function useMultiCellActionButtons(cellIds: CellId[]) {
   });
 
   const toggleSelectedCellsProperty = useEvent(
-    async (cellIds: CellId[], property: keyof CellConfig, value: boolean) => {
-      const configs: Record<CellId, Partial<CellConfig>> = {};
-      cellIds.forEach((cellId) => {
-        configs[cellId] = { [property]: value };
-      });
-
-      await saveCellConfig({ configs });
-
-      cellIds.forEach((cellId) => {
-        updateCellConfig({ cellId, config: configs[cellId] });
+    (cellIds: CellId[], property: keyof CellConfig, value: boolean) => {
+      edit((tx) => {
+        for (const cellId of cellIds) {
+          const config: Partial<CellConfig> = { [property]: value };
+          tx.updateConfig(cellId, config);
+        }
       });
     },
   );

package/src/components/editor/navigation/navigation.ts

@@ -13,7 +13,16 @@ import { mergeProps, useFocusWithin, useKeyboard } from "react-aria";
 import { DATA_FOR_CELL_ID } from "@/components/data-table/cell-utils";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { cellIdsAtom, notebookAtom, useCellActions } from "@/core/cells/cells";
+import { notebookAtom } from "@/core/cells/cells";
+import { hideShownCodeAtom, showHiddenCodeAtom } from "@/core/cells/editor";
+import {
+  columnOf,
+  type LayoutColumn,
+  layoutAtom,
+  visibleNeighbor,
+} from "@/core/cells/layout";
+import { useCellActions } from "@/core/cells/actions";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import { useCellFocusActions } from "@/core/cells/focus";
 import type { CellId } from "@/core/cells/ids";
 import { HTMLCellId } from "@/core/cells/ids";
@@ -36,21 +45,23 @@ import {
 } from "@/core/config/config";
 import type { HotkeyAction } from "@/core/hotkeys/hotkeys";
 import { parseShortcut } from "@/core/hotkeys/shortcuts";
-import { useRequestClient } from "@/core/network/requests";
 import { useSaveNotebook } from "@/core/saving/save-component";
 import { Events } from "@/utils/events";
-import type { CollapsibleTree } from "@/utils/id-tree";
 import type { CellActionsDropdownHandle } from "../cell/cell-actions";
 import { useDeleteManyCellsCallback } from "../cell/useDeleteCell";
 import { useRunCells } from "../cell/useRunCells";
 import { useCellClipboard } from "./clipboard";
-import { focusCell, focusCellEditor, raf2 } from "./focus-utils";
+import {
+  focusCell,
+  focusCellEditor,
+  keepFocusedCellInView,
+  raf2,
+} from "./focus-utils";
 import {
   getSelectedCells,
   useCellSelectionActions,
   useIsCellSelected,
 } from "./selection";
-import { useTemporarilyShownCodeActions } from "./state";
 import { handleVimKeybinding } from "./vim-bindings";
 
 interface HotkeyHandler {
@@ -104,8 +115,7 @@ function useCellFocusProps(
   editorView: React.RefObject<EditorView | null>,
 ) {
   const focusActions = useCellFocusActions();
-  const actions = useCellActions();
-  const temporarilyShownCodeActions = useTemporarilyShownCodeActions();
+  const hideShownCode = useSetAtom(hideShownCodeAtom);
 
   // This occurs at the cell level and descedants.
   const { focusWithinProps } = useFocusWithin({
@@ -130,9 +140,11 @@ function useCellFocusProps(
         return;
       }
 
-      // On blur, hide the code if it was temporarily shown.
-      temporarilyShownCodeActions.remove(cellId);
-      actions.markTouched({ cellId });
+      // On blur, hide the code if it was showing for now. That shifts the
+      // layout, so the cell that took the focus is kept in view.
+      if (hideShownCode({ cellId })) {
+        keepFocusedCellInView();
+      }
       focusActions.blurCell();
       // Close signature help when clicking outside the cell
       if (editorView.current) {
@@ -189,11 +201,12 @@ export function useCellNavigationProps(
   },
 ) {
   const { saveOrNameNotebook } = useSaveNotebook();
-  const { saveCellConfig } = useRequestClient();
+  const edit = useNotebookDocumentEdit();
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
   const actions = useCellActions();
+  const { createCell, moveToNextCell } = useCellActions();
   const store = useStore();
-  const temporarilyShownCodeActions = useTemporarilyShownCodeActions();
+  const showHiddenCode = useSetAtom(showHiddenCodeAtom);
   const runCells = useRunCells();
   const keymapPreset = useAtomValue(keymapPresetAtom);
   const { copyCells, pasteAtCell, cutCells } = useCellClipboard();
@@ -266,12 +279,13 @@ export function useCellNavigationProps(
         // Move left across columns
         ArrowLeft: () => {
           if (canMoveX) {
-            const notebook = store.get(notebookAtom);
-            const column = notebook.cellIds.findWithId(cellId);
-            const columnIndex = notebook.cellIds.indexOf(column);
-            const leftColumn = notebook.cellIds.at(columnIndex - 1);
+            const layout = store.get(layoutAtom);
+            const columnIndex = layout.columns.indexOf(
+              columnOf(layout, cellId),
+            );
+            const leftColumn = layout.columns[columnIndex - 1];
 
-            if (leftColumn && leftColumn.length > 0) {
+            if (leftColumn && leftColumn.visible.length > 0) {
               const leftCellId = findClosestAdjacentCell(cellId, leftColumn);
               actions.focusCell({ cellId: leftCellId, where: "exact" });
 
@@ -284,12 +298,13 @@ export function useCellNavigationProps(
         // Move right across columns
         ArrowRight: () => {
           if (canMoveX) {
-            const notebook = store.get(notebookAtom);
-            const column = notebook.cellIds.findWithId(cellId);
-            const columnIndex = notebook.cellIds.indexOf(column);
-            const rightColumn = notebook.cellIds.at(columnIndex + 1);
+            const layout = store.get(layoutAtom);
+            const columnIndex = layout.columns.indexOf(
+              columnOf(layout, cellId),
+            );
+            const rightColumn = layout.columns[columnIndex + 1];
 
-            if (rightColumn && rightColumn.length > 0) {
+            if (rightColumn && rightColumn.visible.length > 0) {
               const rightCellId = findClosestAdjacentCell(cellId, rightColumn);
               actions.focusCell({ cellId: rightCellId, where: "exact" });
 
@@ -302,12 +317,12 @@ export function useCellNavigationProps(
         // Select up
         "Shift+ArrowUp": () => {
           // Select self
-          const allCellIds = store.get(cellIdsAtom);
-          selectionActions.extend({ cellId, allCellIds });
+          const layout = store.get(layoutAtom);
+          selectionActions.extend({ cellId, layout });
           // Select to where focus is going
-          const beforeCellId = allCellIds.findWithId(cellId).before(cellId);
+          const beforeCellId = visibleNeighbor(layout, cellId, -1);
           if (beforeCellId) {
-            selectionActions.extend({ cellId: beforeCellId, allCellIds });
+            selectionActions.extend({ cellId: beforeCellId, layout });
           }
           // Focus the cell
           actions.focusCell({ cellId, where: "before" });
@@ -316,12 +331,12 @@ export function useCellNavigationProps(
         // Select down
         "Shift+ArrowDown": () => {
           // Select self
-          const allCellIds = store.get(cellIdsAtom);
-          selectionActions.extend({ cellId, allCellIds });
+          const layout = store.get(layoutAtom);
+          selectionActions.extend({ cellId, layout });
           // Select to where focus is going
-          const afterCellId = allCellIds.findWithId(cellId).after(cellId);
+          const afterCellId = visibleNeighbor(layout, cellId, 1);
           if (afterCellId) {
-            selectionActions.extend({ cellId: afterCellId, allCellIds });
+            selectionActions.extend({ cellId: afterCellId, layout });
           }
           // Focus the cell
           actions.focusCell({ cellId, where: "after" });
@@ -343,7 +358,7 @@ export function useCellNavigationProps(
         },
         // Enter will focus the cell editor.
         Enter: () => {
-          temporarilyShownCodeActions.add(cellId);
+          showHiddenCode({ cellId });
           focusCellEditor(store, cellId);
           selectionActions.clear();
           return true;
@@ -376,30 +391,28 @@ export function useCellNavigationProps(
         "cell.runAndNewBelow": addSingleHandler((cellIds) => {
           runCells(cellIds);
           const lastCellId = cellIds[cellIds.length - 1];
-          actions.moveToNextCell({ cellId: lastCellId, before: false });
+          moveToNextCell({ cellId: lastCellId, before: false });
           return true;
         }),
         "cell.runAndNewAbove": addSingleHandler((cellIds) => {
           runCells(cellIds);
           const firstCellId = cellIds[0];
-          actions.moveToNextCell({ cellId: firstCellId, before: true });
+          moveToNextCell({ cellId: firstCellId, before: true });
           return true;
         }),
         "cell.createAbove": (cellId) => {
-          actions.createNewCell({ cellId, before: true });
+          createCell({ before: cellId });
           return true;
         },
         "cell.createBelow": (cellId) => {
-          actions.createNewCell({ cellId, before: false });
+          createCell({ after: cellId });
           return true;
         },
         "cell.moveUp": addSingleHandler((cellIds) => {
           // If moving up, make sure the first cell is not at the top of the notebook
           const firstCellId = cellIds[0];
-          const notebook = store.get(notebookAtom);
-          const isFirst =
-            notebook.cellIds.findWithId(firstCellId).first() === firstCellId;
-          if (isFirst) {
+          const layout = store.get(layoutAtom);
+          if (columnOf(layout, firstCellId).visible[0] === firstCellId) {
             return false;
           }
 
@@ -411,10 +424,8 @@ export function useCellNavigationProps(
         "cell.moveDown": addSingleHandler((cellIds) => {
           // If moving down, make sure the last cell is not at the bottom of the notebook
           const lastCellId = cellIds[cellIds.length - 1];
-          const notebook = store.get(notebookAtom);
-          const isLast =
-            notebook.cellIds.findWithId(lastCellId).last() === lastCellId;
-          if (isLast) {
+          const layout = store.get(layoutAtom);
+          if (columnOf(layout, lastCellId).visible.at(-1) === lastCellId) {
             return false;
           }
 
@@ -453,20 +464,12 @@ export function useCellNavigationProps(
             (config) => config?.hide_code,
           );
 
-          // Fire-and-forget
-          void saveCellConfig({
-            configs: Object.fromEntries(
-              cellIds.map((cellId) => [cellId, { hide_code: nextHideCode }]),
-            ),
+          edit((tx) => {
+            for (const cellId of cellIds) {
+              tx.updateConfig(cellId, { hide_code: nextHideCode });
+            }
           });
 
-          for (const cellId of cellIds) {
-            actions.updateCellConfig({
-              cellId,
-              config: { hide_code: nextHideCode },
-            });
-          }
-
           // Only focus if it is a single cell
           if (cellIds.length === 1) {
             const cellId = cellIds[0];
@@ -543,14 +546,14 @@ export function useCellNavigationProps(
           if (Events.hasModifier(evt)) {
             return false;
           }
-          actions.createNewCell({ cellId, before: true, autoFocus: true });
+          createCell({ before: cellId });
           return true;
         },
         "command.createCellAfter": (cellId) => {
           if (Events.hasModifier(evt)) {
             return false;
           }
-          actions.createNewCell({ cellId, before: false, autoFocus: true });
+          createCell({ after: cellId });
           return true;
         },
         "command.createSqlCellAfter": (cellId) => {
@@ -559,14 +562,12 @@ export function useCellNavigationProps(
           }
           maybeAddMarimoImport({
             autoInstantiate: true,
-            createNewCell: actions.createNewCell,
+            createCell: createCell,
             fromCellId: cellId,
             before: true,
           });
-          actions.createNewCell({
-            cellId,
-            before: false,
-            autoFocus: true,
+          createCell({
+            after: cellId,
             code: LanguageAdapters.sql.defaultCode,
           });
           return true;
@@ -635,11 +636,11 @@ export function useCellNavigationProps(
             return true;
           },
           o: () => {
-            actions.createNewCell({ cellId, before: false, autoFocus: true });
+            createCell({ after: cellId });
             return true;
           },
           "shift+o": () => {
-            actions.createNewCell({ cellId, before: true, autoFocus: true });
+            createCell({ before: cellId });
             return true;
           },
           u: () => {
@@ -701,7 +702,7 @@ export function useCellEditorNavigationProps(
   cellId: CellId,
   editorView: React.RefObject<EditorView | null>,
 ) {
-  const temporarilyShownCodeActions = useTemporarilyShownCodeActions();
+  const hideShownCode = useSetAtom(hideShownCodeAtom);
   const keymapPreset = useAtomValue(keymapPresetAtom);
   const hotkeys = useAtomValue(hotkeysAtom);
 
@@ -711,7 +712,7 @@ export function useCellEditorNavigationProps(
   }, [hotkeys]);
 
   const exitToCommandMode = () => {
-    temporarilyShownCodeActions.remove(cellId);
+    hideShownCode({ cellId });
     focusCell(cellId);
     // Scroll to cell in case it is not in view because of layout shifts.
     raf2(() => {
@@ -778,18 +779,18 @@ export function useCellEditorNavigationProps(
 
 function findClosestAdjacentCell(
   currentCellId: CellId,
-  adjacentColumn: CollapsibleTree<CellId>,
+  adjacentColumn: LayoutColumn,
 ): CellId {
   const current = document.getElementById(HTMLCellId.create(currentCellId));
 
   if (!current) {
     // fallback to first
-    return adjacentColumn.first();
+    return adjacentColumn.visible[0];
   }
 
   const currentRect = current.getBoundingClientRect();
   // first to to either overlap or contain the current cell's top edge
-  for (const candidateId of adjacentColumn.topLevelIds) {
+  for (const candidateId of adjacentColumn.visible) {
     const candidate = document.getElementById(HTMLCellId.create(candidateId));
     if (candidate) {
       const candidateRect = candidate.getBoundingClientRect();
@@ -803,7 +804,7 @@ function findClosestAdjacentCell(
   }
 
   // no aligned cells (column beyond other), jump to last element
-  return adjacentColumn.last();
+  return adjacentColumn.visible[adjacentColumn.visible.length - 1];
 }
 
 export function closeSignatureHelp(view: EditorView) {

package/src/components/editor/navigation/selection.ts

@@ -2,8 +2,8 @@
 import { atom, type createStore, useAtomValue } from "jotai";
 import { useMemo } from "react";
 import type { CellId } from "@/core/cells/ids";
+import type { Layout } from "@/core/cells/layout";
 import { createReducerAndAtoms } from "@/utils/createReducer";
-import type { MultiColumn } from "@/utils/id-tree";
 
 export interface CellSelectionState {
   selectionStart: CellId | null;
@@ -34,43 +34,44 @@ const {
     selected: new Set([payload.cellId]),
   }),
 
+  /**
+   * Extend the selection from where it started to `cellId`, over the cells
+   * the page shows between them. Cells a fold hides are stepped over, as
+   * the focus steps over them.
+   */
   extend: (
     state: CellSelectionState,
-    payload: {
-      cellId: CellId;
-      allCellIds: MultiColumn<CellId>;
-    },
+    payload: { cellId: CellId; layout: Layout },
   ): CellSelectionState => {
-    if (!state.selectionStart) {
-      // fallback to single select
-      return {
-        selectionStart: payload.cellId,
-        selectionEnd: payload.cellId,
-        selected: new Set([payload.cellId]),
-      };
-    }
-    const { cellId, allCellIds } = payload;
-
-    try {
-      const column = allCellIds.findWithId(state.selectionStart);
-      const startIdx = column.indexOfOrThrow(state.selectionStart);
-      const endIdx = column.indexOfOrThrow(cellId);
-      const [from, to] =
-        startIdx < endIdx ? [startIdx, endIdx] : [endIdx, startIdx];
-      const selected = column.slice(from, to + 1);
-      return {
-        selectionStart: state.selectionStart,
-        selectionEnd: cellId,
-        selected: new Set(selected),
-      };
-    } catch {
-      // fallback to single select
+    const { cellId, layout } = payload;
+    const start = state.selectionStart;
+    const column =
+      start === null
+        ? undefined
+        : layout.columns.find((column) => column.visible.includes(start));
+    const startIdx =
+      start === null ? -1 : (column?.visible.indexOf(start) ?? -1);
+    const endIdx = column?.visible.indexOf(cellId) ?? -1;
+    if (
+      start === null ||
+      column === undefined ||
+      startIdx === -1 ||
+      endIdx === -1
+    ) {
+      // Nothing to extend from, or the cells are not on the page together.
       return {
         selectionStart: cellId,
         selectionEnd: cellId,
         selected: new Set([cellId]),
       };
     }
+    const [from, to] =
+      startIdx < endIdx ? [startIdx, endIdx] : [endIdx, startIdx];
+    return {
+      selectionStart: start,
+      selectionEnd: cellId,
+      selected: new Set(column.visible.slice(from, to + 1)),
+    };
   },
 
   clear: (state: CellSelectionState): CellSelectionState => {

package/src/components/editor/notebook-banner.tsx

@@ -41,6 +41,7 @@ export const NotebookBanner: React.FC<Props> = ({ width }) => {
               {banner.title}
             </span>
             <Button
+              aria-label={`Dismiss ${banner.title}`}
               data-testid="remove-banner-button"
               variant="text"
               size="icon"