@marimo-team/islands 0.16.5 → 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/network/requests-network.ts

@@ -409,5 +409,19 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponse);
     },
+    clearCache: async () => {
+      return getClient()
+        .POST("/api/cache/clear", {
+          body: {},
+        })
+        .then(handleResponseReturnNull);
+    },
+    getCacheInfo: async () => {
+      return getClient()
+        .POST("/api/cache/info", {
+          body: {},
+        })
+        .then(handleResponseReturnNull);
+    },
   };
 }

package/src/core/network/requests-static.ts

@@ -86,5 +86,7 @@ export function createStaticRequests(): EditRequests & RunRequests {
     listSecretKeys: throwNotInEditMode,
     writeSecret: throwNotInEditMode,
     invokeAiTool: throwNotInEditMode,
+    clearCache: throwNotInEditMode,
+    getCacheInfo: throwNotInEditMode,
   };
 }

package/src/core/network/requests-toasting.ts

@@ -64,6 +64,8 @@ export function createErrorToastingRequests(
     listSecretKeys: "Failed to fetch secrets",
     writeSecret: "Failed to write secret",
     invokeAiTool: "Failed to invoke AI tool",
+    clearCache: "Failed to clear cache",
+    getCacheInfo: "", // No toast
   };
 
   const handlers = {} as EditRequests & RunRequests;

package/src/core/network/types.ts

@@ -99,6 +99,8 @@ export type OpenTutorialRequest = schemas["OpenTutorialRequest"];
 export type TutorialId = OpenTutorialRequest["tutorialId"];
 export type InvokeAiToolRequest = schemas["InvokeAiToolRequest"];
 export type InvokeAiToolResponse = schemas["InvokeAiToolResponse"];
+export type ClearCacheRequest = schemas["ClearCacheRequest"];
+export type GetCacheInfoRequest = schemas["GetCacheInfoRequest"];
 
 /**
  * Requests sent to the BE during run/edit mode.
@@ -188,6 +190,9 @@ export interface EditRequests {
   writeSecret: (request: CreateSecretRequest) => Promise<null>;
   // AI Tool requests
   invokeAiTool: (request: InvokeAiToolRequest) => Promise<InvokeAiToolResponse>;
+  // Cache requests
+  clearCache: () => Promise<null>;
+  getCacheInfo: () => Promise<null>;
 }
 
 export type RequestKey = keyof (EditRequests & RunRequests);

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/saving/save-component.tsx

@@ -96,54 +96,61 @@ export function useSaveNotebook() {
   const store = useStore();
 
   // Save the notebook with the given filename
-  const saveNotebook = useEvent((filename: string, userInitiated: boolean) => {
-    const notebook = getNotebook();
-    const cells = notebookCells(notebook);
-    const cellIds = cells.map((cell) => cell.id);
-    const codes = cells.map((cell) => cell.code);
-    const cellNames = cells.map((cell) => cell.name);
-    const configs = getCellConfigs(notebook);
-    const connection = store.get(connectionAtom);
-    const autoSaveConfig = store.get(autoSaveConfigAtom);
-    const layout = store.get(layoutStateAtom);
-    const kioskMode = store.get(kioskModeAtom);
+  const saveNotebook = useEvent(
+    async (filename: string, userInitiated: boolean) => {
+      const connection = store.get(connectionAtom);
+      const autoSaveConfig = store.get(autoSaveConfigAtom);
+      const kioskMode = store.get(kioskModeAtom);
 
-    if (kioskMode) {
-      return;
-    }
+      if (kioskMode) {
+        return;
+      }
 
-    // Don't save if there are no cells
-    if (codes.length === 0) {
-      return;
-    }
+      // Don't save if we are not connected to a kernel
+      if (connection.state !== WebSocketState.OPEN) {
+        openAlert("Failed to save notebook: not connected to a kernel.");
+        return;
+      }
 
-    // Don't save if we are not connected to a kernel
-    if (connection.state !== WebSocketState.OPEN) {
-      openAlert("Failed to save notebook: not connected to a kernel.");
-      return;
-    }
+      Logger.log("saving to ", filename);
 
-    Logger.log("saving to ", filename);
-    sendSave({
-      cellIds: cellIds,
-      codes,
-      names: cellNames,
-      filename,
-      configs,
-      layout: getSerializedLayout(),
-      persist: true,
-    }).then(() => {
       if (userInitiated && autoSaveConfig.format_on_save) {
-        formatAll();
+        Logger.log("formatting notebook (onSave)");
+        await formatAll();
+      }
+
+      // Grab the latest notebook state, after formatting
+      const notebook = getNotebook();
+      const cells = notebookCells(notebook);
+      const cellIds = cells.map((cell) => cell.id);
+      const codes = cells.map((cell) => cell.code);
+      const cellNames = cells.map((cell) => cell.name);
+      const configs = getCellConfigs(notebook);
+      const layout = store.get(layoutStateAtom);
+
+      // Don't save if there are no cells
+      if (codes.length === 0) {
+        return;
       }
+
+      await sendSave({
+        cellIds: cellIds,
+        codes,
+        names: cellNames,
+        filename,
+        configs,
+        layout: getSerializedLayout(),
+        persist: true,
+      });
+
       setLastSavedNotebook({
         names: cellNames,
         codes,
         configs,
         layout,
       });
-    });
-  });
+    },
+  );
 
   // Save the notebook with the current filename, only if the filename exists
   const saveIfNotebookIsPersistent = useEvent((userInitiated = false) => {

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

@@ -551,6 +551,8 @@ export class PyodideBridge implements RunRequests, EditRequests {
   };
 
   invokeAiTool = throwNotImplemented;
+  clearCache = throwNotImplemented;
+  getCacheInfo = throwNotImplemented;
 
   private async putControlRequest(operation: object) {
     await this.rpc.proxy.request.bridge({

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/core/websocket/useMarimoWebSocket.tsx

@@ -24,6 +24,7 @@ import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
 import { Logger } from "@/utils/Logger";
 import { reloadSafe } from "@/utils/reload-safe";
 import { useAlertActions } from "../alerts/state";
+import { cacheInfoAtom } from "../cache/requests";
 import type { CellId, UIElementId } from "../cells/ids";
 import { useRunsActions } from "../cells/runs";
 import { focusAndScrollCellOutputIntoView } from "../cells/scrollCellIntoView";
@@ -87,6 +88,7 @@ export function useMarimoWebSocket(opts: {
   const setKioskMode = useSetAtom(kioskModeAtom);
   const setCapabilities = useSetAtom(capabilitiesAtom);
   const runtimeManager = useRuntimeManager();
+  const setCacheInfo = useSetAtom(cacheInfoAtom);
 
   const handleMessage = (e: MessageEvent<JsonString<OperationMessage>>) => {
     const msg = jsonParseWithSpecialChar(e.data);
@@ -248,6 +250,12 @@ export function useMarimoWebSocket(opts: {
       case "secret-keys-result":
         SECRETS_REGISTRY.resolve(msg.data.request_id as RequestId, msg.data);
         return;
+      case "cache-info-fetched":
+        setCacheInfo(msg.data);
+        return;
+      case "cache-cleared":
+        // Cache cleared, could refresh cache info if needed
+        return;
       case "data-source-connections":
         addDataSourceConnection({
           connections: msg.data.connections.map((conn) => ({

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/app/codemirror-completions.css

@@ -247,7 +247,7 @@
 
   .cm-completionInfo.cm-completionInfo-right-narrow {
     /* CodeMirror's default makes the completionInfo overlap with the
-      * compeletion symbol */
+      * completion symbol */
     left: 100%;
   }
 

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/custom.d.ts

@@ -16,6 +16,8 @@ interface JSON {
     text: string,
     reviver?: (this: any, key: string, value: any) => any,
   ): unknown;
+
+  rawJSON(value: string): unknown;
 }
 
 // Improve type inference for Array.filter with BooleanConstructor