address comments

Signed-off-by: Joshua Li <joshuali925@gmail.com>

Author: joshuali925

src/plugins/agent_traces/public/application/pages/traces/flow/node_types.test.ts

@@ -14,7 +14,9 @@ describe('nodeTypes', () => {
     expect(nodeTypes.agent).toBe('MockSpanNode');
     expect(nodeTypes.llm).toBe('MockSpanNode');
     expect(nodeTypes.tool).toBe('MockSpanNode');
-    expect(nodeTypes.error).toBe('MockSpanNode');
+    expect(nodeTypes.content).toBe('MockSpanNode');
+    expect(nodeTypes.embeddings).toBe('MockSpanNode');
+    expect(nodeTypes.retrieval).toBe('MockSpanNode');
     expect(nodeTypes.other).toBe('MockSpanNode');
   });
 });

src/plugins/agent_traces/public/application/pages/traces/flow/node_types.ts

@@ -14,6 +14,5 @@ export const nodeTypes: NodeTypes = {
   content: SpanNode,
   embeddings: SpanNode,
   retrieval: SpanNode,
-  error: SpanNode,
   other: SpanNode,
 };

src/plugins/agent_traces/public/application/pages/traces/flow/span_node.test.tsx

@@ -65,7 +65,7 @@ describe('SpanNode', () => {
 
   it('renders category badge', () => {
     render(<SpanNode data={{ span: makeSpan(), totalDuration: 1000 }} />);
-    expect(screen.getByText('AGENT')).toBeInTheDocument();
+    expect(screen.getByText('Agent')).toBeInTheDocument();
   });
 
   it('renders latency', () => {

src/plugins/agent_traces/public/application/pages/traces/flow/span_node.tsx

@@ -13,7 +13,7 @@ import {
   hexToRgba,
   SpanCategory,
 } from '../../../../services/span_categorization';
-import { parseLatencyMs } from '../flyout/tree_helpers';
+import { parseLatencyMs } from '../trace_details/utils/span_timerange_utils';
 import './span_node.scss';
 
 interface SpanNodeProps {

src/plugins/agent_traces/public/application/pages/traces/flyout/timeline_gantt.tsx

@@ -17,6 +17,7 @@ import {
 import { FormattedMessage } from '@osd/i18n/react';
 import { euiThemeVars } from '@osd/ui-shared-deps/theme';
 import { getCategoryMeta, hexToRgba } from '../../../../services/span_categorization';
+import { formatMs } from '../trace_details/utils/span_timerange_utils';
 import { TimelineSpan } from './tree_helpers';
 import './timeline_gantt.scss';
 
@@ -31,11 +32,6 @@ export interface TimelineGanttProps {
   onToggleExpanded: (nodeId: string) => void;
 }
 
-const formatRelativeMs = (ms: number): string => {
-  if (ms >= 1000) return `${(ms / 1000).toFixed(2)} s`;
-  return `${ms.toFixed(2)} ms`;
-};
-
 export const TimelineGantt: React.FC<TimelineGanttProps> = ({
   timelineVisibleSpans,
   timelineRange,
@@ -173,15 +169,13 @@ export const TimelineGantt: React.FC<TimelineGanttProps> = ({
                           <div style={{ lineHeight: 1.6 }}>
                             <div>
                               <strong>Duration:</strong>{' '}
-                              {span.node.latency || formatRelativeMs(span.durationMs)}
+                              {span.node.latency || formatMs(span.durationMs)}
                             </div>
                             <div>
-                              <strong>Start:</strong>{' '}
-                              {formatRelativeMs(span.startMs - timelineRange.minMs)}
+                              <strong>Start:</strong> {formatMs(span.startMs - timelineRange.minMs)}
                             </div>
                             <div>
-                              <strong>End:</strong>{' '}
-                              {formatRelativeMs(span.endMs - timelineRange.minMs)}
+                              <strong>End:</strong> {formatMs(span.endMs - timelineRange.minMs)}
                             </div>
                           </div>
                         }

src/plugins/agent_traces/public/application/pages/traces/flyout/trace_details_flyout.tsx

@@ -22,13 +22,13 @@ import {
 } from '@elastic/eui';
 import { TraceRow } from '../hooks/use_agent_traces';
 import { TraceFlowView } from '../flow/trace_flow_view';
+import { parseLatencyMs } from '../trace_details/utils/span_timerange_utils';
 import {
   TreeNode,
   buildTreeFromTraceRow,
   flattenTree,
   countSpans,
   sumTokens,
-  parseLatencyMs,
   flattenVisibleNodes,
   calculateTimelineRange,
 } from './tree_helpers';

src/plugins/agent_traces/public/application/pages/traces/flyout/tree_helpers.test.ts

@@ -10,7 +10,7 @@ import {
   countSpans,
   sumTokens,
   parseLatencyMs,
-  parseRawTimestampMs,
+  parseTimestampMs,
   extractTimestamps,
   flattenVisibleNodes,
   calculateTimelineRange,
@@ -128,21 +128,21 @@ describe('tree_helpers', () => {
     });
   });
 
-  describe('parseRawTimestampMs', () => {
+  describe('parseTimestampMs', () => {
     it('parses ISO timestamp', () => {
-      const ms = parseRawTimestampMs('2025-01-01T00:00:00Z');
+      const ms = parseTimestampMs('2025-01-01T00:00:00Z');
       expect(ms).toBe(new Date('2025-01-01T00:00:00Z').getTime());
     });
 
     it('normalizes space-separated timestamps', () => {
-      const ms = parseRawTimestampMs('2025-01-01 12:00:00');
+      const ms = parseTimestampMs('2025-01-01 12:00:00');
       expect(ms).toBeGreaterThan(0);
     });
 
     it('returns 0 for invalid input', () => {
-      expect(parseRawTimestampMs(null)).toBe(0);
-      expect(parseRawTimestampMs('')).toBe(0);
-      expect(parseRawTimestampMs(123)).toBe(0);
+      expect(parseTimestampMs(null)).toBe(0);
+      expect(parseTimestampMs('')).toBe(0);
+      expect(parseTimestampMs(123)).toBe(0);
     });
   });
 

src/plugins/agent_traces/public/application/pages/traces/flyout/tree_helpers.ts

@@ -4,12 +4,15 @@
  */
 
 import { TraceRow } from '../hooks/use_agent_traces';
+import { parseLatencyMs, parseTimestampMs } from '../trace_details/utils/span_timerange_utils';
 import {
   getSpanCategory,
   getCategoryMeta,
   SpanCategory,
 } from '../../../../services/span_categorization';
 
+export { parseLatencyMs, parseTimestampMs };
+
 export interface TreeNode {
   label: string;
   id: string;
@@ -81,38 +84,18 @@ export const sumTokens = (nodes: TreeNode[]): number => {
   return total;
 };
 
-export const parseLatencyMs = (latency?: string): number => {
-  if (!latency || latency === '—') return 0;
-  if (latency.endsWith('ms')) return parseFloat(latency) || 0;
-  if (latency.endsWith('s')) return (parseFloat(latency) || 0) * 1000;
-  return 0;
-};
-
-export const parseRawTimestampMs = (ts: unknown): number => {
-  if (!ts || typeof ts !== 'string') return 0;
-  let normalized = ts;
-  if (ts.includes(' ') && !ts.includes('T')) {
-    normalized = ts.replace(' ', 'T');
-    if (!normalized.includes('Z') && !normalized.includes('+')) {
-      normalized += 'Z';
-    }
-  }
-  const ms = new Date(normalized).getTime();
-  return isNaN(ms) ? 0 : ms;
-};
-
 export const extractTimestamps = (node: TreeNode): { startMs: number; endMs: number } => {
   const raw = node.traceRow?.rawDocument;
   if (raw) {
     const rawStart = raw.startTime;
-    const s = parseRawTimestampMs(rawStart);
+    const s = parseTimestampMs(rawStart);
     if (s > 0) {
       // Prefer durationInNanos for sub-ms precision; fall back to endTime.
       const durationNanos = (raw.durationInNanos as number) || 0;
       if (durationNanos > 0) {
         return { startMs: s, endMs: s + durationNanos / 1_000_000 };
       }
-      const e = parseRawTimestampMs(raw.endTime);
+      const e = parseTimestampMs(raw.endTime);
       if (e > 0) return { startMs: s, endMs: e };
     }
   }

src/plugins/agent_traces/public/application/pages/traces/hooks/index.ts

@@ -9,4 +9,9 @@ export { useTraceMetrics, useTraceMetricsContext, TraceMetricsContext } from './
 export type { TraceMetrics } from './use_trace_metrics';
 export { formatDuration, traceHitToAgentSpan } from './span_transforms';
 export type { AgentSpan } from './span_transforms';
+export {
+  formatMs,
+  parseLatencyMs,
+  parseTimestampMs,
+} from '../trace_details/utils/span_timerange_utils';
 export type { BaseRow, LoadingState } from './tree_utils';

src/plugins/agent_traces/public/application/pages/traces/hooks/tree_utils.test.ts

@@ -3,7 +3,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { setLevels, spanToRow, buildFullSpanTree, hitsToAgentSpans, BaseRow } from './tree_utils';
+import {
+  formatTimestamp,
+  setLevels,
+  spanToRow,
+  buildFullSpanTree,
+  hitsToAgentSpans,
+} from './tree_utils';
 import { AgentSpan } from './span_transforms';
 
 const makeSpan = (overrides: Partial<AgentSpan> = {}): AgentSpan => ({
@@ -31,11 +37,29 @@ const makeSpan = (overrides: Partial<AgentSpan> = {}): AgentSpan => ({
 });
 
 describe('tree_utils', () => {
+  describe('formatTimestamp', () => {
+    it('formats UTC timestamp in the given timezone', () => {
+      const result = formatTimestamp('2025-01-01 12:00:00', 'America/New_York');
+      expect(result).toBe('01/01/2025, 7:00:00.000 AM');
+    });
+
+    it('returns dash for empty timestamp', () => {
+      expect(formatTimestamp('', 'UTC')).toBe('—');
+    });
+
+    it('returns dash for invalid timestamp', () => {
+      expect(formatTimestamp('not-a-date', 'UTC')).toBe('—');
+    });
+
+    it('preserves millisecond precision', () => {
+      const result = formatTimestamp('2025-06-15 08:30:45.789', 'UTC');
+      expect(result).toContain('.789');
+    });
+  });
+
   describe('setLevels', () => {
     it('sets levels recursively', () => {
-      const rows: Array<{ level?: number; children?: any[] }> = [
-        { children: [{ children: [{}] }] },
-      ];
+      const rows: any[] = [{ children: [{ children: [{}] }] }];
       setLevels(rows, 0);
       expect(rows[0].level).toBe(0);
       expect(rows[0].children![0].level).toBe(1);
@@ -74,7 +98,7 @@ describe('tree_utils', () => {
         makeSpan({ spanId: 'root', parentSpanId: null }),
         makeSpan({ spanId: 'child', parentSpanId: 'root' }),
       ];
-      const tree = buildFullSpanTree<BaseRow>(spans, (ts) => ts);
+      const tree = buildFullSpanTree(spans, (ts) => ts);
 
       expect(tree).toHaveLength(1);
       expect(tree[0].spanId).toBe('root');
@@ -87,11 +111,65 @@ describe('tree_utils', () => {
         makeSpan({ spanId: 'root', parentSpanId: null }),
         makeSpan({ spanId: 'child', parentSpanId: 'root' }),
       ];
-      const tree = buildFullSpanTree<BaseRow>(spans, (ts) => ts);
+      const tree = buildFullSpanTree(spans, (ts) => ts);
 
       expect(tree[0].level).toBe(0);
       expect(tree[0].children![0].level).toBe(1);
     });
+
+    it('sorts children earliest-first using rawDocument.startTime', () => {
+      const spans = [
+        makeSpan({ spanId: 'root', parentSpanId: null }),
+        makeSpan({
+          spanId: 'late',
+          parentSpanId: 'root',
+          rawDocument: { startTime: '2025-01-01 00:00:10' },
+        }),
+        makeSpan({
+          spanId: 'early',
+          parentSpanId: 'root',
+          rawDocument: { startTime: '2025-01-01 00:00:01' },
+        }),
+        makeSpan({
+          spanId: 'mid',
+          parentSpanId: 'root',
+          rawDocument: { startTime: '2025-01-01 00:00:05' },
+        }),
+      ];
+      const tree = buildFullSpanTree(spans, (ts) => ts);
+      const childIds = tree[0].children!.map((c) => c.spanId);
+      expect(childIds).toEqual(['early', 'mid', 'late']);
+    });
+
+    it('handles spans without rawDocument.startTime in sort', () => {
+      const spans = [
+        makeSpan({ spanId: 'root', parentSpanId: null }),
+        makeSpan({
+          spanId: 'has-time',
+          parentSpanId: 'root',
+          rawDocument: { startTime: '2025-01-01 00:00:01' },
+        }),
+        makeSpan({
+          spanId: 'no-time',
+          parentSpanId: 'root',
+          rawDocument: {},
+        }),
+      ];
+      const tree = buildFullSpanTree(spans, (ts) => ts);
+      // Should not throw; span without startTime sorts to beginning (time 0)
+      expect(tree[0].children).toHaveLength(2);
+      expect(tree[0].children![0].spanId).toBe('no-time');
+      expect(tree[0].children![1].spanId).toBe('has-time');
+    });
+
+    it('treats orphan spans as roots', () => {
+      const spans = [
+        makeSpan({ spanId: 'a', parentSpanId: 'missing' }),
+        makeSpan({ spanId: 'b', parentSpanId: null }),
+      ];
+      const tree = buildFullSpanTree(spans, (ts) => ts);
+      expect(tree).toHaveLength(2);
+    });
   });
 
   describe('hitsToAgentSpans', () => {