From 4ec311e95b953ae45710dd8a8e3ce5128754c6c3 Mon Sep 17 00:00:00 2001
From: moe-berri <moe@berri.ai>
Date: Sat, 3 Oct 2026 20:06:16 -0700
Subject: [PATCH] fix(lens): preserve investigations during activity refresh
 errors

---
 .../lens/LensSetup.integration.test.tsx       | 25 ++++++++++++++++++-
 .../src/components/lens/LensWorkspace.tsx     |  2 +-
 .../src/components/lens/setup/useLensSetup.ts |  4 ++-
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/lens/LensSetup.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/LensSetup.integration.test.tsx
index 8d3b5742334..036b02f04c4 100644
--- a/ui/litellm-dashboard/src/components/lens/LensSetup.integration.test.tsx
+++ b/ui/litellm-dashboard/src/components/lens/LensSetup.integration.test.tsx
@@ -1,4 +1,4 @@
-import { fireEvent, screen, within, waitFor } from "@testing-library/react";
+import { act, fireEvent, screen, within, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { renderWithProviders, testQueryClient } from "@/../tests/test-utils";
@@ -187,6 +187,29 @@ describe("Lens setup journey", () => {
     expect(screen.queryByRole("alert")).not.toBeInTheDocument();
   });
 
+  it("keeps request-only users in investigations when an activity refresh fails", async () => {
+    serve({ requests: true, connected: true });
+    const user = userEvent.setup();
+    renderWithProviders(<LensWorkspace accessToken="setup-token" userRole="Admin" readOnly={false} />, {
+      searchParams: "?tab=investigations",
+    });
+    expect(await screen.findByText("Request logs received")).toBeVisible();
+    const normal = network.getMockImplementation()!;
+    network.mockImplementation((input, init) =>
+      new URL(String(input), "http://localhost").pathname === "/lens/activity/available"
+        ? Promise.resolve(Response.json({ detail: "Activity unavailable" }, { status: 503 }))
+        : normal(input, init),
+    );
+    await act(() => testQueryClient.refetchQueries());
+    expect(await screen.findByRole("alert")).toHaveTextContent("Could not check recorded activity");
+    expect(screen.queryByRole("heading", { name: "Get Lens running" })).not.toBeInTheDocument();
+    expect(screen.getByRole("button", { name: "New investigation" })).toBeDisabled();
+    network.mockImplementation(normal);
+    await user.click(screen.getByRole("button", { name: "Retry" }));
+    expect(await screen.findByText("Request logs received")).toBeVisible();
+    expect(screen.getByRole("button", { name: "New investigation" })).toBeEnabled();
+  });
+
   it("keeps administrator-only setup unavailable to trace viewers", async () => {
     serve({ enabled: true, traces: true });
     const user = userEvent.setup();
diff --git a/ui/litellm-dashboard/src/components/lens/LensWorkspace.tsx b/ui/litellm-dashboard/src/components/lens/LensWorkspace.tsx
index 18aa6439498..63bef720211 100644
--- a/ui/litellm-dashboard/src/components/lens/LensWorkspace.tsx
+++ b/ui/litellm-dashboard/src/components/lens/LensWorkspace.tsx
@@ -275,6 +275,6 @@ function needsSetup(
   if (requested) return true;
   if (!state.missingTraces) return false;
   if (tab === "traces") return true;
-  const hasActivity = state.hasInvestigations || state.requestsReady || selected;
+  const hasActivity = state.hasInvestigations || state.hasRequests || selected;
   return canInvestigate && !hasActivity;
 }
diff --git a/ui/litellm-dashboard/src/components/lens/setup/useLensSetup.ts b/ui/litellm-dashboard/src/components/lens/setup/useLensSetup.ts
index d92fdb9e708..12903e8a15c 100644
--- a/ui/litellm-dashboard/src/components/lens/setup/useLensSetup.ts
+++ b/ui/litellm-dashboard/src/components/lens/setup/useLensSetup.ts
@@ -23,7 +23,8 @@ export function useLensSetup(accessToken: string, enabled: boolean, canInvestiga
   const activity = useQuery(lensQueries.activity(api, enabled && canInvestigate && list.isSuccess));
   const data = list.data ?? { lenses: [], workers: [], tracing_enabled: false };
   const traceState = traceSetupState(traces, data.tracing_enabled);
-  const requestsReady = activity.data?.requests === true && !activity.error;
+  const hasRequests = activity.data?.requests === true;
+  const requestsReady = hasRequests && !activity.error;
   const connected = data.workers.some((worker) => workerConnected(worker, now));
   const hasInvestigations = data.lenses.length > 0;
   const activityReady = traceState.tracesReady || requestsReady;
@@ -42,6 +43,7 @@ export function useLensSetup(accessToken: string, enabled: boolean, canInvestiga
     tracingEnabled: traceState.tracingEnabled,
     tracesReady: traceState.tracesReady,
     requestsReady,
+    hasRequests,
     connected,
     hasInvestigations,
     missingTraces: traceState.missingTraces,