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

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

@@ -19,7 +19,6 @@ import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import type { FileUIPart, UIMessage } from "ai";
 import {
-  codeToCells,
   getAICompletionBody,
   getAICompletionBodyWithAttachments,
   isContextAttachment,
@@ -485,180 +484,3 @@ describe("context attachment stamping", () => {
     expect(isContextAttachment(attachments[0])).toBe(true);
   });
 });
-
-describe("codeToCells", () => {
-  it("should return empty array for empty string", () => {
-    const code = "";
-    const result = codeToCells(code);
-    expect(result).toEqual([]);
-  });
-
-  it("should return empty array for whitespace only", () => {
-    const code = "   \n\t  ";
-    const result = codeToCells(code);
-    expect(result).toEqual([]);
-  });
-
-  it("should convert code without backticks to single python cell", () => {
-    const code = "print('Hello, world!')";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should convert code with single closed backticks to cells", () => {
-    const code = "```python\nprint('Hello, world!')\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should convert code with unclosed backticks to cells", () => {
-    const code = "```python\nprint('Hello, world!')\n";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should convert code with multiple closed cells", () => {
-    const code =
-      "```python\nprint('Hello, world!')\n```\n```sql\nSELECT * FROM users\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-      { language: "sql", code: "SELECT * FROM users" },
-    ]);
-  });
-
-  it("should handle code with no language identifier", () => {
-    const code = "```\nprint('Hello, world!')\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should handle unclosed code with no language identifier", () => {
-    const code = "```\nprint('Hello, world!')\n";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should handle markdown language", () => {
-    const code = "```markdown\n# Hello, world!\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([{ language: "markdown", code: "# Hello, world!" }]);
-  });
-
-  it("should handle sql language", () => {
-    const code = "```sql\nSELECT * FROM users\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([{ language: "sql", code: "SELECT * FROM users" }]);
-  });
-
-  it("should handle unclosed markdown cell", () => {
-    const code = "```markdown\n# Hello, world!\n";
-    const result = codeToCells(code);
-    expect(result).toEqual([{ language: "markdown", code: "# Hello, world!" }]);
-  });
-
-  it("should handle unclosed sql cell", () => {
-    const code = "```sql\nSELECT * FROM users\n";
-    const result = codeToCells(code);
-    expect(result).toEqual([{ language: "sql", code: "SELECT * FROM users" }]);
-  });
-
-  it("should handle empty cells and skip them", () => {
-    const code = "```python\n\n```\n```sql\nSELECT * FROM users\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([{ language: "sql", code: "SELECT * FROM users" }]);
-  });
-
-  it("should handle cells with only whitespace and skip them", () => {
-    const code = "```python\n   \n```\n```sql\nSELECT * FROM users\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([{ language: "sql", code: "SELECT * FROM users" }]);
-  });
-
-  it("should handle code with trailing newlines", () => {
-    const code = "```python\nprint('Hello, world!')\n\n\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should handle unclosed code with trailing newlines", () => {
-    const code = "```python\nprint('Hello, world!')\n\n\n";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should handle multiple cells with different languages", () => {
-    const code =
-      "```python\nprint('Hello, world!')\n```\n```sql\nSELECT * FROM users\n```\n```markdown\n# Title\nThis is markdown\n```";
-
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-      { language: "sql", code: "SELECT * FROM users" },
-      { language: "markdown", code: "# Title\nThis is markdown" },
-    ]);
-  });
-
-  it("should handle complex multiline code", () => {
-    const code =
-      '```python\ndef hello():\n    print("Hello, world!")\n    return "success"\n\nhello()\n```\n```sql\nSELECT \n    id,\n    name,\n    email\nFROM users\nWHERE active = true\nORDER BY name;\n```';
-
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      {
-        language: "python",
-        code: 'def hello():\n    print("Hello, world!")\n    return "success"\n\nhello()',
-      },
-      {
-        language: "sql",
-        code: "SELECT \n    id,\n    name,\n    email\nFROM users\nWHERE active = true\nORDER BY name;",
-      },
-    ]);
-  });
-
-  it("should handle code with backticks in the content", () => {
-    const code = "```python\nprint('```')\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([{ language: "python", code: "print('" }]);
-  });
-
-  it("should handle code with no backticks in the last cell", () => {
-    const code =
-      "```python\nprint('Hello, world!')\n```\n```python\nprint('Hello, world!')";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should handle case insensitive language detection", () => {
-    const code = "```PYTHON\nprint('Hello, world!')\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "python", code: "print('Hello, world!')" },
-    ]);
-  });
-
-  it("should handle unknown language", { fails: true }, () => {
-    const code = "```javascript\nconsole.log('Hello, world!')\n```";
-    const result = codeToCells(code);
-    expect(result).toEqual([
-      { language: "javascript", code: "console.log('Hello, world!')" },
-    ]);
-  });
-});

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

@@ -0,0 +1,106 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { StagedCellSubmissionController } from "../staged-cell-submission";
+
+describe("StagedCellSubmissionController", () => {
+  it("ignores duplicate submissions while preprocessing", async () => {
+    const preparation = new Deferred<string>();
+    const submit = vi.fn(async () => undefined);
+    const onError = vi.fn();
+    const controller = new StagedCellSubmissionController();
+
+    const first = controller.run({
+      prepare: () => preparation.promise,
+      submit,
+      onError,
+    });
+    await controller.run({
+      prepare: vi.fn(async () => "duplicate"),
+      submit,
+      onError,
+    });
+    preparation.resolve("prepared");
+    await first;
+
+    expect(submit).toHaveBeenCalledExactlyOnceWith("prepared");
+    expect(onError).not.toHaveBeenCalled();
+  });
+
+  it("does not submit after cancellation during preprocessing", async () => {
+    const preparation = new Deferred<string>();
+    const submit = vi.fn(async () => undefined);
+    const onError = vi.fn();
+    const controller = new StagedCellSubmissionController();
+
+    const run = controller.run({
+      prepare: () => preparation.promise,
+      submit,
+      onError,
+    });
+    controller.cancel();
+    preparation.resolve("prepared");
+    await run;
+
+    expect(submit).not.toHaveBeenCalled();
+    expect(onError).not.toHaveBeenCalled();
+  });
+
+  it("blocks another submission until a cancelled submission settles", async () => {
+    const submission = new Deferred<void>();
+    const firstSubmit = vi.fn(() => submission.promise);
+    const secondSubmit = vi.fn(async () => undefined);
+    const onError = vi.fn();
+    const controller = new StagedCellSubmissionController();
+
+    const firstRun = controller.run({
+      prepare: async () => "first",
+      submit: firstSubmit,
+      onError,
+    });
+    await vi.waitFor(() => {
+      expect(firstSubmit).toHaveBeenCalledExactlyOnceWith("first");
+    });
+
+    controller.cancel();
+    await controller.run({
+      prepare: async () => "second",
+      submit: secondSubmit,
+      onError,
+    });
+    expect(secondSubmit).not.toHaveBeenCalled();
+
+    submission.resolve();
+    await firstRun;
+    await controller.run({
+      prepare: async () => "second",
+      submit: secondSubmit,
+      onError,
+    });
+
+    expect(secondSubmit).toHaveBeenCalledExactlyOnceWith("second");
+    expect(onError).not.toHaveBeenCalled();
+  });
+
+  it("reports preprocessing and submission errors", async () => {
+    const preprocessingError = new Error("preprocessing failed");
+    const submissionError = new Error("submission failed");
+    const onError = vi.fn();
+    const controller = new StagedCellSubmissionController();
+
+    await controller.run({
+      prepare: async () => Promise.reject(preprocessingError),
+      submit: vi.fn(),
+      onError,
+    });
+    await controller.run({
+      prepare: async () => "prepared",
+      submit: async () => Promise.reject(submissionError),
+      onError,
+    });
+
+    expect(onError).toHaveBeenNthCalledWith(1, preprocessingError);
+    expect(onError).toHaveBeenNthCalledWith(2, submissionError);
+  });
+});

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

@@ -24,8 +24,9 @@ import {
   SparklesIcon,
   XIcon,
 } from "lucide-react";
-import { useMemo, useRef, useState } from "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 {
@@ -37,7 +38,6 @@ import {
 import {
   buildCompletionRequestBody,
   convertToFileUIPart,
-  handleToolCall,
   PROVIDERS_THAT_SUPPORT_ATTACHMENTS,
   useFileState,
 } from "@/components/chat/chat-utils";
@@ -52,13 +52,11 @@ import {
 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 { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
-import type { ToolNotebookContext } from "@/core/ai/tools/base";
-import { useCellActions } from "@/core/cells/cells";
+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 { useRequestClient } from "@/core/network/requests";
 import type { AiCompletionRequest } from "@/core/network/types";
 import { useRuntimeManager } from "@/core/runtime/config";
 import { useTheme } from "@/theme/useTheme";
@@ -73,7 +71,7 @@ import {
   CONTEXT_TRIGGER,
   mentionsCompletionSource,
 } from "./completion-utils";
-import { StreamingChunkTransport } from "./transport/chat-transport";
+import { StagedCellSubmissionController } from "./staged-cell-submission";
 
 // Persist across sessions
 const languageAtom = atomWithStorage<"python" | "sql">(
@@ -95,78 +93,95 @@ export const AddCellWithAI: React.FC<{
   const store = useStore();
   const [input, setInput] = useState("");
 
-  const { deleteAllStagedCells, clearStagedCells, onStream, addStagedCell } =
-    useStagedCells(store);
+  const {
+    acceptOwnedStagedCells,
+    beginStagedCellGeneration,
+    discardOwnedStagedCells,
+    finishStagedCellGeneration,
+    hasOwnedStagedCells,
+    onData,
+  } = useStagedCellGeneration(store);
   const [language, setLanguage] = useAtom(languageAtom);
   const runtimeManager = useRuntimeManager();
-  const { invokeAiTool, sendRun } = useRequestClient();
 
-  const stagedAICells = useAtomValue(stagedAICellsAtom);
   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 { createNewCell, prepareForRun } = useCellActions();
-  const toolContext: ToolNotebookContext = {
-    store,
-    addStagedCell,
-    createNewCell,
-    prepareForRun,
-    sendRun,
-  };
-
-  const { sendMessage, stop, status, addToolOutput } = useChat({
-    throttle: AI_SDK_UI_THROTTLE_MS,
-    transport: new StreamingChunkTransport(
-      {
+  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: AiCompletionRequest = {
-            ...options,
+          const body = {
             ...completionBody,
             code: "",
             prompt: "", // Don't need prompt since we are using messages
             language: language,
-          };
+          } satisfies AiCompletionRequest;
 
           return {
             api: runtimeManager.getAiURL("completion").toString(),
             body: body,
           };
         },
-      },
-      (chunk) => {
-        onStream(chunk);
-      },
-    ),
-    onToolCall: async ({ toolCall }) => {
-      await handleToolCall({
-        invokeAiTool,
-        addToolOutput,
-        toolCall: {
-          toolName: toolCall.toolName,
-          toolCallId: toolCall.toolCallId,
-          input: toolCall.input as Record<string, never>,
-        },
-        toolContext,
-      });
-    },
-    onError: (error) => {
-      toast({
-        title: "Generate with AI failed",
-        description: prettyError(error),
-      });
+      }),
+    [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 = stagedAICells.size > 0;
+  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;
@@ -175,14 +190,19 @@ export const AddCellWithAI: React.FC<{
 
   const submit = async () => {
     if (!isLoading) {
-      if (inputRef.current?.view) {
-        storePrompt(inputRef.current.view);
-      }
-      // TODO: When we have conversations, don't delete existing cells
-      deleteAllStagedCells();
-
-      const fileParts = files ? await convertToFileUIPart(files) : undefined;
-      sendMessage({ text: input, files: fileParts });
+      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,
+      });
     }
   };
 
@@ -225,18 +245,19 @@ export const AddCellWithAI: React.FC<{
   );
 
   const handleAcceptCompletion = () => {
-    clearStagedCells();
-    onClose();
+    if (acceptOwnedStagedCells()) {
+      onClose();
+    }
   };
 
   const handleDeclineCompletion = () => {
-    deleteAllStagedCells();
+    reject();
     // Focus the input so the user can refine the prompt.
     inputRef.current?.view?.focus();
   };
 
   const handleClose = () => {
-    deleteAllStagedCells();
+    reject();
     onClose();
   };