@marimo-team/islands 0.25.3-dev209 → 0.25.3-dev210

package/dist/{vega-component-Blvspl7v.js → vega-component-9NOEITbg.js}

@@ -1,22 +1,23 @@
 import { a as e } from "./rolldown-runtime-Cb76FRpK.js";
-import { g as t, h as n, m as r, o as i } from "./button-BQRdqQi1.js";
+import { a as t, g as n, h as r, m as i } from "./button-B6RjDKJx.js";
 import { t as a } from "./react-sHtZdjeP.js";
 import { t as o } from "./compiler-runtime-CldRdSsv.js";
-import { t as s } from "./tooltip-s4zZxwKe.js";
-import { C as c, a as l } from "./toDate-BZZ64K_m.js";
+import { t as s } from "./tooltip-C3UnvsRe.js";
+import { l as c, t as l } from "./config-DbjzjwxX.js";
 import { t as u } from "./jsx-runtime-3BD1kT0s.js";
-import { n as d } from "./error-banner-DxJ3s8YO.js";
+import { n as d } from "./error-banner-IsonnvA0.js";
 import { t as f } from "./debounce-C3jyATpc.js";
-import { R as p, n as ee } from "./useTheme-TzPgWXwB.js";
+import { n as p } from "./useTheme-B9C1JIy9.js";
+import { D as ee } from "./utils-DXkj82sK.js";
 import { t as te } from "./objects-DCmMvZtn.js";
-import { a as m, t as h } from "./arrays-DoQ-AILx.js";
-import { t as ne } from "./useAsyncData-Dqw22wZQ.js";
-import { a as g, i as re, n as ie, t as ae } from "./formats-Ry7c-aIV.js";
+import { a as ne, i as re, n as ie, t as m } from "./formats-BK4VefeF.js";
+import { a as h, t as g } from "./arrays-DoQ-AILx.js";
+import { t as ae } from "./useAsyncData-DsZMvNXn.js";
 import { t as _ } from "./semaphore-89jsoHtS.js";
 import { D as oe, n as se } from "./vega-loader.browser-B9Pz0EGP.js";
 import { t as ce } from "./react-vega-C4T9Rb_w.js";
-import { a as le, n as v } from "./loader-DOVB0Syl.js";
-import { t as ue } from "./useDeepCompareMemoize-D3gozcjJ.js";
+import { a as le, n as v } from "./loader-BisjOLiW.js";
+import { t as ue } from "./useDeepCompareMemoize-KMBBux8C.js";
 import { t as de } from "./use-vega-container-remeasure-Bz4PZMtt.js";
 //#region src/plugins/impl/vega/fix-relative-url.ts
 var fe = /* @__PURE__ */ e(a(), 1), pe = o();
@@ -102,7 +103,7 @@ var y = {
 	}
 };
 function C() {
-	return i() ? "metaKey" : "ctrlKey";
+	return t() ? "metaKey" : "ctrlKey";
 }
 var w = {
 	highlight() {
@@ -199,9 +200,9 @@ function T(e) {
 function E(e) {
 	if ("params" in e && e.params && e.params.length > 0) {
 		let t = e.params.filter((e) => e != null && "select" in e && e.select !== void 0).map((e) => e.name);
-		return m(t);
+		return h(t);
 	}
-	return "layer" in e ? m(e.layer.flatMap(E)) : "vconcat" in e ? m(e.vconcat.flatMap(E)) : "hconcat" in e ? m(e.hconcat.flatMap(E)) : [];
+	return "layer" in e ? h(e.layer.flatMap(E)) : "vconcat" in e ? h(e.vconcat.flatMap(E)) : "hconcat" in e ? h(e.hconcat.flatMap(E)) : [];
 }
 function D(e) {
 	if (!e || !("mark" in e)) return;
@@ -378,7 +379,7 @@ function he(e, t) {
 		let t = e.params && e.params.length > 0, i = r !== !1 && !t, a = [];
 		if (i) {
 			let t = e.layer.flatMap((e) => "mark" in e ? k(e) : []);
-			a = [...new Set(t)], Array.isArray(r) && (a = h.intersection(a, r));
+			a = [...new Set(t)], Array.isArray(r) && (a = g.intersection(a, r));
 		}
 		let o = e.layer.map((e, t) => {
 			if (!("mark" in e)) return e;
@@ -404,7 +405,7 @@ function he(e, t) {
 function R(e, t) {
 	if (t === !1) return e;
 	let n = k(e);
-	Array.isArray(t) && (n = h.intersection(n, t));
+	Array.isArray(t) && (n = g.intersection(n, t));
 	let r = n.map((e) => w.legend(e)), i = [...e.params || [], ...r];
 	return {
 		...e,
@@ -510,11 +511,11 @@ async function W(e) {
 //#endregion
 //#region src/plugins/impl/vega/vega-component.tsx
 var G = u();
-se("arrow", ae);
+se("arrow", m);
 var K = (e) => {
 	let t = (0, pe.c)(12), { value: n, setValue: r, chartSelection: i, fieldSelection: a, spec: o, embedOptions: s } = e, c, l;
 	t[0] === o ? (c = t[1], l = t[2]) : (c = async () => W(o), l = [o], t[0] = o, t[1] = c, t[2] = l);
-	let { data: u, error: f } = ne(c, l);
+	let { data: u, error: f } = ae(c, l);
 	if (f) {
 		let e;
 		return t[3] === f ? e = t[4] : (e = /* @__PURE__ */ (0, G.jsx)(d, { error: f }), t[3] = f, t[4] = e), e;
@@ -530,10 +531,10 @@ var K = (e) => {
 		embedOptions: s
 	}), t[5] = i, t[6] = s, t[7] = a, t[8] = u, t[9] = r, t[10] = n, t[11] = p) : p = t[11], p;
 }, q = (e) => {
-	let a = (0, pe.c)(50), { value: o, setValue: l, chartSelection: u, fieldSelection: d, spec: m, embedOptions: h } = e, { theme: ne } = ee(), g = (0, fe.useRef)(null), ae = (0, fe.useRef)(void 0), [_, se] = (0, fe.useState)(), v;
+	let a = (0, pe.c)(50), { value: o, setValue: l, chartSelection: u, fieldSelection: d, spec: ne, embedOptions: m } = e, { theme: h } = p(), g = (0, fe.useRef)(null), ae = (0, fe.useRef)(void 0), [_, se] = (0, fe.useState)(), v;
 	bb0: {
-		if (h && "actions" in h) {
-			v = h.actions;
+		if (m && "actions" in m) {
+			v = m.actions;
 			break bb0;
 		}
 		let e;
@@ -542,7 +543,7 @@ var K = (e) => {
 			compiled: !1
 		}, a[0] = e) : e = a[0], v = e;
 	}
-	let y = v, b = ue(m), x;
+	let y = v, b = ue(ne), x;
 	a[1] !== u || a[2] !== d || a[3] !== b ? (x = he(me(b), {
 		chartSelection: u,
 		fieldSelection: d
@@ -556,10 +557,10 @@ var K = (e) => {
 			...e
 		});
 	}, a[7] = l, a[8] = o, a[9] = D) : D = a[9];
-	let O = p(D), k;
-	a[10] === O ? k = a[11] : (k = f((e, n) => {
-		t.debug("[Vega signal]", e, n);
-		let r = te.mapValues(n, _e);
+	let O = ee(D), k;
+	a[10] === O ? k = a[11] : (k = f((e, t) => {
+		n.debug("[Vega signal]", e, t);
+		let r = te.mapValues(t, _e);
 		r = te.mapValues(r, ge), O({ [e]: r });
 	}, 100), a[10] = O, a[11] = k);
 	let A = k, j = ue(T), M;
@@ -572,18 +573,18 @@ var K = (e) => {
 	} else M = a[14];
 	let N = M, P;
 	a[17] === C ? P = a[18] : (P = (e) => {
-		t.error(e), t.debug(C), se(e);
+		n.error(e), n.debug(C), se(e);
 	}, a[17] = C, a[18] = P);
-	let F = p(P), I;
+	let F = ee(P), I;
 	a[19] === Symbol.for("react.memo_cache_sentinel") ? (I = (e) => {
-		t.debug("[Vega view] created", e), ae.current = e, se(void 0);
+		n.debug("[Vega view] created", e), ae.current = e, se(void 0);
 	}, a[19] = I) : I = a[19];
-	let L = p(I), R;
+	let L = ee(I), R;
 	a[20] === T ? R = a[21] : (R = () => {
 		let e = [];
 		if (S.hasPoint(T) && e.push(["Point selection", "click to select a point; hold shift for multi-select"]), S.hasInterval(T) && e.push(["Interval selection", "click and drag to select an interval"]), S.hasLegend(T) && e.push(["Legend selection", "click to select a legend item; hold shift for multi-select"]), S.hasPanZoom(T)) {
-			let t = i() ? "⌘" : "ctrl";
-			e.push(["Pan", `hold the ${t} key and drag`], ["Zoom", `hold the ${t} key and scroll`]);
+			let n = t() ? "⌘" : "ctrl";
+			e.push(["Pan", `hold the ${n} key and drag`], ["Zoom", `hold the ${n} key and scroll`]);
 		}
 		return e.length === 0 ? null : /* @__PURE__ */ (0, G.jsx)(s, {
 			delayDuration: 300,
@@ -595,7 +596,7 @@ var K = (e) => {
 			children: /* @__PURE__ */ (0, G.jsx)(c, { className: "absolute bottom-1 right-0 m-2 h-4 w-4 cursor-help text-muted-foreground hover:text-foreground" })
 		});
 	}, a[20] = T, a[21] = R);
-	let z = R, B = ne === "dark" ? "dark" : void 0, V;
+	let z = R, B = h === "dark" ? "dark" : void 0, V;
 	a[22] !== y || a[23] !== B ? (V = {
 		theme: B,
 		actions: y,
@@ -617,21 +618,21 @@ var K = (e) => {
 	de(g, K === "container");
 	let q, J;
 	a[32] !== U || a[33] !== N ? (q = () => (N.forEach((e) => {
-		let { signalName: n, handler: r } = e;
+		let { signalName: t, handler: r } = e;
 		try {
-			U?.view.addSignalListener(n, r);
+			U?.view.addSignalListener(t, r);
 		} catch (e) {
-			let n = e;
-			t.error(n);
+			let t = e;
+			n.error(t);
 		}
 	}), () => {
 		N.forEach((e) => {
-			let { signalName: n, handler: r } = e;
+			let { signalName: t, handler: r } = e;
 			try {
-				U?.view.removeSignalListener(n, r);
+				U?.view.removeSignalListener(t, r);
 			} catch (e) {
-				let n = e;
-				t.error(n);
+				let t = e;
+				n.error(t);
 			}
 		});
 	}), J = [U, N], a[32] = U, a[33] = N, a[34] = q, a[35] = J) : (q = a[34], J = a[35]), (0, fe.useEffect)(q, J);
@@ -643,29 +644,29 @@ var K = (e) => {
 			children: _.stack
 		})]
 	}), a[36] = _, a[37] = Y);
-	let ye, be;
-	a[38] === Symbol.for("react.memo_cache_sentinel") ? (ye = n("relative"), be = r.stopPropagation(), a[38] = ye, a[39] = be) : (ye = a[38], be = a[39]);
-	let X;
-	a[40] === K ? X = a[41] : (X = /* @__PURE__ */ (0, G.jsx)("div", {
+	let X, ye;
+	a[38] === Symbol.for("react.memo_cache_sentinel") ? (X = r("relative"), ye = i.stopPropagation(), a[38] = X, a[39] = ye) : (X = a[38], ye = a[39]);
+	let Z;
+	a[40] === K ? Z = a[41] : (Z = /* @__PURE__ */ (0, G.jsx)("div", {
 		ref: g,
 		"data-container-width": K
-	}), a[40] = K, a[41] = X);
-	let Z;
-	a[42] === z ? Z = a[43] : (Z = z(), a[42] = z, a[43] = Z);
+	}), a[40] = K, a[41] = Z);
 	let Q;
-	a[44] !== X || a[45] !== Z ? (Q = /* @__PURE__ */ (0, G.jsxs)("div", {
-		className: ye,
-		onPointerDown: be,
-		children: [X, Z]
-	}), a[44] = X, a[45] = Z, a[46] = Q) : Q = a[46];
+	a[42] === z ? Q = a[43] : (Q = z(), a[42] = z, a[43] = Q);
 	let $;
-	return a[47] !== Y || a[48] !== Q ? ($ = /* @__PURE__ */ (0, G.jsxs)(G.Fragment, { children: [Y, Q] }), a[47] = Y, a[48] = Q, a[49] = $) : $ = a[49], $;
+	a[44] !== Z || a[45] !== Q ? ($ = /* @__PURE__ */ (0, G.jsxs)("div", {
+		className: X,
+		onPointerDown: ye,
+		children: [Z, Q]
+	}), a[44] = Z, a[45] = Q, a[46] = $) : $ = a[46];
+	let be;
+	return a[47] !== Y || a[48] !== $ ? (be = /* @__PURE__ */ (0, G.jsxs)(G.Fragment, { children: [Y, $] }), a[47] = Y, a[48] = $, a[49] = be) : be = a[49], be;
 };
 function ge(e) {
 	return e instanceof Set ? [...e] : e;
 }
 function _e(e) {
-	return Array.isArray(e) ? e.map((e) => e instanceof Date && g(e) ? new Date(e).getTime() : e) : e;
+	return Array.isArray(e) ? e.map((e) => e instanceof Date && ne(e) ? new Date(e).getTime() : e) : e;
 }
 function ve(e, t) {
 	return /* @__PURE__ */ (0, G.jsxs)("div", { children: [

package/dist/{vennDiagram-4TSXK5OY-DzCi-lRh.js → vennDiagram-4TSXK5OY-BfzXp4SH.js}

@@ -3,7 +3,7 @@ import { t } from "./src-CA_m4ONl.js";
 import { n } from "./chunk-Y2CYZVJY-DKKS1VyV.js";
 import { $ as r, H as i, K as a, U as o, a as s, b as c, c as l, et as u, f as d, nt as f, tt as p, v as m, w as h, y as g } from "./chunk-DU6HZSFF-BwiyaNpg.js";
 import { t as _ } from "./rough.esm-CtV9AnV1.js";
-import { p as v } from "./mermaid-BrdwbArF.js";
+import { p as v } from "./mermaid-B-OFqheA.js";
 //#region ../node_modules/.pnpm/khroma@2.1.0/node_modules/khroma/dist/methods/transparentize.js
 var y = (e, t) => p(e, "a", -t), b = 1e-10;
 function x(e, t) {

package/dist/{wardleyDiagram-VM6X3IG4-CW48xa0n.js → wardleyDiagram-VM6X3IG4-CiqEE8eq.js}

@@ -2,7 +2,7 @@ import { i as e } from "./chunk-75Z2AOVW-D7gxIs2Q.js";
 import { n as t } from "./src-CA_m4ONl.js";
 import { n } from "./chunk-Y2CYZVJY-DKKS1VyV.js";
 import { D as r, H as i, K as a, U as o, a as s, b as c, c as l, v as u, w as d, x as f, y as p } from "./chunk-DU6HZSFF-BwiyaNpg.js";
-import { p as m } from "./mermaid-BrdwbArF.js";
+import { p as m } from "./mermaid-B-OFqheA.js";
 import { n as h } from "./mermaid-parser.core-DxuSfYwJ.js";
 import { t as g } from "./chunk-JWPE2WC7-DJvv1t8U.js";
 //#region ../node_modules/.pnpm/mermaid@11.17.2/node_modules/mermaid/dist/chunks/mermaid.core/wardleyDiagram-VM6X3IG4.mjs

package/dist/{xychartDiagram-S5SC5T6Z-UPxGGI2B.js → xychartDiagram-S5SC5T6Z-BTIHCVOI.js}

@@ -9,7 +9,7 @@ import { n as s } from "./chunk-Y2CYZVJY-DKKS1VyV.js";
 import { D as c, H as l, K as u, U as d, a as f, b as p, c as m, f as h, v as g, w as _, y as v, z as y } from "./chunk-DU6HZSFF-BwiyaNpg.js";
 import "./chunk-PWAF6VOD-CQw5qATZ.js";
 import { t as b } from "./chunk-GMAD6QVW-BhaJx0X3.js";
-import { p as x } from "./mermaid-BrdwbArF.js";
+import { p as x } from "./mermaid-B-OFqheA.js";
 //#region ../node_modules/.pnpm/d3-scale@4.0.2/node_modules/d3-scale/src/band.js
 function S() {
 	var e = t().unknown(void 0), r = e.domain, i = e.range, o = 0, s = 1, c, l, u = !1, d = 0, f = 0, p = .5;

package/package.json

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

package/src/components/editor/cell/__tests__/StagedAICell.test.tsx

@@ -18,7 +18,7 @@ vi.mock("@/core/cells/cells", () => ({
 vi.mock("@/core/cells/actions", () => ({
   useCellActions: () => ({
     createCell: vi.fn(),
-    updateCellCode: vi.fn(),
+    setCellCode: vi.fn(),
   }),
 }));
 

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

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { historyField } from "@codemirror/commands";
-import { EditorState, StateEffect } from "@codemirror/state";
+import { EditorState, type Extension, StateEffect } from "@codemirror/state";
 import { EditorView, ViewPlugin } from "@codemirror/view";
 import { useIntersectionObserver } from "@uidotdev/usehooks";
 import { useAtom, useAtomValue, useSetAtom } from "jotai";
@@ -8,18 +8,18 @@ import React, { memo, useEffect, useMemo, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { getNotebook } from "@/core/cells/cells";
 import { useCellActions } from "@/core/cells/actions";
 import {
   showHiddenCodeAtom,
   takeRestoredEditorStateAtom,
   useRestoredEditorState,
 } from "@/core/cells/editor";
-import { useNotebookDocumentEdit } from "@/core/notebook/document";
-import { SETUP_CELL_ID } from "@/core/cells/ids";
+import {
+  useNotebookCellCodeExtension,
+  useNotebookDocumentEdit,
+} from "@/core/notebook/document";
 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";
@@ -29,8 +29,6 @@ import {
   reconfigureLanguageEffect,
   switchLanguage,
 } from "@/core/codemirror/language/extension";
-import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
-import { getEditorCodeAsPython } 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 { autoInstantiateAtom, isAiFeatureEnabled } from "@/core/config/config";
@@ -88,6 +86,8 @@ export interface CellEditorProps
    * defaults to `#App`.
    */
   tooltipParentSelector?: string;
+  /** Further editor extensions, for an editor outside the notebook. */
+  extensions?: Extension[];
 }
 
 const CellEditorInternal = ({
@@ -111,6 +111,7 @@ const CellEditorInternal = ({
   inlineAiTooltip,
   outputArea,
   tooltipParentSelector,
+  extensions: extraExtensions,
 }: CellEditorProps) => {
   const [aiCompletionCell, setAiCompletionCell] = useAtom(aiCompletionCellAtom);
   const serializedEditorState = useRestoredEditorState(cellId);
@@ -120,6 +121,7 @@ const CellEditorInternal = ({
   const { saveOrNameNotebook } = useSaveNotebook();
   const pendingDeleteService = usePendingDeleteService();
   const edit = useNotebookDocumentEdit();
+  const codeBinding = useNotebookCellCodeExtension(cellId);
 
   const loading = status === "running" || status === "queued";
   const cellActions = useCellActions();
@@ -209,35 +211,6 @@ const CellEditorInternal = ({
             });
           }
         },
-        addOrAppendSetupCell: (code) => {
-          const notebook = getNotebook();
-          // No setup cell yet: create one with this code at the top.
-          if (!notebook.cellIds.includes(SETUP_CELL_ID)) {
-            cellActions.addSetupCell({ 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) {
-            const update = prepareCellCodeUpdate(view.state, merged);
-            if (update) {
-              view.dispatch(update);
-            }
-          }
-          cellActions.updateCellCode({
-            cellId: SETUP_CELL_ID,
-            code: merged,
-            formattingChange: false,
-          });
-        },
         splitCell,
         toggleHideCode,
         aiCellCompletion: () => {
@@ -270,6 +243,8 @@ const CellEditorInternal = ({
     });
 
     extensions.push(
+      codeBinding,
+      ...(extraExtensions ?? []),
       // Listen to code changes if we can use markdown
       // Also update the language adapter
       ViewPlugin.define((view) => {
@@ -309,6 +284,8 @@ const CellEditorInternal = ({
     return extensions;
   }, [
     cellId,
+    codeBinding,
+    extraExtensions,
     acceptCompletionOnEnter,
     userConfig.keymap,
     userConfig.completion,

package/src/components/editor/chrome/wrapper/__tests__/pending-ai-cells.test.tsx

@@ -19,7 +19,7 @@ vi.mock("@/core/cells/cells", () => ({
 vi.mock("@/core/cells/actions", () => ({
   useCellActions: () => ({
     createCell: vi.fn(),
-    updateCellCode: vi.fn(),
+    setCellCode: vi.fn(),
   }),
 }));
 

package/src/components/editor/errors/__tests__/send-error-report-button.test.tsx

@@ -422,11 +422,17 @@ describe("SendErrorReportButton", () => {
             cellRuntime: { [CELL_ID]: { output } },
           }),
         );
-        actions.updateCellCode({
-          cellId: CELL_ID,
-          code: "print('edited')",
-          formattingChange: false,
-        });
+        store.set(notebookAtom, (state) => ({
+          ...state,
+          cellData: {
+            ...state.cellData,
+            [CELL_ID]: {
+              ...state.cellData[CELL_ID],
+              code: "print('edited')",
+              edited: true,
+            },
+          },
+        }));
         if (queue) {
           actions.prepareForRun({ cellId: CELL_ID });
           actions.handleCellMessage({

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

@@ -209,7 +209,6 @@ function useCellHiddenLogic({
 
 export type CellComponentActions = Pick<
   CellActions,
-  | "updateCellCode"
   | "focusCell"
   | "moveCell"
   | "collapseCell"

package/src/components/scratchpad/scratchpad.test.tsx

@@ -7,26 +7,22 @@ import { type ComponentProps, useEffect, useRef } from "react";
 import { expect, it, vi } from "vitest";
 import type { CellEditor } from "@/components/editor/cell/code/cell-editor";
 import { TooltipProvider } from "@/components/ui/tooltip";
-import { SCRATCH_CELL_ID } from "@/core/cells/ids";
 import { ScratchPad } from "./scratchpad";
 import {
   historyVisibleAtom,
   scratchpadHistoryAtom,
 } from "./scratchpad-history";
+import { scratchpadCodeAtom } from "./state";
 
 const state = vi.hoisted(() => ({
   editor: null as EditorView | null,
-  updateCellCode: vi.fn(),
 }));
 
 vi.mock("@/core/cells/cells", () => ({
   useNotebook: () => ({ cellRuntime: {}, cellData: {} }),
 }));
 vi.mock("@/core/cells/actions", () => ({
-  useCellActions: () => ({
-    createCell: vi.fn(),
-    updateCellCode: state.updateCellCode,
-  }),
+  useCellActions: () => ({ createCell: vi.fn() }),
 }));
 vi.mock("@/core/cells/focus", () => ({ useLastFocusedCellId: () => null }));
 vi.mock("@/core/network/requests", () => ({
@@ -89,9 +85,5 @@ it("focuses the editor after recalling history", () => {
     screen.getByRole("textbox", { name: "Scratchpad editor" }),
   ).toHaveFocus();
   expect(state.editor?.state.doc.toString()).toBe("print(42)");
-  expect(state.updateCellCode).toHaveBeenCalledExactlyOnceWith({
-    cellId: SCRATCH_CELL_ID,
-    code: "print(42)",
-    formattingChange: false,
-  });
+  expect(store.get(scratchpadCodeAtom)).toBe("print(42)");
 });

package/src/components/scratchpad/scratchpad.tsx

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { EditorView } from "@codemirror/view";
+import { EditorView } from "@codemirror/view";
 import { useAtom, useAtomValue, useSetAtom } from "jotai";
 import {
   BetweenHorizontalStartIcon,
@@ -10,7 +10,7 @@ import {
   PlayIcon,
 } from "lucide-react";
 import type React from "react";
-import { Suspense, useRef, useState } from "react";
+import { Suspense, useMemo, useRef, useState } from "react";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
 import useEvent from "react-use-event-hook";
 import { useNotebook } from "@/core/cells/cells";
@@ -19,6 +19,8 @@ import { useLastFocusedCellId } from "@/core/cells/focus";
 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 { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
+import { getEditorCodeAsPython } from "@/core/codemirror/language/utils";
 import { useResolvedMarimoConfig } from "@/core/config/config";
 import { CSSClasses } from "@/core/constants";
 import { useRequestClient } from "@/core/network/requests";
@@ -44,6 +46,7 @@ import {
   historyVisibleAtom,
   scratchpadHistoryAtom,
 } from "./scratchpad-history";
+import { scratchpadCodeAtom } from "./state";
 
 const scratchpadCellConfig: CellConfig = {
   hide_code: false,
@@ -56,7 +59,7 @@ export const ScratchPad: React.FC = () => {
   const { theme } = useTheme();
   const ref = useRef<EditorView | null>(null);
   const lastFocusedCellId = useLastFocusedCellId();
-  const { updateCellCode, createCell } = useCellActions();
+  const { createCell } = useCellActions();
   const { sendRunScratchpad } = useRequestClient();
   const orientation = usePanelOrientation();
   const section = usePanelSection();
@@ -66,8 +69,17 @@ export const ScratchPad: React.FC = () => {
   const output = cellRuntime?.output;
   const status = cellRuntime?.status;
   const consoleOutputs = cellRuntime?.consoleOutputs;
-  const cellData = notebookState.cellData[cellId];
-  const code = cellData?.code ?? "";
+  const [code, setCode] = useAtom(scratchpadCodeAtom);
+  const codeListener = useMemo(
+    () => [
+      EditorView.updateListener.of((update) => {
+        if (update.docChanged) {
+          setCode(getEditorCodeAsPython(update.view));
+        }
+      }),
+    ],
+    [setCode],
+  );
 
   const addToHistory = useSetAtom(addToHistoryAtom);
   const [historyVisible, setHistoryVisible] = useAtom(historyVisibleAtom);
@@ -86,11 +98,7 @@ export const ScratchPad: React.FC = () => {
   });
 
   const handleClearCode = useEvent(() => {
-    updateCellCode({
-      cellId,
-      code: "",
-      formattingChange: false,
-    });
+    setCode("");
     void sendRunScratchpad({ code: "" });
     const ev = ref.current;
     if (ev) {
@@ -106,20 +114,13 @@ export const ScratchPad: React.FC = () => {
 
   const handleSelectHistoryItem = useEvent((item: string) => {
     setHistoryVisible(false);
-    updateCellCode({
-      cellId,
-      code: item,
-      formattingChange: false,
-    });
+    setCode(item);
     const ev = ref.current;
     if (ev) {
-      ev.dispatch({
-        changes: {
-          from: 0,
-          to: ev.state.doc.length,
-          insert: item,
-        },
-      });
+      const update = prepareCellCodeUpdate(ev.state, item);
+      if (update) {
+        ev.dispatch(update);
+      }
       ev.focus();
     }
   });
@@ -269,6 +270,7 @@ export const ScratchPad: React.FC = () => {
                 showHiddenCode={Functions.NOOP}
                 languageAdapter={languageAdapter}
                 setLanguageAdapter={setLanguageAdapter}
+                extensions={codeListener}
               />
             </div>
             {renderHistory()}

package/src/components/scratchpad/state.ts

@@ -0,0 +1,7 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The scratchpad's editor state, kept while its panel is closed. */
+
+import { atom } from "jotai";
+
+export const scratchpadCodeAtom = atom("");

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

@@ -22,7 +22,7 @@ const { stagedGenerationAtom, createActions, reducer, initialState } =
 
 // Mock the dependencies
 const mockCreateCell = vi.fn();
-const mockUpdateCellCode = vi.fn();
+const mockSetCellCode = vi.fn();
 const mockDeleteCellCallback = vi.fn();
 
 // Mock cell handle with editor view
@@ -37,16 +37,13 @@ const mockCellHandle = {
 vi.mock("../../cells/actions", () => ({
   useCellActions: () => ({
     createCell: mockCreateCell,
-    updateCellCode: mockUpdateCellCode,
+    setCellCode: mockSetCellCode,
   }),
 }));
 vi.mock("../../cells/cells", async () => {
   const { atom } = await import("jotai");
   return {
     notebookAtom: atom({ cellData: {}, cellIds: [] }),
-    cellHandleAtom: vi.fn(() => ({
-      read: vi.fn(() => mockCellHandle),
-    })),
     getCellEditorView: vi.fn(() => mockCellHandle.current.editorViewOrNull),
   };
 });
@@ -403,6 +400,7 @@ describe("staged cell generation", () => {
     // No cell or cell update should have been called
     expect(mockCreateCell).not.toHaveBeenCalled();
     expect(prepareCellCodeUpdate).not.toHaveBeenCalled();
+    expect(mockSetCellCode).not.toHaveBeenCalled();
     expect(store.get(stagedGenerationInProgressAtom)).toBe(true);
   });
 
@@ -688,9 +686,10 @@ describe("staged cell generation", () => {
     expect(
       mockCellHandle.current.editorViewOrNull.dispatch,
     ).toHaveBeenCalledTimes(2);
+    expect(mockSetCellCode).not.toHaveBeenCalled();
   });
 
-  it("should update notebook state when the cell editor is not mounted", () => {
+  it("writes through the document when the cell editor is not mounted", () => {
     const { result } = renderGenerationHook();
     result.current.beginStagedCellGeneration();
 
@@ -710,10 +709,9 @@ describe("staged cell generation", () => {
       },
     });
 
-    expect(mockUpdateCellCode).toHaveBeenCalledWith({
+    expect(mockSetCellCode).toHaveBeenCalledWith({
       cellId: cellId("unmounted-cell"),
       code: "value = 1",
-      formattingChange: false,
     });
   });