fix(data-editor): make edit replay robust across structural changes

frontend/plugins.openapi.yaml

@@ -3071,6 +3071,11 @@ components:
                           - unknown
                       - type: string
             - type: "null"
+        columnNames:
+          default: []
+          type: array
+          items:
+            type: string
         editableColumns:
           anyOf:
             - type: array

frontend/src/plugins/impl/DataEditorPlugin.tsx

@@ -1,27 +1,28 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import glideCss from "@glideapps/glide-data-grid/dist/index.css?inline";
-import React, { useState } from "react";
+import React, { useEffect, useState } from "react";
 import { z } from "zod";
 import { inferFieldTypes } from "@/components/data-table/columns";
 import { LoadingTable } from "@/components/data-table/loading-table";
-import { type FieldTypes, toFieldTypes } from "@/components/data-table/types";
+import {
+  type FieldTypesWithExternalType,
+  toFieldTypes,
+} from "@/components/data-table/types";
 import { Alert, AlertTitle } from "@/components/ui/alert";
 import { DelayMount } from "@/components/utils/delay-mount";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import { createPlugin } from "../core/builder";
 import type { Setter } from "../types";
 import { columnToFieldTypesSchema } from "./data-frames/schema";
-import {
-  BulkEdit,
-  type DataEditorProps,
-  type Edits,
-} from "./data-editor/types";
+import { orderColumnFields } from "./data-editor/data-utils";
+import { applyEditorEdits } from "./data-editor/editor-state";
+import type { EditorRow, EditorState, Edits } from "./data-editor/types";
 import { vegaLoadData } from "./vega/loader";
 import { getVegaFieldTypes } from "./vega/utils";
 
 type CsvURL = string;
-type TableData<T> = T[] | CsvURL;
+type TableData = EditorRow[] | CsvURL;
 
 // Lazy load the data editor since it brings in glide-data-grid
 const LazyDataEditor = React.lazy(
@@ -45,6 +46,7 @@ export const DataEditorPlugin = createPlugin<Edits>("marimo-data-editor", {
       label: z.string().nullable(),
       data: z.union([z.string(), z.array(z.object({}).passthrough())]),
       fieldTypes: columnToFieldTypesSchema.nullish(),
+      columnNames: z.array(z.string()).default([]),
       editableColumns: z.union([z.array(z.string()), z.literal("all")]),
       columnSizingMode: z.enum(["auto", "fit"]).default("auto"), // TODO: Remove this
     }),
@@ -55,39 +57,36 @@ export const DataEditorPlugin = createPlugin<Edits>("marimo-data-editor", {
       <LoadingDataEditor
         data={props.data.data}
         fieldTypes={props.data.fieldTypes}
+        columnNames={props.data.columnNames}
         edits={props.value}
         onEdits={props.setValue}
-        host={props.host}
         editableColumns={props.data.editableColumns}
       />
     );
   });
 
-interface Props extends Omit<
-  DataEditorProps<object>,
-  "data" | "onAddEdits" | "onAddRows"
-> {
-  data: TableData<object>;
+interface Props {
+  data: TableData;
+  fieldTypes: FieldTypesWithExternalType | null | undefined;
   edits: Edits;
   onEdits: Setter<Edits>;
-  host: HTMLElement;
   editableColumns: string[] | "all";
+  columnNames: string[];
 }
 
 const LoadingDataEditor = (props: Props) => {
-  const [data, setData] = useState<unknown[]>([]);
-  const [columnFields, setColumnFields] = useState<FieldTypes>(new Map());
+  const [editorState, setEditorState] = useState<EditorState | null>(null);
 
   // Load the data
-  const { error } = useAsyncData(async () => {
+  const { data: loadedState, error } = useAsyncData(async () => {
     const withoutExternalTypes = toFieldTypes(props.fieldTypes ?? []);
 
     // If we already have the data, return it
     // Otherwise, load the data from the URL. Vega's CSV parser takes a
     // plain `Record`; column order doesn't matter for parsing.
     const localData = Array.isArray(props.data)
       ? props.data
-      : await vegaLoadData(
+      : await vegaLoadData<EditorRow>(
           props.data,
           {
             type: "csv",
@@ -96,11 +95,20 @@ const LoadingDataEditor = (props: Props) => {
           { handleBigIntAndNumberLike: true },
         );
 
-    setData(localData);
-    setColumnFields(
-      toFieldTypes(props.fieldTypes ?? inferFieldTypes(localData)),
-    );
-  }, [props.fieldTypes, props.data]);
+    return {
+      data: localData,
+      columnFields: orderColumnFields(
+        toFieldTypes(props.fieldTypes ?? inferFieldTypes(localData)),
+        props.columnNames,
+      ),
+    } satisfies EditorState;
+  }, [props.fieldTypes, props.columnNames, props.data]);
+
+  useEffect(() => {
+    if (loadedState !== undefined) {
+      setEditorState(applyEditorEdits(loadedState, props.edits.edits));
+    }
+  }, [loadedState, props.edits.edits]);
 
   if (error) {
     return (
@@ -113,7 +121,7 @@ const LoadingDataEditor = (props: Props) => {
     );
   }
 
-  if (!data) {
+  if (editorState === null) {
     return (
       <DelayMount milliseconds={200}>
         <LoadingTable pageSize={10} />
@@ -123,54 +131,14 @@ const LoadingDataEditor = (props: Props) => {
 
   return (
     <LazyDataEditor
-      data={data}
-      setData={setData}
-      columnFields={columnFields}
-      setColumnFields={setColumnFields}
+      data={editorState.data}
+      columnFields={editorState.columnFields}
       editableColumns={props.editableColumns}
-      edits={props.edits.edits}
       onAddEdits={(edits) => {
-        props.onEdits((v) => ({ ...v, edits: [...v.edits, ...edits] }));
-      }}
-      onAddRows={(rows) => {
-        const newEdits = rows.flatMap((row, rowIndex) =>
-          Object.entries(row).map(([columnId, value]) => ({
-            rowIdx: data.length + rowIndex,
-            columnId,
-            value,
-          })),
+        setEditorState((state) =>
+          state === null ? null : applyEditorEdits(state, edits),
         );
-        props.onEdits((v) => ({ ...v, edits: [...v.edits, ...newEdits] }));
-      }}
-      onDeleteRows={(rowIndexes) => {
-        props.onEdits((v) => {
-          const newEdits = rowIndexes.map((rowIdx, index) => ({
-            rowIdx: rowIdx - index,
-            type: BulkEdit.Remove,
-          }));
-          return {
-            ...v,
-            edits: [...v.edits, ...newEdits],
-          };
-        });
-      }}
-      onRenameColumn={(columnIdx: number, newName: string) => {
-        props.onEdits((v) => ({
-          ...v,
-          edits: [...v.edits, { columnIdx, newName, type: BulkEdit.Rename }],
-        }));
-      }}
-      onDeleteColumn={(columnIdx: number) => {
-        props.onEdits((v) => ({
-          ...v,
-          edits: [...v.edits, { columnIdx, type: BulkEdit.Remove }],
-        }));
-      }}
-      onAddColumn={(columnIdx: number, newName: string) => {
-        props.onEdits((v) => ({
-          ...v,
-          edits: [...v.edits, { columnIdx, newName, type: BulkEdit.Insert }],
-        }));
+        props.onEdits((v) => ({ ...v, edits: [...v.edits, ...edits] }));
       }}
     />
   );

frontend/src/plugins/impl/__tests__/DataEditorPlugin.test.ts

@@ -1,9 +1,55 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { describe, expect, it } from "vitest";
+import React, { Suspense } from "react";
+import { act, render, screen, waitFor } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { z } from "zod";
+import type { IPluginProps } from "../../types";
+import type { EditorRow, Edits } from "../data-editor/types";
 import { DataEditorPlugin } from "../DataEditorPlugin";
 
+type PluginData = z.infer<typeof DataEditorPlugin.validator>;
+
+const mocks = vi.hoisted(() => ({
+  loadData: vi.fn(),
+}));
+
+vi.mock("../vega/loader", () => ({
+  vegaLoadData: mocks.loadData,
+}));
+
+vi.mock("../data-editor/glide-data-editor", async () => {
+  const React = await import("react");
+  return {
+    default: (props: { data: EditorRow[] }) =>
+      React.createElement(
+        "pre",
+        { "data-testid": "editor-data" },
+        JSON.stringify(props.data),
+      ),
+  };
+});
+
+function renderPlugin(data: PluginData, value: Edits) {
+  const props: IPluginProps<Edits, PluginData> = {
+    host: document.createElement("div"),
+    data,
+    value,
+    setValue: vi.fn(),
+    functions: {},
+  };
+  return React.createElement(
+    Suspense,
+    { fallback: null },
+    DataEditorPlugin.render(props),
+  );
+}
+
 describe("DataEditorPlugin", () => {
+  beforeEach(() => {
+    mocks.loadData.mockReset();
+  });
+
   it("normalizes an unrecognized field type", () => {
     const result = DataEditorPlugin.validator.parse({
       initialValue: { edits: [] },
@@ -27,4 +73,109 @@ describe("DataEditorPlugin", () => {
 
     expect(result.fieldTypes).toEqual([["geom", ["geometry", "geometry"]]]);
   });
+
+  it("applies the latest edits when an async data load completes", async () => {
+    let resolveLoad: ((data: EditorRow[]) => void) | undefined;
+    mocks.loadData.mockImplementationOnce(
+      () =>
+        new Promise<EditorRow[]>((resolve) => {
+          resolveLoad = resolve;
+        }),
+    );
+
+    const data = DataEditorPlugin.validator.parse({
+      initialValue: { edits: [] },
+      label: null,
+      data: "data.csv",
+      fieldTypes: null,
+      columnNames: ["name"],
+      editableColumns: "all",
+    });
+    const result = render(renderPlugin(data, { edits: [] }));
+    result.rerender(
+      renderPlugin(data, {
+        edits: [{ rowIdx: 0, columnId: "name", value: "latest" }],
+      }),
+    );
+    expect(resolveLoad).toBeDefined();
+    resolveLoad?.([{ name: "original" }]);
+
+    await waitFor(() => {
+      expect(screen.getByTestId("editor-data")).toHaveTextContent(
+        JSON.stringify([{ name: "latest" }]),
+      );
+    });
+  });
+
+  it("applies edit updates after data has loaded", async () => {
+    const data = DataEditorPlugin.validator.parse({
+      initialValue: { edits: [] },
+      label: null,
+      data: [{ name: "original" }],
+      fieldTypes: null,
+      columnNames: ["name"],
+      editableColumns: "all",
+    });
+    const result = render(renderPlugin(data, { edits: [] }));
+    await waitFor(() => {
+      expect(screen.getByTestId("editor-data")).toHaveTextContent(
+        JSON.stringify([{ name: "original" }]),
+      );
+    });
+
+    result.rerender(
+      renderPlugin(data, {
+        edits: [{ rowIdx: 0, columnId: "name", value: "updated" }],
+      }),
+    );
+
+    await waitFor(() => {
+      expect(screen.getByTestId("editor-data")).toHaveTextContent(
+        JSON.stringify([{ name: "updated" }]),
+      );
+    });
+  });
+
+  it("ignores a superseded async data load", async () => {
+    let resolveFirst: ((data: EditorRow[]) => void) | undefined;
+    let resolveSecond: ((data: EditorRow[]) => void) | undefined;
+    mocks.loadData.mockImplementation((source: string) => {
+      return new Promise<EditorRow[]>((resolve) => {
+        if (source === "first.csv") {
+          resolveFirst = resolve;
+        } else {
+          resolveSecond = resolve;
+        }
+      });
+    });
+
+    const makeData = (source: string) =>
+      DataEditorPlugin.validator.parse({
+        initialValue: { edits: [] },
+        label: null,
+        data: source,
+        fieldTypes: null,
+        columnNames: ["source"],
+        editableColumns: "all",
+      });
+    const result = render(renderPlugin(makeData("first.csv"), { edits: [] }));
+    result.rerender(renderPlugin(makeData("second.csv"), { edits: [] }));
+
+    expect(resolveFirst).toBeDefined();
+    expect(resolveSecond).toBeDefined();
+    resolveSecond?.([{ source: "second" }]);
+    await waitFor(() => {
+      expect(screen.getByTestId("editor-data")).toHaveTextContent(
+        JSON.stringify([{ source: "second" }]),
+      );
+    });
+
+    await act(async () => {
+      resolveFirst?.([{ source: "first" }]);
+      await Promise.resolve();
+    });
+    expect(screen.getByTestId("editor-data")).toHaveTextContent(
+      JSON.stringify([{ source: "second" }]),
+    );
+  });
 });