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

package/src/core/session/active-notebook.ts

@@ -1,7 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { atom } from "jotai";
+import type { DocumentModel } from "../notebook/document";
 import type { NotebookReplica } from "../notebook/replica";
 
 /** The replica owned by the mounted notebook session. */
 export const activeNotebookReplicaAtom = atom<NotebookReplica | null>(null);
+
+/** The document model owned by the mounted notebook session. */
+export const activeDocumentModelAtom = atom<DocumentModel | null>(null);

package/src/core/session/useNotebookSession.tsx

@@ -2,16 +2,13 @@
 
 import { invalidatePackageData } from "@/core/packages/package-data";
 import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
-import { useLayoutEffect } from "react";
+import { useEffect, useLayoutEffect, useMemo } from "react";
 import { useErrorBoundary } from "react-error-boundary";
 import { toast } from "@/components/ui/use-toast";
 import { getNotebook, useCellActions } from "@/core/cells/cells";
-import { applyTransactionChanges } from "@/core/cells/document-changes";
 import { AUTOCOMPLETER } from "@/core/codemirror/completion/Autocompleter";
-import type {
-  NotificationMessageData,
-  NotificationPayload,
-} from "@/core/kernel/messages";
+import { createDocumentModel } from "@/core/notebook/document";
+import type { NotificationPayload } from "@/core/kernel/messages";
 import { useNotebookReplica } from "@/core/session/useNotebookReplica";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
@@ -75,6 +72,7 @@ import {
 import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
+import { activeDocumentModelAtom } from "./active-notebook";
 import { WebSocketClosedReason, ConnectionState } from "./types";
 
 const SUPPORTS_LAZY_KERNELS = true;
@@ -100,17 +98,7 @@ export function useNotebookSession(opts: {
   const { showBoundary } = useErrorBoundary();
 
   const { handleCellMessage } = useCellActions();
-  const actionsWithoutMiddleware = useCellActions({ skipMiddleware: true });
-
-  const handleDocumentTransaction = (
-    transaction: NotificationMessageData<"notebook-document-transaction">["transaction"],
-  ) => {
-    applyTransactionChanges(
-      transaction.changes,
-      actionsWithoutMiddleware,
-      () => getNotebook().cellIds.inOrderIds,
-    );
-  };
+  const { reconcileDocument } = useCellActions({ skipMiddleware: true });
   const { addCellNotification } = useRunsActions();
   const setKernelState = useSetAtom(kernelStateAtom);
   const setInitialRunCompleted = useSetAtom(initialRunCompletedAtom);
@@ -395,7 +383,7 @@ export function useNotebookSession(opts: {
         );
         return;
       case "notebook-document-transaction":
-        handleDocumentTransaction(msg.data.transaction);
+        // Edits reach the notebook through the host document.
         return;
       case "consumer-capabilities":
         setKioskMode(!msg.data.consumer_capabilities.edit);
@@ -440,6 +428,38 @@ export function useNotebookSession(opts: {
     },
   });
 
+  const documentModel = useMemo(
+    () =>
+      createDocumentModel({
+        store,
+        documentAtom: replica.documentAtom,
+        backendAtom: replica.backendAtom,
+      }),
+    [store, replica],
+  );
+  // Edits outlive reconnects and end with the session. The notebook's cells
+  // follow the document for as long as the session is mounted.
+  useEffect(() => {
+    const controller = new AbortController();
+    documentModel.start({ signal: controller.signal });
+    store.set(activeDocumentModelAtom, documentModel);
+    const follow = () => {
+      const document = store.get(documentModel.documentAtom);
+      if (document !== null) {
+        reconcileDocument(document);
+      }
+    };
+    const unsubscribe = store.sub(documentModel.documentAtom, follow);
+    follow();
+    return () => {
+      unsubscribe();
+      controller.abort();
+      if (store.get(activeDocumentModelAtom) === documentModel) {
+        store.set(activeDocumentModelAtom, null);
+      }
+    };
+  }, [documentModel, store, reconcileDocument]);
+
   // Subscribe before passive connection setup: static/WASM transports can emit synchronously.
   useLayoutEffect(() => {
     const updateConnection = () => {
@@ -485,5 +505,5 @@ export function useNotebookSession(opts: {
     updateStartupProgress,
   ]);
 
-  return { replica, connection, reconnect: replica.reconnect };
+  return { replica, documentModel, connection, reconnect: replica.reconnect };
 }

package/src/core/wasm/__tests__/bridge.test.ts

@@ -101,9 +101,9 @@ vi.mock("@/core/wasm/store", () => ({
   notebookFileStore: { readFile: mockNotebookReadFile, saveFile: vi.fn() },
 }));
 
-vi.mock("@/core/cells/document-changes", () => ({
-  flushDocumentChanges: mockFlush,
-  documentTransactionMiddleware: vi.fn(),
+vi.mock("@/core/cells/document", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/core/cells/document")>()),
+  flushDocument: mockFlush,
 }));
 
 // Import after all mocks are set up

package/src/core/wasm/bridge.ts

@@ -13,7 +13,7 @@ import { once } from "@/utils/once";
 import { generateUUID } from "@/utils/uuid";
 import { createModuleWorker } from "@/utils/worker";
 import { notebookIsRunningAtom } from "../cells/cells";
-import { flushDocumentChanges } from "../cells/document-changes";
+import { flushDocument } from "../cells/document";
 import { createWasmNotebookBackend } from "../transports/wasm";
 import type { ClientCommand } from "../kernel/messages";
 import { getInitialAppMode } from "../mode";
@@ -280,7 +280,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   }
 
   public sendSave: EditRequests["sendSave"] = async ({ layout }) => {
-    await flushDocumentChanges();
+    await flushDocument();
     await this.rpc.proxy.request.setNotebookLayout(layout ?? null);
     await this.saveHostNotebook();
     return null;
@@ -565,7 +565,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   public exportAsHTML: EditRequests["exportAsHTML"] = async (
     request: ExportAsHTMLRequest,
   ) => {
-    await flushDocumentChanges();
+    await flushDocument();
     const response = await this.rpc.proxy.request.bridge({
       functionName: "export_html",
       payload: withDevAssetUrl(request),
@@ -576,7 +576,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   public exportAsMarkdown: EditRequests["exportAsMarkdown"] = async (
     request: ExportAsMarkdownRequest,
   ) => {
-    await flushDocumentChanges();
+    await flushDocument();
     const response = await this.rpc.proxy.request.bridge({
       functionName: "export_markdown",
       payload: request,
@@ -587,7 +587,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   public exportAsScript: EditRequests["exportAsScript"] = async (
     request: ExportAsScriptRequest,
   ) => {
-    await flushDocumentChanges();
+    await flushDocument();
     const response = await this.rpc.proxy.request.bridge({
       functionName: "export_script",
       payload: request,
@@ -658,13 +658,6 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return null;
   };
 
-  public sendDocumentTransaction: EditRequests["sendDocumentTransaction"] =
-    async (request) => {
-      await this.initialized.promise;
-      await this.rpc.proxy.request.documentTransaction(request);
-      return null;
-    };
-
   public addPackage: EditRequests["addPackage"] = async (request) => {
     return this.rpc.proxy.request.addPackage(request);
   };

package/src/core/wasm/worker/worker.ts

@@ -12,7 +12,6 @@ import type { UserConfig } from "@/core/config/config-schema";
 import type {
   ListPackagesResponse,
   PackageOperationResponse,
-  NotebookDocumentTransactionRequest,
   SaveNotebookRequest,
 } from "@/core/network/types";
 import { decodeUtf8 } from "@/utils/strings";
@@ -302,14 +301,6 @@ const requestHandler = createRPCRequestHandler({
     };
   },
 
-  documentTransaction: async (request: NotebookDocumentTransactionRequest) => {
-    const notebook = await getNotebook();
-    try {
-      await notebook.transaction(JSON.stringify(request));
-    } finally {
-      notebook.destroy();
-    }
-  },
   setNotebookLayout: async (layout: SaveNotebookRequest["layout"]) => {
     const notebook = await getNotebook();
     try {

package/src/utils/__tests__/{createReducer.test.ts → createReducer.test.tsx}

@@ -1,4 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+import { renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import type { PropsWithChildren } from "react";
 import { describe, expect, it, vi } from "vitest";
 import { Logger } from "@/utils/Logger";
 import { createReducer, createReducerAndAtoms } from "../createReducer";
@@ -140,6 +143,31 @@ describe("createReducerAndAtoms", () => {
     expect(middleware).toHaveBeenCalledWith(state, { count: 5 }, action);
   });
 
+  it("runs middleware after the write commits, so it sees the new state", () => {
+    const initialState = () => ({ count: 0 });
+    const reducers = {
+      increment: (state: State, payload: number) => ({
+        count: state.count + payload,
+      }),
+    };
+    const seen: number[] = [];
+    const { valueAtom, useActions } = createReducerAndAtoms(
+      initialState,
+      reducers,
+      [(_prev, next) => seen.push(store.get(valueAtom).count, next.count)],
+    );
+    const store = createStore();
+    const wrapper = ({ children }: PropsWithChildren) => (
+      <Provider store={store}>{children}</Provider>
+    );
+    const { result } = renderHook(() => useActions(), { wrapper });
+
+    result.current.increment(5);
+
+    expect(seen).toEqual([5, 5]);
+    expect(store.get(valueAtom)).toEqual({ count: 5 });
+  });
+
   it("should log an error for non-existent action types", () => {
     const initialState = () => ({ count: 0 });
     const reducers = {

package/src/utils/createReducer.ts

@@ -104,49 +104,70 @@ export function createReducerAndAtoms<
   };
   const { reducer, createActions } = createReducer(initialState, reducers);
 
-  const reducerWithMiddleware = (state: State, action: ReducerActionOf<RH>) => {
+  const reduceSafely = (state: State, action: ReducerActionOf<RH>) => {
     try {
-      const newState = reducer(state, action);
-      for (const mw of allMiddleware) {
-        try {
-          mw(state, newState, action);
-        } catch (error) {
-          Logger.error(`Error in middleware for action ${action.type}:`, error);
-        }
-      }
-      return newState;
+      return reducer(state, action);
     } catch (error) {
       Logger.error(`Error in reducer for action ${action.type}:`, error);
       return state;
     }
   };
 
+  const runMiddleware = (
+    prev: State,
+    next: State,
+    action: ReducerActionOf<RH>,
+  ) => {
+    for (const mw of allMiddleware) {
+      try {
+        mw(prev, next, action);
+      } catch (error) {
+        Logger.error(`Error in middleware for action ${action.type}:`, error);
+      }
+    }
+  };
+
+  const reducerWithMiddleware = (state: State, action: ReducerActionOf<RH>) => {
+    const next = reduceSafely(state, action);
+    runMiddleware(state, next, action);
+    return next;
+  };
+
   const valueAtom = atom(initialState());
-  // map of SetAtom => Actions
+  // map of SetAtom => Actions, with and without middleware
   const actionsMap = new WeakMap();
+  const plainActionsMap = new WeakMap();
 
   function useActions(
     options: { skipMiddleware?: boolean } = {},
   ): ReducerActions<RH> {
     const setState = useSetAtom(valueAtom);
+    const skipMiddleware = options.skipMiddleware === true;
+    const cache = skipMiddleware ? plainActionsMap : actionsMap;
 
-    if (options.skipMiddleware === true) {
-      return createActions((action: ReducerActionOf<RH>) => {
-        setState((state: State) => reducer(state, action));
-      });
-    }
-
-    if (!actionsMap.has(setState)) {
-      actionsMap.set(
+    if (!cache.has(setState)) {
+      cache.set(
         setState,
         createActions((action: ReducerActionOf<RH>) => {
-          setState((state: State) => reducerWithMiddleware(state, action));
+          // Reduce inside the write. Middleware runs once the write has
+          // committed and its subscribers have run, so anything it sets in
+          // turn observes the new state rather than the one being replaced.
+          let prev: State | undefined;
+          let next: State | undefined;
+          setState((state: State) => {
+            prev = state;
+            next = reduceSafely(state, action);
+            return next;
+          });
+          if (!skipMiddleware && prev !== undefined && next !== undefined) {
+            runMiddleware(prev, next, action);
+          }
         }),
       );
     }
 
     // oxlint-disable-next-line typescript/no-non-null-assertion
-    return actionsMap.get(setState)!;
+    return cache.get(setState)!;
   }
 
   return {