@marimo-team/islands 0.23.7-dev9 → 0.23.7-dev90

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

@@ -275,14 +275,14 @@ describe("RuntimeManager", () => {
     });
   });
 
-  describe("isHealthy", () => {
+  describe("probeHealth", () => {
     it("should return true for successful health check", async () => {
       global.fetch = vi.fn().mockResolvedValue({
         ok: true,
       });
 
       const runtime = new RuntimeManager(mockConfig);
-      const result = await runtime.isHealthy();
+      const result = await runtime.probeHealth();
 
       expect(result).toBe(true);
       expect(fetch).toHaveBeenCalledWith("https://example.com/health");
@@ -294,7 +294,7 @@ describe("RuntimeManager", () => {
       });
 
       const runtime = new RuntimeManager(mockConfig);
-      const result = await runtime.isHealthy();
+      const result = await runtime.probeHealth();
 
       expect(result).toBe(false);
     });
@@ -304,11 +304,32 @@ describe("RuntimeManager", () => {
       global.fetch = vi.fn().mockRejectedValue(error);
 
       const runtime = new RuntimeManager(mockConfig);
-      const result = await runtime.isHealthy();
+      const result = await runtime.probeHealth();
 
       expect(result).toBe(false);
     });
 
+    it("should not mutate config.url on redirect", async () => {
+      global.fetch = vi.fn().mockResolvedValue({
+        ok: true,
+        redirected: true,
+        url: "https://sandbox.example.com/health?some_value=abc123",
+      });
+
+      const runtime = new RuntimeManager(
+        {
+          ...mockConfig,
+          url: "https://backend.example.com/lazy?some_value=abc123",
+        },
+        true,
+      );
+      const result = await runtime.probeHealth();
+      expect(result).toBe(true);
+      expect(runtime.httpURL.hostname).toBe("backend.example.com");
+    });
+  });
+
+  describe("reconcileFromHealth", () => {
     it("should update config.url on redirect, stripping /health from pathname", async () => {
       global.fetch = vi.fn().mockResolvedValue({
         ok: true,
@@ -323,7 +344,7 @@ describe("RuntimeManager", () => {
         },
         true, // lazy — don't call init() in constructor
       );
-      const result = await runtime.isHealthy();
+      const result = await runtime.reconcileFromHealth();
 
       expect(result).toBe(true);
       // Should strip /health from pathname but preserve query params
@@ -342,7 +363,7 @@ describe("RuntimeManager", () => {
     it("should resolve immediately if healthy", async () => {
       const runtime = new RuntimeManager(mockConfig, true);
 
-      vi.spyOn(runtime, "isHealthy").mockResolvedValue(true);
+      vi.spyOn(runtime, "reconcileFromHealth").mockResolvedValue(true);
       runtime.init();
 
       await expect(runtime.waitForHealthy()).resolves.toBeUndefined();
@@ -351,7 +372,7 @@ describe("RuntimeManager", () => {
     it("should retry and eventually succeed", async () => {
       const runtime = new RuntimeManager(mockConfig, true);
       const healthySpy = vi
-        .spyOn(runtime, "isHealthy")
+        .spyOn(runtime, "reconcileFromHealth")
         .mockResolvedValueOnce(false)
         .mockResolvedValueOnce(false)
         .mockResolvedValueOnce(true);
@@ -364,7 +385,7 @@ describe("RuntimeManager", () => {
 
     it("should throw after max retries", async () => {
       const runtime = new RuntimeManager(mockConfig, true);
-      vi.spyOn(runtime, "isHealthy").mockResolvedValue(false);
+      vi.spyOn(runtime, "reconcileFromHealth").mockResolvedValue(false);
       runtime.init({ disableRetryDelay: true });
 
       await expect(runtime.waitForHealthy()).rejects.toThrow(
@@ -431,7 +452,7 @@ describe("RuntimeManager", () => {
       // Mock failed health check
       global.fetch = vi.fn().mockResolvedValue({ ok: false });
 
-      await runtime.isHealthy();
+      await runtime.reconcileFromHealth();
 
       baseElement = document.querySelector("base");
       expect(baseElement).toBeNull();
@@ -443,7 +464,7 @@ describe("RuntimeManager", () => {
       // Mock successful health check
       global.fetch = vi.fn().mockResolvedValue({ ok: true });
 
-      await runtime.isHealthy();
+      await runtime.reconcileFromHealth();
 
       const baseElement = document.querySelector("base");
       expect(baseElement).toBeTruthy();
@@ -461,7 +482,7 @@ describe("RuntimeManager", () => {
       // Mock successful health check
       global.fetch = vi.fn().mockResolvedValue({ ok: true });
 
-      await runtime.isHealthy();
+      await runtime.reconcileFromHealth();
 
       const baseElement = document.querySelector("base");
       expect(baseElement).toBe(existingBase); // Should be the same element
@@ -483,7 +504,7 @@ describe("RuntimeManager", () => {
       // Mock successful health check
       global.fetch = vi.fn().mockResolvedValue({ ok: true });
 
-      await runtime.isHealthy();
+      await runtime.reconcileFromHealth();
 
       const baseElement = document.querySelector("base");
       expect(baseElement).toBeTruthy();
@@ -524,11 +545,11 @@ describe("RuntimeManager", () => {
       });
 
       const wsUrl = runtime.getWsURL("test" as SessionId);
-      const httpUrl = runtime.formatHttpURL(
-        "api/test",
-        new URLSearchParams(),
-        false,
-      );
+      const httpUrl = runtime.formatHttpURL({
+        path: "api/test",
+        searchParams: new URLSearchParams(),
+        restrictToKnownQueryParams: false,
+      });
 
       // Should preserve base URL query params
       expect(wsUrl.searchParams.get("base_param")).toBe("existing");

package/src/core/runtime/runtime.ts

@@ -44,17 +44,22 @@ export class RuntimeManager {
     return this.httpURL.origin === window.location.origin;
   }
 
+  private get isServerless(): boolean {
+    return isWasm() || isIslands() || isStaticNotebook();
+  }
+
   /**
    * The base URL of the runtime.
    */
-  formatHttpURL(
-    path?: string,
-    searchParams?: URLSearchParams,
+  formatHttpURL({
+    path = "",
+    searchParams,
     restrictToKnownQueryParams = true,
-  ): URL {
-    if (!path) {
-      path = "";
-    }
+  }: {
+    path?: string;
+    searchParams?: URLSearchParams;
+    restrictToKnownQueryParams?: boolean;
+  }): URL {
     // URL may be something like "http://localhost:8000?auth=123"
     const baseUrl = this.httpURL;
     const currentParams = new URLSearchParams(window.location.search);
@@ -84,11 +89,11 @@ export class RuntimeManager {
   formatWsURL(path: string, searchParams?: URLSearchParams): URL {
     // We don't restrict to known query parameters, since mo.query_params()
     // can accept arbitrary parameters.
-    const url = this.formatHttpURL(
+    const url = this.formatHttpURL({
       path,
       searchParams,
-      /* restrictToKnownQueryParams =*/ false,
-    );
+      restrictToKnownQueryParams: false,
+    });
 
     // For cross-origin runtimes, pass the auth token as a query parameter.
     // WebSocket connections cannot send custom headers (no Authorization
@@ -168,45 +173,63 @@ export class RuntimeManager {
   }
 
   getAiURL(path: "completion" | "chat"): URL {
-    return this.formatHttpURL(`/api/ai/${path}`);
+    return this.formatHttpURL({ path: `/api/ai/${path}` });
   }
 
   /**
    * The URL of the health check endpoint.
    */
   healthURL(): URL {
-    return this.formatHttpURL("/health");
+    return this.formatHttpURL({ path: "/health" });
   }
 
-  async isHealthy(): Promise<boolean> {
-    // Always healthy if WASM, Islands, or a static notebook (no server)
-    if (isWasm() || isIslands() || isStaticNotebook()) {
-      return true;
-    }
-
+  private async fetchHealth(): Promise<Response | null> {
     try {
-      const response = await fetch(this.healthURL().toString());
-      // If there is a redirect, update the URL in the config
-      if (response.redirected) {
-        Logger.debug(`Runtime redirected to ${response.url}`);
-        // strip /health from the URL, using URL parsing to handle query params
-        const redirected = new URL(response.url);
-        redirected.pathname = redirected.pathname.replace(/\/health$/, "");
-        this.config.url = redirected.toString();
-      }
-
-      const success = response.ok;
-      if (success) {
-        this.setDOMBaseUri(this.config.url);
-      }
-      return success;
+      return await fetch(this.healthURL().toString());
     } catch (error) {
       Logger.error(
         `Failed to check health: ${error instanceof Error ? error.message : "Unknown error"}`,
         { cause: error },
       );
+      return null;
+    }
+  }
+
+  async reconcileFromHealth(): Promise<boolean> {
+    // Always healthy if WASM, Islands, or a static notebook (no server)
+    if (this.isServerless) {
+      return true;
+    }
+
+    const response = await this.fetchHealth();
+
+    if (!response) {
       return false;
     }
+
+    if (response.redirected) {
+      Logger.debug(`Runtime redirected to ${response.url}`);
+      // strip /health from the URL, using URL parsing to handle query params
+      const redirected = new URL(response.url);
+      redirected.pathname = redirected.pathname.replace(/\/health$/, "");
+      this.config.url = redirected.toString();
+    }
+
+    if (response.ok) {
+      this.setDOMBaseUri(this.config.url);
+    }
+
+    return response.ok;
+  }
+
+  async probeHealth(): Promise<boolean> {
+    // Always healthy if WASM, Islands, or a static notebook (no server)
+    if (this.isServerless) {
+      return true;
+    }
+
+    const response = await this.fetchHealth();
+    return response?.ok ?? false;
   }
 
   /**
@@ -251,7 +274,7 @@ export class RuntimeManager {
     const growthFactor = 1.2;
     const maxDelay = 2000;
 
-    while (!(await this.isHealthy())) {
+    while (!(await this.reconcileFromHealth())) {
       if (retries >= maxRetries) {
         Logger.error(`Failed to connect after ${maxRetries} retries`);
         this.initialHealthyCheck.reject(

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

@@ -0,0 +1,34 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { shouldLoadDuckDBPackages } from "../utils";
+
+describe("shouldLoadDuckDBPackages", () => {
+  it("loads for mo.sql", () => {
+    expect(shouldLoadDuckDBPackages('df = mo.sql("SELECT 1")')).toBe(true);
+  });
+
+  it("loads for duckdb imports and usage", () => {
+    expect(shouldLoadDuckDBPackages("import duckdb")).toBe(true);
+    expect(shouldLoadDuckDBPackages("from duckdb import sql")).toBe(true);
+    expect(shouldLoadDuckDBPackages("import pandas, duckdb")).toBe(true);
+    expect(shouldLoadDuckDBPackages("rows = duckdb.sql('SELECT 1')")).toBe(
+      true,
+    );
+  });
+
+  it("loads when package discovery found duckdb", () => {
+    expect(
+      shouldLoadDuckDBPackages("print('hello')", new Set(["duckdb"])),
+    ).toBe(true);
+  });
+
+  it("does not load for incidental duckdb text", () => {
+    expect(shouldLoadDuckDBPackages("name = 'duckdb'")).toBe(false);
+    expect(shouldLoadDuckDBPackages("# import duckdb")).toBe(false);
+  });
+
+  it("does not load without mo.sql, duckdb usage, or discovery", () => {
+    expect(shouldLoadDuckDBPackages("print('hello')")).toBe(false);
+  });
+});

package/src/core/wasm/utils.ts

@@ -10,3 +10,17 @@ export function isWasm(): boolean {
     document.querySelector("marimo-wasm") !== null
   );
 }
+
+const DUCKDB_USAGE_PATTERN =
+  /(^|\n)\s*(?:import\s+[^\n#]*\bduckdb\b|from\s+duckdb\b|[^\n#]*\bduckdb\s*\.)/;
+
+export function shouldLoadDuckDBPackages(
+  code: string,
+  foundPackages?: ReadonlySet<string>,
+): boolean {
+  return (
+    code.includes("mo.sql") ||
+    DUCKDB_USAGE_PATTERN.test(code) ||
+    foundPackages?.has("duckdb") === true
+  );
+}

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

@@ -9,6 +9,7 @@ import { WasmFileSystem } from "./fs";
 import { getMarimoWheel } from "./getMarimoWheel";
 import { t } from "./tracer";
 import type { SerializedBridge, WasmController } from "./types";
+import { shouldLoadDuckDBPackages } from "../utils";
 
 const MAKE_SNAPSHOT = false;
 
@@ -163,8 +164,8 @@ export class DefaultWasmController implements WasmController {
   private async loadNotebookDeps(code: string, foundPackages: Set<string>) {
     const pyodide = this.requirePyodide;
 
-    if (code.includes("mo.sql")) {
-      // We need pandas and duckdb for mo.sql
+    if (shouldLoadDuckDBPackages(code, foundPackages)) {
+      // We need pandas and duckdb for mo.sql and for remote duckdb sources
       code = `import pandas\n${code}`;
       code = `import duckdb\n${code}`;
       code = `import sqlglot\n${code}`;

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

@@ -34,6 +34,7 @@ import type {
   SerializedBridge,
   WasmController,
 } from "./types";
+import { shouldLoadDuckDBPackages } from "../utils";
 
 /**
  * Web worker responsible for running the notebook.
@@ -141,8 +142,8 @@ const requestHandler = createRPCRequestHandler({
     const span = t.startSpan("loadPackages");
     await pyodideReadyPromise; // Make sure loading is done
 
-    if (code.includes("mo.sql")) {
-      // Add pandas and duckdb to the code
+    if (shouldLoadDuckDBPackages(code)) {
+      // Add pandas and duckdb to the code for mo.sql and for remote duckdb sources
       code = `import pandas\n${code}`;
       code = `import duckdb\n${code}`;
       code = `import sqlglot\n${code}`;

package/src/core/websocket/__tests__/useMarimoKernelConnection.hook.test.tsx

@@ -0,0 +1,156 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider as JotaiProvider } from "jotai";
+import type React from "react";
+import { ErrorBoundary } from "react-error-boundary";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("@/core/websocket/useWebSocket", async () => {
+  const actual =
+    await vi.importActual<typeof import("../useWebSocket")>("../useWebSocket");
+  return {
+    ...actual,
+    useConnectionTransport: vi.fn(),
+  };
+});
+
+vi.mock("@/core/runtime/config", async () => {
+  const actual = await vi.importActual<typeof import("@/core/runtime/config")>(
+    "@/core/runtime/config",
+  );
+  return {
+    ...actual,
+    useRuntimeManager: vi.fn(),
+  };
+});
+
+import { useRuntimeManager } from "@/core/runtime/config";
+import { connectionAtom } from "../../network/connection";
+import type { SessionId } from "../../kernel/session";
+import { WebSocketClosedReason, WebSocketState } from "../types";
+import { useMarimoKernelConnection } from "../useMarimoKernelConnection";
+import { useConnectionTransport } from "../useWebSocket";
+
+interface MockTransport {
+  readyState: 0 | 1 | 2 | 3;
+  reconnect: ReturnType<typeof vi.fn>;
+  close: ReturnType<typeof vi.fn>;
+  send: ReturnType<typeof vi.fn>;
+  addEventListener: ReturnType<typeof vi.fn>;
+  removeEventListener: ReturnType<typeof vi.fn>;
+}
+
+function makeTransport(
+  readyState: 0 | 1 | 2 | 3 = WebSocket.CLOSED,
+): MockTransport {
+  return {
+    readyState,
+    reconnect: vi.fn(),
+    close: vi.fn(),
+    send: vi.fn(),
+    addEventListener: vi.fn(),
+    removeEventListener: vi.fn(),
+  };
+}
+
+function makeRuntimeManager(
+  reconcileFromHealth = vi.fn().mockResolvedValue(true),
+) {
+  return {
+    reconcileFromHealth,
+    probeHealth: vi.fn().mockResolvedValue(true),
+    getWsURL: () => new URL("ws://localhost/ws"),
+    waitForHealthy: vi.fn().mockResolvedValue(undefined),
+    isSameOrigin: true,
+  };
+}
+
+describe("useMarimoKernelConnection.reconnect()", () => {
+  let transport: MockTransport;
+  let isHealthy: ReturnType<typeof vi.fn>;
+  let store: ReturnType<typeof createStore>;
+
+  beforeEach(() => {
+    transport = makeTransport(WebSocket.CLOSED);
+    isHealthy = vi.fn().mockResolvedValue(true);
+    store = createStore();
+    store.set(connectionAtom, {
+      state: WebSocketState.CLOSED,
+      code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+      reason: "kernel not found",
+    });
+    vi.mocked(useConnectionTransport).mockReturnValue(transport);
+    vi.mocked(useRuntimeManager).mockReturnValue(
+      makeRuntimeManager(isHealthy) as unknown as ReturnType<
+        typeof useRuntimeManager
+      >,
+    );
+  });
+
+  function renderUseHook() {
+    const wrapper: React.FC<React.PropsWithChildren> = ({ children }) => (
+      <JotaiProvider store={store}>
+        <ErrorBoundary fallback={null}>{children}</ErrorBoundary>
+      </JotaiProvider>
+    );
+    return renderHook(
+      () =>
+        useMarimoKernelConnection({
+          sessionId: "test-session" as SessionId,
+          autoInstantiate: false,
+          setCells: () => {},
+        }),
+      { wrapper },
+    );
+  }
+
+  it("is a no-op when the transport is already OPEN", async () => {
+    transport.readyState = WebSocket.OPEN;
+    const { result } = renderUseHook();
+    await act(async () => {
+      await result.current.reconnect();
+    });
+    expect(isHealthy).not.toHaveBeenCalled();
+    expect(transport.reconnect).not.toHaveBeenCalled();
+  });
+
+  it("is a no-op when the transport is already CONNECTING", async () => {
+    transport.readyState = WebSocket.CONNECTING;
+    const { result } = renderUseHook();
+    await act(async () => {
+      await result.current.reconnect();
+    });
+    expect(isHealthy).not.toHaveBeenCalled();
+    expect(transport.reconnect).not.toHaveBeenCalled();
+  });
+
+  it("probes /health and reconnects when the runtime is healthy", async () => {
+    isHealthy.mockResolvedValue(true);
+    const { result } = renderUseHook();
+    await act(async () => {
+      await result.current.reconnect();
+    });
+    expect(isHealthy).toHaveBeenCalledOnce();
+    expect(transport.reconnect).toHaveBeenCalledOnce();
+    expect(store.get(connectionAtom)).toEqual({
+      state: WebSocketState.CONNECTING,
+    });
+  });
+
+  it("transitions to CLOSED and does not call ws.reconnect when the probe fails", async () => {
+    isHealthy.mockResolvedValue(false);
+    const { result } = renderUseHook();
+    await act(async () => {
+      await result.current.reconnect();
+    });
+    expect(isHealthy).toHaveBeenCalledOnce();
+    expect(transport.reconnect).not.toHaveBeenCalled();
+    expect(store.get(connectionAtom)).toEqual({
+      state: WebSocketState.CLOSED,
+      code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+      reason: "kernel not found",
+    });
+  });
+});

package/src/core/websocket/__tests__/useMarimoKernelConnection.test.ts

@@ -0,0 +1,101 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { Logger } from "@/utils/Logger";
+import { WebSocketClosedReason, WebSocketState } from "../types";
+import { classifyCloseEvent } from "../useMarimoKernelConnection";
+
+function classify(reason: string | undefined) {
+  return classifyCloseEvent({ reason });
+}
+
+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: WebSocketState.CONNECTING });
+    });
+
+    it("treats unknown reason strings as transient and logs a warning", () => {
+      const logger = vi.spyOn(Logger, "warn").mockImplementation(() => {});
+      const decision = classify("something-else");
+      expect(decision.kind).toBe("retry");
+      expect(logger).toHaveBeenCalled();
+      logger.mockRestore();
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe("terminal closes (server-initiated)", () => {
+    it("MARIMO_ALREADY_CONNECTED → terminal + closeTransport, with takeover", () => {
+      const decision = classify("MARIMO_ALREADY_CONNECTED");
+      expect(decision.kind).toBe("terminal");
+      expect(decision.status).toMatchObject({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.ALREADY_RUNNING,
+        canTakeover: true,
+      });
+      if (decision.kind === "terminal") {
+        expect(decision.closeTransport).toBe(true);
+      }
+    });
+
+    it.each([
+      "MARIMO_WRONG_KERNEL_ID",
+      "MARIMO_NO_FILE_KEY",
+      "MARIMO_NO_SESSION_ID",
+      "MARIMO_NO_SESSION",
+      "MARIMO_SHUTDOWN",
+    ])("%s → terminal with KERNEL_DISCONNECTED, closes transport", (reason) => {
+      const decision = classify(reason);
+      expect(decision.kind).toBe("terminal");
+      expect(decision.status).toMatchObject({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+      });
+      if (decision.kind === "terminal") {
+        expect(decision.closeTransport).toBe(true);
+      }
+    });
+
+    it("MARIMO_MALFORMED_QUERY → terminal but does NOT close transport", () => {
+      const decision = classify("MARIMO_MALFORMED_QUERY");
+      expect(decision.kind).toBe("terminal");
+      expect(decision.status).toMatchObject({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.MALFORMED_QUERY,
+      });
+      if (decision.kind === "terminal") {
+        expect(decision.closeTransport).toBe(false);
+      }
+    });
+
+    it("MARIMO_KERNEL_STARTUP_ERROR → terminal + closeTransport", () => {
+      const decision = classify("MARIMO_KERNEL_STARTUP_ERROR");
+      expect(decision.kind).toBe("terminal");
+      expect(decision.status).toMatchObject({
+        state: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+      });
+      if (decision.kind === "terminal") {
+        expect(decision.closeTransport).toBe(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: WebSocketState.CLOSED,
+        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+        reason: "kernel not found",
+      });
+    });
+  });
+});