@marimo-team/islands 0.25.3-dev129 → 0.25.3-dev134

package/src/plugins/impl/data-frames/__tests__/sidebar.test.tsx

@@ -0,0 +1,160 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, fireEvent, render, screen } from "@testing-library/react";
+import { useState } from "react";
+import { expect, it, vi } from "vitest";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { Sidebar, TransformPanel } from "../panel";
+import { type TransformType, TransformTypeSchema } from "../schema";
+
+const transform: TransformType = { type: "select_columns", column_ids: [] };
+
+it("selects and deletes transforms through independent buttons", () => {
+  const onSelect = vi.fn();
+  const onDelete = vi.fn();
+  render(
+    <Sidebar
+      items={[{ id: "first", transform }]}
+      selected={0}
+      onSelect={onSelect}
+      onDelete={onDelete}
+      onAdd={vi.fn()}
+    />,
+  );
+  const select = screen.getByRole("button", { name: "Select Columns" });
+  expect(select).toHaveAttribute("aria-pressed", "true");
+  fireEvent.click(select);
+  expect(onSelect).toHaveBeenCalledExactlyOnceWith(0);
+  const remove = screen.getByRole("button", {
+    name: "Delete Select Columns transform 1",
+  });
+  remove.focus();
+  expect(remove).toHaveFocus();
+  fireEvent.click(remove);
+  expect(onDelete).toHaveBeenCalledExactlyOnceWith(0);
+  expect(onSelect).toHaveBeenCalledOnce();
+});
+
+it("preserves focused rows across value edits and removal of earlier duplicates", () => {
+  const first = { id: "first", transform };
+  const second = { id: "second", transform };
+  const props = {
+    selected: 1,
+    onSelect: vi.fn(),
+    onDelete: vi.fn(),
+    onAdd: vi.fn(),
+  };
+  const { rerender } = render(<Sidebar {...props} items={[first, second]} />);
+  const focused = screen.getAllByRole("button", { name: "Select Columns" })[1];
+  act(() => focused.focus());
+  const edited = {
+    ...second,
+    transform: TransformTypeSchema.parse({
+      ...transform,
+      column_ids: ["name"],
+    }),
+  };
+  rerender(<Sidebar {...props} items={[first, edited]} />);
+  expect(screen.getAllByRole("button", { name: "Select Columns" })[1]).toBe(
+    focused,
+  );
+  expect(focused).toHaveFocus();
+
+  rerender(<Sidebar {...props} items={[edited]} selected={0} />);
+  expect(screen.getByRole("button", { name: "Select Columns" })).toBe(focused);
+  expect(focused).toHaveFocus();
+  fireEvent.click(focused);
+  expect(props.onSelect).toHaveBeenLastCalledWith(0);
+});
+
+it.each([
+  { count: 3, removed: 0, destination: 1 },
+  { count: 3, removed: 1, destination: 0 },
+  { count: 3, removed: 2, destination: 1 },
+  { count: 1, removed: 0, destination: undefined },
+])(
+  "moves focus after deleting row $removed of $count",
+  ({ count, removed, destination }) => {
+    function Example() {
+      const [items, setItems] = useState(
+        Array.from({ length: count }, (_, index) => ({
+          id: `transform-${index}`,
+          transform,
+        })),
+      );
+      return (
+        <Sidebar
+          items={items}
+          selected={undefined}
+          onSelect={vi.fn()}
+          onDelete={(index) =>
+            setItems((current) => current.filter((_, i) => i !== index))
+          }
+          onAdd={vi.fn()}
+        />
+      );
+    }
+    render(<Example />);
+    const selectButtons = screen.getAllByRole("button", {
+      name: "Select Columns",
+    });
+    const target =
+      destination === undefined
+        ? screen.getByRole("button", { name: "Add" })
+        : selectButtons[destination];
+    const remove = screen.getByRole("button", {
+      name: `Delete Select Columns transform ${removed + 1}`,
+    });
+    act(() => remove.focus());
+    fireEvent.click(remove, { detail: 0 });
+    expect(remove).not.toBeInTheDocument();
+    expect(target).toHaveFocus();
+  },
+);
+
+it("uses persistent field-array IDs in the transformation panel", () => {
+  render(
+    <TooltipProvider>
+      <TransformPanel
+        initialValue={{ transforms: [transform, transform] }}
+        columns={new Map()}
+        lazy={true}
+        onChange={vi.fn()}
+        onInvalidChange={vi.fn()}
+        getColumnValues={async () => ({ values: [], too_many_values: false })}
+      />
+    </TooltipProvider>,
+  );
+  const survivor = screen.getAllByRole("button", { name: "Select Columns" })[1];
+  const remove = screen.getByRole("button", {
+    name: "Delete Select Columns transform 1",
+  });
+  act(() => remove.focus());
+  fireEvent.click(remove, { detail: 0 });
+  expect(screen.getByRole("button", { name: "Select Columns" })).toBe(survivor);
+  expect(survivor).toHaveFocus();
+  expect(survivor).toHaveAttribute("aria-pressed", "true");
+});
+
+it("does not move focus when deletion starts outside the focused row", () => {
+  render(
+    <>
+      <input aria-label="Transform setting" />
+      <Sidebar
+        items={[{ id: "first", transform }]}
+        selected={0}
+        onSelect={vi.fn()}
+        onDelete={vi.fn()}
+        onAdd={vi.fn()}
+      />
+    </>,
+  );
+  const input = screen.getByRole("textbox", { name: "Transform setting" });
+  act(() => input.focus());
+  fireEvent.click(
+    screen.getByRole("button", {
+      name: "Delete Select Columns transform 1",
+    }),
+  );
+  expect(input).toHaveFocus();
+});

package/src/plugins/impl/data-frames/panel.tsx

@@ -25,6 +25,7 @@ import React, {
   useEffect,
   useImperativeHandle,
   useMemo,
+  useRef,
 } from "react";
 import { useFieldArray, useForm } from "react-hook-form";
 import useEvent from "react-use-event-hook";
@@ -187,7 +188,10 @@ export const TransformPanel: React.FC<Props> = ({
           className="relative flex flex-row max-h-[400px] overflow-hidden bg-background"
         >
           <Sidebar
-            items={form.watch("transforms")}
+            items={transformsField.fields.map((field, index) => ({
+              id: field.id,
+              transform: transforms[index],
+            }))}
             selected={selectedTransform}
             onSelect={(index) => {
               setSelectedTransform(index);
@@ -274,7 +278,7 @@ export const TransformPanel: React.FC<Props> = ({
 };
 
 interface SidebarProps {
-  items: TransformType[];
+  items: { id: string; transform: TransformType }[];
   selected: number | undefined;
   onSelect: (index: number) => void;
   onDelete: (index: number) => void;
@@ -288,34 +292,56 @@ export const Sidebar: React.FC<SidebarProps> = ({
   onSelect,
   onDelete,
 }) => {
+  const selectionButtons = useRef(new Map<string, HTMLButtonElement>());
+  const addButton = useRef<HTMLButtonElement>(null);
+
   return (
     <div className="flex flex-col overflow-y-hidden w-[180px] shadow-xs h-full">
       <div className="flex flex-col overflow-y-auto grow">
-        {items.map((item, idx) => {
+        {items.map(({ id, transform }, idx) => {
           return (
             <div
-              key={`${JSON.stringify(item)}-${idx}`}
-              onClick={() => {
-                onSelect(idx);
-              }}
+              key={id}
               className={cn(
-                "flex flex-row min-h-[40px] items-center px-2 cursor-pointer hover:bg-accent/50 text-sm overflow-hidden hover-actions-parent border border-muted border-l-2 border-l-transparent",
+                "group flex flex-row min-h-[40px] items-center px-2 cursor-pointer hover:bg-accent/50 text-sm overflow-hidden border border-muted border-l-2 border-l-transparent",
                 {
                   "border-l-primary bg-accent text-accent-foreground":
                     selected === idx,
                 },
               )}
             >
-              <div className="grow text-ellipsis">
-                {Strings.startCase(item.type)}
-              </div>
-              <Trash2Icon
-                className="w-3 h-3 hover-action text-muted-foreground hover:text-destructive"
-                onClick={(e) => {
+              <button
+                ref={(button) => {
+                  if (button) {
+                    selectionButtons.current.set(id, button);
+                  } else {
+                    selectionButtons.current.delete(id);
+                  }
+                }}
+                type="button"
+                className="grow text-left text-ellipsis py-2 focus-visible:outline focus-visible:outline-ring"
+                aria-pressed={selected === idx}
+                onClick={() => onSelect(idx)}
+              >
+                {Strings.startCase(transform.type)}
+              </button>
+              <button
+                type="button"
+                aria-label={`Delete ${Strings.startCase(transform.type)} transform ${idx + 1}`}
+                className="opacity-0 group-hover:opacity-100 focus-visible:opacity-100 text-muted-foreground hover:text-destructive focus-visible:outline focus-visible:outline-ring"
+                onClick={(event) => {
+                  if (event.currentTarget.matches(":focus")) {
+                    const neighbor = items[idx - 1] ?? items[idx + 1];
+                    const target = neighbor
+                      ? selectionButtons.current.get(neighbor.id)
+                      : addButton.current;
+                    target?.focus();
+                  }
                   onDelete(idx);
-                  e.stopPropagation();
                 }}
-              />
+              >
+                <Trash2Icon className="w-3 h-3" />
+              </button>
             </div>
           );
         })}
@@ -323,6 +349,7 @@ export const Sidebar: React.FC<SidebarProps> = ({
       <div className="flex flex-row shrink-0">
         <AddTransformDropdown onAdd={onAdd}>
           <Button
+            ref={addButton}
             data-testid="marimo-plugin-data-frames-add-transform"
             variant="text"
             className="w-full rounded-none m-0 hover:text-accent-foreground"

package/src/core/websocket/__tests__/useNotebookClient.hook.test.tsx

@@ -1,228 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-// @vitest-environment jsdom
-
-import { act, renderHook } from "@testing-library/react";
-import { createStore, Provider } from "jotai";
-import type { PropsWithChildren } from "react";
-import { beforeEach, describe, expect, it, vi } from "vitest";
-import type { NotebookTransport } from "@marimo-team/client";
-import { ConnectionState } from "../types";
-import { Deferred } from "@/utils/Deferred";
-import type { NotebookBackend } from "../../host/backend";
-import { createEditorBackend } from "../backend";
-import { useNotebookClient } from "../useNotebookClient";
-import { fixtureMessages } from "../../host/__tests__/fixtures";
-import { agentAttachmentsAtom } from "../../attachments/state";
-import { notebookClientAtom } from "../state";
-
-vi.mock("../backend", () => ({ createEditorBackend: vi.fn() }));
-
-/** A backend that follows `transport` and cannot edit. */
-function backend(transport: NotebookTransport): NotebookBackend {
-  return { transport, notebook: () => ({ edit: vi.fn() }) };
-}
-
-/** A transport whose connections the test drives. */
-function makeTransport() {
-  const opened: {
-    signal: AbortSignal;
-    open(): void;
-    receive(message: unknown): void;
-  }[] = [];
-  const transport: NotebookTransport = ({ signal }, handlers) =>
-    new Promise((resolve) => {
-      opened.push({ signal, open: handlers.open, receive: handlers.message });
-      signal.addEventListener("abort", () => resolve(), { once: true });
-    });
-  return {
-    transport,
-    opened,
-    /** Open the latest connection and deliver a host-framed notification. */
-    notify: () =>
-      opened.at(-1)?.receive({
-        event: "marimo",
-        data: { marimo: { op: "reconnected" } },
-      }),
-  };
-}
-
-function setup(
-  strict = false,
-  readiness: {
-    waitUntilReady?: () => Promise<void>;
-    checkHealth?: () => Promise<boolean>;
-  } = {},
-) {
-  const store = createStore();
-  const options = {
-    static: false,
-    transportType: "websocket" as const,
-    url: () => "ws://localhost/ws",
-    headers: () => ({}),
-    notebookUrl: (id: string) => new URL(`http://localhost/notebooks/${id}`),
-    waitUntilReady: async () => undefined,
-    checkHealth: async () => true,
-    onNotification: vi.fn(),
-    ...readiness,
-  };
-  const wrapper = ({ children }: PropsWithChildren) => (
-    <Provider store={store}>{children}</Provider>
-  );
-  const hook = renderHook((props) => useNotebookClient(props), {
-    initialProps: options,
-    reactStrictMode: strict,
-    wrapper,
-  });
-  return { ...hook, options, store };
-}
-
-beforeEach(() => {
-  vi.mocked(createEditorBackend).mockReset();
-});
-
-describe("useNotebookClient", () => {
-  it.each([true, false])(
-    "rechecks health after failed startup (healthy: %s)",
-    async (healthy) => {
-      const fake = makeTransport();
-      vi.mocked(createEditorBackend).mockImplementation(({ url }) =>
-        backend(async (attempt, handlers) => {
-          if (typeof url === "function") {
-            await url();
-          }
-          if (!attempt.signal.aborted) {
-            return fake.transport(attempt, handlers);
-          }
-        }),
-      );
-      const waitUntilReady = vi
-        .fn()
-        .mockRejectedValue(new Error("Startup gave up"));
-      const checkHealth = vi.fn().mockResolvedValue(healthy);
-      const { result, store, unmount } = setup(false, {
-        waitUntilReady,
-        checkHealth,
-      });
-      await act(async () => Promise.resolve());
-      expect(store.get(result.current.connectionAtom).state).toBe(
-        ConnectionState.CLOSED,
-      );
-
-      await act(async () => result.current.reconnect());
-      expect({
-        healthChecks: checkHealth.mock.calls.length,
-        initialWaits: waitUntilReady.mock.calls.length,
-        connections: fake.opened.length,
-        state: store.get(result.current.connectionAtom).state,
-      }).toEqual({
-        healthChecks: 1,
-        initialWaits: 1,
-        connections: healthy ? 1 : 0,
-        state: healthy ? ConnectionState.CONNECTING : ConnectionState.CLOSED,
-      });
-      act(() => fake.opened.at(-1)?.open());
-      expect(store.get(result.current.connectionAtom).state).toBe(
-        healthy ? ConnectionState.OPEN : ConnectionState.CLOSED,
-      );
-      unmount();
-    },
-  );
-
-  it("waits for initial readiness without reviving an unmounted connection", async () => {
-    const ready = new Deferred<void>();
-    const resolvedUrl = vi.fn();
-    vi.mocked(createEditorBackend).mockImplementation(({ url }) =>
-      backend(async () => {
-        if (typeof url === "function") {
-          await url();
-          resolvedUrl();
-        }
-      }),
-    );
-    const { result, store, unmount } = setup(false, {
-      waitUntilReady: () => ready.promise,
-    });
-    await act(async () => Promise.resolve());
-    expect(resolvedUrl).not.toHaveBeenCalled();
-    expect(store.get(result.current.connectionAtom)).toEqual({
-      state: ConnectionState.CONNECTING,
-    });
-    unmount();
-    await act(async () => ready.resolve());
-    expect(store.get(result.current.connectionAtom)).toEqual({
-      state: ConnectionState.NOT_STARTED,
-    });
-    expect(store.get(notebookClientAtom)).toBeNull();
-  });
-
-  it("exposes host presence to the UI without waiting for kernel-ready", async () => {
-    const fake = makeTransport();
-    vi.mocked(createEditorBackend).mockReturnValue(backend(fake.transport));
-    const { result, store, unmount } = setup();
-    await act(async () => Promise.resolve());
-    expect(store.get(notebookClientAtom)).toBe(result.current);
-    const [connection] = fake.opened;
-    act(() => {
-      connection.open();
-      fixtureMessages("notebook-events.sse").forEach(connection.receive);
-    });
-    expect(store.get(agentAttachmentsAtom)).toEqual([
-      {
-        id: "a2",
-        notebook_id: "n1",
-        kind: "agent",
-        name: "Notebook CLI",
-        since: "2026-10-01T09:20:03Z",
-      },
-    ]);
-    unmount();
-    expect(store.get(notebookClientAtom)).toBeNull();
-    expect(store.get(agentAttachmentsAtom)).toEqual([]);
-  });
-
-  it("uses current callbacks and URL providers without reconnecting on a rerender", async () => {
-    const fake = makeTransport();
-    vi.mocked(createEditorBackend).mockReturnValue(backend(fake.transport));
-    const { rerender, options, unmount } = setup();
-    await act(async () => Promise.resolve());
-    const onNotification = vi.fn();
-    rerender({ ...options, onNotification, url: () => "ws://redirected/ws" });
-    act(() => fake.notify());
-    expect(options.onNotification).not.toHaveBeenCalled();
-    expect(onNotification).toHaveBeenCalledExactlyOnceWith({
-      data: { op: "reconnected" },
-    });
-    expect(createEditorBackend).toHaveBeenCalledOnce();
-    const [{ url }] = vi.mocked(createEditorBackend).mock.calls[0];
-    expect(String(await (url as () => Promise<string>)())).toBe(
-      "ws://redirected/ws",
-    );
-    unmount();
-  });
-
-  it("closes every stream under Strict Mode and only follows the active lifetime", async () => {
-    const fakes: ReturnType<typeof makeTransport>[] = [];
-    vi.mocked(createEditorBackend).mockImplementation(() => {
-      const next = makeTransport();
-      fakes.push(next);
-      return backend(next.transport);
-    });
-    const { result, store, options, unmount } = setup(true);
-    await act(async () => Promise.resolve());
-    expect(fakes).toHaveLength(2);
-    expect(fakes[0].opened.every((attempt) => attempt.signal.aborted)).toBe(
-      true,
-    );
-    const [active] = fakes[1].opened;
-    act(() => active.open());
-    expect(store.get(result.current.connectionAtom)).toEqual({
-      state: ConnectionState.OPEN,
-    });
-    act(() => fakes.forEach((fake) => fake.notify()));
-    expect(options.onNotification).toHaveBeenCalledOnce();
-    unmount();
-    expect(active.signal.aborted).toBe(true);
-    await result.current.reconnect();
-    expect(fakes[1].opened).toHaveLength(1);
-  });
-});

package/src/core/websocket/useNotebookClient.ts

@@ -1,89 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { StreamStopped } from "@marimo-team/client";
-import { useStore } from "jotai";
-import { useEffect, useMemo } from "react";
-import useEvent from "react-use-event-hook";
-import type {
-  NotificationMessage,
-  NotificationPayload,
-} from "../kernel/messages";
-import {
-  createNotebookClient,
-  type NotebookClient,
-} from "../host/notebook-client";
-import { notebookClientAtom } from "./state";
-import { createEditorBackend, type EditorBackendOptions } from "./backend";
-
-type Options = Omit<EditorBackendOptions, "url"> & {
-  /** The session route, read before each connection. */
-  url: () => string;
-  /** Wait until the host can accept a connection. */
-  waitUntilReady: () => Promise<void>;
-  checkHealth: () => Promise<boolean>;
-  onNotification: (notification: NotificationPayload) => void;
-};
-
-export function useNotebookClient(options: Options): NotebookClient {
-  const store = useStore();
-  const { static: isStatic, transportType } = options;
-  const controller = useMemo(() => createNotebookClient(store), [store]);
-  const url = useEvent(options.url);
-  const headers = useEvent(options.headers);
-  const notebookUrl = useEvent(options.notebookUrl);
-  const waitUntilReady = useEvent(options.waitUntilReady);
-  const checkHealth = useEvent(options.checkHealth);
-  const onNotification = useEvent(options.onNotification);
-
-  useEffect(() => {
-    const lifetime = new AbortController();
-    let needsInitialReadiness = true;
-    store.set(notebookClientAtom, controller.client);
-    controller.connect({
-      backend: createEditorBackend({
-        static: isStatic,
-        transportType,
-        url: async () => {
-          if (needsInitialReadiness) {
-            // Explicit retries check current health, even if startup failed.
-            needsInitialReadiness = false;
-            try {
-              await waitUntilReady();
-            } catch (error) {
-              throw new StreamStopped("Healthy connection never made", {
-                cause: error,
-              });
-            }
-          }
-          return url();
-        },
-        headers,
-        notebookUrl,
-      }),
-      signal: lifetime.signal,
-      checkHealth,
-      // The host validates the envelope, not the native notification contents.
-      onNotification: (notification) =>
-        onNotification({ data: notification as NotificationMessage }),
-    });
-    return () => {
-      lifetime.abort();
-      if (store.get(notebookClientAtom) === controller.client) {
-        store.set(notebookClientAtom, null);
-      }
-    };
-  }, [
-    controller,
-    store,
-    isStatic,
-    transportType,
-    url,
-    headers,
-    notebookUrl,
-    waitUntilReady,
-    checkHealth,
-    onNotification,
-  ]);
-
-  return controller.client;
-}