@marimo-team/islands 0.25.3-dev250 → 0.25.3-dev252

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

@@ -1,30 +1,19 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { atom, useAtomValue } from "jotai";
+import { useAtomValue } from "jotai";
 import { selectAtom } from "jotai/utils";
-import { useMemo, useRef } from "react";
+import { useMemo } from "react";
 import { useDeleteCellCallback } from "@/components/editor/cell/useDeleteCell";
 import { CellId } from "@/core/cells/ids";
 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 { getCellEditorView } from "../cells/cells";
-import { type CreateCell, useCellActions } from "../cells/actions";
+import { useCellActions } from "../cells/actions";
 import { prepareCellCodeUpdate } from "../codemirror/language/cell-code";
 import type { JotaiStore } from "../state/jotai";
 import type { EditType } from "./tools/edit-notebook-tool";
 
-/**
- * Pending AI edits shared by completion, chat, and agent workflows.
- * Generation ownership keeps one workflow from accepting or discarding another's
- * staged cells.
- */
+/** Pending AI edits shared by chat and agent workflows. */
 
 export type Edit =
   | { type: Extract<EditType, "update_cell">; previousCode: string }
@@ -33,25 +22,6 @@ 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();
 };
@@ -143,255 +113,9 @@ export function useStagedCells(store: JotaiStore) {
   };
 }
 
-/** Owns the lifecycle of cells created by one structured completion flow. */
-export function useStagedCellGeneration(store: JotaiStore) {
-  const stagedCellActions = useStagedCells(store);
-  const { createCell } = 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 cellIds.toReversed()) {
-      if (stagedAICells.has(cellId)) {
-        stagedCellActions.deleteStagedCell(cellId);
-      }
-    }
-  };
-
-  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,
-        createCell,
-      }),
-    };
-    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 {
-    beginStagedCellGeneration,
-    finishStagedCellGeneration,
-    hasOwnedStagedCells,
-    acceptOwnedStagedCells,
-    discardOwnedStagedCells,
-    onData,
-  };
-}
-
 export { stagedAICellsAtom };
 export const visibleForTesting = {
-  stagedGenerationAtom,
   createActions,
   reducer,
   initialState,
 };
-
-interface CreatedCell {
-  cellId: CellId;
-  code: string;
-}
-
-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;
-  createCell: CreateCell;
-}
-
-class StagedCellReconciler {
-  private createdCells: CreatedCell[] = [];
-  private readonly actions: StagedCellReconcilerActions;
-  private marimoImportCellId: CellId | null = null;
-
-  public constructor(actions: StagedCellReconcilerActions) {
-    this.actions = actions;
-  }
-
-  public reconcile(completionCells: GeneratedCell[]) {
-    for (const [idx, cell] of completionCells.entries()) {
-      if (idx < this.createdCells.length) {
-        const existingCell = this.createdCells[idx];
-        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,
-        });
-      } else {
-        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);
-  }
-
-  public 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;
-  }
-
-  public accept() {
-    for (const cellId of this.stagedCellIds()) {
-      this.actions.removeStagedCell(cellId);
-    }
-  }
-
-  public hasStagedCells(stagedCells: StagedAICells) {
-    return this.stagedCellIds().some((cellId) => stagedCells.has(cellId));
-  }
-
-  public 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,
-      createCell: this.actions.createCell,
-      fromCellId: this.createdCells[0]?.cellId,
-      before: true,
-    });
-    if (cellId) {
-      this.actions.addStagedCell({ cellId, edit: { type: "add_cell" } });
-      this.marimoImportCellId = cellId;
-    }
-  }
-}

package/src/core/ai/state.ts

@@ -31,10 +31,13 @@ export const aiCompletionCellAtom = atom<AiCompletionCell | null>(null);
 export interface PendingAiPrompt {
   prompt: string;
   submit: boolean;
+  append?: boolean;
 }
 
 export const pendingAiPromptAtom = atom<PendingAiPrompt | null>(null);
 
+export const chatGenerationInProgressAtom = atom(false);
+
 const CHAT_OPTIONS_KEY = "marimo:ai:chatOptions";
 export const chatOptionsAtom = atomWithStorage<Required<ChatOptions>>(
   CHAT_OPTIONS_KEY,

package/src/core/codemirror/__tests__/setup.test.ts

@@ -40,6 +40,7 @@ function getOpts() {
     enableAI: false,
     cellActions: {
       toggleHideCode: namedFunction("toggleHideCode"),
+      askAiAboutCell: namedFunction("askAiAboutCell"),
       aiCellCompletion: namedFunction("aiCellCompletion"),
       createManyBelow: namedFunction("createManyBelow"),
       onRun: namedFunction("onRun"),
@@ -240,3 +241,34 @@ test("ai adds more extensions", () => {
     .flat();
   expect(withAI.length - 2).toBe(withoutAI.length);
 });
+
+test("empty-cell AI action opens chat without invoking the local refactor prompt", () => {
+  const askAiAboutCell = vi.fn();
+  const aiCellCompletion = vi.fn();
+  const view = new EditorView({
+    state: EditorState.create({
+      doc: "",
+      extensions: setup({
+        showPlaceholder: true,
+        enableAI: false,
+        enableAIChat: true,
+        cellActions: {
+          ...getOpts().cellActions,
+          askAiAboutCell,
+          aiCellCompletion,
+        },
+      }),
+    }),
+  });
+  try {
+    const button = view.dom.querySelector<HTMLButtonElement>(
+      ".cm-clickable-placeholder",
+    );
+    expect(button?.textContent).toBe("generate with AI");
+    button?.click();
+    expect(askAiAboutCell).toHaveBeenCalledOnce();
+    expect(aiCellCompletion).not.toHaveBeenCalled();
+  } finally {
+    view.destroy();
+  }
+});

package/src/core/codemirror/cells/state.ts

@@ -6,6 +6,7 @@ import type { CellId } from "@/core/cells/ids";
 
 export interface CodemirrorCellActions extends CellActions {
   toggleHideCode: () => boolean;
+  askAiAboutCell: () => void;
   aiCellCompletion: () => boolean;
   createManyBelow: (content: string[]) => void;
   onRun: () => void;

package/src/core/codemirror/cm.ts

@@ -74,6 +74,7 @@ export interface CodeMirrorSetupOpts {
   cellId: CellId;
   showPlaceholder: boolean;
   enableAI: boolean;
+  enableAIChat?: boolean;
   acceptCompletionOnEnter?: boolean;
   cellActions: CodemirrorCellActions;
   completionConfig: CompletionConfig;
@@ -92,7 +93,10 @@ export interface CodeMirrorSetupOpts {
 }
 
 function getPlaceholderType(opts: CodeMirrorSetupOpts) {
-  const { showPlaceholder, enableAI } = opts;
+  const { showPlaceholder, enableAI, enableAIChat } = opts;
+  if (enableAIChat) {
+    return "ai";
+  }
   return showPlaceholder ? "marimo-import" : enableAI ? "ai" : "none";
 }
 

package/src/core/codemirror/config/types.ts

@@ -4,7 +4,7 @@
  * The type of placeholder to use for the editor.
  *
  * - "marimo-import" -> Shows `import marimo as mo`
- * - "ai" -> Shows `Start coding or Generate with AI`
+ * - "ai" -> Shows `Start coding or generate with AI`
  * - "none" -> No placeholder
  */
 export type PlaceholderType = "marimo-import" | "ai" | "none";

package/src/core/codemirror/language/languages/python.ts

@@ -422,11 +422,11 @@ function getPlaceholderExtension(placeholderType: PlaceholderType): Extension {
   if (placeholderType === "ai") {
     return clickablePlaceholderExtension({
       beforeText: "Start coding or ",
-      linkText: "generate",
-      afterText: " with AI.",
+      linkText: "generate with AI",
+      afterText: ".",
       onClick: (ev) => {
         const cellActions = ev.state.facet(cellActionsState);
-        cellActions.aiCellCompletion();
+        cellActions.askAiAboutCell();
       },
     });
   }

package/src/core/codemirror/placeholder/__tests__/extensions.test.ts

@@ -81,16 +81,32 @@ describe("clickablePlaceholderExtension", () => {
 
     // Check text content
     expect(placeholder.textContent).toBe("Click here to continue");
+    expect(
+      [...placeholder.querySelectorAll('[aria-hidden="true"]')].map(
+        (node) => node.textContent,
+      ),
+    ).toEqual(["Click ", " to continue"]);
+    expect(
+      placeholder.querySelector("button")?.closest('[aria-hidden="true"]'),
+    ).toBeNull();
 
     // Find and check the clickable link
     const link = placeholder.querySelector(".cm-clickable-placeholder")!;
     expect(link).toBeTruthy();
     expect(link.textContent).toBe("here");
+    expect(link.tagName).toBe("BUTTON");
+    expect(link.closest('[aria-hidden="true"]')).toBeNull();
+    expect(link.getAttribute("type")).toBe("button");
 
     // Simulate click
     (link as HTMLElement).click();
     expect(onClick).toHaveBeenCalled();
 
+    view.dispatch({ changes: { from: 0, insert: "Hello" } });
+    expect(view.dom.querySelector(".cm-clickable-placeholder")).toBeNull();
+    view.dispatch({ changes: { from: 0, to: 5 } });
+    expect(view.dom.querySelector("button")?.textContent).toBe("here");
+
     // Check theme classes
     expect(link.classList.contains("cm-clickable-placeholder")).toBe(true);
 

package/src/core/codemirror/placeholder/extensions.ts

@@ -1,6 +1,12 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import type { Extension } from "@codemirror/state";
-import { EditorView, keymap, placeholder } from "@codemirror/view";
+import {
+  Decoration,
+  EditorView,
+  keymap,
+  placeholder,
+  WidgetType,
+} from "@codemirror/view";
 import { isEditorReadonly } from "../readonly/extension";
 
 /**
@@ -60,37 +66,69 @@ export function clickablePlaceholderExtension(opts: {
 }): Extension[] {
   const { beforeText, linkText, afterText, onClick } = opts;
 
-  // Create a placeholder
-  // Needs to be a function to keep event listeners
-  // See https://github.com/codemirror/dev/issues/1457
-  const createPlaceholder = (ev: EditorView) => {
-    const placeholderText = document.createElement("span");
-    placeholderText.append(document.createTextNode(beforeText));
-    const link = document.createElement("span");
-    link.textContent = linkText;
-    link.classList.add("cm-clickable-placeholder");
-    link.onclick = (evt) => {
-      evt.stopPropagation();
-      onClick(ev);
-    };
-    placeholderText.append(link);
-    placeholderText.append(document.createTextNode(afterText));
-    return placeholderText;
-  };
+  class ActionPlaceholder extends WidgetType {
+    public toDOM(ev: EditorView) {
+      const wrapper = document.createElement("span");
+      wrapper.className = "cm-placeholder";
+      wrapper.contentEditable = "false";
+      const before = document.createElement("span");
+      before.textContent = beforeText;
+      before.setAttribute("aria-hidden", "true");
+      wrapper.append(before);
+      const button = document.createElement("button");
+      button.type = "button";
+      button.textContent = linkText;
+      button.className = "cm-clickable-placeholder";
+      button.onclick = (event) => {
+        event.stopPropagation();
+        onClick(ev);
+      };
+      const after = document.createElement("span");
+      after.textContent = afterText;
+      after.setAttribute("aria-hidden", "true");
+      wrapper.append(button, after);
+      return wrapper;
+    }
+
+    public override ignoreEvent(event: Event) {
+      return (
+        event.target instanceof Element &&
+        event.target.closest(".cm-clickable-placeholder") !== null
+      );
+    }
+  }
 
   return [
-    placeholder(createPlaceholder),
+    EditorView.decorations.compute(["doc"], (state) =>
+      state.doc.length === 0
+        ? Decoration.set([
+            Decoration.widget({
+              widget: new ActionPlaceholder(),
+              side: 1,
+            }).range(0),
+          ])
+        : Decoration.none,
+    ),
     EditorView.theme({
       ".cm-placeholder": {
         color: "var(--slate-8)",
       },
       ".cm-clickable-placeholder": {
         cursor: "pointer !important",
+        font: "inherit",
+        background: "transparent",
+        border: "none",
+        padding: "0",
         pointerEvents: "auto",
         color: "var(--slate-9)",
         textDecoration: "underline",
         textUnderlineOffset: "0.2em",
       },
+      ".cm-clickable-placeholder:focus-visible": {
+        outline: "1px solid var(--muted-foreground)",
+        outlineOffset: "2px",
+        borderRadius: "2px",
+      },
       ".cm-clickable-placeholder:hover": {
         color: "var(--sky-11)",
       },

package/src/core/hotkeys/hotkeys.ts

@@ -221,7 +221,8 @@ const DEFAULT_HOT_KEY = {
     key: "Mod-h",
   },
   "cell.aiCompletion": {
-    name: "AI completion",
+    name: "Refactor cell",
+    additionalKeywords: ["AI", "AI completion"],
     group: "Editing",
     key: "Mod-Shift-e",
   },

package/src/core/notebook/document/__tests__/codemirror.test.ts

@@ -63,6 +63,7 @@ function languages(id: CellId) {
     cellActionsState.of({
       ...createNotebookActions(vi.fn()),
       toggleHideCode: () => false,
+      askAiAboutCell: () => undefined,
       aiCellCompletion: () => false,
       createManyBelow: vi.fn(),
       addOrAppendSetupCell: vi.fn(),

package/src/plugins/impl/chat/__tests__/chat-prompts.test.tsx

@@ -7,7 +7,7 @@ import type { PopoverContent } from "@/components/ui/popover";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import { Chatbot } from "../chat-ui";
 
-vi.mock("@/components/editor/ai/add-cell-with-ai", () => ({
+vi.mock("@/components/editor/ai/prompt-input", () => ({
   PromptInput: ({ value }: { value: string }) => (
     <textarea aria-label="Chat prompt" value={value} readOnly={true} />
   ),

package/src/plugins/impl/chat/chat-ui.tsx

@@ -39,7 +39,7 @@ import {
 import {
   type AdditionalCompletions,
   PromptInput,
-} from "@/components/editor/ai/add-cell-with-ai";
+} from "@/components/editor/ai/prompt-input";
 import { CopyClipboardIcon } from "@/components/icons/copy-icon";
 import { Spinner } from "@/components/icons/spinner";
 import { Button } from "@/components/ui/button";