@marimo-team/islands 0.25.3-dev28 → 0.25.3-dev281

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -1,506 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { useChat } from "@ai-sdk/react";
-import {
-  autocompletion,
-  type Completion,
-  type CompletionContext,
-  type CompletionSource,
-} from "@codemirror/autocomplete";
-import { markdown } from "@codemirror/lang-markdown";
-import { Prec } from "@codemirror/state";
-import { promptHistory, storePrompt } from "@marimo-team/codemirror-ai";
-import ReactCodeMirror, {
-  EditorView,
-  keymap,
-  minimalSetup,
-  type ReactCodeMirrorRef,
-} from "@uiw/react-codemirror";
-import { useAtom, useAtomValue, useStore } from "jotai";
-import { atomWithStorage } from "jotai/utils";
-import {
-  ChevronsUpDown,
-  DatabaseIcon,
-  SparklesIcon,
-  XIcon,
-} from "lucide-react";
-import { useEffect, useMemo, useRef, useState } from "react";
-import useEvent from "react-use-event-hook";
-import { DefaultChatTransport } from "ai";
-import { z } from "zod";
-import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
-import { PairWithAgentButton } from "@/components/editor/actions/pair-with-agent";
-import {
-  AddContextButton,
-  AttachFileButton,
-  FileAttachmentPill,
-  SendButton,
-} from "@/components/chat/chat-components";
-import {
-  buildCompletionRequestBody,
-  convertToFileUIPart,
-  PROVIDERS_THAT_SUPPORT_ATTACHMENTS,
-  useFileState,
-} from "@/components/chat/chat-utils";
-import { Button } from "@/components/ui/button";
-import {
-  DropdownMenu,
-  DropdownMenuContent,
-  DropdownMenuItem,
-  DropdownMenuSeparator,
-  DropdownMenuTrigger,
-} from "@/components/ui/dropdown-menu";
-import { toast } from "@/components/ui/use-toast";
-import { AiModelId } from "@/core/ai/ids/ids";
-import { AI_SDK_UI_THROTTLE_MS } from "@/core/ai/constants";
-import type { CompletionUIMessage } from "@/core/ai/completion-output";
-import { useStagedCellGeneration } from "@/core/ai/staged-cells";
-import { resourceExtension } from "@/core/codemirror/ai/resources";
-import { aiAtom } from "@/core/config/config";
-import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
-import type { AiCompletionRequest } from "@/core/network/types";
-import { useRuntimeManager } from "@/core/runtime/config";
-import { useTheme } from "@/theme/useTheme";
-import { cn } from "@/utils/cn";
-import { prettyError } from "@/utils/errors";
-import { jotaiJsonStorage } from "@/utils/storage/jotai";
-import { ZodLocalStorage } from "@/utils/storage/typed";
-import { PythonIcon } from "../cell/code/icons";
-import { createAiCompletionOnKeydown } from "./completion-handlers";
-import {
-  addContextCompletion,
-  CONTEXT_TRIGGER,
-  mentionsCompletionSource,
-} from "./completion-utils";
-import { StagedCellSubmissionController } from "./staged-cell-submission";
-
-// Persist across sessions
-const languageAtom = atomWithStorage<"python" | "sql">(
-  "marimo:ai-language",
-  "python",
-  jotaiJsonStorage,
-);
-
-const KEY = "marimo:ai-prompt-history";
-// Store the prompt history in local storage
-const promptHistoryStorage = new ZodLocalStorage(z.array(z.string()), () => []);
-
-/**
- * Add a cell with AI.
- */
-export const AddCellWithAI: React.FC<{
-  onClose: () => void;
-}> = ({ onClose }) => {
-  const store = useStore();
-  const [input, setInput] = useState("");
-
-  const {
-    acceptOwnedStagedCells,
-    beginStagedCellGeneration,
-    discardOwnedStagedCells,
-    finishStagedCellGeneration,
-    hasOwnedStagedCells,
-    onData,
-  } = useStagedCellGeneration(store);
-  const [language, setLanguage] = useAtom(languageAtom);
-  const runtimeManager = useRuntimeManager();
-
-  const inputRef = useRef<ReactCodeMirrorRef>(null);
-  const submissionController = useRef(
-    new StagedCellSubmissionController(),
-  ).current;
-
-  const fileInputRef = useRef<HTMLInputElement>(null);
-  const { files, addFiles, removeFile } = useFileState();
-  const aiConfig = useAtomValue(aiAtom);
-
-  const transport = useMemo(
-    () =>
-      new DefaultChatTransport<CompletionUIMessage>({
-        api: runtimeManager.getAiURL("completion").toString(),
-        headers: () => runtimeManager.headers(),
-        prepareSendMessagesRequest: async (options) => {
-          const completionBody = await buildCompletionRequestBody(
-            options.messages,
-          );
-          const body = {
-            ...completionBody,
-            id: options.id,
-            code: "",
-            prompt: "", // Don't need prompt since we are using messages
-            language: language,
-          } satisfies AiCompletionRequest;
-
-          return {
-            api: runtimeManager.getAiURL("completion").toString(),
-            body: body,
-          };
-        },
-      }),
-    [language, runtimeManager],
-  );
-
-  const handleGenerationError = useEvent((error: unknown) => {
-    finishStagedCellGeneration(false);
-    toast({
-      title: "Generate with AI failed",
-      description: prettyError(error),
-    });
-  });
-
-  const {
-    sendMessage,
-    stop: stopChat,
-    status,
-  } = useChat<CompletionUIMessage>({
-    throttle: AI_SDK_UI_THROTTLE_MS,
-    transport,
-    onData,
-    onError: handleGenerationError,
-    onFinish: ({ isAbort, isDisconnect, isError, finishReason }) => {
-      finishStagedCellGeneration(
-        !isAbort && !isDisconnect && !isError && finishReason === "stop",
-      );
-    },
-  });
-
-  const isLoading = status === "streaming" || status === "submitted";
-  const hasCompletion = hasOwnedStagedCells();
-  const stop = useEvent(() => {
-    submissionController.cancel();
-    finishStagedCellGeneration(false);
-    void stopChat();
-  });
-
-  const reject = useEvent(() => {
-    submissionController.cancel();
-    discardOwnedStagedCells();
-    void stopChat();
-  });
-
-  // Parent-driven unmounts bypass the close handlers, so cancel in-flight work here.
-  useEffect(() => {
-    return () => {
-      stop();
-    };
-  }, [stop]);
-
-  const currentModel = aiConfig?.models?.edit_model || DEFAULT_AI_MODEL;
-  const currentProvider = AiModelId.parse(currentModel).providerId;
-  const isAttachmentSupported =
-    PROVIDERS_THAT_SUPPORT_ATTACHMENTS.has(currentProvider);
-
-  const submit = async () => {
-    if (!isLoading) {
-      await submissionController.run({
-        prepare: async () => {
-          if (inputRef.current?.view) {
-            storePrompt(inputRef.current.view);
-          }
-          return files ? await convertToFileUIPart(files) : undefined;
-        },
-        submit: async (fileParts) => {
-          beginStagedCellGeneration();
-          await sendMessage({ text: input, files: fileParts });
-        },
-        onError: handleGenerationError,
-      });
-    }
-  };
-
-  const pythonIcon = (
-    <>
-      <PythonIcon className="size-4 mr-2" />
-      Python
-    </>
-  );
-
-  const sqlIcon = (
-    <>
-      <DatabaseIcon className="size-4 mr-2" />
-      SQL
-    </>
-  );
-
-  const languageDropdown = (
-    <DropdownMenu modal={false}>
-      <DropdownMenuTrigger
-        className="flex items-center justify-between h-7 text-xs px-2 py-0.5 border rounded-md hover:text-accent-foreground"
-        data-testid="language-button"
-      >
-        {language === "python" ? pythonIcon : sqlIcon}
-        <ChevronsUpDown className="ml-1 h-3.5 w-3.5 text-muted-foreground/70" />
-      </DropdownMenuTrigger>
-      <DropdownMenuContent align="center">
-        <div className="px-2 py-1 text-sm text-muted-foreground">
-          Select language
-        </div>
-        <DropdownMenuSeparator />
-        <DropdownMenuItem onClick={() => setLanguage("python")}>
-          {pythonIcon}
-        </DropdownMenuItem>
-        <DropdownMenuItem onClick={() => setLanguage("sql")}>
-          {sqlIcon}
-        </DropdownMenuItem>
-      </DropdownMenuContent>
-    </DropdownMenu>
-  );
-
-  const handleAcceptCompletion = () => {
-    if (acceptOwnedStagedCells()) {
-      onClose();
-    }
-  };
-
-  const handleDeclineCompletion = () => {
-    reject();
-    // Focus the input so the user can refine the prompt.
-    inputRef.current?.view?.focus();
-  };
-
-  const handleClose = () => {
-    reject();
-    onClose();
-  };
-
-  const inputComponent = (
-    <div className="flex items-center px-3">
-      <SparklesIcon className="size-4 text-(--blue-11) mr-2" />
-      <PromptInput
-        inputRef={inputRef}
-        onClose={handleClose}
-        value={input}
-        onChange={(newValue) => {
-          setInput(newValue);
-        }}
-        onSubmit={submit}
-        onKeyDown={createAiCompletionOnKeydown({
-          handleAcceptCompletion,
-          handleDeclineCompletion,
-          isLoading,
-          hasCompletion,
-        })}
-      />
-      <Button
-        variant="text"
-        size="sm"
-        className="mb-0 px-1"
-        onClick={handleClose}
-      >
-        <XIcon className="size-4" />
-      </Button>
-    </div>
-  );
-
-  const footerComponent = (
-    <div className="px-3 pt-1 flex flex-row items-center justify-between">
-      <div className="flex items-center gap-2">
-        <AIModelDropdown
-          triggerClassName="h-7 text-xs max-w-64"
-          iconSize="small"
-          forRole="edit"
-          showAddCustomModelDocs={true}
-        />
-        {languageDropdown}
-        <PairWithAgentButton />
-      </div>
-      <div className="flex flex-row items-center">
-        {files.length > 0 && (
-          <div className="flex flex-row gap-1 flex-wrap pr-1">
-            {files?.map((file, index) => (
-              <FileAttachmentPill
-                file={file}
-                key={`${file.name}-${index}`}
-                onRemove={() => removeFile(file)}
-              />
-            ))}
-          </div>
-        )}
-        <AddContextButton
-          handleAddContext={() => addContextCompletion(inputRef)}
-          isLoading={isLoading}
-        />
-        {isAttachmentSupported && (
-          <AttachFileButton
-            fileInputRef={fileInputRef}
-            isLoading={isLoading}
-            onAddFiles={addFiles}
-          />
-        )}
-        <SendButton
-          isLoading={isLoading}
-          onStop={stop}
-          onSendClick={submit}
-          isEmpty={!input.trim()}
-          showStopLabel={true}
-        />
-      </div>
-    </div>
-  );
-
-  return (
-    <div className={cn("flex flex-col w-full py-2")}>
-      {inputComponent}
-      {footerComponent}
-    </div>
-  );
-};
-
-export interface AdditionalCompletions {
-  triggerCompletionRegex: RegExp;
-  completions: Completion[];
-}
-
-interface PromptInputProps {
-  inputRef?: React.RefObject<ReactCodeMirrorRef | null>;
-  placeholder?: string;
-  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,
-  inputRef,
-  className,
-  onChange,
-  onSubmit,
-  onKeyDown,
-  onClose,
-  onAddFiles,
-  additionalCompletions,
-  maxHeight,
-}: PromptInputProps) => {
-  const handleSubmit = onSubmit;
-  const handleEscape = 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 [
-      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) {
-                const cursorPosition = view.state.selection.main.from;
-                // Insert a new line
-                view.dispatch({
-                  changes: {
-                    from: cursorPosition,
-                    to: cursorPosition,
-                    insert: "\n",
-                  },
-                  selection: {
-                    anchor: cursorPosition + 1,
-                    head: cursorPosition + 1,
-                  },
-                });
-                event.preventDefault();
-                event.stopPropagation();
-                return true;
-              }
-
-              return false;
-            },
-          },
-        ]),
-      ),
-      keymap.of([
-        {
-          key: "Escape",
-          preventDefault: true,
-          stopPropagation: true,
-          run: () => {
-            handleEscape();
-            return true;
-          },
-        },
-      ]),
-    ];
-  }, [
-    store,
-    onAddFiles,
-    additionalCompletionsSource,
-    handleSubmit,
-    handleEscape,
-  ]);
-
-  return (
-    <ReactCodeMirror
-      ref={inputRef}
-      className={cn("flex-1 font-sans 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={
-        placeholder || `Generate with AI, ${CONTEXT_TRIGGER} to include context`
-      }
-    />
-  );
-};

package/src/components/editor/ai/staged-cell-submission.ts

@@ -1,46 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-interface SubmissionCallbacks<T> {
-  prepare: () => Promise<T>;
-  submit: (prepared: T) => Promise<void>;
-  onError: (error: unknown) => void;
-}
-
-interface SubmissionAttempt {
-  cancelled: boolean;
-}
-
-/** Coordinates preprocessing and submission without resuming cancelled work. */
-export class StagedCellSubmissionController {
-  private activeAttempt: SubmissionAttempt | null = null;
-
-  public cancel() {
-    if (this.activeAttempt) {
-      this.activeAttempt.cancelled = true;
-    }
-  }
-
-  public async run<T>({ prepare, submit, onError }: SubmissionCallbacks<T>) {
-    if (this.activeAttempt !== null) {
-      return;
-    }
-
-    const attempt: SubmissionAttempt = { cancelled: false };
-    this.activeAttempt = attempt;
-
-    try {
-      const prepared = await prepare();
-      if (attempt.cancelled) {
-        return;
-      }
-
-      await submit(prepared);
-    } catch (error) {
-      if (!attempt.cancelled) {
-        onError(error);
-      }
-    } finally {
-      this.activeAttempt = null;
-    }
-  }
-}

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

@@ -1,64 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { render, screen } from "@testing-library/react";
-import { createStore, Provider } from "jotai";
-import { describe, expect, it, vi } from "vitest";
-import { cellId } from "@/__tests__/branded";
-import { TooltipProvider } from "@/components/ui/tooltip";
-import { stagedAICellsAtom, visibleForTesting } from "@/core/ai/staged-cells";
-import { StagedAICellFooter } from "../StagedAICell";
-
-vi.mock("@/components/editor/cell/useDeleteCell", () => ({
-  useDeleteCellCallback: () => vi.fn(),
-}));
-
-vi.mock("@/core/cells/cells", () => ({
-  getCellEditorView: vi.fn(),
-  useCellActions: () => ({
-    createNewCell: vi.fn(),
-    updateCellCode: vi.fn(),
-  }),
-}));
-
-vi.mock("../useRunCells", () => ({
-  useRunCell: () => undefined,
-}));
-
-describe("StagedAICellFooter", () => {
-  it("disables cell actions until generation is complete", () => {
-    const store = createStore();
-    const generatedCellId = cellId("generated-cell");
-    store.set(
-      stagedAICellsAtom,
-      new Map([[generatedCellId, { type: "add_cell" }]]),
-    );
-    store.set(visibleForTesting.stagedGenerationAtom, {
-      id: Symbol("staged-cell-generation"),
-      status: "in_progress",
-      cellIds: [generatedCellId],
-    });
-
-    const { rerender } = render(
-      <Provider store={store}>
-        <TooltipProvider>
-          <StagedAICellFooter cellId={generatedCellId} />
-        </TooltipProvider>
-      </Provider>,
-    );
-
-    expect(screen.getByRole("button", { name: "Keep cell" })).toBeDisabled();
-    expect(screen.getByRole("button", { name: "Discard cell" })).toBeDisabled();
-
-    store.set(visibleForTesting.stagedGenerationAtom, null);
-    rerender(
-      <Provider store={store}>
-        <TooltipProvider>
-          <StagedAICellFooter cellId={generatedCellId} />
-        </TooltipProvider>
-      </Provider>,
-    );
-
-    expect(screen.getByRole("button", { name: "Keep cell" })).toBeEnabled();
-    expect(screen.getByRole("button", { name: "Discard cell" })).toBeEnabled();
-  });
-});

package/src/components/editor/cell/code/__tests__/rtc-editor-config.test.ts

@@ -1,50 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { EditorState } from "@codemirror/state";
-import { describe, expect, it } from "vitest";
-import { composeRtcEditorConfig } from "../cell-editor";
-
-describe("composeRtcEditorConfig", () => {
-  it("does not accumulate RTC extensions across reconfigurations", () => {
-    const baseExtension = EditorState.readOnly.of(false);
-    const baseExtensions = [baseExtension];
-    const firstRtc = {
-      code: "first RTC code",
-      extension: EditorState.tabSize.of(2),
-    };
-    const secondRtc = {
-      code: "second RTC code",
-      extension: EditorState.tabSize.of(4),
-    };
-
-    const firstConfig = composeRtcEditorConfig(
-      "local code",
-      baseExtensions,
-      firstRtc,
-    );
-    const secondConfig = composeRtcEditorConfig(
-      "local code",
-      baseExtensions,
-      secondRtc,
-    );
-
-    expect(baseExtensions).toEqual([baseExtension]);
-    expect(firstConfig).toEqual({
-      code: firstRtc.code,
-      extensions: [baseExtension, firstRtc.extension],
-    });
-    expect(secondConfig).toEqual({
-      code: secondRtc.code,
-      extensions: [baseExtension, secondRtc.extension],
-    });
-  });
-
-  it("reuses the base configuration when RTC is disabled", () => {
-    const extensions = [EditorState.readOnly.of(false)];
-
-    const config = composeRtcEditorConfig("local code", extensions, undefined);
-
-    expect(config).toEqual({ code: "local code", extensions });
-    expect(config.extensions).toBe(extensions);
-  });
-});