@marimo-team/islands 0.25.3-dev14 → 0.25.3-dev140

package/src/core/runtime/__tests__/createWsUrl.test.ts

@@ -13,7 +13,9 @@ describe("RuntimeManager.getWsURL", () => {
     });
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
-    expect(url.toString()).toBe("wss://marimo.app/ws?session_id=1234");
+    expect(url.toString()).toBe(
+      "wss://marimo.app/ws?query_params=&session_id=1234",
+    );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
 
@@ -24,7 +26,9 @@ describe("RuntimeManager.getWsURL", () => {
     });
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
-    expect(url.toString()).toBe("ws://marimo.app/ws?session_id=1234");
+    expect(url.toString()).toBe(
+      "ws://marimo.app/ws?query_params=&session_id=1234",
+    );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
 
@@ -35,7 +39,9 @@ describe("RuntimeManager.getWsURL", () => {
     });
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
-    expect(url.toString()).toBe("ws://marimo.app/nested/ws?session_id=1234");
+    expect(url.toString()).toBe(
+      "ws://marimo.app/nested/ws?query_params=&session_id=1234",
+    );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
 
@@ -48,7 +54,7 @@ describe("RuntimeManager.getWsURL", () => {
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
     expect(url.toString()).toBe(
-      "ws://marimo.app/nested/ws?foo=bar&file=test.py&session_id=1234",
+      "ws://marimo.app/nested/ws?foo=bar&file=test.py&query_params=foo%3Dbar%26file%3Dtest.py&session_id=1234",
     );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
@@ -62,11 +68,75 @@ describe("RuntimeManager.getWsURL", () => {
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
     expect(url.toString()).toBe(
-      "ws://marimo.app/ws?file=test.py&search=test&other=last&session_id=1234",
+      "ws://marimo.app/ws?file=test.py&query_params=file%3Dtest.py%26search%3Dtest%26other%3Dlast&session_id=1234",
     );
     expect(url.searchParams.get("file")).toBe("test.py");
-    expect(url.searchParams.get("search")).toBe("test");
-    expect(url.searchParams.get("other")).toBe("last");
+    expect(url.searchParams.get("query_params")).toBe(
+      "file=test.py&search=test&other=last",
+    );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
 });
+
+describe.each(["getWsURL", "getSseURL"] as const)(
+  "RuntimeManager.%s",
+  (method) => {
+    it("preserves runtime parameters and repeated page values", () => {
+      window.history.pushState({}, "", "/?tag=page-one&tag=page-two");
+      const runtime = new RuntimeManager(
+        {
+          url: `${window.location.origin}/nested?proxy_token=secret&tag=base&base=one&base=two`,
+          lazy: true,
+        },
+        true,
+      );
+      const url = runtime[method]("session" as SessionId);
+      expect([...url.searchParams]).toEqual([
+        ["proxy_token", "secret"],
+        ["tag", "base"],
+        ["base", "one"],
+        ["base", "two"],
+        [
+          "query_params",
+          "proxy_token=secret&base=one&base=two&tag=page-one&tag=page-two",
+        ],
+        ["session_id", "session"],
+      ]);
+    });
+
+    it("keeps notebook parameters out of the stream options", () => {
+      const notebookParams = new URLSearchParams([
+        ["types", "scatter"],
+        ["cursor", "page-2"],
+        ["run_id", "experiment"],
+        ["query_params", "nested=value&other=1"],
+        ["tag", "one"],
+        ["tag", "two"],
+        ["empty", ""],
+        ["text", "a+b & café"],
+        ["file", "test.py"],
+        ["kiosk", "true"],
+      ]);
+      window.history.pushState(
+        {},
+        "",
+        `/?${notebookParams}&session_id=page-session`,
+      );
+      const runtime = new RuntimeManager(
+        { url: window.location.origin, lazy: true },
+        true,
+      );
+      const url = runtime[method]("actual-session" as SessionId);
+
+      expect(Object.fromEntries(url.searchParams)).toEqual({
+        file: "test.py",
+        kiosk: "true",
+        session_id: "actual-session",
+        query_params: notebookParams.toString(),
+      });
+      expect([
+        ...new URLSearchParams(url.searchParams.get("query_params") ?? ""),
+      ]).toEqual([...notebookParams]);
+    });
+  },
+);

package/src/core/runtime/__tests__/runtime.test.ts

@@ -2,11 +2,7 @@
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { Mocks } from "@/__mocks__/common";
-import {
-  generateSessionId,
-  isSessionId,
-  type SessionId,
-} from "@/core/kernel/session";
+import type { SessionId } from "@/core/kernel/session";
 import { Logger } from "@/utils/Logger";
 import { RuntimeManager } from "../runtime";
 import type { RuntimeConfig } from "../types";
@@ -56,38 +52,6 @@ describe("RuntimeManager", () => {
   });
 
   describe("getWsURL", () => {
-    it("uses host notebook events with the existing session and auth parameters", () => {
-      const original = window.location.href;
-      window.history.replaceState({}, "", "/?types=run");
-      try {
-        const runtime = new RuntimeManager({
-          ...mockConfig,
-          url: "https://example.com/nested?file=notebook.py",
-          notebookId: "nb_123",
-        });
-        const sessionId = "s_123456";
-        if (!isSessionId(sessionId)) {
-          throw new Error("Invalid test session ID");
-        }
-        const url = runtime.getWsURL(sessionId);
-        expect({
-          endpoint: `${url.origin}${url.pathname}`,
-          params: Object.fromEntries(url.searchParams),
-        }).toEqual({
-          endpoint:
-            "wss://example.com/nested/api/marimo/v1/notebooks/nb_123/events/ws",
-          params: {
-            file: "notebook.py",
-            types: "marimo",
-            session_id: sessionId,
-            access_token: "test-token",
-          },
-        });
-      } finally {
-        window.history.replaceState({}, "", original);
-      }
-    });
-
     it("should return WebSocket URL with session ID for https", () => {
       const runtime = new RuntimeManager(mockConfig);
       const sessionId = "1234" as SessionId;
@@ -256,22 +220,24 @@ describe("RuntimeManager", () => {
     });
   });
 
-  describe("getSseURL", () => {
-    it.each([undefined, "nb_123"])(
-      "uses the legacy session stream when notebookId is %s",
-      (notebookId) => {
-        const runtime = new RuntimeManager({
-          ...mockConfig,
-          url: "https://example.com/nested?file=notebook.py",
-          notebookId,
-        });
-        const sessionId = generateSessionId();
-        expect(runtime.getSseURL(sessionId).toString()).toBe(
-          `https://example.com/nested/sse?file=notebook.py&session_id=${sessionId}`,
-        );
-      },
-    );
+  describe("getNotebookURL", () => {
+    it("addresses the notebook below the host API without the auth token", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com/base",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getNotebookURL("n/1");
 
+      expect(url.pathname).toBe("/base/api/marimo/v1/notebooks/n%2F1");
+      expect(url.toString()).not.toContain("my-secret-token");
+    });
+  });
+
+  describe("getSseURL", () => {
     it("should never put the auth token in the URL, even cross-origin", () => {
       // SSE requests can send headers, so auth travels in the
       // Authorization header; a token in the URL would leak into
@@ -300,7 +266,8 @@ describe("RuntimeManager", () => {
         const url = runtime.getSseURL("s_123" as SessionId);
 
         expect(url.searchParams.get("access_token")).toBeNull();
-        expect(url.searchParams.get("foo")).toBe("bar");
+        expect(url.searchParams.get("query_params")).toBe("foo=bar");
+        expect(url.toString()).not.toContain("page-token");
       } finally {
         window.history.pushState({}, "", original);
       }
@@ -720,9 +687,16 @@ describe("RuntimeManager", () => {
       expect(httpUrl.searchParams.get("base_param")).toBe("existing");
 
       // Should preserve all window-level query params (including custom ones)
-      expect(wsUrl.searchParams.get("custom_param")).toBe("value123");
-      expect(wsUrl.searchParams.get("user_token")).toBe("abc");
-      expect(wsUrl.searchParams.get("theme")).toBe("dark");
+      expect(
+        Object.fromEntries(
+          new URLSearchParams(wsUrl.searchParams.get("query_params") ?? ""),
+        ),
+      ).toEqual({
+        base_param: "existing",
+        custom_param: "value123",
+        user_token: "abc",
+        theme: "dark",
+      });
       expect(httpUrl.searchParams.get("custom_param")).toBe("value123");
       expect(httpUrl.searchParams.get("user_token")).toBe("abc");
       expect(httpUrl.searchParams.get("theme")).toBe("dark");
@@ -781,7 +755,7 @@ describe("RuntimeManager", () => {
       );
     });
 
-    it("should throw when creating WebSocket URLs from blob URLs", () => {
+    it("warns when creating WebSocket URLs from blob URLs", () => {
       const runtime = new RuntimeManager({
         url: "blob:https://example.com/12345678-1234-1234-1234-123456789abc",
         lazy: true,
@@ -789,7 +763,7 @@ describe("RuntimeManager", () => {
       const sessionId = "test" as SessionId;
 
       expect(runtime.getWsURL(sessionId).toString()).toMatchInlineSnapshot(
-        `"blob:https://example.com/12345678-1234-1234-1234-123456789abc?session_id=test"`,
+        `"blob:https://example.com/12345678-1234-1234-1234-123456789abc?query_params=&session_id=test"`,
       );
       expect(Logger.warn).toHaveBeenCalledOnce();
     });

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,7 +10,7 @@ import { isStaticNotebook } from "../static/static-state";
 import { isWasm } from "../wasm/utils";
 import type { RuntimeConfig } from "./types";
 
-const HOST_EVENT_TYPES = ["marimo"];
+const knownQueryParams = new Set(Object.values(KnownQueryParams));
 
 export class RuntimeManager {
   private initialHealthyCheck = new Deferred<void>();
@@ -73,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);
@@ -131,43 +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 {
-    const searchParams = this.getSessionSearchParams(sessionId);
-    // TODO(MO-8039): Rediscover hosts and resolve notebooks on reconnect.
-    const notebookId = this.config.notebookId;
-    const path = notebookId
-      ? `/api/marimo/v1/notebooks/${encodeURIComponent(notebookId)}/events/ws`
-      : "/ws";
-    const url = this.formatWsURL(path, searchParams);
-    if (notebookId) {
-      // TODO(MO-8034): Separate host options from notebook query params.
-      url.searchParams.set("types", HOST_EVENT_TYPES.join(","));
-    }
-    return url;
+    const url = this.getSessionURL("/ws", sessionId);
+    return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
   /**
@@ -177,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.
@@ -189,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/runtime/types.ts

@@ -1,8 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 export interface RuntimeConfig {
-  /** The host notebook for WebSocket events. SSE uses the legacy session stream. */
-  readonly notebookId?: string;
   /**
    * The URL of the runtime server.
    */

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

@@ -11,6 +11,9 @@ const {
   mockSaveNotebook,
   mockBootstrap,
   rpcListeners,
+  mockNotebookInfo,
+  mockNotebookSubscribe,
+  mockNotebookUnsubscribe,
 } = vi.hoisted(() => ({
   mockBridge: vi.fn(),
   mockNotebookReadFile: vi.fn(),
@@ -18,6 +21,9 @@ const {
   mockSaveNotebook: vi.fn(),
   mockBootstrap: vi.fn(),
   rpcListeners: {} as Record<string, () => void>,
+  mockNotebookInfo: vi.fn(),
+  mockNotebookSubscribe: vi.fn(),
+  mockNotebookUnsubscribe: vi.fn(),
 }));
 
 // Mock browser globals before any imports
@@ -50,8 +56,10 @@ vi.mock("@/core/wasm/rpc", () => ({
         readFile: vi.fn(),
         readNotebook: mockReadNotebook,
         saveNotebook: mockSaveNotebook,
+        notebookInfo: mockNotebookInfo,
+        notebookSubscribe: mockNotebookSubscribe,
+        notebookUnsubscribe: mockNotebookUnsubscribe,
       },
-      send: { consumerReady: vi.fn() },
     },
     addMessageListener: (event: string, cb: () => void) => {
       rpcListeners[event] = cb;
@@ -218,6 +226,39 @@ describe("PyodideBridge.sendSave", () => {
   });
 });
 
+it("cancels a pending notebook subscription before its RPC response", async () => {
+  vi.clearAllMocks();
+  store.set(initialModeAtom, "read");
+  rpcListeners.initialized();
+  const pending = new Deferred<void>();
+  mockNotebookInfo.mockResolvedValue({ id: "n1" });
+  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", () => {
   afterEach(() => {
     delete (window as unknown as { __MARIMO_HAS_WASM_CONTROLLER__?: boolean })