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

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/cell/StagedAICell.tsx

@@ -5,10 +5,10 @@ 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 { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
@@ -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"

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

@@ -31,7 +31,11 @@ import {
 } from "@/core/codemirror/language/extension";
 import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages/markdown";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
-import { autoInstantiateAtom, isAiFeatureEnabled } from "@/core/config/config";
+import {
+  autoInstantiateAtom,
+  isAiEnabled,
+  isAiFeatureEnabled,
+} from "@/core/config/config";
 import type { UserConfig } from "@/core/config/config-schema";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { useSaveNotebook } from "@/core/saving/save-component";
@@ -39,6 +43,7 @@ import type { Theme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { invariant } from "@/utils/invariant";
 import { mergeRefs } from "@/utils/mergeRefs";
+import { useAskAiAboutCell } from "../../chrome/wrapper/useAskAiAboutCell";
 import { AiCompletionEditor } from "../../ai/ai-completion-editor";
 import {
   closeSignatureHelp,
@@ -189,12 +194,15 @@ const CellEditorInternal = ({
   });
 
   const aiFeaturesEnabled = isAiFeatureEnabled(userConfig);
+  const askAiAboutCell = useAskAiAboutCell();
+  const aiChatEnabled = isAiEnabled(userConfig);
 
   const extensions = useMemo(() => {
     const extensions = setupCodeMirror({
       cellId,
       showPlaceholder,
       enableAI: aiFeaturesEnabled,
+      enableAIChat: aiChatEnabled,
       cellActions: {
         ...cellActions,
         afterToggleMarkdown,
@@ -213,6 +221,9 @@ const CellEditorInternal = ({
         },
         splitCell,
         toggleHideCode,
+        askAiAboutCell: () => {
+          void askAiAboutCell(cellId);
+        },
         aiCellCompletion: () => {
           if (!aiFeaturesEnabled) {
             return false;
@@ -296,6 +307,8 @@ const CellEditorInternal = ({
     inlineAiTooltip,
     tooltipParentSelector,
     aiFeaturesEnabled,
+    aiChatEnabled,
+    askAiAboutCell,
     theme,
     showPlaceholder,
     cellActions,

package/src/components/editor/chrome/panels/panel-header.tsx

@@ -0,0 +1,13 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { createContext, useContext, type PropsWithChildren } from "react";
+import { createPortal } from "react-dom";
+
+export const PanelHeaderActionsContext = createContext<HTMLElement | null>(
+  null,
+);
+
+export function PanelHeaderActions({ children }: PropsWithChildren) {
+  const target = useContext(PanelHeaderActionsContext);
+  return target ? createPortal(children, target) : children;
+}

package/src/components/editor/chrome/types.ts

@@ -97,7 +97,7 @@ export const PANELS: PanelDescriptor[] = [
     type: "ai",
     Icon: BotIcon,
     label: "AI",
-    tooltip: "Chat & Agents",
+    tooltip: "Chat",
     defaultSection: "sidebar",
     additionalKeywords: ["chat", "copilot", "assistant"],
   },

package/src/components/editor/chrome/wrapper/__tests__/pending-ai-cells.test.tsx

@@ -1,12 +1,15 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { act, fireEvent, render, screen, within } from "@testing-library/react";
+import { fireEvent, render, screen } from "@testing-library/react";
 import { createStore, Provider } from "jotai";
 import { describe, expect, it, vi } from "vitest";
+import type { ChatStatus } from "ai";
+import { useChatGeneration } from "@/components/chat/use-chat-generation";
 import { cellId } from "@/__tests__/branded";
 import { TooltipProvider } from "@/components/ui/tooltip";
-import { stagedAICellsAtom, visibleForTesting } from "@/core/ai/staged-cells";
+import { stagedAICellsAtom } from "@/core/ai/staged-cells";
 import { scrollAndHighlightCell } from "../../../links/cell-link";
+import { StagedAICellFooter } from "../../../cell/StagedAICell";
 import { PendingAICells } from "../pending-ai-cells";
 
 vi.mock("@/components/editor/cell/useDeleteCell", () => ({
@@ -25,6 +28,7 @@ vi.mock("@/core/cells/actions", () => ({
 
 vi.mock("@/components/editor/cell/useRunCells", () => ({
   useRunCells: () => vi.fn(),
+  useRunCell: () => vi.fn(),
 }));
 
 vi.mock("@/components/editor/links/cell-link", () => ({
@@ -87,50 +91,54 @@ describe("PendingAICells", () => {
     expect(scrollAndHighlightCell).toHaveBeenLastCalledWith(third, "focus");
     expect(screen.getByText("3 / 3")).toBeInTheDocument();
   });
+});
 
-  it("disables bulk actions while cells are still generating", () => {
-    const store = createStore();
-    const generationId = Symbol("staged-cell-generation");
-    store.set(
-      stagedAICellsAtom,
-      new Map([[cellId("generated-cell"), { type: "add_cell" }]]),
-    );
-    store.set(visibleForTesting.stagedGenerationAtom, {
-      id: generationId,
-      status: "in_progress",
-      cellIds: [cellId("generated-cell")],
-    });
-
-    const view = (
-      <Provider store={store}>
-        <TooltipProvider>
-          <PendingAICells />
-        </TooltipProvider>
-      </Provider>
-    );
-    const { rerender } = render(view);
-
-    const acceptButton = screen.getByRole("button", { name: "Keep all" });
-    const acceptButtons = within(
-      acceptButton.parentElement as HTMLElement,
-    ).getAllByRole("button");
-    expect(acceptButtons).toHaveLength(2);
-    for (const button of acceptButtons) {
-      expect(button).toBeDisabled();
+function ChatGeneration({ status }: { status: ChatStatus }) {
+  useChatGeneration(status);
+  return null;
+}
+
+it("locks staged edits during generation and unlocks after completion, stop, error, or unmount", () => {
+  const store = createStore();
+  const pendingCell = cellId("pending-cell");
+  store.set(stagedAICellsAtom, new Map([[pendingCell, { type: "add_cell" }]]));
+  const ui = (status?: ChatStatus) => (
+    <Provider store={store}>
+      <TooltipProvider>
+        {status && <ChatGeneration status={status} />}
+        <PendingAICells />
+        <StagedAICellFooter cellId={pendingCell} />
+      </TooltipProvider>
+    </Provider>
+  );
+  const { rerender } = render(ui("ready"));
+  const keep = screen.getByRole("button", { name: "Keep all" });
+  const discard = screen.getByRole("button", { name: "Discard all" });
+  const keepCell = screen.getByRole("button", { name: "Keep cell" });
+  const discardCell = screen.getByRole("button", { name: "Discard cell" });
+  const keepAndRun = keep.parentElement!.querySelectorAll("button")[1];
+  for (const status of [
+    "submitted",
+    "streaming",
+    "ready",
+    "submitted",
+    "ready",
+    "submitted",
+    "error",
+    "streaming",
+    undefined,
+  ] as const) {
+    rerender(ui(status));
+    const isGenerating = status === "submitted" || status === "streaming";
+    for (const button of [keep, discard, keepAndRun, keepCell, discardCell]) {
+      expect(button.hasAttribute("disabled")).toBe(isGenerating);
+      if (isGenerating) {
+        fireEvent.click(button);
+      }
     }
-    expect(screen.getByRole("button", { name: "Discard all" })).toBeDisabled();
+    expect(store.get(stagedAICellsAtom).has(pendingCell)).toBe(true);
+  }
 
-    act(() => {
-      store.set(visibleForTesting.stagedGenerationAtom, {
-        id: generationId,
-        status: "complete",
-      });
-    });
-    rerender(view);
-
-    for (const button of acceptButtons) {
-      expect(button).toBeEnabled();
-    }
-    expect(screen.getByRole("button", { name: "Discard all" })).toBeEnabled();
-  });
+  fireEvent.click(keep);
+  expect(store.get(stagedAICellsAtom).size).toBe(0);
 });

package/src/components/editor/chrome/wrapper/__tests__/useAskAiAboutCell.test.tsx

@@ -0,0 +1,58 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import type { PropsWithChildren } from "react";
+import { beforeEach, expect, it, vi } from "vitest";
+import { pendingAiPromptAtom } from "@/core/ai/state";
+import { CellId } from "@/core/cells/ids";
+import { aiPanelTabAtom } from "../useAiPanel";
+import { useAskAiAboutCell } from "../useAskAiAboutCell";
+
+const { aiAtom, openApplication, saveModeChange } = await vi.hoisted(
+  async () => {
+    const { atom } = await import("jotai");
+    return {
+      aiAtom: atom<{ mode?: "ask" | "agent" | "code_mode" }>({}),
+      openApplication: vi.fn(),
+      saveModeChange: vi.fn().mockResolvedValue(undefined),
+    };
+  },
+);
+vi.mock("@/core/config/config", async (importOriginal) => ({
+  ...(await importOriginal<object>()),
+  aiAtom,
+}));
+vi.mock("@/core/ai/config", () => ({
+  useAIConfigActions: () => ({ saveModeChange }),
+}));
+vi.mock("../../state", () => ({
+  useChromeActions: () => ({ openApplication }),
+}));
+
+beforeEach(() => vi.clearAllMocks());
+it.each([undefined, "ask", "agent", "code_mode"] as const)(
+  "opens chat from %s mode without submitting or replacing the draft",
+  async (mode) => {
+    const store = createStore();
+    store.set(aiAtom, { mode });
+    store.set(aiPanelTabAtom, "agents");
+    const cellId = CellId.create();
+    const { result } = renderHook(useAskAiAboutCell, {
+      wrapper: ({ children }: PropsWithChildren) => (
+        <Provider store={store}>{children}</Provider>
+      ),
+    });
+    await act(() => result.current(cellId));
+    expect(openApplication).toHaveBeenCalledWith("ai");
+    expect(store.get(aiPanelTabAtom)).toBe("chat");
+    expect(saveModeChange.mock.calls).toEqual(
+      mode === "agent" || mode === "code_mode" ? [] : [["agent"]],
+    );
+    expect(store.get(pendingAiPromptAtom)).toEqual({
+      prompt: `@cell://${cellId} `,
+      submit: false,
+      append: true,
+    });
+  },
+);