Merge pull request #41374 from BerriAI/litellm_fix_guardrail_lifecycle_untimed_entries

fix(ui): keep untimed guardrail entries on the request lifecycle

ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx

@@ -1,7 +1,7 @@
 import React from "react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import userEvent from "@testing-library/user-event";
-import { renderWithProviders, screen, waitFor } from "../../../../tests/test-utils";
+import { renderWithProviders, screen, waitFor, within } from "../../../../tests/test-utils";
 import {
   GuardrailInformation,
   makeBedrockResponse,
@@ -24,6 +24,42 @@ const skippedPreCall: Partial<GuardrailInformation> = {
   duration: null,
 };
 
+const untimedPreCall: Partial<GuardrailInformation> = {
+  guardrail_name: "conduct",
+  guardrail_status: "success",
+  guardrail_mode: "pre_call",
+  start_time: null,
+  end_time: null,
+  duration: null,
+};
+
+const timedPreCall: Partial<GuardrailInformation> = {
+  guardrail_name: "timed-pre-rail",
+  guardrail_status: "success",
+  guardrail_mode: "pre_call",
+  start_time: 1_700_000_000,
+  end_time: 1_700_000_000.1,
+  duration: 0.1,
+};
+
+const latePreCall: Partial<GuardrailInformation> = {
+  guardrail_name: "late-pre-rail",
+  guardrail_status: "success",
+  guardrail_mode: "pre_call",
+  start_time: 1_700_000_500,
+  end_time: 1_700_000_500.1,
+  duration: 0.1,
+};
+
+const untimedPostCall: Partial<GuardrailInformation> = {
+  guardrail_name: "untimed-post-rail",
+  guardrail_status: "success",
+  guardrail_mode: "post_call",
+  start_time: null,
+  end_time: null,
+  duration: null,
+};
+
 const ranPostCall: Partial<GuardrailInformation> = {
   guardrail_name: "ran-rail",
   guardrail_status: "success",
@@ -98,6 +134,67 @@ describe("GuardrailViewer", () => {
     expect(screen.getByText("—")).toBeInTheDocument();
   });
 
+  it("keeps a guardrail that ran without any timing on the lifecycle", () => {
+    renderWithProviders(<GuardrailViewer data={makeGuardrailInformation(untimedPreCall)} />);
+
+    expect(screen.getByText("Request received")).toBeInTheDocument();
+    expect(screen.getByText(/Pre-call guardrail: conduct/)).toBeInTheDocument();
+    expect(screen.getByText("LLM call")).toBeInTheDocument();
+    expect(screen.getByText("Response returned")).toBeInTheDocument();
+    expect(screen.queryByText(/^T\+/)).not.toBeInTheDocument();
+  });
+
+  it("keeps an untimed guardrail ahead of a timed one recorded after it in the same phase", () => {
+    const untimed = makeGuardrailInformation(untimedPreCall);
+    const timedPre = makeGuardrailInformation(timedPreCall);
+    renderWithProviders(<GuardrailViewer data={[untimed, timedPre]} />);
+
+    const rows = screen.getAllByTestId("lifecycle-row");
+    const rowIndex = (label: RegExp): number => rows.findIndex((r) => within(r).queryByText(label) !== null);
+    const untimedIndex = rowIndex(/Pre-call guardrail: conduct/);
+    const timedIndex = rowIndex(/Pre-call guardrail: timed-pre-rail/);
+
+    expect(untimedIndex).toBeGreaterThanOrEqual(0);
+    expect(timedIndex).toBeGreaterThanOrEqual(0);
+    expect(untimedIndex).toBeLessThan(timedIndex);
+  });
+
+  it("orders each phase on its own clock when a later pre-call outlives an earlier post-call", () => {
+    const latePre = makeGuardrailInformation(latePreCall);
+    const untimedPost = makeGuardrailInformation(untimedPostCall);
+    const earlyPost = makeGuardrailInformation(ranPostCall);
+    renderWithProviders(<GuardrailViewer data={[latePre, untimedPost, earlyPost]} />);
+
+    const rows = screen.getAllByTestId("lifecycle-row");
+    const rowIndex = (label: RegExp): number => rows.findIndex((r) => within(r).queryByText(label) !== null);
+    const untimedIndex = rowIndex(/Post-call guardrail: untimed-post-rail/);
+    const earlyIndex = rowIndex(/Post-call guardrail: ran-rail/);
+
+    expect(untimedIndex).toBeGreaterThanOrEqual(0);
+    expect(earlyIndex).toBeGreaterThanOrEqual(0);
+    expect(untimedIndex).toBeLessThan(earlyIndex);
+  });
+
+  it("anchors offsets on the timed entries and gives the untimed one no fabricated offset", () => {
+    const untimed = makeGuardrailInformation(untimedPreCall);
+    const ran = makeGuardrailInformation(ranPostCall);
+    renderWithProviders(<GuardrailViewer data={[untimed, ran]} />);
+
+    const lifecycleRow = (label: string | RegExp): HTMLElement => {
+      const row = screen.getAllByTestId("lifecycle-row").find((r) => within(r).queryByText(label) !== null);
+      if (row === undefined) throw new Error(`no lifecycle row labelled ${label}`);
+      return row;
+    };
+
+    expect(within(lifecycleRow("Request received")).getByText("T+0ms")).toBeInTheDocument();
+    expect(within(lifecycleRow(/Post-call guardrail: ran-rail/)).getByText("T+250ms")).toBeInTheDocument();
+    expect(within(lifecycleRow("Response returned")).getByText("T+251ms")).toBeInTheDocument();
+
+    const untimedRow = within(lifecycleRow(/Pre-call guardrail: conduct/));
+    expect(untimedRow.getByText("—")).toBeInTheDocument();
+    expect(untimedRow.queryByText(/^T\+/)).not.toBeInTheDocument();
+  });
+
   it("calculates and displays masked entity totals", async () => {
     const user = userEvent.setup();
     const data = makeGuardrailInformation({

ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx

@@ -361,7 +361,7 @@ const GenericGuardrailResponse = ({ response }: { response: any }) => {
 interface TimelineEntry {
   type: "request" | "guardrail" | "llm" | "response";
   label: string;
-  offsetMs: number;
+  offsetMs: number | null;
   outcome?: EntryOutcome;
 }
 
@@ -370,73 +370,85 @@ type TimedGuardrailInformation = GuardrailInformation & { start_time: number; en
 const isTimed = (e: GuardrailInformation): e is TimedGuardrailInformation =>
   typeof e.start_time === "number" && typeof e.end_time === "number";
 
+const belongsOnLifecycle = (e: GuardrailInformation): boolean => isTimed(e) || getEntryOutcome(e) !== "not_run";
+
+// Sorts a phase's timed entries by start time while leaving its untimed entries in the
+// slots they were recorded in. Applied per phase, never globally: an entry can land in
+// more than one phase bucket, so a global pass can reorder one phase by another's clock.
+const orderWithinPhase = (group: GuardrailInformation[]): GuardrailInformation[] => {
+  const byStart = group.filter(isTimed).sort((a, b) => a.start_time - b.start_time);
+  const timedSlots = new Map(group.flatMap((e, i) => (isTimed(e) ? [i] : [])).map((slot, k) => [slot, byStart[k]]));
+  return group.map((e, i) => timedSlots.get(i) ?? e);
+};
+
 const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => {
-  const sorted = useMemo(() => entries.filter(isTimed).sort((a, b) => a.start_time - b.start_time), [entries]);
+  const sorted = useMemo(() => entries.filter(belongsOnLifecycle), [entries]);
 
   const timeline = useMemo(() => {
     if (sorted.length === 0) return [];
 
-    const baseTime = sorted[0].start_time;
+    const timed = sorted.filter(isTimed);
+    const baseTime = timed.length > 0 ? Math.min(...timed.map((e) => e.start_time)) : null;
+    const offsetOf = (e: GuardrailInformation): number | null =>
+      baseTime === null || !isTimed(e) ? null : Math.round((e.end_time - baseTime) * 1000);
     const items: TimelineEntry[] = [];
 
     // Request received
-    items.push({ type: "request", label: "Request received", offsetMs: 0 });
+    items.push({ type: "request", label: "Request received", offsetMs: baseTime === null ? null : 0 });
 
     // Pre-call guardrails — use modeMatches so array modes (e.g. ["pre_call", "post_call"])
     // place the entry in every matching bucket.
-    const preCalls = sorted.filter((e) => modeMatches(e.guardrail_mode, "pre_call"));
-    const postCalls = sorted.filter(
-      (e) => modeMatches(e.guardrail_mode, "post_call") || modeMatches(e.guardrail_mode, "logging_only"),
+    const preCalls = orderWithinPhase(sorted.filter((e) => modeMatches(e.guardrail_mode, "pre_call")));
+    const postCalls = orderWithinPhase(
+      sorted.filter((e) => modeMatches(e.guardrail_mode, "post_call") || modeMatches(e.guardrail_mode, "logging_only")),
     );
-    const duringCalls = sorted.filter((e) => modeMatches(e.guardrail_mode, "during_call"));
+    const duringCalls = orderWithinPhase(sorted.filter((e) => modeMatches(e.guardrail_mode, "during_call")));
 
     for (const e of preCalls) {
-      const offsetMs = Math.round((e.end_time - baseTime) * 1000);
       items.push({
         type: "guardrail",
         label: `Pre-call guardrail: ${getDisplayName(e)}`,
-        offsetMs,
+        offsetMs: offsetOf(e),
         outcome: getEntryOutcome(e),
       });
     }
 
     // LLM call — infer from gap between pre-call end and post-call start
-    const lastPreEnd = preCalls.length > 0 ? Math.max(...preCalls.map((e) => e.end_time)) : baseTime;
-    const firstPostStart = postCalls.length > 0 ? Math.min(...postCalls.map((e) => e.start_time)) : undefined;
-    const llmEndTime = firstPostStart ?? lastPreEnd + 1;
-    const llmOffsetMs = Math.round((llmEndTime - baseTime) * 1000);
+    const timedPre = preCalls.filter(isTimed);
+    const timedPost = postCalls.filter(isTimed);
+    const lastPreEnd = timedPre.length > 0 ? Math.max(...timedPre.map((e) => e.end_time)) : baseTime;
+    const firstPostStart = timedPost.length > 0 ? Math.min(...timedPost.map((e) => e.start_time)) : undefined;
+    const llmEndTime = firstPostStart ?? (lastPreEnd === null ? null : lastPreEnd + 1);
 
     items.push({
       type: "llm",
       label: "LLM call",
-      offsetMs: llmOffsetMs,
+      offsetMs: llmEndTime === null || baseTime === null ? null : Math.round((llmEndTime - baseTime) * 1000),
     });
 
     // During-call guardrails (rare)
     for (const e of duringCalls) {
-      const offsetMs = Math.round((e.end_time - baseTime) * 1000);
       items.push({
         type: "guardrail",
         label: `During-call guardrail: ${getDisplayName(e)}`,
-        offsetMs,
+        offsetMs: offsetOf(e),
         outcome: getEntryOutcome(e),
       });
     }
 
     // Post-call guardrails
     for (const e of postCalls) {
-      const offsetMs = Math.round((e.end_time - baseTime) * 1000);
       items.push({
         type: "guardrail",
         label: `Post-call guardrail: ${getDisplayName(e)}`,
-        offsetMs,
+        offsetMs: offsetOf(e),
         outcome: getEntryOutcome(e),
       });
     }
 
     // Response returned
-    const maxEnd = Math.max(...sorted.map((e) => e.end_time));
-    const responseOffsetMs = Math.round((maxEnd - baseTime) * 1000) + 1;
+    const maxEnd = timed.length > 0 ? Math.max(...timed.map((e) => e.end_time)) : null;
+    const responseOffsetMs = maxEnd === null || baseTime === null ? null : Math.round((maxEnd - baseTime) * 1000) + 1;
     items.push({ type: "response", label: "Response returned", offsetMs: responseOffsetMs });
 
     return items;
@@ -447,7 +459,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => {
       <h4 className="text-xs font-semibold text-muted-foreground uppercase tracking-wider mb-4">Request Lifecycle</h4>
       <div className="relative">
         {timeline.map((item, idx) => (
-          <div key={idx} className="flex items-start gap-3 relative">
+          <div key={idx} data-testid="lifecycle-row" className="flex items-start gap-3 relative">
             {/* Vertical line */}
             <div className="flex flex-col items-center">
               <div className="shrink-0">
@@ -475,7 +487,9 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => {
                     {OUTCOME_LABEL[item.outcome]}
                   </span>
                 )}
-                <span className="text-xs text-muted-foreground font-mono ml-auto shrink-0">T+{item.offsetMs}ms</span>
+                <span className="text-xs text-muted-foreground font-mono ml-auto shrink-0">
+                  {item.offsetMs === null ? "—" : `T+${item.offsetMs}ms`}
+                </span>
               </div>
             </div>
           </div>