@marimo-team/islands 0.22.5-dev9 → 0.22.6-dev1

package/src/components/editor/__tests__/Output.test.tsx

@@ -1,7 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { render, screen } from "@testing-library/react";
 import { describe, expect, it } from "vitest";
-import { OutputRenderer } from "../Output";
+import { cellId } from "@/__tests__/branded";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { OutputArea, OutputRenderer } from "../Output";
 
 describe("OutputRenderer renderFallback prop", () => {
   it("should use renderFallback for unsupported mimetypes", () => {
@@ -65,6 +67,39 @@ describe("OutputRenderer renderFallback prop", () => {
   });
 });
 
+describe("OutputArea null/undefined handling", () => {
+  it("should render null when output is null", () => {
+    const { container } = render(
+      <TooltipProvider>
+        <OutputArea
+          output={null}
+          cellId={cellId("test")}
+          stale={false}
+          loading={false}
+          allowExpand={true}
+        />
+      </TooltipProvider>,
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should render null when output is undefined", () => {
+    const { container } = render(
+      <TooltipProvider>
+        <OutputArea
+          // @ts-expect-error -- testing runtime safety for undefined output
+          output={undefined}
+          cellId={cellId("test")}
+          stale={false}
+          loading={false}
+          allowExpand={true}
+        />
+      </TooltipProvider>,
+    );
+    expect(container.innerHTML).toBe("");
+  });
+});
+
 describe("OutputRenderer image and SVG rendering", () => {
   const plainSvgString =
     '<svg><rect x="0" y="0" width="10" height="10"></rect></svg>';

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

@@ -1209,6 +1209,47 @@ describe("cell reducer", () => {
     ]);
   });
 
+  it("does not crash when setStdinResponse has out-of-bounds outputIndex", () => {
+    const STDOUT: OutputMessage = {
+      channel: "stdout",
+      mimetype: "text/plain",
+      data: "hello!",
+      timestamp: 1,
+    };
+
+    // Set the cell to running with a console output
+    actions.prepareForRun({ cellId: firstCellId });
+    actions.handleCellMessage({
+      cell_id: firstCellId,
+      output: undefined,
+      console: null,
+      status: "running",
+      stale_inputs: null,
+      timestamp: new Date(20).getTime() as Seconds,
+    });
+    actions.handleCellMessage({
+      cell_id: firstCellId,
+      output: undefined,
+      console: STDOUT,
+      status: undefined,
+      stale_inputs: null,
+      timestamp: new Date(22).getTime() as Seconds,
+    });
+
+    // Try to set stdin response with an out-of-bounds index
+    // This should not crash - it should return state unchanged
+    actions.setStdinResponse({
+      response: "test",
+      cellId: firstCellId,
+      outputIndex: 999,
+    });
+
+    // Cell state should be unchanged
+    const cell = cells[0];
+    expect(cell.consoleOutputs).toHaveLength(1);
+    expect(cell.consoleOutputs[0]).toMatchObject(STDOUT);
+  });
+
   it("can receive console when the cell is idle and will clear when starts again", () => {
     const OLD_STDOUT: OutputMessage = {
       channel: "stdout",

package/src/core/cells/__tests__/collapseConsoleOutputs.test.ts

@@ -241,6 +241,44 @@ describe("collapseConsoleOutputs", () => {
     expect(result[2].data).toBe("<pre>E\nF\nG\nH\n</pre>");
   });
 
+  it("should not crash when truncating with a single output at the limit boundary", () => {
+    // Create outputs that push truncation to the exact boundary
+    const consoleOutputs: OutputMessage[] = [
+      {
+        mimetype: "text/html",
+        channel: "output",
+        data: "<div>html1</div>",
+        timestamp: 0,
+      },
+      {
+        mimetype: "text/html",
+        channel: "output",
+        data: "<div>html2</div>",
+        timestamp: 0,
+      },
+    ];
+    // With limit=1, truncation must handle edge cases gracefully
+    const result = collapseConsoleOutputs(consoleOutputs, 1);
+    expect(result[0].data).toContain("Streaming output truncated");
+  });
+
+  it("should handle truncation when cutoff indexes past the end of the array", () => {
+    // With maxLines=0, the truncation loop never runs, causing cutoff
+    // to index past the array. This exercises the `output == null`
+    // defensive branch in truncateHead().
+    const consoleOutputs: OutputMessage[] = [
+      {
+        mimetype: "text/html",
+        channel: "output",
+        data: "<div>content</div>",
+        timestamp: 0,
+      },
+    ];
+    const result = collapseConsoleOutputs(consoleOutputs, 0);
+    expect(result).toHaveLength(1);
+    expect(result[0].data).toContain("Streaming output truncated");
+  });
+
   describe("ANSI escape sequences", () => {
     it("should handle cursor movement with collapse", () => {
       const consoleOutputs: OutputMessage[] = [

package/src/core/cells/cells.ts

@@ -947,7 +947,7 @@ const {
       cellReducer: (cell) => {
         const consoleOutputs = [...cell.consoleOutputs];
         const stdinOutput = consoleOutputs[outputIndex];
-        if (stdinOutput.channel !== "stdin") {
+        if (stdinOutput == null || stdinOutput.channel !== "stdin") {
           Logger.warn("Expected stdin output");
           return cell;
         }

package/src/core/cells/collapseConsoleOutputs.tsx

@@ -108,6 +108,9 @@ function truncateHead(consoleOutputs: OutputMessage[], limit: number) {
     timestamp: -1,
   };
   const output = consoleOutputs[cutoff];
+  if (output == null) {
+    return [warningOutput, ...consoleOutputs.slice(cutoff + 1)];
+  }
   if (output.mimetype === "text/plain") {
     invariant(typeof output.data === "string", "expected string");
     const outputLines = output.data.split("\n");

package/src/core/cells/document-changes.ts

@@ -25,6 +25,7 @@ import type { NotebookDocumentTransactionRequest } from "../network/types";
 import { store } from "../state/jotai";
 import type { CellActions, NotebookState } from "./cells";
 import type { CellId } from "./ids";
+import { SCRATCH_CELL_ID } from "./ids";
 import type { CellData } from "./types";
 
 export type DocumentChange =
@@ -572,10 +573,21 @@ const flushChanges = debounce(() => {
   void getRequestClient().sendDocumentTransaction({ changes });
 }, 400);
 
+function isScratchChange(change: DocumentChange): boolean {
+  if ("cellId" in change && change.cellId === SCRATCH_CELL_ID) {
+    return true;
+  }
+  return false;
+}
+
 function enqueue(change: DocumentChange) {
   if (store.get(kioskModeAtom)) {
     return;
   }
+  // The scratchpad cell is local-only — don't sync it to the document.
+  if (isScratchChange(change)) {
+    return;
+  }
   pendingChanges.push(change);
   flushChanges();
 }

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)
     );
   });