From fe29c55d11ab9ef26bf605be104eea372fddef0d Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 6 Oct 2026 13:51:58 -0700 Subject: [PATCH] fix(lens): preserve chronological order when grouping trace steps (#44895) * fix(lens): preserve chronological order when grouping trace steps * test(lens): cover independent paging of repeated trace groups --- .../src/components/lens/traces/utils.test.ts | 78 +++++++++++++++++++ .../src/components/lens/traces/utils.ts | 34 +++----- 2 files changed, 88 insertions(+), 24 deletions(-) diff --git a/ui/litellm-dashboard/src/components/lens/traces/utils.test.ts b/ui/litellm-dashboard/src/components/lens/traces/utils.test.ts index d074a5c56fe..c5dbd88d292 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/utils.test.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/utils.test.ts @@ -133,6 +133,51 @@ describe("buildVisibleTree / isFrameworkSpan", () => { }); describe("buildTreeRows", () => { + it("preserves alternating model and tool calls instead of collecting nonadjacent siblings", () => { + const calls = Array.from({ length: 24 }, (_, index) => + span({ + span_id: `step-${index}`, + parent_span_id: "root", + start_offset_ms: index, + name: index % 2 === 0 ? "model" : "Bash", + type: index % 2 === 0 ? "llm" : "tool", + }), + ); + const rows = buildTreeRows([span({ span_id: "root", type: "agent" }), ...calls.toReversed()], STATE); + expect(groups(rows)).toHaveLength(0); + expect( + spanRows(rows) + .slice(1) + .map((row) => row.span.span_id), + ).toEqual(calls.map((call) => call.span_id)); + }); + + it("keeps separated repetitions independently expandable and reveals a selected later span", () => { + const calls = Array.from({ length: 13 }, (_, index) => + span({ + span_id: `step-${index}`, + parent_span_id: "root", + start_offset_ms: index, + name: index === 6 ? "model" : "Bash", + type: index === 6 ? "llm" : "tool", + }), + ); + const spans = [span({ span_id: "root", type: "agent" }), ...calls]; + const folded = buildTreeRows(spans, STATE); + expect(groups(folded).map((group) => group.members.map((member) => member.span_id))).toEqual([ + calls.slice(0, 6).map((call) => call.span_id), + calls.slice(7).map((call) => call.span_id), + ]); + expect(groups(folded)[0].id).not.toBe(groups(folded)[1].id); + const revealed = buildTreeRows(spans, revealSpanInState(spans, STATE, "step-12")); + expect(groups(revealed).map((group) => group.expanded)).toEqual([false, true]); + expect(spanRows(revealed).map((row) => row.span.span_id)).toEqual([ + "root", + "step-6", + ...calls.slice(7).map((call) => call.span_id), + ]); + }); + it("starts the tree at the root span", () => { const rows = buildTreeRows(research.spans, STATE); expect(rows[0]).toMatchObject({ @@ -204,6 +249,39 @@ describe("buildTreeRows", () => { expect(groups(buildTreeRows([parent, ...kids], STATE))).toHaveLength(0); }); + it("pages a later repeated group independently while an earlier group stays collapsed", () => { + const calls = Array.from({ length: 61 }, (_, index) => + span({ + span_id: `step-${index}`, + parent_span_id: "root", + start_offset_ms: index, + name: index === 30 ? "model" : "Bash", + type: index === 30 ? "llm" : "tool", + }), + ); + const spans = [span({ span_id: "root", type: "agent" }), ...calls]; + const laterId = groupRowId("root", calls[31]); + const state = { ...STATE, expandedGroupIds: new Set([laterId]) }; + const firstPage = buildTreeRows(spans, state); + expect(groups(firstPage).map((group) => group.expanded)).toEqual([false, true]); + expect(firstPage.filter((row) => row.kind === "load-more")).toEqual([ + expect.objectContaining({ groupId: laterId, remaining: 10 }), + ]); + expect(spanRows(firstPage).map((row) => row.span.span_id)).toEqual([ + "root", + "step-30", + ...calls.slice(31, 51).map((call) => call.span_id), + ]); + const nextPage = buildTreeRows(spans, { ...state, groupRevealCounts: { [laterId]: 40 } }); + expect(groups(nextPage).map((group) => group.expanded)).toEqual([false, true]); + expect(nextPage.filter((row) => row.kind === "load-more")).toHaveLength(0); + expect(spanRows(nextPage).map((row) => row.span.span_id)).toEqual([ + "root", + "step-30", + ...calls.slice(31).map((call) => call.span_id), + ]); + }); + it("pages expanded groups 20 at a time with a load-more row", () => { const parent = span({ span_id: "p", type: "agent" }); const kids = Array.from({ length: 45 }, (_, i) => { diff --git a/ui/litellm-dashboard/src/components/lens/traces/utils.ts b/ui/litellm-dashboard/src/components/lens/traces/utils.ts index a71fdb329eb..63ae58486cc 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/utils.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/utils.ts @@ -115,35 +115,21 @@ type GroupOrSpan = Span | { group: Span[] }; const groupKey = (span: Pick): string => `${span.agent}|${span.name}|${span.type}`; -/** Siblings sharing agent + name + type fold into one group once there are enough of them (or enough failures). */ function groupChildren(children: readonly Span[]): GroupOrSpan[] { - const byKey = new Map(); + const runs: Span[][] = []; for (const child of children) { - const key = groupKey(child); - const list = byKey.get(key); - if (list) list.push(child); - else byKey.set(key, [child]); + const previous = runs.at(-1); + if (previous && groupKey(previous[0]) === groupKey(child)) previous.push(child); + else runs.push([child]); } - const emitted = new Set(); - const out: GroupOrSpan[] = []; - for (const child of children) { - const key = groupKey(child); - const group = byKey.get(key) ?? [child]; - const failed = group.filter((s) => s.status === "error").length; - if (group.length >= GROUP_THRESHOLD_OK || failed >= GROUP_THRESHOLD_ERROR) { - if (!emitted.has(key)) { - emitted.add(key); - out.push({ group }); - } - } else { - out.push(child); - } - } - return out; + return runs.flatMap((group): GroupOrSpan[] => { + const failed = group.filter((span) => span.status === "error").length; + return group.length >= GROUP_THRESHOLD_OK || failed >= GROUP_THRESHOLD_ERROR ? [{ group }] : group; + }); } -export const groupRowId = (parentKey: string, span: Pick): string => - `grp::${parentKey}::${groupKey(span)}`; +export const groupRowId = (parentKey: string, span: Pick): string => + `grp::${parentKey}::${span.span_id}`; interface RowContext { children: ChildrenMap;