@marimo-team/islands 0.17.2 → 0.17.3

package/src/core/codemirror/lsp/transports.ts

@@ -10,10 +10,12 @@ import { getRuntimeManager } from "../../runtime/config";
  * before connecting to the LSP server.
  *
  * @param serverName - The name of the LSP server.
+ * @param onReconnect - Optional callback to call after reconnection (e.g., to resync documents).
  * @returns The transport.
  */
 export function createTransport(
   serverName: "pylsp" | "basedpyright" | "copilot" | "ty",
+  onReconnect?: () => Promise<void>,
 ) {
   const runtimeManager = getRuntimeManager();
   return new ReconnectingWebSocketTransport({
@@ -21,5 +23,6 @@ export function createTransport(
     waitForConnection: async () => {
       await waitForConnectionOpen();
     },
+    onReconnect,
   });
 }

package/src/core/datasets/__tests__/data-source.test.ts

@@ -270,6 +270,7 @@ describe("add table list", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "public",
+      dialect: "sqlite",
     });
 
     const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
@@ -283,6 +284,7 @@ describe("add table list", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "public",
+      dialect: "sqlite",
     };
 
     const tableList: DataTable[] = [
@@ -344,6 +346,7 @@ describe("add table list", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "non_existent",
+      dialect: "sqlite",
     });
 
     const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
@@ -407,6 +410,7 @@ describe("add table", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "public",
+      dialect: "sqlite",
     });
 
     const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
@@ -420,6 +424,7 @@ describe("add table", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "public",
+      dialect: "sqlite",
     };
 
     const table: DataTable = {
@@ -476,6 +481,7 @@ describe("add table", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "non_existent",
+      dialect: "sqlite",
     });
 
     const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);

package/src/core/datasets/data-source-connections.ts

@@ -47,6 +47,7 @@ export interface SQLTableContext {
   engine: string;
   database: string;
   schema: string;
+  dialect: string;
   defaultSchema?: string | null;
   defaultDatabase?: string | null;
 }

package/src/core/dom/__tests__/outline.test.ts

@@ -137,6 +137,49 @@ describe("parseOutline", () => {
     `);
   });
 
+  it("can parse markdown outline", () => {
+    const markdown = `
+    <span class="markdown">
+      <h1 id="introduction">Introduction</h1>
+      <h2 id="getting-started">Getting Started</h2>
+      <h3 id="installation">Installation</h3>
+    </span>
+    `;
+    const outline = parseOutline({
+      mimetype: "text/markdown",
+      timestamp: 0,
+      channel: "output",
+      data: markdown,
+    });
+    expect(outline).toMatchInlineSnapshot(`
+      {
+        "items": [
+          {
+            "by": {
+              "id": "introduction",
+            },
+            "level": 1,
+            "name": "Introduction",
+          },
+          {
+            "by": {
+              "id": "getting-started",
+            },
+            "level": 2,
+            "name": "Getting Started",
+          },
+          {
+            "by": {
+              "id": "installation",
+            },
+            "level": 3,
+            "name": "Installation",
+          },
+        ],
+      }
+    `);
+  });
+
   it("can handle non-html outline", () => {
     const html = "foo";
     const outline = parseOutline({

package/src/core/dom/outline.ts

@@ -68,7 +68,7 @@ export function parseOutline(output: OutputMessage | null): Outline | null {
     return null;
   }
 
-  if (output.mimetype !== "text/html") {
+  if (output.mimetype !== "text/html" && output.mimetype !== "text/markdown") {
     return null;
   }
 

package/src/core/islands/main.ts

@@ -36,6 +36,7 @@ import {
 } from "../kernel/handlers";
 import { queryParamHandlers } from "../kernel/queryParamHandlers";
 import { RuntimeState } from "../kernel/RuntimeState";
+import { initialModeAtom } from "../mode";
 import type { RequestId } from "../network/DeferredRequestRegistry";
 import { requestClientAtom } from "../network/requests";
 import { store } from "../state/jotai";
@@ -57,6 +58,7 @@ import { dismissIslandsLoadingToast, toastIslandsLoading } from "./toast";
 export async function initialize() {
   // Setup networking
   store.set(requestClientAtom, IslandsPyodideBridge.INSTANCE);
+  store.set(initialModeAtom, "read");
 
   // This will display all the static HTML content.
   initializePlugins();

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

@@ -230,4 +230,61 @@ describe("ReconnectingWebSocketTransport", () => {
 
     expect(WebSocketTransport).toHaveBeenCalledTimes(1);
   });
+
+  it("should call onReconnect callback after reconnection", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const onReconnect = vi.fn().mockResolvedValue(undefined);
+    const transport = new ReconnectingWebSocketTransport({
+      getWsUrl,
+      onReconnect,
+    });
+
+    // First connection - callback should not be called
+    await transport.connect();
+    expect(onReconnect).not.toHaveBeenCalled();
+
+    // Simulate connection loss
+    mockConnection.readyState = WebSocket.CLOSED;
+
+    // Reconnect - callback should be called this time
+    const data: any = { method: "test", params: [] };
+    await transport.sendData(data, 5000);
+
+    expect(onReconnect).toHaveBeenCalledTimes(1);
+  });
+
+  it("should not call onReconnect callback on first connection", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const onReconnect = vi.fn().mockResolvedValue(undefined);
+    const transport = new ReconnectingWebSocketTransport({
+      getWsUrl,
+      onReconnect,
+    });
+
+    await transport.connect();
+
+    expect(onReconnect).not.toHaveBeenCalled();
+  });
+
+  it("should handle onReconnect callback errors gracefully", async () => {
+    const getWsUrl = vi.fn(() => mockWsUrl);
+    const reconnectError = new Error("Reconnect callback failed");
+    const onReconnect = vi.fn().mockRejectedValue(reconnectError);
+    const transport = new ReconnectingWebSocketTransport({
+      getWsUrl,
+      onReconnect,
+    });
+
+    // First connection
+    await transport.connect();
+
+    // Simulate connection loss
+    mockConnection.readyState = WebSocket.CLOSED;
+
+    // Reconnect - should propagate the error from onReconnect
+    const data: any = { method: "test", params: [] };
+    await expect(transport.sendData(data, 5000)).rejects.toThrow(
+      "Reconnect callback failed",
+    );
+  });
 });

package/src/core/lsp/transport.ts

@@ -15,6 +15,12 @@ export interface ReconnectingWebSocketTransportOptions {
    * This is useful for ensuring dependencies (like the runtime) are ready.
    */
   waitForConnection?: () => Promise<void>;
+
+  /**
+   * Optional callback that is called after a successful reconnection.
+   * This allows the LSP client to re-synchronize state (e.g., re-send document open notifications).
+   */
+  onReconnect?: () => Promise<void>;
 }
 
 /**
@@ -26,6 +32,7 @@ export class ReconnectingWebSocketTransport extends Transport {
   private readonly options: ReconnectingWebSocketTransportOptions;
   private connectionPromise: Promise<void> | undefined;
   private isClosed = false;
+  private hasConnectedBefore = false;
 
   constructor(options: ReconnectingWebSocketTransportOptions) {
     super();
@@ -98,6 +105,15 @@ export class ReconnectingWebSocketTransport extends Transport {
         // Connect the delegate
         await delegate.connect();
         Logger.log("WebSocket transport connected successfully");
+
+        // If this is a reconnection, call the onReconnect callback
+        const isReconnection = this.hasConnectedBefore;
+        this.hasConnectedBefore = true;
+
+        if (isReconnection && this.options.onReconnect) {
+          Logger.log("Calling onReconnect callback to re-synchronize state");
+          await this.options.onReconnect();
+        }
       } catch (error) {
         Logger.error("WebSocket transport connection failed", error);
         // Clear the delegate on failure so we create a new one on retry

package/src/plugins/core/RenderHTML.tsx

@@ -13,6 +13,7 @@ import React, {
   useRef,
 } from "react";
 import { CopyClipboardIcon } from "@/components/icons/copy-icon";
+import { QueryParamPreservingLink } from "@/components/ui/query-param-preserving-link";
 import { sanitizeHtml, useSanitizeHtml } from "./sanitize";
 
 type ReplacementFn = NonNullable<HTMLReactParserOptions["replace"]>;
@@ -106,6 +107,30 @@ const replaceSrcScripts = (domNode: DOMNode): JSX.Element | undefined => {
   }
 };
 
+const preserveQueryParamsInAnchorLinks: TransformFn = (
+  reactNode: ReactNode,
+  domNode: DOMNode,
+): JSX.Element | undefined => {
+  if (domNode instanceof Element && domNode.name === "a") {
+    const href = domNode.attribs.href;
+    // Only handle anchor links (starting with #)
+    if (href?.startsWith("#") && !href.startsWith("#code/")) {
+      // Get the children from the parsed React node
+      let children: ReactNode = null;
+      if (isValidElement(reactNode) && "props" in reactNode) {
+        const props = reactNode.props as { children?: ReactNode };
+        children = props.children;
+      }
+
+      return (
+        <QueryParamPreservingLink href={href} {...domNode.attribs}>
+          {children}
+        </QueryParamPreservingLink>
+      );
+    }
+  }
+};
+
 // Add copy button to codehilite blocks
 const addCopyButtonToCodehilite: TransformFn = (
   reactNode: ReactNode,
@@ -198,6 +223,7 @@ function parseHtml({
 
   const transformFunctions: TransformFn[] = [
     addCopyButtonToCodehilite,
+    preserveQueryParamsInAnchorLinks,
     removeWrappingBodyTags,
     removeWrappingHtmlTags,
   ];

package/src/plugins/core/sanitize.ts

@@ -18,7 +18,14 @@ const sanitizeHtmlAtom = atom<boolean>((get) => {
     return false;
   }
 
-  const isInAppMode = getInitialAppMode() === "read";
+  let isInAppMode = true;
+  try {
+    isInAppMode = getInitialAppMode() === "read";
+  } catch {
+    // If it fails to get the initial app mode, we default to sanitizing.
+    return true;
+  }
+
   // Apps need to run javascript and load external resources.
   if (isInAppMode) {
     return false;

package/src/plugins/impl/DataTablePlugin.tsx

@@ -57,7 +57,6 @@ import { Alert, AlertTitle } from "@/components/ui/alert";
 import { Button } from "@/components/ui/button";
 import { DelayMount } from "@/components/utils/delay-mount";
 import { type CellId, findCellId } from "@/core/cells/ids";
-import { getFeatureFlag } from "@/core/config/feature-flag";
 import { slotsController } from "@/core/slots/slots";
 import { store } from "@/core/state/jotai";
 import { isStaticNotebook } from "@/core/static/static-state";
@@ -85,15 +84,12 @@ import {
 type CsvURL = string;
 export type TableData<T> = T[] | CsvURL;
 
-interface ColumnSummariesArgs {
-  precompute: boolean;
-}
-
 interface ColumnSummaries<T = unknown> {
   data: TableData<T> | null | undefined;
   stats: Record<ColumnName, ColumnHeaderStats>;
   bin_values: Record<ColumnName, BinValues>;
   value_counts: Record<ColumnName, ValueCounts>;
+  show_charts: boolean;
   is_disabled?: boolean;
 }
 
@@ -198,9 +194,7 @@ interface Data<T> {
 // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
 type DataTableFunctions = {
   download_as: DownloadAsArgs;
-  get_column_summaries: <T>(
-    opts: ColumnSummariesArgs,
-  ) => Promise<ColumnSummaries<T>>;
+  get_column_summaries: <T>(opts: {}) => Promise<ColumnSummaries<T>>;
   search: <T>(req: {
     sort?: {
       by: string;
@@ -281,19 +275,16 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
   )
   .withFunctions<DataTableFunctions>({
     download_as: DownloadAsSchema,
-    get_column_summaries: rpc
-      .input(z.object({ precompute: z.boolean() }))
-      .output(
-        z.object({
-          data: z
-            .union([z.string(), z.array(z.object({}).passthrough())])
-            .nullable(),
-          stats: z.record(z.string(), columnStats),
-          bin_values: z.record(z.string(), binValues),
-          value_counts: z.record(z.string(), valueCounts),
-          is_disabled: z.boolean().optional(),
-        }),
-      ),
+    get_column_summaries: rpc.input(z.looseObject({})).output(
+      z.object({
+        data: z.union([z.string(), z.array(z.looseObject({}))]).nullable(),
+        stats: z.record(z.string(), columnStats),
+        bin_values: z.record(z.string(), binValues),
+        value_counts: z.record(z.string(), valueCounts),
+        show_charts: z.boolean(),
+        is_disabled: z.boolean().optional(),
+      }),
+    ),
     search: rpc
       .input(
         z.object({
@@ -598,8 +589,6 @@ export const LoadingDataTableComponent = memo(
       );
     }, [data?.totalRows]);
 
-    const precompute = getFeatureFlag("performant_table_charts");
-
     // Column summaries
     const { data: columnSummaries, error: columnSummariesError } = useAsyncData<
       ColumnSummaries<T>
@@ -607,9 +596,15 @@ export const LoadingDataTableComponent = memo(
       // TODO: props.get_column_summaries is always true,
       // so we are unable to detect if the function is registered
       if (props.totalRows === 0 || !props.showColumnSummaries) {
-        return { data: null, stats: {}, bin_values: {}, value_counts: {} };
+        return {
+          data: null,
+          stats: {},
+          bin_values: {},
+          value_counts: {},
+          show_charts: false,
+        };
       }
-      return props.get_column_summaries({ precompute });
+      return props.get_column_summaries({});
     }, [
       props.get_column_summaries,
       props.showColumnSummaries,
@@ -778,8 +773,7 @@ const DataTableComponent = ({
       columnSummaries.bin_values,
       columnSummaries.value_counts,
       {
-        includeCharts: Boolean(columnSummaries.data),
-        usePreComputedValues: getFeatureFlag("performant_table_charts"),
+        includeCharts: columnSummaries.show_charts,
       },
     );
   }, [fieldTypes, columnSummaries]);

package/src/plugins/impl/data-explorer/ConnectedDataExplorerComponent.tsx

@@ -4,14 +4,15 @@ import * as cql from "compassql/build/src";
 import { createStore, Provider, useAtomValue } from "jotai";
 import { ListFilterIcon } from "lucide-react";
 import React, { type JSX, useMemo } from "react";
-import { VegaLite } from "react-vega";
-import type { VegaLiteProps } from "react-vega/lib/VegaLite";
+import { VegaEmbed, type VegaEmbedProps } from "react-vega";
+import { tooltipHandler } from "@/components/charts/tooltip";
+import { augmentSpecWithData } from "@/components/data-table/charts/chart-spec/spec";
 import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
 import { Tooltip } from "@/components/ui/tooltip";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import { useOnMount } from "@/hooks/useLifecycle";
-import { useTheme } from "@/theme/useTheme";
+import { type ResolvedTheme, useTheme } from "@/theme/useTheme";
 import { Objects } from "@/utils/objects";
 import { ErrorBanner } from "../common/error-banner";
 import { vegaLoadData } from "../vega/loader";
@@ -69,14 +70,20 @@ interface DataTableProps extends DataExplorerData {
   setValue: (value: DataExplorerState) => void;
 }
 
-const ACTIONS: VegaLiteProps["actions"] = {
-  export: { svg: true, png: true },
-  source: false,
-  compiled: false,
-  editor: false,
-};
-
-const PADDING = { left: 20, right: 20, top: 20, bottom: 20 };
+function chartOptions(theme: ResolvedTheme): VegaEmbedProps["options"] {
+  return {
+    padding: { left: 20, right: 20, top: 20, bottom: 20 },
+    actions: {
+      export: { svg: true, png: true },
+      source: false,
+      compiled: false,
+      editor: false,
+    },
+    theme: theme === "dark" ? "dark" : undefined,
+    tooltip: tooltipHandler.call,
+    renderer: "canvas",
+  };
+}
 
 export default ConnectedDataExplorerComponent;
 
@@ -128,15 +135,14 @@ export const DataExplorerComponent = ({
       return <ColumnSummary schema={schema} />;
     }
 
+    const spec = mainPlot.spec;
+    const responsiveSpec = makeResponsive(spec);
+    // TODO: We can optimize by updating the data dynamically. https://github.com/vega/react-vega?tab=readme-ov-file#recipes
+    const augmentedSpec = augmentSpecWithData(responsiveSpec, chartData);
+
     return (
       <div className="flex overflow-y-auto justify-center items-center flex-1 w-[90%]">
-        <VegaLite
-          data={{ source: chartData }}
-          padding={PADDING}
-          actions={ACTIONS}
-          spec={makeResponsive(mainPlot.spec)}
-          theme={theme === "dark" ? "dark" : undefined}
-        />
+        <VegaEmbed spec={augmentedSpec} options={chartOptions(theme)} />
       </div>
     );
   };
@@ -205,12 +211,10 @@ export const DataExplorerComponent = ({
                 </Tooltip>
               }
             >
-              <VegaLite
-                data={{ source: chartData }}
+              <VegaEmbed
+                options={chartOptions(theme)}
                 key={idx}
-                actions={false}
                 spec={plot.spec}
-                theme={theme === "dark" ? "dark" : undefined}
               />
             </HorizontalCarouselItem>
           ))}

package/src/plugins/impl/data-explorer/functions/function.ts

@@ -1,8 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 /* eslint-disable @typescript-eslint/no-base-to-string */
 import type { FieldQuery } from "compassql/build/src/query/encoding";
-import { isAggregateOp } from "vega-lite/build/src/aggregate";
-import type { FieldFunction, TimeUnitOp } from "./types";
+import { type FieldFunction, isAggregateOp, type TimeUnitOp } from "./types";
 
 // This code is adapted and simplified from https://github.com/vega/voyager
 

package/src/plugins/impl/data-explorer/functions/types.ts

@@ -1,19 +1,70 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import { isString } from "lodash-es";
+import type {
+  ArgmaxDef,
+  ArgminDef,
+} from "vega-lite/types_unstable/aggregate.d.ts";
+import type { isAggregateOp as isAggregateOpVega } from "vega-lite/types_unstable/aggregate.js";
+import type { AggregateOp } from "vega-typings";
+
+// Vega doesn't expose the constant, so we define all here
+const AGGREGATE_OPS: readonly AggregateOp[] = [
+  "argmax",
+  "argmin",
+  "average",
+  "count",
+  "distinct",
+  "exponential",
+  "exponentialb",
+  "product",
+  "max",
+  "mean",
+  "median",
+  "min",
+  "missing",
+  "q1",
+  "q3",
+  "ci0",
+  "ci1",
+  "stderr",
+  "stdev",
+  "stdevp",
+  "sum",
+  "valid",
+  "values",
+  "variance",
+  "variancep",
+] as const;
+
 // Subset of aggregate operations that we support
-type AggregateOp =
-  | "average"
-  | "count"
-  | "distinct"
-  | "max"
-  | "mean"
-  | "median"
-  | "min"
-  | "q1"
-  | "q3"
-  | "stderr"
-  | "stdev"
-  | "sum";
+const SUPPORTED_AGGREGATE_OPS: readonly AggregateOp[] = [
+  "average",
+  "count",
+  "distinct",
+  "max",
+  "mean",
+  "median",
+  "min",
+  "q1",
+  "q3",
+  "stderr",
+  "stdev",
+  "sum",
+];
+type SupportedAggregateOp = (typeof SUPPORTED_AGGREGATE_OPS)[number];
+
+// We implement our own isAggregateOp because
+// The vega-lite types_unstable import path fails in Vite (module resolution issue)
+export const isAggregateOp: typeof isAggregateOpVega = (
+  a: string | ArgminDef | ArgmaxDef,
+): a is AggregateOp => {
+  if (!isString(a)) {
+    return false;
+  }
+
+  return AGGREGATE_OPS.includes(a as AggregateOp);
+};
 
 // Subset of time units that we support
 export type TimeUnitOp =
@@ -29,4 +80,4 @@ export type TimeUnitOp =
   | "yearmonthdate"
   | "monthdate";
 
-export type FieldFunction = AggregateOp | "bin" | TimeUnitOp;
+export type FieldFunction = SupportedAggregateOp | "bin" | TimeUnitOp;

package/src/plugins/impl/data-explorer/queries/queries.ts

@@ -11,8 +11,7 @@ import type { SpecQuery } from "compassql/build/src/query/spec";
 import { getTopResultTreeItem, isResultTree } from "compassql/build/src/result";
 import type { Schema } from "compassql/build/src/schema";
 import { contains } from "vega-lite";
-import { NONPOSITION_SCALE_CHANNELS } from "vega-lite/build/src/channel";
-import type { NamedData } from "vega-lite/build/src/data";
+import type { NamedData } from "vega-lite/types_unstable/data.js";
 import { type EncodingChannel, fromFieldQuery } from "../encoding";
 import { toSpecQuery } from "../spec";
 import type { ChartSpec } from "../state/types";
@@ -21,13 +20,14 @@ import {
   addQuantitativeField,
   addTemporalField,
 } from "./field-suggestion";
-import type {
-  PlotFieldInfo,
-  QueryCreator,
-  Result,
-  ResultingCharts,
-  ResultPlot,
-  TopLevelFacetedUnitSpec,
+import {
+  NONPOSITION_SCALE_CHANNELS,
+  type PlotFieldInfo,
+  type QueryCreator,
+  type Result,
+  type ResultingCharts,
+  type ResultPlot,
+  type TopLevelFacetedUnitSpec,
 } from "./types";
 import { hasWildcards, isQueryEmpty } from "./utils";