@marimo-team/frontend 0.24.1-dev80 → 0.24.1-dev82

package/src/core/cells/__tests__/document-changes.test.ts

@@ -33,14 +33,20 @@ import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { requestClientAtom } from "@/core/network/requests";
 import type { EditRequests, RunRequests } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
+import { Deferred } from "@/utils/Deferred";
 import { MultiColumn } from "@/utils/id-tree";
 import { exportedForTesting, type NotebookState } from "../cells";
 import {
   type CellAction,
+  abortDocumentResync,
+  beginDocumentResync,
   coalesceChanges,
+  completeDocumentResync,
   type DocumentChange,
+  flushDocumentChanges,
   exportedForTesting as middlewareExports,
   toDocumentChanges,
+  withDocumentSave,
 } from "../document-changes";
 import { CellId } from "../ids";
 
@@ -698,15 +704,16 @@ describe("coalesceChanges", () => {
   });
 });
 
-/**
- * End-to-end of the debounced flush path: dispatch through the real middleware
- * reducer, let the debounce fire, and inspect the batch that reaches
- * sendDocumentTransaction. Coalescing only matters once edits actually share a
- * debounce window, which drainChanges-based tests can't observe.
- */
-describe("flush path coalescing", () => {
+describe("document transaction middleware", () => {
   let sent: DocumentChange[][];
 
+  const updateCode = (cellId: CellId, code: string) => {
+    state = dispatch(state, {
+      type: "updateCellCode",
+      payload: { cellId, code, formattingChange: false },
+    });
+  };
+
   beforeEach(() => {
     vi.useFakeTimers();
     sent = [];
@@ -732,54 +739,364 @@ describe("flush path coalescing", () => {
     vi.useRealTimers();
   });
 
-  it("coalesces edit + delete of the same cell within one debounce window", () => {
+  it("coalesces edit + delete of the same cell within one debounce window", async () => {
     setup("x = 1");
     const [x] = state.cellIds.inOrderIds;
     // Discard the create-cell from setup; it would flush as its own batch.
     middlewareExports.cancelPendingChanges();
 
-    // Same window: edit the cell, then delete it, with no timer advance.
-    state = dispatch(state, {
-      type: "updateCellCode",
-      payload: { cellId: x, code: "x = 2", formattingChange: false },
-    });
+    updateCode(x, "x = 2");
     state = dispatch(state, {
       type: "deleteCell",
       payload: { cellId: x },
     });
 
-    // Debounced: nothing on the wire yet.
     expect(sent).toEqual([]);
 
-    vi.advanceTimersByTime(400);
+    await vi.advanceTimersByTimeAsync(400);
 
-    // One transaction, and it is just the delete — the conflicting set-code
-    // has been dropped, so the server never sees edit+delete of one cell.
     expect(sent).toHaveLength(1);
     expect(sent[0]).toEqual([{ type: "delete-cell", cellId: x }]);
   });
 
-  it("sends edit and delete as separate clean transactions across windows", () => {
+  it("sends edit and delete as separate clean transactions across windows", async () => {
     setup("x = 1");
     const [x] = state.cellIds.inOrderIds;
     middlewareExports.cancelPendingChanges();
 
-    state = dispatch(state, {
-      type: "updateCellCode",
-      payload: { cellId: x, code: "x = 2", formattingChange: false },
-    });
-    vi.advanceTimersByTime(400);
+    updateCode(x, "x = 2");
+    await vi.advanceTimersByTimeAsync(400);
 
     state = dispatch(state, {
       type: "deleteCell",
       payload: { cellId: x },
     });
-    vi.advanceTimersByTime(400);
+    await vi.advanceTimersByTimeAsync(400);
 
-    // Two separate, individually-valid transactions: neither conflicts, which
-    // is why the bug only ever surfaced when both landed in the same window.
     expect(sent).toHaveLength(2);
     expect(sent[0]).toEqual([{ type: "set-code", cellId: x, code: "x = 2" }]);
     expect(sent[1]).toEqual([{ type: "delete-cell", cellId: x }]);
   });
+
+  it("flushes pending changes before a dependent operation", async () => {
+    setup("x = 1");
+    const [x] = state.cellIds.inOrderIds;
+    middlewareExports.cancelPendingChanges();
+    updateCode(x, "x = 2");
+
+    await flushDocumentChanges();
+
+    expect(sent).toEqual([[{ type: "set-code", cellId: x, code: "x = 2" }]]);
+  });
+
+  it("keeps a failed transaction sticky before a dependent operation", async () => {
+    const failingClient: Pick<
+      EditRequests & RunRequests,
+      "sendDocumentTransaction"
+    > = {
+      sendDocumentTransaction: async ({ changes }) => {
+        sent.push(changes);
+        throw new Error("connection lost");
+      },
+    };
+    store.set(
+      requestClientAtom,
+      failingClient as unknown as EditRequests & RunRequests,
+    );
+
+    setup("x = 1");
+    const [x] = state.cellIds.inOrderIds;
+    middlewareExports.cancelPendingChanges();
+    updateCode(x, "x = 2");
+    await vi.advanceTimersByTimeAsync(400);
+
+    updateCode(x, "x = 3");
+    await expect(flushDocumentChanges()).rejects.toThrow("connection lost");
+
+    expect(sent).toEqual([[{ type: "set-code", cellId: x, code: "x = 2" }]]);
+  });
+
+  it("resumes after a full document save establishes a fresh baseline", async () => {
+    let attempt = 0;
+    const recoveringClient: Pick<
+      EditRequests & RunRequests,
+      "sendDocumentTransaction"
+    > = {
+      sendDocumentTransaction: async ({ changes }) => {
+        sent.push(changes);
+        attempt += 1;
+        if (attempt === 1) {
+          throw new Error("connection lost");
+        }
+        return null;
+      },
+    };
+    store.set(
+      requestClientAtom,
+      recoveringClient as unknown as EditRequests & RunRequests,
+    );
+
+    setup("x = 1");
+    const [x] = state.cellIds.inOrderIds;
+    middlewareExports.cancelPendingChanges();
+    updateCode(x, "x = 2");
+    await vi.advanceTimersByTimeAsync(400);
+
+    updateCode(x, "x = 3");
+    await withDocumentSave(async () => {
+      updateCode(x, "x = 4");
+    });
+
+    await vi.advanceTimersByTimeAsync(400);
+    expect(sent).toEqual([
+      [{ type: "set-code", cellId: x, code: "x = 2" }],
+      [{ type: "set-code", cellId: x, code: "x = 4" }],
+    ]);
+  });
+
+  it("restores unsent edits when a full document save fails", async () => {
+    const firstStarted = new Deferred<void>();
+    const firstRelease = new Deferred<void>();
+    const failingClient: Pick<
+      EditRequests & RunRequests,
+      "sendDocumentTransaction"
+    > = {
+      sendDocumentTransaction: async ({ changes }) => {
+        sent.push(changes);
+        if (sent.length === 1) {
+          firstStarted.resolve();
+          await firstRelease.promise;
+          throw new Error("connection lost");
+        }
+        return null;
+      },
+    };
+    store.set(
+      requestClientAtom,
+      failingClient as unknown as EditRequests & RunRequests,
+    );
+
+    setup("x = 1");
+    const [x] = state.cellIds.inOrderIds;
+    middlewareExports.cancelPendingChanges();
+    updateCode(x, "x = 2");
+    vi.advanceTimersByTime(400);
+    await firstStarted.promise;
+
+    updateCode(x, "x = 3");
+    await vi.advanceTimersByTimeAsync(400);
+    firstRelease.resolve();
+    await expect(flushDocumentChanges()).rejects.toThrow("connection lost");
+
+    updateCode(x, "x = 4");
+    const firstResync = beginDocumentResync();
+    abortDocumentResync(firstResync);
+    const secondResync = beginDocumentResync();
+    expect(secondResync?.includedTransactions).toEqual([
+      [{ type: "set-code", cellId: x, code: "x = 3" }],
+    ]);
+    expect(secondResync?.includedChanges).toEqual([
+      { type: "set-code", cellId: x, code: "x = 4" },
+    ]);
+    await completeDocumentResync(secondResync);
+
+    updateCode(x, "x = 5");
+    await vi.advanceTimersByTimeAsync(400);
+    expect(sent.at(-1)).toEqual([
+      { type: "set-code", cellId: x, code: "x = 5" },
+    ]);
+  });
+
+  it("keeps edits debounced while an earlier transaction is in flight", async () => {
+    const firstRelease = new Deferred<void>();
+    const firstStarted = new Deferred<void>();
+    const delayedClient: Pick<
+      EditRequests & RunRequests,
+      "sendDocumentTransaction"
+    > = {
+      sendDocumentTransaction: async ({ changes }) => {
+        sent.push(changes);
+        if (sent.length === 1) {
+          firstStarted.resolve();
+          await firstRelease.promise;
+        }
+        return null;
+      },
+    };
+    store.set(
+      requestClientAtom,
+      delayedClient as unknown as EditRequests & RunRequests,
+    );
+
+    setup("x = 1");
+    const [x] = state.cellIds.inOrderIds;
+    middlewareExports.cancelPendingChanges();
+    updateCode(x, "x = 2");
+    vi.advanceTimersByTime(400);
+    await firstStarted.promise;
+
+    updateCode(x, "x = 3");
+    firstRelease.resolve();
+    await vi.advanceTimersByTimeAsync(399);
+    expect(sent).toHaveLength(1);
+
+    await vi.advanceTimersByTimeAsync(1);
+    expect(sent).toEqual([
+      [{ type: "set-code", cellId: x, code: "x = 2" }],
+      [{ type: "set-code", cellId: x, code: "x = 3" }],
+    ]);
+  });
+
+  it("waits for every transaction queued while flushing", async () => {
+    const releases = Array.from({ length: 3 }, () => new Deferred<void>());
+    const starts = Array.from({ length: 3 }, () => new Deferred<void>());
+    const orderedClient: Pick<
+      EditRequests & RunRequests,
+      "sendDocumentTransaction"
+    > = {
+      sendDocumentTransaction: async ({ changes }) => {
+        const index = sent.length;
+        sent.push(changes);
+        starts[index].resolve();
+        await releases[index].promise;
+        return null;
+      },
+    };
+    store.set(
+      requestClientAtom,
+      orderedClient as unknown as EditRequests & RunRequests,
+    );
+
+    setup("x = 1");
+    const [x] = state.cellIds.inOrderIds;
+    middlewareExports.cancelPendingChanges();
+    updateCode(x, "x = 2");
+    vi.advanceTimersByTime(400);
+    await starts[0].promise;
+
+    const flushed = flushDocumentChanges();
+
+    updateCode(x, "x = 3");
+    await vi.advanceTimersByTimeAsync(400);
+    expect(sent).toHaveLength(1);
+
+    releases[0].resolve();
+    await starts[1].promise;
+
+    updateCode(x, "x = 4");
+    await vi.advanceTimersByTimeAsync(400);
+    expect(sent).toHaveLength(2);
+
+    let flushResolved = false;
+    void flushed.then(() => {
+      flushResolved = true;
+    });
+    releases[1].resolve();
+    await starts[2].promise;
+    expect(flushResolved).toBe(false);
+
+    releases[2].resolve();
+    await flushed;
+    expect(sent).toEqual([
+      [{ type: "set-code", cellId: x, code: "x = 2" }],
+      [{ type: "set-code", cellId: x, code: "x = 3" }],
+      [{ type: "set-code", cellId: x, code: "x = 4" }],
+    ]);
+  });
+
+  it("orders a full save between earlier and later transactions", async () => {
+    const firstStarted = new Deferred<void>();
+    const firstRelease = new Deferred<void>();
+    const saveStarted = new Deferred<void>();
+    const saveRelease = new Deferred<void>();
+    const events: string[] = [];
+    const delayedClient: Pick<
+      EditRequests & RunRequests,
+      "sendDocumentTransaction"
+    > = {
+      sendDocumentTransaction: async ({ changes }) => {
+        const change = changes[0];
+        events.push("code" in change ? change.code : change.type);
+        if (events.length === 1) {
+          firstStarted.resolve();
+          await firstRelease.promise;
+        }
+        return null;
+      },
+    };
+    store.set(
+      requestClientAtom,
+      delayedClient as unknown as EditRequests & RunRequests,
+    );
+
+    setup("x = 1");
+    const [x] = state.cellIds.inOrderIds;
+    middlewareExports.cancelPendingChanges();
+    updateCode(x, "x = 2");
+    vi.advanceTimersByTime(400);
+    await firstStarted.promise;
+
+    const saving = withDocumentSave(async () => {
+      events.push("save");
+      saveStarted.resolve();
+      await saveRelease.promise;
+    });
+    expect(events).toEqual(["x = 2"]);
+
+    firstRelease.resolve();
+    await saveStarted.promise;
+    updateCode(x, "x = 3");
+    await vi.advanceTimersByTimeAsync(400);
+    expect(events).toEqual(["x = 2", "save"]);
+
+    saveRelease.resolve();
+    await saving;
+    await flushDocumentChanges();
+    expect(events).toEqual(["x = 2", "save", "x = 3"]);
+  });
+
+  it("ignores a failure from a transaction reset while in flight", async () => {
+    const firstStarted = new Deferred<void>();
+    const firstRelease = new Deferred<void>();
+    const firstFinished = new Deferred<void>();
+    const delayedClient: Pick<
+      EditRequests & RunRequests,
+      "sendDocumentTransaction"
+    > = {
+      sendDocumentTransaction: async ({ changes }) => {
+        sent.push(changes);
+        if (sent.length === 1) {
+          firstStarted.resolve();
+          try {
+            await firstRelease.promise;
+          } finally {
+            firstFinished.resolve();
+          }
+        }
+        return null;
+      },
+    };
+    store.set(
+      requestClientAtom,
+      delayedClient as unknown as EditRequests & RunRequests,
+    );
+
+    setup("x = 1");
+    const [x] = state.cellIds.inOrderIds;
+    middlewareExports.cancelPendingChanges();
+    updateCode(x, "x = 2");
+    vi.advanceTimersByTime(400);
+    await firstStarted.promise;
+
+    middlewareExports.cancelPendingChanges();
+    firstRelease.reject(new Error("stale connection failure"));
+    await firstFinished.promise;
+    await Promise.resolve();
+
+    updateCode(x, "x = 3");
+    await flushDocumentChanges();
+    expect(sent.at(-1)).toEqual([
+      { type: "set-code", cellId: x, code: "x = 3" },
+    ]);
+  });
 });

package/src/core/cells/document-changes.ts

@@ -17,6 +17,7 @@
 import { debounce } from "lodash-es";
 import { assertNever } from "@/utils/assertNever";
 import type { DispatchedActionOf } from "@/utils/createReducer";
+import { Deferred } from "@/utils/Deferred";
 import { Logger } from "@/utils/Logger";
 import type { NotificationMessageData } from "../kernel/messages";
 import { kioskModeAtom } from "../mode";
@@ -632,16 +633,184 @@ export function coalesceChanges(changes: DocumentChange[]): DocumentChange[] {
 // ---------------------------------------------------------------------------
 
 let pendingChanges: DocumentChange[] = [];
+let stagedTransactions: DocumentChange[][] = [];
+let transactionQueue: Promise<void> = Promise.resolve();
+let transactionGeneration = 0;
+let activeDocumentSave: Deferred<void> | null = null;
+
+type TransactionSyncState =
+  | { status: "synchronized" }
+  | { status: "failed"; error: unknown };
+
+interface DocumentResync {
+  failure: Extract<TransactionSyncState, { status: "failed" }>;
+  includedChanges: DocumentChange[];
+  includedTransactions: DocumentChange[][];
+}
 
-const flushChanges = debounce(() => {
+let transactionSyncState: TransactionSyncState = { status: "synchronized" };
+
+function stagePendingTransaction(): void {
+  if (transactionSyncState.status === "failed") {
+    return;
+  }
   const changes = coalesceChanges(pendingChanges);
   pendingChanges = [];
-  if (changes.length === 0) {
+  if (changes.length > 0) {
+    stagedTransactions.push(changes);
+  }
+}
+
+async function sendStagedTransactions(generation: number): Promise<void> {
+  if (generation !== transactionGeneration) {
     return;
   }
-  void getRequestClient().sendDocumentTransaction({ changes });
+  if (transactionSyncState.status === "failed") {
+    throw transactionSyncState.error;
+  }
+
+  while (stagedTransactions.length > 0) {
+    const changes = stagedTransactions[0];
+
+    try {
+      await getRequestClient().sendDocumentTransaction({ changes });
+    } catch (error) {
+      if (generation !== transactionGeneration) {
+        return;
+      }
+      // A transport failure is ambiguous: the server may have applied the
+      // transaction before the response was lost. Keep the failure sticky so
+      // an export cannot claim synchronization without a fresh document.
+      transactionSyncState = { status: "failed", error };
+      stagedTransactions = stagedTransactions.slice(1);
+      throw error;
+    }
+    if (generation !== transactionGeneration) {
+      return;
+    }
+    stagedTransactions.shift();
+  }
+}
+
+function queuePendingTransactions(): Promise<void> {
+  const generation = transactionGeneration;
+  const saveBarrier = activeDocumentSave?.promise ?? Promise.resolve();
+  const queued = Promise.all([
+    transactionQueue.catch(() => undefined),
+    saveBarrier,
+  ]).then(() => sendStagedTransactions(generation));
+  transactionQueue = queued;
+  return queued;
+}
+
+const flushChanges = debounce(() => {
+  stagePendingTransaction();
+  void queuePendingTransactions().catch(() => undefined);
 }, 400);
 
+export async function flushDocumentChanges(): Promise<void> {
+  if (transactionSyncState.status === "failed") {
+    throw transactionSyncState.error;
+  }
+  flushChanges.cancel();
+  stagePendingTransaction();
+  let awaitedQueue = queuePendingTransactions();
+  while (true) {
+    await awaitedQueue;
+    flushChanges.cancel();
+    stagePendingTransaction();
+    if (
+      pendingChanges.length === 0 &&
+      stagedTransactions.length === 0 &&
+      transactionQueue === awaitedQueue
+    ) {
+      return;
+    }
+    awaitedQueue =
+      transactionQueue === awaitedQueue
+        ? queuePendingTransactions()
+        : transactionQueue;
+  }
+}
+
+export function beginDocumentResync(): DocumentResync | null {
+  if (transactionSyncState.status !== "failed") {
+    return null;
+  }
+  flushChanges.cancel();
+  const includedChanges = pendingChanges;
+  const includedTransactions = stagedTransactions;
+  pendingChanges = [];
+  stagedTransactions = [];
+  return {
+    failure: transactionSyncState,
+    includedChanges,
+    includedTransactions,
+  };
+}
+
+export function abortDocumentResync(resync: DocumentResync | null): void {
+  if (resync === null || transactionSyncState !== resync.failure) {
+    return;
+  }
+  pendingChanges = [...resync.includedChanges, ...pendingChanges];
+  stagedTransactions = [...resync.includedTransactions, ...stagedTransactions];
+}
+
+export async function completeDocumentResync(
+  resync: DocumentResync | null,
+): Promise<void> {
+  if (resync === null) {
+    return;
+  }
+  await transactionQueue.catch(() => undefined);
+  if (transactionSyncState !== resync.failure) {
+    return;
+  }
+  transactionSyncState = { status: "synchronized" };
+  transactionQueue = Promise.resolve();
+  if (pendingChanges.length > 0) {
+    flushChanges();
+  }
+}
+
+function releaseDocumentSave(barrier: Deferred<void>): void {
+  if (activeDocumentSave === barrier) {
+    activeDocumentSave = null;
+  }
+  barrier.resolve();
+}
+
+/**
+ * Runs a full save after earlier document transactions and holds later
+ * transactions until the save settles. A failed transaction makes the save
+ * the authoritative document resync boundary.
+ */
+export async function withDocumentSave<T>(save: () => Promise<T>): Promise<T> {
+  let resync: DocumentResync | null = null;
+  try {
+    await flushDocumentChanges();
+  } catch (error) {
+    resync = beginDocumentResync();
+    if (resync === null) {
+      throw error;
+    }
+  }
+
+  const barrier = new Deferred<void>();
+  activeDocumentSave = barrier;
+  try {
+    const result = await save();
+    releaseDocumentSave(barrier);
+    await completeDocumentResync(resync);
+    return result;
+  } catch (error) {
+    abortDocumentResync(resync);
+    releaseDocumentSave(barrier);
+    throw error;
+  }
+}
+
 function isScratchChange(change: DocumentChange): boolean {
   if ("cellId" in change && change.cellId === SCRATCH_CELL_ID) {
     return true;
@@ -658,7 +827,9 @@ function enqueue(change: DocumentChange) {
     return;
   }
   pendingChanges.push(change);
-  flushChanges();
+  if (transactionSyncState.status === "synchronized") {
+    flushChanges();
+  }
 }
 
 /**
@@ -746,7 +917,13 @@ export function applyTransactionChanges(
 export const exportedForTesting = {
   cancelPendingChanges: () => {
     flushChanges.cancel();
+    activeDocumentSave?.resolve();
+    activeDocumentSave = null;
+    transactionGeneration += 1;
     pendingChanges = [];
+    stagedTransactions = [];
+    transactionSyncState = { status: "synchronized" };
+    transactionQueue = Promise.resolve();
   },
   drainChanges: (): DocumentChange[] => {
     flushChanges.cancel();

package/src/core/export/hooks.ts

@@ -11,6 +11,7 @@ import { Objects } from "@/utils/objects";
 import { ProgressState } from "@/utils/progress";
 import { cellsRuntimeAtom } from "../cells/cells";
 import type { CellId } from "../cells/ids";
+import { getExportLayout } from "./layout";
 import { connectionAtom } from "../network/connection";
 import { useRequestClient } from "../network/requests";
 import type { UpdateCellOutputsRequest } from "../network/types";
@@ -56,6 +57,7 @@ export function useAutoExport() {
         download: false,
         includeCode: true,
         files: VirtualFileTracker.INSTANCE.filenames(),
+        layout: await getExportLayout(),
       });
     },
     // Run every 5 seconds, or when the document becomes visible

package/src/core/export/layout.ts

@@ -0,0 +1,9 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { flushDocumentChanges } from "@/core/cells/document-changes";
+import { getSerializedLayout } from "@/core/layout/layout";
+
+export async function getExportLayout() {
+  await flushDocumentChanges();
+  return getSerializedLayout();
+}