@marimo-team/frontend 0.22.5-dev9 → 0.22.5

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

@@ -86,6 +86,74 @@ describe("RuntimeManager", () => {
     });
   });
 
+  describe("cross-origin auth token in WS URLs", () => {
+    it("should add access_token to WS URL when cross-origin with authToken", () => {
+      // example.com is cross-origin relative to the test environment (localhost)
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getWsURL("s_123" as SessionId);
+
+      expect(url.searchParams.get("access_token")).toBe("my-secret-token");
+      expect(url.searchParams.get("session_id")).toBe("s_123");
+    });
+
+    it("should not add access_token to WS URL when same-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: window.location.origin,
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getWsURL("s_123" as SessionId);
+
+      expect(url.searchParams.get("access_token")).toBeNull();
+    });
+
+    it("should not add access_token when no authToken is configured", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com",
+          lazy: true,
+        },
+        true,
+      );
+      const url = runtime.getWsURL("s_123" as SessionId);
+
+      expect(url.searchParams.get("access_token")).toBeNull();
+    });
+
+    it("should add access_token to all WS URL types when cross-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+
+      const wsUrl = runtime.getWsURL("s_123" as SessionId);
+      const wsSyncUrl = runtime.getWsSyncURL("s_123" as SessionId);
+      const terminalUrl = runtime.getTerminalWsURL();
+
+      expect(wsUrl.searchParams.get("access_token")).toBe("my-secret-token");
+      expect(wsSyncUrl.searchParams.get("access_token")).toBe(
+        "my-secret-token",
+      );
+      expect(terminalUrl.searchParams.get("access_token")).toBe(
+        "my-secret-token",
+      );
+    });
+  });
+
   describe("getWsSyncURL", () => {
     it("should return WebSocket Sync URL", () => {
       const runtime = new RuntimeManager(mockConfig);
@@ -117,12 +185,52 @@ describe("RuntimeManager", () => {
       expect(url.pathname).toBe("/lsp/pylsp");
     });
 
-    it("should return copilot URL", () => {
-      const runtime = new RuntimeManager(mockConfig);
+    it("should return copilot URL without non-auth query params", () => {
+      const runtime = new RuntimeManager({
+        url: "https://example.com?foo=bar&baz=qux",
+        lazy: true,
+      });
       const url = runtime.getLSPURL("copilot");
 
       expect(url.protocol).toBe("wss:");
       expect(url.pathname).toBe("/lsp/copilot");
+      expect(url.searchParams.get("foo")).toBeNull();
+      expect(url.searchParams.get("baz")).toBeNull();
+    });
+
+    it("should preserve access_token on copilot URL when cross-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com?foo=bar",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getLSPURL("copilot");
+
+      expect(url.protocol).toBe("wss:");
+      expect(url.pathname).toBe("/lsp/copilot");
+      expect(url.searchParams.get("access_token")).toBe("my-secret-token");
+      // Other params should be stripped
+      expect(url.searchParams.get("foo")).toBeNull();
+    });
+
+    it("should not have access_token on copilot URL when same-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: window.location.origin,
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getLSPURL("copilot");
+
+      expect(url.protocol).toBe("ws:");
+      expect(url.pathname).toBe("/lsp/copilot");
+      expect(url.searchParams.get("access_token")).toBeNull();
+      expect(url.search).toBe("");
     });
   });
 
@@ -200,6 +308,34 @@ describe("RuntimeManager", () => {
 
       expect(result).toBe(false);
     });
+
+    it("should update config.url on redirect, stripping /health from pathname", 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, // lazy — don't call init() in constructor
+      );
+      const result = await runtime.isHealthy();
+
+      expect(result).toBe(true);
+      // Should strip /health from pathname but preserve query params
+      const wsUrl = runtime.getWsURL("s_test" as SessionId);
+      expect(wsUrl.pathname).toBe("/ws");
+      expect(wsUrl.hostname).toBe("sandbox.example.com");
+      expect(wsUrl.searchParams.get("some_value")).toBe("abc123");
+
+      // Clean up side effects
+      document.querySelectorAll("base").forEach((el) => el.remove());
+      global.fetch = vi.fn().mockResolvedValue({ ok: false });
+    });
   });
 
   describe("waitForHealthy", () => {

package/src/core/runtime/runtime.ts

@@ -88,6 +88,15 @@ export class RuntimeManager {
       searchParams,
       /* restrictToKnownQueryParams =*/ false,
     );
+
+    // For cross-origin runtimes, pass the auth token as a query parameter.
+    // WebSocket connections cannot send custom headers (no Authorization
+    // header), and cross-origin cookies are blocked by browsers, so the
+    // access_token query param is the only way to authenticate.
+    if (!this.isSameOrigin && this.config.authToken) {
+      url.searchParams.set(KnownQueryParams.accessToken, this.config.authToken);
+    }
+
     return asWsUrl(url.toString());
   }
 
@@ -143,9 +152,15 @@ export class RuntimeManager {
    */
   getLSPURL(lsp: "pylsp" | "basedpyright" | "copilot" | "ty" | "pyrefly"): URL {
     if (lsp === "copilot") {
-      // For copilot, don't include any query parameters
+      // For copilot, strip all query parameters except the auth token.
+      // Copilot doesn't understand arbitrary query params, but we still
+      // need access_token for cross-origin authentication.
       const url = this.formatWsURL(`/lsp/${lsp}`);
+      const accessToken = url.searchParams.get(KnownQueryParams.accessToken);
       url.search = "";
+      if (accessToken) {
+        url.searchParams.set(KnownQueryParams.accessToken, accessToken);
+      }
       return url;
     }
     return this.formatWsURL(`/lsp/${lsp}`);
@@ -173,9 +188,10 @@ export class RuntimeManager {
       // 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
-        const baseUrl = response.url.replace(/\/health$/, "");
-        this.config.url = baseUrl;
+        // 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;
@@ -183,7 +199,11 @@ export class RuntimeManager {
         this.setDOMBaseUri(this.config.url);
       }
       return success;
-    } catch {
+    } catch (error) {
+      Logger.error(
+        `Failed to check health: ${error instanceof Error ? error.message : "Unknown error"}`,
+        { cause: error },
+      );
       return false;
     }
   }

package/src/core/saving/file-state.ts

@@ -16,6 +16,22 @@ export const filenameAtom = atom<string | null>(getFilenameFromDOM());
  */
 export const cwdAtom = atom<string | null>(null);
 
+/**
+ * LSP workspace information from the backend.
+ * Contains the project root and the document's file URI.
+ */
+export interface LspWorkspace {
+  rootUri: string;
+  documentUri: string;
+}
+
+/**
+ * Atom for storing the LSP workspace information.
+ * This is populated during active notebook sessions
+ * and null for other pages.
+ */
+export const lspWorkspaceAtom = atom<LspWorkspace | null>(null);
+
 /**
  * Set for static notebooks.
  */

package/src/hooks/useAsyncData.ts

@@ -341,7 +341,7 @@ export function useAsyncData<T>(
     };
     setResult((prevResult) => {
       // If we have previous data, show reloading state
-      if (prevResult.status === "success") {
+      if (prevResult.status === "success" || prevResult.status === "loading") {
         return Result.loading(prevResult.data);
       }
       // Otherwise, show initial loading state

package/src/mount.tsx

@@ -36,7 +36,12 @@ import {
   DEFAULT_RUNTIME_CONFIG,
   runtimeConfigAtom,
 } from "./core/runtime/config";
-import { codeAtom, cwdAtom, filenameAtom } from "./core/saving/file-state";
+import {
+  codeAtom,
+  cwdAtom,
+  filenameAtom,
+  lspWorkspaceAtom,
+} from "./core/saving/file-state";
 import { store } from "./core/state/jotai";
 import { patchFetch, patchVegaLoader } from "./core/static/files";
 import {
@@ -150,6 +155,16 @@ const mountOptionsSchema = z.object({
    * absolute working directory of the notebook
    */
   cwd: z.string().nullish().default(null),
+  /**
+   * LSP workspace information
+   */
+  lspWorkspace: z
+    .object({
+      rootUri: z.string(),
+      documentUri: z.string(),
+    })
+    .nullish()
+    .default(null),
   /**
    * notebook code
    */
@@ -287,6 +302,7 @@ function initStore(options: unknown) {
   // Files
   store.set(filenameAtom, parsedOptions.data.filename);
   store.set(cwdAtom, parsedOptions.data.cwd ?? null);
+  store.set(lspWorkspaceAtom, parsedOptions.data.lspWorkspace);
   store.set(codeAtom, parsedOptions.data.code);
   store.set(initialModeAtom, mode);
 

package/src/plugins/impl/DataTablePlugin.tsx

@@ -704,7 +704,7 @@ export const LoadingDataTableComponent = memo(
           <LoadingTable
             pageSize={
               props.totalRows !== TOO_MANY_ROWS && props.totalRows > 0
-                ? props.totalRows
+                ? Math.min(props.totalRows, props.pageSize)
                 : props.pageSize
             }
           />

package/src/plugins/impl/plotly/__tests__/selection.test.ts

@@ -117,6 +117,14 @@ describe("shouldHandleClickSelection", () => {
     expect(shouldHandleClickSelection([linePoint])).toBe(true);
   });
 
+  it("accepts waterfall clicks", () => {
+    const waterfallPoint = createPlotDatum({
+      data: { type: "waterfall" },
+    });
+
+    expect(shouldHandleClickSelection([waterfallPoint])).toBe(true);
+  });
+
   it("rejects non-line scatter marker clicks", () => {
     const markerPoint = createPlotDatum({
       data: { type: "scatter", mode: "markers" },
@@ -196,4 +204,18 @@ describe("extractPoints", () => {
 
     expect(extractPoints([point])).toEqual([{ x: 1, y: 2, z: 3 }]);
   });
+
+  it("returns x/y/pointIndex for waterfall clicks", () => {
+    const point = createPlotDatum({
+      x: "Revenue",
+      y: 400,
+      pointIndex: 1,
+      curveNumber: 0,
+      data: { type: "waterfall" },
+    });
+
+    expect(extractPoints([point])).toEqual([
+      { x: "Revenue", y: 400, pointIndex: 1, curveNumber: 0 },
+    ]);
+  });
 });

package/src/plugins/impl/plotly/selection.ts

@@ -227,6 +227,7 @@ export function shouldHandleClickSelection(
       type === "bar" ||
       type === "heatmap" ||
       type === "histogram" ||
+      type === "waterfall" ||
       isLinePoint(point)
     );
   });