fix: setup cell can be focused / enter command mode

frontend/src/components/editor/__tests__/data-attributes.test.tsx

@@ -12,7 +12,7 @@ import type { UserConfig } from "@/core/config/config-schema";
 import type { OutputMessage } from "@/core/kernel/messages";
 import type { AppMode } from "@/core/mode";
 import { requestClientAtom } from "@/core/network/requests";
-import { Cell } from "../Cell";
+import { Cell } from "../notebook-cell";
 import { OutputArea } from "../Output";
 
 function createTestWrapper() {

frontend/src/components/editor/actions/useNotebookActions.tsx

@@ -102,7 +102,7 @@ export function useNotebookActions() {
     updateCellConfig,
     undoDeleteCell,
     clearAllCellOutputs,
-    upsertSetupCell,
+    addSetupCellIfDoesntExist,
     collapseAllCells,
     expandAllCells,
   } = useCellActions();
@@ -401,9 +401,7 @@ export function useNotebookActions() {
       icon: <DiamondPlusIcon size={14} strokeWidth={1.5} />,
       label: "Add setup cell",
       handle: () => {
-        upsertSetupCell({
-          code: "# Initialization code that runs before all other cells",
-        });
+        addSetupCellIfDoesntExist({});
       },
     },
     {

frontend/src/components/editor/cell/CreateCellButton.tsx

@@ -1,5 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import { DatabaseIcon, PlusIcon } from "lucide-react";
+import { DatabaseIcon, DiamondPlusIcon, PlusIcon } from "lucide-react";
 import { Button } from "@/components/editor/inputs/Inputs";
 import {
   ContextMenu,
@@ -69,7 +69,7 @@ const CreateCellButtonContextMenu = (props: {
   children: React.ReactNode;
 }) => {
   const { children, onClick } = props;
-  const { createNewCell } = useCellActions();
+  const { createNewCell, addSetupCellIfDoesntExist } = useCellActions();
 
   if (!onClick) {
     return children;
@@ -125,6 +125,18 @@ const CreateCellButtonContextMenu = (props: {
           </div>
           SQL cell
         </ContextMenuItem>
+        <ContextMenuItem
+          key="setup"
+          onSelect={(evt) => {
+            evt.stopPropagation();
+            addSetupCellIfDoesntExist({});
+          }}
+        >
+          <div className="mr-3 text-muted-foreground">
+            <DiamondPlusIcon size={13} strokeWidth={1.5} />
+          </div>
+          Setup cell
+        </ContextMenuItem>
       </ContextMenuContent>
     </ContextMenu>
   );

frontend/src/components/editor/Cell.tsx → frontend/src/components/editor/notebook-cell.tsx

@@ -1051,12 +1051,16 @@ const SetupCellComponent = ({
           data-status={cellRuntime.status}
           ref={cellRef}
           {...mergeProps(navigationProps, {
-            className,
+            className: cn(
+              className,
+              "focus:ring-1 focus:ring-(--blue-7) focus:ring-offset-0",
+            ),
             onBlur: closeCompletionHandler,
             onKeyDown: resumeCompletionHandler,
           })}
           {...cellDomProps(cellId, cellData.name)}
           title={renderCellTitle()}
+          tabIndex={-1}
           data-setup-cell={true}
         >
           <div className={cn("tray")} data-hidden={!isCellCodeShown}>

frontend/src/components/editor/renderers/CellArray.tsx → frontend/src/components/editor/renderers/cell-array.tsx

@@ -14,7 +14,7 @@ import {
 } from "lucide-react";
 import { useEffect } from "react";
 import { StartupLogsAlert } from "@/components/editor/alerts/startup-logs-alert";
-import { Cell } from "@/components/editor/Cell";
+import { Cell } from "@/components/editor/notebook-cell";
 import { PackageAlert } from "@/components/editor/package-alert";
 import { SortableCellsProvider } from "@/components/sort/SortableCellsProvider";
 import { Button } from "@/components/ui/button";

frontend/src/core/cells/__tests__/cells.test.ts

@@ -12,7 +12,7 @@ import {
   it,
   vi,
 } from "vitest";
-import type { CellHandle } from "@/components/editor/Cell";
+import type { CellHandle } from "@/components/editor/notebook-cell";
 import { CellId } from "@/core/cells/ids";
 import { foldAllBulk, unfoldAllBulk } from "@/core/codemirror/editing/commands";
 import { adaptiveLanguageConfiguration } from "@/core/codemirror/language/extension";
@@ -2094,9 +2094,9 @@ describe("cell reducer", () => {
     `);
   });
 
-  it("can create and update a setup cell", () => {
+  it("can create and noop-update a setup cell", () => {
     // Create the setup cell
-    actions.upsertSetupCell({ code: "# Setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Setup code" });
 
     // Check that setup cell was created
     expect(state.cellData[SETUP_CELL_ID].id).toBe(SETUP_CELL_ID);
@@ -2106,17 +2106,17 @@ describe("cell reducer", () => {
     expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
 
     // Update the setup cell
-    actions.upsertSetupCell({ code: "# Updated setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Updated setup code" });
 
     // Check that the same setup cell was updated, not duplicated
-    expect(state.cellData[SETUP_CELL_ID].code).toBe("# Updated setup code");
+    expect(state.cellData[SETUP_CELL_ID].code).toBe("# Setup code");
     expect(state.cellData[SETUP_CELL_ID].edited).toBe(true);
     expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
   });
 
   it("can delete and undelete the setup cell", () => {
     // Create the setup cell
-    actions.upsertSetupCell({ code: "# Setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Setup code" });
 
     // Check that setup cell was created
     expect(state.cellData[SETUP_CELL_ID].id).toBe(SETUP_CELL_ID);

frontend/src/core/cells/cells.ts

@@ -5,7 +5,7 @@ import { type Atom, atom, useAtom, useAtomValue } from "jotai";
 import { atomFamily, selectAtom, splitAtom } from "jotai/utils";
 import { isEqual, zip } from "lodash-es";
 import { createRef, type ReducerWithoutAction } from "react";
-import type { CellHandle } from "@/components/editor/Cell";
+import type { CellHandle } from "@/components/editor/notebook-cell";
 import {
   type CellColumnId,
   type CellIndex,
@@ -1324,21 +1324,19 @@ const {
       cellRuntime: newCellRuntime,
     };
   },
-  upsertSetupCell: (state, action: { code: string }) => {
-    const { code } = action;
+  addSetupCellIfDoesntExist: (state, action: { code?: string }) => {
+    let { code } = action;
+    if (code == null) {
+      code = "# Initialization code that runs before all other cells";
+    }
 
     // First check if setup cell already exists
     if (SETUP_CELL_ID in state.cellData) {
-      // Update existing setup cell
-      return updateCellData({
-        state,
-        cellId: SETUP_CELL_ID,
-        cellReducer: (cell) => ({
-          ...cell,
-          code,
-          edited: code.trim() !== cell.lastCodeRun?.trim(),
-        }),
-      });
+      // Just focus on the existing setup cell
+      return {
+        ...state,
+        scrollKey: SETUP_CELL_ID,
+      };
     }
 
     return {
@@ -1365,6 +1363,7 @@ const {
         ...state.cellHandles,
         [SETUP_CELL_ID]: createRef(),
       },
+      scrollKey: SETUP_CELL_ID,
     };
   },
 });

frontend/src/core/cells/scrollCellIntoView.ts

@@ -1,10 +1,10 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import type { RefObject } from "react";
-import type { CellHandle } from "@/components/editor/Cell";
 import {
   isAnyCellFocused,
   tryFocus,
 } from "@/components/editor/navigation/focus-utils";
+import type { CellHandle } from "@/components/editor/notebook-cell";
 import { retryWithTimeout } from "@/utils/timeout";
 import { Logger } from "../../utils/Logger";
 import { goToVariableDefinition } from "../codemirror/go-to-definition/commands";

frontend/src/core/edit-app.tsx

@@ -15,7 +15,7 @@ import {
   useRunAllCells,
   useRunStaleCells,
 } from "../components/editor/cell/useRunCells";
-import { CellArray } from "../components/editor/renderers/CellArray";
+import { CellArray } from "../components/editor/renderers/cell-array";
 import { CellsRenderer } from "../components/editor/renderers/cells-renderer";
 import { useHotkey } from "../hooks/useHotkey";
 import {

frontend/src/stories/cell.stories.tsx

@@ -16,7 +16,7 @@ import type { CellConfig } from "@/core/network/types";
 import { WebSocketState } from "@/core/websocket/types";
 import { MultiColumn } from "@/utils/id-tree";
 import type { Milliseconds, Seconds } from "@/utils/time";
-import { Cell as EditorCell } from "../components/editor/Cell";
+import { Cell as EditorCell } from "../components/editor/notebook-cell";
 import { TooltipProvider } from "../components/ui/tooltip";
 import type { CellId } from "../core/cells/ids";
 

frontend/src/stories/layout/vertical/one-column.stories.tsx

@@ -13,7 +13,7 @@ import { resolveRequestClient } from "@/core/network/resolve";
 import { WebSocketState } from "@/core/websocket/types";
 import { MultiColumn } from "@/utils/id-tree";
 import type { Milliseconds, Seconds } from "@/utils/time";
-import { CellArray } from "../../../components/editor/renderers/CellArray";
+import { CellArray } from "../../../components/editor/renderers/cell-array";
 import { CellsRenderer } from "../../../components/editor/renderers/cells-renderer";
 import { TooltipProvider } from "../../../components/ui/tooltip";
 import type { CellId } from "../../../core/cells/ids";