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

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/bridge.ts

@@ -11,6 +11,7 @@ import { reloadSafe } from "@/utils/reload-safe";
 import { generateUUID } from "@/utils/uuid";
 import { createModuleWorker } from "@/utils/worker";
 import { notebookIsRunningAtom } from "../cells/cells";
+import { createWasmNotebookBackend } from "../host/wasm";
 import type { ClientCommand } from "../kernel/messages";
 import { getInitialAppMode } from "../mode";
 import { API } from "../network/api";
@@ -38,8 +39,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";
@@ -69,6 +68,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   private rpc!: ReturnType<typeof getWorkerRPC<WorkerSchema>>;
   private saveRpc: SaveWorker | undefined;
   private pendingSessionSave: Promise<unknown> = Promise.resolve();
+  private notebookConsumers = new Map<string, (event: unknown) => void>();
   private interruptBuffer?: Uint8Array;
   private messageConsumer:
     | ((message: MessageEvent<string>) => void)
@@ -119,6 +119,10 @@ export class PyodideBridge implements RunRequests, EditRequests {
     this.rpc = getWorkerRPC<WorkerSchema>(worker);
     this.rpc.send.bootstrap(runtimeConfig);
 
+    this.rpc.addMessageListener("notebookEvent", ({ id, event }) => {
+      this.notebookConsumers.get(id)?.(JSON.parse(event));
+    });
+
     // Listeners
     this.rpc.addMessageListener("ready", () => {
       this.startSession();
@@ -215,6 +219,36 @@ export class PyodideBridge implements RunRequests, EditRequests {
     this.rpc.proxy.send.consumerReady({});
   }
 
+  public async createNotebookBackend() {
+    await this.initialized.promise;
+    const { id } = await this.rpc.proxy.request.notebookInfo();
+    return createWasmNotebookBackend({
+      notebookId: id,
+      subscribe: async (cursor, emit) => {
+        await this.initialized.promise;
+        const id = generateUUID();
+        this.notebookConsumers.set(id, emit);
+        const close = () => {
+          this.notebookConsumers.delete(id);
+          void this.rpc.proxy.request
+            .notebookUnsubscribe(id)
+            .catch(Logger.error);
+        };
+        try {
+          await this.rpc.proxy.request.notebookSubscribe({ id, cursor });
+          return close;
+        } catch (error) {
+          close();
+          throw error;
+        }
+      },
+      call: async (call) => {
+        await this.initialized.promise;
+        return this.rpc.proxy.request.notebookCall(call);
+      },
+    });
+  }
+
   public sendRename: EditRequests["sendRename"] = async ({ filename }) => {
     if (filename === null) {
       return null;
@@ -409,7 +443,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 +487,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return null;
   };
 
-  public sendInstantiate: RunRequests["sendInstantiate"] = async (request) => {
+  public sendInstantiate: RunRequests["sendInstantiate"] = async () => {
     return null;
   };
 
@@ -719,12 +753,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",
@@ -198,6 +198,7 @@ export class DefaultWasmController implements WasmController {
           if task is not None:
             await asyncio.gather(task, return_exceptions=True)
           if bridge is not None:
+            bridge.close()
             bridge.session = None
           task = None
           kernel_task = None

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/wasm/worker/types.ts

@@ -121,4 +121,15 @@ export type SerializedBridge = {
     : RawBridge[P];
 } & {
   save(payload: string): Promise<void>;
+  notebook: {
+    readonly id: string;
+    subscribe(
+      id: string,
+      cursor: string | null,
+      emit: (event: string) => void,
+    ): Promise<void>;
+    unsubscribe(id: string): void;
+    call(call: string): Promise<string>;
+    destroy(): void;
+  };
 };

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

@@ -1,5 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { unstable_NotebookCall } from "@marimo-team/client";
 import type { PyodideInterface } from "pyodide";
 import {
   createRPC,
@@ -81,7 +82,53 @@ const bridgeReady = new Deferred<SerializedBridge>();
 let started = false;
 
 // Handle RPC requests
+const notebookSubscriptions = new Map<string, Promise<void>>();
+async function getNotebook() {
+  const bridge = await bridgeReady.promise;
+  return bridge.notebook;
+}
+
 const requestHandler = createRPCRequestHandler({
+  notebookInfo: async () => {
+    const notebook = await getNotebook();
+    try {
+      return { id: notebook.id };
+    } finally {
+      notebook.destroy();
+    }
+  },
+  notebookSubscribe: (request: { id: string; cursor: string | null }) => {
+    const subscribing = (async () => {
+      const notebook = await getNotebook();
+      try {
+        await notebook.subscribe(request.id, request.cursor, (event) => {
+          rpc.send.notebookEvent({ id: request.id, event });
+        });
+      } finally {
+        notebook.destroy();
+      }
+    })();
+    notebookSubscriptions.set(request.id, subscribing);
+    return subscribing.finally(() => notebookSubscriptions.delete(request.id));
+  },
+  notebookUnsubscribe: async (id: string) => {
+    // An RPC timeout can close a reader while its subscription is still opening.
+    await notebookSubscriptions.get(id)?.catch(() => undefined);
+    const notebook = await getNotebook();
+    try {
+      notebook.unsubscribe(id);
+    } finally {
+      notebook.destroy();
+    }
+  },
+  notebookCall: async (call: unstable_NotebookCall): Promise<unknown> => {
+    const notebook = await getNotebook();
+    try {
+      return JSON.parse(await notebook.call(JSON.stringify(call)));
+    } finally {
+      notebook.destroy();
+    }
+  },
   /**
    * Start the session
    */
@@ -335,6 +382,7 @@ const requestHandler = createRPCRequestHandler({
 export type WorkerSchema = RPCSchema<
   {
     messages: {
+      notebookEvent: { id: string; event: string };
       // Emitted when the worker is ready
       ready: {};
       // Emitted when the kernel sends a message

package/src/core/websocket/__tests__/connection-state.test.ts

@@ -64,16 +64,4 @@ describe("classifyCloseEvent", () => {
       });
     });
   });
-
-  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",
-      });
-    });
-  });
 });