@marimo-team/frontend 0.23.16-dev9 → 0.23.17-dev0

package/src/core/codemirror/code-lens/__tests__/extension.test.ts

@@ -0,0 +1,325 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { python } from "@codemirror/lang-python";
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { cellId, variableName } from "@/__tests__/branded";
+import { languageAdapterState } from "@/core/codemirror/language/extension";
+import { LanguageAdapters } from "@/core/codemirror/language/LanguageAdapters";
+import type { ExperimentalFeatures } from "@/core/config/feature-flag";
+import { getFeatureFlag } from "@/core/config/feature-flag";
+import {
+  type DataSourceConnection,
+  dataSourceConnectionsAtom,
+} from "@/core/datasets/data-source-connections";
+import { type ConnectionName, DUCKDB_ENGINE } from "@/core/datasets/engines";
+import { datasetsAtom } from "@/core/datasets/state";
+import type { QualifiedColumn } from "@/core/datasets/types";
+import type { DataTable } from "@/core/kernel/messages";
+import { store } from "@/core/state/jotai";
+import { storageAtom } from "@/core/storage/state";
+import type { StorageNamespace } from "@/core/storage/types";
+import { variablesAtom } from "@/core/variables/state";
+import type { Variables } from "@/core/variables/types";
+import { openLensTarget } from "../actions";
+import { CODE_LENS_HOVER_DELAY_MS, codeLensBundle } from "../extension";
+import { mountLensPopover } from "../popover";
+
+vi.mock("@/core/config/feature-flag", () => ({
+  getFeatureFlag: vi.fn(),
+}));
+
+vi.mock("../actions", () => ({
+  openLensTarget: vi.fn(),
+}));
+
+vi.mock("../popover", () => ({
+  mountLensPopover: vi.fn(() => vi.fn()),
+}));
+
+function mockFlags(flags: Partial<ExperimentalFeatures>) {
+  vi.mocked(getFeatureFlag).mockImplementation((flag) => flags[flag] ?? false);
+}
+
+const DF_TABLE: DataTable = {
+  name: "df",
+  source: "memory",
+  source_type: "local",
+  num_rows: 1,
+  num_columns: 1,
+  columns: [],
+  variable_name: variableName("df"),
+};
+
+const BUCKET_NAMESPACE: StorageNamespace = {
+  name: variableName("bucket"),
+  displayName: "bucket",
+  protocol: "s3",
+  rootPath: "s3://bucket",
+  backendType: "obstore",
+  storageEntries: [],
+};
+
+const ENGINE_CONNECTION: DataSourceConnection = {
+  name: "engine" as ConnectionName,
+  source: "postgres",
+  dialect: "postgresql",
+  display_name: "engine (postgres)",
+  databases: [],
+};
+
+function seedStore(opts: {
+  tables?: DataTable[];
+  namespaces?: StorageNamespace[];
+  connections?: DataSourceConnection[];
+  variables?: Variables;
+}) {
+  store.set(datasetsAtom, {
+    tables: opts.tables ?? [],
+    expandedTables: new Set<string>(),
+    expandedColumns: new Set<QualifiedColumn>(),
+    columnsPreviews: new Map(),
+  });
+  store.set(storageAtom, {
+    namespaces: opts.namespaces ?? [],
+    entriesByPath: new Map(),
+    pageMetadataByPath: new Map(),
+  });
+  store.set(dataSourceConnectionsAtom, {
+    latestEngineSelected: DUCKDB_ENGINE,
+    connectionsMap: new Map(
+      (opts.connections ?? []).map((connection) => [
+        connection.name,
+        connection,
+      ]),
+    ),
+  });
+  store.set(variablesAtom, opts.variables ?? {});
+}
+
+describe("codeLensBundle", () => {
+  let view: EditorView | null = null;
+
+  beforeEach(() => {
+    vi.useFakeTimers();
+    mockFlags({ editor_code_lens: true, cache_panel: true });
+    seedStore({});
+  });
+
+  afterEach(() => {
+    if (view) {
+      view.destroy();
+      view = null;
+    }
+    document.body.innerHTML = "";
+    vi.useRealTimers();
+    vi.clearAllMocks();
+  });
+
+  async function mount(
+    code: string,
+    language: "python" | "sql" | "markdown" = "python",
+  ): Promise<EditorView> {
+    const state = EditorState.create({
+      doc: code,
+      extensions: [
+        python(),
+        languageAdapterState.init(() => LanguageAdapters[language]),
+        codeLensBundle(cellId("cell1")),
+      ],
+    });
+    view = new EditorView({ state, parent: document.body });
+    // Fire the debounced analysis, then flush the deferred dispatch
+    vi.advanceTimersByTime(300);
+    await Promise.resolve();
+    view.dispatch({});
+    return view;
+  }
+
+  function lenses(v: EditorView): HTMLElement[] {
+    return [...v.dom.querySelectorAll<HTMLElement>(".mo-code-lens")];
+  }
+
+  it("returns an empty extension when the flag is off", () => {
+    mockFlags({ editor_code_lens: false });
+    expect(codeLensBundle(cellId("cell1"))).toEqual([]);
+  });
+
+  it("renders a lens at a dataframe declaration", async () => {
+    seedStore({ tables: [DF_TABLE] });
+    const v = await mount("df = load()\nx = 1");
+
+    const found = lenses(v);
+    expect(found).toHaveLength(1);
+    expect(found[0].getAttribute("aria-label")).toBe("Open in data sources");
+    // The native tooltip is replaced by the hover popover
+    expect(found[0].getAttribute("title")).toBeNull();
+  });
+
+  it("opens the datasources panel on click", async () => {
+    seedStore({ tables: [DF_TABLE] });
+    const v = await mount("df = load()");
+
+    lenses(v)[0].dispatchEvent(new MouseEvent("click", { bubbles: true }));
+
+    expect(openLensTarget).toHaveBeenCalledWith("table");
+  });
+
+  it("is keyboard focusable and activatable", async () => {
+    seedStore({ tables: [DF_TABLE] });
+    const v = await mount("df = load()");
+    const lens = lenses(v)[0];
+
+    expect(lens.getAttribute("role")).toBe("button");
+    expect(lens.tabIndex).toBe(0);
+
+    lens.dispatchEvent(
+      new KeyboardEvent("keydown", { key: "Enter", bubbles: true }),
+    );
+    expect(openLensTarget).toHaveBeenCalledWith("table");
+
+    lens.dispatchEvent(
+      new KeyboardEvent("keydown", { key: " ", bubbles: true }),
+    );
+    expect(openLensTarget).toHaveBeenCalledTimes(2);
+  });
+
+  it("renders a lens at a bucket declaration", async () => {
+    seedStore({ namespaces: [BUCKET_NAMESPACE] });
+    const v = await mount("bucket = get_bucket()");
+
+    const found = lenses(v);
+    expect(found).toHaveLength(1);
+    expect(found[0].getAttribute("aria-label")).toBe("Open in remote storage");
+
+    found[0].dispatchEvent(new MouseEvent("click", { bubbles: true }));
+    expect(openLensTarget).toHaveBeenCalledWith("bucket");
+  });
+
+  it("renders a lens at a connection declaration", async () => {
+    seedStore({ connections: [ENGINE_CONNECTION] });
+    const v = await mount("engine = create_engine()");
+
+    const found = lenses(v);
+    expect(found).toHaveLength(1);
+    expect(found[0].getAttribute("aria-label")).toBe("Open in data sources");
+
+    found[0].dispatchEvent(new MouseEvent("click", { bubbles: true }));
+    expect(openLensTarget).toHaveBeenCalledWith("connection");
+  });
+
+  it("renders a lens at mo.cache and mo.persistent_cache sites", async () => {
+    const v = await mount(
+      '@mo.cache\ndef f():\n    return 1\n\nwith mo.persistent_cache("k"):\n    pass',
+    );
+
+    const found = lenses(v);
+    expect(found).toHaveLength(2);
+    expect(found[0].getAttribute("aria-label")).toBe("Open cache panel");
+
+    found[0].dispatchEvent(new MouseEvent("click", { bubbles: true }));
+    expect(openLensTarget).toHaveBeenCalledWith("cache");
+  });
+
+  it("shows a popover on hover and hides it on leave", async () => {
+    seedStore({ tables: [DF_TABLE] });
+    const v = await mount("df = load()");
+    const lens = lenses(v)[0];
+
+    lens.dispatchEvent(new MouseEvent("mouseenter"));
+    vi.advanceTimersByTime(CODE_LENS_HOVER_DELAY_MS);
+
+    expect(v.dom.querySelector(".mo-cm-tooltip")).not.toBeNull();
+    expect(mountLensPopover).toHaveBeenCalledWith(
+      expect.any(HTMLElement),
+      expect.objectContaining({ kind: "table", name: "df" }),
+    );
+
+    lens.dispatchEvent(new MouseEvent("mouseleave"));
+    expect(v.dom.querySelector(".mo-cm-tooltip")).toBeNull();
+  });
+
+  it("does not show a popover before the hover delay", async () => {
+    seedStore({ tables: [DF_TABLE] });
+    const v = await mount("df = load()");
+    const lens = lenses(v)[0];
+
+    lens.dispatchEvent(new MouseEvent("mouseenter"));
+    vi.advanceTimersByTime(CODE_LENS_HOVER_DELAY_MS - 100);
+    lens.dispatchEvent(new MouseEvent("mouseleave"));
+    vi.advanceTimersByTime(CODE_LENS_HOVER_DELAY_MS);
+
+    expect(v.dom.querySelector(".mo-cm-tooltip")).toBeNull();
+  });
+
+  it("passes cache context to the popover", async () => {
+    const v = await mount(
+      '@mo.cache\ndef add(a, b):\n    return a + b\n\nwith mo.persistent_cache("k"):\n    pass',
+    );
+    const [decoratorLens, withLens] = lenses(v);
+
+    decoratorLens.dispatchEvent(new MouseEvent("mouseenter"));
+    vi.advanceTimersByTime(CODE_LENS_HOVER_DELAY_MS);
+    expect(mountLensPopover).toHaveBeenCalledWith(
+      expect.any(HTMLElement),
+      expect.objectContaining({
+        kind: "cache",
+        cache: { boundName: "add", cacheName: null },
+      }),
+    );
+
+    decoratorLens.dispatchEvent(new MouseEvent("mouseleave"));
+    withLens.dispatchEvent(new MouseEvent("mouseenter"));
+    vi.advanceTimersByTime(CODE_LENS_HOVER_DELAY_MS);
+    expect(mountLensPopover).toHaveBeenCalledWith(
+      expect.any(HTMLElement),
+      expect.objectContaining({
+        kind: "cache",
+        cache: { boundName: null, cacheName: "k" },
+      }),
+    );
+  });
+
+  it.each(["sql", "markdown"] as const)(
+    "does not render lenses in %s cells",
+    async (language) => {
+      // A `mo.cache` inside a SQL string (or markdown) must not get an icon.
+      seedStore({ tables: [DF_TABLE] });
+      const v = await mount("df = load()\nmo.cache(f)", language);
+
+      expect(lenses(v)).toHaveLength(0);
+    },
+  );
+
+  it("does not render cache lenses when the cache panel is disabled", async () => {
+    mockFlags({ editor_code_lens: true, cache_panel: false });
+    const v = await mount("@mo.cache\ndef f():\n    return 1");
+
+    expect(lenses(v)).toHaveLength(0);
+  });
+
+  it("only decorates the declaring cell", async () => {
+    const dfName = variableName("df");
+    seedStore({
+      tables: [DF_TABLE],
+      variables: {
+        [dfName]: {
+          name: dfName,
+          declaredBy: [cellId("other-cell")],
+          usedBy: [cellId("cell1")],
+        },
+      },
+    });
+    const v = await mount("df = load()");
+
+    expect(lenses(v)).toHaveLength(0);
+  });
+
+  it("decorates only the declaration site, not later references", async () => {
+    seedStore({ tables: [DF_TABLE] });
+    const v = await mount("df = load()\ndf.head()");
+
+    expect(lenses(v)).toHaveLength(1);
+  });
+});

package/src/core/codemirror/code-lens/__tests__/perf.test.ts

@@ -0,0 +1,121 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/**
+ * Regression guards for code-lens hot paths: assert wall time stays within a
+ * generous budget for a large, synthetic notebook.
+ */
+
+import { python } from "@codemirror/lang-python";
+import { EditorState } from "@codemirror/state";
+import { describe, expect, it } from "vitest";
+import { cellId, variableName } from "@/__tests__/branded";
+import type { CellId } from "@/core/cells/ids";
+import { dataSourceConnectionsAtom } from "@/core/datasets/data-source-connections";
+import { DUCKDB_ENGINE } from "@/core/datasets/engines";
+import { datasetsAtom } from "@/core/datasets/state";
+import type { QualifiedColumn } from "@/core/datasets/types";
+import type { DataTable } from "@/core/kernel/messages";
+import { store } from "@/core/state/jotai";
+import { storageAtom } from "@/core/storage/state";
+import { variablesAtom } from "@/core/variables/state";
+import type { Variables } from "@/core/variables/types";
+import { findCacheSites, findDeclarationSites } from "../analyzer";
+import { getLensEntities } from "../entities";
+
+function pythonState(code: string): EditorState {
+  return EditorState.create({ doc: code, extensions: [python()] });
+}
+
+describe("code-lens performance", () => {
+  it("getLensEntities stays fast across many cells in a large notebook", () => {
+    const DATASOURCE_COUNT = 2000;
+    const CELL_COUNT = 1000;
+
+    const tables: DataTable[] = [];
+    const variables: Variables = {};
+    for (let i = 0; i < DATASOURCE_COUNT; i++) {
+      const name = variableName(`table_${i}`);
+      tables.push({
+        name,
+        source: "memory",
+        source_type: "local",
+        num_rows: 1,
+        num_columns: 1,
+        columns: [],
+        variable_name: name,
+      });
+      // Every datasource is owned by a distinct cell, like a real notebook.
+      variables[name] = {
+        name,
+        declaredBy: [cellId(`owner-${i}`)],
+        usedBy: [],
+      };
+    }
+    store.set(datasetsAtom, {
+      tables,
+      expandedTables: new Set<string>(),
+      expandedColumns: new Set<QualifiedColumn>(),
+      columnsPreviews: new Map(),
+    });
+    store.set(storageAtom, {
+      namespaces: [],
+      entriesByPath: new Map(),
+      pageMetadataByPath: new Map(),
+    });
+    store.set(dataSourceConnectionsAtom, {
+      latestEngineSelected: DUCKDB_ENGINE,
+      connectionsMap: new Map(),
+    });
+    store.set(variablesAtom, variables);
+
+    const cellIds: CellId[] = Array.from({ length: CELL_COUNT }, (_, i) =>
+      cellId(`cell-${i}`),
+    );
+
+    const start = performance.now();
+    for (const id of cellIds) {
+      getLensEntities(id);
+    }
+    const elapsed = performance.now() - start;
+
+    // Rescanning all 2000 datasources per cell (instead of an O(1) lookup
+    // into a precomputed per-cell index) would take on the order of seconds
+    // here; 200ms leaves generous headroom for slow CI while still catching
+    // that kind of regression.
+    expect(elapsed).toBeLessThan(200);
+  });
+
+  it("findCacheSites stays fast for a large cell, valid or not", () => {
+    const lines: string[] = [];
+    for (let i = 0; i < 500; i++) {
+      lines.push(`@mo.cache`, `def fn_${i}(x):`, `    return x + ${i}`, "");
+    }
+    const validCode = lines.join("\n");
+    // Trailing unclosed paren makes the whole tree unparsable, mimicking a
+    // cell mid-edit — this must bail out quickly rather than scanning anyway.
+    const brokenCode = `${validCode}\ndf = (`;
+
+    const start = performance.now();
+    findCacheSites(pythonState(validCode));
+    findCacheSites(pythonState(brokenCode));
+    const elapsed = performance.now() - start;
+
+    expect(elapsed).toBeLessThan(200);
+  });
+
+  it("findDeclarationSites stays fast for a large cell", () => {
+    const lines: string[] = [];
+    const names = new Set<string>();
+    for (let i = 0; i < 2000; i++) {
+      lines.push(`var_${i} = compute_${i}()`);
+      names.add(`var_${i}`);
+    }
+    const state = pythonState(lines.join("\n"));
+
+    const start = performance.now();
+    findDeclarationSites({ state, names });
+    const elapsed = performance.now() - start;
+
+    expect(elapsed).toBeLessThan(200);
+  });
+});

package/src/core/codemirror/code-lens/__tests__/popover.test.ts

@@ -0,0 +1,93 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { CodeLensSpec } from "../entities";
+
+// Stable mock reference so every (re)import of popover.tsx after
+// vi.resetModules() sees the same spy — the module-level fetch-throttle
+// state in popover.tsx (`lastCacheInfoFetchAt`) needs to start fresh per
+// test, which requires resetting the whole module, not just the mock.
+const { getCacheInfoMock } = vi.hoisted(() => ({
+  getCacheInfoMock: vi.fn().mockResolvedValue(null),
+}));
+
+vi.mock("@/core/network/requests", () => ({
+  getRequestClient: () => ({ getCacheInfo: getCacheInfoMock }),
+}));
+
+const CACHE_SPEC: CodeLensSpec = {
+  pos: 0,
+  kind: "cache",
+  name: "my_cache",
+  // `boundName: null` means the popover has no per-cache stats to fall back
+  // on, so it always attempts the notebook-wide `getCacheInfo` fetch.
+  cache: { boundName: null, cacheName: "my_cache" },
+};
+
+describe("CachePopover getCacheInfo throttling", () => {
+  let mountLensPopover: typeof import("../popover").mountLensPopover;
+  let now: number;
+
+  beforeEach(async () => {
+    vi.resetModules();
+    getCacheInfoMock.mockClear();
+    getCacheInfoMock.mockResolvedValue(null);
+    ({ mountLensPopover } = await import("../popover"));
+    // Control `Date.now()` directly (rather than `vi.useFakeTimers()`)
+    // so React's own scheduler — which effect-flushing inside `act()`
+    // depends on — keeps using real timers.
+    now = 0;
+    vi.spyOn(Date, "now").mockImplementation(() => now);
+  });
+
+  afterEach(() => {
+    document.body.innerHTML = "";
+    vi.restoreAllMocks();
+  });
+
+  function mount() {
+    const dom = document.createElement("div");
+    document.body.append(dom);
+    let dispose: () => void;
+    act(() => {
+      dispose = mountLensPopover(dom, CACHE_SPEC);
+    });
+    return () => act(() => dispose());
+  }
+
+  it("does not refetch on every hover, only after the throttle interval elapses", () => {
+    const disposeFirst = mount();
+    expect(getCacheInfoMock).toHaveBeenCalledTimes(1);
+    disposeFirst();
+
+    // Hovering again shortly after should reuse whatever is already in
+    // flight/cached rather than firing a second request.
+    const disposeSecond = mount();
+    expect(getCacheInfoMock).toHaveBeenCalledTimes(1);
+    disposeSecond();
+
+    now += 5001;
+
+    const disposeThird = mount();
+    expect(getCacheInfoMock).toHaveBeenCalledTimes(2);
+    disposeThird();
+  });
+
+  it("allows an immediate retry after a failed fetch instead of waiting out the interval", async () => {
+    getCacheInfoMock.mockRejectedValueOnce(new Error("kernel not ready"));
+
+    const disposeFirst = mount();
+    expect(getCacheInfoMock).toHaveBeenCalledTimes(1);
+    // Flush the microtask queue so the effect's `.catch()` runs and resets
+    // the throttle before the next hover.
+    await act(async () => {
+      await Promise.resolve();
+    });
+    disposeFirst();
+
+    const disposeSecond = mount();
+    expect(getCacheInfoMock).toHaveBeenCalledTimes(2);
+    disposeSecond();
+  });
+});

package/src/core/codemirror/code-lens/actions.ts

@@ -0,0 +1,33 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  expandAccordionSection,
+  fileExplorerPanelAtom,
+  sessionPanelAtom,
+} from "@/components/editor/chrome/panels/panel-accordion-state";
+import { openPanel } from "@/components/editor/chrome/state";
+import { assertNever } from "@/utils/assertNever";
+import type { CodeLensKind } from "./entities";
+
+/**
+ * Opens the panel associated with a code lens target, expanding the accordion
+ * section that holds it.
+ */
+export function openLensTarget(kind: CodeLensKind): void {
+  switch (kind) {
+    case "table":
+    case "connection":
+      openPanel("variables");
+      expandAccordionSection(sessionPanelAtom, "datasources");
+      return;
+    case "bucket":
+      openPanel("files");
+      expandAccordionSection(fileExplorerPanelAtom, "remote-storage");
+      return;
+    case "cache":
+      openPanel("cache");
+      return;
+    default:
+      assertNever(kind);
+  }
+}