feat(trace-tree-node): Replacing TraceTreeNode with BaseNode

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

@@ -12,8 +12,6 @@ import {IconTool} from 'sentry/icons/iconTool';
 import {space} from 'sentry/styles/space';
 import getDuration from 'sentry/utils/duration/getDuration';
 import {LLMCosts} from 'sentry/views/insights/agents/components/llmCosts';
-import {getIsAiRunNode} from 'sentry/views/insights/agents/utils/aiTraceNodes';
-import {getNodeId} from 'sentry/views/insights/agents/utils/getNodeId';
 import {
   getIsAiCreateAgentSpan,
   getIsAiGenerationSpan,
@@ -29,8 +27,9 @@ 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 {BaseNode} from 'sentry/views/performance/newTraceDetails/traceModels/traceTreeNode/baseNode';
+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;
@@ -52,11 +51,9 @@ function getNodeTimeBounds(node: AITraceSpanNode | AITraceSpanNode[]) {
   } else {
     if (!node.value) return {startTime: 0, endTime: 0, duration: 0};
 
-    startTime = node.value.start_timestamp;
-    if (isTransactionNode(node) || isSpanNode(node)) {
-      endTime = node.value.timestamp;
-    } else if (isEAPSpanNode(node)) {
-      endTime = node.value.end_timestamp;
+    if (node.startTimestamp && node.endTimestamp) {
+      startTime = node.startTimestamp;
+      endTime = node.endTimestamp;
     }
   }
 
@@ -69,16 +66,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 +76,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(isTransactionNodeEquivalent) as
+        | TransactionNode
+        | EapSpanNode;
       if (!transaction) {
         continue;
       }
@@ -107,7 +93,7 @@ export function AISpanList({
   return (
     <TraceListContainer>
       {nodesByTransaction.entries().map(([transaction, transactionNodes]) => (
-        <Fragment key={getNodeId(transaction)}>
+        <Fragment key={transaction.id}>
           <TransactionWrapper
             canCollapse={nodesByTransaction.size > 1}
             transaction={transaction}
@@ -132,19 +118,19 @@ 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();
   const colors = [...theme.chart.getColorPalette(5), theme.red300];
   const timeBounds = getNodeTimeBounds(nodes);
 
-  const nodeAiRunParentsMap = useMemo<Record<string, AITraceSpanNode>>(() => {
-    const parents: Record<string, AITraceSpanNode> = {};
+  const nodeAiRunParentsMap = useMemo<Record<string, BaseNode>>(() => {
+    const parents: Record<string, BaseNode> = {};
     for (const node of nodes) {
-      const parent = getClosestNode(node, getIsAiRunNode);
+      const parent = node.findParent(p => getIsAiRunSpan({op: p.op}));
       if (parent) {
-        parents[getNodeId(node)] = parent;
+        parents[node.id] = parent;
       }
     }
     return parents;
@@ -171,12 +157,12 @@ function TransactionWrapper({
       </TransactionButton>
       {isExpanded &&
         nodes.map(node => {
-          const aiRunNode = nodeAiRunParentsMap[getNodeId(node)];
+          const aiRunNode = nodeAiRunParentsMap[node.id];
 
           // Only indent if the node is not the ai run node
           const shouldIndent = aiRunNode && aiRunNode !== node;
 
-          const uniqueKey = getNodeId(node);
+          const uniqueKey = node.id;
           return (
             <TraceListItem
               indent={shouldIndent ? 1 : 0}
@@ -242,22 +228,16 @@ interface TraceBounds {
 }
 
 function calculateRelativeTiming(
-  node: TraceTreeNode<TraceTree.NodeValue>,
+  node: BaseNode,
   traceBounds: TraceBounds
 ): {leftPercent: number; widthPercent: number} {
   if (!node.value) return {leftPercent: 0, widthPercent: 0};
 
   let startTime: number, endTime: number;
 
-  if (isTransactionNode(node)) {
-    startTime = node.value.start_timestamp;
-    endTime = node.value.timestamp;
-  } else if (isSpanNode(node)) {
-    startTime = node.value.start_timestamp;
-    endTime = node.value.timestamp;
-  } else if (isEAPSpanNode(node)) {
-    startTime = node.value.start_timestamp;
-    endTime = node.value.end_timestamp;
+  if (node.startTimestamp && node.endTimestamp) {
+    startTime = node.startTimestamp;
+    endTime = node.endTimestamp;
   } else {
     return {leftPercent: 0, widthPercent: 0};
   }

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

@@ -16,7 +16,6 @@ import {useLocationSyncedState} from 'sentry/views/insights/agents/hooks/useLoca
 import {useNodeDetailsLink} from 'sentry/views/insights/agents/hooks/useNodeDetailsLink';
 import {useUrlTraceDrawer} from 'sentry/views/insights/agents/hooks/useUrlTraceDrawer';
 import {getDefaultSelectedNode} from 'sentry/views/insights/agents/utils/getDefaultSelectedNode';
-import {getNodeId} from 'sentry/views/insights/agents/utils/getNodeId';
 import type {AITraceSpanNode} from 'sentry/views/insights/agents/utils/types';
 import {DrawerUrlParams} from 'sentry/views/insights/agents/utils/urlParams';
 import {TraceTreeNodeDetails} from 'sentry/views/performance/newTraceDetails/traceDrawer/tabs/traceTreeNodeDetails';
@@ -53,7 +52,7 @@ const TraceViewDrawer = memo(function TraceViewDrawer({
 
   const handleSelectNode = useCallback(
     (node: AITraceSpanNode) => {
-      const uniqueKey = getNodeId(node);
+      const uniqueKey = node.id;
       setSelectedNodeKey(uniqueKey);
 
       trackAnalytics('agent-monitoring.drawer.span-select', {
@@ -64,7 +63,7 @@ const TraceViewDrawer = memo(function TraceViewDrawer({
   );
 
   const selectedNode =
-    (selectedNodeKey && nodes.find(node => getNodeId(node) === selectedNodeKey)) ||
+    (selectedNodeKey && nodes.find(node => node.id === selectedNodeKey)) ||
     getDefaultSelectedNode(nodes);
 
   const nodeDetailsLink = useNodeDetailsLink({
@@ -183,7 +182,7 @@ function AITraceView({
         <SpansHeader>{t('AI Spans')}</SpansHeader>
         <AISpanList
           nodes={nodes}
-          selectedNodeKey={getNodeId(selectedNode!)}
+          selectedNodeKey={selectedNode!.id}
           onSelectNode={onSelectNode}
         />
       </LeftPanel>

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

@@ -2,18 +2,16 @@ import {useEffect, useState} from 'react';
 
 import useApi from 'sentry/utils/useApi';
 import useOrganization from 'sentry/utils/useOrganization';
-import {getIsAiNode} from 'sentry/views/insights/agents/utils/aiTraceNodes';
+import {getIsAiSpan} from 'sentry/views/insights/agents/utils/query';
 import type {AITraceSpanNode} from 'sentry/views/insights/agents/utils/types';
 import {SpanFields} from 'sentry/views/insights/types';
 import {useTrace} from 'sentry/views/performance/newTraceDetails/traceApi/useTrace';
 import {
   isEAPSpanNode,
   isSpanNode,
-  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 {DEFAULT_TRACE_VIEW_PREFERENCES} from 'sentry/views/performance/newTraceDetails/traceState/tracePreferences';
 import {useTraceQueryParams} from 'sentry/views/performance/newTraceDetails/useTraceQueryParams';
 
@@ -69,29 +67,19 @@ 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 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, {
+        const fetchableNodes = tree.root.findAllChildren(node => node.canFetchChildren);
+        const fetchPromises = fetchableNodes.map(node =>
+          tree.fetchNodeSubTree(true, node, {
             api,
             organization,
             preferences: DEFAULT_TRACE_VIEW_PREFERENCES,
           })
         );
 
-        await Promise.all(zoomPromises);
+        await Promise.all(fetchPromises);
 
         // Keep only transactions that include AI spans and the AI spans themselves
-        const flattenedNodes = TraceTree.FindAll(tree.root, node => {
+        const flattenedNodes = tree.root.findAllChildren(node => {
           if (
             !isTransactionNodeEquivalent(node) &&
             !isSpanNode(node) &&
@@ -100,7 +88,7 @@ export function useAITrace(traceSlug: string): UseAITraceResult {
             return false;
           }
 
-          return getIsAiNode(node);
+          return getIsAiSpan({op: node.op});
         }) as AITraceSpanNode[];
 
         setNodes(flattenedNodes);

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

@@ -1,20 +1,12 @@
 import type {EventTransaction} from 'sentry/types/event';
 import {prettifyAttributeName} from 'sentry/views/explore/components/traceItemAttributes/utils';
 import type {TraceItemResponseAttribute} from 'sentry/views/explore/hooks/useTraceItemDetails';
-import {
-  getIsAiGenerationSpan,
-  getIsAiRunSpan,
-  getIsAiSpan,
-  getIsExecuteToolSpan,
-} from 'sentry/views/insights/agents/utils/query';
-import type {AITraceSpanNode} from 'sentry/views/insights/agents/utils/types';
 import {
   isEAPSpanNode,
   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 +29,7 @@ function getAttributeValue(
 }
 
 export function ensureAttributeObject(
-  node: TraceTreeNode<TraceTree.NodeValue>,
+  node: BaseNode,
   event?: EventTransaction,
   attributes?: TraceItemResponseAttribute[]
 ) {
@@ -66,26 +58,10 @@ export function ensureAttributeObject(
 
 export function getTraceNodeAttribute(
   name: string,
-  node: TraceTreeNode<TraceTree.NodeValue>,
+  node: BaseNode,
   event?: EventTransaction,
   attributes?: TraceItemResponseAttribute[]
 ): string | number | boolean | undefined {
   const attributeObject = ensureAttributeObject(node, event, attributes);
   return attributeObject?.[name];
 }
-
-function createGetIsAiNode(predicate: ({op}: {op?: string}) => boolean) {
-  return (node: TraceTreeNode<TraceTree.NodeValue>): node is AITraceSpanNode => {
-    if (!isTransactionNode(node) && !isSpanNode(node) && !isEAPSpanNode(node)) {
-      return false;
-    }
-
-    const op = isTransactionNode(node) ? node.value?.['transaction.op'] : node.value?.op;
-    return predicate({op});
-  };
-}
-
-export const getIsAiNode = createGetIsAiNode(getIsAiSpan);
-export const getIsAiRunNode = createGetIsAiNode(getIsAiRunSpan);
-export const getIsAiGenerationNode = createGetIsAiNode(getIsAiGenerationSpan);
-export const getIsExecuteToolNode = createGetIsAiNode(getIsExecuteToolSpan);

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

@@ -1,7 +1,8 @@
-import {getIsAiNode} from 'sentry/views/insights/agents/utils/aiTraceNodes';
 import type {AITraceSpanNode} from 'sentry/views/insights/agents/utils/types';
 
+import {getIsAiSpan} from './query';
+
 export function getDefaultSelectedNode(nodes: AITraceSpanNode[]) {
-  const firstAiSpan = nodes.find(getIsAiNode);
+  const firstAiSpan = nodes.find(node => getIsAiSpan({op: node.op}));
   return firstAiSpan ?? nodes[0];
 }

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 = SpanNode | EapSpanNode | TransactionNode;

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

@@ -3,14 +3,12 @@ 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;
   }
 
-  const op = isTransactionNode(node) ? node.value['transaction.op'] : node.value.op;
-  return op?.startsWith('mcp.');
+  return node.op?.startsWith('mcp.');
 }