write top span id on trace and re-query by it

app-server/src/db/trace.rs

@@ -67,7 +67,8 @@ pub async fn update_trace_attributes(
             session_id,
             trace_type,
             metadata,
-            has_browser_session
+            has_browser_session,
+            top_span_id
         )
         VALUES (
             $1,
@@ -83,7 +84,8 @@ pub async fn update_trace_attributes(
             $11,
             COALESCE($12, 'DEFAULT'::trace_type),
             $13,
-            $14
+            $14,
+            $15
         )
         ON CONFLICT(id) DO
         UPDATE
@@ -96,10 +98,11 @@ pub async fn update_trace_attributes(
             cost = traces.cost + COALESCE($8, 0),
             start_time = CASE WHEN traces.start_time IS NULL OR traces.start_time > $9 THEN $9 ELSE traces.start_time END,
             end_time = CASE WHEN traces.end_time IS NULL OR traces.end_time < $10 THEN $10 ELSE traces.end_time END,
-            session_id = CASE WHEN traces.session_id IS NULL THEN $11 ELSE traces.session_id END,
+            session_id = COALESCE(traces.session_id, $11),
             trace_type = CASE WHEN $12 IS NULL THEN traces.trace_type ELSE COALESCE($12, 'DEFAULT'::trace_type) END,
             metadata = COALESCE($13, traces.metadata),
-            has_browser_session = COALESCE($14, traces.has_browser_session)
+            has_browser_session = COALESCE($14, traces.has_browser_session),
+            top_span_id = COALESCE(traces.top_span_id, $15)
         "
     )
     .bind(attributes.id)
@@ -116,6 +119,7 @@ pub async fn update_trace_attributes(
     .bind(&attributes.trace_type)
     .bind(&serde_json::to_value(&attributes.metadata).unwrap())
     .bind(attributes.has_browser_session)
+    .bind(attributes.top_span_id)
     .execute(pool)
     .await?;
     Ok(())

app-server/src/traces/attributes.rs

@@ -22,6 +22,7 @@ pub struct TraceAttributes {
     pub trace_type: Option<TraceType>,
     pub metadata: Option<HashMap<String, String>>,
     pub has_browser_session: Option<bool>,
+    pub top_span_id: Option<Uuid>,
 }
 
 impl TraceAttributes {
@@ -82,4 +83,8 @@ impl TraceAttributes {
     pub fn set_has_browser_session(&mut self, has_browser_session: bool) {
         self.has_browser_session = Some(has_browser_session);
     }
+
+    pub fn set_top_span_id(&mut self, top_span_id: Uuid) {
+        self.top_span_id = Some(top_span_id);
+    }
 }

app-server/src/traces/utils.rs

@@ -129,6 +129,10 @@ pub async fn record_span_to_db(
             }
         }
     });
+    // Once we've set the parent span id, check if it's the top span
+    if span.parent_span_id.is_none() {
+        trace_attributes.set_top_span_id(span.span_id);
+    }
     span_attributes.update_path();
     span.set_attributes(&span_attributes);
 

frontend/app/api/projects/[projectId]/spans/[spanId]/basic-info/route.ts

@@ -0,0 +1,30 @@
+// Returns span basic info that is shown in the traces table
+// when traces arrive realtime
+
+import { and, eq } from 'drizzle-orm';
+import { NextResponse } from 'next/server';
+
+import { db } from '@/lib/db/drizzle';
+import { spans } from '@/lib/db/migrations/schema';
+
+export async function GET(
+  req: Request,
+  props: { params: Promise<{ projectId: string; spanId: string }> }
+): Promise<Response> {
+  const params = await props.params;
+  const projectId = params.projectId;
+  const spanId = params.spanId;
+
+
+  const span = await db.query.spans.findFirst({
+    where: and(eq(spans.spanId, spanId), eq(spans.projectId, projectId)),
+    columns: {
+      spanType: true,
+      name: true,
+      inputPreview: true,
+      outputPreview: true,
+    }
+  });
+
+  return NextResponse.json(span);
+}

frontend/app/api/projects/[projectId]/traces/route.ts

@@ -76,6 +76,8 @@ export async function GET(
       latency: sql<number>`EXTRACT(EPOCH FROM (end_time - start_time))`.as("latency"),
     }).from(traces).leftJoin(
       topLevelSpans,
+      // We could as well join on eq(traces.topSpanId, topLevelSpans.id),
+      // but this is more performant, as spans are indexed by traceId
       eq(traces.id, topLevelSpans.traceId)
     ).where(
       and(
@@ -204,6 +206,14 @@ export async function DELETE(
         )
       );
 
+    await db.delete(spans)
+      .where(
+        and(
+          inArray(traces.id, traceId),
+          eq(traces.projectId, projectId)
+        )
+      );
+
     return new Response('Traces deleted successfully', { status: 200 });
   } catch (error) {
     return new Response('Error deleting traces', { status: 500 });

frontend/components/traces/span-card.tsx

@@ -1,4 +1,4 @@
-import { ChevronDown, ChevronRight, CircleX, X } from 'lucide-react';
+import { ChevronDown, ChevronRight, X } from 'lucide-react';
 import React, { useEffect, useRef, useState } from 'react';
 
 import { getDuration, getDurationString } from '@/lib/flow/utils';

frontend/components/traces/traces-table.tsx

@@ -8,7 +8,7 @@ import DeleteSelectedRows from '@/components/ui/DeleteSelectedRows';
 import { useProjectContext } from '@/contexts/project-context';
 import { useUserContext } from '@/contexts/user-context';
 import { useToast } from '@/lib/hooks/use-toast';
-import { SpanType, Trace, TraceWithSpans } from '@/lib/traces/types';
+import { SpanType, Trace } from '@/lib/traces/types';
 import { isStringDateOld } from '@/lib/traces/utils';
 import { DatatableFilter, PaginatedResponse } from '@/lib/types';
 import { getFilterFromUrlParams } from '@/lib/utils';
@@ -151,6 +151,7 @@ export default function TracesTable({ onRowClick }: TracesTableProps) {
     cost: row.cost,
     metadata: row.metadata,
     hasBrowserSession: row.has_browser_session,
+    topSpanId: row.top_span_id,
     topSpanInputPreview: null,
     topSpanOutputPreview: null,
     topSpanName: null,
@@ -159,16 +160,19 @@ export default function TracesTable({ onRowClick }: TracesTableProps) {
   });
 
   // TODO: maybe also query top span input and output previews?
-  const getTraceTopSpanInfo = async (id: string): Promise<{
-    topSpanName: string | null;
-    topSpanType: SpanType | null;
+  const getTraceTopSpanInfo = async (spanId: string): Promise<{
+    topSpanName: string | null,
+    topSpanType: SpanType | null,
+    topSpanInputPreview: any | null,
+    topSpanOutputPreview: any | null,
   }> => {
-    const response = await fetch(`/api/projects/${projectId}/traces/${id}`);
-    const trace = await response.json() as TraceWithSpans;
-    const topSpan = trace.spans.find(span => span.parentSpanId === null);
+    const response = await fetch(`/api/projects/${projectId}/spans/${spanId}/basic-info`);
+    const span = await response.json();
     return {
-      topSpanName: topSpan?.name ?? null,
-      topSpanType: topSpan?.spanType ?? null,
+      topSpanName: span?.name ?? null,
+      topSpanType: span?.spanType ?? null,
+      topSpanInputPreview: span?.inputPreview ?? null,
+      topSpanOutputPreview: span?.outputPreview ?? null,
     };
   };
 
@@ -182,10 +186,10 @@ export default function TracesTable({ onRowClick }: TracesTableProps) {
       const insertIndex = currentTraces?.findIndex(trace => trace.startTime <= newObj.start_time);
       const newTraces = currentTraces ? [...currentTraces] : [];
       const rtEventTrace = dbTraceRowToTrace(newObj);
-      const { topSpanType, topSpanName, ...rest } = rtEventTrace;
-      const newTrace = (rtEventTrace.topSpanType === null)
+      const { topSpanType, topSpanName, topSpanInputPreview, topSpanOutputPreview, ...rest } = rtEventTrace;
+      const newTrace = (rtEventTrace.topSpanType === null && rtEventTrace.topSpanId != null)
         ? {
-          ...(await getTraceTopSpanInfo(rtEventTrace.id)),
+          ...(await getTraceTopSpanInfo(rtEventTrace.topSpanId)),
           ...rest,
         }
         : rtEventTrace;
@@ -204,10 +208,10 @@ export default function TracesTable({ onRowClick }: TracesTableProps) {
         const newTraces = [...currentTraces];
         const existingTrace = currentTraces[updateIndex];
         const rtEventTrace = dbTraceRowToTrace(newObj);
-        const { topSpanType, topSpanName, ...rest } = rtEventTrace;
-        if (existingTrace.topSpanType === null) {
+        const { topSpanType, topSpanName, topSpanInputPreview, topSpanOutputPreview, ...rest } = rtEventTrace;
+        if (existingTrace.topSpanType === null && rtEventTrace.topSpanId != null) {
           const newTrace = {
-            ...(await getTraceTopSpanInfo(existingTrace.id)),
+            ...(await getTraceTopSpanInfo(rtEventTrace.topSpanId)),
             ...rest,
           };
           newTraces[updateIndex] = newTrace;
@@ -307,24 +311,24 @@ export default function TracesTable({ onRowClick }: TracesTableProps) {
     }
   };
 
-  const handleAddFilter = (column: string, value: string) => {
-    const newFilter = { column, operator: 'eq', value };
-    const existingFilterIndex = activeFilters.findIndex(
-      (filter) => filter.column === column && filter.value === value
-    );
+  // const handleAddFilter = (column: string, value: string) => {
+  //   const newFilter = { column, operator: 'eq', value };
+  //   const existingFilterIndex = activeFilters.findIndex(
+  //     (filter) => filter.column === column && filter.value === value
+  //   );
 
-    let updatedFilters;
-    if (existingFilterIndex === -1) {
+  //   let updatedFilters;
+  //   if (existingFilterIndex === -1) {
 
-      updatedFilters = [...activeFilters, newFilter];
-    } else {
+  //     updatedFilters = [...activeFilters, newFilter];
+  //   } else {
 
-      updatedFilters = [...activeFilters];
-    }
+  //     updatedFilters = [...activeFilters];
+  //   }
 
-    setActiveFilters(updatedFilters);
-    updateUrlWithFilters(updatedFilters);
-  };
+  //   setActiveFilters(updatedFilters);
+  //   updateUrlWithFilters(updatedFilters);
+  // };
 
   const updateUrlWithFilters = (filters: DatatableFilter[]) => {
     searchParams.delete('filter');

frontend/lib/db/migrations/0022_hesitant_skin.sql

@@ -0,0 +1 @@
+ALTER TABLE "traces" ADD COLUMN "top_span_id" uuid;