feat(studio): add experiment trace comparison view [ASE-186]

web/packages/studio/src/components/IntakeDetail/TraceDetailView.tsx

@@ -12,7 +12,10 @@ import {
   buildTraceSummaryEntries,
 } from '@studio/components/IntakeDetail/IntakeComponents/traceKeyValues';
 import { TraceSummaryHeader } from '@studio/components/IntakeDetail/TraceDetailSummaryHeader';
-import { TraceSpanAccordions } from '@studio/components/IntakeDetail/TraceSpanAccordions';
+import {
+  TraceSpanAccordions,
+  type SlotHeaderRenderProp,
+} from '@studio/components/IntakeDetail/TraceSpanAccordions';
 import { Loading } from '@studio/components/Layouts/Loading';
 import { NotFound } from '@studio/components/Layouts/NotFound';
 import {
@@ -34,6 +37,22 @@ interface IntakeTraceDetailViewProps {
   parentBreadcrumbs?: BreadcrumbsItemProps[];
   /** When true, shows "Test case: <test_case_id>" as the header instead of "Trace <name>". Falls back to "Trace <name>" when test_case_id is absent. */
   showTestCaseTitle?: boolean;
+  /** When true, this instance does not update the breadcrumb trail. Use in compare mode where the parent controls breadcrumbs. */
+  disableBreadcrumbs?: boolean;
+  /** Currently linked (selected) span id, owned by the caller. Single view binds this to ?spanId=; compare view keeps one slot per column. */
+  linkedSpanId: string | null;
+  /** Called when the span selection changes. */
+  onLinkedSpanIdChange: (spanId: string | null) => void;
+  /** Rendered in the PageHeader's slotActions (e.g. a "Compare to…" selector). */
+  slotPageHeaderActions?: React.ReactNode;
+  /** When true, the PageHeader (trace title) is not rendered. Use in compare columns where the column card header serves that role. */
+  hidePageHeader?: boolean;
+  /** When true, forces list view in the span explorer and hides the Tree/List toggle. Pass in compare mode. */
+  forceListView?: boolean;
+  /** When true, opens the first span accordion by default once spans load. Pass in compare mode. */
+  defaultOpenFirst?: boolean;
+  /** Render prop receiving expandAll/collapseAll; replaces TraceSummaryHeader and is passed to TraceSpanAccordions as slotHeader. */
+  slotSpanHeader?: SlotHeaderRenderProp;
 }
 
 /**
@@ -44,6 +63,14 @@ export const IntakeTraceDetailView: FC<IntakeTraceDetailViewProps> = ({
   traceId,
   parentBreadcrumbs,
   showTestCaseTitle,
+  disableBreadcrumbs = false,
+  linkedSpanId,
+  onLinkedSpanIdChange,
+  slotPageHeaderActions,
+  hidePageHeader = false,
+  forceListView = false,
+  defaultOpenFirst = false,
+  slotSpanHeader,
 }) => {
   const {
     data: trace,
@@ -66,11 +93,12 @@ export const IntakeTraceDetailView: FC<IntakeTraceDetailViewProps> = ({
   );
 
   useEffect(() => {
+    if (disableBreadcrumbs) return;
     const parent = parentBreadcrumbs ?? [
       { slotLabel: 'Intake', href: getIntakeTracesRoute(workspace) },
     ];
     setBreadcrumbs([...parent, { slotLabel: `Trace ${traceBreadcrumbLabel}` }]);
-  }, [setBreadcrumbs, traceBreadcrumbLabel, workspace, parentBreadcrumbs]);
+  }, [disableBreadcrumbs, setBreadcrumbs, traceBreadcrumbLabel, workspace, parentBreadcrumbs]);
 
   if (error?.response?.status === 404) {
     return (
@@ -109,9 +137,19 @@ export const IntakeTraceDetailView: FC<IntakeTraceDetailViewProps> = ({
   return (
     <AccessibleTitle title={title}>
       <Stack gap="density-2xl" padding="density-2xl" className="h-full overflow-auto">
-        <PageHeader className="p-0" slotHeading={title} />
-        <TraceSummaryHeader trace={trace} />
-        <TraceSpanAccordions workspace={workspace} trace={trace} />
+        {!hidePageHeader && (
+          <PageHeader className="p-0" slotHeading={title} slotActions={slotPageHeaderActions} />
+        )}
+        {!slotSpanHeader && <TraceSummaryHeader trace={trace} />}
+        <TraceSpanAccordions
+          workspace={workspace}
+          trace={trace}
+          linkedSpanId={linkedSpanId}
+          onLinkedSpanIdChange={onLinkedSpanIdChange}
+          forceListView={forceListView}
+          defaultOpenFirst={defaultOpenFirst}
+          slotHeader={slotSpanHeader}
+        />
         <IntakeAccordion
           variant="section"
           defaultValue={[]}

web/packages/studio/src/components/IntakeDetail/TraceSpanAccordions.tsx

@@ -32,7 +32,6 @@ import {
   spanAccordionId,
 } from '@studio/components/IntakeDetail/traceSpanShared';
 import { SpanTreeView } from '@studio/components/IntakeDetail/TraceSpanTreeView';
-import { QUERY_PARAMETERS } from '@studio/routes/constants';
 import {
   buildSpanHierarchyRows,
   buildSpanTree,
@@ -41,25 +40,57 @@ import {
 import { useQueryClient } from '@tanstack/react-query';
 import { ChevronsDownUp, ChevronsUpDown } from 'lucide-react';
 import { type FC, useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import { useSearchParams } from 'react-router-dom';
+
+/** Render prop passed to TraceSpanAccordions for injecting a custom header above the span list. */
+export type SlotHeaderRenderProp = (props: {
+  expandAll: () => void;
+  collapseAll: () => void;
+}) => React.ReactNode;
 
 const TRACE_SPANS_PAGE_SIZE = 1000;
 
+/** Params for the trace's span-list query; shared with the reload invalidation so the keys always match. */
+const traceSpansParams = (traceId: string) =>
+  ({
+    filter: { trace_id: traceId },
+    mode: 'summary',
+    page: 1,
+    page_size: TRACE_SPANS_PAGE_SIZE,
+    sort: 'started_at',
+  }) as const;
+
 type ViewMode = 'tree' | 'list';
 
 // ── Explorer: toolbar (Tree/List + expand/collapse) over the chosen view ─────
 
 interface TraceSpanAccordionsProps {
   workspace: string;
   trace: Trace;
+  /** Currently linked (selected) span id. Owned by the parent — single view binds it to ?spanId=, compare view keeps one slot per column. */
+  linkedSpanId: string | null;
+  /** Called when the selection changes (span click, accordion open, add-note, session reload). */
+  onLinkedSpanIdChange: (spanId: string | null) => void;
+  /** When true, forces list view and hides the Tree/List toggle. Used in compare mode where columns don't have room for the tree layout. */
+  forceListView?: boolean;
+  /** When true, the first span accordion is opened by default once spans load. Used in compare mode. */
+  defaultOpenFirst?: boolean;
+  /** Render prop called with expandAll/collapseAll; rendered above the span list. When provided, the built-in toolbar is suppressed. */
+  slotHeader?: SlotHeaderRenderProp;
 }
 
-export const TraceSpanAccordions: FC<TraceSpanAccordionsProps> = ({ workspace, trace }) => {
+export const TraceSpanAccordions: FC<TraceSpanAccordionsProps> = ({
+  workspace,
+  trace,
+  linkedSpanId,
+  onLinkedSpanIdChange,
+  forceListView = false,
+  defaultOpenFirst = false,
+  slotHeader,
+}) => {
   const queryClient = useQueryClient();
-  const [searchParams, setSearchParams] = useSearchParams();
-  const linkedSpanId = searchParams.get(QUERY_PARAMETERS.spanId) || null;
-  const [viewMode, setViewMode] = useState<ViewMode>('tree');
+  const [viewMode, setViewMode] = useState<ViewMode>(forceListView ? 'list' : 'tree');
   const [openSpanIds, setOpenSpanIds] = useState<string[]>([]);
+  const didOpenFirstRef = useRef(false);
   // Bumped to broadcast expand/collapse-all to the selected span's sections in
   // tree view (list view drives the span rows via `openSpanIds` instead).
   const [sectionExpandToken, setSectionExpandToken] = useState(0);
@@ -75,17 +106,21 @@ export const TraceSpanAccordions: FC<TraceSpanAccordionsProps> = ({ workspace, t
     data: spansResponse,
     isFetching,
     error,
-  } = useListSpans(workspace, {
-    filter: { trace_id: trace.id },
-    mode: 'summary',
-    page: 1,
-    page_size: TRACE_SPANS_PAGE_SIZE,
-    sort: 'started_at',
-  });
+  } = useListSpans(workspace, traceSpansParams(trace.id));
 
   const spans = spansResponse?.data;
   const spanRows = useMemo(() => buildSpanHierarchyRows(spans ?? []), [spans]);
   const spanTree = useMemo(() => buildSpanTree(spans ?? []), [spans]);
+
+  // In compare mode, open the first span accordion once spans arrive so the
+  // content is immediately visible without a manual click.
+  useEffect(() => {
+    if (!defaultOpenFirst || didOpenFirstRef.current) return;
+    const firstId = spanRows[0]?.span_id;
+    if (!firstId) return;
+    didOpenFirstRef.current = true;
+    setOpenSpanIds([firstId]);
+  }, [defaultOpenFirst, spanRows]);
   const sessionDurationMs = useMemo(
     () => trace.duration_ms ?? getSpansDurationMs(spans ?? []),
     [trace.duration_ms, spans]
@@ -158,20 +193,7 @@ export const TraceSpanAccordions: FC<TraceSpanAccordionsProps> = ({ workspace, t
     return { feedbackBySpan: feedback, annotationCountBySpan: counts, notesBySpan: notes };
   }, [annotationsResponse]);
 
-  const updateLinkedSpanId = useCallback(
-    (spanId: string | null) => {
-      setSearchParams(
-        (prev) => {
-          const next = new URLSearchParams(prev);
-          if (spanId) next.set(QUERY_PARAMETERS.spanId, spanId);
-          else next.delete(QUERY_PARAMETERS.spanId);
-          return next;
-        },
-        { replace: true }
-      );
-    },
-    [setSearchParams]
-  );
+  const updateLinkedSpanId = onLinkedSpanIdChange;
 
   const handleSelectSpan = useCallback(
     (spanId: string) => {
@@ -246,7 +268,11 @@ export const TraceSpanAccordions: FC<TraceSpanAccordionsProps> = ({ workspace, t
     updateLinkedSpanId(null);
     setOpenSpanIds([]);
     void queryClient.invalidateQueries({ queryKey: getGetTraceQueryKey(workspace, trace.id) });
-    void queryClient.invalidateQueries({ queryKey: getListSpansQueryKey(workspace) });
+    // Scoped to this trace's spans query — a workspace-wide key would refetch
+    // every column in compare mode.
+    void queryClient.invalidateQueries({
+      queryKey: getListSpansQueryKey(workspace, traceSpansParams(trace.id)),
+    });
     window.scrollTo({ top: 0, behavior: 'smooth' });
   }, [queryClient, updateLinkedSpanId, workspace, trace.id]);
 
@@ -275,7 +301,14 @@ export const TraceSpanAccordions: FC<TraceSpanAccordionsProps> = ({ workspace, t
         ) : linkedSpanError ? (
           <ErrorMessage message={getErrorMessage(linkedSpanError)} />
         ) : !linkedSpanMatchesTrace ? (
-          <ErrorMessage message="The linked span does not belong to this trace." />
+          <ErrorMessage
+            message="The linked span does not belong to this trace."
+            slotFooter={
+              <Button kind="secondary" onClick={handleReloadSession}>
+                Refresh
+              </Button>
+            }
+          />
         ) : null}
       </div>
     ) : undefined;
@@ -286,42 +319,47 @@ export const TraceSpanAccordions: FC<TraceSpanAccordionsProps> = ({ workspace, t
 
   return (
     <Stack gap="density-lg" className="min-w-0">
-      <Flex align="center" justify="between" gap="density-lg" className="min-w-0">
-        <SegmentedControl
-          size="tiny"
-          value={viewMode}
-          onValueChange={handleViewModeChange}
-          items={[
-            { value: 'tree', children: 'Tree' },
-            { value: 'list', children: 'List' },
-          ]}
-        />
-        {spanRows.length > 0 && (
-          <Flex align="center" gap="density-xs">
-            <Button
-              kind="tertiary"
-              size="tiny"
-              type="button"
-              aria-label="Collapse all"
-              title="Collapse all"
-              onClick={collapseAll}
-            >
-              <ChevronsDownUp size={14} aria-hidden />
-            </Button>
-            <Button
-              kind="tertiary"
+      {slotHeader ? (
+        slotHeader({ expandAll, collapseAll })
+      ) : (
+        <Flex align="center" justify="between" gap="density-lg" className="min-w-0">
+          {!forceListView && (
+            <SegmentedControl
               size="tiny"
-              type="button"
-              aria-label="Expand all"
-              title="Expand all"
-              onClick={expandAll}
-            >
-              <ChevronsUpDown size={14} aria-hidden />
-            </Button>
-          </Flex>
-        )}
-      </Flex>
-
+              value={viewMode}
+              onValueChange={handleViewModeChange}
+              items={[
+                { value: 'tree', children: 'Tree' },
+                { value: 'list', children: 'List' },
+              ]}
+            />
+          )}
+          {spanRows.length > 0 && (
+            <Flex align="center" gap="density-xs">
+              <Button
+                kind="tertiary"
+                size="tiny"
+                type="button"
+                aria-label="Collapse all"
+                title="Collapse all"
+                onClick={collapseAll}
+              >
+                <ChevronsDownUp size={14} aria-hidden />
+              </Button>
+              <Button
+                kind="tertiary"
+                size="tiny"
+                type="button"
+                aria-label="Expand all"
+                title="Expand all"
+                onClick={expandAll}
+              >
+                <ChevronsUpDown size={14} aria-hidden />
+              </Button>
+            </Flex>
+          )}
+        </Flex>
+      )}
       {showSpanLimitMessage && (
         <Text kind="body/regular/sm" className="text-secondary">
           Showing first {TRACE_SPANS_PAGE_SIZE.toLocaleString()} of{' '}