@marimo-team/frontend 0.24.3-dev82 → 0.24.3-dev83

package/src/core/alerts/__tests__/state.test.tsx

@@ -0,0 +1,240 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import { expect, it } from "vitest";
+import type { EnvironmentOperation } from "../environment";
+import { alertAtom, getPackageAlert, useAlertActions } from "../state";
+
+it.each(["install", "remove"] as const)(
+  "restores active %s progress without announcing old successes",
+  (action) => {
+    const store = createStore();
+    const { result } = renderHook(() => useAlertActions(), {
+      wrapper: ({ children }) => <Provider store={store}>{children}</Provider>,
+    });
+    const operation: EnvironmentOperation = {
+      operation_id: "package",
+      action,
+      source: "kernel",
+      status: { kind: "succeeded" },
+      packages: { numpy: "succeeded" },
+      logs: { numpy: "Done\n" },
+    };
+    const snapshot = (value: EnvironmentOperation) =>
+      result.current.setEnvironment({
+        source: value.source,
+        state: { restart_required: false, operations: [value] },
+      });
+    act(() => snapshot(operation));
+    expect(getPackageAlert(store.get(alertAtom))).toBeNull();
+    expect(store.get(alertAtom).environments.kernel.operations).toEqual([
+      operation,
+    ]);
+
+    const active: EnvironmentOperation = {
+      ...operation,
+      status: { kind: "running" },
+      packages: { numpy: "running" },
+      logs: { numpy: "Working\n" },
+    };
+    act(() => snapshot(active));
+    expect(getPackageAlert(store.get(alertAtom))).toEqual({
+      ...active,
+      id: active.operation_id,
+      kind: "environment",
+      restartRequired: false,
+    });
+
+    // Reconnecting can complete work that was already visible on this page.
+    act(() => snapshot(operation));
+    expect(getPackageAlert(store.get(alertAtom))).toEqual({
+      ...operation,
+      id: operation.operation_id,
+      kind: "environment",
+      restartRequired: false,
+    });
+    act(() => result.current.clearPackageAlert(operation.operation_id));
+    act(() => snapshot(operation));
+    expect(getPackageAlert(store.get(alertAtom))).toBeNull();
+
+    // A live result still gives feedback for a new package action.
+    act(() =>
+      result.current.updateEnvironment({
+        ...operation,
+        operation_id: "new",
+        log_mode: "replace",
+      }),
+    );
+    expect(getPackageAlert(store.get(alertAtom))?.id).toBe("new");
+
+    // Earlier changes can still require a restart after a later action succeeds.
+    act(() => result.current.clearPackageAlert("new"));
+    act(() =>
+      result.current.setEnvironment({
+        source: operation.source,
+        state: { restart_required: true, operations: [operation] },
+      }),
+    );
+    expect(getPackageAlert(store.get(alertAtom))).toEqual({
+      ...operation,
+      id: operation.operation_id,
+      kind: "environment",
+      restartRequired: true,
+    });
+  },
+);
+
+it.each(["prepare", "sync"] as const)(
+  "retains %s progress without turning it into a package alert",
+  (action) => {
+    const store = createStore();
+    const { result } = renderHook(() => useAlertActions(), {
+      wrapper: ({ children }) => <Provider store={store}>{children}</Provider>,
+    });
+    const operation: EnvironmentOperation = {
+      operation_id: "sandbox",
+      action,
+      source: "kernel",
+      status: { kind: "succeeded" },
+      packages: {},
+      logs: { environment: "Environment ready\n" },
+    };
+    const environment = { restart_required: false, operations: [operation] };
+
+    // A refreshed browser restores the result without announcing it again.
+    act(() =>
+      result.current.setEnvironment({ source: "kernel", state: environment }),
+    );
+    expect(getPackageAlert(store.get(alertAtom))).toBeNull();
+    expect(store.get(alertAtom).environments.kernel).toEqual(environment);
+
+    const install: EnvironmentOperation = {
+      ...operation,
+      operation_id: "install",
+      action: "install",
+      status: { kind: "running" },
+      packages: { numpy: "running" },
+      logs: { numpy: "Installing\n" },
+    };
+    act(() =>
+      result.current.updateEnvironment({ ...install, log_mode: "replace" }),
+    );
+    const packageAlert = getPackageAlert(store.get(alertAtom));
+    for (const status of [
+      { kind: "running" },
+      { kind: "failed", error: "Sync failed" },
+    ] as const) {
+      act(() =>
+        result.current.updateEnvironment({
+          ...operation,
+          status,
+          log_mode: "replace",
+        }),
+      );
+      expect(getPackageAlert(store.get(alertAtom))).toEqual(packageAlert);
+    }
+
+    act(() => result.current.clearPackageAlert(install.operation_id));
+    act(() =>
+      result.current.updateEnvironment({ ...operation, log_mode: "replace" }),
+    );
+    expect(getPackageAlert(store.get(alertAtom))).toBeNull();
+  },
+);
+
+it.each(["kernel", "server"] as const)(
+  "preserves the selected %s result across snapshots and follows new work",
+  (source) => {
+    const store = createStore();
+    const { result } = renderHook(() => useAlertActions(), {
+      wrapper: ({ children }) => <Provider store={store}>{children}</Provider>,
+    });
+    const selected: EnvironmentOperation = {
+      operation_id: "selected",
+      action: "install",
+      source,
+      status: { kind: "failed", error: "Network unavailable" },
+      packages: { numpy: "failed" },
+      logs: { numpy: "Selected logs\n" },
+    };
+    const other: EnvironmentOperation = {
+      ...selected,
+      operation_id: "other",
+      source: source === "kernel" ? "server" : "kernel",
+      status: { kind: "succeeded" },
+      packages: { numpy: "succeeded" },
+      logs: { numpy: "Other logs\n" },
+    };
+    const snapshot = (operation: EnvironmentOperation) =>
+      result.current.setEnvironment({
+        source: operation.source,
+        state: { restart_required: false, operations: [operation] },
+      });
+    act(() => snapshot(selected));
+    for (const operations of [
+      [selected, other],
+      [other, selected],
+    ]) {
+      act(() => operations.forEach(snapshot));
+      expect(getPackageAlert(store.get(alertAtom))).toEqual({
+        ...selected,
+        id: selected.operation_id,
+        kind: "environment",
+        restartRequired: false,
+      });
+    }
+
+    // Live work in the other environment should still take over the banner.
+    const active: EnvironmentOperation = {
+      ...other,
+      operation_id: "new-work",
+      status: { kind: "running" },
+      packages: { numpy: "running" },
+      logs: { numpy: "Installing\n" },
+    };
+    act(() =>
+      result.current.updateEnvironment({ ...active, log_mode: "replace" }),
+    );
+    act(() => snapshot(selected));
+    expect(getPackageAlert(store.get(alertAtom))?.id).toBe(active.operation_id);
+    act(() =>
+      result.current.updateEnvironment({
+        ...active,
+        status: { kind: "succeeded" },
+        packages: { numpy: "succeeded" },
+        logs: {},
+        log_mode: "append",
+      }),
+    );
+    act(() => snapshot(selected));
+    expect(getPackageAlert(store.get(alertAtom))).toEqual({
+      ...active,
+      id: active.operation_id,
+      kind: "environment",
+      restartRequired: false,
+      status: { kind: "succeeded" },
+      packages: { numpy: "succeeded" },
+    });
+
+    // Replacing that environment with an empty snapshot removes its selection.
+    act(() =>
+      result.current.setEnvironment({
+        source: active.source,
+        state: { restart_required: false, operations: [] },
+      }),
+    );
+    expect(getPackageAlert(store.get(alertAtom))?.id).toBe(
+      selected.operation_id,
+    );
+    act(() =>
+      result.current.updateEnvironment({
+        ...other,
+        operation_id: "latest-result",
+        log_mode: "replace",
+      }),
+    );
+    expect(getPackageAlert(store.get(alertAtom))?.id).toBe("latest-result");
+  },
+);

package/src/core/alerts/environment.ts

@@ -0,0 +1,57 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { NotificationMessageData, schemas } from "@/core/kernel/messages";
+
+export type EnvironmentState = schemas["EnvironmentState"];
+export type EnvironmentOperation = schemas["EnvironmentOperation"];
+export type EnvironmentSource = EnvironmentOperation["source"];
+
+export function emptyEnvironmentState(): EnvironmentState {
+  return { restart_required: false, operations: [] };
+}
+
+/** Apply a live delta after the connection's authoritative snapshot. */
+export function reduceEnvironmentState(
+  state: EnvironmentState,
+  update: NotificationMessageData<"environment-operation">,
+): EnvironmentState {
+  const previous = state.operations.find(
+    (operation) => operation.operation_id === update.operation_id,
+  );
+  const logs = { ...previous?.logs };
+  for (const [name, content] of Object.entries(update.logs)) {
+    logs[name] =
+      update.log_mode === "replace" ? content : (logs[name] ?? "") + content;
+  }
+
+  const operation: EnvironmentOperation = {
+    operation_id: update.operation_id,
+    action: update.action,
+    status: update.status,
+    packages: update.packages,
+    logs,
+    source: update.source,
+  };
+  const operations = state.operations.filter(
+    (item) =>
+      (previous && update.status.kind === "running") ||
+      item.status.kind === "running" ||
+      (item.action === "prepare" && update.action !== "prepare"),
+  );
+  const index = operations.findIndex(
+    (item) => item.operation_id === update.operation_id,
+  );
+  if (index === -1) {
+    operations.push(operation);
+  } else {
+    operations[index] = operation;
+  }
+
+  return {
+    restart_required:
+      state.restart_required ||
+      update.status.kind === "restart-required" ||
+      Object.values(update.packages).includes("restart-required"),
+    operations,
+  };
+}

package/src/core/alerts/state.ts

@@ -1,9 +1,16 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { useAtomValue } from "jotai";
-import type { PackageInstallationStatus } from "@/core/kernel/messages";
+import type { NotificationMessageData } from "@/core/kernel/messages";
 import { createReducerAndAtoms } from "@/utils/createReducer";
 import { generateUUID } from "@/utils/uuid";
+import {
+  emptyEnvironmentState,
+  type EnvironmentOperation,
+  type EnvironmentSource,
+  type EnvironmentState,
+  reduceEnvironmentState,
+} from "./environment";
 
 type Identified<T> = { id: string } & T;
 
@@ -14,12 +21,18 @@ export interface MissingPackageAlert {
   source?: "kernel" | "server";
 }
 
-export interface InstallingPackageAlert {
-  kind: "installing";
-  packages: PackageInstallationStatus;
-  logs?: { [key: string]: string } | null;
-  log_status?: "append" | "start" | "done" | null;
-  source?: "kernel" | "server";
+export interface EnvironmentOperationAlert extends EnvironmentOperation {
+  action: "install" | "remove";
+  kind: "environment";
+  restartRequired: boolean;
+}
+
+function isPackageOperation(
+  operation: EnvironmentOperation,
+): operation is EnvironmentOperation & {
+  action: EnvironmentOperationAlert["action"];
+} {
+  return operation.action === "install" || operation.action === "remove";
 }
 
 export interface StartupLogsAlert {
@@ -28,118 +41,161 @@ export interface StartupLogsAlert {
 }
 
 export function isMissingPackageAlert(
-  alert: MissingPackageAlert | InstallingPackageAlert,
+  alert: MissingPackageAlert | EnvironmentOperationAlert,
 ): alert is MissingPackageAlert {
   return alert.kind === "missing";
 }
 
-export function isInstallingPackageAlert(
-  alert: MissingPackageAlert | InstallingPackageAlert,
-): alert is InstallingPackageAlert {
-  return alert.kind === "installing";
+export function isEnvironmentOperationAlert(
+  alert: MissingPackageAlert | EnvironmentOperationAlert,
+): alert is EnvironmentOperationAlert {
+  return alert.kind === "environment";
 }
 
-/** Prominent alerts.
- *
- * Right now we only have one type of alert.
- */
 interface AlertState {
   packageAlert:
     | Identified<MissingPackageAlert>
-    | Identified<InstallingPackageAlert>
+    | { kind: "environment"; id: string; source: EnvironmentSource }
     | null;
+  environments: Record<EnvironmentSource, EnvironmentState>;
   startupLogsAlert: StartupLogsAlert | null;
-  packageLogs: { [packageName: string]: string };
 }
 
-const { valueAtom: alertAtom, useActions } = createReducerAndAtoms(
-  () =>
-    ({
-      packageAlert: null,
-      startupLogsAlert: null,
-      packageLogs: {},
-    }) as AlertState,
-  {
-    addPackageAlert: (
-      state,
-      alert: MissingPackageAlert | InstallingPackageAlert,
-    ) => {
-      const newPackageLogs = { ...state.packageLogs };
-
-      // Handle streaming logs for installing package alerts
-      if (isInstallingPackageAlert(alert) && alert.logs && alert.log_status) {
-        for (const [packageName, newContent] of Object.entries(alert.logs)) {
-          switch (alert.log_status) {
-            case "start":
-              // Start new log for this package
-              newPackageLogs[packageName] = newContent;
-
-              break;
-
-            case "append": {
-              // Append to existing log
-              const prevContent = newPackageLogs[packageName] || "";
-              newPackageLogs[packageName] = prevContent + newContent;
-
-              break;
-            }
-            case "done": {
-              // Append final content and mark as done
-              const prevContent = newPackageLogs[packageName] || "";
-              newPackageLogs[packageName] = prevContent + newContent;
-
-              break;
-            }
-            // No default
-          }
+function setEnvironment(
+  state: AlertState,
+  source: EnvironmentSource,
+  environment: EnvironmentState,
+): AlertState {
+  // Sandbox preparation and sync use the existing sandbox UI.
+  const environments = { ...state.environments, [source]: environment };
+  const operations = Object.values(environments)
+    .flatMap((item) => item.operations)
+    .filter(isPackageOperation);
+  const incoming = environment.operations.filter(isPackageOperation);
+  const alert = state.packageAlert;
+  const selected =
+    alert?.kind === "environment"
+      ? operations.find(
+          (item) =>
+            item.operation_id === alert.id && item.source === alert.source,
+        )
+      : undefined;
+  // A snapshot restores work and unresolved issues, not old success banners.
+  const needsAttention = (operation: EnvironmentOperation) =>
+    operation.status.kind !== "succeeded" ||
+    environments[operation.source].restart_required;
+  const operation =
+    (selected?.status.kind === "running" ? selected : undefined) ??
+    incoming.findLast((item) => item.status.kind === "running") ??
+    operations.findLast((item) => item.status.kind === "running") ??
+    selected ??
+    incoming.findLast(needsAttention) ??
+    operations.findLast(needsAttention);
+  return {
+    ...state,
+    environments,
+    packageAlert: operation
+      ? {
+          kind: "environment",
+          id: operation.operation_id,
+          source: operation.source,
         }
-      }
-
-      const existingAlert = state.packageAlert;
-      const alertId = existingAlert?.id || generateUUID();
+      : state.packageAlert?.kind === "missing"
+        ? state.packageAlert
+        : null,
+  };
+}
 
-      return {
+export const { valueAtom: alertAtom, useActions: useAlertActions } =
+  createReducerAndAtoms(
+    (): AlertState => ({
+      packageAlert: null,
+      startupLogsAlert: null,
+      environments: {
+        kernel: emptyEnvironmentState(),
+        server: emptyEnvironmentState(),
+      },
+    }),
+    {
+      addMissingPackageAlert: (state, alert: MissingPackageAlert) => ({
         ...state,
-        packageAlert: { id: alertId, ...alert },
-        packageLogs: newPackageLogs,
-      };
-    },
-
-    clearPackageAlert: (state, id: string) => {
-      return state.packageAlert !== null && state.packageAlert.id === id
-        ? { ...state, packageAlert: null, packageLogs: {} }
-        : state;
-    },
-
-    addStartupLog: (
-      state,
-      logData: { content: string; status: "append" | "start" | "done" },
-    ) => {
-      const prevContent = state.startupLogsAlert?.content || "";
-      return {
+        packageAlert: { id: generateUUID(), ...alert },
+      }),
+
+      updateEnvironment: (
+        state,
+        update: NotificationMessageData<"environment-operation">,
+      ) => {
+        const source = update.source;
+        const environment = reduceEnvironmentState(
+          state.environments[source],
+          update,
+        );
+        if (!isPackageOperation(update)) {
+          return {
+            ...state,
+            environments: { ...state.environments, [source]: environment },
+          };
+        }
+        return setEnvironment(
+          {
+            ...state,
+            packageAlert: {
+              kind: "environment",
+              id: update.operation_id,
+              source,
+            },
+          },
+          source,
+          environment,
+        );
+      },
+
+      setEnvironment: (
+        state,
+        snapshot: NotificationMessageData<"environment-state">,
+      ) => setEnvironment(state, snapshot.source, snapshot.state),
+
+      clearPackageAlert: (state, id: string) =>
+        state.packageAlert?.id === id
+          ? { ...state, packageAlert: null }
+          : state,
+
+      addStartupLog: (
+        state,
+        logData: { content: string; status: "append" | "start" | "done" },
+      ) => ({
         ...state,
         startupLogsAlert: {
-          ...state.startupLogsAlert,
-          content: prevContent + logData.content,
+          content: (state.startupLogsAlert?.content ?? "") + logData.content,
           status: logData.status,
         },
-      };
-    },
+      }),
 
-    clearStartupLogsAlert: (state) => {
-      return { ...state, startupLogsAlert: null };
+      clearStartupLogsAlert: (state) => ({ ...state, startupLogsAlert: null }),
     },
-  },
-);
-
-/**
- * React hook to get the Alert state.
- */
-export const useAlerts = () => useAtomValue(alertAtom);
-
-/**
- * React hook to get the Alerts actions.
- */
-export function useAlertActions() {
-  return useActions();
+  );
+
+export function getPackageAlert(
+  state: AlertState,
+): Identified<MissingPackageAlert | EnvironmentOperationAlert> | null {
+  const alert = state.packageAlert;
+  if (alert?.kind !== "environment") {
+    return alert;
+  }
+  const environment = state.environments[alert.source];
+  const operation = environment.operations.find(
+    (item) => item.operation_id === alert.id,
+  );
+  return operation && isPackageOperation(operation)
+    ? { ...operation, ...alert, restartRequired: environment.restart_required }
+    : null;
+}
+
+export function useAlerts() {
+  const state = useAtomValue(alertAtom);
+  return {
+    packageAlert: getPackageAlert(state),
+    startupLogsAlert: state.startupLogsAlert,
+  };
 }

package/src/core/islands/bootstrap.ts

@@ -181,7 +181,8 @@ function handleMessage(
       // Unsupported operations in islands mode
       case "banner":
       case "missing-package-alert":
-      case "installing-package-alert":
+      case "environment-operation":
+      case "environment-state":
       case "completion-result":
       case "reload":
       case "focus-cell":

package/src/core/kernel/messages.ts

@@ -36,8 +36,8 @@ export type CompletionResultMessage =
 export type HumanReadableStatus = schemas["HumanReadableStatus"];
 export type FunctionCallResultMessage =
   NotificationMessageData<"function-call-result">;
-export type PackageInstallationStatus =
-  schemas["InstallingPackageAlertNotification"]["packages"];
+export type PackageOperationStatus =
+  schemas["EnvironmentOperationNotification"]["packages"];
 export type DataColumnPreview = NotificationMessageData<"data-column-preview">;
 export type SQLTablePreview = NotificationMessageData<"sql-table-preview">;
 export type SQLTableListPreview =