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

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

@@ -4,26 +4,177 @@ import type { Atom } from "jotai";
 import { type Edge, MarkerType, type Node, type NodeProps } from "reactflow";
 import { getNotebook } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import type { CellData } from "@/core/cells/types";
+import type { CellData, CellRuntimeState } from "@/core/cells/types";
 import { store } from "@/core/state/jotai";
-import type { Variables } from "@/core/variables/types";
+import type { Variable, VariableName, Variables } from "@/core/variables/types";
 import { Arrays } from "@/utils/arrays";
+import { extractCellPreview } from "./utils/cell-preview";
+
+/** The edge type registered for dependency edges (see `custom-edge.tsx`). */
+export const DEPENDENCY_EDGE_TYPE = "dependency";
+
+// Stroke/marker colors shared by the dependency edge and its arrow marker.
+export const STATE_FLOW_EDGE_COLOR = "var(--amber-10)";
+export const DEFAULT_EDGE_COLOR = "var(--gray-8)";
+
+export interface DependencyEdgeData {
+  /** Variables carried across this edge. */
+  variables: VariableName[];
+  /** True when any carried variable is a mutable `State`/`SetFunctor` flow. */
+  isStateFlow: boolean;
+}
+
+// A `State`/`SetFunctor` (from `mo.state`) is mutable mid-run, so it is drawn
+// distinctly. `UIElement` is intentionally excluded: it is frontend-driven and
+// immutable during a run.
+const STATE_DATATYPES = new Set(["State", "SetFunctor"]);
+
+export function isStateFlowVariable(variable: Variable): boolean {
+  return variable.dataType != null && STATE_DATATYPES.has(variable.dataType);
+}
 
 export interface NodeData {
   atom: Atom<CellData>;
-  forceWidth?: number;
+  /** Variables this cell defines — the node's primary identity when collapsed. */
+  defs: VariableName[];
+  /** When true the node renders its full code instead of the compact ref list. */
+  expanded?: boolean;
 }
 export type CustomNodeProps = NodeProps<NodeData>;
 
+const LINE_HEIGHT = 11; // matches TinyCode.css
+
 export function getNodeHeight(linesOfCode: number) {
-  const LINE_HEIGHT = 11; // matches TinyCode.css
   return Math.min(linesOfCode * LINE_HEIGHT + 35, 200);
 }
 
+// Compact-node sizing. A collapsed node wraps its defined variables into
+// comma-separated lines at a character budget, so it grows taller rather than
+// wider; the same budget drives the estimated width so dagre and the DOM agree.
+export const NODE_HEADER_HEIGHT = 22;
+const DEF_LINE_HEIGHT = 16;
+const BODY_V_PADDING = 10;
+const EXPANDED_NODE_WIDTH = 300;
+const CHAR_WIDTH = 6.5; // text-xs monospace-ish
+const NODE_H_PADDING = 22;
+const MIN_NODE_WIDTH = 90;
+const MAX_NODE_WIDTH = 240;
+// Wrap defs to the longest single name, but never narrower than this — short
+// names pack several per line, a long one gets a line to itself.
+const WRAP_MIN_CHARS = 12;
+// ... and never wider than fits MAX_NODE_WIDTH, or packed lines would be
+// clipped by the width clamp. A single name over this budget still gets its
+// own line; the renderer truncates it with an ellipsis and a title.
+const WRAP_MAX_CHARS = Math.floor(
+  (MAX_NODE_WIDTH - NODE_H_PADDING) / CHAR_WIDTH,
+);
+
+export function collapsedNodeWidth(label: string): number {
+  const raw = label.length * CHAR_WIDTH + NODE_H_PADDING;
+  return Math.min(Math.max(raw, MIN_NODE_WIDTH), MAX_NODE_WIDTH);
+}
+
+/** The per-line character budget used to wrap a node's defs. */
+function defsWrapBudget(defs: VariableName[]): number {
+  const longest = defs.reduce(
+    (max, def) => Math.max(max, def.length),
+    WRAP_MIN_CHARS,
+  );
+  return Math.min(longest, WRAP_MAX_CHARS);
+}
+
+/**
+ * Greedily pack `defs` into comma-separated lines, each no wider than the
+ * budget (the longest name, or `WRAP_MIN_CHARS`). No trailing commas.
+ */
+export function wrapDefs(defs: VariableName[]): string[] {
+  const budget = defsWrapBudget(defs);
+  const lines: string[] = [];
+  let current = "";
+  for (const def of defs) {
+    if (current === "") {
+      current = def;
+    } else if (current.length + 2 + def.length <= budget) {
+      current = `${current}, ${def}`;
+    } else {
+      lines.push(current);
+      current = def;
+    }
+  }
+  if (current !== "") {
+    lines.push(current);
+  }
+  return lines;
+}
+
+/**
+ * The lines shown in a collapsed node's body: its defs wrapped to a budget, or
+ * a single code-preview line when the cell defines nothing.
+ */
+export function nodeBodyLines(data: NodeData): string[] {
+  if (data.defs.length > 0) {
+    return wrapDefs(data.defs);
+  }
+  const preview = extractCellPreview(store.get(data.atom).code).text;
+  return preview ? [preview] : [];
+}
+
+export function nodeDimensions(data: NodeData): {
+  width: number;
+  height: number;
+} {
+  if (data.expanded) {
+    const lines = store.get(data.atom).code.trim().split("\n").length;
+    return { width: EXPANDED_NODE_WIDTH, height: getNodeHeight(lines) };
+  }
+  const lines = nodeBodyLines(data);
+  const widthChars =
+    data.defs.length > 0
+      ? defsWrapBudget(data.defs)
+      : lines.reduce((max, line) => Math.max(max, line.length), 0);
+  const bodyLines = Math.max(lines.length, 1);
+  return {
+    width: collapsedNodeWidth("x".repeat(widthChars)),
+    height: NODE_HEADER_HEIGHT + bodyLines * DEF_LINE_HEIGHT + BODY_V_PADDING,
+  };
+}
+
 // The nodes must have the same handle IDs to ensure edges connect correctly
 export const OUTPUTS_HANDLE_ID = "outputs";
 export const INPUTS_HANDLE_ID = "inputs";
 
+/**
+ * Cells hidden by the reusable-functions filter: those the kernel marked as
+ * valid top-level definitions (serialization hint `Valid`).
+ */
+export function reusableCellIds(
+  cellIds: CellId[],
+  cellRuntime: Record<CellId, CellRuntimeState>,
+): Set<CellId> {
+  const reusable = new Set<CellId>();
+  for (const cellId of cellIds) {
+    if (cellRuntime[cellId]?.serialization?.toLowerCase() === "valid") {
+      reusable.add(cellId);
+    }
+  }
+  return reusable;
+}
+
+/** Map each cell to the names of the variables it declares. */
+export function computeDefsByCell(
+  variables: Variables,
+): Map<CellId, VariableName[]> {
+  const defsByCell = new Map<CellId, VariableName[]>();
+  for (const variable of Object.values(variables)) {
+    for (const cellId of variable.declaredBy) {
+      const defs = defsByCell.get(cellId) ?? [];
+      defs.push(variable.name);
+      defsByCell.set(cellId, defs);
+    }
+  }
+  return defsByCell;
+}
+
 interface ElementsBuilder {
   createElements: (
     cellIds: CellId[],
@@ -59,16 +210,17 @@ export class VerticalElementsBuilder implements ElementsBuilder {
   private createNode(
     id: string,
     atom: Atom<CellData>,
+    defs: VariableName[],
     prevY: number,
   ): Node<NodeData> {
-    const linesOfCode = store.get(atom).code.trim().split("\n").length;
-    const height = getNodeHeight(linesOfCode);
+    const data: NodeData = { atom, defs };
+    const { width, height } = nodeDimensions(data);
     return {
       id: id,
-      data: { atom },
-      width: 250,
+      data,
+      width,
       type: "custom",
-      height: height,
+      height,
       position: { x: 0, y: prevY + 20 },
     };
   }
@@ -83,8 +235,14 @@ export class VerticalElementsBuilder implements ElementsBuilder {
     let prevY = 0;
     const nodes: Node<NodeData>[] = [];
     const edges: Edge[] = [];
+    const defsByCell = computeDefsByCell(variables);
     for (const [cellId, cellAtom] of Arrays.zip(cellIds, cellAtoms)) {
-      const node = this.createNode(cellId, cellAtom, prevY);
+      const node = this.createNode(
+        cellId,
+        cellAtom,
+        defsByCell.get(cellId) ?? [],
+        prevY,
+      );
       nodes.push(node);
       prevY = node.position.y + (node.height || 0);
     }
@@ -110,16 +268,24 @@ export class VerticalElementsBuilder implements ElementsBuilder {
   }
 }
 
+interface AggregatedEdge {
+  source: CellId;
+  target: CellId;
+  variables: VariableName[];
+  isStateFlow: boolean;
+}
+
 export class TreeElementsBuilder implements ElementsBuilder {
-  private createEdge(source: CellId, target: CellId): Edge {
+  private createEdge(edge: AggregatedEdge): Edge<DependencyEdgeData> {
+    const { source, target, variables, isStateFlow } = edge;
     return {
-      animated: true,
+      type: DEPENDENCY_EDGE_TYPE,
       markerEnd: {
         type: MarkerType.ArrowClosed,
+        color: isStateFlow ? STATE_FLOW_EDGE_COLOR : DEFAULT_EDGE_COLOR,
       },
       id: `${source}-${target}`,
-      // Make thicker
-      style: { strokeWidth: 2 },
+      data: { variables, isStateFlow },
       source: source,
       // Use the same handle ids as the custom node
       sourceHandle: OUTPUTS_HANDLE_ID,
@@ -128,15 +294,19 @@ export class TreeElementsBuilder implements ElementsBuilder {
     };
   }
 
-  private createNode(id: string, atom: Atom<CellData>): Node<NodeData> {
-    const linesOfCode = store.get(atom).code.trim().split("\n").length;
-    const height = getNodeHeight(linesOfCode);
+  private createNode(
+    id: string,
+    atom: Atom<CellData>,
+    defs: VariableName[],
+  ): Node<NodeData> {
+    const data: NodeData = { atom, defs };
+    const { width, height } = nodeDimensions(data);
     return {
       id: id,
-      data: { atom, forceWidth: 300 },
-      width: 300,
+      data,
+      width,
       type: "custom",
-      height: height,
+      height,
       position: { x: 0, y: 0 },
     };
   }
@@ -149,10 +319,19 @@ export class TreeElementsBuilder implements ElementsBuilder {
     hideReusableFunctions: boolean,
   ) {
     const nodes: Node<NodeData>[] = [];
-    const edges: Edge[] = [];
+
+    const cellRuntime = getNotebook().cellRuntime;
+    // Hidden reusable cells are dropped along with their edges — a reusable
+    // function is nearly always referenced somewhere, so gating on "no edges"
+    // would make the filter a no-op.
+    const hiddenReusable = hideReusableFunctions
+      ? reusableCellIds(cellIds, cellRuntime)
+      : new Set<CellId>();
 
     const nodesWithEdges = new Set<CellId>();
-    const visited = new Set<string>();
+    // Dedupe cell→cell edges, aggregating every variable that crosses them so a
+    // single edge can carry multiple refs and flag whether any is a state flow.
+    const edgesByPair = new Map<string, AggregatedEdge>();
 
     for (const variable of Object.values(variables)) {
       // Skip marimo, since likely every cell uses it
@@ -160,40 +339,55 @@ export class TreeElementsBuilder implements ElementsBuilder {
         continue;
       }
 
+      const isState = isStateFlowVariable(variable);
       const { declaredBy, usedBy } = variable;
       for (const fromId of declaredBy) {
         for (const toId of usedBy) {
-          const key = `${fromId}-${toId}`;
-          if (visited.has(key)) {
+          if (hiddenReusable.has(fromId) || hiddenReusable.has(toId)) {
             continue;
           }
-          visited.add(key);
+          const key = `${fromId}-${toId}`;
+          const existing = edgesByPair.get(key);
+          if (existing) {
+            existing.variables.push(variable.name);
+            existing.isStateFlow ||= isState;
+          } else {
+            edgesByPair.set(key, {
+              source: fromId,
+              target: toId,
+              variables: [variable.name],
+              isStateFlow: isState,
+            });
+          }
           nodesWithEdges.add(fromId);
           nodesWithEdges.add(toId);
-          edges.push(this.createEdge(fromId, toId));
         }
       }
     }
 
-    const cellRuntime = getNotebook().cellRuntime;
+    const edges = [...edgesByPair.values()].map((edge) =>
+      this.createEdge(edge),
+    );
+
+    const defsByCell = computeDefsByCell(variables);
 
     for (const [cellId, cellAtom] of Arrays.zip(cellIds, cellAtoms)) {
       const code = store.get(cellAtom).code.trim();
       const hasEdge = nodesWithEdges.has(cellId);
       const isMarkdown = code.startsWith("mo.md");
-      const runtime = cellRuntime[cellId];
-      const isReusable = runtime?.serialization?.toLowerCase() === "valid";
 
       // Apply filters
       if (hidePureMarkdown && isMarkdown && !hasEdge) {
         continue;
       }
-      if (hideReusableFunctions && isReusable && !hasEdge) {
+      if (hiddenReusable.has(cellId)) {
         continue;
       }
 
       // Show every cell that wasn't filtered out
-      nodes.push(this.createNode(cellId, cellAtom));
+      nodes.push(
+        this.createNode(cellId, cellAtom, defsByCell.get(cellId) ?? []),
+      );
     }
 
     return { nodes, edges };

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

@@ -4,6 +4,7 @@ import {
   ArrowRightFromLineIcon,
   ArrowRightIcon,
   ArrowRightToLineIcon,
+  ExternalLinkIcon,
   MoreVerticalIcon,
   NetworkIcon,
   SettingsIcon,
@@ -13,17 +14,18 @@ import {
 } from "lucide-react";
 import React, { memo, useId } from "react";
 import { type Edge, Panel } from "reactflow";
-import { getCellEditorView } from "@/core/cells/cells";
+import { getCellEditorView, useCellActions } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { goToVariableDefinition } from "@/core/codemirror/go-to-definition/commands";
 import type { Variable, Variables } from "@/core/variables/types";
 import { ConnectionCellActionsDropdown } from "../editor/cell/cell-actions";
-import { CellLink } from "../editor/links/cell-link";
+import { CellLink, scrollAndHighlightCell } from "../editor/links/cell-link";
 import { CellLinkList } from "../editor/links/cell-link-list";
 import { Button } from "../ui/button";
 import { Checkbox } from "../ui/checkbox";
 import { Label } from "../ui/label";
 import { Popover, PopoverContent, PopoverTrigger } from "../ui/popover";
+import { Tooltip } from "../ui/tooltip";
 import { VariableName } from "../variables/common";
 import type { GraphSelection, GraphSettings, LayoutDirection } from "./types";
 
@@ -122,6 +124,8 @@ export const GraphSelectionPanel: React.FC<{
   edges: Edge[];
   variables: Variables;
 }> = memo(({ selection, variables, onClearSelection }) => {
+  const { showCellIfHidden } = useCellActions();
+
   if (!selection) {
     return null;
   }
@@ -194,6 +198,23 @@ export const GraphSelectionPanel: React.FC<{
             <SquareFunction className="w-5 h-5" />
             <CellLink cellId={selection.id} />
             <div className="flex-1" />
+            <Tooltip content="Open cell in editor" delayDuration={200}>
+              <Button
+                aria-label="Open cell in editor"
+                variant="text"
+                size="icon"
+                onClick={() => {
+                  // The cell may sit in a collapsed section; reveal it first
+                  // or the scroll target won't be in the DOM.
+                  showCellIfHidden({ cellId: selection.id });
+                  requestAnimationFrame(() => {
+                    scrollAndHighlightCell(selection.id, "focus");
+                  });
+                }}
+              >
+                <ExternalLinkIcon className="w-4 h-4" />
+              </Button>
+            </Tooltip>
             <ConnectionCellActionsDropdown cellId={selection.id}>
               <Button variant="ghost" size="icon">
                 <MoreVerticalIcon className="w-4 h-4" />

package/src/components/dependency-graph/utils/__tests__/layout.test.ts

@@ -0,0 +1,78 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { Node } from "reactflow";
+import { describe, expect, it } from "vitest";
+import type { NodeData } from "../../elements";
+import { boundingBox, placeIslands } from "../layout";
+
+function node(
+  id: string,
+  x: number,
+  y: number,
+  width: number,
+  height: number,
+): Node<NodeData> {
+  return {
+    id,
+    position: { x, y },
+    width,
+    height,
+    data: { atom: {} as NodeData["atom"], defs: [] },
+  };
+}
+
+describe("boundingBox", () => {
+  it("is all-zero for an empty set", () => {
+    expect(boundingBox([])).toEqual({ minX: 0, minY: 0, maxX: 0, maxY: 0 });
+  });
+
+  it("spans position + size across nodes", () => {
+    const box = boundingBox([
+      node("a", 0, 0, 50, 50),
+      node("b", 10, 20, 100, 40),
+    ]);
+    expect(box).toEqual({ minX: 0, minY: 0, maxX: 110, maxY: 60 });
+  });
+
+  it("reads size from style.width/height when width/height are absent", () => {
+    const groupNode: Node = {
+      id: "g",
+      position: { x: 5, y: 5 },
+      data: {},
+      style: { width: 200, height: 100 },
+    };
+    expect(boundingBox([groupNode])).toEqual({
+      minX: 5,
+      minY: 5,
+      maxX: 205,
+      maxY: 105,
+    });
+  });
+});
+
+describe("placeIslands", () => {
+  it("returns nothing when there are no islands", () => {
+    expect(placeIslands([], [node("a", 0, 0, 100, 100)])).toEqual([]);
+  });
+
+  it("stacks islands in a column to the right of the placed graph", () => {
+    const placed = [node("main", 0, 0, 200, 300)];
+    const islands = [node("i0", 0, 0, 90, 40), node("i1", 0, 0, 90, 40)];
+    const result = placeIslands(islands, placed);
+    // startX = maxX(200) + margin(40) = 240; startY = minY(0).
+    expect(result[0].position).toEqual({ x: 240, y: 0 });
+    // Second island stacked below the first (height 40 + gap 20).
+    expect(result[1].position).toEqual({ x: 240, y: 60 });
+  });
+
+  it("wraps into a new column when the column gets too tall", () => {
+    const placed = [node("main", 0, 0, 200, 300)];
+    const islands = Array.from({ length: 10 }, (_, i) =>
+      node(`i${i}`, 0, 0, 90, 40),
+    );
+    const result = placeIslands(islands, placed);
+    const firstColumnX = result[0].position.x;
+    const wrapped = result.find((n) => n.position.x > firstColumnX);
+    expect(wrapped).toBeDefined();
+  });
+});

package/src/components/dependency-graph/utils/layout.ts

@@ -1,45 +1,135 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { graphlib, layout } from "@dagrejs/dagre";
 import type { Edge, Node } from "reactflow";
+import type { NodeData } from "../elements";
 import type { LayoutDirection } from "../types";
 
-const g = new graphlib.Graph().setDefaultEdgeLabel(() => ({}));
+// Disconnected cells are stacked in a compact column to the side of the main
+// graph, like TensorBoard's extracted/auxiliary nodes (extractXOffset 15,
+// extractYOffset 20).
+const ISLAND_GAP = 20;
+const ISLAND_MARGIN = 40;
+const ISLAND_MIN_COLUMN_HEIGHT = 320;
 
+/**
+ * Lay out cell nodes with dagre. Disconnected cells are pulled out of the
+ * dagre flow and stacked in a compact column to the side of the main graph.
+ */
 export const layoutElements = ({
   nodes,
   edges,
   direction,
 }: {
-  nodes: Node[];
+  nodes: Node<NodeData>[];
   edges: Edge[];
   direction: LayoutDirection;
-}) => {
+}): { nodes: Node<NodeData>[]; edges: Edge[] } => {
+  const g = new graphlib.Graph().setDefaultEdgeLabel(() => ({}));
   g.setGraph({
     rankdir: direction,
-    nodesep: 150,
-    ranksep: 200,
-    // So far, longest-path seems to give the best results as trees are
-    // generally quite wide.
-    ranker: "longest-path",
+    // Tight spacing mirroring TensorBoard's graph (nodesep 5, ranksep 25,
+    // edgesep 5), nudged up a little since our nodes are wider ref boxes than
+    // TensorBoard's op ellipses. The dagre default ranker (network-simplex)
+    // gives cleaner layered routing with fewer crossings than longest-path.
+    nodesep: 20,
+    ranksep: 30,
+    edgesep: 10,
   });
 
+  // Disconnected cells are pulled out of the dagre flow and stacked to the
+  // side, so they don't stretch the main graph's layout.
+  const connectedIds = new Set<string>();
+  for (const edge of edges) {
+    connectedIds.add(edge.source);
+    connectedIds.add(edge.target);
+  }
+  const isIsland = (node: Node<NodeData>) => !connectedIds.has(node.id);
+  const flowNodes = nodes.filter((node) => !isIsland(node));
+  const islands = nodes.filter(isIsland);
+
   edges.forEach((edge) => g.setEdge(edge.source, edge.target));
-  nodes.forEach((node) =>
+  flowNodes.forEach((node) => {
     g.setNode(node.id, {
       ...node,
       width: node.width ?? 0,
       height: node.height ?? 0,
-    }),
-  );
+    });
+  });
 
   layout(g);
 
-  return {
-    nodes: nodes.map((node) => {
-      const { x, y } = g.node(node.id);
-
-      return { ...node, position: { x, y } };
-    }),
-    edges,
+  // dagre reports node coords as centers; ReactFlow wants top-left.
+  const topLeft = (id: string) => {
+    const { x, y, width, height } = g.node(id);
+    return { x: x - width / 2, y: y - height / 2 };
   };
+
+  const cellNodes = flowNodes.map((node) => ({
+    ...node,
+    position: topLeft(node.id),
+  }));
+
+  const islandNodes = placeIslands(islands, cellNodes);
+  return { nodes: [...cellNodes, ...islandNodes], edges };
 };
+
+/** Bounding box of already-positioned nodes (top-left position + size). */
+export function boundingBox(nodes: Node[]): {
+  minX: number;
+  minY: number;
+  maxX: number;
+  maxY: number;
+} {
+  if (nodes.length === 0) {
+    return { minX: 0, minY: 0, maxX: 0, maxY: 0 };
+  }
+  let minX = Infinity;
+  let minY = Infinity;
+  let maxX = -Infinity;
+  let maxY = -Infinity;
+  for (const node of nodes) {
+    const styleW = node.style?.width;
+    const styleH = node.style?.height;
+    const w = node.width ?? (typeof styleW === "number" ? styleW : 0);
+    const h = node.height ?? (typeof styleH === "number" ? styleH : 0);
+    minX = Math.min(minX, node.position.x);
+    minY = Math.min(minY, node.position.y);
+    maxX = Math.max(maxX, node.position.x + w);
+    maxY = Math.max(maxY, node.position.y + h);
+  }
+  return { minX, minY, maxX, maxY };
+}
+
+/**
+ * Stack disconnected nodes in a compact column (wrapping into more columns when
+ * tall) just to the right of the main graph.
+ */
+export function placeIslands(
+  islands: Node<NodeData>[],
+  placed: Node[],
+): Node<NodeData>[] {
+  if (islands.length === 0) {
+    return [];
+  }
+  const box = boundingBox(placed);
+  const startX = (placed.length === 0 ? 0 : box.maxX) + ISLAND_MARGIN;
+  const startY = placed.length === 0 ? 0 : box.minY;
+  const columnHeight = Math.max(box.maxY - box.minY, ISLAND_MIN_COLUMN_HEIGHT);
+
+  let x = startX;
+  let y = startY;
+  let columnWidth = 0;
+  return islands.map((node) => {
+    const h = node.height ?? 0;
+    if (y > startY && y + h > startY + columnHeight) {
+      // Wrap to the next column.
+      x += columnWidth + ISLAND_MARGIN;
+      y = startY;
+      columnWidth = 0;
+    }
+    const position = { x, y };
+    y += h + ISLAND_GAP;
+    columnWidth = Math.max(columnWidth, node.width ?? 0);
+    return { ...node, position };
+  });
+}