@marimo-team/islands 0.25.3-dev265 → 0.25.3-dev267

package/src/components/chat/user-message.tsx

@@ -0,0 +1,145 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
+import type { FileUIPart, TextUIPart, UIMessage } from "ai";
+import { PencilIcon } from "lucide-react";
+import React from "react";
+import { Button } from "@/components/ui/button";
+import { Tooltip } from "@/components/ui/tooltip";
+import { focusInputAndMoveToEnd } from "@/core/codemirror/utils";
+import { PromptInput } from "../editor/ai/prompt-input";
+import { CopyClipboardIcon } from "../icons/copy-icon";
+import { AttachmentRenderer } from "./chat-components";
+
+interface UserMessageProps {
+  message: UIMessage;
+  canEdit: boolean;
+  onEdit: (newValue: string) => void;
+}
+
+export const UserMessage: React.FC<UserMessageProps> = ({
+  message,
+  canEdit,
+  onEdit,
+}) => {
+  const content = message.parts
+    .filter((p): p is TextUIPart => p.type === "text")
+    .map((p) => p.text)
+    .join("\n");
+  const attachments = message.parts
+    .filter((p): p is FileUIPart => p.type === "file")
+    .map((part, idx) => (
+      // Sent parts never reorder, and a file attached twice repeats its URL.
+      // react-doctor-disable-next-line react-doctor/no-array-index-as-key
+      <AttachmentRenderer attachment={part} key={idx} />
+    ));
+
+  const [draft, setDraft] = React.useState<string | null>(null);
+  const inputRef = React.useRef<ReactCodeMirrorRef>(null);
+  const editButtonRef = React.useRef<HTMLButtonElement>(null);
+  const editing = draft != null;
+
+  const startEditing = () => {
+    setDraft(content);
+    focusInputAndMoveToEnd(inputRef);
+  };
+
+  const cancel = () => {
+    setDraft(null);
+    requestAnimationFrame(() => editButtonRef.current?.focus());
+  };
+
+  const submit = (value: string) => {
+    if (!canEdit || !value.trim()) {
+      return;
+    }
+    onEdit(value);
+    setDraft(null);
+  };
+
+  if (editing) {
+    return (
+      <div className="w-full rounded-[10px] border border-border bg-background">
+        {attachments}
+        <div className="px-2 py-1.5">
+          <PromptInput
+            inputRef={inputRef}
+            className="max-h-[400px]"
+            value={draft}
+            onChange={setDraft}
+            onSubmit={(_e, value) => submit(value)}
+            onClose={cancel}
+            placeholder="Type your message. Use @ to reference variables and cells."
+          />
+        </div>
+        <div className="flex items-center gap-2 border-t border-border/20 px-3 py-1.5">
+          <span className="mr-auto text-xs text-muted-foreground">
+            Responses after this message will be replaced.
+          </span>
+          <Button variant="text" size="xs" onClick={cancel}>
+            Cancel
+          </Button>
+          <Button
+            size="xs"
+            disabled={!canEdit || !draft.trim()}
+            onClick={() => submit(draft)}
+          >
+            Send
+          </Button>
+        </div>
+      </div>
+    );
+  }
+
+  return (
+    <div className="flex max-w-[90%] flex-col items-end gap-1">
+      {attachments}
+      <div className="rounded-[10px] bg-(--slate-3) px-3 py-1.5 whitespace-pre-wrap wrap-break-word">
+        <MentionText text={content} />
+      </div>
+      <div className="flex items-center gap-1 opacity-0 transition-opacity group-hover:opacity-100 group-focus-within:opacity-100">
+        <Tooltip content="Edit">
+          <Button
+            ref={editButtonRef}
+            variant="text"
+            size="icon"
+            aria-label="Edit message"
+            className="chat-control"
+            disabled={!canEdit}
+            onClick={startEditing}
+          >
+            <PencilIcon className="size-3" />
+          </Button>
+        </Tooltip>
+        <CopyClipboardIcon
+          value={content}
+          className="size-3"
+          buttonClassName="chat-control flex size-6 items-center justify-center rounded-md"
+        />
+      </div>
+    </div>
+  );
+};
+
+const MENTION = /([@#])[\w-]+:\/\/([\w./-]+)/g;
+
+const MentionText: React.FC<{ text: string }> = ({ text }) => {
+  const nodes: React.ReactNode[] = [];
+  let last = 0;
+  for (const match of text.matchAll(MENTION)) {
+    nodes.push(text.slice(last, match.index));
+    nodes.push(
+      <span
+        key={match.index}
+        title={match[0].slice(1)}
+        className="rounded-sm bg-(--slate-5) px-1 font-mono text-[0.85em]"
+      >
+        {match[1]}
+        {match[2]}
+      </span>,
+    );
+    last = match.index + match[0].length;
+  }
+  nodes.push(text.slice(last));
+  return nodes;
+};

package/src/components/editor/actions/__tests__/useCellActionButton.test.tsx

@@ -10,6 +10,7 @@ import { partial } from "@/__mocks__/partial";
 import { MockRequestClient } from "@/__mocks__/requests";
 import { ModalProvider } from "@/components/modal/ImperativeModal";
 import { toast } from "@/components/ui/use-toast";
+import { type Edit, stagedAICellsAtom } from "@/core/ai/staged-cells";
 import { CellId } from "@/core/cells/ids";
 import { requestClientAtom } from "@/core/network/requests";
 import type { NotebookReplica } from "@/core/notebook/replica";
@@ -21,8 +22,11 @@ import { useCellActionButtons } from "../useCellActionButton";
 vi.mock("@/utils/copy", () => ({ copyToClipboard: vi.fn() }));
 vi.mock("@/components/ui/use-toast", () => ({ toast: vi.fn() }));
 
-function renderActions(cellId: CellId, name: string) {
+function renderActions(cellId: CellId, name: string, edit?: Edit) {
   const store = createStore();
+  if (edit) {
+    store.set(stagedAICellsAtom, new Map([[cellId, edit]]));
+  }
   const { model, documentAtom, stop } = mockDocumentModel(store);
   onTestFinished(stop);
   store.set(requestClientAtom, MockRequestClient.create());
@@ -111,3 +115,16 @@ describe("useCellActionButtons", () => {
     expect(getAction(groups, "Copy cell ID").disabled).toBeUndefined();
   });
 });
+
+it("offers individual AI review actions in the cell menu only for pending changes", () => {
+  const id = CellId.create();
+  const pending = renderActions(id, "_", { type: "add_cell" });
+  expect(pending[0].map((action) => action.label)).toEqual([
+    "Keep cell",
+    "Discard cell",
+  ]);
+  const ordinary = renderActions(id, "_");
+  expect(ordinary.flat().some((action) => action.label === "Keep cell")).toBe(
+    false,
+  );
+});

package/src/components/editor/actions/useCellActionButton.tsx

@@ -63,6 +63,7 @@ import { copyToClipboard } from "@/utils/copy";
 import { downloadCellOutputAsImage } from "@/utils/download";
 import { MarkdownIcon, PythonIcon } from "../cell/code/icons";
 import { useDeleteCellCallback } from "../cell/useDeleteCell";
+import { useStagedAICellActions } from "../cell/StagedAICell";
 import { useRunCell } from "../cell/useRunCells";
 import { useSplitCellCallback } from "../cell/useSplitCell";
 import { NameCellInput } from "./name-cell-input";
@@ -93,6 +94,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
     clearCellOutput,
     createCell,
   } = useCellActions();
+  const stagedActions = useStagedAICellActions(cell?.cellId);
   const showHiddenCode = useSetAtom(showHiddenCodeAtom);
   const edit = useNotebookDocumentEdit();
   const splitCell = useSplitCellCallback();
@@ -144,6 +146,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
 
   // Actions
   const actions: ActionButton[][] = [
+    stagedActions,
     [
       {
         icon: <PlayIcon size={13} strokeWidth={1.5} />,

package/src/components/editor/ai/completion-handlers.tsx

@@ -46,43 +46,6 @@ export const createAiCompletionOnKeydown = (opts: {
   };
 };
 
-export const CompletionActionsCellFooter: React.FC<{
-  isLoading: boolean;
-  onAccept: () => void;
-  onDecline: () => void;
-  size?: "xs" | "sm";
-  multipleCompletions?: boolean;
-  runCell?: () => void;
-  acceptLabel?: string;
-  declineLabel?: string;
-}> = ({
-  isLoading,
-  onAccept,
-  onDecline,
-  runCell,
-  acceptLabel,
-  declineLabel,
-  size,
-}) => {
-  return (
-    <>
-      <AcceptCompletionButton
-        isLoading={isLoading}
-        onAccept={onAccept}
-        size={size}
-        runCell={runCell}
-        label={acceptLabel}
-      />
-      <RejectCompletionButton
-        onDecline={onDecline}
-        size={size}
-        label={declineLabel}
-        disabled={isLoading}
-      />
-    </>
-  );
-};
-
 export const AcceptCompletionButton: React.FC<{
   isLoading: boolean;
   onAccept: () => void;

package/src/components/editor/cell/StagedAICell.tsx

@@ -1,24 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtomValue, useStore } from "jotai";
-import { SparklesIcon } from "lucide-react";
-import { Tooltip } from "@/components/ui/tooltip";
+import { atom, useAtomValue, useStore } from "jotai";
+import { useMemo } from "react";
 import {
   type Edit,
+  stagedAICellsAtom,
   useStagedAICell,
   useStagedCells,
 } from "@/core/ai/staged-cells";
 import { chatGenerationInProgressAtom } from "@/core/ai/state";
-import { getCellEditorView } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
-import { cn } from "@/utils/cn";
-import { Logger } from "@/utils/Logger";
-import {
-  AI_EDIT_ACTION_LABELS,
-  CompletionActionsCellFooter,
-} from "../ai/completion-handlers";
-import { useRunCell } from "./useRunCells";
+import { AI_EDIT_ACTION_LABELS } from "../ai/completion-handlers";
+import type { ActionButton } from "../actions/types";
 
 const actionLabels: Record<Edit["type"], { accept: string; decline: string }> =
   {
@@ -27,142 +20,56 @@ const actionLabels: Record<Edit["type"], { accept: string; decline: string }> =
     delete_cell: { accept: "Delete cell", decline: "Keep cell" },
   };
 
-export const StagedAICellBackground: React.FC<{
-  cellId: CellId;
-  className?: string;
-}> = ({ cellId, className }) => {
-  const stagedCell = useStagedAICell(cellId);
-
-  if (!stagedCell) {
-    return null;
-  }
-
-  const cellClass =
-    stagedCell.type === "delete_cell"
-      ? "mo-ai-deleted-cell"
-      : "mo-ai-generated-cell";
-
-  return <div className={cn(cellClass, className)} />;
-};
-
-export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
+export const StagedAIDeletionMarker: React.FC<{ cellId: CellId }> = ({
   cellId,
-}) => {
+}) =>
+  useStagedAICell(cellId)?.type === "delete_cell" ? (
+    <div className="mo-ai-deleted-cell" />
+  ) : null;
+
+export function useStagedAICellActions(
+  cellId: CellId | undefined,
+): ActionButton[] {
   const store = useStore();
-  const stagedAiCell = useStagedAICell(cellId);
-  const isGenerating = useAtomValue(chatGenerationInProgressAtom);
-  const runCell = useRunCell(cellId);
-
-  const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
-
-  if (!stagedAiCell) {
-    return null;
-  }
-
-  const isDeletion = stagedAiCell.type === "delete_cell";
-  const labels = actionLabels[stagedAiCell.type];
-
-  const handleCompletion = (type: "accept" | "reject") => {
-    const completionFunc =
-      type === "accept" ? acceptStagedCell : rejectStagedCell;
-    completionFunc(cellId, stagedAiCell, removeStagedCell, deleteStagedCell);
-  };
-
-  const tooltipContent =
-    stagedAiCell.type === "add_cell"
-      ? "AI-generated cell"
-      : isDeletion
-        ? "AI suggests deleting this cell"
-        : "AI changed this cell";
-
-  return (
-    <div className="mo-ai-cell-footer">
-      <div className="flex items-center gap-1.5 text-xs text-muted-foreground">
-        <Tooltip content={tooltipContent}>
-          <span
-            className="inline-flex items-center gap-1.5"
-            aria-label={tooltipContent}
-          >
-            <SparklesIcon
-              aria-hidden="true"
-              className="size-3.5 text-(--blue-11)"
-            />
-            <span>AI</span>
-          </span>
-        </Tooltip>
-      </div>
-      <div className="flex items-center gap-1.5">
-        <CompletionActionsCellFooter
-          isLoading={isGenerating}
-          onAccept={() => handleCompletion("accept")}
-          onDecline={() => handleCompletion("reject")}
-          size="xs"
-          runCell={isDeletion ? undefined : runCell}
-          acceptLabel={labels.accept}
-          declineLabel={labels.decline}
-        />
-      </div>
-    </div>
+  const edit = useStagedAICell(cellId);
+  // Every cell calls this; only cells with a pending edit care about the lock.
+  const isLocked = useAtomValue(
+    useMemo(
+      () =>
+        atom(
+          (get) =>
+            cellId !== undefined &&
+            get(stagedAICellsAtom).has(cellId) &&
+            get(chatGenerationInProgressAtom),
+        ),
+      [cellId],
+    ),
   );
-};
+  const { acceptStagedCell, rejectStagedCell } = useStagedCells();
 
-/**
- * Accept a staged cell and apply the changes.
- */
-export function acceptStagedCell(
-  cellId: CellId,
-  edit: Edit,
-  removeStagedCell: (cellId: CellId) => void,
-  deleteStagedCell: (cellId: CellId) => void,
-): void {
-  switch (edit.type) {
-    case "delete_cell":
-      // For delete cells, the cell is deleted when the completion is accepted
-      deleteStagedCell(cellId);
-      break;
-    default:
-      removeStagedCell(cellId);
-      break;
+  if (!edit || cellId === undefined) {
+    return [];
   }
-}
-
-/**
- * Reject a staged cell and revert the changes.
- */
-export function rejectStagedCell(
-  cellId: CellId,
-  edit: Edit,
-  removeStagedCell: (cellId: CellId) => void,
-  deleteStagedCell: (cellId: CellId) => void,
-): void {
-  switch (edit.type) {
-    case "update_cell": {
-      // Revert cell code
-      const editorView = getCellEditorView(cellId);
-      if (!editorView) {
-        Logger.error("Editor for this cell not found", { cellId });
-        break;
-      }
 
-      const update = prepareCellCodeUpdate(
-        editorView.state,
-        edit.previousCode,
-        { initialize: true },
-      );
-
-      if (update) {
-        editorView.dispatch(update);
-      }
-      removeStagedCell(cellId);
-      break;
-    }
-    case "add_cell":
-      // Delete the cell since it's newly created
-      deleteStagedCell(cellId);
-      break;
-    case "delete_cell":
-      // Just remove the deletion marker - cell stays in notebook
-      removeStagedCell(cellId);
-      break;
-  }
+  const labels = actionLabels[edit.type];
+  return [
+    {
+      label: labels.accept,
+      disabled: isLocked,
+      handle: () => {
+        if (!store.get(chatGenerationInProgressAtom)) {
+          acceptStagedCell(cellId, edit);
+        }
+      },
+    },
+    {
+      label: labels.decline,
+      disabled: isLocked,
+      handle: () => {
+        if (!store.get(chatGenerationInProgressAtom)) {
+          rejectStagedCell(cellId, edit);
+        }
+      },
+    },
+  ];
 }