@marimo-team/islands 0.25.3-dev15 → 0.25.3-dev151

package/src/components/editor/__tests__/dynamic-favicon.test.tsx

@@ -1,17 +1,41 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { render, waitFor } from "@testing-library/react";
+import {
+  act,
+  render as renderComponent,
+  waitFor,
+} from "@testing-library/react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { useCellErrors } from "@/core/cells/cells";
+import { Provider, createStore } from "jotai";
+import type { ReactNode } from "react";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { cellId } from "@/__tests__/branded";
+import { notebookAtom } from "@/core/cells/cells";
 import { DynamicFavicon } from "../dynamic-favicon";
 
-// Mock useCellErrors hook
-vi.mock("@/core/cells/cells", () => ({
-  useCellErrors: vi.fn(),
-}));
-
 describe("DynamicFavicon", () => {
   let favicon: HTMLLinkElement;
+  let store: ReturnType<typeof createStore>;
+  const id = cellId("cell");
+
+  function render(ui: ReactNode) {
+    return renderComponent(ui, {
+      wrapper: ({ children }) => <Provider store={store}>{children}</Provider>,
+    });
+  }
+
+  function setErrors() {
+    store.set(
+      notebookAtom,
+      MockNotebook.notebookStateWithErrors([
+        {
+          cellId: id,
+          cellName: "cell",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]),
+    );
+  }
 
   beforeEach(() => {
     document
@@ -27,15 +51,20 @@ describe("DynamicFavicon", () => {
     // Mock document.hasFocus
     vi.spyOn(document, "hasFocus").mockReturnValue(true);
 
-    // Mock useCellErrors to return no errors by default
-    (useCellErrors as ReturnType<typeof vi.fn>).mockReturnValue([]);
+    store = createStore();
+    store.set(
+      notebookAtom,
+      MockNotebook.notebookState({
+        cellData: { [id]: { code: "x = 1" } },
+      }),
+    );
   });
 
   afterEach(() => {
     document
       .querySelectorAll("link[rel~='icon']")
       .forEach((element) => element.remove());
-    vi.clearAllMocks();
+    vi.restoreAllMocks();
     vi.useRealTimers();
   });
 
@@ -67,9 +96,7 @@ describe("DynamicFavicon", () => {
   });
 
   it("should show error favicon when run completes with errors", async () => {
-    (useCellErrors as ReturnType<typeof vi.fn>).mockReturnValue([
-      { error: "mock error" },
-    ]);
+    setErrors();
 
     const { rerender } = render(<DynamicFavicon isRunning={true} />);
 
@@ -147,6 +174,72 @@ describe("DynamicFavicon", () => {
     });
   });
 
+  it("does not look up or rewrite the favicon when code changes", async () => {
+    const { rerender } = render(<DynamicFavicon isRunning={true} />);
+    await waitFor(() => expect(favicon.href).toContain("circle-play.ico"));
+    rerender(<DynamicFavicon isRunning={false} />);
+    await waitFor(() => expect(favicon.href).toContain("circle-check.ico"));
+    const lookup = vi.spyOn(document, "querySelector");
+    const write = vi.spyOn(favicon, "href", "set");
+
+    await act(async () => {
+      store.set(notebookAtom, (state) => ({
+        ...state,
+        cellData: {
+          ...state.cellData,
+          [id]: { ...state.cellData[id], code: "x = 2" },
+        },
+      }));
+    });
+
+    expect(
+      lookup.mock.calls.filter(
+        ([selector]) => selector === "link[rel~='icon']",
+      ),
+    ).toEqual([]);
+    expect(write).not.toHaveBeenCalled();
+    lookup.mockRestore();
+    write.mockRestore();
+  });
+
+  it("updates the favicon when the error count changes", async () => {
+    const { rerender } = render(<DynamicFavicon isRunning={true} />);
+    await waitFor(() => expect(favicon.href).toContain("circle-play.ico"));
+    rerender(<DynamicFavicon isRunning={false} />);
+    await waitFor(() => expect(favicon.href).toContain("circle-check.ico"));
+
+    act(() => setErrors());
+    await waitFor(() => expect(favicon.href).toContain("circle-x.ico"));
+
+    act(() =>
+      store.set(
+        notebookAtom,
+        MockNotebook.notebookState({ cellData: { [id]: {} } }),
+      ),
+    );
+    await waitFor(() => expect(favicon.href).toContain("circle-check.ico"));
+  });
+
+  it("keeps the completion icon until its own reset timer expires", async () => {
+    vi.useFakeTimers();
+    const { rerender } = render(<DynamicFavicon isRunning={true} />);
+    await act(() => vi.advanceTimersByTimeAsync(1000));
+    rerender(<DynamicFavicon isRunning={false} />);
+    await act(() => vi.advanceTimersByTimeAsync(2000));
+    expect(favicon.href).toContain("circle-check.ico");
+    await act(() => vi.advanceTimersByTimeAsync(1000));
+    expect(favicon.href).toContain("favicon.ico");
+  });
+
+  it("cancels the reset timer when unmounted", async () => {
+    vi.useFakeTimers();
+    const { unmount } = render(<DynamicFavicon isRunning={true} />);
+    await act(() => vi.advanceTimersByTimeAsync(0));
+    expect(vi.getTimerCount()).toBe(1);
+    unmount();
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
   describe("notifications", () => {
     beforeEach(() => {
       vi.spyOn(document, "visibilityState", "get").mockReturnValue("hidden");
@@ -156,7 +249,7 @@ describe("DynamicFavicon", () => {
       global.Notification.permission = "granted";
     });
 
-    it("should send success notification when run completes without errors", () => {
+    it("should send success notification when run completes without errors", async () => {
       // @ts-expect-error ok in tests
       global.Notification = vi
         .fn()
@@ -172,12 +265,11 @@ describe("DynamicFavicon", () => {
 
       const { rerender } = render(<DynamicFavicon isRunning={true} />);
       rerender(<DynamicFavicon isRunning={false} />);
+      await waitFor(() => expect(Notification).toHaveBeenCalledOnce());
     });
 
-    it("should send error notification when run completes with errors", () => {
-      (useCellErrors as ReturnType<typeof vi.fn>).mockReturnValue([
-        { error: "mock error" },
-      ]);
+    it("should send error notification when run completes with errors", async () => {
+      setErrors();
 
       // @ts-expect-error ok in tests
       global.Notification = vi
@@ -194,6 +286,7 @@ describe("DynamicFavicon", () => {
 
       const { rerender } = render(<DynamicFavicon isRunning={true} />);
       rerender(<DynamicFavicon isRunning={false} />);
+      await waitFor(() => expect(Notification).toHaveBeenCalledOnce());
     });
 
     it("should not send notification when document is visible", () => {

package/src/components/editor/__tests__/package-alert.test.tsx

@@ -27,10 +27,12 @@ function mountPackageAlert({
   configuredManager,
   sandboxBackend,
   source = "kernel",
+  packages = ["polars"],
 }: {
   configuredManager: PackageManagerName;
   sandboxBackend: "uv" | "pixi" | null;
   source?: "kernel" | "server";
+  packages?: string[];
 }) {
   const store = createStore();
   const client = MockRequestClient.create();
@@ -55,7 +57,7 @@ function mountPackageAlert({
     packageAlert: {
       id: "missing-package",
       kind: "missing",
-      packages: ["polars"],
+      packages,
       isolated: true,
       source,
     },
@@ -97,6 +99,25 @@ it("keeps the configured package manager selectable outside a sandbox", () => {
   expect(managerSelect).not.toBeDisabled();
 });
 
+it("removes an extra absent from available choices through a keyboard menu item", async () => {
+  const client = mountPackageAlert({
+    configuredManager: "uv",
+    sandboxBackend: null,
+    packages: ["polars[dev]"],
+  });
+  fireEvent.keyDown(screen.getByRole("button", { name: "dev" }), {
+    key: "ArrowDown",
+  });
+  const remove = await screen.findByRole("menuitem", { name: "Remove dev" });
+  fireEvent.keyDown(remove, { key: "Enter" });
+  fireEvent.click(screen.getByTestId("install-packages-button"));
+  expect(client.sendInstallMissingPackages).toHaveBeenCalledExactlyOnceWith({
+    manager: "uv",
+    source: "kernel",
+    versions: { polars: "" },
+  });
+});
+
 it("installs a selected version in a Pixi sandbox", () => {
   const client = mountPackageAlert({
     configuredManager: "uv",

package/src/components/editor/actions/__tests__/pair-with-agent-modal.test.tsx

@@ -314,6 +314,30 @@ describe("PairWithAgentModal", () => {
     expect(screen.getByRole("tabpanel")).not.toHaveTextContent("--notebook");
   });
 
+  it("reverifies the target when host mode is enabled again for the same Session", async () => {
+    await act(async () => {
+      render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
+    });
+    expect(screen.getByRole("tabpanel")).toHaveTextContent(NOTEBOOK_ID);
+    act(() =>
+      store.set(pairConfigAtom, { ...PAIR_CONFIG, hostEnabled: false }),
+    );
+
+    const pending = new Deferred<Response>();
+    vi.mocked(fetch).mockImplementation((url) =>
+      String(url).endsWith("/api/sessions")
+        ? pending.promise
+        : responseFor(url),
+    );
+    act(() => store.set(pairConfigAtom, PAIR_CONFIG));
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeDisabled();
+    expect(screen.getByRole("tabpanel")).not.toHaveTextContent(NOTEBOOK_ID);
+
+    await act(async () => pending.resolve(await responseFor("/api/sessions")));
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeEnabled();
+    expect(screen.getByRole("tabpanel")).toHaveTextContent(NOTEBOOK_ID);
+  });
+
   it("keeps copying disabled when Notebook verification fails", async () => {
     vi.mocked(fetch).mockImplementation((url) =>
       String(url).endsWith("/api/sessions")

package/src/components/editor/actions/pair-with-agent-modal.tsx

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { CheckIcon, CopyIcon } from "lucide-react";
-import React, { useEffect, useState } from "react";
+import React, { useEffect, useMemo, useState } from "react";
 import { useAtomValue } from "jotai";
 import { Button } from "@/components/ui/button";
 import {
@@ -53,21 +53,20 @@ export const PairWithAgentModal: React.FC<{
   const config = useAtomValue(pairConfigAtom);
   const hostEnabled = config?.hostEnabled ?? false;
   const stableSessionId = useAtomValue(stableSessionIdAtom);
+  const lookupIdentity = useMemo(
+    () => ({ hostEnabled, stableSessionId }),
+    [hostEnabled, stableSessionId],
+  );
   const [lookup, setLookup] = useState<
     {
-      sessionId: string;
-    } & (
-      | { status: "ready"; notebookId: string | null }
-      | { status: "loading" | "error" }
-    )
+      identity: typeof lookupIdentity;
+    } & ({ status: "ready"; notebookId: string | null } | { status: "error" })
   >();
   useEffect(() => {
-    if (!hostEnabled || !stableSessionId) {
-      setLookup(undefined);
+    const { hostEnabled, stableSessionId: sessionId } = lookupIdentity;
+    if (!hostEnabled || !sessionId) {
       return;
     }
-    const sessionId = stableSessionId;
-    setLookup({ sessionId, status: "loading" });
     let cancelled = false;
     API.get<Record<string, { session_id: string; notebook_id: string | null }>>(
       "/sessions",
@@ -79,22 +78,26 @@ export const PairWithAgentModal: React.FC<{
           );
           setLookup(
             session
-              ? { sessionId, status: "ready", notebookId: session.notebook_id }
-              : { sessionId, status: "error" },
+              ? {
+                  identity: lookupIdentity,
+                  status: "ready",
+                  notebookId: session.notebook_id,
+                }
+              : { identity: lookupIdentity, status: "error" },
           );
         }
       })
       .catch(() => {
         if (!cancelled) {
-          setLookup({ sessionId, status: "error" });
+          setLookup({ identity: lookupIdentity, status: "error" });
         }
       });
     return () => {
       cancelled = true;
     };
-  }, [hostEnabled, stableSessionId]);
+  }, [lookupIdentity]);
   const currentLookup =
-    lookup?.sessionId === stableSessionId ? lookup : undefined;
+    lookup?.identity === lookupIdentity ? lookup : undefined;
   const verified = !hostEnabled || currentLookup?.status === "ready";
   const notebookId =
     hostEnabled && currentLookup?.status === "ready"

package/src/components/editor/actions/useRestartKernel.tsx

@@ -5,7 +5,7 @@ import { useImperativeModal } from "@/components/modal/ImperativeModal";
 import { AlertDialogDestructiveAction } from "@/components/ui/alert-dialog";
 import { connectionAtom } from "@/core/network/connection";
 import { useRequestClient } from "@/core/network/requests";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { reloadSafe } from "@/utils/reload-safe";
 
 export function useRestartKernel() {

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -106,9 +106,9 @@ export const AddCellWithAI: React.FC<{
   const runtimeManager = useRuntimeManager();
 
   const inputRef = useRef<ReactCodeMirrorRef>(null);
-  const submissionController = useRef(
-    new StagedCellSubmissionController(),
-  ).current;
+  const [submissionController] = useState(
+    () => new StagedCellSubmissionController(),
+  );
 
   const fileInputRef = useRef<HTMLInputElement>(null);
   const { files, addFiles, removeFile } = useFileState();

package/src/components/editor/ai/ai-completion-editor.tsx

@@ -408,6 +408,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
               </div>
             </Tooltip>
             <Button
+              aria-label="Discard AI completion"
               data-testid="decline-completion-button"
               variant="text"
               size="icon"

package/src/components/editor/alerts/__tests__/connection-notice.test.tsx

@@ -11,7 +11,7 @@ import { AppConfigSchema } from "@/core/config/config-schema";
 import { kernelStartupErrorAtom } from "@/core/errors/state";
 import { connectionAtom } from "@/core/network/connection";
 import { sandboxAtom, sandboxSyncAtom } from "@/core/packages/sandbox-state";
-import { WebSocketClosedReason, ConnectionState } from "@/core/websocket/types";
+import { WebSocketClosedReason, ConnectionState } from "@/core/session/types";
 import { chromeAtom } from "../../chrome/state";
 import { runtimeDetailsExpandedAtom } from "../../chrome/panels/runtime-details-state";
 import { ConnectionNotice } from "../connection-notice";

package/src/components/editor/alerts/connection-notice.tsx

@@ -8,7 +8,7 @@ import type { AppConfig } from "@/core/config/config-schema";
 import { connectionAtom } from "@/core/network/connection";
 import { useConnectionNotice } from "@/core/network/useConnectionNotice";
 import { sandboxAtom } from "@/core/packages/sandbox-state";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { useInterval } from "@/hooks/useInterval";
 import { cn } from "@/utils/cn";
 import { runtimeDetailsExpandedAtom } from "../chrome/panels/runtime-details-state";

package/src/components/editor/alerts/startup-logs-alert.tsx

@@ -51,6 +51,7 @@ export const StartupLogsAlert: React.FC = () => {
             {isRunning ? "Initializing..." : "Initialized"}
           </span>
           <Button
+            aria-label="Dismiss startup logs"
             variant="text"
             data-testid="remove-startup-logs-button"
             size="icon"

package/src/components/editor/app-container.tsx

@@ -8,7 +8,7 @@ import {
   type ConnectionStatus,
   WebSocketClosedReason,
   ConnectionState,
-} from "@/core/websocket/types";
+} from "@/core/session/types";
 import { cn } from "@/utils/cn";
 import { DynamicFavicon } from "./dynamic-favicon";
 import { StatusOverlay } from "./header/status";

package/src/components/editor/boundary/__tests__/ErrorBoundary.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 { ErrorBoundary } from "../ErrorBoundary";
 
 vi.mock("@/utils/copy", () => ({

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

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { DatabaseIcon, DiamondPlusIcon, PlusIcon } from "lucide-react";
-import { useState } from "react";
+import { useRef, useState } from "react";
 import { Button } from "@/components/editor/inputs/Inputs";
 import { MinimalHotkeys } from "@/components/shortcuts/renderShortcut";
 import {
@@ -16,8 +16,8 @@ import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages
 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 { cn } from "@/utils/cn";
 import { Tooltip } from "../../ui/tooltip";
 import { MarkdownIcon, PythonIcon } from "./code/icons";
@@ -36,7 +36,8 @@ export const CreateCellButton = ({
   const { createNewCell, addSetupCellIfDoesntExist } = useCellActions();
   const shortcut = `${oneClickShortcut}-Click`;
   const [open, setOpen] = useState(false);
-  const [justOpened, setJustOpened] = useState(false);
+  // Time the dropdown last opened, to ignore the click that opened it
+  const openedAt = useRef(0);
 
   const baseTooltipContent =
     getConnectionTooltip(connectionState) || tooltipContent;
@@ -83,9 +84,7 @@ export const CreateCellButton = ({
 
   const openDropdown = () => {
     setOpen(true);
-    setJustOpened(true);
-    // Allow interactions after a brief delay to prevent the dropdown items immediately being clicked
-    setTimeout(() => setJustOpened(false), 200);
+    openedAt.current = Date.now();
   };
 
   // We use onPointerDownCapture (not onPointerDown) to intercept events in
@@ -121,15 +120,14 @@ export const CreateCellButton = ({
   const handleOpenChange = (isOpen: boolean) => {
     setOpen(isOpen);
     if (isOpen) {
-      setJustOpened(true);
-      // Allow interactions after a brief delay
-      setTimeout(() => setJustOpened(false), 200);
+      openedAt.current = Date.now();
     }
   };
 
   const handleFirstItemClick = (e: React.MouseEvent) => {
-    // Hack to prevent the first item from being clicked when the dropdown is opened
-    if (justOpened) {
+    // Hack to prevent the first item from being clicked when the dropdown is
+    // opened: allow interactions only after a brief delay
+    if (Date.now() - openedAt.current < 200) {
       e.preventDefault();
       e.stopPropagation();
       return;

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

@@ -5,8 +5,8 @@ import type { 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 { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
 import { Button } from "../../ui/button";
@@ -37,6 +37,7 @@ export const DeleteButton = (props: {
   return (
     <Tooltip content={tooltipMsg} usePortal={false}>
       <Button
+        aria-label="Delete cell"
         variant="ghost"
         size="icon"
         onClick={onClick}

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/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__/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

@@ -41,7 +41,7 @@ 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";
@@ -572,7 +572,7 @@ const CellEditorInternal = ({
           <div className="absolute top-1 right-5">
             <LanguageToggles
               code={code}
-              editorView={editorViewRef.current}
+              editorViewRef={editorViewRef}
               currentLanguageAdapter={languageAdapter}
               onAfterToggle={afterToggleMarkdown}
             />