From fb554d34f39a8b58ff7e3da88c38047ec3370bc9 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 6 Oct 2026 12:10:19 -0700 Subject: [PATCH] fix(lens): paginate visible conversation entries (#44878) * fix(lens): paginate visible conversation entries * fix(lens): stop pagination at exhausted subagents * refactor(lens): derive pending branches without mutation * fix(lens): bound pending ancestry work for deep traces --- .../TraceConversation.integration.test.tsx | 231 +++++++++++++++++- .../detail/conversation/TraceConversation.tsx | 231 +++++++++++++----- .../detail/conversation/conversation.test.ts | 69 ++++++ .../detail/conversation/conversation.ts | 61 ++++- .../conversation/useConversationDetails.ts | 48 ++++ .../lens/traces/detail/run/RunBody.tsx | 6 +- .../lens/traces/detail/run/RunView.tsx | 5 + 7 files changed, 578 insertions(+), 73 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/lens/traces/detail/conversation/useConversationDetails.ts diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.integration.test.tsx index 3db85bc22da..d39c486025d 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.integration.test.tsx @@ -120,7 +120,7 @@ describe("TraceConversation", () => { expect(within(result).getByText(/"exit_code": 1/, { selector: "pre" })).toHaveTextContent('"error": null'); }); - it("loads only twenty full steps at a time and fetches the remainder on demand", async () => { + it("loads twenty span details initially and pages conversation entries on demand", async () => { const user = userEvent.setup(); const spans = [ root, @@ -128,7 +128,7 @@ describe("TraceConversation", () => { ]; const long = { ...trace, spans } as Trace; renderWithProviders(); - const more = await screen.findByRole("button", { name: "Load next 11 steps" }); + const more = await screen.findByRole("button", { name: "Load next 20 entries" }); await waitFor(() => expect(more).toBeEnabled()); expect(agentTraceSpanCall).toHaveBeenCalledTimes(20); expect(screen.queryByText("The release is ready")).not.toBeInTheDocument(); @@ -150,6 +150,200 @@ describe("TraceConversation", () => { expect(screen.queryByText("End of conversation")).not.toBeInTheDocument(); }); + it("adds twenty visible entries across hidden subagents, supplemental spans, and server pages", async () => { + const user = userEvent.setup(); + const child = { ...root, span_id: "child", parent_span_id: "root", name: "Reviewer", start_offset_ms: 1 }; + const childSpans = Array.from({ length: 50 }, (_, index) => [ + { + ...tool, + span_id: `child-${index}`, + parent_span_id: "child", + name: `child tool ${index}`, + start_offset_ms: 2 + index * 2, + }, + { + ...tool, + span_id: `supplement-${index}`, + parent_span_id: "child", + name: "claude_code.tool_result", + framework: "claude-code", + type: "event", + start_offset_ms: 3 + index * 2, + }, + ]).flat(); + const parentSpans = Array.from({ length: 45 }, (_, index) => [ + { ...tool, span_id: `parent-${index}`, name: `parent tool ${index}`, start_offset_ms: 102 + index * 2 }, + { ...tool, span_id: `empty-${index}`, type: "llm", start_offset_ms: 103 + index * 2 }, + ]).flat(); + const spans = [root, child, ...childSpans, ...parentSpans] as Trace["spans"]; + const first: Trace = { ...trace, spans: spans.slice(0, 120), next_cursor: "next-page" }; + const last: Trace = { ...trace, spans: spans.slice(120), next_cursor: null }; + vi.mocked(agentTraceCall).mockImplementation(async (_token, _trace, _ref, cursor) => (cursor ? last : first)); + vi.mocked(agentTraceSpanCall).mockImplementation(async (_token, _trace, id) => { + if (id === "root" || id === "child") + return { ...rootDetail, span_id: id, input: id === "root" ? "Main task" : "Child task", output: "" }; + if (id.startsWith("empty-") || id.startsWith("supplement-")) + return { span_id: id, input: "", output: "", attributes: {} }; + return { ...toolDetail, span_id: id }; + }); + renderWithProviders( + , + ); + await user.click(await screen.findByRole("tab", { name: "Conversation" })); + expect(await screen.findByText("2 entries shown")).toBeVisible(); + expect(agentTraceSpanCall).toHaveBeenCalledTimes(20); + await user.click(screen.getByRole("button", { name: "Load next 20 entries", exact: true })); + expect(await screen.findByText("22 entries shown")).toBeVisible(); + expect(screen.getAllByRole("region", { name: /Conversation step parent tool/ })).toHaveLength(20); + expect( + screen.queryByRole("button", { name: "Expand parent tool 20 tool call", exact: true }), + ).not.toBeInTheDocument(); + expect(screen.getByRole("region", { name: "Conversation step child tool 0", exact: true })).not.toBeVisible(); + expect(agentTraceCall).toHaveBeenCalledWith("test", trace.summary.trace_id, undefined, "next-page"); + + await user.click(screen.getByText("Subagent: Reviewer", { exact: true })); + expect(screen.getAllByRole("region", { name: /Conversation step child tool/ })).toHaveLength(19); + await user.click(screen.getByRole("button", { name: "Load next 20 entries in Reviewer", exact: true })); + expect(screen.getAllByRole("region", { name: /Conversation step child tool/ })).toHaveLength(39); + expect(screen.getAllByRole("region", { name: /Conversation step parent tool/ })).toHaveLength(20); + expect(screen.queryByText("End of conversation")).not.toBeInTheDocument(); + }); + + it.each([ + { duration: 10, nextCursor: null }, + { duration: 1000, nextCursor: null }, + { duration: 10, nextCursor: "unrelated-page" }, + ])("hides an exhausted subagent control with unrelated work remaining (%j)", async ({ duration, nextCursor }) => { + const user = userEvent.setup(); + const child = { + ...root, + span_id: "child", + parent_span_id: "root", + name: "Reviewer", + start_offset_ms: 1, + duration_ms: duration, + }; + const spans = [ + root, + child, + { ...tool, span_id: "child-tool", parent_span_id: "child", start_offset_ms: 2, duration_ms: 1 }, + ...Array.from({ length: 40 }, (_, index) => ({ + ...tool, + span_id: `parent-${index}`, + start_offset_ms: 100 + index, + duration_ms: 1, + })), + ] as Trace["spans"]; + vi.mocked(agentTraceSpanCall).mockImplementation(async (_token, _trace, id) => { + if (id === "child") return { ...rootDetail, span_id: id, input: "Review task", output: "Review finished" }; + return id === "root" ? rootDetail : { ...toolDetail, span_id: id }; + }); + const loadMore = vi.fn(); + renderWithProviders( + , + ); + const more = await screen.findByRole("button", { name: "Load next 20 entries", exact: true }); + await waitFor(() => expect(more).toBeEnabled()); + await user.click(screen.getByText("Subagent: Reviewer", { exact: true })); + expect(screen.getByText("Review finished")).toBeVisible(); + expect(screen.queryByRole("button", { name: /entries in Reviewer/ })).not.toBeInTheDocument(); + expect(agentTraceSpanCall).toHaveBeenCalledTimes(20); + expect(loadMore).not.toHaveBeenCalled(); + }); + + it.each([false, true])("stops loading at the selected subagent's end (server paging: %s)", async (paged) => { + const user = userEvent.setup(); + const child = { + ...root, + span_id: "child", + parent_span_id: "root", + name: "Reviewer", + start_offset_ms: 1, + duration_ms: 60, + }; + const spans = [ + root, + child, + ...Array.from({ length: 25 }, (_, index) => ({ + ...tool, + span_id: `child-${index}`, + parent_span_id: "child", + name: `child tool ${index}`, + start_offset_ms: 2 + index, + duration_ms: 1, + })), + ...Array.from({ length: 40 }, (_, index) => ({ + ...tool, + span_id: `parent-${index}`, + name: `parent tool ${index}`, + start_offset_ms: 100 + index, + duration_ms: 1, + })), + ] as Trace["spans"]; + const first: Trace = { + ...trace, + spans: paged ? spans.slice(0, 20) : spans, + next_cursor: paged ? "child-page" : null, + }; + const next: Trace = { ...trace, spans: spans.slice(20, 40), next_cursor: "unrelated-page" }; + vi.mocked(agentTraceCall).mockImplementation(async (_token, _trace, _ref, cursor) => (cursor ? next : first)); + vi.mocked(agentTraceSpanCall).mockImplementation(async (_token, _trace, id) => { + if (id === "child") return { ...rootDetail, span_id: id, input: "Review task", output: "Review finished" }; + return id === "root" ? rootDetail : { ...toolDetail, span_id: id }; + }); + renderWithProviders( + , + ); + await user.click(await screen.findByRole("tab", { name: "Conversation" })); + await user.click(await screen.findByText("Subagent: Reviewer", { exact: true })); + const more = await screen.findByRole("button", { name: "Load next 20 entries in Reviewer" }); + await waitFor(() => expect(more).toBeEnabled()); + await user.click(more); + expect(await screen.findByRole("button", { name: "Expand child tool 24 tool call" })).toBeVisible(); + await waitFor(() => expect(screen.queryByRole("button", { name: /entries in Reviewer/ })).not.toBeInTheDocument()); + expect(screen.getByText("Review finished")).toBeVisible(); + expect(screen.getAllByRole("region", { name: /Conversation step child tool/ })).toHaveLength(25); + expect(screen.getByRole("button", { name: "Load next 20 entries", exact: true })).toBeEnabled(); + expect(screen.queryByText("Loading conversation…")).not.toBeInTheDocument(); + expect(agentTraceSpanCall).toHaveBeenCalledTimes(40); + expect(agentTraceCall).toHaveBeenCalledTimes(paged ? 2 : 1); + expect(agentTraceCall).not.toHaveBeenCalledWith("test", trace.summary.trace_id, undefined, "unrelated-page"); + }); + + it("stops at a failed server page and resumes the requested entries after retry", async () => { + const user = userEvent.setup(); + const first: Trace = { ...trace, next_cursor: "next-page" }; + const last: Trace = { + ...trace, + spans: [{ ...tool, span_id: "later", name: "later tool", start_offset_ms: 2 }], + next_cursor: null, + }; + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(first) + .mockRejectedValueOnce(new Error("temporarily unavailable")) + .mockResolvedValue(last); + renderWithProviders( + , + ); + await user.click(await screen.findByRole("tab", { name: "Conversation" })); + const more = await screen.findByRole("button", { name: "Load next 20 entries", exact: true }); + await waitFor(() => expect(more).toBeEnabled()); + await user.click(more); + expect(await screen.findByRole("alert")).toHaveTextContent("Could not load more conversation entries"); + expect(more).toBeDisabled(); + expect(agentTraceCall).toHaveBeenCalledTimes(2); + expect(screen.getByText("Read the release notes")).toBeVisible(); + await user.click(screen.getByRole("button", { name: "Retry", exact: true })); + expect(await screen.findByRole("button", { name: "Expand later tool tool call" })).toBeVisible(); + expect(await screen.findByText("End of conversation")).toBeVisible(); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + it.each([false, true])( "checks for missing replies after the final trace page and details load (reply: %s)", async (hasReply) => { @@ -240,7 +434,7 @@ describe("TraceConversation", () => { ); expect(screen.getAllByText("Agent exceeded its execution limit")).toHaveLength(1); await act(async () => rootFetch.resolve({ ...rootDetail, output })); - const more = screen.getByRole("button", { name: "Load next 5 steps" }); + const more = screen.getByRole("button", { name: "Load next 20 entries" }); await waitFor(() => expect(more).toBeEnabled()); expect(screen.getAllByText("Agent exceeded its execution limit")).toHaveLength(1); expect(screen.queryByText("The release is ready")).not.toBeInTheDocument(); @@ -262,6 +456,35 @@ describe("TraceConversation", () => { expect(screen.queryByRole("alert")).not.toBeInTheDocument(); }); + it("pauses a requested page at a failed prefetched detail until it is retried", async () => { + const user = userEvent.setup(); + const spans = [ + root, + ...Array.from({ length: 45 }, (_, index) => ({ + ...tool, + span_id: `tool-${index}`, + name: `check ${index}`, + start_offset_ms: index + 1, + })), + ] as Trace["spans"]; + const retry = vi.fn().mockRejectedValueOnce(new Error("unavailable")).mockResolvedValue(toolDetail); + vi.mocked(agentTraceSpanCall).mockImplementation(async (_token, _trace, id) => { + if (id === "tool-21") return retry(); + return id === "root" ? rootDetail : { ...toolDetail, span_id: id }; + }); + renderWithProviders(); + const more = screen.getByRole("button", { name: "Load next 20 entries", exact: true }); + await waitFor(() => expect(more).toBeEnabled()); + await user.click(more); + expect(await screen.findByRole("alert")).toHaveTextContent("Could not load check 21"); + expect(retry).toHaveBeenCalledTimes(1); + expect(screen.queryByRole("button", { name: "Expand check 22 tool call" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Retry step" })); + expect(await screen.findByText("40 entries shown")).toBeVisible(); + expect(screen.getAllByRole("region", { name: /Conversation step check/ })).toHaveLength(39); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + it("holds later turns and the final answer behind a failed step until retry succeeds", async () => { const user = userEvent.setup(); const first = { ...tool, span_id: "first", name: "First response", type: "llm", start_offset_ms: 1 }; @@ -293,7 +516,7 @@ describe("TraceConversation", () => { expect(await screen.findByText("Checking the release")).toBeVisible(); expect(screen.queryByText("The release is ready")).not.toBeInTheDocument(); expect(screen.queryByText("End of conversation")).not.toBeInTheDocument(); - expect(screen.getByText("2 of 4 steps loaded")).toBeVisible(); + expect(screen.getByText("2 entries shown")).toBeVisible(); await user.click(screen.getByRole("button", { name: "Retry step" })); expect(await screen.findByText("The release is ready")).toBeVisible(); diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.tsx index 17df054d2e9..93fb6456d28 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.tsx @@ -1,18 +1,17 @@ "use client"; -import { useQueries } from "@tanstack/react-query"; -import { useState } from "react"; +import { useEffect, useState } from "react"; import { ChevronRight, Wrench } from "lucide-react"; import { cn } from "@/lib/cva.config"; import CopyButton from "@/components/shared/CopyButton"; import { ToolArguments, ToolOutput } from "../content/ToolContent"; import { toolSummary } from "../content/payload"; -import { useTracesApi } from "../../api"; import { Button } from "@/components/ui/button"; import { buildConversation, - conversationSteps, conversationWarnings, + conversationBranchGroups, + pendingConversationBranches, groupConversation, CONVERSATION_PAGE_SIZE, type ConversationItem, @@ -21,44 +20,70 @@ import { import { ErrorBlock } from "../content/SpanError"; import { Markdown } from "../content/Markdown"; import { ToolCallBlock, ToolResultCard } from "../content/Messages"; -import type { SpanDetail, Trace, TraceMessage } from "../../types"; +import type { Trace, TraceMessage } from "../../types"; import { fmtMs } from "../../utils"; +import { useConversationDetails } from "./useConversationDetails"; + +export interface ConversationTracePaging { + loading: boolean; + failed: boolean; + loadMore: () => void; +} + +function hasRemainingSource(branchId: string | null, sourceRemaining: boolean, pendingBranches: ReadonlySet) { + return sourceRemaining && (branchId === null || pendingBranches.has(branchId)); +} export function TraceConversation({ trace, accessToken, onOpenStep, + paging, }: { trace: Trace; accessToken: string; onOpenStep: (id: string) => void; + paging?: ConversationTracePaging; }) { - const traces = useTracesApi(accessToken); - const [limit, setLimit] = useState(CONVERSATION_PAGE_SIZE); - const steps = conversationSteps(trace.spans); - const visible = steps.slice(0, limit); - const { trace_id: traceId, trace_ref: traceRef } = trace.summary; - const queries = useQueries({ - queries: visible.map((span) => ({ - queryKey: ["agentTraceSpan", traceId, traceRef, span.span_id, accessToken], - queryFn: (): Promise => traces.span(traceId, span.span_id, traceRef), - staleTime: Infinity, - retry: false, - })), - }); - const loading = queries.some((query) => query.isPending); - const failed = queries.some((query) => query.isError); - const unresolvedIndex = queries.findIndex((query) => !query.isSuccess); - const loadedCount = unresolvedIndex < 0 ? queries.length : unresolvedIndex; - const details = new Map( - queries.slice(0, loadedCount).map((query, index) => [visible[index].span_id, query.data!] as const), - ); - const complete = loadedCount === steps.length; + const [pages, setPages] = useState>(new Map()); + const [requestedBranch, setRequestedBranch] = useState(); + const { details, entries, complete, loading, failed, hasMore, loadMore } = useConversationDetails(trace, accessToken); const traceComplete = complete && !trace.next_cursor; - const items = buildConversation(trace.spans, details, complete); + const pendingBranches = pendingConversationBranches(trace.spans, details, Boolean(trace.next_cursor)); + const items = buildConversation(trace.spans, details, complete, pendingBranches); + const groups = groupConversation(items, trace.spans); + const requestedId = requestedBranch ?? null; + const requestedGroups = conversationBranchGroups(groups, requestedId); + const requestedLimit = pages.get(requestedId) ?? CONVERSATION_PAGE_SIZE; + const sourceRemaining = hasMore || Boolean(trace.next_cursor && paging); + const requestedSourceRemaining = hasRemainingSource(requestedId, sourceRemaining, pendingBranches); + const needsMore = + requestedBranch !== undefined && requestedGroups.length < requestedLimit && requestedSourceRemaining; + const busy = loading || Boolean(paging?.loading); + const blocked = failed || Boolean(paging?.failed); + const loadTracePage = paging?.loadMore; + useEffect(() => { + if (!needsMore || busy || blocked) return; + if (hasMore) { + const controller = new AbortController(); + void loadMore(controller.signal); + return () => controller.abort(); + } + if (trace.next_cursor) loadTracePage?.(); + }, [needsMore, busy, blocked, hasMore, loadMore, trace.next_cursor, loadTracePage]); + const fillingPage = needsMore && sourceRemaining && !blocked; + const loadEntries = (branchId: string | null, shown: number) => { + setPages((current) => new Map([...current, [branchId, shown + CONVERSATION_PAGE_SIZE]])); + setRequestedBranch(branchId); + }; const warnings = conversationWarnings(details, traceComplete); const multipleAgents = new Set(items.map((item) => item.agentId).filter(Boolean)).size > 1; - const inlineErrorIds = new Set(items.filter((item) => item.showError).map((item) => item.span.span_id)); + const rootLimit = pages.get(null) ?? CONVERSATION_PAGE_SIZE; + const inlineErrorIds = new Set( + groups + .slice(0, rootLimit) + .flatMap((group) => (group.kind === "item" && group.item.showError ? [group.item.span.span_id] : [])), + ); const rootErrors = trace.spans.filter((span) => { const failedRoot = span.parent_span_id === null && span.status === "error" && span.type !== "tool"; return failedRoot && !inlineErrorIds.has(span.span_id); @@ -75,59 +100,104 @@ export function TraceConversation({

))} + - {queries.map( - (query, index) => - query.isError && ( -
- Could not load {visible[index].name}. Retry this step to continue the conversation. - -
- ), - )} - {loading && ( -

- Loading conversation… -

- )} - {traceComplete && items.length === 0 && ( -

No conversation content recorded.

- )} -
- {traceComplete ? "End of conversation" : `${loadedCount} of ${steps.length} steps loaded`} - {visible.length < steps.length && ( - - )} -
); } +function ConversationFeedback({ + entries, + loading, + complete, + empty, + paging, +}: { + entries: ReturnType["entries"]; + loading: boolean; + complete: boolean; + empty: boolean; + paging?: ConversationTracePaging; +}) { + return ( + <> + {entries.map( + ({ query, span }) => + query.isError && ( +
+ Could not load {span.name}. Retry this step to continue the conversation. + +
+ ), + )} + {loading && ( +

+ Loading conversation… +

+ )} + {complete && empty &&

No conversation content recorded.

} + {paging?.failed && ( +

+ Could not load more conversation entries. Use Retry or Refresh trace above to continue. +

+ )} + + ); +} + function ConversationGroups({ groups, multipleAgents, onOpenStep, + branchId = null, + name, + pages, + sourceRemaining, + pendingBranches, + disabled, + onLoadEntries, + complete, }: { groups: ConversationGroup[]; multipleAgents: boolean; onOpenStep: (id: string) => void; + branchId?: string | null; + name?: string; + pages: ReadonlyMap; + sourceRemaining: boolean; + pendingBranches: ReadonlySet; + disabled: boolean; + onLoadEntries: (branchId: string | null, shown: number) => void; + complete: boolean; }) { + const limit = pages.get(branchId) ?? CONVERSATION_PAGE_SIZE; + const visible = groups.slice(0, limit); + const branchRemaining = hasRemainingSource(branchId, sourceRemaining, pendingBranches); + const hasMore = visible.length < groups.length || branchRemaining; + const nextCount = branchRemaining + ? CONVERSATION_PAGE_SIZE + : Math.min(CONVERSATION_PAGE_SIZE, groups.length - visible.length); return ( <> - {groups.map((group) => + {visible.map((group) => group.kind === "item" ? ( ) : (
- Subagent: {group.name} + + Subagent: {group.name} + + {group.children.length} entries loaded + +
- +
), )} +
+ + {complete && !hasMore && branchId === null ? "End of conversation" : `${visible.length} entries shown`} + + {hasMore && ( + + )} +
); } diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.test.ts b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.test.ts index f06e6ae57df..e4d47a94f90 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.test.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.test.ts @@ -5,6 +5,7 @@ import { newConversationMessages, groupConversation, conversationWarnings, + pendingConversationBranches, } from "./conversation"; import type { Span, SpanDetail, TraceMessage } from "../../types"; import research from "../../__fixtures__/research_trace.json"; @@ -26,6 +27,74 @@ const detail = (span_id: string, input: unknown, output: unknown): SpanDetail => }); describe("trace conversation", () => { + it("keeps missing nested details pending through framework and native agent ancestors", () => { + const agent = { + ...root, + span_id: "agent", + parent_span_id: "root", + type: "tool", + name: "Agent", + framework: "claude-code", + }; + const framework = { ...root, span_id: "framework", parent_span_id: "agent", type: "framework" }; + const child = { ...root, span_id: "child", parent_span_id: "framework" }; + const other = { ...root, span_id: "other", parent_span_id: "root" }; + const spans = [root, agent, framework, child, other] as Span[]; + const details = new Map( + spans.filter((span) => span !== child).map((span) => [span.span_id, detail(span.span_id, [], [])]), + ); + expect([...pendingConversationBranches(spans, details, false)].toSorted()).toEqual([ + "agent", + "child", + "framework", + "root", + ]); + const completeDetails = new Map([...details, [child.span_id, detail(child.span_id, [], [])]]); + expect(pendingConversationBranches(spans, completeDetails, false).size).toBe(0); + }); + + it.each(["child", "missing-parent"])("stops pending ancestry at cycles or missing parents (%s)", (parentId) => { + const agent = { ...root, span_id: "agent", parent_span_id: parentId }; + const child = { ...root, span_id: "child", parent_span_id: "agent" }; + const spans = [root, agent, child]; + const details = new Map([ + [root.span_id, detail(root.span_id, [], [])], + [agent.span_id, detail(agent.span_id, [], [])], + ]); + expect([...pendingConversationBranches(spans, details, false)].toSorted()).toEqual(["agent", "child"]); + }); + + it.each([false, true])("handles 20,000-level pending ancestry without overflowing (cycle: %s)", (cycle) => { + const rootParent = cycle ? "deep-19999" : null; + const chain = Array.from({ length: 20_000 }, (_, index) => ({ + ...root, + span_id: `deep-${index}`, + parent_span_id: index === 0 ? rootParent : `deep-${index - 1}`, + })); + const spans = [...chain, root]; + const details = new Map( + spans.filter((span) => span.span_id !== "deep-19999").map((span) => [span.span_id, detail(span.span_id, [], [])]), + ); + const pending = pendingConversationBranches(spans, details, false); + expect(pending.size).toBe(chain.length); + expect(chain.every((span) => pending.has(span.span_id))).toBe(true); + expect(pending.has(root.span_id)).toBe(false); + }); + + it.each([ + { boundary: 10, morePages: true, pending: true }, + { boundary: 20, morePages: true, pending: true }, + { boundary: 21, morePages: true, pending: false }, + { boundary: 10, morePages: false, pending: false }, + ])("only waits for pages that could contain later branch children (%j)", ({ boundary, morePages, pending }) => { + const agent = { ...root, span_id: "agent", parent_span_id: "root", start_offset_ms: 1, duration_ms: 5 }; + const child = { ...agent, span_id: "child", parent_span_id: "agent", start_offset_ms: 2, duration_ms: 18 }; + const other = { ...root, span_id: "other", parent_span_id: "root", start_offset_ms: boundary }; + const spans = [root, agent, child, other]; + const details = new Map(spans.map((span) => [span.span_id, detail(span.span_id, [], [])])); + expect(pendingConversationBranches(spans, details, morePages).has("agent")).toBe(pending); + }); + it.each(["reviewer", "__proto__", "constructor"])("keeps repeated agent name %s stable as steps load", (name) => { const first = { ...root, span_id: "first", name, start_offset_ms: 1 }; const second = { ...first, span_id: "second", start_offset_ms: 2 }; diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.ts b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.ts index d69f9dd091f..573f9eb9f9b 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.ts @@ -17,6 +17,48 @@ export function conversationSteps(spans: readonly Span[]): Span[] { .sort((a, b) => a.start_offset_ms - b.start_offset_ms); } +export function pendingConversationBranches( + spans: readonly Span[], + details: ReadonlyMap, + hasMoreSpans: boolean, +): ReadonlySet { + const byId = new Map(spans.map((span) => [span.span_id, span])); + const steps = new Set(conversationSteps(spans).map((span) => span.span_id)); + const pageBoundary = spans.reduce((latest, span) => Math.max(latest, span.start_offset_ms), -Infinity); + const pending = spans.filter((span) => { + const missingDetail = steps.has(span.span_id) && !details.has(span.span_id); + const mayHaveLaterChildren = hasMoreSpans && span.start_offset_ms + span.duration_ms >= pageBoundary; + return missingDetail || mayHaveLaterChildren; + }); + const initial = { + pending: new Set(pending.map((span) => span.span_id)), + ancestors: new Map( + spans.flatMap((span) => + span.parent_span_id && byId.has(span.parent_span_id) ? [[span.span_id, span.parent_span_id] as const] : [], + ), + ), + }; + const doublingPasses = Math.ceil(Math.log2(Math.max(1, spans.length))); + return Array.from({ length: doublingPasses }).reduce((state) => { + if (!state.pending.size || !state.ancestors.size) return state; + return { + pending: new Set([ + ...state.pending, + ...[...state.pending].flatMap((id) => { + const ancestor = state.ancestors.get(id); + return ancestor ? [ancestor] : []; + }), + ]), + ancestors: new Map( + [...state.ancestors].flatMap(([id, ancestor]) => { + const next = state.ancestors.get(ancestor); + return next ? [[id, next] as const] : []; + }), + ), + }; + }, initial).pending; +} + function contentText(value: string, content?: UIContent): string { if (content?.kind === "text") return content.text; if (content?.kind === "fields") @@ -147,7 +189,7 @@ function conversationEvents( steps: readonly Span[], byId: ReadonlyMap, details: ReadonlyMap, - complete: boolean, + { complete, pendingBranches }: { complete: boolean; pendingBranches?: ReadonlySet }, ): ConversationEvent[] { const missingIndex = steps.findIndex((span) => !details.has(span.span_id)); const loaded = missingIndex < 0 ? steps : steps.slice(0, missingIndex); @@ -170,7 +212,8 @@ function conversationEvents( const start = { span, time: messageTime(span, loaded, details), output: false }; if (span.type === "tool" || (span.type !== "agent" && span.parent_span_id !== null)) return [start]; const end = ends.get(span.span_id)!; - return (complete && missingIndex < 0) || end < boundary ? [start, { span, time: end, output: true }] : [start]; + const branchComplete = (complete && missingIndex < 0) || pendingBranches?.has(span.span_id) === false; + return branchComplete || end < boundary ? [start, { span, time: end, output: true }] : [start]; }) .sort( (a, b) => @@ -252,6 +295,7 @@ export function buildConversation( spans: readonly Span[], recordedDetails: ReadonlyMap, complete: boolean, + pendingBranches?: ReadonlySet, ): ConversationItem[] { const details = claudeToolDetails(spans, recordedDetails); const byId = new Map(spans.map((span) => [span.span_id, span])); @@ -259,7 +303,7 @@ export function buildConversation( const completedOutputs = new Map(); const pendingCalls = new Map(); const items: ConversationItem[] = []; - const events = conversationEvents(conversationSteps(spans), byId, details, complete); + const events = conversationEvents(conversationSteps(spans), byId, details, { complete, pendingBranches }); const branch = (span: Span): string => conversationBranch(span, byId); for (const event of events) { const { span } = event; @@ -354,6 +398,17 @@ export type ConversationGroup = | { kind: "item"; item: ConversationItem } | { kind: "branch"; id: string; name: string; children: ConversationGroup[] }; +export function conversationBranchGroups(groups: ConversationGroup[], branchId: string | null): ConversationGroup[] { + if (branchId === null) return groups; + for (const group of groups) { + if (group.kind === "item") continue; + if (group.id === branchId) return group.children; + const nested = conversationBranchGroups(group.children, branchId); + if (nested.length) return nested; + } + return []; +} + export function groupConversation(items: readonly ConversationItem[], spans: readonly Span[]): ConversationGroup[] { const byId = new Map(spans.map((span) => [span.span_id, span])); const parentById = new Map( diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/useConversationDetails.ts b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/useConversationDetails.ts new file mode 100644 index 00000000000..e94bb82cd98 --- /dev/null +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/useConversationDetails.ts @@ -0,0 +1,48 @@ +import { useQueries, useQueryClient } from "@tanstack/react-query"; +import { useCallback, useMemo, useState } from "react"; +import { useTracesApi } from "../../api"; +import type { Span, SpanDetail, Trace } from "../../types"; +import { CONVERSATION_PAGE_SIZE, conversationSteps } from "./conversation"; + +export function useConversationDetails(trace: Trace, accessToken: string) { + const traces = useTracesApi(accessToken); + const queryClient = useQueryClient(); + const [limit, setLimit] = useState(CONVERSATION_PAGE_SIZE); + const steps = useMemo(() => conversationSteps(trace.spans), [trace.spans]); + const visible = steps.slice(0, limit); + const { trace_id: traceId, trace_ref: traceRef } = trace.summary; + const spanQuery = useCallback( + (span: Span) => ({ + queryKey: ["agentTraceSpan", traceId, traceRef, span.span_id, accessToken], + queryFn: (): Promise => traces.span(traceId, span.span_id, traceRef), + staleTime: Infinity, + retry: false as const, + retryOnMount: false, + }), + [traceId, traceRef, accessToken, traces], + ); + const queries = useQueries({ queries: visible.map(spanQuery) }); + const unresolvedIndex = queries.findIndex((query) => !query.isSuccess); + const loadedCount = unresolvedIndex < 0 ? queries.length : unresolvedIndex; + const details = new Map( + queries.slice(0, loadedCount).map((query, index) => [visible[index].span_id, query.data!] as const), + ); + const loadMore = useCallback( + async (signal: AbortSignal) => { + const nextLimit = limit + CONVERSATION_PAGE_SIZE; + const batch = steps.slice(limit, nextLimit); + await Promise.all(batch.map((span) => queryClient.prefetchQuery(spanQuery(span)))); + if (!signal.aborted) setLimit(nextLimit); + }, + [limit, steps, queryClient, spanQuery], + ); + return { + details, + entries: visible.map((span, index) => ({ span, query: queries[index] })), + complete: loadedCount === steps.length, + loading: queries.some((query) => query.isPending), + failed: queries.some((query) => query.isError), + hasMore: visible.length < steps.length, + loadMore, + }; +} diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunBody.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunBody.tsx index 84f6e199072..b7019d07fae 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunBody.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunBody.tsx @@ -10,7 +10,7 @@ import { TabsContent } from "@/components/ui/tabs"; import type { RunSelection } from "../../routing"; import type { Trace } from "../../types"; -import { TraceConversation } from "../conversation/TraceConversation"; +import { TraceConversation, type ConversationTracePaging } from "../conversation/TraceConversation"; import { DetailPane } from "../span/DetailPane"; import { SpanTree } from "../tree/SpanTree"; import type { TreeLayout } from "../tree/TreeRows"; @@ -25,10 +25,11 @@ interface RunBodyProps { selection: RunSelection; embedded: boolean; stale: boolean; + conversationPaging: ConversationTracePaging; } /** Tree + detail pane for one loaded run. Arrows move and fold steps; J/K also move unless the drawer owns them. */ -export function RunBody({ trace, accessToken, selection, embedded, stale }: RunBodyProps) { +export function RunBody({ trace, accessToken, selection, embedded, stale, conversationPaging }: RunBodyProps) { const { view, setView, stepQuery, setStepQuery, errorsOnly, setErrorsOnly } = selection; const tree = useRunTree(trace, selection); const [detailOpen, setDetailOpen] = useState(true); @@ -58,6 +59,7 @@ export function RunBody({ trace, accessToken, selection, embedded, stale }: RunB { select(id); setView("steps"); diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunView.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunView.tsx index 4d78b136c17..d00481498f1 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunView.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunView.tsx @@ -158,6 +158,11 @@ function LoadedRun({ selection={selection} embedded={embedded} stale={switching} + conversationPaging={{ + loading: isFetching, + failed: isError, + loadMore: fetchNextPage, + }} /> );