fx_viewer: label budget, bbox shrink, info-panel dedup + LRU

Three interrelated fixes for per-layer-accuracy usability in Observatory:

**Build-time label cap (exporter.py)**
- Add `_MAX_LABEL_EXTENSIONS = 2` (cap on label rows reserved in the layout
  slot) and `_layout_constants_payload()` which emits line_height, y_padding,
  font sizes, and `max_label_extensions` into `payload.layout` so the JS
  runtime reads the same values Python used for layout.
- `_select_top_label_extensions()`: for each node, picks the 2 extensions
  with the largest label text to reserve bbox space, keeping the reserved slot
  bounded to 2 extra rows regardless of how many label-bearing layers exist.
- `_compute_layout` and `relayout_payload_base` both use the new selector.
- `GraphExtensionPayload.has_label_formatter` flag (models.py + extension.py)
  lets JS identify which extensions are label-bearing without inspecting nodes.
- `GraphPayload.layout` field carries the constants block.

**JS LRU queue (graph_data_store.js + view_controller.js + ui_manager.js)**
- `GraphDataStore` gets `labelLru` (Set, insertion-order LRU, cap from
  `layoutConstants.max_label_extensions`) and `recordLabelActivation /
  recordLabelDeactivation / extensionHasLabels` methods.
- `upsertExtension` preserves `has_label_formatter` and `sync_keys`.
- `computeActiveGraph` gates `label_append` on LRU membership; non-LRU
  labeled extensions still contribute color/info/tooltip.
- `ViewerController.setState` diffs prev/next active set, calling
  `recordLabelActivation` for newly-activated label-bearing exts (LRU evicts
  oldest on overflow) and `recordLabelDeactivation` for newly-removed.
- Seeded at construction time so default layers render labels on first paint.
- Layer panel shows a small "L" badge on the ≤2 extensions whose labels are
  currently visible; badge synced in `syncControlsFromState`.

**Canvas visual-rect shrink + edge re-anchor (canvas_renderer.js)**
- `_visualHeight(node)`: `Math.max(floor, Math.min(desired, node.height))`.
  Hard-clamped to Python-reserved `node.height` so no rendering can exceed
  layout. Floor ensures a rect is always drawn.
- `_visualBottom(node)`: `node.y - node.height/2 + _visualHeight(node)`.
  Top of rect is fixed at the layout-slot top; bottom floats.
- `_effectiveSourceStart(edge)`: substitutes the polyline's first point y
  with source's `_visualBottom`, re-anchoring outgoing edges to the visible
  bottom of the drawn rect. Target-side (incoming) endpoints unchanged.
- All node draw calls (fillRect/strokeRect), label positioning, highlight
  overlays, and hit-testing updated to use `_visualHeight`/`_visualBottom`/
  `_effectiveSourceStart`.
- `edge.bounds.minY` widened at `_initTopology` time by the worst-case
  source-endpoint substitution delta so AABB hover-test covers the new range.
- JS magic literals (16, 14px, 12px) replaced with `layoutConstants.*`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: quic-boyuc

devtools/fx_viewer/exporter.py

@@ -53,11 +53,18 @@ def __init__(self, graph_module: torch.fx.GraphModule):
     _NODE_X_PADDING = 20
     _NODE_LINE_HEIGHT = 16
     _NODE_Y_PADDING = 20
+    _NODE_BASE_FONT_PX = 14
+    _NODE_EXT_FONT_PX = 12
     _LAYOUT_XSPACE = 50
     _LAYOUT_YSPACE = 30
     _DUMMY_SIZE_X = 100  # dummy nodes (from fast-sugiyama) occupy no real width/height
     _DUMMY_SIZE_Y = 30  # dummy nodes (from fast-sugiyama) occupy no real width/height
     _SPINE_COHESION_ITER = 20
+    # JS canvas LRU enforces the same cap. Layout reserves vertical space for
+    # the two extensions whose label formatters produce the most text per node,
+    # so any active subset of size <= _MAX_LABEL_EXTENSIONS fits without
+    # overlap. The constant is also emitted to the JS via _layout_constants_payload.
+    _MAX_LABEL_EXTENSIONS = 2
 
     def _default_base_label(self, node: GraphNode) -> str:
         target = str(node.info.get("target") or node.info.get("op") or "")
@@ -237,6 +244,55 @@ def _ext_label_lines_for_layout(self, extension: GraphExtension, node_id: str) -
             context=f"extension '{extension.id}' label formatter(node='{node_id}')",
         )
 
+    @classmethod
+    def _select_top_label_extensions(
+        cls,
+        per_ext_lines: dict[str, list[str]],
+    ) -> list[str]:
+        """Pick label lines from at most ``_MAX_LABEL_EXTENSIONS`` extensions.
+
+        For layout sizing only. Selects the extensions whose label formatter
+        produces the largest text (sum of character lengths across the
+        extension's lines for that node) and concatenates their lines. The JS
+        canvas LRU caps the same number of extensions at run time, so any
+        active subset fits within the reserved bbox.
+        """
+        if not per_ext_lines:
+            return []
+        if len(per_ext_lines) <= cls._MAX_LABEL_EXTENSIONS:
+            flat: list[str] = []
+            for lines in per_ext_lines.values():
+                flat.extend(lines)
+            return flat
+        ranked = sorted(
+            per_ext_lines.values(),
+            key=lambda lns: sum(len(s) for s in lns),
+            reverse=True,
+        )[: cls._MAX_LABEL_EXTENSIONS]
+        flat = []
+        for lines in ranked:
+            flat.extend(lines)
+        return flat
+
+    @classmethod
+    def _layout_constants_payload(cls) -> dict[str, Any]:
+        """Layout/font constants embedded into ``payload.layout`` for the JS runtime.
+
+        Single source of truth for sizing. Any drift between Python build-time
+        layout and JS render-time drawing is eliminated by reading these values
+        on the JS side instead of hard-coding them in the canvas renderer.
+        """
+        return {
+            "line_height": cls._NODE_LINE_HEIGHT,
+            "y_padding": cls._NODE_Y_PADDING,
+            "x_padding": cls._NODE_X_PADDING,
+            "char_width": cls._NODE_CHAR_WIDTH,
+            "min_width": cls._NODE_MIN_WIDTH,
+            "base_font_px": cls._NODE_BASE_FONT_PX,
+            "ext_font_px": cls._NODE_EXT_FONT_PX,
+            "max_label_extensions": cls._MAX_LABEL_EXTENSIONS,
+        }
+
     @classmethod
     def _compute_node_box_size(
         cls,
@@ -619,10 +675,14 @@ def _segment_crosses_aabb(
     def _compute_layout(self, nodes: dict[str, GraphNode], edges: list[GraphEdge]) -> None:
         ext_label_lines_by_node: dict[str, list[str]] = {}
         for node_id in nodes:
-            ext_lines: list[str] = []
+            per_ext_lines: dict[str, list[str]] = {}
             for ext in self.extensions:
-                ext_lines.extend(self._ext_label_lines_for_layout(ext, node_id))
-            ext_label_lines_by_node[node_id] = ext_lines
+                lines = self._ext_label_lines_for_layout(ext, node_id)
+                if lines:
+                    per_ext_lines[ext.id] = lines
+            ext_label_lines_by_node[node_id] = self._select_top_label_extensions(
+                per_ext_lines
+            )
 
         self._compute_layout_with_ext_lines(
             nodes,
@@ -695,7 +755,9 @@ def relayout_payload_base(
         else:
             active_layer_ids = [layer_id for layer_id in include_layers if layer_id in ext_payloads]
 
-        ext_label_lines_by_node: dict[str, list[str]] = {node_id: [] for node_id in nodes}
+        per_node_per_ext: dict[str, dict[str, list[str]]] = {
+            node_id: {} for node_id in nodes
+        }
         for layer_id in active_layer_ids:
             layer_payload = ext_payloads.get(layer_id)
             if not isinstance(layer_payload, dict):
@@ -704,11 +766,16 @@ def relayout_payload_base(
             if not isinstance(layer_nodes, dict):
                 continue
             for node_id, node_payload in layer_nodes.items():
-                if node_id not in ext_label_lines_by_node or not isinstance(node_payload, dict):
+                if node_id not in per_node_per_ext or not isinstance(node_payload, dict):
                     continue
-                ext_label_lines_by_node[node_id].extend(
-                    cls._coerce_str_lines(node_payload.get("label_append"))
-                )
+                lines = cls._coerce_str_lines(node_payload.get("label_append"))
+                if lines:
+                    per_node_per_ext[node_id][layer_id] = lines
+
+        ext_label_lines_by_node: dict[str, list[str]] = {
+            node_id: cls._select_top_label_extensions(per_ext_lines)
+            for node_id, per_ext_lines in per_node_per_ext.items()
+        }
 
         cls._compute_layout_with_ext_lines(
             nodes,
@@ -769,7 +836,11 @@ def generate_json_payload(self) -> Dict[str, Any]:
         self._compute_layout(nodes, edges)
         base_payload = self._build_base_payload(nodes, edges)
         extensions_payload = self._build_extensions_payload()
-        payload = GraphPayload(base=base_payload, extensions=extensions_payload)
+        payload = GraphPayload(
+            base=base_payload,
+            extensions=extensions_payload,
+            layout=self._layout_constants_payload(),
+        )
         return asdict(payload)
 
     def export_json(self, output_path: str):

devtools/fx_viewer/extension.py

@@ -129,6 +129,7 @@ def build_payload(self) -> GraphExtensionPayload:
             legend=legend,
             nodes=compiled_nodes,
             sync_keys=list(self.sync_keys),
+            has_label_formatter=self.label_formatter is not None,
         )
 
     def build(self) -> Dict[str, Any]:

devtools/fx_viewer/models.py

@@ -57,9 +57,11 @@ class GraphExtensionPayload:
     legend: list[dict[str, str]] = field(default_factory=list)
     nodes: dict[str, GraphExtensionNodePayload] = field(default_factory=dict)
     sync_keys: list[str] = field(default_factory=list)
+    has_label_formatter: bool = False
 
 
 @dataclass
 class GraphPayload:
     base: BaseGraphPayload
     extensions: dict[str, GraphExtensionPayload]
+    layout: dict[str, Any] = field(default_factory=dict)

devtools/fx_viewer/templates/canvas_renderer.js

@@ -135,18 +135,20 @@ class CanvasRenderer {
     detectHover(graphX, graphY) {
         let nearestNode = null;
         let nearestEdge = null;
-        
+
         for (let i = 0; i < this.viewer.store.baseData.nodes.length; i++) {
             const node = this.viewer.store.baseData.nodes[i];
+            const active = this.viewer.store.activeNodeMap.get(node.id) || node;
             const w = node.width;
-            const h = node.height;
+            const h = this._visualHeight(active);
+            const top = node.y - node.height / 2;
             if (graphX >= node.x - w/2 && graphX <= node.x + w/2 &&
-                graphY >= node.y - h/2 && graphY <= node.y + h/2) {
+                graphY >= top && graphY <= top + h) {
                 nearestNode = node.id;
-                break; 
+                break;
             }
         }
-        
+
         if (!nearestNode) {
             const transform = this.viewer.controller.transform;
             const hoverDist = 5 / transform.k;
@@ -160,7 +162,11 @@ class CanvasRenderer {
                 const w = this.viewer.store.activeNodeMap.get(edge.w);
                 let min_d = Infinity;
                 if (edge.points && edge.points.length > 0) {
-                    for (let j = 0; j < edge.points.length - 1; j++) {
+                    const start = this._effectiveSourceStart(edge) || edge.points[0];
+                    const a = start;
+                    const b = edge.points.length > 1 ? edge.points[1] : edge.points[0];
+                    min_d = Math.min(min_d, this.distToSegment({x: graphX, y: graphY}, a, b));
+                    for (let j = 1; j < edge.points.length - 1; j++) {
                         const d = this.distToSegment({x: graphX, y: graphY}, edge.points[j], edge.points[j+1]);
                         min_d = Math.min(min_d, d);
                     }
@@ -185,6 +191,42 @@ class CanvasRenderer {
         return Math.hypot(p.x - (v.x + t * (w.x - v.x)), p.y - (v.y + t * (w.y - v.y)));
     }
 
+    /**
+     * Visual height of a node's drawn rect. Sized for the actual count of
+     * label_append entries (LRU-filtered upstream in graph_data_store), but
+     * clamped to never exceed the Python-reserved layout slot (`node.height`).
+     * Floor at one base label line so empty-label nodes still draw a sane rect.
+     */
+    _visualHeight(node) {
+        const layout = this.viewer.store.layoutConstants;
+        const labelLines = (node && node.label_append) ? node.label_append.length : 0;
+        const desired = (1 + labelLines) * layout.line_height + layout.y_padding;
+        const floor = layout.line_height + layout.y_padding;
+        const ceiling = node && node.height ? node.height : desired;
+        return Math.max(floor, Math.min(desired, ceiling));
+    }
+
+    /**
+     * Y-coordinate of the bottom edge of the drawn rect. Top of the rect is
+     * the layout-authoritative `node.y - node.height/2`; bottom floats with
+     * the active label count. Used to re-anchor source-side edge endpoints.
+     */
+    _visualBottom(node) {
+        return node.y - node.height / 2 + this._visualHeight(node);
+    }
+
+    /**
+     * Source-side edge attach point: the polyline's first waypoint with the y
+     * shifted to the source node's current visual bottom. Target-side
+     * endpoint is unchanged (top of next-layer node is fixed).
+     */
+    _effectiveSourceStart(edge) {
+        if (!edge.points || edge.points.length === 0) return null;
+        const source = this.viewer.store.activeNodeMap.get(edge.v);
+        if (!source) return { x: edge.points[0].x, y: edge.points[0].y };
+        return { x: edge.points[0].x, y: this._visualBottom(source) };
+    }
+
     render() {
         const dpr = window.devicePixelRatio || 1;
         const ctx = this.ctx;
@@ -274,7 +316,8 @@ class CanvasRenderer {
             ctx.beginPath();
             let midX = 0, midY = 0;
             if (edge.points && edge.points.length > 0) {
-                ctx.moveTo(edge.points[0].x, edge.points[0].y);
+                const start = this._effectiveSourceStart(edge) || edge.points[0];
+                ctx.moveTo(start.x, start.y);
                 for (let i = 1; i < edge.points.length; i++) {
                     ctx.lineTo(edge.points[i].x, edge.points[i].y);
                 }
@@ -387,8 +430,11 @@ class CanvasRenderer {
                 ctx.fillStyle = renderedFill;
             }
 
-            ctx.fillRect(node.x - node.width/2, node.y - node.height/2, node.width, node.height);
-            
+            const layout = this.viewer.store.layoutConstants;
+            const visH = this._visualHeight(node);
+            const rectTop = node.y - node.height / 2;
+            ctx.fillRect(node.x - node.width/2, rectTop, node.width, visH);
+
             if (isSelected || isPreview || isHovered) {
                 ctx.strokeStyle = theme.edgeHover;
                 if (isSelected || isPreview) {
@@ -401,24 +447,27 @@ class CanvasRenderer {
                 } else {
                     ctx.setLineDash([]);
                 }
-                ctx.strokeRect(node.x - node.width/2, node.y - node.height/2, node.width, node.height);
+                ctx.strokeRect(node.x - node.width/2, rectTop, node.width, visH);
                 ctx.setLineDash([]);
             }
-            
+
             ctx.fillStyle = node.fill_color
                 ? (fxReadableTextColor(renderedFill) || theme.text)
                 : theme.text;
             let allLines = [node.label || node.id];
             if (node.label_append && node.label_append.length > 0) {
                 allLines = allLines.concat(node.label_append);
             }
-            
-            const lineHeight = 16;
-            const startY = node.y - ((allLines.length - 1) * lineHeight) / 2;
-            
+
+            const lineHeight = layout.line_height;
+            // Anchor label rows to the top of the visible rect; previous code
+            // centered around node.y, which drifted off the visible rect once
+            // the rect started shrinking from full layout slot.
+            const startY = rectTop + (layout.y_padding / 2) + lineHeight / 2;
+
             for (let i = 0; i < allLines.length; i++) {
-                if (i === 0) ctx.font = 'bold 14px sans-serif';
-                else ctx.font = '12px sans-serif';
+                if (i === 0) ctx.font = `bold ${layout.base_font_px}px sans-serif`;
+                else ctx.font = `${layout.ext_font_px}px sans-serif`;
                 ctx.fillText(allLines[i], node.x, startY + (i * lineHeight));
             }
 
@@ -441,11 +490,13 @@ class CanvasRenderer {
                 nodeIds.forEach((id) => {
                     const node = this.viewer.store.activeNodeMap.get(id);
                     if (!node) return;
+                    const visH = this._visualHeight(node);
+                    const rectTop = node.y - node.height / 2;
                     ctx.strokeRect(
                         node.x - node.width / 2 - outerOffset,
-                        node.y - node.height / 2 - outerOffset,
+                        rectTop - outerOffset,
                         node.width + outerOffset * 2,
-                        node.height + outerOffset * 2
+                        visH + outerOffset * 2
                     );
                 });
             });
@@ -534,7 +585,7 @@ class CanvasRenderer {
 
         const padding = 8 / transform.k;
         const tw = maxW + padding * 2;
-        const lineHeight = 16 / transform.k;
+        const lineHeight = this.viewer.store.layoutConstants.line_height / transform.k;
         const th = (tooltipLines.length * lineHeight) + padding * 2;
 
         const viewLeft = -transform.x / transform.k;