@marimo-team/islands 0.25.3-dev169 → 0.25.3-dev170

package/src/core/notebook/document/codemirror.ts

@@ -28,34 +28,24 @@
  * coarsely instead of blocking the editor.
  */
 
-import { diff } from "@codemirror/merge";
 import {
-  type ChangeSpec,
   type EditorState,
   type Extension,
-  type StateEffect,
   Transaction,
 } from "@codemirror/state";
 import { type EditorView, ViewPlugin, type ViewUpdate } from "@codemirror/view";
 import type { CellId } from "@marimo-team/client/protocol";
-import { dequal } from "dequal";
 import type { Atom, createStore } from "jotai/vanilla";
 import {
-  languageAdapterFromCode,
-  languageAdapterState,
-  switchLanguage,
-} from "../../codemirror/language/extension";
-import { initializeSQLDialect } from "../../codemirror/language/languages/sql/sql";
-import {
-  languageMetadataField,
-  setLanguageMetadata,
-} from "../../codemirror/language/metadata";
+  editorShowsCellCode,
+  getCellCode,
+  prepareCellCodeUpdate,
+} from "../../codemirror/language/cell-code";
+import { languageAdapterState } from "../../codemirror/language/extension";
+import { languageMetadataField } from "../../codemirror/language/metadata";
 import type { DocumentModel } from "./model";
 import { selectCellField } from "./select";
 
-/** The diff scan limit, matching the default of the merge view. */
-const SCAN_LIMIT = 500;
-
 interface Binding {
   readonly store: ReturnType<typeof createStore>;
   readonly model: DocumentModel;
@@ -94,17 +84,11 @@ export function cellCodeBinding({
 }
 
 class CellCodeSync {
+  // Only the initial load infers a mode; subsequent edits preserve it.
+  #initialized = false;
   #scheduled = false;
   /** A pull arrived during IME composition and runs once it ends. */
   #deferred = false;
-  /**
-   * A pull is dispatching to the editor.
-   *
-   * A language switch dispatches transactions of its own before the text, and
-   * their in-between states represent the wrong code. Until the switch can be
-   * part of the pull's one remote transaction, this keeps them from pushing.
-   */
-  #pulling = false;
   /**
    * The code the editor and the model last agreed on.
    *
@@ -130,7 +114,10 @@ class CellCodeSync {
   }
 
   public update(update: ViewUpdate) {
-    if (!this.#pulling && changesCode(update)) {
+    if (
+      changesCode(update) &&
+      !update.transactions.every((tr) => tr.annotation(Transaction.remote))
+    ) {
       this.#push(update.state);
     }
     if (this.#deferred && !update.view.composing) {
@@ -164,11 +151,12 @@ class CellCodeSync {
       this.#synced = undefined;
       return;
     }
-    if (shows(state, current)) {
+    this.#initialized = true;
+    if (editorShowsCellCode(state, current)) {
       this.#synced = current;
       return;
     }
-    const next = codeOf(state);
+    const next = getCellCode(state);
     this.#synced = next;
     store.set(model.editAtom, (tx) => tx.setCode(cellId, next));
   }
@@ -185,15 +173,21 @@ class CellCodeSync {
       return;
     }
     this.#deferred = false;
-    if (next !== this.#synced && !shows(this.#view.state, next)) {
-      this.#pulling = true;
-      try {
-        replaceCode(this.#view, next);
-      } finally {
-        this.#pulling = false;
+    if (next !== this.#synced) {
+      const update = prepareCellCodeUpdate(this.#view.state, next, {
+        initialize: !this.#initialized,
+      });
+      if (update) {
+        this.#view.dispatch(update, {
+          annotations: [
+            Transaction.remote.of(true),
+            Transaction.addToHistory.of(false),
+          ],
+        });
       }
     }
     this.#synced = next;
+    this.#initialized = true;
   }
 
   /** Pull once after the current update, however many changes arrive. */
@@ -205,44 +199,6 @@ class CellCodeSync {
   }
 }
 
-/**
- * The cell code an editor state represents.
- *
- * With a language adapter, this is the adapter's Python for the editor text.
- */
-function codeOf(state: EditorState): string {
-  const text = state.doc.toString();
-  const adapter = state.field(languageAdapterState, false);
-  return adapter
-    ? adapter.transformOut(
-        text,
-        state.field(languageMetadataField, false) ?? {},
-      )[0]
-    : text;
-}
-
-/**
- * Whether an editor state shows the given cell code.
- *
- * An adapter can spell the same code more than one way, such as `mo.md("hi")`
- * and its canonical form, so a state shows code that reads into its text and
- * metadata as well as the code its own spelling produces.
- */
-function shows(state: EditorState, code: string): boolean {
-  if (codeOf(state) === code) {
-    return true;
-  }
-  const adapter = state.field(languageAdapterState, false);
-  if (!adapter) {
-    return false;
-  }
-  const [text, , metadata] = adapter.transformIn(code);
-  return (
-    text === state.doc.toString() &&
-    dequal(metadata, state.field(languageMetadataField, false))
-  );
-}
-
 /**
  * Whether an update could change the code an editor represents.
  *
@@ -258,54 +214,3 @@ function changesCode(update: ViewUpdate): boolean {
       update.startState.field(languageMetadataField, false)
   );
 }
-
-/**
- * Show cell code in the language it implies.
- *
- * The language switch happens first, in its own transactions, and clears undo
- * history because its entries edit the previous language's text. The text and
- * metadata follow in one remote transaction that undo history leaves out. An
- * SQL dialect depends on the engine in the metadata, so it is set again after.
- */
-function replaceCode(view: EditorView, code: string) {
-  let text = code;
-  const effects: StateEffect<unknown>[] = [];
-  let sql = false;
-  const adapter = view.state.field(languageAdapterState, false);
-  if (adapter) {
-    const target = languageAdapterFromCode(code);
-    sql = target.type === "sql";
-    if (target.type !== adapter.type) {
-      switchLanguage(view, { language: target.type, keepCodeAsIs: true });
-    }
-    const [inner, , metadata] = target.transformIn(code);
-    text = inner;
-    if (!dequal(metadata, view.state.field(languageMetadataField, false))) {
-      effects.push(setLanguageMetadata.of(metadata));
-    }
-  }
-  const current = view.state.doc.toString();
-  if (current === text && effects.length === 0) {
-    return;
-  }
-  view.dispatch({
-    changes: changesBetween(current, text),
-    effects,
-    annotations: [
-      Transaction.remote.of(true),
-      Transaction.addToHistory.of(false),
-    ],
-  });
-  if (sql && effects.length > 0) {
-    initializeSQLDialect(view);
-  }
-}
-
-/** The changes that turn one text into another, one per edited range. */
-function changesBetween(from: string, to: string): ChangeSpec {
-  return diff(from, to, { scanLimit: SCAN_LIMIT }).map((change) => ({
-    from: change.fromA,
-    to: change.toA,
-    insert: to.slice(change.fromB, change.toB),
-  }));
-}

package/src/core/notebook/document/hooks.tsx

@@ -12,8 +12,8 @@ const DocumentContext = createContext<DocumentModel | null>(null);
 
 /**
  * Expose a session-owned model to React. Construct the model with the same store
- * as the enclosing Jotai Provider. The session owner controls its AbortSignal;
- * mounting or unmounting this provider does not change the model's lifetime.
+ * as the enclosing Jotai Provider. The session owner starts it; mounting or
+ * unmounting this provider does not change the model's lifetime.
  *
  * ```tsx
  * <Provider store={store}>

package/src/core/notebook/document/index.ts

@@ -12,4 +12,12 @@ export {
   useNotebookDocumentEdit,
   useNotebookDocumentSync,
 } from "./hooks";
-export { createDocumentModel, type DocumentModel } from "./model";
+export {
+  createDocumentModel,
+  type DocumentEdit,
+  type DocumentEditor,
+  type DocumentModel,
+  type DocumentSyncStatus,
+  type EditableCell,
+  type EditableDocument,
+} from "./model";

package/src/core/notebook/document/model.ts

@@ -1,9 +1,19 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type * as Protocol from "@marimo-team/client/protocol";
-import type { Atom, createStore } from "jotai/vanilla";
+import {
+  type Atom,
+  atom,
+  type createStore,
+  type WritableAtom,
+} from "jotai/vanilla";
 import type { Document, NotebookBackend } from "../replica";
-import { createDocumentState, type DocumentState } from "./state";
+import type { EditableDocument } from "./draft";
+import {
+  createDocumentState,
+  type DocumentEdit,
+  type DocumentState,
+} from "./state";
 import { createDocumentSync, type DocumentSyncStatus } from "./sync";
 
 export type { DocumentEdit, DocumentEditor } from "./state";
@@ -11,45 +21,85 @@ export type { EditableCell, EditableDocument } from "./draft";
 export type { DocumentSyncStatus } from "./sync";
 
 /** The editable document and its synchronization with the host. */
-export interface DocumentModel extends Pick<
-  DocumentState,
-  "documentAtom" | "editAtom"
-> {
+export interface DocumentModel {
+  /** Confirmed host state with local drafts applied, null while unavailable or stopped. */
+  readonly documentAtom: Atom<EditableDocument | null>;
+  /** Ignored while stopped. */
+  readonly editAtom: WritableAtom<null, [edit: DocumentEdit], void>;
   readonly syncStatusAtom: Atom<DocumentSyncStatus>;
   /** Drain local edits and return the confirmed revision. */
   flush(): Promise<Protocol.Document>;
   /** Retry a rejection against fresh state, or repeat an uncertain request exactly. */
   retry(): Promise<Protocol.Document>;
+  /** Keep edits local until the returned release is called. A no-op while stopped. */
+  hold(): () => void;
+  /**
+   * Edit and synchronize until the signal aborts, which discards local drafts.
+   * The model can start again once stopped.
+   */
+  start(options: { signal: AbortSignal }): void;
 }
 
-/** Create once per editing session; abort to stop editing and synchronization together. */
+interface Running {
+  readonly state: DocumentState;
+  readonly sync: ReturnType<typeof createDocumentSync>;
+}
+
+/** Create without subscribing to anything; editing begins on start. */
 export function createDocumentModel({
   store,
-  documentAtom,
+  documentAtom: source,
   backendAtom,
-  signal,
 }: {
   store: ReturnType<typeof createStore>;
   /** The document the host last confirmed. */
   documentAtom: Atom<Document | null>;
   /** The backend that confirmed it, whose commands reach its host. */
   backendAtom: Atom<NotebookBackend | null>;
-  signal: AbortSignal;
 }): DocumentModel {
-  signal.throwIfAborted();
-  const state = createDocumentState({ store, source: documentAtom, signal });
-  const sync = createDocumentSync({
-    store,
-    documentAtom,
-    backendAtom,
-    state,
-    signal,
-  });
+  const runningAtom = atom<Running | null>(null);
+
+  function stopped(): Promise<never> {
+    return Promise.reject(new Error("The document model is not started"));
+  }
+
   return {
-    documentAtom: state.documentAtom,
-    editAtom: state.editAtom,
-    syncStatusAtom: sync.statusAtom,
-    flush: sync.flush,
-    retry: sync.retry,
+    documentAtom: atom((get) => {
+      const running = get(runningAtom);
+      return running ? get(running.state.documentAtom) : null;
+    }),
+    editAtom: atom(null, (get, set, edit: DocumentEdit) => {
+      const running = get(runningAtom);
+      if (running) {
+        set(running.state.editAtom, edit);
+      }
+    }),
+    syncStatusAtom: atom((get) => {
+      const running = get(runningAtom);
+      return running ? get(running.sync.statusAtom) : { status: "unavailable" };
+    }),
+    flush: () => store.get(runningAtom)?.sync.flush() ?? stopped(),
+    retry: () => store.get(runningAtom)?.sync.retry() ?? stopped(),
+    hold: () => store.get(runningAtom)?.sync.hold() ?? (() => undefined),
+    start({ signal }) {
+      signal.throwIfAborted();
+      if (store.get(runningAtom)) {
+        throw new Error("The document model is already started");
+      }
+      const state = createDocumentState({ store, source, signal });
+      const sync = createDocumentSync({
+        store,
+        documentAtom: source,
+        backendAtom,
+        state,
+        signal,
+      });
+      const stop = () => store.set(runningAtom, null);
+      signal.addEventListener("abort", stop, { once: true });
+      store.set(runningAtom, { state, sync });
+      if (signal.aborted) {
+        stop();
+      }
+    },
   };
 }

package/src/core/notebook/document/sync.ts

@@ -97,6 +97,13 @@ export function createDocumentSync({
   let failure: Extract<DocumentSyncStatus, { status: "error" }> | null = null;
   let timer: ReturnType<typeof setTimeout> | undefined;
   let refreshing = false;
+  // Holds keep edits local, as a legacy save needs while it writes the
+  // notebook the flush before it produced.
+  let holds = 0;
+  let releaseWaiters: Array<{
+    resolve: () => void;
+    reject: (reason: unknown) => void;
+  }> = [];
 
   function clearTimer() {
     clearTimeout(timer);
@@ -139,9 +146,11 @@ export function createDocumentSync({
         store.set(statusAtom, { status: "idle" });
       } else {
         store.set(statusAtom, { status: "pending" });
-        timer = setTimeout(() => {
-          void flush().catch(() => undefined);
-        }, DEBOUNCE_MS);
+        if (holds === 0) {
+          timer = setTimeout(() => {
+            void flush().catch(() => undefined);
+          }, DEBOUNCE_MS);
+        }
       }
     } finally {
       refreshing = false;
@@ -314,6 +323,11 @@ export function createDocumentSync({
     if (running) {
       return running;
     }
+    if (holds > 0) {
+      return new Promise<void>((resolve, reject) => {
+        releaseWaiters.push({ resolve, reject });
+      }).then(flush);
+    }
     const notebookId =
       active?.submission.notebookId ?? store.get(source)?.revision.notebook_id;
     if (notebookId === undefined) {
@@ -353,17 +367,52 @@ export function createDocumentSync({
     return flush();
   }
 
+  /**
+   * Keep edits local until the returned release is called. Flushes
+   * requested meanwhile wait for it. Holds nest, and a release is
+   * idempotent.
+   *
+   * TODO(MO-8028): Only the legacy save needs this, since the host takes
+   * the notebook it writes as the document. Saving through the host
+   * removes both.
+   */
+  function hold(): () => void {
+    holds += 1;
+    let held = true;
+    return () => {
+      if (!held) {
+        return;
+      }
+      held = false;
+      holds -= 1;
+      if (holds > 0) {
+        return;
+      }
+      const waiters = releaseWaiters;
+      releaseWaiters = [];
+      for (const waiter of waiters) {
+        waiter.resolve();
+      }
+      refresh();
+    };
+  }
+
   const unsubscribe = store.sub(observed, refresh);
   function stop() {
     clearTimer();
     unsubscribe();
     signal.removeEventListener("abort", stop);
     active?.controller.abort(signal.reason);
+    const waiters = releaseWaiters;
+    releaseWaiters = [];
+    for (const waiter of waiters) {
+      waiter.reject(signal.reason);
+    }
   }
   signal.addEventListener("abort", stop, { once: true });
   if (signal.aborted) {
     stop();
   }
   refresh();
-  return { statusAtom, flush, retry };
+  return { statusAtom, flush, retry, hold };
 }

package/src/core/run-app.tsx

@@ -17,6 +17,7 @@ import {
   useCellActions,
 } from "./cells/cells";
 import type { AppConfig } from "./config/config-schema";
+import { DocumentProvider } from "./notebook/document";
 import { RuntimeState } from "./kernel/RuntimeState";
 import { getSessionId } from "./kernel/session";
 import { connectionAtom } from "./network/connection";
@@ -51,7 +52,7 @@ export const RunApp: React.FC<AppProps> = ({ appConfig, hideHeader }) => {
     };
   }, [sendComponentValues]);
 
-  const { connection, reconnect } = useNotebookSession({
+  const { connection, reconnect, documentModel } = useNotebookSession({
     autoInstantiate: true,
     setCells: setCells,
     sessionId: getSessionId(),
@@ -73,47 +74,49 @@ export const RunApp: React.FC<AppProps> = ({ appConfig, hideHeader }) => {
   })();
 
   return (
-    <AppContainer
-      connection={connection}
-      isRunning={isRunning}
-      width={appConfig.width}
-      onReconnect={reconnect}
-    >
-      <KernelStartupErrorModal />
-      <AppHeader
+    <DocumentProvider model={documentModel}>
+      <AppContainer
         connection={connection}
-        className={hideHeader ? "hidden" : "sm:pt-8"}
+        isRunning={isRunning}
+        width={appConfig.width}
+        onReconnect={reconnect}
       >
-        {galleryHref && (
-          <div className="flex items-center px-6 pt-4 sm:-mt-8">
-            <a
-              href={galleryHref}
-              aria-label="Back to gallery"
-              className={cn(
-                buttonVariants({ variant: "text", size: "sm" }),
-                "gap-2 px-0 text-muted-foreground hover:text-foreground",
-              )}
-            >
-              <ArrowLeftIcon className="size-4" aria-hidden={true} />
-              <span>Back</span>
-            </a>
-          </div>
-        )}
-      </AppHeader>
-      <ProgressiveBoundary
-        requires={canPaintRunAppAtom}
-        delay={2000}
-        fallback={
-          <>
-            <Spinner className="mx-auto" />
-            <p className="text-center text-sm text-muted-foreground mt-2">
-              Connecting...
-            </p>
-          </>
-        }
-      >
-        <CellsRenderer appConfig={appConfig} mode="read" />
-      </ProgressiveBoundary>
-    </AppContainer>
+        <KernelStartupErrorModal />
+        <AppHeader
+          connection={connection}
+          className={hideHeader ? "hidden" : "sm:pt-8"}
+        >
+          {galleryHref && (
+            <div className="flex items-center px-6 pt-4 sm:-mt-8">
+              <a
+                href={galleryHref}
+                aria-label="Back to gallery"
+                className={cn(
+                  buttonVariants({ variant: "text", size: "sm" }),
+                  "gap-2 px-0 text-muted-foreground hover:text-foreground",
+                )}
+              >
+                <ArrowLeftIcon className="size-4" aria-hidden={true} />
+                <span>Back</span>
+              </a>
+            </div>
+          )}
+        </AppHeader>
+        <ProgressiveBoundary
+          requires={canPaintRunAppAtom}
+          delay={2000}
+          fallback={
+            <>
+              <Spinner className="mx-auto" />
+              <p className="text-center text-sm text-muted-foreground mt-2">
+                Connecting...
+              </p>
+            </>
+          }
+        >
+          <CellsRenderer appConfig={appConfig} mode="read" />
+        </ProgressiveBoundary>
+      </AppContainer>
+    </DocumentProvider>
   );
 };

package/src/core/saving/save-component.tsx

@@ -21,10 +21,7 @@ import { Label } from "../../components/ui/label";
 import { useEvent } from "../../hooks/useEvent";
 import { Logger } from "../../utils/Logger";
 import { getCellConfigs, getNotebook, useNotebook } from "../cells/cells";
-import {
-  flushDocumentChanges,
-  withDocumentSave,
-} from "../cells/document-changes";
+import { flushDocument, withDocumentSave } from "../cells/document";
 import { notebookCells } from "../cells/utils";
 import { formatAll } from "../codemirror/format";
 import { autoSaveConfigAtom } from "../config/config";
@@ -165,7 +162,7 @@ export function useSaveNotebook() {
       // TODO(MO-7955): Remove this conditional once all frontend saves use
       // the host protocol. Host saves cannot recover edits from an editor snapshot.
       const savedNotebook = isWasm()
-        ? await flushDocumentChanges().then(save)
+        ? await flushDocument().then(save)
         : await withDocumentSave(save);
       setLastSavedNotebook(savedNotebook);
     }),

package/src/core/session/__tests__/useNotebookSession.hook.test.tsx

@@ -52,7 +52,10 @@ import { RuntimeManager } from "../../runtime/runtime";
 import { WebSocketClosedReason, ConnectionState } from "../types";
 import { classifyCloseEvent } from "../../transports/connection-state";
 import { createEditorBackend } from "../backend";
-import { fakeBackend } from "../../notebook/replica/__tests__/fixtures";
+import {
+  fakeBackend,
+  fixtureMessages,
+} from "../../notebook/replica/__tests__/fixtures";
 import { useNotebookSession } from "../useNotebookSession";
 import { type NotebookTransport, StreamStopped } from "@marimo-team/client";
 
@@ -142,7 +145,10 @@ function makeRuntimeManager(
   return runtime;
 }
 
-function renderConnectionHook(store: ReturnType<typeof createStore>) {
+function renderConnectionHook(
+  store: ReturnType<typeof createStore>,
+  { reactStrictMode = false } = {},
+) {
   const wrapper: React.FC<React.PropsWithChildren> = ({ children }) => (
     <JotaiProvider store={store}>
       <ErrorBoundary fallback={null}>{children}</ErrorBoundary>
@@ -155,7 +161,7 @@ function renderConnectionHook(store: ReturnType<typeof createStore>) {
         autoInstantiate: false,
         setCells: vi.fn(),
       }),
-    { wrapper },
+    { wrapper, reactStrictMode },
   );
 }
 
@@ -228,6 +234,60 @@ describe("useNotebookSession.reconnect()", () => {
   });
 });
 
+describe("useNotebookSession document model", () => {
+  it("keeps drafts while the connection drops and discards them when the session ends", async () => {
+    const fake = withTransport();
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
+    const store = createStore();
+    const { result, unmount } = renderConnectionHook(store, {
+      reactStrictMode: true,
+    });
+    await started();
+    act(() => {
+      connection(fake).open();
+      fixtureMessages("notebook-events.sse").forEach(connection(fake).receive);
+    });
+    const model = result.current.documentModel;
+    const code = () => store.get(model.documentAtom)?.cells.get("c1")?.code;
+    act(() => store.set(model.editAtom, (tx) => tx.setCode("c1", "local")));
+    expect(code()).toBe("local");
+
+    await act(async () => connection(fake).drop());
+    expect(result.current.documentModel).toBe(model);
+    expect(code()).toBe("local");
+
+    unmount();
+    expect(code()).toBeUndefined();
+    await expect(model.flush()).rejects.toThrow("not started");
+  });
+
+  it("follows the host document into the notebook's cells", async () => {
+    const fake = withTransport();
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
+    const store = createStore();
+    const { result, unmount } = renderConnectionHook(store);
+    await started();
+    act(() => {
+      connection(fake).open();
+      fixtureMessages("notebook-events.sse").forEach(connection(fake).receive);
+    });
+
+    let notebook = store.get(notebookAtom);
+    expect(notebook.cellIds.inOrderIds).toEqual(["c1", "c3", "k7f2"]);
+    expect(notebook.cellData[cellId("k7f2")]).toMatchObject({
+      code: "y = x * 2",
+      config: { hide_code: true },
+    });
+
+    const model = result.current.documentModel;
+    act(() => store.set(model.editAtom, (tx) => tx.setCode("c3", "x = 42")));
+    notebook = store.get(notebookAtom);
+    expect(notebook.cellData[cellId("c3")].code).toBe("x = 42");
+
+    unmount();
+  });
+});
+
 describe("useNotebookSession messages", () => {
   let store: ReturnType<typeof createStore>;
   let fake: FakeTransport;