feat(webapp): select runs and spans in the session Trace tab, with deep links

Behind a feature flag, selecting a run or span in the session Trace tab
shows it in the side panel, and the selection is kept in the URL so it
can be shared. Links to older runs open with that run expanded and the
span selected.

Mono-RevId: a5c0191b681a71dbf8f9b91ef514a5f1402f6737

Author: Mariefay

apps/webapp/app/components/primitives/TreeView/TreeView.tsx

@@ -162,6 +162,8 @@ type TreeStateHookProps<TData, TFilterValue> = {
   filter?: Filter<TData, TFilterValue>;
   /** Nodes that start collapsed when they first appear, initially or as the tree grows. */
   isCollapsedByDefault?: (node: FlatTreeItem<TData>) => boolean;
+  /** Rows arrow keys can land on; others are skipped. Every row when unset. */
+  canSelect?: (node: FlatTreeItem<TData>) => boolean;
   /** Whether a node can be expanded or collapsed. Defaults to having children. */
   canExpand?: (node: FlatTreeItem<TData>) => boolean;
   /** Called when a single node is expanded or collapsed (click, keyboard, toggle). */
@@ -200,7 +202,7 @@ export type UseTreeStateOutput = {
   selectNextVisibleNode: (scrollToNode?: boolean) => void;
   selectPreviousVisibleNode: (scrollToNode?: boolean) => void;
   selectParentNode: (scrollToNode?: boolean) => void;
-  scrollToNode: (id: string) => void;
+  scrollToNode: (id: string, align?: "auto" | "center") => void;
 };
 
 export function useTree<TData, TFilterValue>({
@@ -212,6 +214,7 @@ export function useTree<TData, TFilterValue>({
   estimatedRowHeight,
   filter,
   isCollapsedByDefault,
+  canSelect,
   canExpand = (node) => node.hasChildren,
   onExpandedChange,
   structureKey,
@@ -344,11 +347,11 @@ export function useTree<TData, TFilterValue>({
     overscan: 50,
   });
 
-  const scrollToNodeFn = (id: string) => {
+  const scrollToNodeFn = (id: string, align: "auto" | "center" = "auto") => {
     const itemIndex = state.visibleNodeIds.findIndex((nodeId) => nodeId === id);
 
     if (itemIndex !== -1) {
-      virtualizer.scrollToIndex(itemIndex, { align: "auto" });
+      virtualizer.scrollToIndex(itemIndex, { align });
     }
   };
 
@@ -396,35 +399,35 @@ export function useTree<TData, TFilterValue>({
   const selectFirstVisibleNode = (scrollToNode = true) => {
     dispatch({
       type: "SELECT_FIRST_VISIBLE_NODE",
-      payload: { scrollToNode, scrollToNodeFn },
+      payload: { scrollToNode, scrollToNodeFn, canSelect },
     });
   };
 
   const selectLastVisibleNode = (scrollToNode = true) => {
     dispatch({
       type: "SELECT_LAST_VISIBLE_NODE",
-      payload: { scrollToNode, scrollToNodeFn },
+      payload: { scrollToNode, scrollToNodeFn, canSelect },
     });
   };
 
   const selectNextVisibleNode = (scrollToNode = true) => {
     dispatch({
       type: "SELECT_NEXT_VISIBLE_NODE",
-      payload: { scrollToNode, scrollToNodeFn },
+      payload: { scrollToNode, scrollToNodeFn, canSelect },
     });
   };
 
   const selectPreviousVisibleNode = (scrollToNode = true) => {
     dispatch({
       type: "SELECT_PREVIOUS_VISIBLE_NODE",
-      payload: { scrollToNode, scrollToNodeFn },
+      payload: { scrollToNode, scrollToNodeFn, canSelect },
     });
   };
 
   const selectParentNode = (scrollToNode = true) => {
     dispatch({
       type: "SELECT_PARENT_NODE",
-      payload: { scrollToNode, scrollToNodeFn },
+      payload: { scrollToNode, scrollToNodeFn, canSelect },
     });
   };
 

apps/webapp/app/components/primitives/TreeView/reducer.ts

@@ -40,6 +40,11 @@ type WithScrollToNode = {
   scrollToNodeFn: ScrollToNodeFn;
 };
 
+/** Rows the visible-node moves can land on; every row when unset. */
+type CanSelectNode = (node: FlatTree<any>[number]) => boolean;
+
+type WithCanSelect = { canSelect?: CanSelectNode };
+
 type UpdateTreeAction = {
   type: "UPDATE_TREE";
   payload: {
@@ -132,27 +137,27 @@ type ToggleExpandLevelAction = {
 
 type SelectFirstVisibleNodeAction = {
   type: "SELECT_FIRST_VISIBLE_NODE";
-  payload: {} & WithScrollToNode;
+  payload: {} & WithScrollToNode & WithCanSelect;
 };
 
 type SelectLastVisibleNodeAction = {
   type: "SELECT_LAST_VISIBLE_NODE";
-  payload: {} & WithScrollToNode;
+  payload: {} & WithScrollToNode & WithCanSelect;
 };
 
 type SelectNextVisibleNodeAction = {
   type: "SELECT_NEXT_VISIBLE_NODE";
-  payload: {} & WithScrollToNode;
+  payload: {} & WithScrollToNode & WithCanSelect;
 };
 
 type SelectPreviousVisibleNodeAction = {
   type: "SELECT_PREVIOUS_VISIBLE_NODE";
-  payload: {} & WithScrollToNode;
+  payload: {} & WithScrollToNode & WithCanSelect;
 };
 
 type SelectParentNodeAction = {
   type: "SELECT_PARENT_NODE";
-  payload: {} & WithScrollToNode;
+  payload: {} & WithScrollToNode & WithCanSelect;
 };
 
 type UpdateFilterAction = {
@@ -409,7 +414,10 @@ export function reducer(state: TreeState, action: Action): TreeState {
       });
     }
     case "SELECT_FIRST_VISIBLE_NODE": {
-      const node = firstVisibleNode(state.tree, state.filteredNodes);
+      const { canSelect } = action.payload;
+      const node = canSelect
+        ? visibleNodes(state.tree, state.filteredNodes).find(canSelect)
+        : firstVisibleNode(state.tree, state.filteredNodes);
       if (node) {
         return reducer(state, {
           type: "SELECT_NODE",
@@ -424,7 +432,10 @@ export function reducer(state: TreeState, action: Action): TreeState {
       return state;
     }
     case "SELECT_LAST_VISIBLE_NODE": {
-      const node = lastVisibleNode(state.tree, state.filteredNodes);
+      const { canSelect } = action.payload;
+      const node = canSelect
+        ? findLastNode(visibleNodes(state.tree, state.filteredNodes), canSelect)
+        : lastVisibleNode(state.tree, state.filteredNodes);
       if (node) {
         return reducer(state, {
           type: "SELECT_NODE",
@@ -446,13 +457,17 @@ export function reducer(state: TreeState, action: Action): TreeState {
           payload: {
             scrollToNode: action.payload.scrollToNode,
             scrollToNodeFn: action.payload.scrollToNodeFn,
+            canSelect: action.payload.canSelect,
           },
         });
       }
 
       const visible = visibleNodes(state.tree, state.filteredNodes);
       const selectedIndex = visible.findIndex((node) => node.id === selected);
-      const nextNode = visible[selectedIndex + 1];
+      const { canSelect } = action.payload;
+      const nextNode = canSelect
+        ? visible.slice(selectedIndex + 1).find(canSelect)
+        : visible[selectedIndex + 1];
       if (nextNode) {
         return reducer(state, {
           type: "SELECT_NODE",
@@ -475,13 +490,20 @@ export function reducer(state: TreeState, action: Action): TreeState {
           payload: {
             scrollToNode: action.payload.scrollToNode,
             scrollToNodeFn: action.payload.scrollToNodeFn,
+            canSelect: action.payload.canSelect,
           },
         });
       }
 
       const visible = visibleNodes(state.tree, state.filteredNodes);
       const selectedIndex = visible.findIndex((node) => node.id === selected);
-      const previousNode = visible[Math.max(0, selectedIndex - 1)];
+      const { canSelect } = action.payload;
+      // A hidden selection (its parent collapsed) starts again from the top, as without canSelect.
+      const previousNode = canSelect
+        ? selectedIndex === -1
+          ? visible.find(canSelect)
+          : findLastNode(visible.slice(0, selectedIndex), canSelect)
+        : visible[Math.max(0, selectedIndex - 1)];
       if (previousNode) {
         return reducer(state, {
           type: "SELECT_NODE",
@@ -504,6 +526,7 @@ export function reducer(state: TreeState, action: Action): TreeState {
           payload: {
             scrollToNode: action.payload.scrollToNode,
             scrollToNodeFn: action.payload.scrollToNodeFn,
+            canSelect: action.payload.canSelect,
           },
         });
       }
@@ -513,7 +536,13 @@ export function reducer(state: TreeState, action: Action): TreeState {
         return state;
       }
 
-      const parentNode = state.tree.find((node) => node.id === selectedNode.parentId);
+      // With canSelect, skip parents that can't be selected.
+      const { canSelect } = action.payload;
+      let parentNode = state.tree.find((node) => node.id === selectedNode.parentId);
+      while (parentNode && canSelect && !canSelect(parentNode)) {
+        const parentId: string | undefined = parentNode.parentId;
+        parentNode = state.tree.find((node) => node.id === parentId);
+      }
       if (parentNode) {
         return reducer(state, {
           type: "SELECT_NODE",
@@ -556,3 +585,10 @@ export function reducer(state: TreeState, action: Action): TreeState {
 
   throw new Error(`Unhandled action type: ${(action as any).type}`);
 }
+
+function findLastNode(nodes: FlatTree<any>, canSelect: CanSelectNode) {
+  for (let i = nodes.length - 1; i >= 0; i--) {
+    if (canSelect(nodes[i])) return nodes[i];
+  }
+  return undefined;
+}

apps/webapp/app/components/sessions/v1/trace/RunTraceSource.tsx

@@ -9,6 +9,7 @@ import { useProject } from "~/hooks/useProject";
 import type { RunTraceStart } from "~/presenters/v3/TraceChunkPresenter.server";
 import { v3RunTraceChunkPath } from "~/utils/pathBuilder";
 import { isFinalRunStatus } from "~/v3/taskStatus";
+import type { RunPanelData } from "./sessionPanelStore";
 import type { SessionRunLoad, SessionRunLoadOutcome } from "./sessionRunLoads";
 import type { SessionTraceRun } from "./sessionTraceTypes";
 
@@ -28,15 +29,18 @@ type RunTraceSourceProps = {
   run: SessionTraceRun & { status: TaskRunStatus };
   load: SessionRunLoad;
   live: boolean;
+  /** A deep-linked span: the first chunk includes it and its parents even past the first page. */
+  selectedSpanId?: string;
   onSettle: (runFriendlyId: string, loadId: number, outcome: SessionRunLoadOutcome) => void;
   onSnapshot: (runFriendlyId: string, snapshot: RunTraceSnapshot | undefined) => void;
   onSpansChanged: (runFriendlyId: string, hasSpans: boolean) => void;
+  onPanelData: (runFriendlyId: string, data: RunPanelData | undefined) => void;
   onRunDataStale: () => void;
 };
 
 /** Loads one run's trace for the session tree. Renders nothing. */
 export function RunTraceSource(props: RunTraceSourceProps) {
-  const { run, load, onSettle } = props;
+  const { run, load, selectedSpanId, onSettle } = props;
   const organization = useOrganization();
   const project = useProject();
   const environment = useEnvironment();
@@ -49,7 +53,7 @@ export function RunTraceSource(props: RunTraceSourceProps) {
   useEffect(() => {
     if (fetchLoadId === undefined) return;
     const controller = new AbortController();
-    void fetchTraceStart(chunkPath, controller.signal).then((result) => {
+    void fetchTraceStart(chunkPath, selectedSpanId, controller.signal).then((result) => {
       if (controller.signal.aborted) return;
       if (result.kind === "progressive" || result.kind === "empty") {
         // Same JSON payload the run page loader sends; the hook reads its dates as strings.
@@ -59,7 +63,7 @@ export function RunTraceSource(props: RunTraceSourceProps) {
       onSettle(run.friendlyId, fetchLoadId, toOutcome(result));
     });
     return () => controller.abort();
-  }, [fetchLoadId, chunkPath, run.friendlyId, onSettle]);
+  }, [fetchLoadId, chunkPath, selectedSpanId, run.friendlyId, onSettle]);
 
   if (!started) return null;
   return <RunTraceLive {...props} started={started} />;
@@ -73,6 +77,7 @@ function RunTraceLive({
   onSettle,
   onSnapshot,
   onSpansChanged,
+  onPanelData,
   onRunDataStale,
 }: RunTraceSourceProps & { started: StartedTrace }) {
   const trace = useRunTrace({
@@ -102,6 +107,16 @@ function RunTraceLive({
 
   useEffect(() => () => onSnapshot(run.friendlyId, undefined), [run.friendlyId, onSnapshot]);
 
+  const { inspectorRefreshKey, overridesBySpanId, linkedRunIdBySpanId } = trace;
+  useEffect(() => {
+    onPanelData(run.friendlyId, {
+      refreshKey: inspectorRefreshKey,
+      overridesBySpanId,
+      linkedRunIdBySpanId,
+    });
+  }, [run.friendlyId, inspectorRefreshKey, overridesBySpanId, linkedRunIdBySpanId, onPanelData]);
+  useEffect(() => () => onPanelData(run.friendlyId, undefined), [run.friendlyId, onPanelData]);
+
   // The run's root span is hidden, so only spans below it count as on screen.
   const hasSpans = trace.events.some((event) => event.level > 0);
   useEffect(() => {
@@ -139,9 +154,14 @@ export function settledOutcome(
 
 type TraceStartResult = RunTraceStart | { kind: "failed"; message: string };
 
-async function fetchTraceStart(chunkPath: string, signal: AbortSignal): Promise<TraceStartResult> {
+async function fetchTraceStart(
+  chunkPath: string,
+  selectedSpanId: string | undefined,
+  signal: AbortSignal
+): Promise<TraceStartResult> {
   const url = new URL(chunkPath, window.location.origin);
   url.searchParams.set("start", "1");
+  if (selectedSpanId) url.searchParams.set("span", selectedSpanId);
 
   const [error, response] = await tryCatch(fetch(url, { signal }));
   if (error) return { kind: "failed", message: "Couldn't load this run's trace." };