@marimo-team/islands 0.21.2-dev57 → 0.21.2-dev59

package/src/core/network/types.ts

@@ -80,6 +80,8 @@ export interface SetCellConfigRequest {
 export type UpdateUIElementRequest = schemas["UpdateUIElementRequest"];
 export type ModelRequest = schemas["ModelRequest"];
 export type UpdateCellIdsRequest = schemas["UpdateCellIdsRequest"];
+export type NotebookDocumentTransactionRequest =
+  schemas["NotebookDocumentTransactionRequest"];
 export type UpdateUserConfigRequest = schemas["UpdateUserConfigRequest"];
 export type ShutdownSessionRequest = schemas["ShutdownSessionRequest"];
 export type Snippet = schemas["Snippet"];
@@ -138,6 +140,9 @@ export interface EditRequests {
   saveCellConfig: (request: SetCellConfigRequest) => Promise<null>;
   sendRestart: () => Promise<null>;
   syncCellIds: (request: UpdateCellIdsRequest) => Promise<null>;
+  sendDocumentTransaction: (
+    request: NotebookDocumentTransactionRequest,
+  ) => Promise<null>;
   sendInstallMissingPackages: (
     request: InstallPackagesRequest,
   ) => Promise<null>;

package/src/core/wasm/bridge.ts

@@ -549,6 +549,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   };
 
   syncCellIds = () => Promise.resolve(null);
+  sendDocumentTransaction = () => Promise.resolve(null);
 
   addPackage: EditRequests["addPackage"] = async (request) => {
     return this.rpc.proxy.request.addPackage(request);

package/src/core/websocket/useMarimoKernelConnection.tsx

@@ -5,8 +5,12 @@ import { useRef } 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 { NotificationPayload } from "@/core/kernel/messages";
+import type {
+  NotificationMessageData,
+  NotificationPayload,
+} from "@/core/kernel/messages";
 import { useConnectionTransport } from "@/core/websocket/useWebSocket";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
@@ -93,6 +97,17 @@ export function useMarimoKernelConnection(opts: {
   const { showBoundary } = useErrorBoundary();
 
   const { handleCellMessage, setCellCodes, setCellIds } = useCellActions();
+  const actionsWithoutMiddleware = useCellActions({ skipMiddleware: true });
+
+  const handleDocumentTransaction = (
+    transaction: NotificationMessageData<"notebook-document-transaction">["transaction"],
+  ) => {
+    applyTransactionChanges(
+      transaction.changes,
+      actionsWithoutMiddleware,
+      () => getNotebook().cellIds.inOrderIds,
+    );
+  };
   const { addCellNotification } = useRunsActions();
   const setKernelState = useSetAtom(kernelStateAtom);
   const setAppConfig = useSetAppConfig();
@@ -321,6 +336,9 @@ export function useMarimoKernelConnection(opts: {
       case "update-cell-ids":
         setCellIds({ cellIds: msg.data.cell_ids });
         return;
+      case "notebook-document-transaction":
+        handleDocumentTransaction(msg.data.transaction);
+        return;
       default:
         logNever(msg.data);
     }

package/src/utils/createReducer.ts

@@ -27,6 +27,13 @@ type ReducerActions<RH extends ReducerHandlers<any>> = {
     : never;
 };
 
+/** Helper for typing middleware that receives dispatched actions. */
+export type DispatchedActionOf<T> = {
+  [Key in keyof T]: T[Key] extends (payload: infer P) => any
+    ? { type: Key; payload: P }
+    : never;
+}[keyof T & string];
+
 export interface ReducerCreatorResult<
   State,
   RH extends ReducerHandlers<State>,
@@ -80,21 +87,20 @@ export function createReducerAndAtoms<
   State,
   RH extends ReducerHandlers<NoInfer<State>>,
 >(initialState: () => State, reducers: RH, middleware?: Middleware<State>[]) {
+  const allMiddleware = [...(middleware ?? [])];
+  const addMiddleware = (mw: Middleware<State>) => {
+    allMiddleware.push(mw);
+  };
   const { reducer, createActions } = createReducer(initialState, reducers);
 
   const reducerWithMiddleware = (state: State, action: ReducerAction<any>) => {
     try {
       const newState = reducer(state, action);
-      if (middleware) {
-        for (const mw of middleware) {
-          try {
-            mw(state, newState, action);
-          } catch (error) {
-            Logger.error(
-              `Error in middleware for action ${action.type}:`,
-              error,
-            );
-          }
+      for (const mw of allMiddleware) {
+        try {
+          mw(state, newState, action);
+        } catch (error) {
+          Logger.error(`Error in middleware for action ${action.type}:`, error);
         }
       }
       return newState;
@@ -108,9 +114,17 @@ export function createReducerAndAtoms<
   // map of SetAtom => Actions
   const actionsMap = new WeakMap();
 
-  function useActions(): ReducerActions<RH> {
+  function useActions(
+    options: { skipMiddleware?: boolean } = {},
+  ): ReducerActions<RH> {
     const setState = useSetAtom(valueAtom);
 
+    if (options.skipMiddleware === true) {
+      return createActions((action: ReducerAction<any>) => {
+        setState((state: State) => reducer(state, action));
+      });
+    }
+
     if (!actionsMap.has(setState)) {
       actionsMap.set(
         setState,
@@ -126,6 +140,7 @@ export function createReducerAndAtoms<
 
   return {
     reducer: reducerWithMiddleware,
+    addMiddleware,
     createActions,
     valueAtom,
     useActions,