@marimo-team/islands 0.25.3-dev29 → 0.25.3-dev290

package/src/components/editor/chrome/wrapper/__tests__/attachment-status.test.tsx

@@ -14,7 +14,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { PAIR_CONFIG } from "@/__tests__/fixtures/pair-config";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import { toast } from "@/components/ui/use-toast";
-import type { Attachment } from "@/core/attachments/state";
+import type { Attachment } from "@/core/notebook/attachments";
 import { pairConfigAtom } from "@/core/config/pair";
 import {
   type StableSessionId,
@@ -22,8 +22,7 @@ import {
 } from "@/core/kernel/session";
 import { API } from "@/core/network/api";
 import { connectionAtom } from "@/core/network/connection";
-import { notebookClientAtom } from "@/core/websocket/state";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { Deferred } from "@/utils/Deferred";
 import { AttachmentStatus } from "../footer-items/attachment-status";
 
@@ -31,6 +30,10 @@ vi.mock("@/components/ui/use-toast", () => ({ toast: vi.fn() }));
 vi.mock("@/core/network/api", () => ({ API: { post: vi.fn() } }));
 
 const attachmentsAtom = atom<Attachment[]>([]);
+vi.mock("@/core/session/active-notebook", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/core/session/active-notebook")>()),
+  activeNotebookReplicaAtom: atom(() => ({ attachmentsAtom })),
+}));
 
 const AGENT: Attachment = {
   id: "a1",
@@ -38,18 +41,13 @@ const AGENT: Attachment = {
   name: "Pi",
   since: new Date().toISOString(),
   notebook_id: "nb-one",
+  identity: null,
 };
 
 function renderStatus(attachments: Attachment[] = [AGENT], preview = true) {
   const store = createStore();
   store.set(connectionAtom, { state: ConnectionState.OPEN });
   store.set(pairConfigAtom, { ...PAIR_CONFIG, hostEnabled: preview });
-  store.set(notebookClientAtom, {
-    connectionAtom: atom({ state: ConnectionState.OPEN }),
-    documentAtom: atom(null),
-    attachmentsAtom,
-    reconnect: async () => undefined,
-  });
   store.set(attachmentsAtom, attachments);
   store.set(stableSessionIdAtom, "stable/test" as StableSessionId);
   return {

package/src/components/editor/chrome/wrapper/__tests__/document-sync-status.test.tsx

@@ -0,0 +1,92 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, fireEvent, render, screen } from "@testing-library/react";
+import { atom, createStore, Provider } from "jotai";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { createDocumentModel } from "@/core/notebook/document";
+import { document } from "@/core/notebook/document/__tests__/fixtures";
+import {
+  type FakeEdit,
+  fakeBackend,
+} from "@/core/notebook/replica/__tests__/fixtures";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { DocumentSyncStatus } from "../footer-items/document-sync-status";
+
+const lifetimes: AbortController[] = [];
+
+afterEach(() => {
+  act(() => lifetimes.splice(0).forEach((controller) => controller.abort()));
+});
+
+const conflict = {
+  ok: false as const,
+  status: 409,
+  problem: {
+    type: "tag:marimo.io,2026:host/conflict",
+    title: "Conflict",
+    status: 409,
+    detail: "cell c1 has changed since it was read.",
+  },
+};
+
+function setup(edit: FakeEdit) {
+  const store = createStore();
+  const model = createDocumentModel({
+    store,
+    documentAtom: atom(document()),
+    backendAtom: atom(fakeBackend({ edit })),
+  });
+  const controller = new AbortController();
+  lifetimes.push(controller);
+  model.start({ signal: controller.signal });
+  store.set(activeDocumentModelAtom, model);
+  render(
+    <Provider store={store}>
+      <TooltipProvider>
+        <DocumentSyncStatus />
+      </TooltipProvider>
+    </Provider>,
+  );
+  return { store, model };
+}
+
+describe("DocumentSyncStatus", () => {
+  it("renders nothing while edits reach the server", () => {
+    setup(vi.fn<FakeEdit>());
+    expect(screen.queryByTestId("footer-document-sync")).toBeNull();
+  });
+
+  it("shows a rejected edit and resubmits it on click", async () => {
+    const accepted = document().revision;
+    const edit = vi
+      .fn<FakeEdit>()
+      .mockResolvedValueOnce(conflict)
+      .mockResolvedValueOnce({
+        ok: true,
+        status: 200,
+        data: {
+          document: { ...accepted, version: 2, etag: "document-2" },
+          cells: {},
+        },
+      });
+    const { store, model } = setup(edit);
+
+    await act(async () => {
+      store.set(model.editAtom, (tx) => tx.setCode("c1", "x = 2"));
+      await model.flush().catch(() => undefined);
+    });
+    const item = screen.getByTestId("footer-document-sync");
+    expect(item).toHaveAttribute(
+      "aria-label",
+      "Edits are not reaching the server",
+    );
+
+    await act(async () => {
+      fireEvent.click(item);
+      await Promise.resolve();
+    });
+    expect(edit).toHaveBeenCalledTimes(2);
+    expect(screen.queryByTestId("footer-document-sync")).toBeNull();
+  });
+});

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

@@ -1,77 +1,305 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { act, 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 {
+  isChatBusy,
+  useChatGeneration,
+} from "@/components/chat/use-chat-generation";
 import { cellId } from "@/__tests__/branded";
+import { HTMLCellId } from "@/core/cells/ids";
+import { scrollToCell } from "@/core/cells/scrollCellIntoView";
 import { TooltipProvider } from "@/components/ui/tooltip";
-import { stagedAICellsAtom, visibleForTesting } from "@/core/ai/staged-cells";
+import { stagedAICellsAtom } from "@/core/ai/staged-cells";
 import { PendingAICells } from "../pending-ai-cells";
+import { PendingAIChangesBadge } from "../pending-ai-changes-badge";
+
+const deleteCell = vi.hoisted(() => vi.fn());
+const setCellCode = vi.hoisted(() => vi.fn());
 
 vi.mock("@/components/editor/cell/useDeleteCell", () => ({
-  useDeleteCellCallback: () => vi.fn(),
+  useDeleteCellCallback: () => deleteCell,
 }));
 
 vi.mock("@/core/cells/cells", () => ({
   getCellEditorView: vi.fn(),
+}));
+vi.mock("@/core/cells/actions", () => ({
   useCellActions: () => ({
-    createNewCell: vi.fn(),
-    updateCellCode: vi.fn(),
+    createCell: vi.fn(),
+    setCellCode,
   }),
 }));
 
 vi.mock("@/components/editor/cell/useRunCells", () => ({
   useRunCells: () => vi.fn(),
+  useRunCell: () => vi.fn(),
 }));
 
-vi.mock("@/components/editor/links/cell-link", () => ({
-  scrollAndHighlightCell: vi.fn(),
+vi.mock("@/core/cells/scrollCellIntoView", () => ({
+  scrollToCell: vi.fn(),
 }));
 
 describe("PendingAICells", () => {
-  it("disables bulk actions while cells are still generating", () => {
+  it("navigates between pending cells through named controls", () => {
     const store = createStore();
-    const generationId = Symbol("staged-cell-generation");
+    const first = cellId("first-pending-cell");
+    const second = cellId("second-pending-cell");
+    const third = cellId("third-pending-cell");
     store.set(
       stagedAICellsAtom,
-      new Map([[cellId("generated-cell"), { type: "add_cell" }]]),
+      new Map([
+        [first, { type: "add_cell" }],
+        [second, { type: "add_cell" }],
+        [third, { type: "add_cell" }],
+      ]),
     );
-    store.set(visibleForTesting.stagedGenerationAtom, {
-      id: generationId,
-      status: "in_progress",
-      cellIds: [cellId("generated-cell")],
-    });
 
-    const view = (
+    render(
       <Provider store={store}>
         <TooltipProvider>
           <PendingAICells />
         </TooltipProvider>
-      </Provider>
+      </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",
-      });
+    const next = screen.getByRole("button", { name: "Next pending AI cell" });
+    const previous = screen.getByRole("button", {
+      name: "Previous pending AI cell",
     });
-    rerender(view);
 
-    for (const button of acceptButtons) {
-      expect(button).toBeEnabled();
+    fireEvent.click(screen.getByRole("button", { name: "Cell 1 of 3" }));
+    expect(scrollToCell).toHaveBeenLastCalledWith(first);
+    expect(screen.getByText("Cell 1 of 3")).toBeInTheDocument();
+
+    fireEvent.click(next);
+    expect(scrollToCell).toHaveBeenLastCalledWith(second);
+    expect(screen.getByText("Cell 2 of 3")).toBeInTheDocument();
+
+    fireEvent.click(previous);
+    expect(scrollToCell).toHaveBeenLastCalledWith(first);
+    expect(screen.getByText("Cell 1 of 3")).toBeInTheDocument();
+
+    fireEvent.click(next);
+    expect(scrollToCell).toHaveBeenLastCalledWith(second);
+    expect(screen.getByText("Cell 2 of 3")).toBeInTheDocument();
+
+    fireEvent.click(next);
+    expect(scrollToCell).toHaveBeenLastCalledWith(third);
+    expect(screen.getByText("Cell 3 of 3")).toBeInTheDocument();
+
+    fireEvent.click(next);
+    expect(scrollToCell).toHaveBeenLastCalledWith(first);
+    expect(screen.getByText("Cell 1 of 3")).toBeInTheDocument();
+
+    fireEvent.click(previous);
+    expect(scrollToCell).toHaveBeenLastCalledWith(third);
+    expect(screen.getByText("Cell 3 of 3")).toBeInTheDocument();
+  });
+});
+
+function ChatGeneration({ status }: { status: ChatStatus }) {
+  useChatGeneration(isChatBusy(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" }],
+      [cellId("other-pending-cell"), { type: "add_cell" }],
+    ]),
+  );
+  const ui = (status?: ChatStatus) => (
+    <Provider store={store}>
+      <TooltipProvider>
+        {status && <ChatGeneration status={status} />}
+        <PendingAICells />
+      </TooltipProvider>
+    </Provider>
+  );
+  const { rerender } = render(ui("ready"));
+  for (const status of [
+    "submitted",
+    "streaming",
+    "ready",
+    "submitted",
+    "ready",
+    "submitted",
+    "error",
+    "streaming",
+    undefined,
+  ] as const) {
+    rerender(ui(status));
+    const isGenerating = status === "submitted" || status === "streaming";
+    const keep = screen.getByRole("button", { name: "Keep all" });
+    expect(keep.hasAttribute("disabled")).toBe(isGenerating);
+    expect(
+      screen.queryByRole("button", { name: "Keep all" }),
+    ).toBeInTheDocument();
+    const keepCell = screen.getByRole("button", { name: "Keep" });
+    const discardCell = screen.getByRole("button", { name: "Discard" });
+    for (const button of [keepCell, discardCell]) {
+      expect(button.hasAttribute("disabled")).toBe(isGenerating);
+      if (isGenerating) {
+        fireEvent.click(button);
+      }
     }
-    expect(screen.getByRole("button", { name: "Discard all" })).toBeEnabled();
+    expect(store.get(stagedAICellsAtom).has(pendingCell)).toBe(true);
+    expect(screen.queryByRole("button", { name: "Keep" })).toBeInTheDocument();
+  }
+
+  fireEvent.click(screen.getByRole("button", { name: "Keep all" }));
+  expect(store.get(stagedAICellsAtom).size).toBe(0);
+});
+
+it.each(["Delete cell", "Keep cell"])(
+  "resolves a proposed deletion with %s and advances review",
+  (label) => {
+    deleteCell.mockClear();
+    const store = createStore();
+    const first = cellId("deleted-cell");
+    const second = cellId("added-cell");
+    store.set(
+      stagedAICellsAtom,
+      new Map([
+        [first, { type: "delete_cell", previousCode: "x = 1" }],
+        [second, { type: "add_cell" }],
+      ]),
+    );
+    render(
+      <Provider store={store}>
+        <PendingAICells />
+      </Provider>,
+    );
+    fireEvent.click(screen.getByRole("button", { name: label }));
+    expect(store.get(stagedAICellsAtom).has(first)).toBe(false);
+    expect(deleteCell).toHaveBeenCalledTimes(label === "Delete cell" ? 1 : 0);
+    expect(scrollToCell).toHaveBeenLastCalledWith(second);
+    expect(screen.getByText("1 cell changed")).toBeInTheDocument();
+    expect(
+      screen.queryByRole("button", { name: "Keep all" }),
+    ).not.toBeInTheDocument();
+    expect(
+      screen.queryByRole("button", { name: "Discard all" }),
+    ).not.toBeInTheDocument();
+    fireEvent.click(screen.getByRole("button", { name: "Keep" }));
+    expect(
+      screen.queryByRole("region", { name: "Review AI changes" }),
+    ).not.toBeInTheDocument();
+  },
+);
+
+it("keeps batch discard visible and locks it during generation", () => {
+  deleteCell.mockClear();
+  const store = createStore();
+  const pendingCell = cellId("discarded-cell");
+  store.set(
+    stagedAICellsAtom,
+    new Map([
+      [pendingCell, { type: "add_cell" }],
+      [cellId("other-pending-cell"), { type: "add_cell" }],
+    ]),
+  );
+  const ui = (status: ChatStatus) => (
+    <Provider store={store}>
+      <ChatGeneration status={status} />
+      <PendingAICells />
+    </Provider>
+  );
+  const { rerender } = render(ui("streaming"));
+  const discard = screen.getByRole("button", { name: "Discard all" });
+  expect(discard).toBeVisible();
+  expect(discard).toBeDisabled();
+  fireEvent.click(discard);
+  expect(store.get(stagedAICellsAtom).has(pendingCell)).toBe(true);
+  rerender(ui("ready"));
+  expect(discard).toBeEnabled();
+  fireEvent.click(discard);
+  expect(deleteCell).toHaveBeenCalledTimes(2);
+  expect(deleteCell).toHaveBeenCalledWith({ cellId: pendingCell });
+  expect(store.get(stagedAICellsAtom).size).toBe(0);
+});
+
+it("clears the panel badge when all pending changes are accepted", () => {
+  const store = createStore();
+  store.set(
+    stagedAICellsAtom,
+    new Map([[cellId("pending"), { type: "add_cell" }]]),
+  );
+  render(
+    <Provider store={store}>
+      <PendingAIChangesBadge />
+      <PendingAICells />
+    </Provider>,
+  );
+  expect(screen.getByText("1 pending AI change")).toBeInTheDocument();
+  fireEvent.click(screen.getByRole("button", { name: "Keep" }));
+  expect(screen.queryByText("1 pending AI change")).not.toBeInTheDocument();
+});
+
+it("discards an edit to a cell whose editor is not mounted", () => {
+  setCellCode.mockClear();
+  const store = createStore();
+  const collapsed = cellId("collapsed-cell");
+  store.set(
+    stagedAICellsAtom,
+    new Map([[collapsed, { type: "update_cell", previousCode: "x = 1" }]]),
+  );
+  render(
+    <Provider store={store}>
+      <PendingAICells />
+    </Provider>,
+  );
+  fireEvent.click(screen.getByRole("button", { name: "Discard" }));
+  expect(setCellCode).toHaveBeenCalledWith({
+    cellId: collapsed,
+    code: "x = 1",
+  });
+  expect(store.get(stagedAICellsAtom).size).toBe(0);
+});
+
+it("outlines only the reviewed cell, moving the outline with the selection", () => {
+  const store = createStore();
+  const first = cellId("outlined-first");
+  const second = cellId("outlined-second");
+  store.set(
+    stagedAICellsAtom,
+    new Map([
+      [first, { type: "add_cell" }],
+      [second, { type: "add_cell" }],
+    ]),
+  );
+  const cells = [first, second].map((id) => {
+    const element = document.createElement("div");
+    element.id = HTMLCellId.create(id);
+    document.body.append(element);
+    return element;
   });
+  const { unmount } = render(
+    <Provider store={store}>
+      <TooltipProvider>
+        <PendingAICells />
+      </TooltipProvider>
+    </Provider>,
+  );
+  const outlined = () =>
+    cells.map((cell) => cell.hasAttribute("data-ai-review-selected"));
+
+  expect(outlined()).toEqual([false, false]);
+  fireEvent.click(screen.getByRole("button", { name: "Cell 1 of 2" }));
+  expect(outlined()).toEqual([true, false]);
+  fireEvent.click(screen.getByRole("button", { name: "Next pending AI cell" }));
+  expect(outlined()).toEqual([false, true]);
+
+  unmount();
+  expect(outlined()).toEqual([false, false]);
+  for (const cell of cells) {
+    cell.remove();
+  }
 });

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,
+    });
+  },
+);