diff --git a/frontend/app/api/projects/[projectId]/traces/[traceId]/route.ts b/frontend/app/api/projects/[projectId]/traces/[traceId]/route.ts index 26575c20..81ea5f50 100644 --- a/frontend/app/api/projects/[projectId]/traces/[traceId]/route.ts +++ b/frontend/app/api/projects/[projectId]/traces/[traceId]/route.ts @@ -1,10 +1,8 @@ -import { and, asc, eq, inArray, sql } from 'drizzle-orm'; +import { and, eq } from 'drizzle-orm'; import { NextRequest, NextResponse } from 'next/server'; -import { searchSpans } from '@/lib/clickhouse/spans'; -import { TimeRange } from '@/lib/clickhouse/utils'; import { db } from '@/lib/db/drizzle'; -import { events, spans, traces } from '@/lib/db/migrations/schema'; +import { traces } from '@/lib/db/migrations/schema'; export async function GET( req: NextRequest, @@ -13,65 +11,14 @@ export async function GET( const params = await props.params; const projectId = params.projectId; const traceId = params.traceId; - const searchQuery = req.nextUrl.searchParams.get("search"); - let searchSpanIds = null; - if (searchQuery) { - const timeRange = { pastHours: 'all' } as TimeRange; - const searchResult = await searchSpans(projectId, searchQuery, timeRange); - searchSpanIds = Array.from(searchResult.spanIds); - } - - const traceQuery = db.query.traces.findFirst({ + const trace = await db.query.traces.findFirst({ where: and(eq(traces.id, traceId), eq(traces.projectId, projectId)), }); - const spanEventsQuery = db.$with('span_events').as( - db.select({ - spanId: events.spanId, - projectId: events.projectId, - events: sql`jsonb_agg(jsonb_build_object( - 'id', events.id, - 'spanId', events.span_id, - 'timestamp', events.timestamp, - 'name', events.name, - 'attributes', events.attributes - ))`.as('events') - }) - .from(events) - .groupBy(events.spanId, events.projectId) - ); - - const spansQuery = db.with(spanEventsQuery).select({ - // inputs and outputs are ignored on purpose - spanId: spans.spanId, - startTime: spans.startTime, - endTime: spans.endTime, - traceId: spans.traceId, - parentSpanId: spans.parentSpanId, - name: spans.name, - attributes: spans.attributes, - spanType: spans.spanType, - events: sql`COALESCE(${spanEventsQuery.events}, '[]'::jsonb)`.as('events'), - }) - .from(spans) - .leftJoin(spanEventsQuery, - and( - eq(spans.spanId, spanEventsQuery.spanId), - eq(spans.projectId, spanEventsQuery.projectId) - ) - ) - .where( - and( - eq(spans.traceId, traceId), - eq(spans.projectId, projectId), - ...(searchSpanIds ? [inArray(spans.spanId, searchSpanIds)] : []) - ) - ) - .orderBy(asc(spans.startTime)); - - const [trace, spanItems] = await Promise.all([traceQuery, spansQuery]); + if (!trace) { + return NextResponse.json({ error: 'Trace not found' }, { status: 404 }); + } - return NextResponse.json({ ...trace, spans: spanItems }); + return NextResponse.json(trace); } - diff --git a/frontend/app/api/projects/[projectId]/traces/[traceId]/spans/route.ts b/frontend/app/api/projects/[projectId]/traces/[traceId]/spans/route.ts new file mode 100644 index 00000000..9758792f --- /dev/null +++ b/frontend/app/api/projects/[projectId]/traces/[traceId]/spans/route.ts @@ -0,0 +1,71 @@ +import { and, asc, eq, inArray, sql } from 'drizzle-orm'; +import { NextRequest, NextResponse } from 'next/server'; + +import { searchSpans } from '@/lib/clickhouse/spans'; +import { TimeRange } from '@/lib/clickhouse/utils'; +import { db } from '@/lib/db/drizzle'; +import { events, spans } from '@/lib/db/migrations/schema'; + +export async function GET( + req: NextRequest, + props: { params: Promise<{ projectId: string; traceId: string }> } +): Promise { + const params = await props.params; + const projectId = params.projectId; + const traceId = params.traceId; + const searchQuery = req.nextUrl.searchParams.get("search"); + + let searchSpanIds = null; + if (searchQuery) { + const timeRange = { pastHours: 'all' } as TimeRange; + const searchResult = await searchSpans(projectId, searchQuery, timeRange); + searchSpanIds = Array.from(searchResult.spanIds); + } + + const spanEventsQuery = db.$with('span_events').as( + db.select({ + spanId: events.spanId, + projectId: events.projectId, + events: sql`jsonb_agg(jsonb_build_object( + 'id', events.id, + 'spanId', events.span_id, + 'timestamp', events.timestamp, + 'name', events.name, + 'attributes', events.attributes + ))`.as('events') + }) + .from(events) + .groupBy(events.spanId, events.projectId) + ); + + const spanItems = await db.with(spanEventsQuery).select({ + // inputs and outputs are ignored on purpose + spanId: spans.spanId, + startTime: spans.startTime, + endTime: spans.endTime, + traceId: spans.traceId, + parentSpanId: spans.parentSpanId, + name: spans.name, + attributes: spans.attributes, + spanType: spans.spanType, + events: sql`COALESCE(${spanEventsQuery.events}, '[]'::jsonb)`.as('events'), + }) + .from(spans) + .leftJoin(spanEventsQuery, + and( + eq(spans.spanId, spanEventsQuery.spanId), + eq(spans.projectId, spanEventsQuery.projectId) + ) + ) + .where( + and( + eq(spans.traceId, traceId), + eq(spans.projectId, projectId), + ...(searchSpanIds ? [inArray(spans.spanId, searchSpanIds)] : []) + ) + ) + .orderBy(asc(spans.startTime)); + + + return NextResponse.json(spanItems); +} diff --git a/frontend/components/traces/spans-table.tsx b/frontend/components/traces/spans-table.tsx index 7eed9f43..d6b9c151 100644 --- a/frontend/components/traces/spans-table.tsx +++ b/frontend/components/traces/spans-table.tsx @@ -163,8 +163,8 @@ export default function SpansTable({ onRowClick }: SpansTableProps) { inputPreview: row.input_preview, outputPreview: row.output_preview, events: [], - inputUrl: null, - outputUrl: null, + inputUrl: row.input_url, + outputUrl: row.output_url, model: row.attributes['gen_ai.response.model'] ?? row.attributes['gen_ai.request.model'] ?? null, }); diff --git a/frontend/components/traces/trace-view.tsx b/frontend/components/traces/trace-view.tsx index 419f3d69..ebdda0d8 100644 --- a/frontend/components/traces/trace-view.tsx +++ b/frontend/components/traces/trace-view.tsx @@ -1,12 +1,11 @@ import { ChartNoAxesGantt, ChevronsRight, Disc } from 'lucide-react'; import { usePathname, useRouter, useSearchParams } from 'next/navigation'; import React, { useEffect, useRef, useState } from 'react'; -import useSWR from 'swr'; import { useProjectContext } from '@/contexts/project-context'; import { useUserContext } from '@/contexts/user-context'; -import { Span, SpanType, TraceWithSpans } from '@/lib/traces/types'; -import { cn, swrFetcher } from '@/lib/utils'; +import { Span, SpanType, Trace } from '@/lib/traces/types'; +import { cn } from '@/lib/utils'; import { Button } from '../ui/button'; import MonoWithCopy from '../ui/mono-with-copy'; @@ -28,6 +27,8 @@ export default function TraceView({ traceId, onClose }: TraceViewProps) { const searchParams = new URLSearchParams(useSearchParams().toString()); const router = useRouter(); const pathName = usePathname(); + const { projectId } = useProjectContext(); + const container = useRef(null); // containerHeight refers to the height of the trace view container const [containerHeight, setContainerHeight] = useState(0); @@ -37,18 +38,23 @@ export default function TraceView({ traceId, onClose }: TraceViewProps) { // here timelineWidth refers to the width of the trace tree panel AND waterfall timeline const [timelineWidth, setTimelineWidth] = useState(0); const [traceTreePanelWidth, setTraceTreePanelWidth] = useState(0); - const { projectId } = useProjectContext(); const [hasBrowserSession, setHasBrowserSession] = useState(false); const [showBrowserSession, setShowBrowserSession] = useState(false); const browserSessionRef = useRef(null); - const { data: trace, isLoading, mutate } = useSWR( - `/api/projects/${projectId}/traces/${traceId}`, - swrFetcher - ); + + const [trace, setTrace] = useState(null); + + const [spans, setSpans] = useState([]); + const spansRef = useRef([]); + + // Keep ref updated + useEffect(() => { + spansRef.current = spans; + }, [spans]); const [childSpans, setChildSpans] = useState<{ [key: string]: Span[] }>({}); const [topLevelSpans, setTopLevelSpans] = useState([]); - const [spans, setSpans] = useState([]); + const [selectedSpan, setSelectedSpan] = useState( searchParams.get('spanId') ? spans.find( @@ -64,12 +70,24 @@ export default function TraceView({ traceId, onClose }: TraceViewProps) { const [browserSessionTime, setBrowserSessionTime] = useState(null); useEffect(() => { - if (!trace) { - return; - } + const fetchTrace = async () => { + const trace = await fetch(`/api/projects/${projectId}/traces/${traceId}`); + return await trace.json(); + }; - const spans = enrichSpansWithPending(trace.spans); + fetchTrace().then((trace) => { + setTrace(trace); + if (trace.hasBrowserSession) { + if (!hasBrowserSession) { + // if we previously didn't have a browser session, show it + setShowBrowserSession(true); + } + setHasBrowserSession(true); + } + }); + }, [traceId, spans, projectId]); + useEffect(() => { const childSpans = {} as { [key: string]: Span[] }; const topLevelSpans = spans.filter((span: Span) => !span.parentSpanId); @@ -85,39 +103,36 @@ export default function TraceView({ traceId, onClose }: TraceViewProps) { setChildSpans(childSpans); setTopLevelSpans(topLevelSpans); - setSpans(spans); - - // If there's only one span, select it automatically - if (spans.length === 1) { - const singleSpan = spans[0]; - setSelectedSpan(singleSpan); - searchParams.set('spanId', singleSpan.spanId); - router.push(`${pathName}?${searchParams.toString()}`); - } else { - // Otherwise, use the spanId from URL if present - setSelectedSpan( - searchParams.get('spanId') - ? spans.find( - (span: Span) => span.spanId === searchParams.get('spanId') - ) || null - : null - ); - } - if (trace?.hasBrowserSession) { - if (!hasBrowserSession) { - // if we previously didn't have a browser session, show it - setShowBrowserSession(true); - } - setHasBrowserSession(true); - } - }, [trace]); + }, [spans]); useEffect(() => { - if (trace?.hasBrowserSession) { - setHasBrowserSession(true); - setShowBrowserSession(true); - } - }, []); + const fetchSpans = async () => { + const response = await fetch(`/api/projects/${projectId}/traces/${traceId}/spans`); + const results = await response.json(); + return enrichSpansWithPending(results); + }; + + fetchSpans().then((spans) => { + setSpans(spans); + + // If there's only one span, select it automatically + if (spans.length === 1) { + const singleSpan = spans[0]; + setSelectedSpan(singleSpan); + searchParams.set('spanId', singleSpan.spanId); + router.push(`${pathName}?${searchParams.toString()}`); + } else { + // Otherwise, use the spanId from URL if present + setSelectedSpan( + searchParams.get('spanId') + ? spans.find( + (span: Span) => span.spanId === searchParams.get('spanId') + ) || null + : null + ); + } + }); + }, [traceId, projectId]); useEffect(() => { setSelectedSpan( @@ -166,11 +181,30 @@ export default function TraceView({ traceId, onClose }: TraceViewProps) { } }, [containerWidth, selectedSpan, traceTreePanel.current, collapsedSpans]); + const dbSpanRowToSpan = (row: Record): Span => ({ + spanId: row.span_id, + parentSpanId: row.parent_span_id, + traceId: row.trace_id, + spanType: row.span_type, + name: row.name, + path: row.attributes['lmnr.span.path'] ?? "", + startTime: row.start_time, + endTime: row.end_time, + attributes: row.attributes, + input: null, + output: null, + inputPreview: row.input_preview, + outputPreview: row.output_preview, + events: [], + inputUrl: row.input_url, + outputUrl: row.output_url, + model: row.attributes['gen_ai.response.model'] ?? row.attributes['gen_ai.request.model'] ?? null, + }); const { supabaseClient: supabase } = useUserContext(); useEffect(() => { - if (!supabase) { + if (!supabase || !projectId) { return; } @@ -188,8 +222,19 @@ export default function TraceView({ traceId, onClose }: TraceViewProps) { }, (payload) => { if (payload.eventType === 'INSERT') { - // just mutate for now - mutate(); + // Use a functional state update to avoid race conditions + setSpans(currentSpans => { + const rtEventSpan = dbSpanRowToSpan(payload.new); + const newSpans = [...currentSpans]; + const index = newSpans.findIndex(span => span.spanId === rtEventSpan.spanId); + if (index !== -1 && newSpans[index].pending) { + newSpans[index] = rtEventSpan; + } else { + newSpans.push(rtEventSpan); + } + + return enrichSpansWithPending(newSpans); + }); } } ) @@ -199,7 +244,7 @@ export default function TraceView({ traceId, onClose }: TraceViewProps) { return () => { supabase.removeAllChannels(); }; - }, []); + }, [supabase, projectId]); return (
@@ -250,7 +295,7 @@ export default function TraceView({ traceId, onClose }: TraceViewProps) {
- {isLoading && ( + {(!trace && spans.length === 0) && (
@@ -424,14 +469,11 @@ const enrichSpansWithPending = (existingSpans: Span[]): Span[] => { const parentSpanIds = span.attributes['lmnr.span.ids_path'] as string[] | undefined; const parentSpanNames = span.attributes['lmnr.span.path'] as string[] | undefined; - if (parentSpanIds === undefined || parentSpanNames === undefined) { - continue; - } - - if (parentSpanIds.length === 0 || parentSpanNames.length === 0) { - continue; - } - if (parentSpanIds.length !== parentSpanNames.length) { + if ( + parentSpanIds === undefined || parentSpanNames === undefined || + parentSpanIds.length === 0 || parentSpanNames.length === 0 || + parentSpanIds.length !== parentSpanNames.length + ) { continue; } @@ -446,6 +488,11 @@ const enrichSpansWithPending = (existingSpans: Span[]): Span[] => { } if (pendingSpans.has(spanId)) { + // TODO: looks like this is not working for realtime inserts and is + // unreachable because of the existingSpanIds check above. + // If the check above is removed, the pending spans are inserted + // repeatedly, many times over. + // if the pending span is already present, just update the start and end time // to span over all its children const existingStartTime = new Date(pendingSpans.get(spanId)!.startTime); diff --git a/frontend/lib/traces/types.ts b/frontend/lib/traces/types.ts index 0f32a508..7404ba27 100644 --- a/frontend/lib/traces/types.ts +++ b/frontend/lib/traces/types.ts @@ -71,22 +71,6 @@ export type Span = { pending?: boolean; }; -export type TraceWithSpans = { - id: string; - startTime: string; - endTime: string; - inputTokenCount: number; - outputTokenCount: number; - totalTokenCount: number; - inputCost: number | null; - outputCost: number | null; - cost: number | null; - metadata: Record | null; - hasBrowserSession: boolean | null; - projectId: string; - spans: Span[]; -}; - export type Trace = { startTime: string; endTime: string;