fix add setup cell overwriting original cell on refresh (#7061)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->
Fixes #7052 . Adds new logic to detect whether a setup cell exists.

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [x] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

Author: Light2Dark

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

@@ -39,12 +39,8 @@ import { Switch } from "@/components/ui/switch";
 import { toast } from "@/components/ui/use-toast";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import {
-  hasOnlyOneCellAtom,
-  SETUP_CELL_ID,
-  useCellActions,
-} from "@/core/cells/cells";
-import type { CellId } from "@/core/cells/ids";
+import { hasOnlyOneCellAtom, useCellActions } from "@/core/cells/cells";
+import { type CellId, SETUP_CELL_ID } from "@/core/cells/ids";
 import type { CellData } from "@/core/cells/types";
 import { formatEditorViews } from "@/core/codemirror/format";
 import { toggleToLanguage } from "@/core/codemirror/language/commands";

frontend/src/components/editor/links/cell-link.tsx

@@ -1,13 +1,8 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import type { JSX } from "react"; /* Copyright 2024 Marimo. All rights reserved. */
-import {
-  SCRATCH_CELL_ID,
-  useCellActions,
-  useCellIds,
-  useCellNames,
-} from "@/core/cells/cells";
-import { type CellId, HTMLCellId } from "@/core/cells/ids";
+import { useCellActions, useCellIds, useCellNames } from "@/core/cells/cells";
+import { type CellId, HTMLCellId, SCRATCH_CELL_ID } from "@/core/cells/ids";
 import { displayCellName } from "@/core/cells/names";
 import { goToCellLine } from "@/core/codemirror/go-to-definition/utils";
 import { useFilename } from "@/core/saving/filename";

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

@@ -41,13 +41,12 @@ import type { Milliseconds, Seconds } from "@/utils/time";
 import {
   type CellActions,
   createUntouchedCellAtom,
-  SETUP_CELL_ID,
   useCellActions,
   useCellData,
   useCellHandle,
   useCellRuntime,
 } from "../../core/cells/cells";
-import type { CellId } from "../../core/cells/ids";
+import { type CellId, SETUP_CELL_ID } from "../../core/cells/ids";
 import { isUninstantiated } from "../../core/cells/utils";
 import type { UserConfig } from "../../core/config/config-schema";
 import {

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

@@ -20,6 +20,7 @@ import { SortableCellsProvider } from "@/components/sort/SortableCellsProvider";
 import { Button } from "@/components/ui/button";
 import { Tooltip } from "@/components/ui/tooltip";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
+import { SETUP_CELL_ID } from "@/core/cells/ids";
 import { LanguageAdapters } from "@/core/codemirror/language/LanguageAdapters";
 import { aiEnabledAtom } from "@/core/config/config";
 import { isConnectedAtom } from "@/core/network/connection";
@@ -30,7 +31,6 @@ import type { CellColumnId } from "@/utils/id-tree";
 import { invariant } from "@/utils/invariant";
 import {
   columnIdsAtom,
-  SETUP_CELL_ID,
   useCellActions,
   useCellIds,
   useScrollKey,

frontend/src/components/scratchpad/scratchpad.tsx

@@ -11,13 +11,9 @@ import {
 import type React from "react";
 import { Suspense, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
-import {
-  SCRATCH_CELL_ID,
-  useCellActions,
-  useNotebook,
-} from "@/core/cells/cells";
+import { useCellActions, useNotebook } from "@/core/cells/cells";
 import { useLastFocusedCellId } from "@/core/cells/focus";
-import { HTMLCellId } from "@/core/cells/ids";
+import { HTMLCellId, SCRATCH_CELL_ID } from "@/core/cells/ids";
 import { DEFAULT_CELL_NAME } from "@/core/cells/names";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
 import { useResolvedMarimoConfig } from "@/core/config/config";

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

@@ -13,7 +13,7 @@ import {
   vi,
 } from "vitest";
 import type { CellHandle } from "@/components/editor/notebook-cell";
-import { CellId } from "@/core/cells/ids";
+import { CellId, SETUP_CELL_ID } from "@/core/cells/ids";
 import { foldAllBulk, unfoldAllBulk } from "@/core/codemirror/editing/commands";
 import { adaptiveLanguageConfiguration } from "@/core/codemirror/language/extension";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
@@ -24,7 +24,6 @@ import {
   exportedForTesting,
   flattenTopLevelNotebookCells,
   type NotebookState,
-  SETUP_CELL_ID,
 } from "../cells";
 import {
   focusAndScrollCellIntoView,
@@ -2108,8 +2107,8 @@ describe("cell reducer", () => {
     // Update the setup cell
     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");
+    // Check that the setup cell did not change, since it already exists
+    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);
   });

frontend/src/core/cells/cells.ts

@@ -29,7 +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 { CellId } from "./ids";
+import { CellId, SCRATCH_CELL_ID, SETUP_CELL_ID } from "./ids";
 import { type CellLog, getCellLogsForMessage } from "./logs";
 import {
   focusAndScrollCellIntoView,
@@ -52,13 +52,6 @@ import {
   notebookQueueOrRunningCount,
 } from "./utils";
 
-export const SCRATCH_CELL_ID = "__scratch__" as CellId;
-export const SETUP_CELL_ID = "setup" as CellId;
-
-export function isSetupCell(cellId: CellId): boolean {
-  return cellId === SETUP_CELL_ID;
-}
-
 /**
  * The state of the notebook.
  */
@@ -1335,8 +1328,7 @@ const {
       code = "# Initialization code that runs before all other cells";
     }
 
-    // First check if setup cell already exists
-    if (SETUP_CELL_ID in state.cellIds) {
+    if (state.cellIds.setupCellExists()) {
       // Just focus on the existing setup cell
       return {
         ...state,

frontend/src/core/cells/focus.ts

@@ -4,8 +4,9 @@ import type { EditorView } from "@codemirror/view";
 import { atom, useAtomValue } from "jotai";
 import { createReducerAndAtoms } from "@/utils/createReducer";
 import type { CellConfig, RuntimeState } from "../network/types";
-import { type NotebookState, notebookAtom, SCRATCH_CELL_ID } from "./cells";
+import { type NotebookState, notebookAtom } from "./cells";
 import type { CellId } from "./ids";
+import { SCRATCH_CELL_ID } from "./ids";
 export interface CellFocusState {
   focusedCellId: CellId | null;
   lastFocusedCellId: CellId | null;

frontend/src/core/cells/ids.ts

@@ -9,6 +9,9 @@ const uppercase = lowercase.toUpperCase();
 const alphabet = lowercase + uppercase;
 const seen = new Set<CellId>();
 
+export const SCRATCH_CELL_ID = "__scratch__" as CellId;
+export const SETUP_CELL_ID = "setup" as CellId;
+
 /**
  * A typed CellId
  */

frontend/src/core/codemirror/cells/extensions.ts

@@ -3,8 +3,8 @@
 import { closeCompletion, completionStatus } from "@codemirror/autocomplete";
 import { type Extension, Prec } from "@codemirror/state";
 import { EditorView, type KeyBinding, keymap } from "@codemirror/view";
-import { createTracebackInfoAtom, SCRATCH_CELL_ID } from "@/core/cells/cells";
-import { type CellId, HTMLCellId } from "@/core/cells/ids";
+import { createTracebackInfoAtom } from "@/core/cells/cells";
+import { type CellId, HTMLCellId, SCRATCH_CELL_ID } from "@/core/cells/ids";
 import type { KeymapConfig } from "@/core/config/config-schema";
 import type { HotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { store } from "@/core/state/jotai";