@marimo-team/frontend 0.24.1-dev58 → 0.24.1-dev59

package/src/components/editor/ai/ai-completion-editor.tsx

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useCompletion } from "@ai-sdk/react";
+import { useChat } from "@ai-sdk/react";
 import { EditorView } from "@codemirror/view";
 import {
   CircleCheckIcon,
@@ -10,6 +10,7 @@ import {
 } from "lucide-react";
 import React, { useCallback, useEffect, useId, useState } from "react";
 import CodeMirrorMerge from "react-codemirror-merge";
+import { DefaultChatTransport } from "ai";
 import { Button } from "@/components/ui/button";
 import { customPythonLanguageSupport } from "@/core/codemirror/language/languages/python";
 
@@ -27,14 +28,19 @@ import { Label } from "@/components/ui/label";
 import { Switch } from "@/components/ui/switch";
 import { Tooltip } from "@/components/ui/tooltip";
 import { toast } from "@/components/ui/use-toast";
+import {
+  CELL_COMPLETION_DATA_TYPE,
+  cellCompletionSchema,
+  type CompletionUIMessage,
+} from "@/core/ai/completion-output";
 import { AI_SDK_UI_THROTTLE_MS } from "@/core/ai/constants";
-import { stripWrappingBackticks } from "@/core/ai/strip-wrapping-backticks";
 import { type AiCompletionCell, includeOtherCellsAtom } from "@/core/ai/state";
 import type { CellId } from "@/core/cells/ids";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
 import { selectAllText } from "@/core/codemirror/utils";
 import { useRuntimeManager } from "@/core/runtime/config";
+import type { AiCompletionRequest } from "@/core/network/types";
 import { useTheme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { prettyError } from "@/utils/errors";
@@ -90,7 +96,6 @@ export const AiCompletionEditor: React.FC<Props> = ({
   children,
 }) => {
   const [showInputPrompt, setShowInputPrompt] = useState(false);
-  const [completionBody, setCompletionBody] = useState<object>({});
 
   const [includeOtherCells, setIncludeOtherCells] = useAtom(
     includeOtherCellsAtom,
@@ -112,50 +117,84 @@ export const AiCompletionEditor: React.FC<Props> = ({
     previousCellCode = updatedCell.previousCode;
   }
 
+  const [completion, setCompletion] = useState("");
+  const [input, setInput] = useState(initialPrompt ?? "");
+  const transport = React.useMemo(
+    () =>
+      new DefaultChatTransport<CompletionUIMessage>({
+        api: runtimeManager.getAiURL("completion").toString(),
+        headers: () => runtimeManager.headers(),
+        prepareSendMessagesRequest: ({ body }) => ({ body: body ?? {} }),
+      }),
+    [runtimeManager],
+  );
   const {
-    completion: untrimmedCompletion,
-    input,
-    stop,
-    isLoading,
-    setCompletion,
-    setInput,
-    handleSubmit,
-    complete,
-  } = useCompletion({
-    api: runtimeManager.getAiURL("completion").toString(),
-    headers: runtimeManager.headers(),
-    initialInput: initialPrompt,
+    sendMessage,
+    stop: stopChat,
+    status,
+  } = useChat<CompletionUIMessage>({
     throttle: AI_SDK_UI_THROTTLE_MS,
-    body: {
-      ...(Object.keys(completionBody).length > 0
-        ? completionBody
-        : initialPrompt
-          ? getAICompletionBody({ input: initialPrompt })
-          : {}),
-      includeOtherCode: includeOtherCells ? getCodes(currentCode) : "",
-      code: currentCode,
-      language: currentLanguageAdapter,
+    transport,
+    onData: (part) => {
+      if (part.type === CELL_COMPLETION_DATA_TYPE) {
+        const completion = cellCompletionSchema.parse(part.data);
+        setCompletion(completion.code.trimEnd());
+      }
     },
     onError: (error) => {
+      setCompletion("");
       toast({
         title: "Completion failed",
         description: prettyError(error),
       });
     },
-    onFinish: (_prompt, completion) => {
-      setCompletion(stripWrappingBackticks(completion).trimEnd());
+    onFinish: ({ isAbort, isDisconnect, isError, finishReason }) => {
+      if (isAbort || isDisconnect || isError || finishReason !== "stop") {
+        setCompletion("");
+      }
     },
   });
+  const isLoading = status === "submitted" || status === "streaming";
+
+  const stop = useCallback(() => {
+    setCompletion("");
+    void stopChat();
+  }, [stopChat]);
+
+  const complete = useCallback(
+    async (prompt: string) => {
+      await stopChat();
+      setCompletion("");
+
+      const body = {
+        ...getAICompletionBody({ input: prompt }),
+        prompt,
+        includeOtherCode: includeOtherCells ? getCodes(currentCode) : "",
+        code: currentCode,
+        language: currentLanguageAdapter ?? "python",
+      } satisfies AiCompletionRequest;
+
+      await sendMessage({ text: prompt }, { body });
+    },
+    [
+      currentCode,
+      currentLanguageAdapter,
+      includeOtherCells,
+      sendMessage,
+      stopChat,
+    ],
+  );
+
+  const handleSubmit = useCallback(() => {
+    void complete(input);
+  }, [complete, input]);
 
   const inputRef = React.useRef<ReactCodeMirrorRef>(null);
-  const completion = stripWrappingBackticks(untrimmedCompletion, {
-    streaming: isLoading,
-  }).trimEnd();
 
   const initialSubmit = useCallback(() => {
     if (triggerImmediately && !isLoading && initialPrompt) {
       // Use complete to pass the prompt directly, else input might be empty
-      complete(initialPrompt);
+      void complete(initialPrompt);
     }
     // oxlint-disable-next-line react-hooks/exhaustive-deps
   }, [triggerImmediately]);
@@ -197,7 +236,6 @@ export const AiCompletionEditor: React.FC<Props> = ({
   // Reject discards the suggestion but keeps the prompt open for refinement.
   const handleDeclineCompletion = () => {
     stop();
-    setCompletion("");
     setShowInputPrompt(true);
     inputRef.current?.view?.focus();
   };
@@ -299,14 +337,11 @@ export const AiCompletionEditor: React.FC<Props> = ({
               inputRef={inputRef}
               className="h-full my-0 py-2 flex items-center"
               onClose={() => {
+                stop();
                 declineChange();
-                setCompletion("");
               }}
               value={input}
-              onChange={(newValue) => {
-                setInput(newValue);
-                setCompletionBody(getAICompletionBody({ input: newValue }));
-              }}
+              onChange={setInput}
               onSubmit={() => {
                 if (!isLoading) {
                   if (inputRef.current?.view) {
@@ -378,7 +413,6 @@ export const AiCompletionEditor: React.FC<Props> = ({
               onClick={() => {
                 stop();
                 declineChange();
-                setCompletion("");
               }}
             >
               <XIcon className="text-(--red-10)" size={16} />

package/src/components/editor/ai/completion-handlers.tsx

@@ -77,6 +77,7 @@ export const CompletionActionsCellFooter: React.FC<{
         onDecline={onDecline}
         size={size}
         label={declineLabel}
+        disabled={isLoading}
       />
     </>
   );
@@ -169,6 +170,7 @@ export const AcceptCompletionButton: React.FC<{
 
 export const RejectCompletionButton: React.FC<{
   onDecline: () => void;
+  disabled?: boolean;
   multipleCompletions?: boolean;
   size?: "xs" | "sm";
   className?: string;
@@ -177,6 +179,7 @@ export const RejectCompletionButton: React.FC<{
   label?: string;
 }> = ({
   onDecline,
+  disabled = false,
   multipleCompletions = false,
   size = "sm",
   className,
@@ -188,6 +191,7 @@ export const RejectCompletionButton: React.FC<{
     <Button
       variant="text"
       size={size}
+      disabled={disabled}
       onClick={onDecline}
       className={`h-6 text-(--red-10) bg-(--red-3)/60 hover:bg-(--red-3)
     dark:bg-(--red-4)/80 dark:hover:bg-(--red-3) rounded px-3 font-semibold

package/src/components/editor/ai/completion-utils.ts

@@ -11,7 +11,6 @@ import type { DataUIPart, FileUIPart, UIMessage } from "ai";
 import { getAIContextRegistry } from "@/core/ai/context/context";
 import type { ContextLocatorId } from "@/core/ai/context/registry";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
-import type { LanguageAdapterType } from "@/core/codemirror/language/types";
 import type { AiCompletionRequest } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
 import { Logger } from "@/utils/Logger";
@@ -249,94 +248,3 @@ export function addContextCompletion(
     startCompletion(inputRef.current.view);
   }
 }
-
-export interface AiCompletion {
-  language: LanguageAdapterType;
-  code: string;
-}
-
-/**
- * Extracts code blocks (delimited by triple backticks) and their language ("python", "sql", "markdown").
- * Defaults to "python" if no language is specified or no code blocks are found.
- * Returns an array of AiCompletion objects.
- */
-export function codeToCells(code: string): AiCompletion[] {
-  if (code.trim().length === 0) {
-    return [];
-  }
-
-  // If there are no backticks, assume code is in 1 cell and python
-  if (!code.includes("```")) {
-    return [{ language: "python", code: code }];
-  }
-
-  // If code has opening backticks, get the code after it
-  const cells: AiCompletion[] = [];
-  let start = 0;
-
-  let openIndex = code.indexOf("```", start);
-  while (openIndex !== -1) {
-    const newlineIndex = code.indexOf("\n", openIndex);
-    if (newlineIndex === -1) {
-      // If there's no newline after opening backticks, treat everything after as code
-      const remaining = code.slice(openIndex + 3);
-      const firstSpace = remaining.indexOf(" ");
-      const language =
-        firstSpace === -1 ? remaining : remaining.slice(0, firstSpace);
-      const finalLanguage =
-        language === "markdown"
-          ? "markdown"
-          : language === "sql"
-            ? "sql"
-            : "python";
-      // Extract code after the language identifier
-      const codeContent =
-        firstSpace === -1 ? "" : remaining.slice(firstSpace + 1);
-      if (codeContent.trim()) {
-        cells.push({ language: finalLanguage, code: codeContent.trim() });
-      }
-      break;
-    }
-
-    let language = code.slice(openIndex + 3, newlineIndex).trim() || "";
-    language =
-      language === "markdown"
-        ? "markdown"
-        : language === "sql"
-          ? "sql"
-          : "python";
-    const codeStart = newlineIndex + 1;
-
-    const closeIndex = code.indexOf("```", codeStart);
-    if (closeIndex === -1) {
-      // If there's no closing backticks, treat everything after the opening as code
-      const codeContent = code.slice(codeStart).replace(/\n+$/, "");
-      if (codeContent.trim()) {
-        cells.push({
-          language: language as LanguageAdapterType,
-          code: codeContent,
-        });
-      }
-      break;
-    }
-
-    // Remove trailing newlines
-    const codeContent = code.slice(codeStart, closeIndex).replace(/\n+$/, "");
-    if (codeContent.trim()) {
-      cells.push({
-        language: language as LanguageAdapterType,
-        code: codeContent,
-      });
-    }
-
-    start = closeIndex + 3;
-    openIndex = code.indexOf("```", start);
-  }
-
-  // If no cells found, assume code is in 1 cell and python
-  if (cells.length === 0) {
-    cells.push({ language: "python", code: code });
-  }
-
-  return cells;
-}

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

@@ -0,0 +1,46 @@
+/* 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;
+
+  cancel() {
+    if (this.activeAttempt) {
+      this.activeAttempt.cancelled = true;
+    }
+  }
+
+  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/StagedAICell.tsx

@@ -1,10 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useStore } from "jotai";
+import { useAtomValue, useStore } from "jotai";
 import { SparklesIcon } from "lucide-react";
 import { Tooltip } from "@/components/ui/tooltip";
 import {
   type Edit,
+  stagedGenerationInProgressAtom,
   useStagedAICell,
   useStagedCells,
 } from "@/core/ai/staged-cells";
@@ -48,6 +49,7 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
   cellId,
 }) => {
   const store = useStore();
+  const generationInProgress = useAtomValue(stagedGenerationInProgressAtom);
   const stagedAiCell = useStagedAICell(cellId);
   const runCell = useRunCell(cellId);
 
@@ -91,7 +93,7 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
       </div>
       <div className="flex items-center gap-1.5">
         <CompletionActionsCellFooter
-          isLoading={false}
+          isLoading={generationInProgress}
           onAccept={() => handleCompletion("accept")}
           onDecline={() => handleCompletion("reject")}
           size="xs"

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

@@ -0,0 +1,64 @@
+/* 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/chrome/wrapper/__tests__/pending-ai-cells.test.tsx

@@ -0,0 +1,77 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, render, screen, within } 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 { PendingAICells } from "../pending-ai-cells";
+
+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("@/components/editor/cell/useRunCells", () => ({
+  useRunCells: () => vi.fn(),
+}));
+
+vi.mock("@/components/editor/links/cell-link", () => ({
+  scrollAndHighlightCell: vi.fn(),
+}));
+
+describe("PendingAICells", () => {
+  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();
+    }
+    expect(screen.getByRole("button", { name: "Discard all" })).toBeDisabled();
+
+    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();
+  });
+});

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

@@ -4,7 +4,11 @@ import { useAtomValue, useStore } from "jotai";
 import { ChevronDown, ChevronUp, SparklesIcon } from "lucide-react";
 import { useState } from "react";
 import { Button } from "@/components/ui/button";
-import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
+import {
+  stagedAICellsAtom,
+  stagedGenerationInProgressAtom,
+  useStagedCells,
+} from "@/core/ai/staged-cells";
 import type { CellId } from "@/core/cells/ids";
 import { getNextIndex } from "@/utils/arrays";
 import { cn } from "@/utils/cn";
@@ -20,6 +24,7 @@ export const PendingAICells: React.FC = () => {
   const [currentIndex, setCurrentIndex] = useState<number | null>(null);
 
   const stagedAiCells = useAtomValue(stagedAICellsAtom);
+  const generationInProgress = useAtomValue(stagedGenerationInProgressAtom);
   const listStagedCells = [...stagedAiCells.keys()];
   const store = useStore();
   const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
@@ -93,7 +98,7 @@ export const PendingAICells: React.FC = () => {
         <AcceptCompletionButton
           multipleCompletions={true}
           onAccept={acceptAllCompletions}
-          isLoading={false}
+          isLoading={generationInProgress}
           label="Keep all"
           size="xs"
           buttonStyles="h-6.5"
@@ -103,6 +108,7 @@ export const PendingAICells: React.FC = () => {
         <RejectCompletionButton
           multipleCompletions={true}
           onDecline={rejectAllCompletions}
+          disabled={generationInProgress}
           size="xs"
           label="Discard all"
           className="h-6.5"