@marimo-team/islands 0.25.3-dev18 → 0.25.3-dev182

package/src/components/editor/cell/RunButton.tsx

@@ -5,8 +5,8 @@ import type { CellConfig, RuntimeState } from "@/core/network/types";
 import {
   getConnectionTooltip,
   isAppInteractionDisabled,
-} from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+} from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { renderShortcut } from "../../shortcuts/renderShortcut";
 import { ToolbarItem } from "./toolbar";
 

package/src/components/editor/cell/StagedAICell.tsx

@@ -11,7 +11,7 @@ import {
 } from "@/core/ai/staged-cells";
 import { getCellEditorView } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
+import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
 import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
 import {
@@ -144,7 +144,15 @@ export function rejectStagedCell(
         break;
       }
 
-      updateEditorCodeFromPython(editorView, edit.previousCode);
+      const update = prepareCellCodeUpdate(
+        editorView.state,
+        edit.previousCode,
+        { initialize: true },
+      );
+
+      if (update) {
+        editorView.dispatch(update);
+      }
       removeStagedCell(cellId);
       break;
     }

package/src/components/editor/cell/StopButton.tsx

@@ -3,8 +3,8 @@ import { SquareIcon } from "lucide-react";
 import type { JSX } from "react";
 import { useRequestClient } from "@/core/network/requests";
 import type { RuntimeState } from "@/core/network/types";
-import { isAppInteractionDisabled } from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+import { isAppInteractionDisabled } from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { Functions } from "@/utils/functions";
 import { renderShortcut } from "../../shortcuts/renderShortcut";
 import { ToolbarItem } from "./toolbar";

package/src/components/editor/cell/__tests__/StagedAICell.test.tsx

@@ -14,8 +14,10 @@ vi.mock("@/components/editor/cell/useDeleteCell", () => ({
 
 vi.mock("@/core/cells/cells", () => ({
   getCellEditorView: vi.fn(),
+}));
+vi.mock("@/core/cells/actions", () => ({
   useCellActions: () => ({
-    createNewCell: vi.fn(),
+    createCell: vi.fn(),
     updateCellCode: vi.fn(),
   }),
 }));

package/src/components/editor/cell/__tests__/useShouldShowInterrupt.test.ts

@@ -0,0 +1,39 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { useShouldShowInterrupt } from "../useShouldShowInterrupt";
+
+describe("useShouldShowInterrupt", () => {
+  beforeEach(() => vi.useFakeTimers());
+  afterEach(() => vi.useRealTimers());
+
+  it("waits for each run and hides immediately when it stops", () => {
+    const { result, rerender } = renderHook(useShouldShowInterrupt, {
+      initialProps: false,
+    });
+    expect(result.current).toBe(false);
+    rerender(true);
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(200));
+    expect(result.current).toBe(true);
+    rerender(false);
+    expect(result.current).toBe(false);
+    rerender(true);
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(199));
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(1));
+    expect(result.current).toBe(true);
+  });
+
+  it("cancels the delay for a short run", () => {
+    const { result, rerender } = renderHook(useShouldShowInterrupt, {
+      initialProps: true,
+    });
+    act(() => vi.advanceTimersByTime(100));
+    rerender(false);
+    act(() => vi.advanceTimersByTime(200));
+    expect(result.current).toBe(false);
+  });
+});

package/src/components/editor/cell/code/__tests__/language-toggle.test.tsx

@@ -0,0 +1,54 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { fireEvent, render, screen } from "@testing-library/react";
+import { createRef } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { switchLanguage } from "@/core/codemirror/language/extension";
+import { LanguageToggle } from "../language-toggle";
+
+vi.mock("@/core/codemirror/language/extension", () => ({
+  switchLanguage: vi.fn(),
+}));
+
+describe("LanguageToggle", () => {
+  afterEach(() => {
+    vi.mocked(switchLanguage).mockClear();
+  });
+
+  it("uses the editor mounted after the button rendered", () => {
+    const editorViewRef = createRef<EditorView>();
+    const onAfterToggle = vi.fn();
+    render(
+      <TooltipProvider>
+        <LanguageToggle
+          editorViewRef={editorViewRef}
+          currentLanguageAdapter="python"
+          canSwitchToLanguage={true}
+          toType="markdown"
+          displayName="Markdown"
+          icon="Markdown"
+          onAfterToggle={onAfterToggle}
+        />
+      </TooltipProvider>,
+    );
+
+    fireEvent.click(screen.getByRole("button", { name: "Markdown" }));
+    expect(switchLanguage).not.toHaveBeenCalled();
+    expect(onAfterToggle).not.toHaveBeenCalled();
+
+    const view = new EditorView({ state: EditorState.create() });
+    editorViewRef.current = view;
+    try {
+      fireEvent.click(screen.getByRole("button", { name: "Markdown" }));
+      expect(switchLanguage).toHaveBeenCalledWith(view, {
+        language: "markdown",
+      });
+      expect(onAfterToggle).toHaveBeenCalledOnce();
+    } finally {
+      view.destroy();
+    }
+  });
+});

package/src/components/editor/cell/code/cell-editor.tsx

@@ -10,7 +10,9 @@ import { Button } from "@/components/ui/button";
 import { DelayMount } from "@/components/utils/delay-mount";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { getNotebook, useCellActions } from "@/core/cells/cells";
+import { getNotebook } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import { SETUP_CELL_ID } from "@/core/cells/ids";
 import { usePendingDeleteService } from "@/core/cells/pending-delete-service";
 import type { CellData, CellRuntimeState } from "@/core/cells/types";
@@ -24,10 +26,8 @@ import {
   reconfigureLanguageEffect,
   switchLanguage,
 } from "@/core/codemirror/language/extension";
-import {
-  getEditorCodeAsPython,
-  updateEditorCodeFromPython,
-} from "@/core/codemirror/language/utils";
+import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
+import { getEditorCodeAsPython } from "@/core/codemirror/language/utils";
 import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages/markdown";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
 import {
@@ -38,10 +38,9 @@ import { autoInstantiateAtom, isAiFeatureEnabled } from "@/core/config/config";
 import type { UserConfig } from "@/core/config/config-schema";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { connectionAtom } from "@/core/network/connection";
-import { useRequestClient } from "@/core/network/requests";
 import { canUseRtc, isRtcEnabled } from "@/core/rtc/state";
 import { useSaveNotebook } from "@/core/saving/save-component";
-import { isAppConnecting } from "@/core/websocket/connection-utils";
+import { isAppConnecting } from "@/core/session/connection-utils";
 import type { Theme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { invariant } from "@/utils/invariant";
@@ -137,7 +136,7 @@ const CellEditorInternal = ({
   const deleteCell = useDeleteCellCallback();
   const { saveOrNameNotebook } = useSaveNotebook();
   const pendingDeleteService = usePendingDeleteService();
-  const { saveCellConfig } = useRequestClient();
+  const edit = useNotebookDocumentEdit();
 
   const loading = status === "running" || status === "queued";
   const cellActions = useCellActions();
@@ -177,9 +176,7 @@ const CellEditorInternal = ({
   const toggleHideCode = useEvent(() => {
     // Use cellConfig.hide_code instead of hidden, since it may be temporarily shown
     const nextHidden = !cellConfig.hide_code;
-    // Fire-and-forget save
-    void saveCellConfig({ configs: { [cellId]: { hide_code: nextHidden } } });
-    cellActions.updateCellConfig({ cellId, config: { hide_code: nextHidden } });
+    edit((tx) => tx.updateConfig(cellId, { hide_code: nextHidden }));
     return nextHidden;
   });
 
@@ -188,17 +185,13 @@ const CellEditorInternal = ({
   const afterToggleMarkdown = useEvent(() => {
     maybeAddMarimoImport({
       autoInstantiate,
-      createNewCell: cellActions.createNewCell,
+      createCell: cellActions.createCell,
     });
     // Code stays visible until the user blurs the cell
     if (!cellConfig.hide_code && MARKDOWN_INITIAL_HIDE_CODE) {
-      void saveCellConfig({
-        configs: { [cellId]: { hide_code: MARKDOWN_INITIAL_HIDE_CODE } },
-      });
-      cellActions.updateCellConfig({
-        cellId,
-        config: { hide_code: MARKDOWN_INITIAL_HIDE_CODE },
-      });
+      edit((tx) =>
+        tx.updateConfig(cellId, { hide_code: MARKDOWN_INITIAL_HIDE_CODE }),
+      );
       cellActions.markUntouched({ cellId });
     }
   });
@@ -206,7 +199,7 @@ const CellEditorInternal = ({
   const afterToggleSQL = useEvent(() => {
     maybeAddMarimoImport({
       autoInstantiate,
-      createNewCell: cellActions.createNewCell,
+      createCell: cellActions.createCell,
     });
   });
 
@@ -226,11 +219,9 @@ const CellEditorInternal = ({
         saveNotebook: saveOrNameNotebook,
         createManyBelow: (cells) => {
           for (const code of cells.toReversed()) {
-            cellActions.createNewCell({
+            cellActions.createCell({
+              after: cellId,
               code,
-              before: false,
-              cellId: cellId,
-              // If the code already exists, skip creation
               skipIfCodeExists: true,
             });
           }
@@ -239,7 +230,7 @@ const CellEditorInternal = ({
           const notebook = getNotebook();
           // No setup cell yet: create one with this code at the top.
           if (!notebook.cellIds.setupCellExists()) {
-            cellActions.addSetupCellIfDoesntExist({ code });
+            cellActions.addSetupCell({ code });
             return;
           }
           // Setup cell exists: append the pasted code to it. Prefer updating
@@ -253,7 +244,10 @@ const CellEditorInternal = ({
             ? `${existing.trimEnd()}\n\n${code}`
             : code;
           if (view) {
-            updateEditorCodeFromPython(view, merged);
+            const update = prepareCellCodeUpdate(view.state, merged);
+            if (update) {
+              view.dispatch(update);
+            }
           }
           cellActions.updateCellCode({
             cellId: SETUP_CELL_ID,
@@ -363,15 +357,7 @@ const CellEditorInternal = ({
       return composeRtcEditorConfig(code, extensions, undefined);
     }
 
-    const rtc = realTimeCollaboration(
-      cellId,
-      (code) => {
-        // It's not really a formatting change,
-        // but this means it won't be marked as stale
-        cellActions.updateCellCode({ cellId, code, formattingChange: true });
-      },
-      initialCode,
-    );
+    const rtc = realTimeCollaboration(cellId, initialCode);
     return composeRtcEditorConfig(code, extensions, rtc);
   };
 
@@ -572,7 +558,7 @@ const CellEditorInternal = ({
           <div className="absolute top-1 right-5">
             <LanguageToggles
               code={code}
-              editorView={editorViewRef.current}
+              editorViewRef={editorViewRef}
               currentLanguageAdapter={languageAdapter}
               onAfterToggle={afterToggleMarkdown}
             />

package/src/components/editor/cell/code/language-toggle.tsx

@@ -13,7 +13,7 @@ import { Functions } from "@/utils/functions";
 import { MarkdownIcon, PythonIcon } from "./icons";
 
 interface LanguageTogglesProps {
-  editorView: EditorView | null;
+  editorViewRef: React.RefObject<EditorView | null>;
   code: string;
   currentLanguageAdapter: LanguageAdapter["type"] | undefined;
   onAfterToggle: () => void;
@@ -25,7 +25,7 @@ interface LanguageTogglesProps {
 }
 
 export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
-  editorView,
+  editorViewRef,
   code,
   currentLanguageAdapter,
   onAfterToggle,
@@ -43,7 +43,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
   return (
     <div className={className}>
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={canUseSQL && currentLanguageAdapter === "python"}
         icon={
@@ -58,7 +58,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
         onAfterToggle={onAfterToggle}
       />
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={
           canUseMarkdown && currentLanguageAdapter === "python"
@@ -75,7 +75,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
         onAfterToggle={onAfterToggle}
       />
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={true}
         icon={
@@ -95,7 +95,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
 
 interface Props {
   className?: string;
-  editorView: EditorView | null;
+  editorViewRef: React.RefObject<EditorView | null>;
   canSwitchToLanguage: boolean;
   currentLanguageAdapter: LanguageAdapter["type"] | undefined;
   toType: LanguageAdapter["type"];
@@ -105,7 +105,7 @@ interface Props {
 }
 
 export const LanguageToggle: React.FC<Props> = ({
-  editorView,
+  editorViewRef,
   currentLanguageAdapter,
   canSwitchToLanguage,
   icon,
@@ -114,6 +114,7 @@ export const LanguageToggle: React.FC<Props> = ({
   onAfterToggle,
 }) => {
   const handleClick = () => {
+    const editorView = editorViewRef.current;
     if (!editorView) {
       return;
     }

package/src/components/editor/cell/collapse.tsx

@@ -67,10 +67,12 @@ export const CollapsedCellBanner: React.FC<{
   const states = getDescendantsStatus(notebook, cellId);
 
   return (
-    <div
+    <button
+      type="button"
+      aria-label={`Expand ${count} collapsed ${count === 1 ? "cell" : "cells"}`}
       onClick={onClick}
       className={cn(
-        "flex items-center justify-between w-[calc(100%-2rem)] h-9 bg-muted rounded-b mx-4 opacity-80 hover:opacity-100 cursor-pointer",
+        "flex items-center justify-between w-[calc(100%-2rem)] h-9 bg-muted rounded-b mx-4 opacity-80 hover:opacity-100 cursor-pointer focus-visible:outline focus-visible:outline-ring",
       )}
     >
       <Toolbar
@@ -102,7 +104,7 @@ export const CollapsedCellBanner: React.FC<{
           </>
         }
       />
-    </div>
+    </button>
   );
 });
 CollapsedCellBanner.displayName = "CollapsedCellBanner";

package/src/components/editor/cell/useAddCell.ts

@@ -2,28 +2,24 @@
 
 import { useAtomValue } from "jotai";
 import { maybeAddAltairImport } from "@/core/cells/add-missing-import";
-import { useCellActions } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import { useLastFocusedCellId } from "@/core/cells/focus";
 import { autoInstantiateAtom } from "@/core/config/config";
 
 export function useAddCodeToNewCell(): (code: string) => void {
   const autoInstantiate = useAtomValue(autoInstantiateAtom);
   const lastFocusedCellId = useLastFocusedCellId();
-  const { createNewCell } = useCellActions();
+  const { createCell } = useCellActions();
 
   return (code: string) => {
     if (code.includes("alt")) {
       maybeAddAltairImport({
         autoInstantiate,
-        createNewCell,
+        createCell,
         fromCellId: lastFocusedCellId,
       });
     }
 
-    createNewCell({
-      code: code,
-      before: false,
-      cellId: lastFocusedCellId ?? "__end__",
-    });
+    createCell({ after: lastFocusedCellId, code });
   };
 }

package/src/components/editor/cell/useDeleteCell.tsx

@@ -3,35 +3,22 @@
 import useEvent from "react-use-event-hook";
 import { UndoButton } from "@/components/buttons/undo-button";
 import { toast } from "@/components/ui/use-toast";
-import {
-  hasOnlyOneCellAtom,
-  notebookAtom,
-  useCellActions,
-} from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
+import { notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import { useRequestClient } from "@/core/network/requests";
 import { store } from "@/core/state/jotai";
 
 export function useDeleteCellCallback() {
-  const { deleteCell, undoDeleteCell } = useCellActions();
-  const { sendDeleteCell } = useRequestClient();
+  const { deleteCells, undoDeleteCell } = useCellActions();
 
   return useEvent((opts: { cellId: CellId }) => {
-    // Can't delete the last cell
-    if (store.get(hasOnlyOneCellAtom)) {
-      return;
-    }
-
     const { cellId } = opts;
     const notebook = store.get(notebookAtom);
     const isEmptyCell = (notebook.cellData[cellId]?.code ?? "").trim() === "";
 
-    // Optimistic update
-    deleteCell({ cellId });
-    sendDeleteCell({ cellId: cellId }).catch(() => {
-      // Fall back on failure
-      undoDeleteCell();
-    });
+    if (!deleteCells({ cellIds: [cellId] })) {
+      return;
+    }
 
     if (!isEmptyCell) {
       const { dismiss } = toast({
@@ -53,33 +40,20 @@ export function useDeleteCellCallback() {
 }
 
 export function useDeleteManyCellsCallback() {
-  const { deleteCell, undoDeleteCell } = useCellActions();
-  const { sendDeleteCell } = useRequestClient();
+  const { deleteCells, undoDeleteCell } = useCellActions();
 
-  return useEvent(async (opts: { cellIds: CellId[] }) => {
-    // Can't delete the last cell
-    if (store.get(hasOnlyOneCellAtom)) {
+  return useEvent((opts: { cellIds: CellId[] }) => {
+    if (!deleteCells({ cellIds: opts.cellIds })) {
       return;
     }
 
-    const { cellIds } = opts;
-    for (const cellId of cellIds) {
-      await sendDeleteCell({ cellId }).then(() => {
-        deleteCell({ cellId });
-      });
-    }
-
     const { dismiss } = toast({
       title: "Cells deleted",
       action: (
         <UndoButton
           data-testid="undo-delete-button"
           onClick={() => {
-            for (const _cellId of cellIds) {
-              // This function does not take a cellId,
-              // so we just undo the number of cells that were deleted
-              undoDeleteCell();
-            }
+            undoDeleteCell();
             dismiss();
           }}
         />

package/src/components/editor/cell/useRunCells.ts

@@ -3,11 +3,8 @@
 import { closeCompletion } from "@codemirror/autocomplete";
 import { atom, useSetAtom } from "jotai";
 import useEvent from "react-use-event-hook";
-import {
-  getNotebook,
-  type NotebookState,
-  useCellActions,
-} from "@/core/cells/cells";
+import { getNotebook, type NotebookState } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
 import { enabledCellIds, staleCellIds } from "@/core/cells/utils";
 import { getCurrentLanguageAdapter } from "@/core/codemirror/language/commands";

package/src/components/editor/cell/useShouldShowInterrupt.ts

@@ -1,5 +1,5 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { useEffect, useState } from "react";
+import { useDelayElapsed } from "@/hooks/useDelayElapsed";
 
 const DELAY_MS = 200;
 
@@ -10,17 +10,7 @@ const DELAY_MS = 200;
 export function useShouldShowInterrupt(running: boolean) {
   // Start a timer when the run starts.
   // After 200ms, show the interrupt button to avoid flickering.
-  const [hasRunLongEnough, setHasRunLongEnough] = useState(false);
-  useEffect(() => {
-    if (!running) {
-      return;
-    }
-    setHasRunLongEnough(false);
-    const timeout = setTimeout(() => {
-      setHasRunLongEnough(true);
-    }, DELAY_MS);
-    return () => clearTimeout(timeout);
-  }, [running]);
+  const hasRunLongEnough = useDelayElapsed(running ? DELAY_MS : 0);
 
   return running && hasRunLongEnough;
 }

package/src/components/editor/cell/useSplitCell.tsx

@@ -3,7 +3,8 @@
 import useEvent from "react-use-event-hook";
 import { UndoButton } from "@/components/buttons/undo-button";
 import { toast } from "@/components/ui/use-toast";
-import { getCellEditorView, useCellActions } from "@/core/cells/cells";
+import { getCellEditorView } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
 import { getEditorCodeAsPython } from "@/core/codemirror/language/utils";
 

package/src/components/editor/chrome/components/__tests__/feedback-button.test.tsx

@@ -14,7 +14,7 @@ import { requestClientAtom } from "@/core/network/requests";
 import type { EnvironmentInfo } from "@/core/network/types";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import * as copyModule from "@/utils/copy";
 import { FeedbackModal } from "../feedback-button";
 

package/src/components/editor/chrome/components/feedback-button.tsx

@@ -39,7 +39,7 @@ import { connectionAtom } from "@/core/network/connection";
 import { requestClientAtom } from "@/core/network/requests";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import { useLocalStorage } from "@/hooks/useLocalStorage";
 import { cn } from "@/utils/cn";

package/src/components/editor/chrome/panels/__tests__/packages-panel.test.tsx

@@ -18,11 +18,14 @@ import type {
   DependencyTreeNode,
   DependencyTreeResponse,
 } from "@/core/network/types";
-import { withPackageInvalidation } from "@/core/packages/package-data";
+import {
+  invalidatePackageData,
+  withPackageInvalidation,
+} from "@/core/packages/package-data";
 import { useInstallPackages } from "@/core/packages/useInstallPackage";
 import { sandboxAtom, sandboxSyncAtom } from "@/core/packages/sandbox-state";
 import { store } from "@/core/state/jotai";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import PackagesPanel from "../packages-panel";
 import { PanelSectionProvider } from "../panel-context";
 import { RuntimeStatusToggle } from "../runtime-status-toggle";
@@ -143,6 +146,32 @@ describe("PackagesPanel", () => {
     expect(screen.queryByText("cycle")).not.toBeInTheDocument();
   });
 
+  it("collapses expanded nodes when the dependency tree is refreshed", async () => {
+    const { client } = renderPanel(
+      { kind: "sandbox", backend: "pixi" },
+      populatedTree,
+    );
+    fireEvent.click(await screen.findByRole("treeitem", { name: /polars/ }));
+    expect(screen.getByText("already in tree")).toBeInTheDocument();
+
+    vi.mocked(client.getDependencyTree).mockResolvedValue({
+      tree: structuredClone(populatedTree),
+      context: { kind: "sandbox", backend: "pixi" },
+    });
+    act(() => invalidatePackageData());
+
+    await waitFor(() =>
+      expect(client.getDependencyTree).toHaveBeenCalledTimes(2),
+    );
+    await waitFor(() =>
+      expect(screen.queryByText("already in tree")).not.toBeInTheDocument(),
+    );
+    expect(screen.getByRole("treeitem", { name: /polars/ })).toHaveAttribute(
+      "aria-expanded",
+      "false",
+    );
+  });
+
   it("keeps the configured package manager changeable outside a sandbox", async () => {
     const { getPackageList } = renderPanel({
       kind: "package-manager",

package/src/components/editor/chrome/panels/__tests__/runtime-details.test.tsx

@@ -11,7 +11,7 @@ import { kernelStartupErrorAtom } from "@/core/errors/state";
 import { connectionAtom, startupProgressAtom } from "@/core/network/connection";
 import { requestClientAtom } from "@/core/network/requests";
 import { sandboxAtom } from "@/core/packages/sandbox-state";
-import { WebSocketClosedReason, ConnectionState } from "@/core/websocket/types";
+import { WebSocketClosedReason, ConnectionState } from "@/core/session/types";
 import PackagesPanel from "../packages-panel";
 import { PanelSectionProvider } from "../panel-context";
 import { RuntimeStatusToggle } from "../runtime-status-toggle";

package/src/components/editor/chrome/panels/__tests__/sandbox-recovery.test.tsx

@@ -24,7 +24,7 @@ import {
 } from "@/core/packages/sandbox-state";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
-import { WebSocketClosedReason, ConnectionState } from "@/core/websocket/types";
+import { WebSocketClosedReason, ConnectionState } from "@/core/session/types";
 import { chromeAtom } from "../../state";
 import PackagesPanel from "../packages-panel";
 import { PanelSectionProvider } from "../panel-context";