@marimo-team/islands 0.25.3-dev26 → 0.25.3-dev262

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

@@ -12,7 +12,7 @@ const hasBlockingStdinAtom = atom((get) => {
   const notebook = get(notebookAtom);
 
   // 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") {

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/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

@@ -5,8 +5,8 @@ import type { CellConfig, 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 { renderShortcut } from "../../shortcuts/renderShortcut";
 import { ToolbarItem } from "./toolbar";
 

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

@@ -5,13 +5,13 @@ import { SparklesIcon } from "lucide-react";
 import { Tooltip } from "@/components/ui/tooltip";
 import {
   type Edit,
-  stagedGenerationInProgressAtom,
   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 { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
+import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
 import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
 import {
@@ -49,8 +49,8 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
   cellId,
 }) => {
   const store = useStore();
-  const generationInProgress = useAtomValue(stagedGenerationInProgressAtom);
   const stagedAiCell = useStagedAICell(cellId);
+  const isGenerating = useAtomValue(chatGenerationInProgressAtom);
   const runCell = useRunCell(cellId);
 
   const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
@@ -93,7 +93,7 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
       </div>
       <div className="flex items-center gap-1.5">
         <CompletionActionsCellFooter
-          isLoading={generationInProgress}
+          isLoading={isGenerating}
           onAccept={() => handleCompletion("accept")}
           onDecline={() => handleCompletion("reject")}
           size="xs"
@@ -144,7 +144,15 @@ export function rejectStagedCell(
         break;
       }
 
-      updateEditorCodeFromPython(editorView, edit.previousCode);
+      const update = prepareCellCodeUpdate(
+        editorView.state,
+        edit.previousCode,
+        { initialize: true },
+      );
+
+      if (update) {
+        editorView.dispatch(update);
+      }
       removeStagedCell(cellId);
       break;
     }

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

@@ -3,8 +3,8 @@ import { SquareIcon } from "lucide-react";
 import type { JSX } from "react";
 import { useRequestClient } from "@/core/network/requests";
 import type { RuntimeState } from "@/core/network/types";
-import { isAppInteractionDisabled } from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+import { isAppInteractionDisabled } from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { Functions } from "@/utils/functions";
 import { renderShortcut } from "../../shortcuts/renderShortcut";
 import { ToolbarItem } from "./toolbar";

package/src/components/editor/cell/__tests__/useShouldShowInterrupt.test.ts

@@ -0,0 +1,39 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { useShouldShowInterrupt } from "../useShouldShowInterrupt";
+
+describe("useShouldShowInterrupt", () => {
+  beforeEach(() => vi.useFakeTimers());
+  afterEach(() => vi.useRealTimers());
+
+  it("waits for each run and hides immediately when it stops", () => {
+    const { result, rerender } = renderHook(useShouldShowInterrupt, {
+      initialProps: false,
+    });
+    expect(result.current).toBe(false);
+    rerender(true);
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(200));
+    expect(result.current).toBe(true);
+    rerender(false);
+    expect(result.current).toBe(false);
+    rerender(true);
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(199));
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(1));
+    expect(result.current).toBe(true);
+  });
+
+  it("cancels the delay for a short run", () => {
+    const { result, rerender } = renderHook(useShouldShowInterrupt, {
+      initialProps: true,
+    });
+    act(() => vi.advanceTimersByTime(100));
+    rerender(false);
+    act(() => vi.advanceTimersByTime(200));
+    expect(result.current).toBe(false);
+  });
+});

package/src/components/editor/cell/cell-actions.tsx

@@ -60,6 +60,7 @@ const CellActionsDropdownInternal = (
   ref: React.Ref<CellActionsDropdownHandle>,
 ) => {
   const [open, setOpen] = useState(false);
+  const pendingAction = useRef<(() => void) | null>(null);
   const closePopover = () => setOpen(false);
   const buttonRef = useRef<HTMLButtonElement>(null);
   const actions = useCellActionButtons({ cell: props, closePopover });
@@ -107,6 +108,16 @@ const CellActionsDropdownInternal = (
       className="w-[300px] p-0 pt-1 overflow-auto"
       scrollable={true}
       {...restoreFocus}
+      onCloseAutoFocus={(event) => {
+        restoreFocus.onCloseAutoFocus(event);
+        const action = pendingAction.current;
+        pendingAction.current = null;
+        if (action) {
+          // Restore focus before the action moves it or hides the editor.
+          event.preventDefault();
+          action();
+        }
+      }}
     >
       <Command>
         <CommandInput
@@ -153,7 +164,7 @@ const CellActionsDropdownInternal = (
                         if (action.disableClick || action.disabled) {
                           return;
                         }
-                        action.handle();
+                        pendingAction.current = action.handle;
                         setOpen(false);
                       }}
                       variant={action.variant}

package/src/components/editor/cell/code/__tests__/language-toggle.test.tsx

@@ -0,0 +1,54 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { fireEvent, render, screen } from "@testing-library/react";
+import { createRef } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { switchLanguage } from "@/core/codemirror/language/extension";
+import { LanguageToggle } from "../language-toggle";
+
+vi.mock("@/core/codemirror/language/extension", () => ({
+  switchLanguage: vi.fn(),
+}));
+
+describe("LanguageToggle", () => {
+  afterEach(() => {
+    vi.mocked(switchLanguage).mockClear();
+  });
+
+  it("uses the editor mounted after the button rendered", () => {
+    const editorViewRef = createRef<EditorView>();
+    const onAfterToggle = vi.fn();
+    render(
+      <TooltipProvider>
+        <LanguageToggle
+          editorViewRef={editorViewRef}
+          currentLanguageAdapter="python"
+          canSwitchToLanguage={true}
+          toType="markdown"
+          displayName="Markdown"
+          icon="Markdown"
+          onAfterToggle={onAfterToggle}
+        />
+      </TooltipProvider>,
+    );
+
+    fireEvent.click(screen.getByRole("button", { name: "Markdown" }));
+    expect(switchLanguage).not.toHaveBeenCalled();
+    expect(onAfterToggle).not.toHaveBeenCalled();
+
+    const view = new EditorView({ state: EditorState.create() });
+    editorViewRef.current = view;
+    try {
+      fireEvent.click(screen.getByRole("button", { name: "Markdown" }));
+      expect(switchLanguage).toHaveBeenCalledWith(view, {
+        language: "markdown",
+      });
+      expect(onAfterToggle).toHaveBeenCalledOnce();
+    } finally {
+      view.destroy();
+    }
+  });
+});