@marimo-team/islands 0.25.3-dev141 → 0.25.3-dev143

package/src/core/runtime/runtime.ts

@@ -1,5 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { NotebookId } from "@marimo-team/client/protocol";
 import { Deferred } from "@/utils/Deferred";
 import { Logger } from "@/utils/Logger";
 import { KnownQueryParams } from "../constants";
@@ -9,6 +10,8 @@ import { isStaticNotebook } from "../static/static-state";
 import { isWasm } from "../wasm/utils";
 import type { RuntimeConfig } from "./types";
 
+const knownQueryParams = new Set(Object.values(KnownQueryParams));
+
 export class RuntimeManager {
   private initialHealthyCheck = new Deferred<void>();
   private config: RuntimeConfig;
@@ -71,10 +74,7 @@ export class RuntimeManager {
     }
 
     for (const [key, value] of currentParams.entries()) {
-      if (
-        restrictToKnownQueryParams &&
-        !Object.values(KnownQueryParams).includes(key)
-      ) {
+      if (restrictToKnownQueryParams && !knownQueryParams.has(key)) {
         continue;
       }
       baseUrl.searchParams.set(key, value);
@@ -129,32 +129,32 @@ export class RuntimeManager {
     return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
-  /**
-   * Search params for a session connection: the base URL's params, merged
-   * with the current page's, plus the session id.
-   */
-  private getSessionSearchParams(sessionId: SessionId): URLSearchParams {
-    const baseUrl = new URL(this.config.url);
-    const searchParams = new URLSearchParams(baseUrl.search);
-
-    // Merge in current page's query parameters
-    const currentParams = new URLSearchParams(window.location.search);
-    currentParams.forEach((value, key) => {
-      // Don't override base URL params
-      if (!searchParams.has(key)) {
-        searchParams.set(key, value);
-      }
-    });
-
-    searchParams.set(KnownQueryParams.sessionId, sessionId);
-    return searchParams;
+  /** Keep notebook page parameters separate from connection options. */
+  private getSessionURL(path: string, sessionId: SessionId): URL {
+    const url = this.formatHttpURL({ path });
+    const notebookParams = new URLSearchParams(this.httpURL.search);
+    const pageParams = new URLSearchParams(window.location.search);
+    for (const key of pageParams.keys()) {
+      notebookParams.delete(key);
+    }
+    for (const [key, value] of pageParams) {
+      notebookParams.append(key, value);
+    }
+    // These parameters are internal and are not exposed by mo.query_params().
+    notebookParams.delete(KnownQueryParams.accessToken);
+    notebookParams.delete("refresh_token");
+    notebookParams.delete(KnownQueryParams.sessionId);
+    url.searchParams.set("query_params", notebookParams.toString());
+    url.searchParams.set(KnownQueryParams.sessionId, sessionId);
+    return url;
   }
 
   /**
    * The WebSocket URL of the runtime.
    */
   public getWsURL(sessionId: SessionId): URL {
-    return this.formatWsURL("/ws", this.getSessionSearchParams(sessionId));
+    const url = this.getSessionURL("/ws", sessionId);
+    return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
   /**
@@ -164,11 +164,7 @@ export class RuntimeManager {
    * (see `headers()`) instead of the URL.
    */
   public getSseURL(sessionId: SessionId): URL {
-    const url = this.formatHttpURL({
-      path: "/sse",
-      searchParams: this.getSessionSearchParams(sessionId),
-      restrictToKnownQueryParams: false,
-    });
+    const url = this.getSessionURL("/sse", sessionId);
     // The page URL can still carry access_token (merged in from the
     // current page's params, before auth cleanup strips it); never
     // forward it — a token in the URL leaks into proxy/server logs.
@@ -176,11 +172,24 @@ export class RuntimeManager {
     return url;
   }
 
+  /**
+   * The host API URL of a notebook, below which its operations are served.
+   * Requests send auth in headers (see `headers()`), never in the URL.
+   */
+  public getNotebookURL(notebookId: NotebookId): URL {
+    const url = this.formatHttpURL({
+      path: `/api/marimo/v1/notebooks/${encodeURIComponent(notebookId)}`,
+    });
+    url.searchParams.delete(KnownQueryParams.accessToken);
+    return url;
+  }
+
   /**
    * The WebSocket Sync URL of the runtime, for real-time updates.
    */
   public getWsSyncURL(sessionId: SessionId): URL {
-    return this.formatWsURL("/ws_sync", this.getSessionSearchParams(sessionId));
+    const url = this.getSessionURL("/ws_sync", sessionId);
+    return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
   /**

package/src/core/saving/save-component.tsx

@@ -21,7 +21,10 @@ import { Label } from "../../components/ui/label";
 import { useEvent } from "../../hooks/useEvent";
 import { Logger } from "../../utils/Logger";
 import { getCellConfigs, getNotebook, useNotebook } from "../cells/cells";
-import { withDocumentSave } from "../cells/document-changes";
+import {
+  flushDocumentChanges,
+  withDocumentSave,
+} from "../cells/document-changes";
 import { notebookCells } from "../cells/utils";
 import { formatAll } from "../codemirror/format";
 import { autoSaveConfigAtom } from "../config/config";
@@ -32,6 +35,7 @@ import { kioskModeAtom } from "../mode";
 import { connectionAtom } from "../network/connection";
 import { useRequestClient } from "../network/requests";
 import { ConnectionState } from "../websocket/types";
+import { isWasm } from "../wasm/utils";
 import { filenameAtom } from "./file-state";
 import { useFilename, useUpdateFilename } from "./filename";
 import { lastSavedNotebookAtom, needsSaveAtom } from "./state";
@@ -136,7 +140,7 @@ export function useSaveNotebook() {
         return;
       }
 
-      const savedNotebook = await withDocumentSave(async () => {
+      const save = async () => {
         // Grab the latest notebook state after formatting and after every
         // earlier document transaction has reached the server.
         const notebook = getNotebook();
@@ -157,7 +161,12 @@ export function useSaveNotebook() {
           persist: true,
         });
         return { names: cellNames, codes, configs, layout };
-      });
+      };
+      // TODO(MO-7955): Remove this conditional once all frontend saves use
+      // the host protocol. Host saves cannot recover edits from an editor snapshot.
+      const savedNotebook = isWasm()
+        ? await flushDocumentChanges().then(save)
+        : await withDocumentSave(save);
       setLastSavedNotebook(savedNotebook);
     }),
   );

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

@@ -0,0 +1,108 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { atom, Provider } from "jotai";
+import type React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { store, useJotaiEffect, waitFor } from "../jotai";
+
+const wrapper = ({ children }: { children: React.ReactNode }) => (
+  <Provider store={store}>{children}</Provider>
+);
+
+describe("waitFor function", () => {
+  it("should resolve when the atom satisfies the predicate", async () => {
+    const testAtom = atom(0);
+    store.set(testAtom, 10);
+
+    const result = await waitFor(testAtom, (value) => value === 10);
+    expect(result).toBe(10);
+  });
+
+  it("should resolve when the atom changes to satisfy the predicate", async () => {
+    const testAtom = atom(0);
+
+    setTimeout(() => store.set(testAtom, 15), 5);
+    setTimeout(() => store.set(testAtom, 20), 10);
+
+    const result = await waitFor(testAtom, (value) => value === 20);
+    expect(result).toBe(20);
+  });
+});
+
+describe("useJotaiEffect", () => {
+  it("calls the effect with the new and previous values", () => {
+    const testAtom = atom(0);
+    const effect = vi.fn();
+    renderHook(() => useJotaiEffect(testAtom, effect), { wrapper });
+
+    act(() => store.set(testAtom, 1));
+    act(() => store.set(testAtom, 2));
+
+    expect(effect.mock.calls).toEqual([
+      [1, 0],
+      [2, 1],
+    ]);
+  });
+
+  it("calls only the latest callback after a rerender", () => {
+    const testAtom = atom(0);
+    const first = vi.fn();
+    const second = vi.fn();
+    const { rerender } = renderHook(
+      ({ effect }) => useJotaiEffect(testAtom, effect),
+      { initialProps: { effect: first }, wrapper },
+    );
+
+    rerender({ effect: second });
+    act(() => store.set(testAtom, 1));
+
+    expect(first).not.toHaveBeenCalled();
+    expect(second).toHaveBeenCalledTimes(1);
+  });
+
+  it("subscribes once per rerender with an inline callback", () => {
+    const testAtom = atom(0);
+    const calls: number[] = [];
+    const { rerender } = renderHook(
+      () => useJotaiEffect(testAtom, (value) => calls.push(value)),
+      { wrapper },
+    );
+
+    rerender();
+    rerender();
+    act(() => store.set(testAtom, 1));
+
+    expect(calls).toEqual([1]);
+  });
+
+  it("stops listening to the previous atom when the atom changes", () => {
+    const firstAtom = atom(0);
+    const secondAtom = atom(0);
+    const effect = vi.fn();
+    const { rerender } = renderHook(
+      ({ target }) => useJotaiEffect(target, effect),
+      { initialProps: { target: firstAtom }, wrapper },
+    );
+
+    rerender({ target: secondAtom });
+    act(() => store.set(firstAtom, 1));
+    expect(effect).not.toHaveBeenCalled();
+
+    act(() => store.set(secondAtom, 1));
+    expect(effect).toHaveBeenCalledTimes(1);
+  });
+
+  it("unsubscribes on unmount", () => {
+    const testAtom = atom(0);
+    const effect = vi.fn();
+    const { unmount } = renderHook(() => useJotaiEffect(testAtom, effect), {
+      wrapper,
+    });
+
+    unmount();
+    act(() => store.set(testAtom, 1));
+
+    expect(effect).not.toHaveBeenCalled();
+  });
+});

package/src/core/state/jotai.ts

@@ -42,7 +42,7 @@ export function useJotaiEffect<T>(
   const store = useStore();
   useEffect(() => {
     let prevValue = store.get(atom);
-    store.sub(atom, () => {
+    return store.sub(atom, () => {
       const value = store.get(atom);
       effect(value, prevValue);
       prevValue = value;

package/src/core/static/__tests__/files.test.ts

@@ -654,19 +654,47 @@ describe("error handling and edge cases", () => {
     unpatch();
   });
 
-  it("should handle malformed URLs gracefully", async () => {
-    const virtualFiles = {};
-    const unpatch = patchFetch(virtualFiles);
+  it("forwards relative URLs and options to the original fetch", async () => {
+    const response = new Response("relative content");
+    const originalFetch = vi.spyOn(window, "fetch").mockResolvedValue(response);
+    const unpatch = patchFetch({});
+    const options = { headers: { Accept: "text/plain" } };
+
+    try {
+      expect(await window.fetch("relative-file.txt", options)).toBe(response);
+      expect(originalFetch).toHaveBeenCalledExactlyOnceWith(
+        "relative-file.txt",
+        options,
+      );
+    } finally {
+      unpatch();
+      originalFetch.mockRestore();
+    }
+  });
 
-    // Mock Logger.error to avoid test output
+  it("preserves the original fetch rejection for malformed URLs", async () => {
+    const error = new TypeError("Invalid URL");
+    const originalFetch = vi.spyOn(window, "fetch").mockRejectedValue(error);
     const loggerSpy = vi
       .spyOn(console, "error")
       .mockImplementation(Functions.NOOP);
-
-    // This should catch the error and fallback to original fetch
-    await window.fetch("not-a-valid-url").catch(() => undefined);
-
-    unpatch();
-    loggerSpy.mockRestore();
+    const unpatch = patchFetch({});
+    const options = { headers: { Accept: "text/plain" } };
+
+    try {
+      await expect(window.fetch("http://[", options)).rejects.toBe(error);
+      expect(originalFetch).toHaveBeenCalledExactlyOnceWith(
+        "http://[",
+        options,
+      );
+      expect(loggerSpy).toHaveBeenCalledWith(
+        "Error parsing URL",
+        expect.any(TypeError),
+      );
+    } finally {
+      unpatch();
+      originalFetch.mockRestore();
+      loggerSpy.mockRestore();
+    }
   });
 });

package/src/core/static/files.ts

@@ -125,6 +125,8 @@ function withoutLeadingDot(path: string): string {
 /**
  * Resolve a URL to a blob URL if it's a virtual file, for use with dynamic import().
  * Unlike fetch, import() can't be patched, so we need to convert data URLs to blob URLs.
+ * The caller owns a returned URL that differs from the input and must revoke it
+ * once no module consumers or their cleanup callbacks need it.
  *
  * @returns The original URL if not a virtual file, or a blob URL if it is
  */

package/src/core/wasm/__tests__/bridge.test.ts

@@ -8,16 +8,36 @@ const {
   mockBridge,
   mockNotebookReadFile,
   mockReadNotebook,
-  mockSaveNotebook,
+  mockWriteNotebook,
+  mockNotebookCall,
+  mockSetNotebookLayout,
+  mockFlush,
+  parentRequests,
   mockBootstrap,
   rpcListeners,
+  mockNotebookSubscribe,
+  mockNotebookUnsubscribe,
 } = vi.hoisted(() => ({
   mockBridge: vi.fn(),
   mockNotebookReadFile: vi.fn(),
   mockReadNotebook: vi.fn(),
-  mockSaveNotebook: vi.fn(),
+  mockWriteNotebook: vi.fn(),
+  mockNotebookCall: vi.fn(),
+  mockSetNotebookLayout: vi.fn(),
+  mockFlush: vi.fn(),
+  parentRequests: {} as {
+    persistNotebook?: (snapshot: {
+      filename: string;
+      contents: string;
+    }) => Promise<void>;
+  },
   mockBootstrap: vi.fn(),
-  rpcListeners: {} as Record<string, () => void>,
+  rpcListeners: {} as Record<
+    string,
+    (message?: { id: string; event: string }) => void
+  >,
+  mockNotebookSubscribe: vi.fn(),
+  mockNotebookUnsubscribe: vi.fn(),
 }));
 
 // Mock browser globals before any imports
@@ -41,22 +61,31 @@ class MockURL extends URL {
 vi.stubGlobal("URL", MockURL);
 
 vi.mock("@/core/wasm/rpc", () => ({
-  getWorkerRPC: () => ({
-    send: { bootstrap: mockBootstrap },
-    proxy: {
-      request: {
-        bridge: mockBridge,
-        startSession: vi.fn(),
-        readFile: vi.fn(),
-        readNotebook: mockReadNotebook,
-        saveNotebook: mockSaveNotebook,
+  getWorkerRPC: (_worker: Worker, handlers?: typeof parentRequests) => {
+    Object.assign(parentRequests, handlers);
+    return {
+      send: { bootstrap: mockBootstrap },
+      proxy: {
+        request: {
+          bridge: mockBridge,
+          startSession: vi.fn(),
+          readFile: vi.fn(),
+          readNotebook: mockReadNotebook,
+          writeNotebook: mockWriteNotebook,
+          notebookSubscribe: mockNotebookSubscribe,
+          notebookUnsubscribe: mockNotebookUnsubscribe,
+          notebookCall: mockNotebookCall,
+          setNotebookLayout: mockSetNotebookLayout,
+        },
       },
-      send: { consumerReady: vi.fn() },
-    },
-    addMessageListener: (event: string, cb: () => void) => {
-      rpcListeners[event] = cb;
-    },
-  }),
+      addMessageListener: (
+        event: string,
+        cb: (typeof rpcListeners)[string],
+      ) => {
+        rpcListeners[event] = cb;
+      },
+    };
+  },
 }));
 
 vi.mock("@/core/meta/globals", () => ({
@@ -72,14 +101,25 @@ vi.mock("@/core/wasm/store", () => ({
   notebookFileStore: { readFile: mockNotebookReadFile, saveFile: vi.fn() },
 }));
 
+vi.mock("@/core/cells/document-changes", () => ({
+  flushDocumentChanges: mockFlush,
+  documentTransactionMiddleware: vi.fn(),
+}));
+
 // Import after all mocks are set up
 import { store } from "@/core/state/jotai";
 import { initialModeAtom } from "@/core/mode";
+import { fixtureEvents } from "@/core/host/__tests__/fixtures";
+import { createNotebookReplica } from "@/core/host";
+import { notebookReplicaAtom } from "@/core/websocket/state";
+import { createEditorBackend } from "@/core/websocket/backend";
+import { fallbackFileStore, notebookFileStore } from "../store";
 import { getWasmWorkerName, PyodideBridge } from "../bridge";
 
-// Access INSTANCE once at module level so the constructor runs (and
-// addMessageListener populates rpcListeners) before any test executes.
-void PyodideBridge.INSTANCE;
+beforeEach(() => {
+  Reflect.deleteProperty(window, "_marimo_private_PyodideBridge");
+  void PyodideBridge.INSTANCE;
+});
 
 describe("PyodideBridge.readCode", () => {
   beforeEach(() => {
@@ -161,6 +201,7 @@ describe("PyodideBridge.exportAsScript", () => {
 });
 
 describe("PyodideBridge.sendSave", () => {
+  let disconnect: () => void;
   const request = {
     cellIds: [cellId("cell-1")],
     codes: ['value = "NEW"'],
@@ -171,51 +212,173 @@ describe("PyodideBridge.sendSave", () => {
     persist: true,
   };
 
-  beforeEach(() => {
+  beforeEach(async () => {
     vi.clearAllMocks();
     store.set(initialModeAtom, "edit");
     rpcListeners.initialized();
-    mockSaveNotebook.mockResolvedValue(null);
-    mockReadNotebook.mockResolvedValue(
-      "import marimo\napp = marimo.App()\n# NEW",
+    mockFlush.mockResolvedValue(undefined);
+    mockWriteNotebook.mockResolvedValue(undefined);
+    mockNotebookUnsubscribe.mockResolvedValue(undefined);
+    const document = fixtureEvents("notebook-events.sse").find(
+      (event) => event.event === "document",
+    )?.data.document;
+    mockNotebookCall.mockResolvedValue({
+      status: 200,
+      body: { document, cells: {} },
+    });
+    mockNotebookSubscribe.mockImplementation(async ({ id }) => {
+      for (const event of fixtureEvents("notebook-events.sse")) {
+        rpcListeners.notebookEvent({ id, event: JSON.stringify(event) });
+      }
+    });
+    const replica = createNotebookReplica(store);
+    store.set(notebookReplicaAtom, replica);
+    const backend = createEditorBackend({
+      static: false,
+      transportType: "sse",
+      url: () => "http://localhost/",
+      headers: () => ({}),
+      notebookUrl: (id) => new URL(`http://localhost/notebooks/${id}`),
+      waitUntilReady: async () => undefined,
+      checkHealth: async () => true,
+    });
+    disconnect = replica.connect(backend, { onNotification: vi.fn() });
+    await vi.waitFor(() =>
+      expect(store.get(replica.notebookAtom)?.id).toBe("n1"),
     );
   });
 
   afterEach(() => {
+    disconnect();
+    store.set(notebookReplicaAtom, null);
     store.set(initialModeAtom, undefined);
   });
 
-  it("waits for the session save before generating exports", async () => {
-    const mainSave = new Deferred<void>();
-    mockSaveNotebook
-      .mockResolvedValueOnce(null)
-      .mockReturnValueOnce(mainSave.promise);
+  it("saves through the connected backend using the stream's notebook ID", async () => {
+    const bridge = PyodideBridge.INSTANCE;
+    const createBackend = vi.spyOn(bridge, "createNotebookBackend");
+    await bridge.sendSave(request);
+    expect(createBackend).not.toHaveBeenCalled();
+    expect(mockNotebookCall).toHaveBeenCalledWith(
+      expect.objectContaining({
+        operation: "notebook.save",
+        params: { notebook_id: "n1" },
+        body: {},
+      }),
+    );
+  });
 
-    await PyodideBridge.INSTANCE.sendSave(request);
-    expect(mockSaveNotebook).toHaveBeenCalledTimes(2);
-    expect(mockSaveNotebook).toHaveBeenNthCalledWith(1, request);
-    expect(mockSaveNotebook).toHaveBeenNthCalledWith(2, request);
+  it.each(["unmounted", "not ready", "disconnected"])(
+    "does not save when the notebook is %s",
+    async (state) => {
+      if (state === "unmounted") {
+        store.set(notebookReplicaAtom, null);
+      } else if (state === "not ready") {
+        store.set(notebookReplicaAtom, createNotebookReplica(store));
+      } else {
+        disconnect();
+      }
+      await expect(PyodideBridge.INSTANCE.sendSave(request)).rejects.toThrow(
+        "The notebook is not connected",
+      );
+      expect(mockNotebookCall).not.toHaveBeenCalled();
+    },
+  );
 
-    const markdown = PyodideBridge.INSTANCE.exportAsMarkdown({
-      download: false,
-    });
-    const script = PyodideBridge.INSTANCE.exportAsScript({
-      download: false,
-    });
+  it("flushes edits before asking the host to save without an editor snapshot", async () => {
+    const flushed = new Deferred<void>();
+    mockFlush.mockReturnValueOnce(flushed.promise);
+    const saving = PyodideBridge.INSTANCE.sendSave(request);
     await Promise.resolve();
-    expect(mockBridge).not.toHaveBeenCalled();
+    expect(mockNotebookCall).not.toHaveBeenCalled();
+    flushed.resolve();
+    await saving;
+    expect(mockSetNotebookLayout).toHaveBeenCalledWith(null);
+    expect(mockNotebookCall).toHaveBeenCalledWith(
+      expect.objectContaining({ operation: "notebook.save", body: {} }),
+    );
+    expect(mockWriteNotebook).not.toHaveBeenCalled();
+  });
 
-    mainSave.resolve();
-    await Promise.all([markdown, script]);
-    expect(mockBridge).toHaveBeenCalledWith({
-      functionName: "export_markdown",
-      payload: { download: false },
-    });
-    expect(mockBridge).toHaveBeenCalledWith({
-      functionName: "export_script",
-      payload: { download: false },
+  it("does not save when flushing fails", async () => {
+    mockFlush.mockRejectedValueOnce(new Error("unsynchronized edit"));
+    await expect(PyodideBridge.INSTANCE.sendSave(request)).rejects.toThrow(
+      "unsynchronized edit",
+    );
+    expect(mockNotebookCall).not.toHaveBeenCalled();
+  });
+
+  it("reports host persistence failures", async () => {
+    mockNotebookCall.mockResolvedValueOnce({
+      status: 409,
+      body: {
+        type: "tag:marimo.io,2026:host/conflict",
+        title: "Conflict",
+        status: 409,
+        detail: "storage full",
+      },
     });
+    await expect(PyodideBridge.INSTANCE.sendSave(request)).rejects.toThrow(
+      "storage full",
+    );
   });
+
+  it("persists exactly the host's content before acknowledging storage", async () => {
+    const durable = new Deferred<void>();
+    const snapshot = { filename: "notebook.py", contents: "# host revision" };
+    mockWriteNotebook.mockReturnValueOnce(durable.promise);
+    const writing = parentRequests.persistNotebook!(snapshot);
+    await Promise.resolve();
+    expect(mockWriteNotebook).toHaveBeenCalledWith(snapshot);
+    expect(notebookFileStore.saveFile).not.toHaveBeenCalled();
+    durable.resolve();
+    await writing;
+    expect(notebookFileStore.saveFile).toHaveBeenCalledWith(snapshot.contents);
+    expect(fallbackFileStore.saveFile).toHaveBeenCalledWith(snapshot.contents);
+  });
+
+  it("propagates storage failure without updating the URL", async () => {
+    mockWriteNotebook.mockRejectedValueOnce(new Error("IndexedDB failed"));
+    await expect(
+      parentRequests.persistNotebook!({
+        filename: "notebook.py",
+        contents: "# host",
+      }),
+    ).rejects.toThrow("IndexedDB failed");
+    expect(notebookFileStore.saveFile).not.toHaveBeenCalled();
+  });
+});
+
+it("cancels a pending notebook subscription before its RPC response", async () => {
+  vi.clearAllMocks();
+  store.set(initialModeAtom, "read");
+  rpcListeners.initialized();
+  const pending = new Deferred<void>();
+  mockNotebookSubscribe.mockReturnValue(pending.promise);
+  mockNotebookUnsubscribe.mockResolvedValue(undefined);
+  const controller = new AbortController();
+  try {
+    const backend = await PyodideBridge.INSTANCE.createNotebookBackend();
+    const following = backend.transport(
+      { cursor: null, signal: controller.signal },
+      { open: vi.fn(), message: vi.fn() },
+    );
+    await vi.waitFor(() =>
+      expect(mockNotebookSubscribe).toHaveBeenCalledOnce(),
+    );
+    const [{ id }] = mockNotebookSubscribe.mock.calls[0];
+    controller.abort();
+    await following;
+    expect(mockNotebookUnsubscribe.mock.calls).toEqual([[id]]);
+
+    pending.resolve();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(mockNotebookUnsubscribe.mock.calls).toEqual([[id]]);
+  } finally {
+    controller.abort();
+    pending.resolve();
+    store.set(initialModeAtom, undefined);
+  }
 });
 
 describe("getWasmWorkerName", () => {