@marimo-team/islands 0.23.13-dev0 → 0.23.13-dev1

package/dist/{code-visibility-D9IipVFG.js → code-visibility-DwD7769p.js}

@@ -36001,7 +36001,7 @@ ${d}`,
       return Logger.warn("Failed to get version from mount config"), null;
     }
   }
-  marimoVersionAtom = atom(getVersionFromMountConfig() || "0.23.13-dev0");
+  marimoVersionAtom = atom(getVersionFromMountConfig() || "0.23.13-dev1");
   showCodeInRunModeAtom = atom(true);
   atom(null);
   var import_compiler_runtime = require_compiler_runtime();

package/dist/main.js

@@ -26,7 +26,7 @@ import { $ as reducer, B as safeExtractSetUIElementMessageBuffers, Bn as CircleA
 import { __tla as __tla_1 } from "./chunk-5FQGJX7Z-BbqSm5gU.js";
 import { o as useSize, s as Root$1, u as createLucideIcon } from "./dist--2Bqjvs0.js";
 import { A as SquareFunction, C as DEFAULT_COLOR_SCHEME, D as SCALE_TYPE_DESCRIPTIONS, E as EMPTY_VALUE$1, O as TIME_UNIT_DESCRIPTIONS, S as DEFAULT_AGGREGATION, T as DEFAULT_TIME_UNIT, _ as AGGREGATION_TYPE_DESCRIPTIONS, a as AGGREGATION_FNS$1, b as COLOR_SCHEMES, c as COLOR_BY_FIELDS, d as NONE_VALUE, f as SELECTABLE_DATA_TYPES, g as TIME_UNITS, h as STRING_AGGREGATION_FNS, i as convertDataTypeToSelectable, j as ChartColumn, k as escapeFieldName, l as COMBINED_TIME_UNITS, m as SORT_TYPES, n as createSpecWithoutData, o as BIN_AGGREGATION, p as SINGLE_TIME_UNITS, r as isFieldSet, s as CHART_TYPES, t as augmentSpecWithData, u as ChartType, v as AGGREGATION_TYPE_ICON, w as DEFAULT_MAX_BINS_FACET, x as COUNT_FIELD, y as CHART_TYPE_ICON } from "./spec-Bv-XlYiv.js";
-import { $ as contextAwarePanelOpen, $t as EmotionCacheProvider, A as prettifyRowCount, At as SELECT_COLUMN_ID, B as DatePicker, Bt as dateToLocalISOTime, C as downloadSizeLimitAtom, Ct as DelayMount, D as ErrorState, Dt as loadTableAndRawData, E as EmptyState, Et as getPageIndexForRow, F as ContextMenuSeparator, Ft as isRecord, G as CommandEmpty, Gt as ChartErrorState, H as Combobox, Ht as TabsContent, I as ContextMenuTrigger, It as isNullishFilter, J as CommandList, Jt as LazyVegaEmbed, K as CommandInput, Kt as ChartInfoState, L as useInternalStateWithSync, Lt as Maps, M as ContextMenu, Mt as toFieldTypes, N as ContextMenuContent, Nt as getMimeValues, O as LoadingState, Ot as loadTableData, P as ContextMenuItem, Pt as hasFunctionProperty, Q as PANEL_TYPES, Qt as HtmlOutput, R as useSelectList, Rt as dateToLocalISODate, S as Filenames, St as ColumnChartSpecModel, T as ColumnPreviewContainer, Tt as usePrevious$1, U as ComboboxItem, Ut as TabsList, V as DateRangePicker, Vt as Tabs, W as Command, Wt as TabsTrigger, X as smartMatch, Xt as RenderTextWithLinks, Y as CommandSeparator, Yt as useOverflowDetection, Z as ContextAwarePanelItem, Zt as Kbd, _ as ADD_PRINTING_CLASS, _t as NAMELESS_COLUMN_PREFIX, an as EyeOff, at as Toggle, b as downloadHTMLAsImage, bt as renderCellValue, c as Slide, cn as Download, d as RadioGroupItem, dn as ChevronsRight, dt as Table, en as $fae977aafc393c5c$export$588937bcd60ade55, et as contextAwarePanelOwner, f as JsonOutput, fn as ChevronsLeft, ft as TableBody, g as InstallPackageButton, gt as TableRow, h as DataTable, hn as ArrowDownWideNarrow, ht as TableHeader, in as Funnel, it as slotsController, j as getColumnCountForDisplay, jt as TOO_MANY_ROWS, k as prettifyRowColumnCount, kt as INDEX_COLUMN_NAME, l as Switch, ln as Code, lt as Fill, m as OutputRenderer, mn as ChevronLeft, mt as TableHead, n as marimoVersionAtom, nn as TextWrap, nt as isCellAwareAtom, o as SLIDE_TYPE_OPTIONS_BY_VALUE, p as OutputArea, pn as ChevronsDownUp, pt as TableCell, q as CommandItem, qt as ChartLoadingState, r as showCodeInRunModeAtom, rn as GripHorizontal, rt as SlotNames, sn as Ellipsis, t as useNotebookCodeAvailable, tn as $fae977aafc393c5c$export$6b862160d295c8e, tt as contextAwarePanelType, u as RadioGroup, un as ChevronsUpDown, ut as Provider$1, v as downloadBlob, vt as generateColumns, w as ColumnName, wt as useIntersectionObserver, x as Progress, xt as ColumnChartContext, y as downloadByURL, yt as inferFieldTypes, z as CompactChipRow, zt as dateToLocalISODateTime, __tla as __tla_2 } from "./code-visibility-D9IipVFG.js";
+import { $ as contextAwarePanelOpen, $t as EmotionCacheProvider, A as prettifyRowCount, At as SELECT_COLUMN_ID, B as DatePicker, Bt as dateToLocalISOTime, C as downloadSizeLimitAtom, Ct as DelayMount, D as ErrorState, Dt as loadTableAndRawData, E as EmptyState, Et as getPageIndexForRow, F as ContextMenuSeparator, Ft as isRecord, G as CommandEmpty, Gt as ChartErrorState, H as Combobox, Ht as TabsContent, I as ContextMenuTrigger, It as isNullishFilter, J as CommandList, Jt as LazyVegaEmbed, K as CommandInput, Kt as ChartInfoState, L as useInternalStateWithSync, Lt as Maps, M as ContextMenu, Mt as toFieldTypes, N as ContextMenuContent, Nt as getMimeValues, O as LoadingState, Ot as loadTableData, P as ContextMenuItem, Pt as hasFunctionProperty, Q as PANEL_TYPES, Qt as HtmlOutput, R as useSelectList, Rt as dateToLocalISODate, S as Filenames, St as ColumnChartSpecModel, T as ColumnPreviewContainer, Tt as usePrevious$1, U as ComboboxItem, Ut as TabsList, V as DateRangePicker, Vt as Tabs, W as Command, Wt as TabsTrigger, X as smartMatch, Xt as RenderTextWithLinks, Y as CommandSeparator, Yt as useOverflowDetection, Z as ContextAwarePanelItem, Zt as Kbd, _ as ADD_PRINTING_CLASS, _t as NAMELESS_COLUMN_PREFIX, an as EyeOff, at as Toggle, b as downloadHTMLAsImage, bt as renderCellValue, c as Slide, cn as Download, d as RadioGroupItem, dn as ChevronsRight, dt as Table, en as $fae977aafc393c5c$export$588937bcd60ade55, et as contextAwarePanelOwner, f as JsonOutput, fn as ChevronsLeft, ft as TableBody, g as InstallPackageButton, gt as TableRow, h as DataTable, hn as ArrowDownWideNarrow, ht as TableHeader, in as Funnel, it as slotsController, j as getColumnCountForDisplay, jt as TOO_MANY_ROWS, k as prettifyRowColumnCount, kt as INDEX_COLUMN_NAME, l as Switch, ln as Code, lt as Fill, m as OutputRenderer, mn as ChevronLeft, mt as TableHead, n as marimoVersionAtom, nn as TextWrap, nt as isCellAwareAtom, o as SLIDE_TYPE_OPTIONS_BY_VALUE, p as OutputArea, pn as ChevronsDownUp, pt as TableCell, q as CommandItem, qt as ChartLoadingState, r as showCodeInRunModeAtom, rn as GripHorizontal, rt as SlotNames, sn as Ellipsis, t as useNotebookCodeAvailable, tn as $fae977aafc393c5c$export$6b862160d295c8e, tt as contextAwarePanelType, u as RadioGroup, un as ChevronsUpDown, ut as Provider$1, v as downloadBlob, vt as generateColumns, w as ColumnName, wt as useIntersectionObserver, x as Progress, xt as ColumnChartContext, y as downloadByURL, yt as inferFieldTypes, z as CompactChipRow, zt as dateToLocalISODateTime, __tla as __tla_2 } from "./code-visibility-DwD7769p.js";
 import { c as Calendar, i as createReducerAndAtoms, n as useOnUnmount, o as ToggleLeft, t as useOnMount } from "./useLifecycle-AHlswLw-.js";
 import { t as Check } from "./check-C9OoNtR4.js";
 import { A as Icon, C as logNever, D as $18f2051aff69b9bf$export$a54013f0d02a8f82, E as $18f2051aff69b9bf$export$43bb16f9c6d9e3f7, F as createCollection, I as X, M as clamp$2, N as usePrevious$2, P as useDirection, R as ChevronDown, S as assertNever, a as SelectGroup, c as SelectSeparator, d as NativeSelect, i as SelectContent, j as Trigger$1, l as SelectTrigger, n as capitalize, o as SelectItem, r as Select, s as SelectLabel, t as Strings, u as SelectValue, w as $a916eb452884faea$export$b7a616150fdb9f44 } from "./strings-Dq_j3Rxw.js";
@@ -36235,7 +36235,7 @@ ${c}
   function _temp2$2(e) {
     e.target === e.currentTarget && e.key === "Enter" && (e.preventDefault(), e.stopPropagation(), e.currentTarget.click());
   }
-  var LazySlidesComponent = import_react.lazy(() => import("./reveal-component-Dk32fyu2.js"));
+  var LazySlidesComponent = import_react.lazy(() => import("./reveal-component-KlAClS-s.js"));
   const SlidesLayoutRenderer = ({ layout: e, setLayout: r, cells: c, mode: l }) => {
     var _a3;
     let u = useAtomValue(kioskModeAtom), d = l === "read" || u, f = useAtomValue(numColumnsAtom) > 1, [p, m] = (0, import_react.useState)(null), { slideCells: h, skippedIds: g, noOutputIds: _, slideTypes: v, startCellIndex: y } = (0, import_react.useMemo)(() => computeSlideCellsInfo(c, e), [

package/dist/{reveal-component-Dk32fyu2.js → reveal-component-KlAClS-s.js}

@@ -9,7 +9,7 @@ import { t as require_compiler_runtime } from "./compiler-runtime-CEbnTgxf.js";
 import { lt as kioskModeAtom } from "./html-to-image-DXwLcQ6l.js";
 import "./chunk-5FQGJX7Z-BbqSm5gU.js";
 import { u as createLucideIcon } from "./dist--2Bqjvs0.js";
-import { a as DEFAULT_SLIDE_TYPE, an as EyeOff, c as Slide, ct as PanelResizeHandle, i as DEFAULT_DECK_TRANSITION, ln as Code, on as Expand, ot as Panel, s as SlideSidebar, st as PanelGroup, t as useNotebookCodeAvailable } from "./code-visibility-D9IipVFG.js";
+import { a as DEFAULT_SLIDE_TYPE, an as EyeOff, c as Slide, ct as PanelResizeHandle, i as DEFAULT_DECK_TRANSITION, ln as Code, on as Expand, ot as Panel, s as SlideSidebar, st as PanelGroup, t as useNotebookCodeAvailable } from "./code-visibility-DwD7769p.js";
 import { X as useDebouncedCallback } from "./input-CbEz_aj_.js";
 import "./toDate-D-l5s8nn.js";
 import "./react-dom-BTJzcVJ9.js";

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/islands",
-  "version": "0.23.13-dev0",
+  "version": "0.23.13-dev1",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",

package/src/components/editor/cell/code/cell-editor.tsx

@@ -10,9 +10,11 @@ import { Button } from "@/components/ui/button";
 import { DelayMount } from "@/components/utils/delay-mount";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { useCellActions } from "@/core/cells/cells";
+import { getNotebook, useCellActions } from "@/core/cells/cells";
+import { SETUP_CELL_ID } from "@/core/cells/ids";
 import { usePendingDeleteService } from "@/core/cells/pending-delete-service";
 import type { CellData, CellRuntimeState } from "@/core/cells/types";
+import { notebookCellEditorViews } from "@/core/cells/utils";
 import { setupCodeMirror } from "@/core/codemirror/cm";
 import { acceptCompletionOnEnterAtom } from "@/core/codemirror/completion/accept-on-enter-atom";
 import { editorMountScheduler } from "@/core/codemirror/editor-mount-scheduler";
@@ -22,6 +24,10 @@ import {
   reconfigureLanguageEffect,
   switchLanguage,
 } from "@/core/codemirror/language/extension";
+import {
+  getEditorCodeAsPython,
+  updateEditorCodeFromPython,
+} from "@/core/codemirror/language/utils";
 import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages/markdown";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
 import {
@@ -214,6 +220,32 @@ const CellEditorInternal = ({
             });
           }
         },
+        addOrAppendSetupCell: (code) => {
+          const notebook = getNotebook();
+          // No setup cell yet: create one with this code at the top.
+          if (!notebook.cellIds.setupCellExists()) {
+            cellActions.addSetupCellIfDoesntExist({ code });
+            return;
+          }
+          // Setup cell exists: append the pasted code to it. Prefer updating
+          // the mounted editor (keeps the view and state in sync, like
+          // formatting does); fall back to a plain state update otherwise.
+          const view = notebookCellEditorViews(notebook)[SETUP_CELL_ID];
+          const existing = view
+            ? getEditorCodeAsPython(view)
+            : notebook.cellData[SETUP_CELL_ID].code;
+          const merged = existing.trim()
+            ? `${existing.trimEnd()}\n\n${code}`
+            : code;
+          if (view) {
+            updateEditorCodeFromPython(view, merged);
+          }
+          cellActions.updateCellCode({
+            cellId: SETUP_CELL_ID,
+            code: merged,
+            formattingChange: false,
+          });
+        },
         splitCell,
         toggleHideCode,
         aiCellCompletion: () => {

package/src/core/codemirror/cells/state.ts

@@ -8,6 +8,11 @@ export interface CodemirrorCellActions extends CellActions {
   toggleHideCode: () => boolean;
   aiCellCompletion: () => boolean;
   createManyBelow: (content: string[]) => void;
+  /**
+   * Create the setup cell with the given code, or append the code to the
+   * existing setup cell if one already exists.
+   */
+  addOrAppendSetupCell: (code: string) => void;
   onRun: () => void;
   deleteCell: () => void;
   afterToggleMarkdown: () => void;

package/src/core/codemirror/misc/__tests__/paste.test.ts

@@ -7,7 +7,7 @@ import {
   type CodemirrorCellActions,
   cellActionsState,
 } from "../../cells/state";
-import { extractCells, pasteBundle } from "../paste";
+import { extractCells, extractMarimoApp, pasteBundle } from "../paste";
 
 describe("extractCells", () => {
   it("returns empty array for non-marimo text", () => {
@@ -158,6 +158,154 @@ def _(mo, px):
     ]);
   });
 
+  it("preserves decorators on nested functions", () => {
+    const input = `
+@app.cell
+def _():
+    @functools.cache
+    def fib(n):
+        return n
+    return fib
+`;
+    expect(extractCells(input)).toEqual([
+      "@functools.cache\ndef fib(n):\n    return n",
+    ]);
+  });
+
+  it("preserves decorated methods inside a class", () => {
+    const input = `
+@app.cell
+def _():
+    class A:
+        @property
+        def x(self):
+            return self._x
+    return A
+`;
+    expect(extractCells(input)).toEqual([
+      "class A:\n    @property\n    def x(self):\n        return self._x",
+    ]);
+  });
+
+  it("handles async cells with multi-line args", () => {
+    const input = `
+@app.cell
+async def _(
+    a,
+    b,
+):
+    x = await foo(a, b)
+    return x
+`;
+    expect(extractCells(input)).toEqual(["x = await foo(a, b)"]);
+  });
+
+  it("strips multi-line returns using brackets", () => {
+    const input = `
+@app.cell
+def _():
+    a = 1
+    b = 2
+    return [
+        a,
+        b,
+    ]
+`;
+    expect(extractCells(input)).toEqual(["a = 1\nb = 2"]);
+  });
+
+  it("strips multi-line returns using braces", () => {
+    const input = `
+@app.cell
+def _():
+    a = 1
+    return {
+        "a": a,
+    }
+`;
+    expect(extractCells(input)).toEqual(["a = 1"]);
+  });
+
+  it("does not corrupt the following cell after a bracketed return", () => {
+    const input = `
+@app.cell
+def _():
+    a = 1
+    return [
+        a,
+    ]
+
+@app.cell
+def _():
+    b = 2
+    return b
+`;
+    expect(extractCells(input)).toEqual(["a = 1", "b = 2"]);
+  });
+
+  it("preserves comments in the cell body", () => {
+    const input = `
+@app.cell
+def _():
+    # leading comment
+    x = 1
+    return x
+`;
+    expect(extractCells(input)).toEqual(["# leading comment\nx = 1"]);
+  });
+
+  it("extracts the setup block separately from cells", () => {
+    const input = `
+import marimo
+app = marimo.App()
+
+with app.setup(hide_code=True):
+    import marimo as mo
+    import numpy as np
+    from numpy.linalg import eigh
+`;
+    expect(extractMarimoApp(input)).toEqual({
+      setup:
+        "import marimo as mo\nimport numpy as np\nfrom numpy.linalg import eigh",
+      cells: [],
+    });
+    // The convenience wrapper excludes the setup block.
+    expect(extractCells(input)).toEqual([]);
+  });
+
+  it("extracts @app.function definitions, keeping their return", () => {
+    const input = `
+@app.function
+def laplacian_matrix(adjacency_matrix):
+    degree_matrix = np.diag(np.sum(adjacency_matrix, axis=1))
+    L = degree_matrix - adjacency_matrix
+    return L
+`;
+    expect(extractCells(input)).toEqual([
+      "def laplacian_matrix(adjacency_matrix):\n    degree_matrix = np.diag(np.sum(adjacency_matrix, axis=1))\n    L = degree_matrix - adjacency_matrix\n    return L",
+    ]);
+  });
+
+  it("extracts setup separately while keeping cells and functions in order", () => {
+    const input = `
+with app.setup:
+    import numpy as np
+
+@app.cell
+def _():
+    a = 1
+    return (a,)
+
+@app.function
+def double(x):
+    return x * 2
+`;
+    expect(extractMarimoApp(input)).toEqual({
+      setup: "import numpy as np",
+      cells: ["a = 1", "def double(x):\n    return x * 2"],
+    });
+  });
+
   it("handles cells with config", () => {
     const input = `
 @app.cell(hide_code=True, column=2)
@@ -254,4 +402,45 @@ def _():
 
     expect(createManyBelow).not.toHaveBeenCalled();
   });
+
+  it("routes the setup block to the setup cell, not a normal cell", () => {
+    const createManyBelow = vi.fn();
+    const addOrAppendSetupCell = vi.fn();
+    const extension = pasteBundle();
+    const view = new EditorView({
+      state: EditorState.create({
+        doc: "",
+        extensions: [
+          extension,
+          cellActionsState.of({
+            createManyBelow,
+            addOrAppendSetupCell,
+          } as never as CodemirrorCellActions),
+        ],
+      }),
+    });
+
+    const clipboardData = new MockDataTransfer();
+    clipboardData.setData(
+      "text/plain",
+      `
+with app.setup:
+    import numpy as np
+
+@app.cell
+def _():
+    x = 1
+    return x
+`,
+    );
+    const event = new MockClipboardEvent("paste", { clipboardData });
+
+    // HACK: manually add the paste event listener
+    // @ts-expect-error extension not typed
+    view.dom.onpaste = (evt) => extension[0].domEventHandlers.paste(evt, view);
+    view.dom.dispatchEvent(event);
+
+    expect(addOrAppendSetupCell).toHaveBeenCalledWith("import numpy as np");
+    expect(createManyBelow).toHaveBeenCalledWith(["x = 1"]);
+  });
 });

package/src/core/codemirror/misc/paste.ts

@@ -1,6 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import type { Extension } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
+import type { SyntaxNode } from "@lezer/common";
+import { parser } from "@lezer/python";
 import { cellActionsState } from "../cells/state";
 
 export function pasteBundle(): Extension[] {
@@ -8,152 +10,166 @@ export function pasteBundle(): Extension[] {
     EditorView.domEventHandlers({
       paste: (event: ClipboardEvent, view: EditorView) => {
         const text = event.clipboardData?.getData("text/plain");
-        if (!text?.includes("@app.cell")) {
+        if (!text || !looksLikeMarimoApp(text)) {
           return false;
         }
 
-        const cells = extractCells(text);
-        if (cells.length === 0) {
+        const { setup, cells } = extractMarimoApp(text);
+        if (setup === null && cells.length === 0) {
           return false;
         }
 
         const actions = view.state.facet(cellActionsState);
-        actions.createManyBelow(cells);
+        // The `with app.setup` block belongs in the special setup cell, not in
+        // an ordinary cell.
+        if (setup !== null) {
+          actions.addOrAppendSetupCell(setup);
+        }
+        if (cells.length > 0) {
+          actions.createManyBelow(cells);
+        }
         return true;
       },
     }),
   ];
 }
 
+/** Whether the pasted text looks like marimo app source worth parsing. */
+function looksLikeMarimoApp(text: string): boolean {
+  return /@app\.cell|@app\.function|\bapp\.setup\b/.test(text);
+}
+
+export interface ExtractedApp {
+  /** The `with app.setup` block's body, if present. Belongs in the setup cell. */
+  setup: string | null;
+  /** Ordinary cells, in source order. */
+  cells: string[];
+}
+
 /**
- * Extract the cells from a marimo app.
+ * Extract the cells of a marimo app from its source.
+ *
+ * Parses the pasted source with the Python grammar (the same one that powers
+ * the editor) and turns each marimo construct into a cell:
+ * - `with app.setup(...)` contributes the setup block's body, returned
+ *   separately as `setup` because it belongs in the dedicated setup cell.
+ * - `@app.cell` functions contribute their body, minus the trailing
+ *   auto-generated `return`.
+ * - `@app.function` definitions contribute the whole function (the decorator
+ *   is dropped, but the function — including its `return` — is kept).
+ *
+ * Using the real syntax tree means decorators, `async def`, multi-line
+ * signatures, nested functions, and multi-line returns of any bracket type are
+ * all handled structurally rather than by line-based heuristics.
  */
-export function extractCells(text: string): string[] {
-  // Quick check if this looks like a marimo app
-  if (!text.includes("@app.cell")) {
-    return [];
+export function extractMarimoApp(text: string): ExtractedApp {
+  if (!looksLikeMarimoApp(text)) {
+    return { setup: null, cells: [] };
   }
 
   const cells: string[] = [];
-  const lines = text.split("\n");
-  let currentCell: string[] = [];
-  let inCell = false;
-  let skipLines = 0;
-  let inMultilineArgs = false;
-  let inMultilineReturn = false;
-  let parenCount = 0;
-  let cellBaseIndent: number | null = null;
-
-  // Pre-compile regex patterns
-  const leadingParenRegex = /\(/g;
-  const trailingParenRegex = /\)/g;
-  const cellEndMarkers = new Set(["@"]);
-
-  function countParens(line: string): number {
-    return (
-      (line.match(leadingParenRegex) || []).length -
-      (line.match(trailingParenRegex) || []).length
-    );
-  }
-
-  function getIndent(line: string): number {
-    const match = line.match(/^\s*/);
-    return match ? match[0].length : 0;
-  }
-
-  function finalizeCellIfNeeded() {
-    if (currentCell.length === 0) {
-      return;
+  const setupBlocks: string[] = [];
+  const tree = parser.parse(text);
+
+  // A node's `from` starts after its line's leading indentation; extend back to
+  // the start of the line so `dedent` sees consistent indentation, then dedent.
+  const blockText = (from: number, to: number): string => {
+    const lineStart = text.lastIndexOf("\n", from - 1) + 1;
+    return dedent(text.slice(lineStart, to));
+  };
+
+  const pushCell = (from: number, to: number) => {
+    const cell = blockText(from, to);
+    if (cell.trim()) {
+      cells.push(cell);
     }
-
-    // Only add non-empty cells
-    if (currentCell.some((l) => l.trim() !== "")) {
-      cells.push(dedent(currentCell.join("\n")));
+  };
+
+  // The text of a block's body (the statements after its `:`), optionally
+  // dropping a trailing auto-generated `return`. Returns null if empty.
+  const bodyText = (
+    body: SyntaxNode,
+    stripTrailingReturn: boolean,
+  ): string | null => {
+    const children: SyntaxNode[] = [];
+    for (let child = body.firstChild; child; child = child.nextSibling) {
+      if (child.name !== ":") {
+        children.push(child);
+      }
     }
-    currentCell = [];
-  }
-
-  for (const line of lines) {
-    const trimmed = line.trim();
-
-    // Skip empty lines between cells
-    if (!trimmed && !inCell) {
-      continue;
+    if (stripTrailingReturn && children.at(-1)?.name === "ReturnStatement") {
+      children.pop();
     }
-
-    // Start of a new cell
-    if (trimmed.startsWith("@app.cell")) {
-      finalizeCellIfNeeded();
-      inCell = true;
-      skipLines = 1; // Skip the def line
-      cellBaseIndent = null;
-      continue;
+    const first = children[0];
+    const last = children.at(-1);
+    if (!first || !last) {
+      return null;
     }
-
-    // Handle function definition and args
-    if (skipLines > 0) {
-      if (
-        trimmed.startsWith("def") &&
-        trimmed.includes("(") &&
-        !trimmed.includes("):")
-      ) {
-        inMultilineArgs = true;
-        parenCount = countParens(trimmed);
+    const cell = blockText(first.from, last.to);
+    return cell.trim() ? cell : null;
+  };
+
+  tree.iterate({
+    enter: (node) => {
+      // Setup cell: `with app.setup:` / `with app.setup(...):`
+      if (node.name === "WithStatement") {
+        const body = node.node.getChild("Body");
+        if (body && text.slice(node.from, body.from).includes("app.setup")) {
+          const setup = bodyText(body, false);
+          if (setup !== null) {
+            setupBlocks.push(setup);
+          }
+        }
+        return false;
       }
-      skipLines--;
-      continue;
-    }
 
-    // Track multi-line arguments
-    if (inMultilineArgs) {
-      parenCount += countParens(trimmed);
-      if (parenCount === 0) {
-        inMultilineArgs = false;
+      if (node.name !== "DecoratedStatement") {
+        // Keep descending until we reach the decorated definitions.
+        return true;
       }
-      continue;
-    }
-
-    if (!inCell) {
-      continue;
-    }
-
-    // Handle cell content
-    if (cellEndMarkers.has(trimmed[0]) || trimmed.startsWith("if __name__")) {
-      finalizeCellIfNeeded();
-      inCell = trimmed.startsWith("@");
-      continue;
-    }
-
-    // Detect base indentation of cell body from first content line
-    if (cellBaseIndent === null && trimmed) {
-      cellBaseIndent = getIndent(line);
-    }
 
-    // Handle return statements — only strip cell-level returns
-    if (trimmed.startsWith("return") && getIndent(line) === cellBaseIndent) {
-      if (trimmed.includes("(") && !trimmed.endsWith(")")) {
-        inMultilineReturn = true;
-        parenCount = countParens(trimmed);
+      const decorator = node.node.getChild("Decorator");
+      const fn = node.node.getChild("FunctionDefinition");
+      if (!decorator || !fn) {
+        return false;
+      }
+      const decoratorText = text.slice(decorator.from, decorator.to).trim();
+
+      // `@app.cell` / `@app.cell(...)`: the body becomes a cell, minus the
+      // trailing auto-generated return.
+      if (decoratorText.startsWith("@app.cell")) {
+        const body = fn.getChild("Body");
+        if (body) {
+          const cell = bodyText(body, true);
+          if (cell !== null) {
+            cells.push(cell);
+          }
+        }
+        return false;
       }
-      continue;
-    }
 
-    if (inMultilineReturn) {
-      parenCount += countParens(trimmed);
-      if (parenCount === 0) {
-        inMultilineReturn = false;
+      // `@app.function`: the whole function definition is the cell (the
+      // decorator is dropped, but the function — including its return — stays).
+      if (decoratorText.startsWith("@app.function")) {
+        pushCell(fn.from, fn.to);
       }
-      continue;
-    }
 
-    // Add line to current cell
-    currentCell.push(line);
-  }
+      return false;
+    },
+  });
 
-  // Handle last cell
-  finalizeCellIfNeeded();
+  return {
+    setup: setupBlocks.length > 0 ? setupBlocks.join("\n\n") : null,
+    cells,
+  };
+}
 
-  return cells;
+/**
+ * Convenience wrapper returning only the ordinary cells (excluding the setup
+ * block). Prefer {@link extractMarimoApp} when the setup block matters.
+ */
+export function extractCells(text: string): string[] {
+  return extractMarimoApp(text).cells;
 }
 
 function dedent(text: string): string {