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

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) => {

ui/litellm-dashboard/src/components/lens/traces/utils.ts

@@ -115,35 +115,21 @@ type GroupOrSpan = Span | { group: Span[] };
 
 const groupKey = (span: Pick<Span, "name" | "type" | "agent">): 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<string, Span[]>();
+  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<string>();
-  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<Span, "name" | "type" | "agent">): string =>
-  `grp::${parentKey}::${groupKey(span)}`;
+export const groupRowId = (parentKey: string, span: Pick<Span, "span_id">): string =>
+  `grp::${parentKey}::${span.span_id}`;
 
 interface RowContext {
   children: ChildrenMap;