@marimo-team/islands 0.24.3-dev9 → 0.25.0

package/src/core/packages/__tests__/usePackageAction.test.tsx

@@ -0,0 +1,129 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { toast } from "@/components/ui/use-toast";
+import { requestClientAtom } from "@/core/network/requests";
+import type { PackageOperationResponse } from "@/core/network/types";
+import { usePackageAction } from "../usePackageAction";
+
+vi.mock("@/components/ui/use-toast", () => ({ toast: vi.fn() }));
+
+function setup(action: "upgrade" | "remove") {
+  const client = MockRequestClient.create();
+  const store = createStore();
+  store.set(requestClientAtom, client);
+  return {
+    request: vi.mocked(
+      action === "upgrade" ? client.addPackage : client.removePackage,
+    ),
+    ...renderHook(
+      () =>
+        usePackageAction(action, "numpy", [{ kind: "group", value: "dev" }]),
+      {
+        wrapper: ({ children }) => (
+          <Provider store={store}>{children}</Provider>
+        ),
+      },
+    ),
+  };
+}
+
+describe("usePackageAction", () => {
+  beforeEach(() => vi.clearAllMocks());
+
+  it.each([
+    { action: "upgrade", title: "Package upgraded", upgrade: true },
+    { action: "remove", title: "Package removed" },
+  ] as const)(
+    "preserves dependency groups and prevents duplicate $action requests",
+    async ({ action, title, ...options }) => {
+      const { result, request } = setup(action);
+      const { promise, resolve } =
+        Promise.withResolvers<PackageOperationResponse>();
+      request.mockReturnValue(promise);
+      let pending: Promise<void>;
+      act(() => {
+        pending = result.current.run();
+      });
+      expect(result.current.loading).toBe(true);
+      await act(async () => {
+        await result.current.run();
+      });
+      expect(request).toHaveBeenCalledExactlyOnceWith({
+        package: "numpy",
+        group: "dev",
+        ...options,
+      });
+      await act(async () => {
+        resolve({ success: true });
+        await pending;
+      });
+      expect(result.current.loading).toBe(false);
+      expect(toast).toHaveBeenCalledExactlyOnceWith(
+        expect.objectContaining({ title }),
+      );
+    },
+  );
+
+  it("reports saved changes requiring a restart without claiming success", async () => {
+    const { result, request } = setup("upgrade");
+    request.mockResolvedValue({ success: false, restartRequired: true });
+    await act(async () => {
+      await result.current.run();
+    });
+    expect(toast).toHaveBeenCalledExactlyOnceWith(
+      expect.objectContaining({
+        title: "Changes saved — restart required",
+      }),
+    );
+  });
+
+  it.each([
+    {
+      response: { success: false, error: "resolver failed" },
+      description: "resolver failed",
+    },
+    { response: { success: false }, description: expect.any(String) },
+  ])(
+    "reports failed responses as errors: $response",
+    async ({ response, description }) => {
+      const { result, request } = setup("remove");
+      request.mockResolvedValue(response);
+      await act(async () => {
+        await result.current.run();
+      });
+      expect(toast).toHaveBeenCalledExactlyOnceWith(
+        expect.objectContaining({
+          title: "Failed to remove package",
+          description,
+          variant: "danger",
+        }),
+      );
+    },
+  );
+
+  it("reports request errors and allows a retry", async () => {
+    const { result, request } = setup("remove");
+    request.mockRejectedValueOnce(new Error("offline"));
+    await act(async () => {
+      await result.current.run();
+    });
+    expect(toast).toHaveBeenCalledExactlyOnceWith(
+      expect.objectContaining({
+        title: "Failed to remove package",
+        description: "offline",
+        variant: "danger",
+      }),
+    );
+    expect(result.current.loading).toBe(false);
+    request.mockResolvedValue({ success: true });
+    await act(async () => {
+      await result.current.run();
+    });
+    expect(toast).toHaveBeenLastCalledWith(
+      expect.objectContaining({ title: "Package removed" }),
+    );
+  });
+});

package/src/core/packages/__tests__/useSandboxController.test.tsx

@@ -0,0 +1,246 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { Provider } from "jotai";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MockRequestClient } from "@/__mocks__/requests";
+import {
+  emptyEnvironmentState,
+  type EnvironmentOperation,
+} from "@/core/alerts/environment";
+import { alertAtom } from "@/core/alerts/state";
+import { kernelStartupErrorAtom } from "@/core/errors/state";
+import { connectionAtom } from "@/core/network/connection";
+import { requestClientAtom } from "@/core/network/requests";
+import type {
+  SandboxResponse,
+  SyncSandboxResponse,
+} from "@/core/network/types";
+import { filenameAtom } from "@/core/saving/file-state";
+import { store } from "@/core/state/jotai";
+import { WebSocketState } from "@/core/websocket/types";
+import { HTTPError } from "@/utils/errors";
+import { packageDataVersionAtom } from "../package-data";
+import {
+  sandboxActionsAtom,
+  sandboxAtom,
+  sandboxSyncAtom,
+} from "../sandbox-state";
+import { useSandboxController } from "../useSandboxController";
+
+vi.mock("../toast-components", () => ({
+  showPackageRestartToast: vi.fn(),
+  showSandboxSyncToast: vi.fn(),
+  RESTART_REQUIRED_DESCRIPTION:
+    "Restart the kernel to use the updated environment.",
+}));
+
+const manifest = 'dependencies = ["numpy"]';
+const sandbox: SandboxResponse = {
+  backend: "uv",
+  manifest,
+  filename: "notebook.py",
+};
+
+beforeEach(() => {
+  store.set(alertAtom, (value) => ({
+    ...value,
+    environments: {
+      kernel: emptyEnvironmentState(),
+      server: emptyEnvironmentState(),
+    },
+  }));
+  store.set(sandboxAtom, null);
+  store.set(sandboxSyncAtom, { kind: "succeeded" });
+  store.set(sandboxActionsAtom, null);
+  store.set(connectionAtom, { state: WebSocketState.OPEN });
+  store.set(kernelStartupErrorAtom, null);
+  store.set(filenameAtom, "notebook.py");
+  store.set(packageDataVersionAtom, 0);
+});
+
+function setup() {
+  const client = MockRequestClient.create({
+    getSandbox: vi.fn().mockResolvedValue(sandbox),
+    updateManifest: vi.fn(async ({ contents }) => ({
+      ...sandbox,
+      manifest: contents,
+    })),
+    syncSandbox: vi.fn().mockResolvedValue({ success: true, reconnect: false }),
+  });
+  store.set(requestClientAtom, client);
+  const reconnect = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
+  return {
+    client,
+    reconnect,
+    ...renderHook(() => useSandboxController(reconnect), {
+      wrapper: ({ children }) => <Provider store={store}>{children}</Provider>,
+    }),
+  };
+}
+
+describe("useSandboxController", () => {
+  it("preserves a restart-required response instead of reporting a failed sync", async () => {
+    const { client, result } = setup();
+    vi.mocked(client.syncSandbox).mockResolvedValue({
+      success: false,
+      restartRequired: true,
+      error: "Python version changed",
+    });
+    await act(async () => {
+      expect(await store.get(sandboxActionsAtom)?.sync()).toBe(false);
+    });
+    expect(store.get(sandboxSyncAtom)).toEqual({
+      kind: "restart-required",
+      reason: "Python version changed",
+    });
+    expect(result.current.pending).toBe(false);
+    expect(result.current.diagnostic).toBe("Python version changed");
+  });
+
+  it("restores sync progress and lets a retry replace the restored failure", async () => {
+    const operation: EnvironmentOperation = {
+      operation_id: "sync",
+      action: "sync",
+      source: "kernel",
+      status: { kind: "running" },
+      packages: {},
+      logs: {},
+    };
+    const restore = (status: EnvironmentOperation["status"]) =>
+      store.set(alertAtom, (value) => ({
+        ...value,
+        environments: {
+          ...value.environments,
+          kernel: {
+            restart_required: false,
+            operations: [{ ...operation, status }],
+          },
+        },
+      }));
+    restore({ kind: "running" });
+    const { client, result } = setup();
+    const actions = store.get(sandboxActionsAtom)!;
+    expect(result.current.pending).toBe(true);
+    await act(async () => {
+      expect(await actions.sync()).toBe(false);
+    });
+    expect(client.syncSandbox).not.toHaveBeenCalled();
+    act(() => restore({ kind: "failed", error: "Dependency conflict" }));
+    expect(result.current.diagnostic).toBe("Dependency conflict");
+
+    const { promise, resolve } = Promise.withResolvers<SyncSandboxResponse>();
+    vi.mocked(client.syncSandbox).mockReturnValue(promise);
+    let pending: Promise<boolean>;
+    act(() => {
+      pending = actions.sync();
+    });
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "running" });
+    act(() => restore({ kind: "running" }));
+    await act(async () => {
+      resolve({ success: true, reconnect: false });
+      expect(await pending).toBe(true);
+    });
+    // The response can arrive before the final websocket notification.
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "running" });
+    act(() => restore({ kind: "succeeded" }));
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "succeeded" });
+  });
+
+  it("keeps a conflicting draft across closing and only discards it on explicit reload", async () => {
+    const { result, client } = setup();
+    vi.mocked(client.updateManifest).mockRejectedValue(
+      new HTTPError(409, "Conflict", { detail: "Manifest changed on disk" }),
+    );
+    await act(async () => {
+      await store.get(sandboxActionsAtom)?.editManifest();
+    });
+    act(() => result.current.setDraft("dependencies = []"));
+    await act(async () => {
+      await result.current.saveAndSync();
+    });
+    expect(client.updateManifest).toHaveBeenCalledExactlyOnceWith({
+      fileKey: "notebook.py",
+      previous: manifest,
+      contents: "dependencies = []",
+    });
+    expect(result.current).toMatchObject({
+      open: true,
+      draft: "dependencies = []",
+      dirty: true,
+      conflict: true,
+      diagnostic: "Manifest changed on disk",
+      pending: false,
+    });
+    expect(client.syncSandbox).not.toHaveBeenCalled();
+    act(() => result.current.setOpen(false));
+    await act(async () => {
+      await store.get(sandboxActionsAtom)?.editManifest();
+    });
+    expect(result.current.draft).toBe("dependencies = []");
+    await act(async () => {
+      await result.current.loadManifest();
+    });
+    expect(result.current).toMatchObject({
+      draft: manifest,
+      dirty: false,
+      conflict: false,
+      diagnostic: null,
+    });
+  });
+
+  it("shares sync progress across callers and refreshes packages without reconnecting", async () => {
+    const { client, reconnect } = setup();
+    const { promise, resolve } = Promise.withResolvers<SyncSandboxResponse>();
+    vi.mocked(client.syncSandbox).mockReturnValue(promise);
+    const actions = store.get(sandboxActionsAtom)!;
+    let pending: Promise<boolean>;
+    act(() => {
+      pending = actions.sync();
+    });
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "running" });
+    await act(async () => {
+      expect(await actions.sync()).toBe(false);
+    });
+    expect(client.syncSandbox).toHaveBeenCalledTimes(1);
+    await act(async () => {
+      resolve({ success: true, reconnect: false });
+      expect(await pending).toBe(true);
+    });
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "succeeded" });
+    expect(store.get(packageDataVersionAtom)).toBe(1);
+    expect(reconnect).not.toHaveBeenCalled();
+  });
+
+  it("reports sync failure and clears it when retrying", async () => {
+    const { client } = setup();
+    const actions = store.get(sandboxActionsAtom)!;
+    vi.mocked(client.syncSandbox).mockRejectedValueOnce(new Error("offline"));
+    await act(async () => {
+      expect(await actions.sync()).toBe(false);
+    });
+    expect(store.get(sandboxSyncAtom)).toEqual({
+      kind: "failed",
+      error: "offline",
+    });
+    await act(async () => {
+      expect(await actions.sync()).toBe(true);
+    });
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "succeeded" });
+  });
+
+  it("ignores stale metadata after a filename change and unregisters actions on unmount", async () => {
+    const { client, unmount } = setup();
+    await waitFor(() => expect(store.get(sandboxAtom)).toEqual(sandbox));
+    const { promise, resolve } = Promise.withResolvers<SandboxResponse>();
+    vi.mocked(client.getSandbox).mockReturnValueOnce(promise);
+    act(() => store.set(filenameAtom, "old.py"));
+    const current = { ...sandbox, filename: "current.py" };
+    vi.mocked(client.getSandbox).mockResolvedValue(current);
+    act(() => store.set(filenameAtom, "current.py"));
+    await waitFor(() => expect(store.get(sandboxAtom)).toEqual(current));
+    await act(async () => resolve({ ...sandbox, filename: "old.py" }));
+    expect(store.get(sandboxAtom)).toEqual(current);
+    unmount();
+    expect(store.get(sandboxActionsAtom)).toBeNull();
+  });
+});

package/src/core/packages/package-data.ts

@@ -0,0 +1,35 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { atom } from "jotai";
+import type {
+  EditRequests,
+  PackageOperationResponse,
+  RunRequests,
+} from "@/core/network/types";
+import { store } from "@/core/state/jotai";
+
+export const packageDataVersionAtom = atom(0);
+export function invalidatePackageData() {
+  store.set(packageDataVersionAtom, (version) => version + 1);
+}
+
+/** Refresh package consumers regardless of which UI initiated a mutation. */
+export function withPackageInvalidation(
+  client: EditRequests & RunRequests,
+): EditRequests & RunRequests {
+  function invalidateAfter<Request, Response extends PackageOperationResponse>(
+    request: (body: Request) => Promise<Response>,
+  ) {
+    return async (body: Request) => {
+      const result = await request(body);
+      if (result.success || result.restartRequired) {
+        invalidatePackageData();
+      }
+      return result;
+    };
+  }
+  return {
+    ...client,
+    addPackage: invalidateAfter(client.addPackage),
+    removePackage: invalidateAfter(client.removePackage),
+  };
+}

package/src/core/packages/sandbox-state.ts

@@ -0,0 +1,67 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { atom } from "jotai";
+import type { EnvironmentOperation } from "@/core/alerts/environment";
+import { alertAtom } from "@/core/alerts/state";
+import type { SandboxResponse } from "@/core/network/types";
+
+export const sandboxAtom = atom<SandboxResponse | null>(null);
+
+export const preparationAtom = atom((get) =>
+  get(alertAtom).environments.kernel.operations.findLast(
+    (operation) => operation.action === "prepare",
+  ),
+);
+
+type SandboxSyncState = EnvironmentOperation["status"] | { kind: "idle" };
+
+export const sandboxSyncOperationAtom = atom((get) => {
+  const operations = get(alertAtom).environments.kernel.operations;
+  return (
+    operations.findLast(
+      (operation) =>
+        operation.action === "sync" && operation.status.kind === "running",
+    ) ?? operations.findLast((operation) => operation.action === "sync")
+  );
+});
+
+const sandboxSyncRequestAtom = atom<{
+  state: SandboxSyncState;
+  operation: EnvironmentOperation | undefined;
+} | null>(null);
+
+export const sandboxSyncAtom = atom(
+  (get): SandboxSyncState => {
+    const operation = get(sandboxSyncOperationAtom);
+    const request = get(sandboxSyncRequestAtom);
+    // Local request status lasts until newer progress arrives from the server.
+    const state: SandboxSyncState =
+      request &&
+      request.operation === operation &&
+      operation?.status.kind !== "running"
+        ? request.state
+        : (operation?.status ?? { kind: "idle" });
+    // Later package operations can replace the sync result, but only a new
+    // environment snapshot can clear an outstanding restart.
+    if (
+      get(alertAtom).environments.kernel.restart_required &&
+      (state.kind === "idle" || state.kind === "succeeded")
+    ) {
+      return {
+        kind: "restart-required",
+        reason:
+          "Dependency changes are saved; restart the kernel to apply them.",
+      };
+    }
+    return state;
+  },
+  (get, set, state: SandboxSyncState) => {
+    set(sandboxSyncRequestAtom, {
+      state,
+      operation: get(sandboxSyncOperationAtom),
+    });
+  },
+);
+export const sandboxActionsAtom = atom<{
+  sync: () => Promise<boolean>;
+  editManifest: () => Promise<void>;
+} | null>(null);

package/src/core/packages/toast-components.tsx

@@ -29,6 +29,17 @@ export const showPackageRestartToast = () => {
   });
 };
 
+export const showSandboxSyncToast = () => {
+  toast({
+    title: "Sandbox environment synced",
+    description: (
+      <div className="text-xs text-muted-foreground">
+        Some Python packages may require a kernel restart to see changes.
+      </div>
+    ),
+  });
+};
+
 export const showAddPackageToast = (
   packageName: string | string[],
   error?: string | null,

package/src/core/packages/useInstallPackage.ts

@@ -37,10 +37,10 @@ export function useInstallPackages(): {
         showPackageRestartToast();
       } else if (response.success) {
         showAddPackageToast(packages);
+        onSuccess?.();
       } else {
         showAddPackageToast(packages, response.error);
       }
-      onSuccess?.();
     } catch (error) {
       Logger.error(error);
     } finally {

package/src/core/packages/usePackageAction.ts

@@ -0,0 +1,51 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { useRef, useState } from "react";
+import useEvent from "react-use-event-hook";
+import { useRequestClient } from "@/core/network/requests";
+import type { DependencyTreeNode } from "@/core/network/types";
+import { prettyError } from "@/utils/errors";
+import {
+  showPackageRestartToast,
+  showRemovePackageToast,
+  showUpgradePackageToast,
+} from "./toast-components";
+
+/** Shared mutation lifecycle for package actions in both list and tree views. */
+export function usePackageAction(
+  action: "upgrade" | "remove",
+  packageName: string,
+  tags?: DependencyTreeNode["tags"],
+) {
+  const { addPackage, removePackage } = useRequestClient();
+  const [loading, setLoading] = useState(false);
+  const inFlight = useRef(false);
+  const run = useEvent(async () => {
+    if (inFlight.current) {
+      return;
+    }
+    inFlight.current = true;
+    setLoading(true);
+    const showResult =
+      action === "upgrade" ? showUpgradePackageToast : showRemovePackageToast;
+    try {
+      const group = tags?.find((tag) => tag.kind === "group")?.value;
+      const request = { package: packageName, group };
+      const response = await (action === "upgrade"
+        ? addPackage({ ...request, upgrade: true })
+        : removePackage(request));
+      if (response.restartRequired) {
+        showPackageRestartToast();
+      } else if (response.success) {
+        showResult(packageName);
+      } else {
+        showResult(packageName, response.error || "Package operation failed.");
+      }
+    } catch (error) {
+      showResult(packageName, prettyError(error));
+    } finally {
+      inFlight.current = false;
+      setLoading(false);
+    }
+  });
+  return { loading, run };
+}

package/src/core/packages/usePackageDependencies.ts

@@ -0,0 +1,22 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { useAtomValue } from "jotai";
+import { useResolvedMarimoConfig } from "@/core/config/config";
+import { useRequestClient } from "@/core/network/requests";
+import { useAsyncData } from "@/hooks/useAsyncData";
+import { packageDataVersionAtom } from "./package-data";
+
+/** Load the available package views and refresh them after package mutations. */
+export function usePackageDependencies() {
+  const packageDataVersion = useAtomValue(packageDataVersionAtom);
+  const [config] = useResolvedMarimoConfig();
+  const packageManager = config.package_management.manager;
+  const { getDependencyTree, getPackageList } = useRequestClient();
+
+  return useAsyncData(async () => {
+    const { context, tree } = await getDependencyTree();
+    // Both sandbox endpoints inspect the same environment; only fetch it once.
+    const list =
+      context.kind === "sandbox" ? [] : (await getPackageList()).packages;
+    return { list, context, tree };
+  }, [getDependencyTree, getPackageList, packageManager, packageDataVersion]);
+}