@marimo-team/islands 0.25.3-dev266 → 0.25.3-dev268

package/src/components/chat/turn-timing.ts

@@ -0,0 +1,73 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { UIMessage } from "ai";
+
+export interface TurnTiming {
+  startedAt: number;
+  finishedAt?: number;
+  stopped?: boolean;
+}
+
+export function getTurnTiming(message: UIMessage): TurnTiming | undefined {
+  const metadata = message.metadata;
+  if (typeof metadata !== "object" || metadata === null) {
+    return undefined;
+  }
+  const { startedAt, finishedAt, stopped } = metadata as Record<
+    string,
+    unknown
+  >;
+  if (typeof startedAt !== "number") {
+    return undefined;
+  }
+  return {
+    startedAt,
+    finishedAt: typeof finishedAt === "number" ? finishedAt : undefined,
+    stopped: stopped === true,
+  };
+}
+
+export function withTurnTiming(
+  messages: UIMessage[],
+  turn: { startedAt: number | null; stopped: boolean },
+  now = Date.now(),
+): UIMessage[] {
+  const last = messages.at(-1);
+  if (last?.role !== "assistant") {
+    return messages;
+  }
+  const startedAt = turn.startedAt ?? getTurnTiming(last)?.startedAt;
+  if (startedAt == null) {
+    return messages;
+  }
+  const metadata =
+    typeof last.metadata === "object" && last.metadata !== null
+      ? last.metadata
+      : {};
+  return [
+    ...messages.slice(0, -1),
+    {
+      ...last,
+      metadata: {
+        ...metadata,
+        startedAt,
+        finishedAt: now,
+        stopped: turn.stopped,
+      },
+    },
+  ];
+}
+
+export function formatDuration(ms: number): string {
+  const totalSeconds = Math.max(1, Math.round(ms / 1000));
+  const hours = Math.floor(totalSeconds / 3600);
+  const minutes = Math.floor((totalSeconds % 3600) / 60);
+  const seconds = totalSeconds % 60;
+  if (hours > 0) {
+    return minutes > 0 ? `${hours}h ${minutes}m` : `${hours}h`;
+  }
+  if (minutes > 0) {
+    return seconds > 0 ? `${minutes}m ${seconds}s` : `${minutes}m`;
+  }
+  return `${seconds}s`;
+}

package/src/components/chat/use-assistant-turn.ts

@@ -0,0 +1,65 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { UIMessage } from "ai";
+import { useState } from "react";
+import useEvent from "react-use-event-hook";
+import { willAutoContinue } from "./chat-utils";
+import { getTurnTiming, withTurnTiming } from "./turn-timing";
+
+export function useAssistantTurn() {
+  const [turn, setTurn] = useState<{
+    startedAt: number;
+    ended: boolean;
+  } | null>(null);
+
+  const start = useEvent(() =>
+    setTurn({ startedAt: Date.now(), ended: false }),
+  );
+  const reset = useEvent(() => setTurn(null));
+  const stop = useEvent(() =>
+    setTurn((prev) => prev && { ...prev, ended: true }),
+  );
+
+  // Shift the start forward so time spent waiting on the user isn't counted.
+  const resume = useEvent((reply: UIMessage | undefined) => {
+    const timing = reply && getTurnTiming(reply);
+    const workedMs =
+      timing?.finishedAt == null ? 0 : timing.finishedAt - timing.startedAt;
+    setTurn({ startedAt: Date.now() - workedMs, ended: false });
+  });
+
+  const finish = useEvent(
+    (finished: UIMessage[], opts: { isAbort: boolean; isError: boolean }) => {
+      if (opts.isAbort || opts.isError) {
+        stop();
+      }
+      return withTurnTiming(finished, {
+        startedAt: turn?.startedAt ?? null,
+        stopped: opts.isAbort,
+      });
+    },
+  );
+
+  // Between requests of one turn the SDK briefly reports "ready".
+  const isActive = (state: {
+    isLoading: boolean;
+    hasError: boolean;
+    messages: UIMessage[];
+  }) =>
+    state.isLoading ||
+    (turn != null &&
+      !turn.ended &&
+      !state.hasError &&
+      willAutoContinue(state.messages));
+
+  return {
+    startedAt: turn?.startedAt ?? null,
+    hasEnded: turn?.ended === true,
+    start,
+    resume,
+    stop,
+    reset,
+    finish,
+    isActive,
+  };
+}

package/src/components/chat/use-chat-generation.ts

@@ -5,15 +5,16 @@ import { useSetAtom } from "jotai";
 import { useEffect } from "react";
 import { chatGenerationInProgressAtom } from "@/core/ai/state";
 
-export function useChatGeneration(status: ChatStatus) {
+export function isChatBusy(status: ChatStatus): boolean {
+  return status === "submitted" || status === "streaming";
+}
+
+export function useChatGeneration(isGenerating: boolean) {
   const setGenerating = useSetAtom(chatGenerationInProgressAtom);
-  const isGenerating = status === "submitted" || status === "streaming";
 
   useEffect(() => {
     setGenerating(isGenerating);
     // useChat stops its request when the panel unmounts or becomes hidden.
     return () => setGenerating(false);
   }, [isGenerating, setGenerating]);
-
-  return isGenerating;
 }

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/dependency-graph/minimap-content.tsx

@@ -7,11 +7,12 @@ import {
   cellGraphsAtom,
   isVariableAffectedBySelectedCell,
 } from "@/components/editor/chrome/wrapper/minimap-state";
-import { useCellData, useCellRuntime } from "@/core/cells/cells";
+import { useCellData } from "@/core/cells/cells";
 import { useCellActions } from "@/core/cells/actions";
 import { cellFocusAtom, useCellFocusActions } from "@/core/cells/focus";
 import type { CellId } from "@/core/cells/ids";
 import { useLayout } from "@/core/cells/layout";
+import { useCellExecution } from "@/core/executions";
 import { useVariables } from "@/core/variables/state";
 import { cn } from "@/utils/cn";
 import { extractCellPreview } from "./utils/cell-preview";
@@ -32,7 +33,7 @@ const MinimapCell: React.FC<MinimapCellProps> = (props) => {
     id: props.cellId,
     graph: graphs[props.cellId],
     code: useCellData(props.cellId).code,
-    hasError: useCellRuntime(props.cellId).errored,
+    hasError: useCellExecution(props.cellId).outcome === "error",
   };
 
   let selectedCell: undefined | { id: CellId; graph: CellGraph };

package/src/components/editor/__tests__/data-attributes.test.tsx

@@ -104,16 +104,9 @@ describe("Cell data attributes", () => {
       });
 
       notebook.cellRuntime[cid] = createCellRuntimeState({
-        status: "idle",
         output: null,
         consoleOutputs: [],
-        interrupted: false,
-        errored: false,
-        stopped: false,
         staleInputs: false,
-        runStartTimestamp: null,
-        lastRunStartTimestamp: null,
-        runElapsedTimeMs: null,
         debuggerActive: false,
         outline: null,
       });

package/src/components/editor/__tests__/output-persistence.test.tsx

@@ -4,6 +4,7 @@ import { createStore, Provider } from "jotai";
 import { beforeAll, describe, expect, it } from "vitest";
 import { SetupMocks } from "@/__mocks__/common";
 import { mockDocumentModel } from "@/__mocks__/document";
+import { MockExecutions } from "@/__mocks__/executions";
 import { MockNotebook } from "@/__mocks__/notebook";
 import { MockRequestClient } from "@/__mocks__/requests";
 import { cellId } from "@/__tests__/branded";
@@ -50,7 +51,6 @@ function seedNotebook(
     },
   });
   notebook.cellRuntime[cid] = createCellRuntimeState({
-    status: "idle",
     output: {
       channel: "output",
       mimetype: "text/html",
@@ -207,8 +207,8 @@ describe("output persistence across mode changes", () => {
 
   it("hides errored cells while presenting without unmounting them", () => {
     const { store, wrapper } = createTestWrapper();
+    MockExecutions.seed(store, [MockExecutions.errored(cid)]);
     seedNotebook(store, {
-      errored: true,
       output: {
         channel: "marimo-error",
         mimetype: "application/vnd.marimo+error",

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());
@@ -48,6 +52,7 @@ function renderActions(cellId: CellId, name: string) {
           name,
           config: { disabled: false, hide_code: false, column: null },
           status: "idle",
+          blocked: false,
           hasOutput: false,
           hasConsoleOutput: false,
           getEditorView: () => null,
@@ -111,3 +116,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

@@ -56,13 +56,15 @@ import {
   appWidthAtom,
   autoInstantiateAtom,
 } from "@/core/config/config";
-import type { CellConfig, RuntimeState } from "@/core/network/types";
+import type { CellRunStatus } from "@/core/executions";
+import type { CellConfig } from "@/core/network/types";
 import { canEditAtom } from "@/core/notebook";
 import { canLinkToCell, createCellLink } from "@/utils/cell-urls";
 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";
@@ -73,7 +75,9 @@ export interface CellActionButtonProps extends Pick<
   "name" | "config"
 > {
   cellId: CellId;
-  status: RuntimeState;
+  status: CellRunStatus;
+  /** Whether a cell it depends on is disabled. */
+  blocked: boolean;
   hasOutput: boolean;
   hasConsoleOutput: boolean;
   getEditorView: () => EditorView | null;
@@ -93,6 +97,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();
@@ -119,6 +124,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
     hasOutput,
     hasConsoleOutput,
     status,
+    blocked,
   } = cell;
 
   const toggleDisabled = () => {
@@ -144,6 +150,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
 
   // Actions
   const actions: ActionButton[][] = [
+    stagedActions,
     [
       {
         icon: <PlayIcon size={13} strokeWidth={1.5} />,
@@ -152,7 +159,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         hidden:
           status === "running" ||
           status === "queued" ||
-          status === "disabled-transitively" ||
+          blocked ||
           config.disabled,
         redundant: true,
         handle: () => runCell(),

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/alerts/stdin-blocking-alert.tsx

@@ -2,6 +2,7 @@
 
 import { atom, useAtomValue } from "jotai";
 import { notebookAtom } from "@/core/cells/cells";
+import { cellExecutionInputsAtom, cellExecutionView } from "@/core/executions";
 import { Logger } from "@/utils/Logger";
 import { Button } from "../../ui/button";
 import { Kbd } from "../../ui/kbd";
@@ -10,12 +11,13 @@ import { FloatingAlert } from "./floating-alert";
 // Atom to check if there is a cell with unresolved stdin
 const hasBlockingStdinAtom = atom((get) => {
   const notebook = get(notebookAtom);
+  const executions = get(cellExecutionInputsAtom);
 
   // Check each cell in order
   for (const cellId of notebook.cellIds) {
     // Check if the cell is idle and has unresolved stdin
     const runtime = notebook.cellRuntime[cellId];
-    if (runtime.status === "idle") {
+    if (cellExecutionView(executions, cellId).status === "idle") {
       const hasUnresolvedStdin = runtime.consoleOutputs.some(
         (output) => output.channel === "stdin" && output.response === undefined,
       );

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

@@ -7,24 +7,29 @@ import {
 } from "lucide-react";
 import { useDateFormatter } from "react-aria";
 import { MultiIcon } from "@/components/icons/multi-icon";
+import type { CellRunStatus } from "@/core/executions";
 import { Logger } from "@/utils/Logger";
-import type { CellRuntimeState } from "../../../core/cells/types";
 import { useElapsedTime } from "../../../hooks/useElapsedTime";
 import { Tooltip } from "../../ui/tooltip";
 
 import "./cell-status.css";
 import { formatDistanceToNow } from "date-fns";
-import { formatElapsedTime, Time } from "@/utils/time";
+import { formatElapsedTime, type Milliseconds, Time } from "@/utils/time";
 
-export interface CellStatusComponentProps extends Pick<
-  CellRuntimeState,
-  "status" | "runStartTimestamp" | "interrupted" | "lastRunStartTimestamp"
-> {
+export interface CellStatusComponentProps {
   editing: boolean;
   edited: boolean;
   disabled: boolean;
+  /** Whether a cell it depends on is disabled. */
+  blocked: boolean;
   staleInputs: boolean;
-  elapsedTime: number | null;
+  status: CellRunStatus;
+  /** Whether the cell's last run was interrupted. */
+  interrupted: boolean;
+  /** How long the last run took, once it has ended. */
+  elapsedTime: Milliseconds | null;
+  /** When the current run started, or the last one, since the epoch. */
+  startedAt: Milliseconds | null;
   uninstantiated: boolean;
 }
 
@@ -32,20 +37,21 @@ export const CellStatusComponent: React.FC<CellStatusComponentProps> = ({
   editing,
   status,
   disabled,
+  blocked,
   staleInputs,
   edited,
   interrupted,
   elapsedTime,
-  runStartTimestamp,
-  lastRunStartTimestamp,
+  startedAt,
   uninstantiated,
 }) => {
   if (!editing) {
     return null;
   }
 
-  const start = runStartTimestamp ?? lastRunStartTimestamp;
-  const lastRanTime = start ? <LastRanTime lastRanTime={start} /> : null;
+  const lastRanTime = startedAt ? (
+    <LastRanTime lastRanTime={startedAt} />
+  ) : null;
 
   // unexpected states
   if (disabled && status === "running") {
@@ -103,7 +109,7 @@ export const CellStatusComponent: React.FC<CellStatusComponentProps> = ({
   }
 
   // disabled from parent
-  if (!staleInputs && status === "disabled-transitively") {
+  if (!staleInputs && blocked) {
     return (
       <Tooltip
         content={
@@ -131,7 +137,7 @@ export const CellStatusComponent: React.FC<CellStatusComponentProps> = ({
   }
 
   // stale from parent being disabled
-  if (staleInputs && status === "disabled-transitively") {
+  if (staleInputs && blocked) {
     return (
       <Tooltip
         content={
@@ -175,9 +181,7 @@ export const CellStatusComponent: React.FC<CellStatusComponentProps> = ({
           data-testid="cell-status"
           data-status="running"
         >
-          <CellTimer
-            startTime={Time.fromSeconds(runStartTimestamp) || Time.now()}
-          />
+          <CellTimer startTime={startedAt ?? Time.now().toMilliseconds()} />
         </div>
       </Tooltip>
     );
@@ -307,7 +311,7 @@ export const ElapsedTime = (props: { elapsedTime: string }) => {
   );
 };
 const LastRanTime = (props: { lastRanTime: number }) => {
-  const date = new Date(props.lastRanTime * 1000);
+  const date = new Date(props.lastRanTime);
   const today = new Date();
 
   // Looks like HH:MM:SS.SSS AM/PM
@@ -348,7 +352,7 @@ const LastRanTime = (props: { lastRanTime: number }) => {
   );
 };
 
-const CellTimer = (props: { startTime: Time }) => {
-  const time = useElapsedTime(props.startTime.toMilliseconds());
+const CellTimer = (props: { startTime: Milliseconds }) => {
+  const time = useElapsedTime(props.startTime);
   return <span>{formatElapsedTime(time)}</span>;
 };