@marimo-team/islands 0.24.3-dev81 → 0.24.3-dev83

package/src/core/network/connection-notice.ts

@@ -2,40 +2,69 @@
 import { atom } from "jotai";
 import { kernelStartupErrorAtom } from "@/core/errors/state";
 import { sandboxAtom, sandboxSyncAtom } from "@/core/packages/sandbox-state";
-import { WebSocketClosedReason, WebSocketState } from "@/core/websocket/types";
+import {
+  type ConnectionPhase,
+  WebSocketClosedReason,
+  WebSocketState,
+} from "@/core/websocket/types";
 import { connectionAtom } from "./connection";
 
-export const connectionNoticeAtom = atom((get) => {
+export interface ConnectionNotice {
+  kind: "startup" | "sync" | "connection";
+  phase?: ConnectionPhase;
+  sandbox: boolean;
+  title: string;
+  description: string;
+  pending: boolean;
+  error: string | null;
+}
+
+export const connectionNoticeAtom = atom<ConnectionNotice | null>((get) => {
   const connection = get(connectionAtom);
   const backend = get(sandboxAtom)?.backend;
   const sandbox = Boolean(backend);
   const sync = get(sandboxSyncAtom);
+  const syncing = sync.kind === "running";
+  const restartRequired = sync.kind === "restart-required";
+  const syncError =
+    sync.kind === "failed"
+      ? sync.error
+      : sync.kind === "cancelled"
+        ? "Sandbox sync was interrupted."
+        : restartRequired
+          ? sync.reason
+          : null;
   if (
     sandbox &&
-    (sync.pending || sync.error) &&
+    (syncing || syncError !== null) &&
     connection.state === WebSocketState.OPEN
   ) {
     return {
+      kind: "sync",
       sandbox,
-      title: sync.pending ? "Syncing sandbox…" : "Sandbox sync failed",
-      description: sync.pending
+      title: syncing
+        ? "Syncing sandbox…"
+        : restartRequired
+          ? "Sandbox restart required"
+          : "Sandbox sync failed",
+      description: syncing
         ? "Applying the saved manifest to this notebook’s environment."
-        : "Review the error details, then try again.",
-      pending: sync.pending,
-      error: sync.error,
+        : restartRequired
+          ? "Restart the kernel to use the updated environment."
+          : "Review the error details, then try again.",
+      pending: syncing,
+      error: syncError,
     };
   }
   if (connection.state === WebSocketState.CONNECTING) {
     const messages = {
       "preparing-environment": {
-        title: backend
-          ? `Preparing ${backend} sandbox…`
-          : "Preparing environment…",
+        title: "Preparing environment…",
         description:
           "Getting this notebook’s Python environment ready. The first start can take a few minutes.",
       },
       "starting-kernel": {
-        title: "Starting notebook…",
+        title: "Starting kernel…",
         description:
           "The environment is prepared. Waiting for the Python kernel to connect to this notebook.",
       },
@@ -46,6 +75,8 @@ export const connectionNoticeAtom = atom((get) => {
       },
     };
     return {
+      kind: connection.phase === "reconnecting" ? "connection" : "startup",
+      phase: connection.phase,
       sandbox,
       ...(connection.phase
         ? messages[connection.phase]
@@ -63,6 +94,8 @@ export const connectionNoticeAtom = atom((get) => {
     const environmentFailed =
       startupFailed && connection.phase === "preparing-environment";
     return {
+      kind: startupFailed ? "startup" : "connection",
+      phase: connection.phase,
       sandbox,
       title: environmentFailed
         ? "Sandbox setup failed"

package/src/core/network/connection.ts

@@ -1,5 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { atom } from "jotai";
+import { atomWithReducer } from "jotai/utils";
+import type { NotificationMessageData } from "../kernel/messages";
 import { isNotebookPage } from "../mode";
 import { waitFor } from "../state/jotai";
 import { type ConnectionStatus, WebSocketState } from "../websocket/types";
@@ -12,6 +14,23 @@ export const connectionAtom = atom<ConnectionStatus>({
   state: WebSocketState.NOT_STARTED,
 });
 
+type StartupProgress = NotificationMessageData<"startup-progress">;
+
+export const startupProgressAtom = atomWithReducer<
+  Pick<StartupProgress, "phase" | "logs"> | null,
+  StartupProgress | null
+>(null, (previous, update: StartupProgress | null) =>
+  update === null
+    ? null
+    : {
+        phase: update.phase,
+        logs:
+          update.log_mode === "append" && previous?.phase === update.phase
+            ? previous.logs + update.logs
+            : update.logs,
+      },
+);
+
 export function waitForConnectionOpen() {
   return waitFor(connectionAtom, (value) => {
     return value.state === WebSocketState.OPEN;

package/src/core/network/useConnectionNotice.ts

@@ -0,0 +1,53 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { useAtomValue } from "jotai";
+import { useEffect, useState } from "react";
+import { useDelayElapsed } from "@/hooks/useDelayElapsed";
+import { isConnectedAtom } from "./connection";
+import {
+  type ConnectionNotice,
+  connectionNoticeAtom,
+} from "./connection-notice";
+
+export type DisplayConnectionNotice = ConnectionNotice & { ready: boolean };
+
+/** Briefly acknowledge completion, but only for setup the user actually saw. */
+export function useConnectionNotice(
+  delayMs = 500,
+): DisplayConnectionNotice | null {
+  const notice = useAtomValue(connectionNoticeAtom);
+  const connected = useAtomValue(isConnectedAtom);
+  const elapsed = useDelayElapsed(notice?.pending ? delayMs : 0);
+  const [previous, setPrevious] = useState<ConnectionNotice | null>(null);
+
+  useEffect(() => {
+    if (notice || !connected) {
+      setPrevious(elapsed ? notice : null);
+      return;
+    }
+    const timeout = setTimeout(() => setPrevious(null), 1500);
+    return () => clearTimeout(timeout);
+  }, [notice, connected, elapsed]);
+
+  if (notice) {
+    return elapsed ? { ...notice, ready: false } : null;
+  }
+  if (
+    connected &&
+    previous?.pending &&
+    previous.sandbox &&
+    previous.kind !== "connection"
+  ) {
+    return {
+      ...previous,
+      pending: false,
+      ready: true,
+      title:
+        previous.kind === "sync"
+          ? "Environment synced"
+          : "Your notebook is ready",
+      description: "The notebook is ready to run.",
+      error: null,
+    };
+  }
+  return null;
+}

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

@@ -3,6 +3,11 @@ 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";
@@ -25,6 +30,8 @@ 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"]';
@@ -35,8 +42,15 @@ const sandbox: SandboxResponse = {
 };
 
 beforeEach(() => {
+  store.set(alertAtom, (value) => ({
+    ...value,
+    environments: {
+      kernel: emptyEnvironmentState(),
+      server: emptyEnvironmentState(),
+    },
+  }));
   store.set(sandboxAtom, null);
-  store.set(sandboxSyncAtom, { pending: false, error: null });
+  store.set(sandboxSyncAtom, { kind: "succeeded" });
   store.set(sandboxActionsAtom, null);
   store.set(connectionAtom, { state: WebSocketState.OPEN });
   store.set(kernelStartupErrorAtom, null);
@@ -65,6 +79,73 @@ function setup() {
 }
 
 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(
@@ -116,7 +197,7 @@ describe("useSandboxController", () => {
     act(() => {
       pending = actions.sync();
     });
-    expect(store.get(sandboxSyncAtom)).toEqual({ pending: true, error: null });
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "running" });
     await act(async () => {
       expect(await actions.sync()).toBe(false);
     });
@@ -125,7 +206,7 @@ describe("useSandboxController", () => {
       resolve({ success: true, reconnect: false });
       expect(await pending).toBe(true);
     });
-    expect(store.get(sandboxSyncAtom)).toEqual({ pending: false, error: null });
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "succeeded" });
     expect(store.get(packageDataVersionAtom)).toBe(1);
     expect(reconnect).not.toHaveBeenCalled();
   });
@@ -138,13 +219,13 @@ describe("useSandboxController", () => {
       expect(await actions.sync()).toBe(false);
     });
     expect(store.get(sandboxSyncAtom)).toEqual({
-      pending: false,
+      kind: "failed",
       error: "offline",
     });
     await act(async () => {
       expect(await actions.sync()).toBe(true);
     });
-    expect(store.get(sandboxSyncAtom)).toEqual({ pending: false, error: null });
+    expect(store.get(sandboxSyncAtom)).toEqual({ kind: "succeeded" });
   });
 
   it("ignores stale metadata after a filename change and unregisters actions on unmount", async () => {

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

@@ -1,12 +1,64 @@
 /* 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 sandboxSyncAtom = atom<{ pending: boolean; error: string | null }>(
-  {
-    pending: false,
-    error: 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<{

package/src/core/packages/useSandboxController.ts

@@ -17,6 +17,7 @@ import {
   sandboxSyncAtom,
 } from "./sandbox-state";
 import {
+  RESTART_REQUIRED_DESCRIPTION,
   showPackageRestartToast,
   showSandboxSyncToast,
 } from "./toast-components";
@@ -63,15 +64,24 @@ export function useSandboxController(onReconnect: () => Promise<void>) {
   }, [requests, filename, connection.state, setSandbox]);
 
   const sync = useEvent(async () => {
-    if (inFlight.current || connection.state === WebSocketState.CONNECTING) {
+    if (
+      inFlight.current ||
+      operation.kind === "running" ||
+      connection.state === WebSocketState.CONNECTING
+    ) {
       return false;
     }
     inFlight.current = true;
-    setOperation((value) => ({ ...value, pending: true, error: null }));
+    setOperation({ kind: "running" });
     try {
       const result = await requests.syncSandbox({ fileKey: filename });
       if (result.restartRequired) {
+        setOperation({
+          kind: "restart-required",
+          reason: result.error ?? RESTART_REQUIRED_DESCRIPTION,
+        });
         showPackageRestartToast();
+        return false;
       }
       if (!result.success) {
         throw new Error(result.error ?? "Sandbox sync failed.");
@@ -85,6 +95,7 @@ export function useSandboxController(onReconnect: () => Promise<void>) {
             value.state === WebSocketState.CLOSED,
         );
         if (status.state === WebSocketState.CLOSED) {
+          setOperation({ kind: "idle" });
           return false;
         }
       }
@@ -92,13 +103,13 @@ export function useSandboxController(onReconnect: () => Promise<void>) {
         invalidatePackageData();
       }
       showSandboxSyncToast();
+      setOperation({ kind: "succeeded" });
       return true;
     } catch (cause) {
-      setOperation((value) => ({ ...value, error: prettyError(cause) }));
+      setOperation({ kind: "failed", error: prettyError(cause) });
       return false;
     } finally {
       inFlight.current = false;
-      setOperation((value) => ({ ...value, pending: false }));
     }
   });
 
@@ -135,7 +146,7 @@ export function useSandboxController(onReconnect: () => Promise<void>) {
   }, [sync, editManifest, setActions]);
 
   const saveAndSync = async () => {
-    if (!document || loading || saving || operation.pending) {
+    if (!document || loading || saving || operation.kind === "running") {
       return;
     }
     setSaving(true);
@@ -165,9 +176,18 @@ export function useSandboxController(onReconnect: () => Promise<void>) {
   const pending =
     saving ||
     loading ||
-    operation.pending ||
+    operation.kind === "running" ||
     connection.state === WebSocketState.CONNECTING;
-  const diagnostic = error ?? operation.error ?? startupError;
+  const diagnostic =
+    error ??
+    (operation.kind === "failed"
+      ? operation.error
+      : operation.kind === "restart-required"
+        ? operation.reason
+        : operation.kind === "cancelled"
+          ? "Sandbox sync was interrupted."
+          : null) ??
+    startupError;
   return {
     filename,
     open,