@marimo-team/islands 0.24.1-dev113 → 0.24.1-dev115

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}

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" />