@marimo-team/frontend 0.24.1-dev99 → 0.24.2-dev129

package/src/components/data-table/renderers.tsx

@@ -71,11 +71,7 @@ export function renderTableHeader<TData>(
         {renderHeaderGroup(table.getCenterHeaderGroups())}
         {renderHeaderGroup(table.getRightHeaderGroups())}
         {table.getAllColumns().length <= AUTO_WIDTH_MAX_COLUMNS && (
-          <th
-            className="w-full border-0"
-            aria-hidden="true"
-            role="presentation"
-          />
+          <th className="w-full border-0" aria-hidden="true" />
         )}
       </TableRow>
     </TableHeader>
@@ -217,7 +213,7 @@ export const DataTableBody = <TData,>({
         {renderCells(row.getCenterVisibleCells())}
         {renderCells(row.getRightVisibleCells())}
         {columns.length <= AUTO_WIDTH_MAX_COLUMNS && (
-          <td className="border-0" aria-hidden="true" role="presentation" />
+          <td className="border-0" aria-hidden="true" />
         )}
       </TableRow>
     );
@@ -247,7 +243,7 @@ export const DataTableBody = <TData,>({
               data-virtual-spacer=""
               style={{ height: virtualItems[0].start }}
             >
-              <td colSpan={totalColSpan} />
+              <td aria-hidden="true" colSpan={totalColSpan} />
             </tr>
           )}
           {virtualItems.map((vItem) => renderRow(rows[vItem.index]))}
@@ -258,7 +254,7 @@ export const DataTableBody = <TData,>({
                 height: totalSize - (virtualItems.at(-1)?.end ?? totalSize),
               }}
             >
-              <td colSpan={totalColSpan} />
+              <td aria-hidden="true" colSpan={totalColSpan} />
             </tr>
           )}
         </>

package/src/components/data-table/show-only-column-button.tsx

@@ -0,0 +1,58 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+"use no memo";
+
+import type { Table } from "@tanstack/react-table";
+import { ScanEyeIcon } from "lucide-react";
+import { Button } from "@/components/ui/button";
+import { Tooltip } from "@/components/ui/tooltip";
+import { cn } from "@/utils/cn";
+import { applyShowOnlyColumns } from "./hooks/use-column-visibility";
+
+export const SHOW_ONLY_COLUMN_LABEL = "Show only this column";
+
+interface ShowOnlyColumnButtonProps<TData> {
+  table: Table<TData>;
+  columnIds: string[];
+  disabled: boolean;
+  className?: string;
+  iconClassName?: string;
+}
+
+export function ShowOnlyColumnButton<TData>({
+  table,
+  columnIds,
+  disabled,
+  className,
+  iconClassName = "h-3 w-3",
+}: ShowOnlyColumnButtonProps<TData>) {
+  return (
+    <Tooltip content={SHOW_ONLY_COLUMN_LABEL} delayDuration={400}>
+      <Button
+        type="button"
+        variant="text"
+        size="icon"
+        aria-label={SHOW_ONLY_COLUMN_LABEL}
+        aria-disabled={disabled}
+        tabIndex={disabled ? -1 : 0}
+        className={cn(
+          "h-6 w-6 hover:bg-muted text-muted-foreground hover:text-primary",
+          disabled && "opacity-50",
+          className,
+        )}
+        onPointerDown={(event) => {
+          event.stopPropagation();
+        }}
+        onClick={(event) => {
+          event.preventDefault();
+          event.stopPropagation();
+          if (disabled) {
+            return;
+          }
+          applyShowOnlyColumns(table, columnIds);
+        }}
+      >
+        <ScanEyeIcon className={iconClassName} strokeWidth={2.5} />
+      </Button>
+    </Tooltip>
+  );
+}

package/src/components/dependency-graph/__tests__/elements.test.ts

@@ -0,0 +1,151 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import type { CellId } from "@/core/cells/ids";
+import type { Variable, VariableName, Variables } from "@/core/variables/types";
+import { createCellRuntimeState } from "@/core/cells/types";
+import {
+  collapsedNodeWidth,
+  computeDefsByCell,
+  isStateFlowVariable,
+  reusableCellIds,
+  wrapDefs,
+} from "../elements";
+
+const names = (list: string[]) => list as VariableName[];
+
+function variable(
+  partial: Omit<Partial<Variable>, "name"> & { name: string },
+): Variable {
+  return {
+    declaredBy: [],
+    usedBy: [],
+    ...partial,
+    name: partial.name as VariableName,
+  };
+}
+
+describe("collapsedNodeWidth", () => {
+  it("clamps short labels to the minimum width", () => {
+    expect(collapsedNodeWidth("")).toBe(90);
+    expect(collapsedNodeWidth("x")).toBe(90);
+  });
+
+  it("clamps long labels to the maximum width", () => {
+    expect(collapsedNodeWidth("a".repeat(100))).toBe(240);
+  });
+
+  it("scales with label length between the bounds", () => {
+    const width = collapsedNodeWidth("df, summary");
+    expect(width).toBeGreaterThan(90);
+    expect(width).toBeLessThan(240);
+  });
+});
+
+describe("isStateFlowVariable", () => {
+  it("flags State and SetFunctor (from mo.state)", () => {
+    expect(
+      isStateFlowVariable(variable({ name: "g", dataType: "State" })),
+    ).toBe(true);
+    expect(
+      isStateFlowVariable(variable({ name: "s", dataType: "SetFunctor" })),
+    ).toBe(true);
+  });
+
+  it("does not flag UI elements or plain values", () => {
+    // UIElement is frontend-driven / immutable during a run — not a state flow.
+    expect(
+      isStateFlowVariable(variable({ name: "sl", dataType: "slider" })),
+    ).toBe(false);
+    expect(isStateFlowVariable(variable({ name: "n", dataType: "int" }))).toBe(
+      false,
+    );
+    expect(isStateFlowVariable(variable({ name: "u", dataType: null }))).toBe(
+      false,
+    );
+    expect(isStateFlowVariable(variable({ name: "u2" }))).toBe(false);
+  });
+});
+
+describe("wrapDefs", () => {
+  it("packs short names to the 12-char minimum budget", () => {
+    expect(
+      wrapDefs(names(["np", "math", "time", "io", "numpy", "pandas"])),
+    ).toEqual(["np, math", "time, io", "numpy", "pandas"]);
+  });
+
+  it("gives a long name its own line, then packs the rest to its width", () => {
+    expect(
+      wrapDefs(
+        names([
+          "very_very_very_long_variable_name",
+          "short_var",
+          "shr_v",
+          "io",
+          "numpy",
+          "wev",
+        ]),
+      ),
+    ).toEqual([
+      "very_very_very_long_variable_name",
+      "short_var, shr_v, io, numpy, wev",
+    ]);
+  });
+
+  it("returns an empty list for no defs", () => {
+    expect(wrapDefs(names([]))).toEqual([]);
+  });
+
+  it("caps the budget so packed lines fit the max node width", () => {
+    // One name far beyond the width cap must not widen every other line.
+    const huge = "x".repeat(60);
+    const lines = wrapDefs(
+      names([huge, "aaaaaaaaaa", "bbbbbbbbbb", "cccccccccc"]),
+    );
+    expect(lines[0]).toBe(huge);
+    for (const line of lines.slice(1)) {
+      expect(line.length).toBeLessThanOrEqual(33);
+    }
+  });
+});
+
+describe("computeDefsByCell", () => {
+  it("maps each cell to the variables it declares", () => {
+    const c1 = "c1" as CellId;
+    const c2 = "c2" as CellId;
+    const a = "a" as VariableName;
+    const b = "b" as VariableName;
+    const c = "c" as VariableName;
+    const variables: Variables = {
+      [a]: variable({ name: a, declaredBy: [c1], usedBy: [c2] }),
+      [b]: variable({ name: b, declaredBy: [c1] }),
+      [c]: variable({ name: c, declaredBy: [c2] }),
+    };
+    const defs = computeDefsByCell(variables);
+    expect(defs.get(c1)).toEqual([a, b]);
+    expect(defs.get(c2)).toEqual([c]);
+  });
+
+  it("returns an empty map when there are no variables", () => {
+    expect(computeDefsByCell({}).size).toBe(0);
+  });
+});
+
+describe("reusableCellIds", () => {
+  const c1 = "c1" as CellId;
+  const c2 = "c2" as CellId;
+  const c3 = "c3" as CellId;
+
+  it("collects cells whose serialization hint is valid, case-insensitively", () => {
+    const runtime = {
+      [c1]: createCellRuntimeState({ serialization: "Valid" }),
+      [c2]: createCellRuntimeState({ serialization: "Unparsable" }),
+      [c3]: createCellRuntimeState(),
+    };
+    expect(reusableCellIds([c1, c2, c3], runtime)).toEqual(new Set([c1]));
+  });
+
+  it("ignores cells with no runtime entry", () => {
+    expect(reusableCellIds([c1], {})).toEqual(new Set());
+  });
+});

package/src/components/dependency-graph/custom-edge.tsx

@@ -0,0 +1,68 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { memo } from "react";
+import { BaseEdge, type EdgeProps, getBezierPath } from "reactflow";
+import {
+  DEFAULT_EDGE_COLOR,
+  type DependencyEdgeData,
+  STATE_FLOW_EDGE_COLOR,
+} from "./elements";
+
+// Gentle curvature keeps edges reading as clean directional splines
+// (TensorBoard-style) rather than hard right-angle elbows.
+const EDGE_CURVATURE = 0.35;
+
+/**
+ * A dependency edge, drawn as a smooth curved spline that flows out of the
+ * source handle and into the target — the TensorBoard graph aesthetic. Mutable
+ * `State`/setter flows get a distinct dashed accent style.
+ */
+export const DependencyEdge = memo((props: EdgeProps<DependencyEdgeData>) => {
+  const {
+    sourceX,
+    sourceY,
+    targetX,
+    targetY,
+    sourcePosition,
+    targetPosition,
+    markerEnd,
+    data,
+    selected,
+  } = props;
+
+  const [path] = getBezierPath({
+    sourceX,
+    sourceY,
+    targetX,
+    targetY,
+    sourcePosition,
+    targetPosition,
+    curvature: EDGE_CURVATURE,
+  });
+
+  const isStateFlow = data?.isStateFlow ?? false;
+  // The inline stroke replaces React Flow's stylesheet-driven edge colors, so
+  // the selected state must be restyled here too or it becomes invisible.
+  const stroke = selected
+    ? "var(--primary)"
+    : isStateFlow
+      ? STATE_FLOW_EDGE_COLOR
+      : DEFAULT_EDGE_COLOR;
+
+  return (
+    <BaseEdge
+      path={path}
+      markerEnd={markerEnd}
+      style={{
+        strokeWidth: selected ? 2.5 : isStateFlow ? 2 : 1.5,
+        stroke,
+        strokeDasharray: isStateFlow ? "5 3" : undefined,
+      }}
+    />
+  );
+});
+DependencyEdge.displayName = "DependencyEdge";
+
+export const edgeTypes = {
+  dependency: DependencyEdge,
+};

package/src/components/dependency-graph/custom-node.tsx

@@ -2,25 +2,20 @@
 
 import { useAtomValue } from "jotai";
 import React, { memo, use } from "react";
-import { Handle, Position, useStore } from "reactflow";
+import { Handle, Position } from "reactflow";
 import { TinyCode } from "@/components/editor/cell/TinyCode";
 import { useCellIds } from "@/core/cells/cells";
-import { displayCellName } from "@/core/cells/names";
+import { displayCellName, isInternalCellName } from "@/core/cells/names";
 import { cn } from "@/utils/cn";
 import {
   type CustomNodeProps,
-  getNodeHeight,
   INPUTS_HANDLE_ID,
+  nodeBodyLines,
+  nodeDimensions,
   OUTPUTS_HANDLE_ID,
 } from "./elements";
 import type { LayoutDirection } from "./types";
-
-function getWidth(canvasWidth: number) {
-  const minWidth = 100;
-  const maxWidth = 400;
-  const padding = 50;
-  return Math.min(Math.max(canvasWidth - padding * 2, minWidth), maxWidth);
-}
+import { extractCellPreview } from "./utils/cell-preview";
 
 export const EdgeMarkerContext = React.createContext<LayoutDirection>("LR");
 
@@ -36,13 +31,19 @@ export const CustomNode = memo((props: CustomNodeProps) => {
   const { data, selected } = props; // must match the equality check
   const cell = useAtomValue(data.atom);
   const cellIndex = useCellIds().inOrderIds.indexOf(cell.id);
-  const nonSelectedColor = "var(--gray-3)";
-  const selectedColor = "var(--gray-9)";
-  const color = selected ? selectedColor : nonSelectedColor;
-  const reactFlowWidth = useStore(({ width }) => width);
+  const handleColor = selected ? "var(--gray-9)" : "var(--gray-3)";
   const edgeMarkers = use(EdgeMarkerContext);
 
-  const linesOfCode = cell.code.split("\n").length;
+  const { width, height } = nodeDimensions(data);
+  const isNamed = !isInternalCellName(cell.name);
+  const name = displayCellName(cell.name, cellIndex);
+  // Only collapsed nodes without defs render the preview; skip the parse
+  // otherwise, since it scales with notebook size.
+  const preview =
+    !data.expanded && data.defs.length === 0
+      ? extractCellPreview(cell.code).text
+      : undefined;
+
   return (
     <div>
       <Handle
@@ -50,43 +51,70 @@ export const CustomNode = memo((props: CustomNodeProps) => {
         id={INPUTS_HANDLE_ID}
         data-testid="input-one"
         position={edgeMarkers === "LR" ? Position.Left : Position.Top}
-        style={{ background: color }}
+        style={{ background: handleColor }}
       />
       <Handle
         type="source"
         id={INPUTS_HANDLE_ID}
         data-testid="input-two"
         position={edgeMarkers === "LR" ? Position.Left : Position.Top}
-        style={{ background: color }}
+        style={{ background: handleColor }}
       />
       <div
         className={cn(
           "flex flex-col bg-card border border-input/50 rounded-md mx-[2px] overflow-hidden",
           selected && "border-primary",
         )}
-        style={{
-          height: getNodeHeight(linesOfCode),
-          width: data.forceWidth || getWidth(reactFlowWidth),
-        }}
+        style={{ height, width }}
       >
-        <div className="text-muted-foreground font-semibold text-xs py-1 px-2 bg-muted border-b">
-          {displayCellName(cell.name, cellIndex)}
+        <div className="text-xs py-0.5 px-2 bg-muted border-b flex items-center justify-between gap-2 shrink-0">
+          <span
+            className={cn(
+              "truncate",
+              isNamed
+                ? "font-semibold text-foreground"
+                : "text-muted-foreground/70",
+            )}
+          >
+            {name}
+          </span>
         </div>
-        <TinyCode code={cell.code} />
+        {data.expanded ? (
+          <TinyCode code={cell.code} />
+        ) : (
+          <div className="flex-1 min-h-0 flex flex-col justify-center px-2 py-1 font-mono text-xs leading-4">
+            {data.defs.length > 0 ? (
+              // Defs wrapped into comma-separated lines — taller, not wider.
+              nodeBodyLines(data).map((line) => (
+                <span
+                  key={line}
+                  className="text-foreground truncate"
+                  title={line}
+                >
+                  {line}
+                </span>
+              ))
+            ) : (
+              <span className="truncate italic text-muted-foreground/60">
+                {preview ?? "empty"}
+              </span>
+            )}
+          </div>
+        )}
       </div>
       <Handle
         type="source"
         id={OUTPUTS_HANDLE_ID}
         data-testid="output-one"
         position={edgeMarkers === "LR" ? Position.Right : Position.Bottom}
-        style={{ background: color }}
+        style={{ background: handleColor }}
       />
       <Handle
         type="target"
         id={OUTPUTS_HANDLE_ID}
         data-testid="output-two"
         position={edgeMarkers === "LR" ? Position.Right : Position.Bottom}
-        style={{ background: color }}
+        style={{ background: handleColor }}
       />
     </div>
   );

package/src/components/dependency-graph/dependency-graph-tree.tsx

@@ -1,8 +1,13 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { Atom } from "jotai";
+import { type Atom, atom, useAtomValue } from "jotai";
 import { MapPinIcon } from "lucide-react";
-import React, { type PropsWithChildren, useEffect, useState } from "react";
+import React, {
+  type PropsWithChildren,
+  useEffect,
+  useMemo,
+  useState,
+} from "react";
 import useEvent from "react-use-event-hook";
 import ReactFlow, {
   Background,
@@ -15,6 +20,7 @@ import ReactFlow, {
   useNodesState,
   useReactFlow,
 } from "reactflow";
+import { edgeTypes } from "@/components/dependency-graph/custom-edge";
 import {
   EdgeMarkerContext,
   nodeTypes,
@@ -25,11 +31,16 @@ import type { CellData } from "@/core/cells/types";
 import { store } from "@/core/state/jotai";
 import type { Variables } from "@/core/variables/types";
 import { Events } from "@/utils/events";
-import { scrollAndHighlightCell } from "../editor/links/cell-link";
 import { Tooltip } from "../ui/tooltip";
-import { type NodeData, TreeElementsBuilder } from "./elements";
+import {
+  computeDefsByCell,
+  type NodeData,
+  nodeDimensions,
+  TreeElementsBuilder,
+} from "./elements";
 import { GraphSelectionPanel } from "./panels";
 import type { GraphSelection, GraphSettings, LayoutDirection } from "./types";
+import { extractCellPreview } from "./utils/cell-preview";
 import { layoutElements } from "./utils/layout";
 import { useFitToViewOnDimensionChange } from "./utils/useFitToViewOnDimensionChange";
 
@@ -43,6 +54,24 @@ interface Props {
 
 const elementsBuilder = new TreeElementsBuilder();
 
+/**
+ * Apply the current expand/collapse selection to freshly-built nodes, resizing
+ * each so the layout engine and the DOM agree on node dimensions.
+ */
+function withExpansion(
+  nodes: Node<NodeData>[],
+  expandedIds: Set<CellId>,
+): Node<NodeData>[] {
+  return nodes.map((node) => {
+    const data: NodeData = {
+      ...node.data,
+      expanded: expandedIds.has(node.id as CellId),
+    };
+    const { width, height } = nodeDimensions(data);
+    return { ...node, data, width, height };
+  });
+}
+
 export const DependencyGraphTree: React.FC<PropsWithChildren<Props>> = ({
   cellIds,
   variables,
@@ -51,22 +80,23 @@ export const DependencyGraphTree: React.FC<PropsWithChildren<Props>> = ({
   layoutDirection,
   settings,
 }) => {
+  // Cells whose node is expanded to show its full code (toggled by double-click).
+  const [expandedIds, setExpandedIds] = useState<Set<CellId>>(() => new Set());
+
   // oxlint-disable-next-line react/hook-use-state
   const [initial] = useState(() => {
-    let elements = elementsBuilder.createElements(
+    const elements = elementsBuilder.createElements(
       cellIds,
       cellAtoms,
       variables,
       settings.hidePureMarkdown,
       settings.hideReusableFunctions,
     );
-    elements = layoutElements({
-      nodes: elements.nodes,
+    return layoutElements({
+      nodes: withExpansion(elements.nodes, expandedIds),
       edges: elements.edges,
       direction: layoutDirection,
     });
-
-    return elements;
     // Only run once
   });
 
@@ -82,22 +112,68 @@ export const DependencyGraphTree: React.FC<PropsWithChildren<Props>> = ({
         edges: elements.edges,
         direction: layoutDirection,
       });
-      setNodes(result.nodes);
-      setEdges(result.edges);
+      // Rebuilt elements are fresh objects; carry selection over so a
+      // re-layout (e.g. toggling expansion) doesn't silently clear it while
+      // the selection panel still targets the element.
+      setNodes((prev) => {
+        const selected = new Set(
+          prev.filter((node) => node.selected).map((node) => node.id),
+        );
+        return result.nodes.map((node) =>
+          selected.has(node.id) ? { ...node, selected: true } : node,
+        );
+      });
+      setEdges((prev) => {
+        const selected = new Set(
+          prev.filter((edge) => edge.selected).map((edge) => edge.id),
+        );
+        return result.edges.map((edge) =>
+          selected.has(edge.id) ? { ...edge, selected: true } : edge,
+        );
+      });
     },
   );
 
-  // If the cellIds change, update the nodes.
+  // Node sizes depend on live cell code: the line count when expanded, or the
+  // preview when a cell defines nothing. The cell atoms keep their identity
+  // across edits, so the rebuild effect below can't see code changes through
+  // its deps — subscribe to the size-relevant inputs explicitly.
+  const defsByCell = useMemo(() => computeDefsByCell(variables), [variables]);
+  const sizeSignature = useAtomValue(
+    useMemo(
+      () =>
+        atom((get) =>
+          cellIds
+            .map((cellId, index) => {
+              const code = get(cellAtoms[index]).code;
+              if (expandedIds.has(cellId)) {
+                return `e${code.trim().split("\n").length}`;
+              }
+              if ((defsByCell.get(cellId) ?? []).length > 0) {
+                // Collapsed nodes with defs size off the defs, not the code.
+                return "d";
+              }
+              return `p${extractCellPreview(code).text?.length ?? 0}`;
+            })
+            .join(),
+        ),
+      [cellIds, cellAtoms, expandedIds, defsByCell],
+    ),
+  );
+
+  // Rebuild + re-layout when the graph inputs or the expand/collapse set change.
   useEffect(() => {
-    syncChanges(
-      elementsBuilder.createElements(
-        cellIds,
-        cellAtoms,
-        variables,
-        settings.hidePureMarkdown,
-        settings.hideReusableFunctions,
-      ),
+    const elements = elementsBuilder.createElements(
+      cellIds,
+      cellAtoms,
+      variables,
+      settings.hidePureMarkdown,
+      settings.hideReusableFunctions,
     );
+    syncChanges({
+      nodes: withExpansion(elements.nodes, expandedIds),
+      edges: elements.edges,
+    });
   }, [
     cellIds,
     variables,
@@ -105,6 +181,8 @@ export const DependencyGraphTree: React.FC<PropsWithChildren<Props>> = ({
     syncChanges,
     settings.hidePureMarkdown,
     settings.hideReusableFunctions,
+    expandedIds,
+    sizeSignature,
   ]);
 
   const [selection, setSelection] = useState<GraphSelection>();
@@ -120,6 +198,7 @@ export const DependencyGraphTree: React.FC<PropsWithChildren<Props>> = ({
         nodes={nodes}
         edges={edges}
         nodeTypes={nodeTypes}
+        edgeTypes={edgeTypes}
         minZoom={0.2}
         fitViewOptions={{
           minZoom: 0.5,
@@ -137,7 +216,17 @@ export const DependencyGraphTree: React.FC<PropsWithChildren<Props>> = ({
           });
         }}
         onNodeDoubleClick={(_event, node) => {
-          scrollAndHighlightCell(node.id as CellId, "focus");
+          // Expand/collapse the node to reveal its full code in place.
+          const id = node.id as CellId;
+          setExpandedIds((prev) => {
+            const next = new Set(prev);
+            if (next.has(id)) {
+              next.delete(id);
+            } else {
+              next.add(id);
+            }
+            return next;
+          });
         }}
         fitView={true}
         onNodesChange={onNodesChange}