Add Show spans/Show logs visibility toggles to trace waterfall panel

Users can now independently show/hide spans and log events in the trace
waterfall using dedicated checkboxes. The 'Show spans' checkbox is always
visible; the 'Show logs' checkbox appears when a log source is configured.

Filtering is applied client-side after the DAG flattening step, so the
toggle state has no effect on data fetching. Counts in the header update
to reflect only visible rows.

Co-authored-by: Mike Shi <mike@hyperdx.io>

Author: cursoragent

packages/app/src/components/DBTraceWaterfallChart.tsx

@@ -610,6 +610,8 @@ export function DBTraceWaterfallChartContainer({
 
   const [collapsedIds, setCollapsedIds] = useState<Set<string>>(new Set());
   const [showSpanEvents, setShowSpanEvents] = useState(true);
+  const [showSpans, setShowSpans] = useState(true);
+  const [showLogs, setShowLogs] = useState(true);
 
   const { nodesMap, flattenedNodes } = useMemo(() => {
     const rootNodes: Node[] = [];
@@ -737,8 +739,16 @@ export function DBTraceWaterfallChartContainer({
     [nodesMap],
   );
 
-  const spanCount = flattenedNodes.length;
-  const errorCount = flattenedNodes.filter(
+  const visibleNodes = useMemo(() => {
+    if (showSpans && showLogs) return flattenedNodes;
+    return flattenedNodes.filter(node => {
+      if (node.type === SourceKind.Log) return showLogs;
+      return showSpans;
+    });
+  }, [flattenedNodes, showSpans, showLogs]);
+
+  const spanCount = visibleNodes.length;
+  const errorCount = visibleNodes.filter(
     node =>
       node.StatusCode === 'Error' ||
       node.SeverityText?.toLowerCase() === 'error',
@@ -763,7 +773,7 @@ export function DBTraceWaterfallChartContainer({
 
   const timelineRows = useMemo(
     () =>
-      flattenedNodes.map((result, i) => {
+      visibleNodes.map((result, i) => {
         const tookMs = (result.Duration || 0) * 1000;
         const startOffset = new Date(result.Timestamp).getTime();
         const start = startOffset - minOffset;
@@ -940,7 +950,7 @@ export function DBTraceWaterfallChartContainer({
       }),
     [
       collapsedIds,
-      flattenedNodes,
+      visibleNodes,
       formatTime,
       highlightedRowWhere,
       isFilterActive,
@@ -952,8 +962,7 @@ export function DBTraceWaterfallChartContainer({
       setCollapseTooltipShown,
     ],
   );
-  // TODO: Highlighting support
-  const initialScrollRowIndex = flattenedNodes.findIndex(v => {
+  const initialScrollRowIndex = visibleNodes.findIndex(v => {
     return v.id === highlightedRowWhere;
   });
 
@@ -1013,6 +1022,22 @@ export function DBTraceWaterfallChartContainer({
               {errorCountString}
             </span>
           </Text>
+          <Checkbox
+            size="xs"
+            label="Show spans"
+            checked={showSpans}
+            onChange={() => setShowSpans(!showSpans)}
+            data-testid="show-spans-checkbox"
+          />
+          {logTableSource && (
+            <Checkbox
+              size="xs"
+              label="Show logs"
+              checked={showLogs}
+              onChange={() => setShowLogs(!showLogs)}
+              data-testid="show-logs-checkbox"
+            />
+          )}
           <Checkbox
             size="xs"
             label="Show span events"
@@ -1071,7 +1096,7 @@ export function DBTraceWaterfallChartContainer({
           <div>
             An unknown error occurred. <ContactSupportText />
           </div>
-        ) : flattenedNodes.length === 0 ? (
+        ) : visibleNodes.length === 0 ? (
           (emptyState ?? (
             <div className="my-3">No matching spans or logs found</div>
           ))

packages/app/src/components/__tests__/DBTraceWaterfallChart.test.tsx

@@ -6,6 +6,7 @@ import {
 } from '@hyperdx/common-utils/dist/types';
 import { screen, waitFor } from '@testing-library/react';
 import { renderHook } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
 
 import { TimelineChart } from '@/components/TimelineChart';
 import useOffsetPaginatedQuery from '@/hooks/useOffsetPaginatedQuery';
@@ -310,6 +311,56 @@ describe('DBTraceWaterfallChartContainer', () => {
       screen.getByText('http span https://api.example.com/users'),
     ).toBeInTheDocument();
   });
+
+  it('renders Show spans and Show logs checkboxes when log source is present', async () => {
+    setupQueryMocks({ traceData: mockTraceData, logData: mockLogData });
+    renderComponent();
+    await waitForLoading();
+
+    expect(screen.getByTestId('show-spans-checkbox')).toBeInTheDocument();
+    expect(screen.getByTestId('show-logs-checkbox')).toBeInTheDocument();
+  });
+
+  it('does not render Show logs checkbox when no log source', async () => {
+    setupQueryMocks({ traceData: mockTraceData });
+    renderComponent(null);
+    await waitForLoading();
+
+    expect(screen.getByTestId('show-spans-checkbox')).toBeInTheDocument();
+    expect(screen.queryByTestId('show-logs-checkbox')).not.toBeInTheDocument();
+  });
+
+  it('hides log rows when Show logs is unchecked', async () => {
+    const user = userEvent.setup();
+    setupQueryMocks({ traceData: mockTraceData, logData: mockLogData });
+    renderComponent();
+    await waitForLoading();
+
+    expect(MockTimelineChart.latestProps.rows.length).toBe(2);
+
+    const showLogsCheckbox = screen.getByTestId('show-logs-checkbox');
+    await user.click(showLogsCheckbox);
+
+    await waitFor(() => {
+      expect(MockTimelineChart.latestProps.rows.length).toBe(1);
+    });
+  });
+
+  it('hides span rows when Show spans is unchecked', async () => {
+    const user = userEvent.setup();
+    setupQueryMocks({ traceData: mockTraceData, logData: mockLogData });
+    renderComponent();
+    await waitForLoading();
+
+    expect(MockTimelineChart.latestProps.rows.length).toBe(2);
+
+    const showSpansCheckbox = screen.getByTestId('show-spans-checkbox');
+    await user.click(showSpansCheckbox);
+
+    await waitFor(() => {
+      expect(MockTimelineChart.latestProps.rows.length).toBe(1);
+    });
+  });
 });
 
 describe('useEventsAroundFocus', () => {