diff --git a/app/src/components/SpanTree.tsx b/app/src/components/SpanTree.tsx index efd7659..d0c2837 100644 --- a/app/src/components/SpanTree.tsx +++ b/app/src/components/SpanTree.tsx @@ -55,14 +55,113 @@ function typeInfo(span: Span) { return TYPE_LABEL.INTERNAL; } -function SpanRow({ span, depth, minTime, totalDur, selected, flashing, onClick, onContextMenu, annotations, freshIds, onClearFresh }: { - span: Span; depth: number; minTime: number; totalDur: number; - selected: boolean; flashing: boolean; onClick: () => void; +interface SpanTreeModel { + spanMap: Map; + children: Map; + roots: Span[]; + descendantCounts: Map; +} + +interface VisibleSpanRow { + span: Span; + depth: number; +} + +function buildSpanTreeModel(spans: Span[]): SpanTreeModel { + const spanMap = new Map(spans.map(s => [s.id, s])); + const children = new Map(); + const roots: Span[] = []; + + for (const span of spans) { + if (!span.parent_span_id || !spanMap.has(span.parent_span_id)) { + roots.push(span); + continue; + } + + const siblings = children.get(span.parent_span_id) ?? []; + siblings.push(span); + children.set(span.parent_span_id, siblings); + } + + const descendantCounts = new Map(); + const countDescendants = (span: Span): number => { + let count = 0; + for (const child of children.get(span.id) ?? []) { + count += 1 + countDescendants(child); + } + descendantCounts.set(span.id, count); + return count; + }; + + for (const root of roots) countDescendants(root); + + return { spanMap, children, roots, descendantCounts }; +} + +function getVisibleSpanRows(model: SpanTreeModel, collapsedIds: Set): VisibleSpanRow[] { + const rows: VisibleSpanRow[] = []; + + const walk = (span: Span, depth: number) => { + rows.push({ span, depth }); + if (collapsedIds.has(span.id)) return; + for (const child of model.children.get(span.id) ?? []) walk(child, depth + 1); + }; + + for (const root of model.roots) walk(root, 0); + return rows; +} + +function getAncestorIds(spanId: string, spanMap: Map): string[] { + const ancestors: string[] = []; + let parentId = spanMap.get(spanId)?.parent_span_id ?? null; + + while (parentId) { + ancestors.push(parentId); + parentId = spanMap.get(parentId)?.parent_span_id ?? null; + } + + return ancestors; +} + +function isDescendantOf(spanId: string, ancestorId: string, spanMap: Map): boolean { + return getAncestorIds(spanId, spanMap).includes(ancestorId); +} + +interface SpanRowProps { + span: Span; + depth: number; + minTime: number; + totalDur: number; + selected: boolean; + flashing: boolean; + hasChildren: boolean; + collapsed: boolean; + hiddenDescendantCount: number; + onClick: () => void; + onToggleCollapse: (e: React.MouseEvent) => void; onContextMenu?: (e: React.MouseEvent, span: Span) => void; annotations: Annotation[]; freshIds: Set; onClearFresh: (id: string) => void; -}) { +} + +function SpanRow({ + span, + depth, + minTime, + totalDur, + selected, + flashing, + hasChildren, + collapsed, + hiddenDescendantCount, + onClick, + onToggleCollapse, + onContextMenu, + annotations, + freshIds, + onClearFresh, +}: SpanRowProps) { const info = typeInfo(span); const color = info.color; const isErr = span.status === "ERROR"; @@ -84,12 +183,32 @@ function SpanRow({ span, depth, minTime, totalDur, selected, flashing, onClick, onContextMenu={onContextMenu ? (e) => { e.preventDefault(); onContextMenu(e, span); } : undefined} >
+ {hasChildren ? ( + + ) : ( + + )} {info.label} {span.name} + {collapsed && hiddenDescendantCount > 0 && ( + + +{hiddenDescendantCount} + + )} {annotations.map((a) => ( (null); const [addingForSpan, setAddingForSpan] = useState(null); const [flashId, setFlashId] = useState(null); + const [collapsedIds, setCollapsedIds] = useState>(() => new Set()); const runId = spans[0]?.run_id ?? null; - const reportedSelectedId = selectedId && spans.some((s) => s.id === selectedId) ? selectedId : null; + const treeModel = useMemo(() => buildSpanTreeModel(spans), [spans]); + const reportedSelectedId = selectedId && treeModel.spanMap.has(selectedId) ? selectedId : null; useEffect(() => { - if (controlled || !runId || autoSelectedRunRef.current === runId) return; + if (!runId || autoSelectedRunRef.current === runId) return; autoSelectedRunRef.current = runId; - setInternalSelectedId(spans[0]?.id ?? null); + if (!controlled) setInternalSelectedId(spans[0]?.id ?? null); + setCollapsedIds(new Set()); }, [controlled, runId, spans]); useEffect(() => { @@ -321,6 +443,13 @@ export function SpanTree({ const handler = (ev: Event) => { const spanId = (ev as CustomEvent).detail?.spanId as string | undefined; if (!spanId || !spanIds.has(spanId)) return; + setCollapsedIds((prev) => { + const next = new Set(prev); + for (const parentId of getAncestorIds(spanId, treeModel.spanMap)) { + next.delete(parentId); + } + return next; + }); setInternalSelectedId(spanId); setFlashId(spanId); requestAnimationFrame(() => { @@ -331,10 +460,17 @@ export function SpanTree({ }; window.addEventListener("workshop:deep-link-span", handler); return () => window.removeEventListener("workshop:deep-link-span", handler); - }, [controlled, spans]); + }, [controlled, spans, treeModel.spanMap]); useEffect(() => { if (!selectedSpanId) return; + setCollapsedIds((prev) => { + const next = new Set(prev); + for (const parentId of getAncestorIds(selectedSpanId, treeModel.spanMap)) { + next.delete(parentId); + } + return next; + }); setFlashId(selectedSpanId); requestAnimationFrame(() => { const el = document.querySelector(`[data-span-row="${selectedSpanId}"]`); @@ -342,7 +478,7 @@ export function SpanTree({ }); const timeout = window.setTimeout(() => setFlashId(null), 1500); return () => window.clearTimeout(timeout); - }, [selectedSpanId]); + }, [selectedSpanId, treeModel.spanMap]); // Dismiss context menu on scroll / outside click / escape useEffect(() => { @@ -370,26 +506,26 @@ export function SpanTree({ return map; }, [annotations]); - const spanMap = new Map(spans.map(s => [s.id, s])); - const children = new Map(); - const roots: Span[] = []; - for (const s of spans) { - if (!s.parent_span_id || !spanMap.has(s.parent_span_id)) roots.push(s); - else { const c = children.get(s.parent_span_id) ?? []; c.push(s); children.set(s.parent_span_id, c); } - } - - const flat: { span: Span; depth: number }[] = []; - function walk(span: Span, depth: number) { - flat.push({ span, depth }); - for (const kid of children.get(span.id) ?? []) walk(kid, depth + 1); - } - for (const r of roots) walk(r, 0); + const flat = useMemo(() => getVisibleSpanRows(treeModel, collapsedIds), [treeModel, collapsedIds]); const minTime = flat.length > 0 ? Math.min(...flat.map(f => f.span.start_time_ms)) : 0; const maxTime = flat.length > 0 ? Math.max(...flat.map(f => f.span.end_time_ms)) : 0; const totalDur = maxTime - minTime || 1; - const selectedSpan = selectedId ? spanMap.get(selectedId) : null; + const selectedSpan = selectedId ? treeModel.spanMap.get(selectedId) : null; + const toggleCollapse = (span: Span) => (e: React.MouseEvent) => { + e.stopPropagation(); + const isCollapsing = !collapsedIds.has(span.id); + if (isCollapsing && selectedId && isDescendantOf(selectedId, span.id, treeModel.spanMap)) { + setSelectedId(span.id); + } + setCollapsedIds((prev) => { + const next = new Set(prev); + if (next.has(span.id)) next.delete(span.id); + else next.add(span.id); + return next; + }); + }; if (flat.length === 0) return
No spans
; @@ -411,7 +547,11 @@ export function SpanTree({ minTime={minTime} totalDur={totalDur} selected={span.id === selectedId} flashing={span.id === flashId} + hasChildren={(treeModel.children.get(span.id)?.length ?? 0) > 0} + collapsed={collapsedIds.has(span.id)} + hiddenDescendantCount={treeModel.descendantCounts.get(span.id) ?? 0} onClick={() => setSelectedId(span.id === selectedId ? null : span.id)} + onToggleCollapse={toggleCollapse(span)} onContextMenu={onCreateAnnotation ? (e, s) => setContextMenu({ spanId: s.id, x: e.clientX, y: e.clientY }) : undefined} annotations={annotationsBySpan.get(span.id) ?? []} freshIds={freshIds}