@marimo-team/frontend 0.25.3-dev10 → 0.25.3-dev100

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

@@ -8,8 +8,9 @@ import { RuntimeManager } from "../runtime";
 import type { RuntimeConfig } from "../types";
 
 // Mock the session module
-vi.mock("@/core/kernel/session", () => ({
-  getSessionId: () => "test-session-id" as SessionId,
+vi.mock("@/core/kernel/session", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/core/kernel/session")>()),
+  getSessionId: () => "test-session-id",
 }));
 
 vi.mock("@/utils/Logger", () => ({ Logger: Mocks.quietLogger() }));
@@ -59,6 +60,7 @@ describe("RuntimeManager", () => {
       expect(url.protocol).toBe("wss:");
       expect(url.hostname).toBe("example.com");
       expect(url.pathname).toBe("/ws");
+      expect(url.searchParams.has("types")).toBe(false);
       expect(url.searchParams.get("session_id")).toBe(sessionId);
     });
 
@@ -219,15 +221,6 @@ describe("RuntimeManager", () => {
   });
 
   describe("getSseURL", () => {
-    it("should return an http(s) URL with the session ID", () => {
-      const runtime = new RuntimeManager(mockConfig);
-      const url = runtime.getSseURL("1234" as SessionId);
-
-      expect(url.protocol).toBe("https:");
-      expect(url.pathname).toBe("/sse");
-      expect(url.searchParams.get("session_id")).toBe("1234");
-    });
-
     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
@@ -256,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);
       }
@@ -676,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");
@@ -737,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,
@@ -745,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,6 +9,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 +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);
@@ -129,32 +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 {
-    return this.formatWsURL("/ws", this.getSessionSearchParams(sessionId));
+    const url = this.getSessionURL("/ws", sessionId);
+    return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
   /**
@@ -164,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.
@@ -180,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/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/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/bridge.ts

@@ -38,8 +38,6 @@ import type {
 } from "../network/types";
 import { filenameAtom } from "../saving/file-state";
 import { store } from "../state/jotai";
-import { BasicTransport } from "../websocket/transports/basic";
-import type { IConnectionTransport } from "../websocket/transports/transport";
 import { PyodideRouter } from "./router";
 import { getWorkerRPC } from "./rpc";
 import { getWasmRuntimeConfig } from "./runtime-config";
@@ -409,7 +407,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return response as Snippets;
   };
 
-  public openFile: EditRequests["openFile"] = async ({ path }) => {
+  public openFile: EditRequests["openFile"] = async () => {
     const url = createShareableLink({
       code: null,
       baseUrl: window.location.origin,
@@ -453,7 +451,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return null;
   };
 
-  public sendInstantiate: RunRequests["sendInstantiate"] = async (request) => {
+  public sendInstantiate: RunRequests["sendInstantiate"] = async () => {
     return null;
   };
 
@@ -719,12 +717,6 @@ export class PyodideBridge implements RunRequests, EditRequests {
   }
 }
 
-export function createPyodideConnection(): IConnectionTransport {
-  return BasicTransport.withProducerCallback((callback) => {
-    PyodideBridge.INSTANCE.attachMessageConsumer(callback);
-  });
-}
-
 // Compose the worker name. The optional "::controller" suffix tells
 // getController.ts that the host page provides a custom /wasm/controller.js
 // and that the dynamic import should be attempted. Hosts opt in by setting

package/src/core/wasm/store.ts

@@ -41,7 +41,7 @@ const urlFileStore: FileStore = {
 };
 
 export const domElementFileStore: FileStore = {
-  saveFile(contents: string) {
+  saveFile(_contents: string) {
     // Do nothing
   },
   readFile() {
@@ -54,7 +54,7 @@ export const domElementFileStore: FileStore = {
 };
 
 const remoteDefaultFileStore: FileStore = {
-  saveFile(contents: string) {
+  saveFile(_contents: string) {
     // Do nothing
   },
   readFile() {
@@ -79,7 +79,7 @@ export const mountConfigFileStore: FileStore = {
 };
 
 const emptyFileStore: FileStore = {
-  saveFile(contents: string) {
+  saveFile(_contents: string) {
     // Do nothing
   },
   readFile() {

package/src/core/wasm/worker/bootstrap.ts

@@ -85,7 +85,7 @@ export class DefaultWasmController implements WasmController {
         packages: [
           "micropip",
           "msgspec",
-          getMarimoWheel(opts.version),
+          getMarimoWheel(),
           "Markdown",
           "pymdown-extensions",
           "narwhals",

package/src/core/wasm/worker/getMarimoWheel.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-export function getMarimoWheel(version: string) {
+export function getMarimoWheel() {
   if (import.meta.env.DEV && import.meta.env.VITE_WASM_MARIMO_PREBUILT_WHEEL) {
     return "marimo-base";
   }

package/src/core/websocket/__tests__/{useMarimoKernelConnection.test.ts → connection-state.test.ts}

@@ -3,7 +3,7 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { Logger } from "@/utils/Logger";
 import { WebSocketClosedReason, ConnectionState } from "../types";
-import { classifyCloseEvent } from "../useMarimoKernelConnection";
+import { classifyCloseEvent } from "../connection-state";
 
 function classify(reason: string | undefined) {
   return classifyCloseEvent({ reason });
@@ -13,8 +13,10 @@ describe("classifyCloseEvent", () => {
   describe("transient closes (default branch)", () => {
     it("retries on empty/undefined reason", () => {
       const decision = classify(undefined);
-      expect(decision.kind).toBe("retry");
-      expect(decision.status).toEqual({ state: ConnectionState.CONNECTING });
+      expect(decision).toEqual({
+        kind: "retry",
+        status: { state: ConnectionState.CONNECTING },
+      });
     });
 
     it("treats unknown reason strings as transient and logs a warning", () => {
@@ -40,41 +42,25 @@ describe("classifyCloseEvent", () => {
       // terminal or a rejected client would reconnect forever.
       "MARIMO_UNAUTHORIZED",
       "MARIMO_KIOSK_NOT_ALLOWED",
-    ])("%s → terminal with KERNEL_DISCONNECTED, closes transport", (reason) => {
+    ])("%s → terminal with KERNEL_DISCONNECTED", (reason) => {
       const decision = classify(reason);
-      expect(decision.kind).toBe("terminal");
-      expect(decision.status).toMatchObject({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-      });
       expect(decision).toMatchObject({
         kind: "terminal",
-        closeTransport: true,
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+        },
       });
     });
 
-    it("MARIMO_KERNEL_STARTUP_ERROR → terminal + closeTransport", () => {
+    it("MARIMO_KERNEL_STARTUP_ERROR → terminal", () => {
       const decision = classify("MARIMO_KERNEL_STARTUP_ERROR");
-      expect(decision.kind).toBe("terminal");
-      expect(decision.status).toMatchObject({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
-      });
       expect(decision).toMatchObject({
         kind: "terminal",
-        closeTransport: true,
-      });
-    });
-  });
-
-  describe("transport exhaustion", () => {
-    it("MARIMO_TRANSPORT_EXHAUSTED → gave-up with KERNEL_DISCONNECTED", () => {
-      const decision = classify("MARIMO_TRANSPORT_EXHAUSTED");
-      expect(decision.kind).toBe("gave-up");
-      expect(decision.status).toEqual({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+        },
       });
     });
   });