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

packages/app/src/components/DBTraceWaterfallChart.tsx

@@ -17,7 +17,7 @@ import {
   Anchor,
   Box,
   Center,
-  Checkbox,
+  Chip,
   Code,
   Divider,
   Group,
@@ -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,12 +1022,51 @@ export function DBTraceWaterfallChartContainer({
               {errorCountString}
             </span>
           </Text>
-          <Checkbox
-            size="xs"
-            label="Show span events"
-            checked={showSpanEvents}
-            onChange={() => setShowSpanEvents(!showSpanEvents)}
-          />
+          <Group gap="xs" align="center">
+            <Text size="xs" c="dimmed">
+              Show:
+            </Text>
+            <Group gap={4}>
+              <Chip
+                size="xs"
+                color="gray"
+                checked={showSpans}
+                onChange={() => setShowSpans(!showSpans)}
+                data-testid="show-spans-chip"
+                styles={{
+                  label: { paddingInline: 8, height: 22, minHeight: 22 },
+                }}
+              >
+                Spans
+              </Chip>
+              {logTableSource && (
+                <Chip
+                  size="xs"
+                  color="gray"
+                  checked={showLogs}
+                  onChange={() => setShowLogs(!showLogs)}
+                  data-testid="show-logs-chip"
+                  styles={{
+                    label: { paddingInline: 8, height: 22, minHeight: 22 },
+                  }}
+                >
+                  Logs
+                </Chip>
+              )}
+              <Chip
+                size="xs"
+                color="gray"
+                checked={showSpanEvents}
+                onChange={() => setShowSpanEvents(!showSpanEvents)}
+                data-testid="show-span-events-chip"
+                styles={{
+                  label: { paddingInline: 8, height: 22, minHeight: 22 },
+                }}
+              >
+                Span events
+              </Chip>
+            </Group>
+          </Group>
         </Group>
         <span>
           <Anchor
@@ -1071,10 +1119,14 @@ export function DBTraceWaterfallChartContainer({
           <div>
             An unknown error occurred. <ContactSupportText />
           </div>
-        ) : flattenedNodes.length === 0 ? (
-          (emptyState ?? (
-            <div className="my-3">No matching spans or logs found</div>
-          ))
+        ) : visibleNodes.length === 0 ? (
+          flattenedNodes.length > 0 ? (
+            <div className="my-3">All items are hidden by filters</div>
+          ) : (
+            (emptyState ?? (
+              <div className="my-3">No matching spans or logs found</div>
+            ))
+          )
         ) : (
           <TimelineChart
             maxHeight={heightPx}

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 Spans and Logs chips when log source is present', async () => {
+    setupQueryMocks({ traceData: mockTraceData, logData: mockLogData });
+    renderComponent();
+    await waitForLoading();
+
+    expect(screen.getByTestId('show-spans-chip')).toBeInTheDocument();
+    expect(screen.getByTestId('show-logs-chip')).toBeInTheDocument();
+  });
+
+  it('does not render Logs chip when no log source', async () => {
+    setupQueryMocks({ traceData: mockTraceData });
+    renderComponent(null);
+    await waitForLoading();
+
+    expect(screen.getByTestId('show-spans-chip')).toBeInTheDocument();
+    expect(screen.queryByTestId('show-logs-chip')).not.toBeInTheDocument();
+  });
+
+  it('hides log rows when Logs chip is toggled off', async () => {
+    const user = userEvent.setup();
+    setupQueryMocks({ traceData: mockTraceData, logData: mockLogData });
+    renderComponent();
+    await waitForLoading();
+
+    expect(MockTimelineChart.latestProps.rows.length).toBe(2);
+
+    const showLogsChip = screen.getByTestId('show-logs-chip');
+    await user.click(showLogsChip);
+
+    await waitFor(() => {
+      expect(MockTimelineChart.latestProps.rows.length).toBe(1);
+    });
+  });
+
+  it('hides span rows when Spans chip is toggled off', async () => {
+    const user = userEvent.setup();
+    setupQueryMocks({ traceData: mockTraceData, logData: mockLogData });
+    renderComponent();
+    await waitForLoading();
+
+    expect(MockTimelineChart.latestProps.rows.length).toBe(2);
+
+    const showSpansChip = screen.getByTestId('show-spans-chip');
+    await user.click(showSpansChip);
+
+    await waitFor(() => {
+      expect(MockTimelineChart.latestProps.rows.length).toBe(1);
+    });
+  });
 });
 
 describe('useEventsAroundFocus', () => {