@marimo-team/islands 0.25.3-dev27 → 0.25.3-dev271

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/ai/prompt-input.tsx

@@ -0,0 +1,189 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  autocompletion,
+  type Completion,
+  type CompletionContext,
+  type CompletionSource,
+} from "@codemirror/autocomplete";
+import { markdown } from "@codemirror/lang-markdown";
+import { Prec } from "@codemirror/state";
+import { promptHistory } from "@marimo-team/codemirror-ai";
+import ReactCodeMirror, {
+  EditorView,
+  keymap,
+  minimalSetup,
+  type ReactCodeMirrorRef,
+} from "@uiw/react-codemirror";
+import { useStore } from "jotai";
+import { useMemo } from "react";
+import useEvent from "react-use-event-hook";
+import { z } from "zod";
+import { clickablePlaceholderExtension } from "@/core/codemirror/placeholder/extensions";
+import { resourceExtension } from "@/core/codemirror/ai/resources";
+import { useTheme } from "@/theme/useTheme";
+import { cn } from "@/utils/cn";
+import { ZodLocalStorage } from "@/utils/storage/typed";
+import { CONTEXT_TRIGGER, mentionsCompletionSource } from "./completion-utils";
+
+const KEY = "marimo:ai-prompt-history";
+const promptHistoryStorage = new ZodLocalStorage(z.array(z.string()), () => []);
+
+export interface AdditionalCompletions {
+  triggerCompletionRegex: RegExp;
+  completions: Completion[];
+}
+
+interface PromptInputProps {
+  inputRef?: React.RefObject<ReactCodeMirrorRef | null>;
+  placeholder?: string;
+  placeholderAction?: Parameters<typeof clickablePlaceholderExtension>[0];
+  value: string;
+  className?: string;
+  onKeyDown?: (e: React.KeyboardEvent<HTMLDivElement>) => void;
+  onClose: () => void;
+  onChange: (value: string) => void;
+  onSubmit: (e: KeyboardEvent | undefined, value: string) => void;
+  additionalCompletions?: AdditionalCompletions;
+  maxHeight?: string;
+  onAddFiles?: (files: File[]) => void;
+}
+
+/**
+ * CodeMirror-based input for the AI prompt.
+ *
+ * This is just text (no language support), but we use codemirror to get autocomplete
+ * for @dataframe and @sql_table.
+ */
+export const PromptInput = ({
+  value,
+  placeholder,
+  placeholderAction,
+  inputRef,
+  className,
+  onChange,
+  onSubmit,
+  onKeyDown,
+  onClose,
+  onAddFiles,
+  additionalCompletions,
+  maxHeight,
+}: PromptInputProps) => {
+  const handleSubmit = useEvent(onSubmit);
+  const handleEscape = useEvent(onClose);
+  const store = useStore();
+  const { theme } = useTheme();
+
+  const additionalCompletionsSource: CompletionSource = useEvent(
+    (context: CompletionContext) => {
+      if (!additionalCompletions) {
+        return null;
+      }
+
+      return mentionsCompletionSource(
+        [additionalCompletions.triggerCompletionRegex],
+        additionalCompletions.completions,
+      )(context);
+    },
+  );
+
+  // Changing extensions can be expensive, so
+  // it is worth making sure this is memoized well.
+  const extensions = useMemo(() => {
+    const markdownLanguage = markdown();
+    return [
+      ...(placeholderAction
+        ? clickablePlaceholderExtension(placeholderAction)
+        : []),
+      autocompletion({}),
+      markdownLanguage,
+      resourceExtension({
+        language: markdownLanguage.language,
+        store,
+        onAddFiles,
+      }),
+      markdownLanguage.language.data.of({
+        autocomplete: additionalCompletionsSource,
+      }),
+      promptHistory({
+        storage: {
+          load: () => promptHistoryStorage.get(KEY),
+          save: (prompts) => promptHistoryStorage.set(KEY, prompts),
+        },
+      }),
+      EditorView.lineWrapping,
+      minimalSetup(),
+      Prec.highest(
+        keymap.of([
+          {
+            preventDefault: true,
+            stopPropagation: true,
+            any: (view, event) => {
+              const pressedModOrShift =
+                event.metaKey || event.ctrlKey || event.shiftKey;
+              // If no mod key is pressed, submit
+              if (event.key === "Enter" && !pressedModOrShift) {
+                handleSubmit(event, view.state.doc.toString());
+                event.preventDefault();
+                event.stopPropagation();
+                return true;
+              }
+              // Mod+Enter does add a new line already by codemirror
+              // But Shift+Enter does not, so we need to handle it manually
+              if (event.key === "Enter" && event.shiftKey) {
+                view.dispatch(view.state.replaceSelection("\n"));
+                event.preventDefault();
+                event.stopPropagation();
+                return true;
+              }
+
+              return false;
+            },
+          },
+        ]),
+      ),
+      keymap.of([
+        {
+          key: "Escape",
+          preventDefault: true,
+          stopPropagation: true,
+          run: () => {
+            handleEscape();
+            return true;
+          },
+        },
+      ]),
+    ];
+  }, [
+    store,
+    placeholderAction,
+    onAddFiles,
+    additionalCompletionsSource,
+    handleSubmit,
+    handleEscape,
+  ]);
+
+  return (
+    <ReactCodeMirror
+      ref={inputRef}
+      className={cn(
+        "mo-ai-prompt flex-1 font-sans text-sm leading-6 overflow-auto my-1",
+        className,
+      )}
+      width="100%"
+      maxHeight={maxHeight}
+      value={value}
+      basicSetup={false}
+      extensions={extensions}
+      onChange={onChange}
+      onKeyDown={onKeyDown}
+      theme={theme === "dark" ? "dark" : "light"}
+      placeholder={
+        placeholderAction
+          ? undefined
+          : placeholder ||
+            `Generate with AI, ${CONTEXT_TRIGGER} to include context`
+      }
+    />
+  );
+};

package/src/components/editor/alerts/__tests__/connection-notice.test.tsx

@@ -11,7 +11,7 @@ import { AppConfigSchema } from "@/core/config/config-schema";
 import { kernelStartupErrorAtom } from "@/core/errors/state";
 import { connectionAtom } from "@/core/network/connection";
 import { sandboxAtom, sandboxSyncAtom } from "@/core/packages/sandbox-state";
-import { WebSocketClosedReason, ConnectionState } from "@/core/websocket/types";
+import { WebSocketClosedReason, ConnectionState } from "@/core/session/types";
 import { chromeAtom } from "../../chrome/state";
 import { runtimeDetailsExpandedAtom } from "../../chrome/panels/runtime-details-state";
 import { ConnectionNotice } from "../connection-notice";

package/src/components/editor/alerts/connection-notice.tsx

@@ -8,7 +8,7 @@ import type { AppConfig } from "@/core/config/config-schema";
 import { connectionAtom } from "@/core/network/connection";
 import { useConnectionNotice } from "@/core/network/useConnectionNotice";
 import { sandboxAtom } from "@/core/packages/sandbox-state";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { useInterval } from "@/hooks/useInterval";
 import { cn } from "@/utils/cn";
 import { runtimeDetailsExpandedAtom } from "../chrome/panels/runtime-details-state";

package/src/components/editor/alerts/startup-logs-alert.tsx

@@ -51,6 +51,7 @@ export const StartupLogsAlert: React.FC = () => {
             {isRunning ? "Initializing..." : "Initialized"}
           </span>
           <Button
+            aria-label="Dismiss startup logs"
             variant="text"
             data-testid="remove-startup-logs-button"
             size="icon"

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.inOrderIds) {
+  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/app-container.tsx

@@ -8,7 +8,7 @@ import {
   type ConnectionStatus,
   WebSocketClosedReason,
   ConnectionState,
-} from "@/core/websocket/types";
+} from "@/core/session/types";
 import { cn } from "@/utils/cn";
 import { DynamicFavicon } from "./dynamic-favicon";
 import { StatusOverlay } from "./header/status";

package/src/components/editor/boundary/__tests__/ErrorBoundary.test.tsx

@@ -14,7 +14,7 @@ import { requestClientAtom } from "@/core/network/requests";
 import type { EnvironmentInfo } from "@/core/network/types";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { ErrorBoundary } from "../ErrorBoundary";
 
 vi.mock("@/utils/copy", () => ({

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>;
 };

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

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { DatabaseIcon, DiamondPlusIcon, PlusIcon } from "lucide-react";
-import { useState } from "react";
+import { useRef, useState } from "react";
 import { Button } from "@/components/editor/inputs/Inputs";
 import { MinimalHotkeys } from "@/components/shortcuts/renderShortcut";
 import {
@@ -10,14 +10,14 @@ import {
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { useCellActions } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import { LanguageAdapters } from "@/core/codemirror/language/LanguageAdapters";
 import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages/markdown";
 import {
   getConnectionTooltip,
   isAppInteractionDisabled,
-} from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+} from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { cn } from "@/utils/cn";
 import { Tooltip } from "../../ui/tooltip";
 import { MarkdownIcon, PythonIcon } from "./code/icons";
@@ -33,10 +33,11 @@ export const CreateCellButton = ({
   onClick: ((opts: { code: string; hideCode?: boolean }) => void) | undefined;
   oneClickShortcut: "shift" | "mod";
 }) => {
-  const { createNewCell, addSetupCellIfDoesntExist } = useCellActions();
+  const { createCell, addSetupCell } = useCellActions();
   const shortcut = `${oneClickShortcut}-Click`;
   const [open, setOpen] = useState(false);
-  const [justOpened, setJustOpened] = useState(false);
+  // Time the dropdown last opened, to ignore the click that opened it
+  const openedAt = useRef(0);
 
   const baseTooltipContent =
     getConnectionTooltip(connectionState) || tooltipContent;
@@ -61,7 +62,7 @@ export const CreateCellButton = ({
   // enabled; the user experience is confusing otherwise (how does the user
   // know they need to run import marimo as mo. first?).
   const addMarkdownCell = () => {
-    maybeAddMarimoImport({ autoInstantiate: true, createNewCell });
+    maybeAddMarimoImport({ autoInstantiate: true, createCell });
     onClick?.({
       code: LanguageAdapters.markdown.defaultCode,
       hideCode: MARKDOWN_INITIAL_HIDE_CODE,
@@ -69,23 +70,17 @@ export const CreateCellButton = ({
   };
 
   const addSQLCell = () => {
-    maybeAddMarimoImport({ autoInstantiate: true, createNewCell });
+    maybeAddMarimoImport({ autoInstantiate: true, createCell });
     onClick?.({ code: LanguageAdapters.sql.defaultCode });
   };
 
-  const addSetupCell = () => {
-    addSetupCellIfDoesntExist({});
-  };
-
   const renderIcon = (icon: React.ReactNode) => {
     return <div className="mr-3 text-muted-foreground">{icon}</div>;
   };
 
   const openDropdown = () => {
     setOpen(true);
-    setJustOpened(true);
-    // Allow interactions after a brief delay to prevent the dropdown items immediately being clicked
-    setTimeout(() => setJustOpened(false), 200);
+    openedAt.current = Date.now();
   };
 
   // We use onPointerDownCapture (not onPointerDown) to intercept events in
@@ -121,15 +116,14 @@ export const CreateCellButton = ({
   const handleOpenChange = (isOpen: boolean) => {
     setOpen(isOpen);
     if (isOpen) {
-      setJustOpened(true);
-      // Allow interactions after a brief delay
-      setTimeout(() => setJustOpened(false), 200);
+      openedAt.current = Date.now();
     }
   };
 
   const handleFirstItemClick = (e: React.MouseEvent) => {
-    // Hack to prevent the first item from being clicked when the dropdown is opened
-    if (justOpened) {
+    // Hack to prevent the first item from being clicked when the dropdown is
+    // opened: allow interactions only after a brief delay
+    if (Date.now() - openedAt.current < 200) {
       e.preventDefault();
       e.stopPropagation();
       return;
@@ -172,7 +166,7 @@ export const CreateCellButton = ({
           {renderIcon(<DatabaseIcon size={13} strokeWidth={1.5} />)}
           SQL cell
         </DropdownMenuItem>
-        <DropdownMenuItem onClick={addSetupCell}>
+        <DropdownMenuItem onClick={() => addSetupCell()}>
           {renderIcon(<DiamondPlusIcon size={13} strokeWidth={1.5} />)}
           Setup cell
         </DropdownMenuItem>

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

@@ -5,8 +5,8 @@ import type { RuntimeState } from "@/core/network/types";
 import {
   getConnectionTooltip,
   isAppInteractionDisabled,
-} from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+} from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
 import { Button } from "../../ui/button";
@@ -37,6 +37,7 @@ export const DeleteButton = (props: {
   return (
     <Tooltip content={tooltipMsg} usePortal={false}>
       <Button
+        aria-label="Delete cell"
         variant="ghost"
         size="icon"
         onClick={onClick}

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

@@ -1,12 +1,13 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { HardDriveDownloadIcon, PlayIcon } from "lucide-react";
 import type { JSX } from "react";
-import type { CellConfig, RuntimeState } from "@/core/network/types";
+import type { CellRunStatus } from "@/core/executions";
+import type { CellConfig } from "@/core/network/types";
 import {
   getConnectionTooltip,
   isAppInteractionDisabled,
-} from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+} from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { renderShortcut } from "../../shortcuts/renderShortcut";
 import { ToolbarItem } from "./toolbar";
 
@@ -35,20 +36,29 @@ function computeColor({
 
 export const RunButton = (props: {
   edited: boolean;
-  status: RuntimeState;
+  status: CellRunStatus;
+  /** Whether a cell it depends on is disabled. */
+  blocked: boolean;
   needsRun: boolean;
   connectionState: ConnectionState;
   config: CellConfig;
   onClick?: () => void;
 }): JSX.Element => {
-  const { onClick, connectionState, needsRun, status, config, edited } = props;
+  const {
+    onClick,
+    connectionState,
+    needsRun,
+    status,
+    blocked,
+    config,
+    edited,
+  } = props;
 
-  const blockedStatus = status === "disabled-transitively";
   const loading = status === "running" || status === "queued";
   const inactive =
     isAppInteractionDisabled(connectionState) ||
     loading ||
-    (!config.disabled && blockedStatus && !edited);
+    (!config.disabled && blocked && !edited);
   const variant = computeColor({
     connectionState,
     needsRun,
@@ -69,7 +79,7 @@ export const RunButton = (props: {
       </ToolbarItem>
     );
   }
-  if (!config.disabled && blockedStatus && !edited) {
+  if (!config.disabled && blocked && !edited) {
     return (
       <ToolbarItem
         disabled={inactive}