feat(trace-tree-node): Replacing legacy TraceTreeNode usage with new BaseNode class

knip.config.ts

@@ -57,8 +57,6 @@ const config: KnipConfig = {
     // helper files for stories - it's fine that they are only used in tests
     '!static/app/**/__stories__/*.{js,mjs,ts,tsx}!',
     '!static/app/stories/**/*.{js,mjs,ts,tsx}!',
-    // TEMPORARY! Abdullah Khan: WILL BE REMOVING IN STACKED PRs. Trying to merge PRs in smaller batches.
-    '!static/app/views/performance/newTraceDetails/traceModels/traceTreeNode/**/*.{js,mjs,ts,tsx}!',
   ],
   compilers: {
     mdx: async text => String(await compile(text)),

static/app/components/events/eventEntries.tsx

@@ -139,7 +139,7 @@ function EventEntries({
 // Because replays are not an interface, we need to manually insert the replay section
 // into the array of entries. The long-term solution here is to move the ordering
 // logic to this component, similar to how GroupEventDetailsContent works.
-export function partitionEntriesForReplay(entries: Entry[]) {
+function partitionEntriesForReplay(entries: Entry[]) {
   let replayIndex = 0;
 
   for (const [i, entry] of entries.entries()) {

static/app/views/insights/agents/components/aiSpanList.tsx

@@ -29,8 +29,8 @@ import {
   isTransactionNode,
   isTransactionNodeEquivalent,
 } from 'sentry/views/performance/newTraceDetails/traceGuards';
-import {TraceTree} from 'sentry/views/performance/newTraceDetails/traceModels/traceTree';
-import type {TraceTreeNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode';
+import type {EapSpanNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/eapSpanNode';
+import type {TransactionNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/transactionNode';
 
 function getNodeTimeBounds(node: AITraceSpanNode | AITraceSpanNode[]) {
   let startTime = 0;
@@ -69,16 +69,6 @@ function getNodeTimeBounds(node: AITraceSpanNode | AITraceSpanNode[]) {
   };
 }
 
-function getClosestNode<T extends AITraceSpanNode>(
-  node: AITraceSpanNode,
-  predicate: (node: TraceTreeNode) => node is T
-): T | null {
-  if (predicate(node)) {
-    return node;
-  }
-  return TraceTree.ParentNode(node, predicate) as T | null;
-}
-
 export function AISpanList({
   nodes,
   selectedNodeKey,
@@ -89,12 +79,11 @@ export function AISpanList({
   selectedNodeKey: string | null;
 }) {
   const nodesByTransaction = useMemo(() => {
-    const result: Map<
-      TraceTreeNode<TraceTree.Transaction | TraceTree.EAPSpan>,
-      AITraceSpanNode[]
-    > = new Map();
+    const result: Map<TransactionNode | EapSpanNode, AITraceSpanNode[]> = new Map();
     for (const node of nodes) {
-      const transaction = getClosestNode(node, isTransactionNodeEquivalent);
+      const transaction = node.findParent<TransactionNode | EapSpanNode>(p =>
+        isTransactionNodeEquivalent(p)
+      );
       if (!transaction) {
         continue;
       }
@@ -132,7 +121,7 @@ function TransactionWrapper({
   nodes: AITraceSpanNode[];
   onSelectNode: (node: AITraceSpanNode) => void;
   selectedNodeKey: string | null;
-  transaction: TraceTreeNode<TraceTree.Transaction | TraceTree.EAPSpan>;
+  transaction: TransactionNode | EapSpanNode;
 }) {
   const [isExpanded, setIsExpanded] = useState(true);
   const theme = useTheme();
@@ -142,7 +131,7 @@ function TransactionWrapper({
   const nodeAiRunParentsMap = useMemo<Record<string, AITraceSpanNode>>(() => {
     const parents: Record<string, AITraceSpanNode> = {};
     for (const node of nodes) {
-      const parent = getClosestNode(node, getIsAiRunNode);
+      const parent = node.findParent<AITraceSpanNode>(p => getIsAiRunNode(p));
       if (parent) {
         parents[getNodeId(node)] = parent;
       }
@@ -242,7 +231,7 @@ interface TraceBounds {
 }
 
 function calculateRelativeTiming(
-  node: TraceTreeNode<TraceTree.NodeValue>,
+  node: AITraceSpanNode,
   traceBounds: TraceBounds
 ): {leftPercent: number; widthPercent: number} {
   if (!node.value) return {leftPercent: 0, widthPercent: 0};

static/app/views/insights/agents/hooks/useAITrace.tsx

@@ -13,7 +13,7 @@ import {
   isTransactionNodeEquivalent,
 } from 'sentry/views/performance/newTraceDetails/traceGuards';
 import {TraceTree} from 'sentry/views/performance/newTraceDetails/traceModels/traceTree';
-import type {TraceTreeNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode';
+import type {TransactionNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/transactionNode';
 import {DEFAULT_TRACE_VIEW_PREFERENCES} from 'sentry/views/performance/newTraceDetails/traceState/tracePreferences';
 import {useTraceQueryParams} from 'sentry/views/performance/newTraceDetails/useTraceQueryParams';
 
@@ -69,19 +69,21 @@ export function useAITrace(traceSlug: string): UseAITraceResult {
 
         tree.build();
 
-        const fetchableTransactions = TraceTree.FindAll(tree.root, node => {
-          return isTransactionNode(node) && node.canFetch && node.value !== null;
-        }).filter((node): node is TraceTreeNode<TraceTree.Transaction> =>
-          isTransactionNode(node)
-        );
+        const fetchableTransactions = tree.root
+          .findAllChildren(node => {
+            return (
+              isTransactionNode(node) && node.canFetchChildren && node.value !== null
+            );
+          })
+          .filter((node): node is TransactionNode => isTransactionNode(node));
 
         const uniqueTransactions = fetchableTransactions.filter(
           (node, index, array) =>
             index === array.findIndex(tx => tx.value.event_id === node.value.event_id)
         );
 
         const zoomPromises = uniqueTransactions.map(node =>
-          tree.zoom(node, true, {
+          tree.fetchNodeSubTree(true, node, {
             api,
             organization,
             preferences: DEFAULT_TRACE_VIEW_PREFERENCES,
@@ -91,7 +93,7 @@ export function useAITrace(traceSlug: string): UseAITraceResult {
         await Promise.all(zoomPromises);
 
         // Keep only transactions that include AI spans and the AI spans themselves
-        const flattenedNodes = TraceTree.FindAll(tree.root, node => {
+        const flattenedNodes = tree.root.findAllChildren<AITraceSpanNode>(node => {
           if (
             !isTransactionNodeEquivalent(node) &&
             !isSpanNode(node) &&
@@ -101,7 +103,7 @@ export function useAITrace(traceSlug: string): UseAITraceResult {
           }
 
           return getIsAiNode(node);
-        }) as AITraceSpanNode[];
+        });
 
         setNodes(flattenedNodes);
         setIsLoading(false);

static/app/views/insights/agents/utils/aiTraceNodes.tsx

@@ -13,8 +13,7 @@ import {
   isSpanNode,
   isTransactionNode,
 } from 'sentry/views/performance/newTraceDetails/traceGuards';
-import type {TraceTree} from 'sentry/views/performance/newTraceDetails/traceModels/traceTree';
-import type {TraceTreeNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode';
+import type {BaseNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/baseNode';
 
 // TODO(aknaus): Remove the special handling for tags once the endpoint returns the correct type
 function getAttributeValue(
@@ -37,7 +36,7 @@ function getAttributeValue(
 }
 
 export function ensureAttributeObject(
-  node: TraceTreeNode<TraceTree.NodeValue>,
+  node: AITraceSpanNode,
   event?: EventTransaction,
   attributes?: TraceItemResponseAttribute[]
 ) {
@@ -66,7 +65,7 @@ export function ensureAttributeObject(
 
 export function getTraceNodeAttribute(
   name: string,
-  node: TraceTreeNode<TraceTree.NodeValue>,
+  node: AITraceSpanNode,
   event?: EventTransaction,
   attributes?: TraceItemResponseAttribute[]
 ): string | number | boolean | undefined {
@@ -75,7 +74,7 @@ export function getTraceNodeAttribute(
 }
 
 function createGetIsAiNode(predicate: ({op}: {op?: string}) => boolean) {
-  return (node: TraceTreeNode<TraceTree.NodeValue>): node is AITraceSpanNode => {
+  return (node: BaseNode): node is AITraceSpanNode => {
     if (!isTransactionNode(node) && !isSpanNode(node) && !isEAPSpanNode(node)) {
       return false;
     }

static/app/views/insights/agents/utils/getNodeId.tsx

@@ -7,7 +7,7 @@ import {
 
 export function getNodeId(node: AITraceSpanNode): string {
   if (isEAPSpanNode(node)) {
-    return node.metadata.event_id as string;
+    return node.value.event_id;
   }
   if (isTransactionNode(node)) {
     return node.value.event_id;

static/app/views/insights/agents/utils/types.tsx

@@ -1,6 +1,5 @@
-import type {TraceTree} from 'sentry/views/performance/newTraceDetails/traceModels/traceTree';
-import type {TraceTreeNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode';
+import type {EapSpanNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/eapSpanNode';
+import type {SpanNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/spanNode';
+import type {TransactionNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/transactionNode';
 
-export type AITraceSpanNode = TraceTreeNode<
-  TraceTree.Transaction | TraceTree.EAPSpan | TraceTree.Span
->;
+export type AITraceSpanNode = TransactionNode | EapSpanNode | SpanNode;

static/app/views/insights/mcp/utils/mcpTraceNodes.tsx

@@ -3,10 +3,9 @@ import {
   isSpanNode,
   isTransactionNode,
 } from 'sentry/views/performance/newTraceDetails/traceGuards';
-import type {TraceTree} from 'sentry/views/performance/newTraceDetails/traceModels/traceTree';
-import type {TraceTreeNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode';
+import type {BaseNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/baseNode';
 
-export function getIsMCPNode(node: TraceTreeNode<TraceTree.NodeValue>) {
+export function getIsMCPNode(node: BaseNode) {
   if (!isTransactionNode(node) && !isSpanNode(node) && !isEAPSpanNode(node)) {
     return false;
   }

static/app/views/performance/newTraceDetails/issuesTraceWaterfall.tsx

@@ -27,11 +27,10 @@ import {
   isTransactionNode,
 } from 'sentry/views/performance/newTraceDetails/traceGuards';
 import {IssuesTraceTree} from 'sentry/views/performance/newTraceDetails/traceModels/issuesTraceTree';
-import {TraceTree} from 'sentry/views/performance/newTraceDetails/traceModels/traceTree';
 import {useDividerResizeSync} from 'sentry/views/performance/newTraceDetails/useDividerResizeSync';
 import {useTraceSpaceListeners} from 'sentry/views/performance/newTraceDetails/useTraceSpaceListeners';
 
-import type {TraceTreeNode} from './traceModels/traceTreeNode';
+import type {BaseNode} from './traceModels/traceTreeNode/baseNode';
 import {useTraceState, useTraceStateDispatch} from './traceState/traceStateProvider';
 import {Trace} from './trace';
 import {
@@ -80,9 +79,7 @@ export function IssuesTraceWaterfall(props: IssuesTraceWaterfallProps) {
     });
   }, [props.organization, props.source]);
 
-  const previouslyFocusedNodeRef = useRef<TraceTreeNode<TraceTree.NodeValue> | null>(
-    null
-  );
+  const previouslyFocusedNodeRef = useRef<BaseNode | null>(null);
 
   const {viewManager, traceScheduler, traceView} = useTraceWaterfallModels();
 
@@ -95,17 +92,13 @@ export function IssuesTraceWaterfall(props: IssuesTraceWaterfallProps) {
   }, [props.tree.list.length, traceDispatch]);
 
   const onRowClick = useCallback(
-    (
-      node: TraceTreeNode<TraceTree.NodeValue>,
-      _event: React.MouseEvent<HTMLElement>,
-      index: number
-    ) => {
+    (node: BaseNode, _event: React.MouseEvent<HTMLElement>, index: number) => {
       trackAnalytics('trace.trace_layout.span_row_click', {
         organization,
         num_children: node.children.length,
         type: traceNodeAnalyticsName(node),
         project_platform:
-          projects.find(p => p.slug === node.metadata.project_slug)?.platform || 'other',
+          projects.find(p => p.slug === node.projectSlug)?.platform || 'other',
         ...traceNodeAdjacentAnalyticsProperties(node),
       });
 
@@ -134,7 +127,7 @@ export function IssuesTraceWaterfall(props: IssuesTraceWaterfallProps) {
 
     // Find all the nodes that match the event id from the error so that we can try and
     // link the user to the most specific one.
-    const nodes = IssuesTraceTree.FindAll(props.tree.root, n => {
+    const nodes = props.tree.root.findAllChildren(n => {
       if (isTraceErrorNode(n) || isEAPErrorNode(n)) {
         return n.value.event_id === props.event.eventID;
       }
@@ -200,7 +193,7 @@ export function IssuesTraceWaterfall(props: IssuesTraceWaterfallProps) {
     const index = node ? IssuesTraceTree.EnforceVisibility(props.tree, node) : -1;
 
     if (node) {
-      const preserveNodes: Array<TraceTreeNode<TraceTree.NodeValue>> = [node];
+      const preserveNodes: BaseNode[] = [node];
 
       let start = index;
       while (--start > 0) {

static/app/views/performance/newTraceDetails/issuesTraceWaterfallOverlay.tsx

@@ -21,7 +21,7 @@ import type {VirtualizedViewManager} from './traceRenderers/virtualizedViewManag
 interface RowPosition {
   height: number;
   left: number;
-  pathToNode: ReturnType<typeof TraceTree.PathToNode>;
+  pathToNode: TraceTree.NodePath[];
   top: number;
   width: number;
 }
@@ -90,7 +90,7 @@ export function IssueTraceWaterfallOverlay({
           return;
         }
 
-        const pathToNode = TraceTree.PathToNode(node);
+        const pathToNode = node.pathToNode();
 
         if (!pathToNode) {
           return;
@@ -155,7 +155,7 @@ export function IssueTraceWaterfallOverlay({
 
 export function getTraceLinkForIssue(
   traceTarget: LocationDescriptor,
-  pathToNode?: ReturnType<typeof TraceTree.PathToNode>
+  pathToNode?: TraceTree.NodePath[]
 ) {
   if (typeof traceTarget === 'string') {
     return traceTarget;

static/app/views/performance/newTraceDetails/trace.spec.tsx

@@ -1079,7 +1079,7 @@ describe('trace view', () => {
     });
 
     it('scrolls to sibling autogroup node', async () => {
-      mockQueryString('?node=ag-http0&node=txn-1');
+      mockQueryString('?node=ag-span0&node=txn-1');
 
       const {virtualizedContainer} = await completeTestSetup();
       await within(virtualizedContainer).findAllByText(/Autogrouped/i);