@marimo-team/islands 0.17.0 → 0.17.1

package/src/core/lsp/__tests__/transport.test.ts

@@ -0,0 +1,233 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+import { WebSocketTransport } from "@open-rpc/client-js";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { Mocks } from "@/__mocks__/common";
+import { ReconnectingWebSocketTransport } from "../transport";
+
+// Mock the Logger
+vi.mock("@/utils/Logger", () => ({
+  Logger: Mocks.logger(),
+}));
+
+// Mock the WebSocketTransport
+vi.mock("@open-rpc/client-js", () => {
+  const mockWebSocketTransport = vi.fn();
+  mockWebSocketTransport.prototype.connect = vi.fn();
+  mockWebSocketTransport.prototype.close = vi.fn();
+  mockWebSocketTransport.prototype.sendData = vi.fn();
+
+  return {
+    WebSocketTransport: mockWebSocketTransport,
+  };
+});
+
+describe("ReconnectingWebSocketTransport", () => {
+  const mockWsUrl = "ws://localhost:8080/lsp";
+  let mockConnection: any;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    // Create a mock WebSocket connection with readyState
+    mockConnection = {
+      readyState: WebSocket.OPEN,
+    };
+
+    // Mock the WebSocketTransport constructor to set the connection
+    (WebSocketTransport as any).mockImplementation(function (this: any) {
+      this.connection = mockConnection;
+      this.connect = vi.fn().mockResolvedValue(undefined);
+      this.close = vi.fn();
+      this.sendData = vi.fn().mockResolvedValue({ result: "success" });
+    });
+  });
+
+  it("should create a transport with the provided URL function", () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    expect(transport).toBeDefined();
+    expect(getWsUrl).not.toHaveBeenCalled(); // URL function not called until connect
+  });
+
+  it("should connect successfully", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    await transport.connect();
+
+    expect(getWsUrl).toHaveBeenCalledTimes(1);
+    expect(WebSocketTransport).toHaveBeenCalledWith(mockWsUrl);
+  });
+
+  it("should wait for connection before connecting", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const waitForConnection = vi.fn().mockResolvedValue(undefined);
+    const transport = new ReconnectingWebSocketTransport({
+      getWsUrl,
+      waitForConnection,
+    });
+
+    await transport.connect();
+
+    expect(waitForConnection).toHaveBeenCalledTimes(1);
+    expect(getWsUrl).toHaveBeenCalledTimes(1);
+  });
+
+  it("should reuse the same connection promise if already connecting", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const waitForConnection = vi
+      .fn()
+      .mockImplementation(
+        () => new Promise((resolve) => setTimeout(resolve, 100)),
+      );
+    const transport = new ReconnectingWebSocketTransport({
+      getWsUrl,
+      waitForConnection,
+    });
+
+    // Start two connections concurrently
+    const promise1 = transport.connect();
+    const promise2 = transport.connect();
+
+    await Promise.all([promise1, promise2]);
+
+    // Should only create one delegate
+    expect(WebSocketTransport).toHaveBeenCalledTimes(1);
+    expect(waitForConnection).toHaveBeenCalledTimes(1);
+  });
+
+  it("should send data successfully when connected", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    await transport.connect();
+
+    const data: any = { method: "test", params: [] };
+    const result = await transport.sendData(data, 5000);
+
+    expect(result).toEqual({ result: "success" });
+  });
+
+  it("should reconnect when WebSocket is in CLOSED state", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    // First connection
+    await transport.connect();
+    expect(WebSocketTransport).toHaveBeenCalledTimes(1);
+
+    // Simulate WebSocket closing
+    mockConnection.readyState = WebSocket.CLOSED;
+
+    // Send data should trigger reconnection
+    const data: any = { method: "test", params: [] };
+    await transport.sendData(data, 5000);
+
+    // Should have created a new WebSocketTransport
+    expect(WebSocketTransport).toHaveBeenCalledTimes(2);
+  });
+
+  it("should reconnect when WebSocket is in CLOSING state", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    // First connection
+    await transport.connect();
+    expect(WebSocketTransport).toHaveBeenCalledTimes(1);
+
+    // Simulate WebSocket closing
+    mockConnection.readyState = WebSocket.CLOSING;
+
+    // Send data should trigger reconnection
+    const data: any = { method: "test", params: [] };
+    await transport.sendData(data, 5000);
+
+    // Should have created a new WebSocketTransport
+    expect(WebSocketTransport).toHaveBeenCalledTimes(2);
+  });
+
+  it("should close the transport and prevent reconnection", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    await transport.connect();
+    transport.close();
+
+    // Attempting to connect again should throw
+    await expect(transport.connect()).rejects.toThrow("Transport is closed");
+  });
+
+  it("should close old delegate when creating a new one", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    // First connection
+    await transport.connect();
+    const firstDelegate = (transport as any).delegate;
+    expect(firstDelegate).toBeDefined();
+
+    // Simulate connection closed
+    mockConnection.readyState = WebSocket.CLOSED;
+
+    // Reconnect by sending data
+    const data: any = { method: "test", params: [] };
+    await transport.sendData(data, 5000);
+
+    // Old delegate should have been closed
+    expect(firstDelegate.close).toHaveBeenCalled();
+  });
+
+  it("should handle connection failures gracefully", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const connectionError = new Error("Connection failed");
+
+    // Mock connect to fail
+    (WebSocketTransport as any).mockImplementationOnce(function (this: any) {
+      this.connection = mockConnection;
+      this.connect = vi.fn().mockRejectedValue(connectionError);
+      this.close = vi.fn();
+      this.sendData = vi.fn();
+    });
+
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    await expect(transport.connect()).rejects.toThrow("Connection failed");
+
+    // Delegate should be cleared after failure
+    expect((transport as any).delegate).toBeUndefined();
+  });
+
+  it("should handle waitForConnection failures", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const waitError = new Error("Wait failed");
+    const waitForConnection = vi.fn().mockRejectedValue(waitError);
+
+    const transport = new ReconnectingWebSocketTransport({
+      getWsUrl,
+      waitForConnection,
+    });
+
+    await expect(transport.connect()).rejects.toThrow("Wait failed");
+
+    // Should not have created a delegate
+    expect(WebSocketTransport).not.toHaveBeenCalled();
+  });
+
+  it("should automatically reconnect on sendData after connection loss", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const transport = new ReconnectingWebSocketTransport({ getWsUrl });
+
+    // Don't connect initially
+    // Simulate WebSocket in closed state (no delegate exists)
+    expect((transport as any).delegate).toBeUndefined();
+
+    // Send data should trigger automatic connection
+    const data: any = { method: "test", params: [] };
+    await transport.sendData(data, 5000);
+
+    expect(WebSocketTransport).toHaveBeenCalledTimes(1);
+  });
+});

package/src/core/lsp/transport.ts

@@ -0,0 +1,139 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { WebSocketTransport } from "@open-rpc/client-js";
+import type { JSONRPCRequestData } from "@open-rpc/client-js/build/Request";
+import { Transport } from "@open-rpc/client-js/build/transports/Transport";
+import { Logger } from "@/utils/Logger";
+
+export interface ReconnectingWebSocketTransportOptions {
+  /**
+   * Function that returns the WebSocket URL to connect to.
+   */
+  getWsUrl: () => string;
+
+  /**
+   * Optional function to wait for before attempting to connect.
+   * This is useful for ensuring dependencies (like the runtime) are ready.
+   */
+  waitForConnection?: () => Promise<void>;
+}
+
+/**
+ * A WebSocket transport that automatically reconnects when the connection is lost.
+ * This handles cases like computer sleep/wake or network interruptions.
+ */
+export class ReconnectingWebSocketTransport extends Transport {
+  private delegate: WebSocketTransport | undefined;
+  private readonly options: ReconnectingWebSocketTransportOptions;
+  private connectionPromise: Promise<void> | undefined;
+  private isClosed = false;
+
+  constructor(options: ReconnectingWebSocketTransportOptions) {
+    super();
+    this.options = options;
+    this.delegate = undefined;
+  }
+
+  /**
+   * Create a new WebSocket delegate, replacing any existing one.
+   */
+  private createDelegate(): WebSocketTransport {
+    // Close the old delegate if it exists
+    if (this.delegate) {
+      try {
+        this.delegate.close();
+      } catch (error) {
+        Logger.warn("Error closing old WebSocket delegate", error);
+      }
+    }
+
+    // Create a new delegate
+    this.delegate = new WebSocketTransport(this.options.getWsUrl());
+    return this.delegate;
+  }
+
+  /**
+   * Check if the current delegate's WebSocket is in a closed or closing state.
+   */
+  private isDelegateClosedOrClosing(): boolean {
+    if (!this.delegate) {
+      return true;
+    }
+
+    // Access the internal connection to check its readyState
+    const ws = this.delegate.connection;
+    if (!ws) {
+      return true;
+    }
+
+    // WebSocket.CLOSING = 2, WebSocket.CLOSED = 3
+    return (
+      ws.readyState === WebSocket.CLOSING || ws.readyState === WebSocket.CLOSED
+    );
+  }
+
+  override async connect() {
+    // Don't reconnect if explicitly closed
+    if (this.isClosed) {
+      throw new Error("Transport is closed");
+    }
+
+    // If already connecting, wait for that connection
+    if (this.connectionPromise) {
+      return this.connectionPromise;
+    }
+
+    this.connectionPromise = (async () => {
+      try {
+        // Wait for dependencies to be ready (e.g., runtime connection)
+        if (this.options.waitForConnection) {
+          await this.options.waitForConnection();
+        }
+
+        // Create a new delegate if needed
+        let delegate = this.delegate;
+        if (!delegate || this.isDelegateClosedOrClosing()) {
+          delegate = this.createDelegate();
+        }
+
+        // Connect the delegate
+        await delegate.connect();
+        Logger.log("WebSocket transport connected successfully");
+      } catch (error) {
+        Logger.error("WebSocket transport connection failed", error);
+        // Clear the delegate on failure so we create a new one on retry
+        this.delegate = undefined;
+        throw error;
+      } finally {
+        this.connectionPromise = undefined;
+      }
+    })();
+
+    return this.connectionPromise;
+  }
+
+  override close() {
+    this.isClosed = true;
+    this.delegate?.close();
+    this.delegate = undefined;
+    this.connectionPromise = undefined;
+  }
+
+  override async sendData(
+    data: JSONRPCRequestData,
+    timeout: number | null | undefined,
+  ) {
+    // If the delegate is closed or closing, try to reconnect
+    if (this.isDelegateClosedOrClosing()) {
+      Logger.warn("WebSocket is closed or closing, attempting to reconnect");
+      try {
+        await this.connect();
+      } catch (error) {
+        Logger.error("Failed to reconnect WebSocket", error);
+        throw error;
+      }
+    }
+
+    // Send the data using the delegate
+    return this.delegate?.sendData(data, timeout);
+  }
+}

package/src/core/rtc/state.ts

@@ -1,13 +1,18 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import { atomWithStorage } from "jotai/utils";
 import { once } from "lodash-es";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
 import { getFeatureFlag } from "../config/feature-flag";
 
 /**
  * The username for the current user when using real-time collaboration.
  * This is stored in localStorage.
  */
-export const usernameAtom = atomWithStorage<string>("marimo:rtc:username", "");
+export const usernameAtom = atomWithStorage<string>(
+  "marimo:rtc:username",
+  "",
+  jotaiJsonStorage,
+);
 
 /**
  * Whether RTC is enabled.

package/src/core/vscode/is-in-vscode.ts

@@ -0,0 +1,9 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+/**
+ * Whether the current environment is in the VSCode extension
+ */
+export function isInVscodeExtension(): boolean {
+  // We check if the document has a data-vscode-theme-kind attribute
+  return document.querySelector("[data-vscode-theme-kind]") !== null;
+}

package/src/core/wasm/store.ts

@@ -4,7 +4,7 @@ import {
   compressToEncodedURIComponent,
   decompressFromEncodedURIComponent,
 } from "lz-string";
-import { TypedLocalStorage } from "@/utils/localStorage";
+import { TypedLocalStorage } from "@/utils/storage/typed";
 import { PyodideRouter } from "./router";
 
 export interface FileStore {

package/src/css/app/Cell.css

@@ -34,6 +34,37 @@
     z-index: 1;
   }
 
+  /* Styling for cells marked for deletion by AI */
+  &:has(.mo-ai-deleted-cell) {
+    opacity: 0.7;
+    position: relative;
+
+    .output-area,
+    .cm-gutters,
+    .cm {
+      background-color: var(--red-2);
+    }
+
+    /* Add a red glow to indicate deletion */
+    &::before {
+      content: "";
+      position: absolute;
+      inset: 0;
+      border-radius: 10px;
+      pointer-events: none;
+      opacity: 0.6;
+      box-shadow: 0px 0px 10px 0px var(--red-9);
+      z-index: 1;
+    }
+
+    /* Add strikethrough effect to the code */
+    .cm-content {
+      text-decoration: line-through;
+      text-decoration-color: var(--red-9);
+      text-decoration-thickness: 1px;
+    }
+  }
+
   &:focus-within {
     z-index: 20;
   }

package/src/css/globals.css

@@ -135,6 +135,11 @@
   a.hyperlink:hover {
     text-decoration: underline;
   }
+
+  button:not(:disabled),
+  [role="button"]:not(:disabled) {
+    cursor: pointer;
+  }
 }
 
 @theme {

package/src/hooks/useIframeCapabilities.ts

@@ -0,0 +1,14 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { useMemo } from "react";
+import {
+  getIframeCapabilities,
+  type IframeCapabilities,
+} from "@/utils/capabilities";
+
+/**
+ * React hook to access iframe capabilities
+ */
+export function useIframeCapabilities(): IframeCapabilities {
+  return useMemo(() => getIframeCapabilities(), []);
+}

package/src/hooks/useLocalStorage.ts

@@ -1,7 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { useState } from "react";
-import { TypedLocalStorage } from "@/utils/localStorage";
+import { TypedLocalStorage } from "@/utils/storage/typed";
 
 /**
  * React hook to use localStorage

package/src/plugins/core/RenderHTML.tsx

@@ -5,14 +5,26 @@ import parse, {
   Element,
   type HTMLReactParserOptions,
 } from "html-react-parser";
-import React, { isValidElement, type JSX, type ReactNode, useId } from "react";
+import React, {
+  isValidElement,
+  type JSX,
+  type ReactNode,
+  useMemo,
+  useRef,
+} from "react";
 import { CopyClipboardIcon } from "@/components/icons/copy-icon";
+import { sanitizeHtml, useSanitizeHtml } from "./sanitize";
 
 type ReplacementFn = NonNullable<HTMLReactParserOptions["replace"]>;
 type TransformFn = NonNullable<HTMLReactParserOptions["transform"]>;
 
 interface Options {
   html: string;
+  /**
+   * Whether to sanitize the HTML.
+   * @default true
+   */
+  alwaysSanitizeHtml?: boolean;
   additionalReplacements?: ReplacementFn[];
 }
 
@@ -110,9 +122,9 @@ const addCopyButtonToCodehilite: TransformFn = (
 };
 
 const CopyableCode = ({ children }: { children: ReactNode }) => {
-  const id = useId();
+  const ref = useRef<HTMLDivElement>(null);
   return (
-    <div className="relative group codehilite-wrapper" id={id}>
+    <div className="relative group codehilite-wrapper" ref={ref}>
       {children}
 
       <div className="absolute top-2 right-2 opacity-0 group-hover:opacity-100 transition-opacity">
@@ -120,7 +132,7 @@ const CopyableCode = ({ children }: { children: ReactNode }) => {
           tooltip={false}
           className="p-1"
           value={() => {
-            const codeElement = document.getElementById(id)?.firstChild;
+            const codeElement = ref.current?.firstChild;
             if (codeElement) {
               return codeElement.textContent || "";
             }
@@ -132,7 +144,51 @@ const CopyableCode = ({ children }: { children: ReactNode }) => {
   );
 };
 
-export const renderHTML = ({ html, additionalReplacements = [] }: Options) => {
+/**
+ *
+ * @param html - The HTML to render.
+ * @param additionalReplacements - Additional replacements to apply to the HTML.
+ * @param alwaysSanitizeHtml - Whether to sanitize the HTML.
+ * @returns
+ */
+export const renderHTML = ({
+  html,
+  additionalReplacements = [],
+  alwaysSanitizeHtml = true,
+}: Options) => {
+  return (
+    <RenderHTML
+      html={html}
+      alwaysSanitizeHtml={alwaysSanitizeHtml}
+      additionalReplacements={additionalReplacements}
+    />
+  );
+};
+
+const RenderHTML = ({
+  html,
+  additionalReplacements = [],
+  alwaysSanitizeHtml,
+}: Options) => {
+  const shouldSanitizeHtml = useSanitizeHtml();
+
+  const sanitizedHtml = useMemo(() => {
+    if (alwaysSanitizeHtml || shouldSanitizeHtml) {
+      return sanitizeHtml(html);
+    }
+    return html;
+  }, [html, alwaysSanitizeHtml, shouldSanitizeHtml]);
+
+  return parseHtml({
+    html: sanitizedHtml,
+    additionalReplacements,
+  });
+};
+
+function parseHtml({
+  html,
+  additionalReplacements = [],
+}: Pick<Options, "html" | "additionalReplacements">) {
   const renderFunctions: ReplacementFn[] = [
     replaceValidTags,
     replaceValidIframes,
@@ -166,4 +222,8 @@ export const renderHTML = ({ html, additionalReplacements = [] }: Options) => {
       return reactNode as JSX.Element;
     },
   });
+}
+
+export const visibleForTesting = {
+  parseHtml,
 };