@marimo-team/islands 0.23.15-dev7 → 0.23.15-dev72

package/src/components/editor/__tests__/output-persistence.test.tsx

@@ -0,0 +1,241 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { render } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import { beforeAll, describe, expect, it } from "vitest";
+import { SetupMocks } from "@/__mocks__/common";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { cellId } from "@/__tests__/branded";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { notebookAtom } from "@/core/cells/cells";
+import type { CellRuntimeState } from "@/core/cells/types";
+import { createCellRuntimeState } from "@/core/cells/types";
+import { defaultUserConfig } from "@/core/config/config-schema";
+import type { AppMode } from "@/core/mode";
+import { requestClientAtom } from "@/core/network/requests";
+import { Cell } from "../notebook-cell";
+
+const cid = cellId("test");
+
+function createTestWrapper() {
+  const store = createStore();
+  store.set(requestClientAtom, MockRequestClient.create());
+  const wrapper = ({ children }: { children: React.ReactNode }) => (
+    <Provider store={store}>{children}</Provider>
+  );
+  return { wrapper, store };
+}
+
+function seedNotebook(
+  store: ReturnType<typeof createStore>,
+  runtimeOverrides: Partial<CellRuntimeState> = {},
+) {
+  const notebook = MockNotebook.notebookState({
+    cellData: {
+      [cid]: {
+        code: "1 + 1",
+        name: "test_cell",
+        edited: false,
+        serializedEditorState: null,
+        config: {
+          disabled: false,
+          hide_code: false,
+          column: null,
+        },
+      },
+    },
+  });
+  notebook.cellRuntime[cid] = createCellRuntimeState({
+    status: "idle",
+    output: {
+      channel: "output",
+      mimetype: "text/html",
+      data: "<span>persist me</span>",
+      timestamp: 0,
+    },
+    ...runtimeOverrides,
+  });
+  store.set(notebookAtom, notebook);
+}
+
+function renderCell(mode: AppMode) {
+  return (
+    <TooltipProvider>
+      <Cell
+        cellId={cid}
+        mode={mode}
+        canDelete={true}
+        userConfig={defaultUserConfig()}
+        isCollapsed={false}
+        collapseCount={0}
+        canMoveX={false}
+        theme="light"
+        showPlaceholder={false}
+      />
+    </TooltipProvider>
+  );
+}
+
+beforeAll(() => {
+  SetupMocks.resizeObserver();
+  global.HTMLDivElement.prototype.scrollIntoView = () => {
+    // do nothing
+  };
+});
+
+describe("output persistence across mode changes", () => {
+  it("preserves the output DOM node when toggling edit <-> present", () => {
+    const { store, wrapper } = createTestWrapper();
+    seedNotebook(store);
+
+    const { container, rerender } = render(renderCell("edit"), { wrapper });
+
+    const outputNode = container.querySelector<HTMLElement>(
+      '[data-cell-role="output"]',
+    );
+    expect(outputNode).toBeTruthy();
+    // A property survives only if the exact DOM node instance survives.
+    (outputNode as HTMLElement & { __sentinel?: string }).__sentinel = "alive";
+
+    rerender(renderCell("present"));
+    const presentNode = container.querySelector<HTMLElement>(
+      '[data-cell-role="output"]',
+    );
+    expect(presentNode).toBe(outputNode);
+    expect(presentNode?.isConnected).toBe(true);
+    expect(
+      (presentNode as HTMLElement & { __sentinel?: string }).__sentinel,
+    ).toBe("alive");
+
+    // Edit chrome is hidden but stays mounted while presenting
+    const tray = container.querySelector<HTMLElement>(
+      "[data-testid='cell-tray']",
+    );
+    expect(tray).toBeTruthy();
+    expect(tray?.hidden).toBe(true);
+    // Present styling is applied
+    expect(
+      container
+        .querySelector(`[data-cell-id="${cid}"]`)
+        ?.classList.contains("published"),
+    ).toBe(true);
+
+    rerender(renderCell("edit"));
+    expect(container.querySelector('[data-cell-role="output"]')).toBe(
+      outputNode,
+    );
+    expect(
+      container.querySelector<HTMLElement>("[data-testid='cell-tray']")?.hidden,
+    ).toBe(false);
+    expect(
+      container
+        .querySelector(`[data-cell-id="${cid}"]`)
+        ?.classList.contains("published"),
+    ).toBe(false);
+  });
+
+  it("hides all edit chrome while presenting", () => {
+    const { store, wrapper } = createTestWrapper();
+    seedNotebook(store);
+
+    const { container, rerender } = render(renderCell("edit"), { wrapper });
+
+    // Returns true when the element is not visible: either gone from the DOM or
+    // inside a `hidden` ancestor. Edit chrome may unmount while presenting (only
+    // the cell *output* is contractually kept mounted — see the DOM-node test
+    // above); either way it must not be visible.
+    const isHidden = (testId: string) => {
+      const el = container.querySelector(`[data-testid='${testId}']`);
+      return el === null || el.closest("[hidden]") !== null;
+    };
+
+    // Everything visible in edit mode
+    const editChrome = [
+      "cell-tray",
+      "drag-button",
+      "cell-actions-button",
+      "delete-button",
+      "create-cell-button",
+      "console-output-area",
+    ];
+    for (const testId of editChrome) {
+      expect({ testId, hidden: isHidden(testId) }).toEqual({
+        testId,
+        hidden: false,
+      });
+    }
+
+    rerender(renderCell("present"));
+    for (const testId of editChrome) {
+      expect({ testId, hidden: isHidden(testId) }).toEqual({
+        testId,
+        hidden: true,
+      });
+    }
+
+    rerender(renderCell("edit"));
+    for (const testId of editChrome) {
+      expect({ testId, hidden: isHidden(testId) }).toEqual({
+        testId,
+        hidden: false,
+      });
+    }
+  });
+
+  it("does not dim edited-but-not-run outputs as stale while presenting", () => {
+    const { store, wrapper } = createTestWrapper();
+    seedNotebook(store);
+    // Mark the cell as edited without re-running it
+    const notebook = store.get(notebookAtom);
+    notebook.cellData[cid] = { ...notebook.cellData[cid], edited: true };
+    store.set(notebookAtom, { ...notebook });
+
+    const { container, rerender } = render(renderCell("edit"), { wrapper });
+    expect(container.querySelector(".marimo-output-stale")).toBeTruthy();
+
+    // The read view never dims pending edits; present mode matches it
+    rerender(renderCell("present"));
+    expect(container.querySelector(".marimo-output-stale")).toBeFalsy();
+
+    rerender(renderCell("edit"));
+    expect(container.querySelector(".marimo-output-stale")).toBeTruthy();
+  });
+
+  it("hides errored cells while presenting without unmounting them", () => {
+    const { store, wrapper } = createTestWrapper();
+    seedNotebook(store, {
+      errored: true,
+      output: {
+        channel: "marimo-error",
+        mimetype: "application/vnd.marimo+error",
+        data: [{ type: "exception", exception_type: "ValueError", msg: "!" }],
+        timestamp: 0,
+      },
+    });
+
+    const { container, rerender } = render(renderCell("edit"), { wrapper });
+
+    const outputNode = container.querySelector<HTMLElement>(
+      '[data-cell-role="output"]',
+    );
+    expect(outputNode).toBeTruthy();
+
+    const getSortableWrapper = () =>
+      container.querySelector<HTMLElement>(`[data-cell-id="${cid}"]`)
+        ?.parentElement;
+    expect(getSortableWrapper()?.hidden).toBe(false);
+
+    rerender(renderCell("present"));
+    // Still mounted, but hidden
+    expect(container.querySelector('[data-cell-role="output"]')).toBe(
+      outputNode,
+    );
+    expect(getSortableWrapper()?.hidden).toBe(true);
+
+    rerender(renderCell("edit"));
+    expect(container.querySelector('[data-cell-role="output"]')).toBe(
+      outputNode,
+    );
+    expect(getSortableWrapper()?.hidden).toBe(false);
+  });
+});

package/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts

@@ -0,0 +1,153 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { afterEach, describe, expect, it, vi } from "vitest";
+import {
+  type ConnectionInfo,
+  getMarimoCommand,
+  getRawPrompt,
+  getTerminalCommand,
+  maskToken,
+  shellQuote,
+} from "../pair-with-agent-commands";
+
+const CONNECTION: ConnectionInfo = {
+  url: "http://localhost:8000",
+  sessionId: "s_ab12cd",
+};
+
+describe("shellQuote", () => {
+  it("quotes an empty string", () => {
+    expect(shellQuote("")).toBe("''");
+  });
+
+  it("leaves shell-safe values untouched", () => {
+    expect(shellQuote("http://localhost:8000")).toBe("http://localhost:8000");
+    expect(shellQuote("s_ab12cd")).toBe("s_ab12cd");
+  });
+
+  it("quotes values with shell metacharacters", () => {
+    expect(shellQuote("http://host:8000?a=1&b=2")).toBe(
+      "'http://host:8000?a=1&b=2'",
+    );
+    expect(shellQuote("has space")).toBe("'has space'");
+    expect(shellQuote("$(rm -rf /)")).toBe("'$(rm -rf /)'");
+  });
+
+  it("escapes embedded single quotes without breaking out", () => {
+    // Closes the quote, emits a literal ' via "'", then reopens: '"'"'
+    expect(shellQuote("a'b")).toBe(`'a'"'"'b'`);
+  });
+});
+
+describe("getMarimoCommand", () => {
+  afterEach(() => {
+    vi.unstubAllEnvs();
+  });
+
+  it("uses the local checkout in dev", () => {
+    vi.stubEnv("DEV", true);
+    expect(getMarimoCommand()).toBe("uv run marimo");
+  });
+
+  it("uses uvx outside of dev", () => {
+    vi.stubEnv("DEV", false);
+    expect(getMarimoCommand()).toBe("uvx marimo@latest");
+  });
+});
+
+describe("getTerminalCommand", () => {
+  it("includes the url and session for each agent", () => {
+    expect(getTerminalCommand("claude", CONNECTION, false)).toBe(
+      `claude "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --claude)"`,
+    );
+    expect(getTerminalCommand("codex", CONNECTION, false)).toBe(
+      `codex "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --codex)"`,
+    );
+    expect(getTerminalCommand("opencode", CONNECTION, false)).toBe(
+      `opencode --prompt "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --opencode)"`,
+    );
+  });
+
+  it("always targets the given session, not a random one", () => {
+    const command = getTerminalCommand("claude", CONNECTION, false);
+    expect(command).toContain("--session s_ab12cd");
+  });
+
+  it("shell-escapes a url containing metacharacters", () => {
+    const command = getTerminalCommand(
+      "claude",
+      { url: "http://host:8000?file=a&b", sessionId: "s_ab12cd" },
+      false,
+    );
+    expect(command).toContain("--url 'http://host:8000?file=a&b'");
+  });
+
+  it("adds --with-token before the agent flag when requested", () => {
+    const command = getTerminalCommand("claude", CONNECTION, true);
+    expect(command).toContain("--with-token --claude");
+  });
+
+  it("omits --with-token when not requested", () => {
+    expect(getTerminalCommand("claude", CONNECTION, false)).not.toContain(
+      "--with-token",
+    );
+  });
+});
+
+describe("getRawPrompt", () => {
+  it("references the session-scoped execute-code command", () => {
+    const prompt = getRawPrompt(CONNECTION, null);
+    expect(prompt).toContain(
+      "execute-code.sh --url http://localhost:8000 --session s_ab12cd",
+    );
+    expect(prompt).toContain(
+      "Connect to the notebook at: http://localhost:8000 (session s_ab12cd)",
+    );
+  });
+
+  it("omits the token hint when there is no token", () => {
+    const prompt = getRawPrompt(CONNECTION, null);
+    expect(prompt).not.toContain("--token");
+    expect(prompt).not.toContain("auth token");
+  });
+
+  it("includes a session-scoped token hint when a token is present", () => {
+    const prompt = getRawPrompt(CONNECTION, "secret-token");
+    expect(prompt).toContain(
+      "execute-code.sh --url http://localhost:8000 --session s_ab12cd --token secret-token",
+    );
+  });
+
+  it("shell-escapes a token containing a single quote", () => {
+    const prompt = getRawPrompt(CONNECTION, "tok'en");
+    expect(prompt).toContain(`--token 'tok'"'"'en'`);
+  });
+
+  it("matches the CLI prompt shape", () => {
+    const prompt = getRawPrompt(CONNECTION, null);
+    expect(prompt).toMatchInlineSnapshot(`
+      "Use the /marimo-pair skill to pair-program on a running marimo notebook.
+
+      Connect to the notebook at: http://localhost:8000 (session s_ab12cd)
+
+      Use \`execute-code.sh --url http://localhost:8000 --session s_ab12cd\` from the marimo-pair skill to execute code in the notebook.
+
+      Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair."
+    `);
+  });
+});
+
+describe("maskToken", () => {
+  it("masks short tokens entirely", () => {
+    expect(maskToken("ab")).toBe("****");
+    expect(maskToken("abcd")).toBe("****");
+  });
+
+  it("reveals only the last four characters", () => {
+    expect(maskToken("abcdefgh")).toBe("****efgh");
+  });
+
+  it("caps the number of mask characters at eight", () => {
+    expect(maskToken("0123456789abcdef")).toBe("********cdef");
+  });
+});

package/src/components/editor/actions/pair-with-agent-commands.ts

@@ -0,0 +1,109 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { assertNever } from "@/utils/assertNever";
+
+export type AgentTab = "claude" | "codex" | "opencode" | "prompt";
+
+export const TERMINAL_TABS = ["claude", "codex", "opencode"] as const;
+
+export const AGENT_TABS = ["claude", "codex", "opencode", "prompt"] as const;
+
+export const AGENT_LABELS: Record<AgentTab, string> = {
+  claude: "Claude",
+  codex: "Codex",
+  opencode: "OpenCode",
+  prompt: "Prompt",
+};
+
+export const SKILL_INSTALL = "npx skills add marimo-team/marimo-pair";
+
+/** How to invoke marimo: from the local checkout in dev, else via uvx. */
+export function getMarimoCommand(): string {
+  return import.meta.env.DEV ? "uv run marimo" : "uvx marimo@latest";
+}
+
+/**
+ * POSIX-quote a value for safe embedding in a shell command. These commands are
+ * meant to be copied into a terminal, so a url/token containing shell
+ * metacharacters (`'`, `&`, `$(...)`, ...) must not break out of its argument.
+ *
+ * Mirrors Python's `shlex.quote` (used on the CLI side in
+ * `marimo/_cli/pair/commands.py`) so both sides produce identical commands:
+ * values that are already shell-safe are left as-is for readability, and
+ * anything else is single-quoted with embedded quotes escaped as `'"'"'`.
+ */
+export function shellQuote(value: string): string {
+  if (value === "") {
+    return "''";
+  }
+  // Same "safe" character set as CPython's shlex.quote (ASCII \w plus a few).
+  if (/^[\w@%+=:,./-]+$/.test(value)) {
+    return value;
+  }
+  return `'${value.replaceAll("'", `'"'"'`)}'`;
+}
+
+/** Identifies the specific running notebook to pair on. */
+export interface ConnectionInfo {
+  url: string;
+  /**
+   * The current session id, so agents connect to *this* notebook rather than
+   * guessing when multiple sessions are open on the same server.
+   */
+  sessionId: string;
+}
+
+/**
+ * The shell command that wraps an agent CLI, delegating prompt generation to
+ * `marimo pair prompt` so the terminal and CLI stay in sync.
+ */
+export function getTerminalCommand(
+  agent: Exclude<AgentTab, "prompt">,
+  { url, sessionId }: ConnectionInfo,
+  withToken: boolean,
+): string {
+  const tokenFlag = withToken ? " --with-token" : "";
+  const base = `${getMarimoCommand()} pair prompt --url ${shellQuote(url)} --session ${shellQuote(sessionId)}${tokenFlag}`;
+  switch (agent) {
+    case "claude":
+      return `claude "$(${base} --claude)"`;
+    case "codex":
+      return `codex "$(${base} --codex)"`;
+    case "opencode":
+      return `opencode --prompt "$(${base} --opencode)"`;
+    default:
+      assertNever(agent);
+  }
+}
+
+/**
+ * The raw prompt for the "Prompt" tab. Mirrors the output of
+ * `marimo pair prompt` (see `marimo/_cli/pair/commands.py`) so pasting it into
+ * an agent behaves the same as the terminal commands.
+ */
+export function getRawPrompt(
+  { url, sessionId }: ConnectionInfo,
+  token: string | null,
+): string {
+  const executeCmd = `execute-code.sh --url ${shellQuote(url)} --session ${shellQuote(sessionId)}`;
+  const tokenHint = token
+    ? `\n\nUse this auth token when calling \`execute-code.sh\`: \`${executeCmd} --token ${shellQuote(token)}\`.`
+    : "";
+  return [
+    "Use the /marimo-pair skill to pair-program on a running marimo notebook.",
+    "",
+    `Connect to the notebook at: ${url} (session ${sessionId})`,
+    "",
+    `Use \`${executeCmd}\` from the marimo-pair skill to execute code in the notebook.${tokenHint}`,
+    "",
+    "Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair.",
+  ].join("\n");
+}
+
+/** Mask all but the last 4 chars of a token for display. */
+export function maskToken(token: string): string {
+  if (token.length <= 4) {
+    return "****";
+  }
+  return `${"*".repeat(Math.min(token.length - 4, 8))}${token.slice(-4)}`;
+}

package/src/components/editor/actions/pair-with-agent-modal.tsx

@@ -14,67 +14,20 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { copyToClipboard } from "@/utils/copy";
 import { Events } from "@/utils/events";
 import { Tooltip } from "@/components/ui/tooltip";
-import { assertNever } from "@/utils/assertNever";
 import { asRemoteURL, useRuntimeManager } from "@/core/runtime/config";
 import { API } from "@/core/network/api";
-
-type AgentTab = "claude" | "codex" | "opencode" | "prompt";
-
-const TERMINAL_TABS = ["claude", "codex", "opencode"] as const;
-
-function getMarimoCommand(): string {
-  return import.meta.env.DEV ? "uv run marimo" : "uvx marimo@latest";
-}
-
-function getTerminalCommand(
-  agent: Exclude<AgentTab, "prompt">,
-  url: string,
-  withToken: boolean,
-): string {
-  const tokenFlag = withToken ? " --with-token" : "";
-  const base = `${getMarimoCommand()} pair prompt --url '${url}'${tokenFlag}`;
-  switch (agent) {
-    case "claude":
-      return `claude "$(${base} --claude)"`;
-    case "codex":
-      return `codex "$(${base} --codex)"`;
-    case "opencode":
-      return `opencode --prompt "$(${base} --opencode)"`;
-    default:
-      assertNever(agent);
-  }
-}
-
-function getRawPrompt(url: string, token: string | null): string {
-  const tokenHint = token
-    ? `\n\nUse this auth token when calling \`execute-code.sh\`: \`execute-code.sh --url '${url}' --token '${token}'\`.`
-    : "";
-  return [
-    "Use the /marimo-pair skill to pair-program on a running marimo notebook.",
-    "",
-    `Connect to the notebook at: ${url}`,
-    "",
-    `Use \`execute-code.sh --url ${url}\` from the marimo-pair skill to execute code in the notebook.${tokenHint}`,
-    "",
-    "Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair.",
-  ].join("\n");
-}
-
-function maskToken(token: string): string {
-  if (token.length <= 4) {
-    return "****";
-  }
-  return `${"*".repeat(Math.min(token.length - 4, 8))}${token.slice(-4)}`;
-}
-
-const SKILL_INSTALL = "npx skills add marimo-team/marimo-pair";
-
-const AGENT_LABELS: Record<AgentTab, string> = {
-  claude: "Claude",
-  codex: "Codex",
-  opencode: "OpenCode",
-  prompt: "Prompt",
-};
+import { getSessionId } from "@/core/kernel/session";
+import {
+  AGENT_LABELS,
+  AGENT_TABS,
+  type AgentTab,
+  type ConnectionInfo,
+  getRawPrompt,
+  getTerminalCommand,
+  maskToken,
+  SKILL_INSTALL,
+  TERMINAL_TABS,
+} from "./pair-with-agent-commands";
 
 function useAuthToken(): string | null {
   const [token, setToken] = useState<string | null>(null);
@@ -98,7 +51,10 @@ export const PairWithAgentModal: React.FC<{
   const runtimeManager = useRuntimeManager();
   const authToken = useAuthToken();
   const hasToken = Boolean(authToken);
-  const remoteUrl = runtimeManager.httpURL.toString();
+  const connection: ConnectionInfo = {
+    url: runtimeManager.httpURL.toString(),
+    sessionId: getSessionId(),
+  };
 
   return (
     <DialogContent className="sm:max-w-2xl">
@@ -125,7 +81,7 @@ export const PairWithAgentModal: React.FC<{
           onValueChange={(v) => setActiveTab(v as AgentTab)}
         >
           <TabsList className="w-full">
-            {(["claude", "codex", "opencode", "prompt"] as const).map((tab) => (
+            {AGENT_TABS.map((tab) => (
               <TabsTrigger key={tab} value={tab} className="flex-1">
                 {AGENT_LABELS[tab]}
               </TabsTrigger>
@@ -147,7 +103,7 @@ export const PairWithAgentModal: React.FC<{
               </Step>
               <Step index={2} title="Run in your terminal">
                 <CommandBlock
-                  command={getTerminalCommand(tab, remoteUrl, hasToken)}
+                  command={getTerminalCommand(tab, connection, hasToken)}
                 />
               </Step>
               {hasToken && authToken && (
@@ -179,9 +135,9 @@ export const PairWithAgentModal: React.FC<{
               }
             >
               <CommandBlock
-                command={getRawPrompt(remoteUrl, authToken)}
+                command={getRawPrompt(connection, authToken)}
                 display={getRawPrompt(
-                  remoteUrl,
+                  connection,
                   authToken ? maskToken(authToken) : null,
                 )}
                 multiline={true}

package/src/components/editor/actions/useNotebookActions.tsx

@@ -28,6 +28,7 @@ import {
   KeyboardIcon,
   LayoutTemplateIcon,
   LinkIcon,
+  MessageCircleQuestionIcon,
   MessagesSquareIcon,
   NotebookIcon,
   PanelLeftIcon,
@@ -44,6 +45,7 @@ import {
   settingDialogAtom,
   useOpenSettingsToTab,
 } from "@/components/app-config/state";
+import { FeedbackModal } from "@/components/editor/chrome/components/feedback-button";
 import { MarkdownIcon } from "@/components/editor/cell/code/icons";
 import { GitHubIcon } from "@/components/icons/github";
 import { MarimoPlusIcon } from "@/components/icons/marimo-icons";
@@ -642,6 +644,12 @@ export function useNotebookActions() {
       redundant: true,
       additionalKeywords: ["preferences", "options", "configuration"],
     },
+    {
+      icon: <MessageCircleQuestionIcon size={14} strokeWidth={1.5} />,
+      label: "Report an issue",
+      additionalKeywords: ["feedback", "bug", "issue", "report", "diagnostics"],
+      handle: () => openModal(<FeedbackModal onClose={closeModal} />),
+    },
     {
       icon: <ExternalLinkIcon size={14} strokeWidth={1.5} />,
       label: "Resources",

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

@@ -231,6 +231,13 @@ export const AddCellWithAI: React.FC<{
 
   const handleDeclineCompletion = () => {
     deleteAllStagedCells();
+    // Focus the input so the user can refine the prompt.
+    inputRef.current?.view?.focus();
+  };
+
+  const handleClose = () => {
+    deleteAllStagedCells();
+    onClose();
   };
 
   const inputComponent = (
@@ -238,10 +245,7 @@ export const AddCellWithAI: React.FC<{
       <SparklesIcon className="size-4 text-(--blue-11) mr-2" />
       <PromptInput
         inputRef={inputRef}
-        onClose={() => {
-          deleteAllStagedCells();
-          onClose();
-        }}
+        onClose={handleClose}
         value={input}
         onChange={(newValue) => {
           setInput(newValue);
@@ -254,7 +258,12 @@ export const AddCellWithAI: React.FC<{
           hasCompletion,
         })}
       />
-      <Button variant="text" size="sm" className="mb-0 px-1" onClick={onClose}>
+      <Button
+        variant="text"
+        size="sm"
+        className="mb-0 px-1"
+        onClick={handleClose}
+      >
         <XIcon className="size-4" />
       </Button>
     </div>