@marimo-team/frontend 0.24.1-dev58 → 0.24.1-dev59

package/src/core/ai/__tests__/stream-completion-text.test.ts

@@ -0,0 +1,74 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { streamCompletionText } from "../stream-completion-text";
+
+function streamResponse(...chunks: object[]): Response {
+  const body = [
+    ...chunks.map((chunk) => `data: ${JSON.stringify(chunk)}\n\n`),
+    "data: [DONE]\n\n",
+  ].join("");
+  return new Response(body, {
+    headers: { "Content-Type": "text/event-stream" },
+  });
+}
+
+describe("streamCompletionText", () => {
+  it("returns validated cell code", async () => {
+    const response = streamResponse(
+      { type: "start" },
+      {
+        type: "data-cell-completion",
+        data: { code: "print('```')" },
+      },
+      { type: "finish", finishReason: "stop" },
+    );
+
+    await expect(streamCompletionText(response)).resolves.toBe("print('```')");
+  });
+
+  it("rejects a missing structured completion", async () => {
+    const response = streamResponse(
+      { type: "start" },
+      { type: "finish", finishReason: "stop" },
+    );
+
+    await expect(streamCompletionText(response)).rejects.toThrow(
+      "AI completion returned no cell code",
+    );
+  });
+
+  it("rejects malformed completion data", async () => {
+    const response = streamResponse({
+      type: "data-cell-completion",
+      data: { code: 42 },
+    });
+
+    await expect(streamCompletionText(response)).rejects.toThrow();
+  });
+
+  it("rejects a partial completion without successful final validation", async () => {
+    const response = streamResponse({
+      type: "data-cell-completion",
+      data: { code: "partial" },
+    });
+
+    await expect(streamCompletionText(response)).rejects.toThrow(
+      "AI completion ended before final validation",
+    );
+  });
+
+  it("rejects a completion with an error finish reason", async () => {
+    const response = streamResponse(
+      {
+        type: "data-cell-completion",
+        data: { code: "partial" },
+      },
+      { type: "finish", finishReason: "error" },
+    );
+
+    await expect(streamCompletionText(response)).rejects.toThrow(
+      "AI completion ended before final validation",
+    );
+  });
+});

package/src/core/ai/completion-output.ts

@@ -0,0 +1,56 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { DataUIPart, UIMessage, UIMessageChunk } from "ai";
+import { z } from "zod";
+
+// Mirrored from marimo/_server/ai/completion_output.py; SSE data parts are not
+// represented in OpenAPI.
+const CELL_COMPLETION_DATA_PART = "cell-completion" as const;
+const NOTEBOOK_CELLS_COMPLETION_DATA_PART =
+  "notebook-cells-completion" as const;
+
+export const CELL_COMPLETION_DATA_TYPE =
+  `data-${CELL_COMPLETION_DATA_PART}` as const;
+export const NOTEBOOK_CELLS_COMPLETION_DATA_TYPE =
+  `data-${NOTEBOOK_CELLS_COMPLETION_DATA_PART}` as const;
+
+export const cellCompletionSchema = z.object({
+  code: z.string(),
+});
+
+export const generatedCellSchema = z.object({
+  language: z.enum(["python", "sql", "markdown"]),
+  code: z.string().min(1),
+});
+
+export const notebookCellsCompletionSchema = z.object({
+  cells: z.array(generatedCellSchema).min(1),
+});
+
+export type CellCompletion = z.infer<typeof cellCompletionSchema>;
+export type GeneratedCell = z.infer<typeof generatedCellSchema>;
+export type NotebookCellsCompletion = z.infer<
+  typeof notebookCellsCompletionSchema
+>;
+
+// Keep parsing data parts explicitly at the SSE boundary. AI SDK 7.0.37 types
+// schema keys without `data-`, but its runtime lookup uses the full wire type.
+const completionDataSchemas = {
+  [CELL_COMPLETION_DATA_PART]: cellCompletionSchema,
+  [NOTEBOOK_CELLS_COMPLETION_DATA_PART]: notebookCellsCompletionSchema,
+};
+
+export type CompletionDataParts = {
+  [Name in keyof typeof completionDataSchemas]: z.infer<
+    (typeof completionDataSchemas)[Name]
+  >;
+};
+
+export type CompletionUIMessage = UIMessage<unknown, CompletionDataParts>;
+export type CompletionDataPart = DataUIPart<CompletionDataParts>;
+
+type DataChunk = Extract<UIMessageChunk, { type: `data-${string}` }>;
+
+export function isDataChunk(chunk: UIMessageChunk): chunk is DataChunk {
+  return chunk.type.startsWith("data-");
+}

package/src/core/ai/staged-cells.ts

@@ -1,34 +1,32 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { UIMessageChunk } from "ai";
-import { useAtomValue } from "jotai";
+import { atom, useAtomValue } from "jotai";
 import { selectAtom } from "jotai/utils";
 import { useMemo, useRef } from "react";
-import {
-  type AiCompletion,
-  codeToCells,
-} from "@/components/editor/ai/completion-utils";
 import { useDeleteCellCallback } from "@/components/editor/cell/useDeleteCell";
 import { CellId } from "@/core/cells/ids";
-import { logNever } from "@/utils/assertNever";
 import { createReducerAndAtoms } from "@/utils/createReducer";
 import { Logger } from "@/utils/Logger";
+import {
+  type CompletionDataPart,
+  NOTEBOOK_CELLS_COMPLETION_DATA_TYPE,
+  type GeneratedCell,
+  notebookCellsCompletionSchema,
+} from "./completion-output";
 import { maybeAddMarimoImport } from "../cells/add-missing-import";
 import {
   type CreateNewCellAction,
   getCellEditorView,
   useCellActions,
 } from "../cells/cells";
-import type { LanguageAdapterType } from "../codemirror/language/types";
 import { updateEditorCodeFromPython } from "../codemirror/language/utils";
 import type { JotaiStore } from "../state/jotai";
 import type { EditType } from "./tools/edit-notebook-tool";
 
 /**
- * Cells that are staged for AI completion
- * They function similarly to cells in the notebook, but they can be accepted or rejected by the user.
- * We track edited, new and deleted cells.
- * And we only track one set of staged cells at a time.
+ * Pending AI edits shared by completion, chat, and agent workflows.
+ * Generation ownership keeps one workflow from accepting or discarding another's
+ * staged cells.
  */
 
 export type Edit =
@@ -38,6 +36,25 @@ export type Edit =
 
 export type StagedAICells = Map<CellId, Edit>;
 
+type StagedGeneration =
+  | {
+      id: symbol;
+      status: "in_progress";
+      cellIds: readonly CellId[];
+    }
+  | { id: symbol; status: "complete" };
+
+interface OwnedStagedGeneration {
+  id: symbol;
+  status: "in_progress" | "complete";
+  reconciler: StagedCellReconciler;
+}
+
+const stagedGenerationAtom = atom<StagedGeneration | null>(null);
+export const stagedGenerationInProgressAtom = atom(
+  (get) => get(stagedGenerationAtom)?.status === "in_progress",
+);
+
 const initialState = (): StagedAICells => {
   return new Map();
 };
@@ -60,16 +77,9 @@ const {
     newState.delete(cellId);
     return newState;
   },
-  clearStagedCells: () => {
-    return initialState();
-  },
 });
 
-export {
-  useStagedAICellsActions,
-  createActions as createStagedAICellsActions,
-  reducer as stagedAICellsReducer,
-};
+export { useStagedAICellsActions };
 
 export function useStagedAICell(cellId: CellId): Edit | undefined {
   const stagedCellAtom = useMemo(
@@ -83,20 +93,14 @@ export function useStagedAICell(cellId: CellId): Edit | undefined {
 interface UpdateStagedCellAction {
   cellId: CellId;
   code: string;
-  language?: LanguageAdapterType;
 }
 
-/**
- * Helper functions to create and delete staged cells.
- */
+/** Manage staged cells without assuming which AI workflow owns them. */
 export function useStagedCells(store: JotaiStore) {
-  const { addStagedCell, removeStagedCell, clearStagedCells } =
-    useStagedAICellsActions();
+  const { addStagedCell, removeStagedCell } = useStagedAICellsActions();
   const { createNewCell, updateCellCode } = useCellActions();
   const deleteCellCallback = useDeleteCellCallback();
 
-  const cellCreationStream = useRef<CellCreationStream | null>(null);
-
   const createStagedCell = (code: string): CellId => {
     const newCellId = CellId.create();
     addStagedCell({ cellId: newCellId, edit: { type: "add_cell" } });
@@ -118,8 +122,6 @@ export function useStagedCells(store: JotaiStore) {
       return;
     }
 
-    // Update the editor code if the cell is mounted
-    // Else, update the cell code in the notebook
     const editorView = getCellEditorView(cellId);
     if (editorView) {
       updateEditorCodeFromPython(editorView, code);
@@ -128,190 +130,269 @@ export function useStagedCells(store: JotaiStore) {
     }
   };
 
-  // Delete a staged cell and the corresponding cell in the notebook.
   const deleteStagedCell = (cellId: CellId) => {
     removeStagedCell(cellId);
     deleteCellCallback({ cellId });
   };
 
-  // Delete all staged cells and the corresponding cells in the notebook.
-  const deleteAllStagedCells = () => {
+  return {
+    createStagedCell,
+    updateStagedCell,
+    addStagedCell,
+    removeStagedCell,
+    deleteStagedCell,
+  };
+}
+
+/** Owns the lifecycle of cells created by one structured completion flow. */
+export function useStagedCellGeneration(store: JotaiStore) {
+  const stagedCellActions = useStagedCells(store);
+  const { createNewCell } = useCellActions();
+  const ownedGeneration = useRef<OwnedStagedGeneration | null>(null);
+  // Ownership lives in a ref, so subscribe to its backing atoms to rerender
+  // consumers when completion or per-cell acceptance changes.
+  useAtomValue(stagedGenerationAtom, { store });
+  useAtomValue(stagedAICellsAtom, { store });
+
+  const discardStagedCellIds = (cellIds: readonly CellId[]) => {
     const stagedAICells = store.get(stagedAICellsAtom);
-    for (const cellId of stagedAICells.keys()) {
-      deleteCellCallback({ cellId });
+    for (const cellId of cellIds.toReversed()) {
+      if (stagedAICells.has(cellId)) {
+        stagedCellActions.deleteStagedCell(cellId);
+      }
     }
-    clearStagedCells();
   };
 
-  const onStream = (chunk: UIMessageChunk) => {
-    switch (chunk.type) {
-      case "text-start":
-        // Create stream
-        cellCreationStream.current = new CellCreationStream(
-          createStagedCell,
-          updateStagedCell,
-          addStagedCell,
-          createNewCell,
-        );
-        break;
-      case "text-delta":
-        if (!cellCreationStream.current) {
-          Logger.error("Cell creation stream not found");
-          return;
-        }
-        cellCreationStream.current.stream(chunk);
-        break;
-      case "text-end":
-      case "finish":
-        if (!cellCreationStream.current) {
-          Logger.error("Cell creation stream not found");
-          return;
-        }
-        cellCreationStream.current.stop();
-        break;
-      case "abort":
-      case "error":
-      case "tool-input-error":
-      case "tool-output-error":
-        Logger.error("Error", chunk.type, { chunk });
-        break;
-      case "tool-approval-request":
-        Logger.log("Tool approval request", { chunk });
-        break;
-      case "tool-output-denied":
-        Logger.error("Tool output denied", { chunk });
-        break;
-      // These logs are not useful for debugging
-      case "start":
-      case "start-step":
-      case "finish-step":
-      case "data-reasoning-signature":
-        break;
-      case "message-metadata":
-      case "tool-input-available":
-      case "tool-output-available":
-      case "tool-approval-response":
-      case "reasoning-start":
-      case "reasoning-delta":
-      case "reasoning-end":
-      case "file":
-      case "reasoning-file":
-      case "source-document":
-      case "source-url":
-      case "tool-input-start":
-      case "tool-input-delta":
-      case "custom":
-        Logger.debug(chunk.type, { chunk });
-        break;
-      default:
-        if (isDataChunk(chunk)) {
-          Logger.debug("Data chunk", { chunk });
-          break;
-        }
-        logNever(chunk);
+  const beginStagedCellGeneration = () => {
+    const activeGeneration = store.get(stagedGenerationAtom);
+    if (activeGeneration?.status === "in_progress") {
+      discardStagedCellIds(activeGeneration.cellIds);
+    }
+    ownedGeneration.current?.reconciler.discard(store.get(stagedAICellsAtom));
+    const generationId = Symbol("staged-cell-generation");
+    ownedGeneration.current = {
+      id: generationId,
+      status: "in_progress",
+      reconciler: new StagedCellReconciler({
+        ...stagedCellActions,
+        createNewCell,
+      }),
+    };
+    store.set(stagedGenerationAtom, {
+      id: generationId,
+      status: "in_progress",
+      cellIds: [],
+    });
+  };
+
+  const finishStagedCellGeneration = (successful: boolean) => {
+    const generation = ownedGeneration.current;
+    if (generation?.status !== "in_progress") {
+      return;
+    }
+
+    const stagedGeneration = store.get(stagedGenerationAtom);
+    const ownsActiveGeneration =
+      stagedGeneration?.id === generation.id &&
+      stagedGeneration.status === "in_progress";
+    if (!ownsActiveGeneration || !successful) {
+      generation.reconciler.discard(store.get(stagedAICellsAtom));
+      if (ownsActiveGeneration) {
+        store.set(stagedGenerationAtom, null);
+      }
+      ownedGeneration.current = null;
+      return;
+    }
+
+    generation.status = "complete";
+    store.set(stagedGenerationAtom, {
+      id: generation.id,
+      status: "complete",
+    });
+  };
+
+  const hasOwnedStagedCells = () => {
+    const generation = ownedGeneration.current;
+    const stagedGeneration = store.get(stagedGenerationAtom);
+    return (
+      generation !== null &&
+      generation.status === "complete" &&
+      stagedGeneration?.status === "complete" &&
+      stagedGeneration.id === generation.id &&
+      generation.reconciler.hasStagedCells(store.get(stagedAICellsAtom))
+    );
+  };
+
+  const acceptOwnedStagedCells = () => {
+    if (!hasOwnedStagedCells()) {
+      return false;
+    }
+
+    ownedGeneration.current?.reconciler.accept();
+    store.set(stagedGenerationAtom, null);
+    ownedGeneration.current = null;
+    return true;
+  };
+
+  const discardOwnedStagedCells = () => {
+    const generation = ownedGeneration.current;
+    const stagedGeneration = store.get(stagedGenerationAtom);
+    if (generation === null || stagedGeneration?.id !== generation.id) {
+      return false;
     }
+
+    generation.reconciler.discard(store.get(stagedAICellsAtom));
+    store.set(stagedGenerationAtom, null);
+    ownedGeneration.current = null;
+    return true;
+  };
+
+  const onData = (part: CompletionDataPart) => {
+    if (part.type !== NOTEBOOK_CELLS_COMPLETION_DATA_TYPE) {
+      return;
+    }
+    const generation = ownedGeneration.current;
+    const stagedGeneration = store.get(stagedGenerationAtom);
+    if (
+      generation?.status !== "in_progress" ||
+      stagedGeneration?.id !== generation.id ||
+      stagedGeneration.status !== "in_progress"
+    ) {
+      return;
+    }
+    const completion = notebookCellsCompletionSchema.parse(part.data);
+    generation.reconciler.reconcile(completion.cells);
+    store.set(stagedGenerationAtom, {
+      ...stagedGeneration,
+      cellIds: generation.reconciler.stagedCellIds(),
+    });
   };
 
   return {
-    createStagedCell,
-    updateStagedCell,
-    addStagedCell,
-    removeStagedCell,
-    clearStagedCells,
-    deleteStagedCell,
-    deleteAllStagedCells,
-    onStream,
+    beginStagedCellGeneration,
+    finishStagedCellGeneration,
+    hasOwnedStagedCells,
+    acceptOwnedStagedCells,
+    discardOwnedStagedCells,
+    onData,
   };
 }
 
 export { stagedAICellsAtom };
 export const visibleForTesting = {
+  stagedGenerationAtom,
   createActions,
   reducer,
   initialState,
-  useStagedAICellsActions,
 };
 
-type TextDeltaChunk = Extract<UIMessageChunk, { type: "text-delta" }>;
-
 interface CreatedCell {
   cellId: CellId;
-  cell: AiCompletion;
+  code: string;
 }
 
-class CellCreationStream {
+interface StagedCellReconcilerActions {
+  createStagedCell: (code: string) => CellId;
+  updateStagedCell: (opts: UpdateStagedCellAction) => void;
+  deleteStagedCell: (cellId: CellId) => void;
+  removeStagedCell: (cellId: CellId) => void;
+  addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
+  createNewCell: (opts: CreateNewCellAction) => void;
+}
+
+class StagedCellReconciler {
   private createdCells: CreatedCell[] = [];
-  private buffer = "";
-
-  private onCreateCell: (code: string) => CellId;
-  private onUpdateCell: (opts: UpdateStagedCellAction) => void;
-  private addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
-  private createNewCell: (opts: CreateNewCellAction) => void;
-  private hasMarimoImport = false;
-
-  constructor(
-    onCreateCell: (code: string) => CellId,
-    onUpdateCell: (opts: UpdateStagedCellAction) => void,
-    addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void,
-    createNewCell: (opts: CreateNewCellAction) => void,
-  ) {
-    this.onCreateCell = onCreateCell;
-    this.onUpdateCell = onUpdateCell;
-    this.addStagedCell = addStagedCell;
-    this.createNewCell = createNewCell;
-  }
+  private readonly actions: StagedCellReconcilerActions;
+  private marimoImportCellId: CellId | null = null;
 
-  stream(chunk: TextDeltaChunk) {
-    const delta = chunk.delta;
-    this.buffer += delta;
-    const completionCells = codeToCells(this.buffer);
+  constructor(actions: StagedCellReconcilerActions) {
+    this.actions = actions;
+  }
 
-    // As incoming chunks are appended to the buffer,
-    // we parse the buffer into cells and determine which parts correspond to which cell.
-    // For each parsed cell, we either update an existing staged cell or create a new one.
+  reconcile(completionCells: GeneratedCell[]) {
     for (const [idx, cell] of completionCells.entries()) {
       if (idx < this.createdCells.length) {
-        this.addMarimoImport(cell.language);
         const existingCell = this.createdCells[idx];
-        this.createdCells[idx] = { ...existingCell, cell };
-        this.onUpdateCell({
+        const codeChanged = existingCell.code !== cell.code;
+        this.createdCells[idx] = { ...existingCell, code: cell.code };
+        if (!codeChanged) {
+          continue;
+        }
+        this.actions.updateStagedCell({
           cellId: existingCell.cellId,
           code: cell.code,
-          language: cell.language,
         });
       } else {
-        const newCellId = this.onCreateCell(cell.code);
-        this.createdCells.push({ cellId: newCellId, cell });
+        const newCellId = this.actions.createStagedCell(cell.code);
+        this.createdCells.push({ cellId: newCellId, code: cell.code });
       }
     }
+
+    const removedCells = this.createdCells.splice(completionCells.length);
+    for (const { cellId } of removedCells.toReversed()) {
+      this.actions.deleteStagedCell(cellId);
+    }
+
+    this.syncMarimoImport(completionCells);
   }
 
-  /** Add a marimo import if the cell is SQL or Markdown and we haven't added it yet. */
-  private addMarimoImport(language: LanguageAdapterType) {
-    if (this.hasMarimoImport || language === "python") {
+  discard(stagedCells: StagedAICells) {
+    for (const { cellId } of this.createdCells.toReversed()) {
+      if (stagedCells.has(cellId)) {
+        this.actions.deleteStagedCell(cellId);
+      }
+    }
+    this.createdCells = [];
+
+    if (this.marimoImportCellId && stagedCells.has(this.marimoImportCellId)) {
+      this.actions.deleteStagedCell(this.marimoImportCellId);
+    }
+    this.marimoImportCellId = null;
+  }
+
+  accept() {
+    for (const cellId of this.stagedCellIds()) {
+      this.actions.removeStagedCell(cellId);
+    }
+  }
+
+  hasStagedCells(stagedCells: StagedAICells) {
+    return this.stagedCellIds().some((cellId) => stagedCells.has(cellId));
+  }
+
+  stagedCellIds() {
+    const cellIds = this.createdCells.map(({ cellId }) => cellId);
+    if (this.marimoImportCellId) {
+      cellIds.push(this.marimoImportCellId);
+    }
+    return cellIds;
+  }
+
+  /** Keep the generated marimo import consistent with the latest snapshot. */
+  private syncMarimoImport(completionCells: GeneratedCell[]) {
+    const requiresMarimo = completionCells.some(
+      (cell) => cell.language !== "python",
+    );
+    if (!requiresMarimo) {
+      if (this.marimoImportCellId) {
+        this.actions.deleteStagedCell(this.marimoImportCellId);
+        this.marimoImportCellId = null;
+      }
+      return;
+    }
+    if (this.marimoImportCellId) {
       return;
     }
 
     const cellId = maybeAddMarimoImport({
       autoInstantiate: false,
-      createNewCell: this.createNewCell,
+      createNewCell: this.actions.createNewCell,
       fromCellId: this.createdCells[0]?.cellId,
       before: true,
     });
     if (cellId) {
-      this.addStagedCell({ cellId, edit: { type: "add_cell" } });
+      this.actions.addStagedCell({ cellId, edit: { type: "add_cell" } });
+      this.marimoImportCellId = cellId;
     }
-    this.hasMarimoImport = true;
   }
-
-  stop() {
-    // Clear all state
-    this.buffer = "";
-  }
-}
-
-type DataChunk = Extract<UIMessageChunk, { type: `data-${string}` }>;
-
-function isDataChunk(chunk: UIMessageChunk): chunk is DataChunk {
-  return chunk.type.startsWith("data-");
 }