@marimo-team/islands 0.25.3-dev12 → 0.25.3-dev121

package/src/core/host/wasm.ts

@@ -0,0 +1,92 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  NotebookClient,
+  StreamStopped,
+  type unstable_NotebookCall,
+} from "@marimo-team/client";
+import type { NotebookBackend } from "./backend";
+
+/** The worker owns the notebook; this interface only carries commands and events. */
+export interface WasmNotebookBridge {
+  readonly notebookId: string;
+  subscribe(
+    cursor: string | null,
+    emit: (event: unknown) => void,
+  ): Promise<() => void>;
+  /** Deliver a call to the worker and resolve with its status and body. */
+  call(call: unstable_NotebookCall): Promise<unknown>;
+}
+
+/** Stop waiting on cancellation; a command already sent may still commit. */
+function abortable<T>(
+  send: () => Promise<T>,
+  signal?: AbortSignal,
+): Promise<T> {
+  if (!signal) {
+    return send();
+  }
+  signal.throwIfAborted();
+  return new Promise((resolve, reject) => {
+    const abort = () => reject(signal.reason);
+    signal.addEventListener("abort", abort, { once: true });
+    send()
+      .then(resolve, reject)
+      .finally(() => signal.removeEventListener("abort", abort));
+  });
+}
+
+export function createWasmNotebookBackend(
+  bridge: WasmNotebookBridge,
+): NotebookBackend {
+  return {
+    transport: async ({ cursor, signal }, { open, message }) => {
+      if (signal.aborted) {
+        return;
+      }
+      let unsubscribe: (() => void) | undefined;
+      let finish!: () => void;
+      const closed = new Promise<void>((resolve) => {
+        finish = resolve;
+      });
+      const abort = () => {
+        unsubscribe?.();
+        finish();
+      };
+      signal.addEventListener("abort", abort, { once: true });
+      try {
+        open();
+        // Keep the pending subscription so a close during startup still releases it.
+        await abortable(
+          () =>
+            bridge
+              .subscribe(cursor, (event) => {
+                if (!signal.aborted) {
+                  message(event);
+                }
+              })
+              .then((close) => {
+                unsubscribe = close;
+                if (signal.aborted) {
+                  close();
+                }
+              }),
+          signal,
+        );
+        await closed;
+      } catch (error) {
+        if (!signal.aborted) {
+          unsubscribe?.();
+          throw new StreamStopped("Could not follow the WASM notebook", {
+            cause: error,
+          });
+        }
+      } finally {
+        signal.removeEventListener("abort", abort);
+      }
+    },
+    client: NotebookClient.unstable_over((call) => bridge.call(call), {
+      notebookId: bridge.notebookId,
+    }),
+  };
+}

package/src/core/hotkeys/actions.ts

@@ -2,6 +2,7 @@
 import { atom, useAtomValue, useSetAtom } from "jotai";
 import useEvent from "react-use-event-hook";
 import type { HotkeyAction } from "@/core/hotkeys/hotkeys";
+import { Objects } from "@/utils/objects";
 
 /**
  * Map of registered keyboard shortcuts and their callbacks.
@@ -29,10 +30,7 @@ export function useSetRegisteredAction() {
       },
     ),
     unregisterAction: useEvent((shortcut: HotkeyAction) => {
-      set((actions) => {
-        const { [shortcut]: unused, ...rest } = actions;
-        return rest;
-      });
+      set((actions) => Objects.omit(actions, [shortcut]));
     }),
   };
 }

package/src/core/islands/__tests__/bootstrap.test.ts

@@ -3,8 +3,6 @@
 import { createStore } from "jotai";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { RuntimeState } from "@/core/kernel/RuntimeState";
-import { attachmentsAtom } from "@/core/attachments/state";
-import { Logger } from "@/utils/Logger";
 import { initializeIslands } from "../bootstrap";
 import type { IslandsPyodideBridge } from "../bridge";
 import { islandsHydratedAtom, islandsPendingInitialRunsAtom } from "../state";
@@ -50,30 +48,6 @@ describe("initializeIslands", () => {
     expect(store.get(islandsHydratedAtom)).toBe(false);
     store.set(islandsPendingInitialRunsAtom, new Set([1, 2]));
 
-    const errors = vi.spyOn(Logger, "error");
-    consumeMessage?.({
-      sessionGeneration: 1,
-      message: JSON.stringify({
-        op: "attachments",
-        data: {
-          op: "attachments",
-          attachments: [
-            {
-              id: "a1",
-              kind: "agent",
-              name: "Pi",
-              since: "2026-10-14T00:00:00Z",
-              notebook_id: "nb-test",
-              identity: "pi-thread",
-            },
-          ],
-        },
-      }) as IslandsKernelMessage["message"],
-    });
-    expect(errors).not.toHaveBeenCalled();
-    expect(store.get(attachmentsAtom)).toEqual([]);
-    expect(store.get(islandsPendingInitialRunsAtom)).toEqual(new Set([1, 2]));
-
     consumeMessage?.(completedRun(1));
     expect(store.get(islandsHydratedAtom)).toBe(false);
 

package/src/core/islands/bootstrap.ts

@@ -209,7 +209,6 @@ function handleMessage(
       case "cache-info":
       case "kernel-startup-error":
       case "notebook-document-transaction":
-      case "attachments":
         return;
 
       case "completed-run":

package/src/core/islands/components/__tests__/IslandControls.test.tsx

@@ -2,6 +2,7 @@
 
 import { fireEvent, render, screen } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MockRequestClient } from "@/__mocks__/requests";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import type { CellId } from "@/core/cells/ids";
 import * as requestsModule from "@/core/network/requests";
@@ -18,7 +19,8 @@ vi.mock("@/utils/copy", () => ({
 }));
 
 describe("IslandControls", () => {
-  const mockSendRun = vi.fn();
+  const mockSendRun =
+    vi.fn<ReturnType<typeof requestsModule.useRequestClient>["sendRun"]>();
   const mockCopyToClipboard = vi.fn();
   const mockCodeCallback = vi.fn(() => "print('test code')");
   const cellId = "test-cell-id" as CellId;
@@ -30,11 +32,11 @@ describe("IslandControls", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    mockSendRun.mockResolvedValue(undefined);
+    mockSendRun.mockResolvedValue(null);
 
-    vi.spyOn(requestsModule, "useRequestClient").mockReturnValue({
-      sendRun: mockSendRun,
-    } as any);
+    vi.spyOn(requestsModule, "useRequestClient").mockReturnValue(
+      MockRequestClient.create({ sendRun: mockSendRun }),
+    );
 
     vi.spyOn(copyModule, "copyToClipboard").mockImplementation(
       mockCopyToClipboard,

package/src/core/kernel/__tests__/RuntimeState.test.ts

@@ -25,7 +25,7 @@ describe("RuntimeState", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    mockSendComponentValues = vi.fn((args) => Promise.resolve(null));
+    mockSendComponentValues = vi.fn(() => Promise.resolve(null));
     uiElementRegistry = UIElementRegistry.INSTANCE;
     uiElementRegistry.entries.clear();
 

package/src/core/network/connection.ts

@@ -7,8 +7,10 @@ import { waitFor } from "../state/jotai";
 import { type ConnectionStatus, ConnectionState } from "../websocket/types";
 
 /**
- * Atom for storing the connection status.
- * Initialized to NOT_STARTED.
+ * Shared application state, not the actual event-stream status.
+ *
+ * TODO(MO-8051) Expose read-only connection state and derive UI availability
+ * separately.
  */
 export const connectionAtom = atom<ConnectionStatus>({
   state: ConnectionState.NOT_STARTED,

package/src/core/network/useConnectionNotice.ts

@@ -17,11 +17,17 @@ export function useConnectionNotice(
   const notice = useAtomValue(connectionNoticeAtom);
   const connected = useAtomValue(isConnectedAtom);
   const elapsed = useDelayElapsed(notice?.pending ? delayMs : 0);
-  const [previous, setPrevious] = useState<ConnectionNotice | null>(null);
+  const [previous, setPrevious] = useState<ConnectionNotice | null>(
+    elapsed ? notice : null,
+  );
+
+  const visibleNotice = elapsed ? notice : null;
+  if ((notice || !connected) && previous !== visibleNotice) {
+    setPrevious(visibleNotice);
+  }
 
   useEffect(() => {
     if (notice || !connected) {
-      setPrevious(elapsed ? notice : null);
       return;
     }
     const timeout = setTimeout(() => setPrevious(null), 1500);

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;
@@ -257,21 +221,6 @@ 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}`,
-        );
-      },
-    );
-
     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 +249,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 +670,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 +738,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 +746,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

@@ -9,7 +9,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 +73,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 +128,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 +163,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.
@@ -193,7 +175,8 @@ export class RuntimeManager {
    * 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.
    */