typesafety

This commit is contained in:
Yujong Lee 2026-10-04 12:26:04 -07:00
parent 35437dd7f5
commit df0e426692
15 changed files with 418 additions and 244 deletions

View file

@ -20,6 +20,7 @@ from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
from litellm.proxy.lens.billing import validate_key
from litellm.proxy.lens.inference import Deployment, deployment_prices
from litellm.proxy.lens.models import (
ActivitySelection,
Claim,
Execution,
ExecutionContent,
@ -129,7 +130,7 @@ def required(lens: Lens | None) -> Lens:
return lens
def validate_selection(settings: LensSettings) -> None:
def validate_selection(settings: ActivitySelection) -> None:
for identity in settings.execution_ids:
try:
source, _, _, _ = parse_execution(identity)
@ -391,13 +392,13 @@ async def update_finding(lens_id: str, finding_id: str, body: FindingUpdate, aut
class Preview(BaseModel):
as_of: AwareDatetime | None = None
offset: int = Field(default=0, ge=0)
settings: LensSettings
selection: ActivitySelection
lookback_hours: LookbackHours = 24
@router.post("/preview/sample", response_model=Sample)
async def preview_sample(body: Preview, auth: Auth, storage: StorageDep) -> Sample:
validate_selection(body.settings)
validate_selection(body.selection)
now: Final = min(body.as_of or datetime.now(timezone.utc), datetime.now(timezone.utc))
try:
start: Final = int((now - timedelta(hours=body.lookback_hours)).timestamp() * 1000)
@ -406,7 +407,7 @@ async def preview_sample(body: Preview, auth: Auth, storage: StorageDep) -> Samp
raise HTTPException(422, "Preview window exceeds the supported calendar range") from error
return await source_reader(storage).sample(
user_scope(auth),
body.settings,
body.selection,
start,
end,
offset=body.offset,

View file

@ -45,23 +45,26 @@ class Check(Record):
enabled: bool = True
class LensSettings(Record):
name: str = Field(min_length=1)
context: str = Field(default="")
class ActivitySelection(Record):
source: Literal["traces", "requests", "both"] = "traces"
lookback_hours: LookbackHours = 24
service: str = Field(default="")
agent_name: str = Field(default="")
filters: tuple[MetadataFilter, ...] = Field(default=())
sample_size: int | None = Field(default=None, ge=1)
sample_percent: float = Field(default=100, gt=0, le=100, allow_inf_nan=False)
team_id: str = ""
execution_ids: tuple[str, ...] = ()
class LensSettings(ActivitySelection):
name: str = Field(min_length=1)
context: str = Field(default="")
lookback_hours: LookbackHours = 24
checks: tuple[Check, ...] = ()
model: str = Field(min_length=1)
enabled: bool = True
interval_minutes: IntervalMinutes = 15
sample_size: int | None = Field(default=None, ge=1)
sample_percent: float = Field(default=100, gt=0, le=100, allow_inf_nan=False)
concurrency: int = Field(default=8, ge=1)
team_id: str = ""
execution_ids: tuple[str, ...] = ()
monthly_budget: float = Field(default=100, gt=0, allow_inf_nan=False)
@model_validator(mode="after")

View file

@ -6,10 +6,10 @@ from typing import Final, Protocol, TypeAlias
from pydantic import TypeAdapter
from litellm.proxy.lens.models import (
ActivitySelection,
Evidence,
Execution,
ExecutionContent,
LensSettings,
MetadataFilter,
Sample,
Scope,
@ -73,7 +73,7 @@ class SourceReader:
async def sample(
self,
scope: Scope,
settings: LensSettings,
settings: ActivitySelection,
start: int,
end: int,
offset: int = 0,

View file

@ -18,7 +18,7 @@ from litellm.proxy.lens.endpoints import (
watching,
worker_supports_model,
)
from litellm.proxy.lens.models import Lens, LensSettings, RunRequest, Scope
from litellm.proxy.lens.models import ActivitySelection, Lens, LensSettings, RunRequest, Scope
@pytest.fixture
@ -281,6 +281,25 @@ def test_model_errors_reach_worker_with_status_and_redacted_provider_message(pro
assert error.headers == {"retry-after": "60"}
@pytest.mark.asyncio
async def test_preview_samples_a_selection_without_investigation_settings() -> None:
from litellm.proxy.lens.endpoints import Preview, preview_sample
class SelectionStorage:
async def lens_sample(self, parameters):
assert (parameters.source, parameters.agent_name, parameters.selected_team) == ("requests", "billing", "t1")
assert parameters.preview == 1 and parameters.offset == 3
return []
body: Final = Preview.model_validate(
{"selection": {"source": "requests", "agent_name": "billing", "team_id": "t1"}, "offset": 3}
)
sample: Final = await preview_sample(
body, UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), SelectionStorage()
)
assert sample.eligible == 0 and not sample.executions
@pytest.mark.asyncio
async def test_preview_reports_calendar_overflow_as_a_validation_error() -> None:
from datetime import datetime, timezone
@ -288,7 +307,7 @@ async def test_preview_reports_calendar_overflow_as_a_validation_error() -> None
from litellm.proxy.lens.endpoints import Preview, preview_sample
body: Final = Preview(
settings=LensSettings(name="Calendar regression", model="analysis", context="Read recorded activity"),
selection=ActivitySelection(),
as_of=datetime.min.replace(tzinfo=timezone.utc),
)
with pytest.raises(HTTPException) as error:
@ -323,7 +342,9 @@ async def test_unknown_gateway_release_refuses_registration_and_claims(monkeypat
monkeypatch.setenv("LITELLM_RELEASE_TAG", "")
monkeypatch.setenv("LENS_WORKER_IMAGE", "registry.example/lens-worker:old")
with pytest.raises(HTTPException) as registration_error:
await register_worker(WorkerName(analysis_key_id="a" * 64), UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN))
await register_worker(
WorkerName(analysis_key_id="a" * 64), UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
)
assert registration_error.value.status_code == 503
assert "LITELLM_RELEASE_TAG" in registration_error.value.detail
with pytest.raises(HTTPException) as claim_error:

View file

@ -2,7 +2,7 @@ import { screen } 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";
import { dismissLensIntro } from "@/../tests/lens-test-utils";
import { dismissLensIntro, requestPath } from "@/../tests/lens-test-utils";
import LensPage from "@/app/(dashboard)/lens/page";
const { auth } = vi.hoisted(() => ({ auth: vi.fn() }));
@ -26,7 +26,7 @@ describe("Lens navigation", () => {
vi.stubGlobal(
"fetch",
vi.fn(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/v1/traces") return Response.json({ data: [{}] });
if (path === "/lens") return Response.json({ lenses: [], workers: [], tracing_enabled: true });
return Response.json({ traces: true, requests: false, data: [] });

View file

@ -2,7 +2,7 @@ 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";
import { dismissLensIntro } from "@/../tests/lens-test-utils";
import { dismissLensIntro, readRequest, requestPath } from "@/../tests/lens-test-utils";
import { readStorage } from "@/lib/storage";
import { LENS_INTRO_DISMISSED, LENS_INTRO_SEEN } from "./storage";
import { LensWorkspace } from "./LensWorkspace";
@ -23,14 +23,14 @@ const worker = () => ({
function serve({ enabled = false, traces = false, requests = false, connected = false } = {}) {
list.mockResolvedValue({ lenses: [], workers: connected ? [worker()] : [], tracing_enabled: enabled });
network.mockImplementation(async (input, init) => {
const path = new URL(String(input), "http://localhost").pathname;
const { path, method, body } = await readRequest(input, init);
if (path === "/v1/traces")
return enabled
? Response.json({ data: traces ? [data.runs[0].trace.summary] : [] })
: Response.json({ detail: "Tracing is not enabled" }, { status: 501 });
if (path === "/lens/activity/available") return Response.json({ traces, requests });
if (path === "/lens" && init?.method === "POST") {
const saved = { ...data.lenses[0], settings: { ...data.lenses[0].settings, ...JSON.parse(String(init.body)) } };
if (path === "/lens" && method === "POST") {
const saved = { ...data.lenses[0], settings: { ...data.lenses[0].settings, ...(body as object) } };
list.mockResolvedValue({ lenses: [saved], workers: [worker()], tracing_enabled: true });
return Response.json(saved);
}
@ -154,9 +154,7 @@ describe("Lens setup journey", () => {
serve({ enabled: true, traces: true });
const normal = network.getMockImplementation()!;
network.mockImplementation((input, init) =>
new URL(String(input), "http://localhost").pathname === pendingPath
? new Promise<Response>(() => {})
: normal(input, init),
requestPath(input) === pendingPath ? new Promise<Response>(() => {}) : normal(input, init),
);
renderWorkspace();
expect(await screen.findByRole("table", { name: "Agent runs" })).toBeVisible();
@ -171,9 +169,7 @@ describe("Lens setup journey", () => {
list.mockResolvedValue({ lenses: data.lenses, workers: [worker()], tracing_enabled: false });
const normal = network.getMockImplementation()!;
network.mockImplementation((input, init) =>
new URL(String(input), "http://localhost").pathname === pendingPath
? new Promise<Response>(() => {})
: normal(input, init),
requestPath(input) === pendingPath ? new Promise<Response>(() => {}) : normal(input, init),
);
renderWorkspace({ searchParams: `?lens=${data.lenses[0].id}` });
expect(await screen.findByRole("heading", { name: data.lenses[0].settings.name })).toBeVisible();
@ -247,7 +243,7 @@ describe("Lens setup journey", () => {
await intro.findByRole("button", { name: "Check for traces" });
const normal = network.getMockImplementation()!;
network.mockImplementation(async (input, init) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/v1/traces") return Response.json({ detail: "Trace storage unavailable" }, { status: 503 });
return normal(input, init);
});
@ -268,7 +264,7 @@ describe("Lens setup journey", () => {
expect(await screen.findByRole("button", { name: "New investigation" })).toBeEnabled();
const normal = network.getMockImplementation()!;
network.mockImplementation((input, init) =>
new URL(String(input), "http://localhost").pathname === "/lens/activity/available"
requestPath(input) === "/lens/activity/available"
? Promise.resolve(Response.json({ detail: "Activity unavailable" }, { status: 503 }))
: normal(input, init),
);
@ -288,9 +284,7 @@ describe("Lens setup journey", () => {
const intro = within(await screen.findByRole("dialog"));
expect(await intro.findByText(/A gateway administrator can connect a worker/)).toBeVisible();
expect(intro.getByRole("button", { name: "Connect worker" })).toBeDisabled();
expect(network.mock.calls.some(([input]) => new URL(String(input), "http://localhost").pathname === "/lens")).toBe(
false,
);
expect(network.mock.calls.some(([input]) => requestPath(input) === "/lens")).toBe(false);
});
it.each(["traces", "requests with trace errors", "requests with pending traces", "traces with activity errors"])(
@ -302,7 +296,7 @@ describe("Lens setup journey", () => {
const normal = network.getMockImplementation()!;
const failingPath = scenario === "requests with trace errors" ? "/v1/traces" : "/lens/activity/available";
network.mockImplementation((input, init) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/v1/traces" && scenario === "requests with pending traces")
return new Promise<Response>(() => {});
if (scenario.endsWith("errors") && path === failingPath)
@ -328,13 +322,10 @@ describe("Lens setup journey", () => {
within(screen.getByRole("tablist", { name: "Lens" })).getByRole("tab", { name: "Investigations" }),
).toHaveAttribute("aria-selected", "true");
await waitFor(() => expect(setupParam(onUrlUpdate)).toBeNull());
const create = network.mock.calls.find(
([input, init]) => new URL(String(input), "http://localhost").pathname === "/lens" && init?.method === "POST",
);
const requests = await Promise.all(network.mock.calls.map(([input, init]) => readRequest(input, init)));
const create = requests.find((request) => request.path === "/lens" && request.method === "POST");
expect(create).toBeDefined();
expect(JSON.parse(String(create?.[1]?.body))).toEqual(
expect.objectContaining({ name: "My first review", source }),
);
expect(create?.body).toEqual(expect.objectContaining({ name: "My first review", source }));
},
);
});

View file

@ -2,7 +2,7 @@ import { screen, waitFor, within } 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";
import { dismissLensIntro } from "@/../tests/lens-test-utils";
import { dismissLensIntro, readRequest, requestPath } from "@/../tests/lens-test-utils";
import { LensWorkspace } from "./LensWorkspace";
import { lensKeys } from "./data/queries";
import { createLensDemoData } from "./data/demo/fixtures";
@ -21,7 +21,7 @@ beforeEach(() => {
vi.stubGlobal("fetch", network);
network.mockReset();
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 });
if (path === "/lens") return Response.json({ lenses: [], workers: [], tracing_enabled: false });
return Response.json({ data: [], traces: false, requests: false });
@ -173,7 +173,7 @@ describe("Lens interactive demo", () => {
const data = createLensDemoData();
const saved = data.lenses[0];
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/lens") return Response.json({ lenses: [saved], workers: [], tracing_enabled: true });
if (path.endsWith("/runs")) return Response.json(saved.jobs);
if (path === "/v1/traces") return Response.json({ data: data.runs.map((run) => run.trace.summary) });
@ -194,7 +194,7 @@ describe("Lens interactive demo", () => {
const user = userEvent.setup();
const saved = createLensDemoData().lenses[0];
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/lens") return Response.json({ lenses: [saved], workers: [], tracing_enabled: false });
if (path.endsWith("/runs")) return Response.json(saved.jobs);
if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 });
@ -220,7 +220,7 @@ describe("Lens interactive demo", () => {
});
const lenses = vi.fn(() => [withJob("running")]);
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/lens") return Response.json({ lenses: lenses(), workers: [], tracing_enabled: false });
if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 });
return Response.json({ data: [], traces: false, requests: false });
@ -237,7 +237,7 @@ describe("Lens interactive demo", () => {
const user = userEvent.setup();
const saved = createLensDemoData().lenses[0];
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/lens") return Response.json({ lenses: [saved], workers: [], tracing_enabled: true });
if (path.endsWith("/runs")) return Response.json(saved.jobs);
if (path === "/lens/agents") return Response.json([]);
@ -272,7 +272,7 @@ describe("Lens interactive demo", () => {
};
const workers = vi.fn(() => [worker]);
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/lens") return Response.json({ lenses: [saved], workers: workers(), tracing_enabled: true });
if (path.endsWith("/runs")) return Response.json(saved.jobs);
if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 });
@ -308,7 +308,7 @@ describe("Lens interactive demo", () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn();
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/lens") return Response.json({ lenses: [], workers: [], tracing_enabled: true });
if (path === "/v1/traces") return Response.json({ data: [{}] });
return Response.json({ data: [], traces: true, requests: false });
@ -339,9 +339,9 @@ describe("Lens interactive demo", () => {
};
const workers = vi.fn((): (typeof worker)[] => []);
network.mockImplementation(async (input, init) => {
const path = new URL(String(input), "http://localhost").pathname;
const { path, method } = await readRequest(input, init);
if (path === "/lens") return Response.json({ lenses: [], workers: workers(), tracing_enabled: true });
if (path === "/lens/workers/register" && init?.method === "POST") {
if (path === "/lens/workers/register" && method === "POST") {
workers.mockReturnValue([worker]);
return Response.json({ token: "lens-test-token", image: "lens-worker:v1", worker });
}
@ -405,7 +405,7 @@ describe("Lens interactive demo", () => {
last_seen: new Date().toISOString(),
};
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/lens") return Response.json({ lenses: [saved], workers: [worker], tracing_enabled: true });
if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 });
return Response.json({ data: [], traces: true, requests: false });
@ -434,10 +434,9 @@ describe("Lens interactive demo", () => {
last_seen: new Date(Date.now() - 600_000).toISOString(),
};
const workers = vi.fn(() => [worker]);
const listCalls = () =>
network.mock.calls.filter(([input]) => new URL(String(input), "http://localhost").pathname === "/lens").length;
const listCalls = () => network.mock.calls.filter(([input]) => requestPath(input) === "/lens").length;
network.mockImplementation(async (input) => {
const path = new URL(String(input), "http://localhost").pathname;
const path = requestPath(input);
if (path === "/lens") return Response.json({ lenses: [saved], workers: workers(), tracing_enabled: true });
if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 });
return Response.json({ data: [], traces: true, requests: false });

View file

@ -2,6 +2,7 @@
import { createContext, useContext, useMemo, type ReactNode } from "react";
import { apiClient } from "@/components/networking";
import { fetchClient } from "@/lib/http/api";
import { liveTracesApi, TracesApiContext, type TracesApi } from "@/components/lens/traces/api";
import { liveLensApi, type LensApi } from "./service";
@ -14,7 +15,7 @@ export interface LensServices {
const LensServicesContext = createContext<LensServices | null>(null);
export function liveLensServices(accessToken: string): LensServices {
return { accessToken, lens: liveLensApi(apiClient, accessToken), traces: liveTracesApi(accessToken) };
return { accessToken, lens: liveLensApi(fetchClient, apiClient, accessToken), traces: liveTracesApi(accessToken) };
}
function useLensServices(): LensServices {

View file

@ -1,6 +1,8 @@
import { z } from "zod";
import type { ApiClient } from "@/lib/http/client";
import type { components } from "@/lib/http/schema";
import { getAuthHeaderName } from "@/lib/http/runtime";
import type { Client } from "openapi-fetch";
import type { components, paths } from "@/lib/http/schema";
import type {
ActivitySelection,
AnalysisModelInfo,
@ -42,7 +44,7 @@ export interface LensApi {
/** Partitions query caches between backends (one token, or the demo). */
readonly scope: string;
lenses(): Promise<LensList>;
activity(): Promise<{ traces: boolean; requests: boolean }>;
activity(): Promise<components["schemas"]["ActivityAvailability"]>;
runs(lensId: string, offset: number): Promise<Job[]>;
run(lensId: string, jobId: string): Promise<Job>;
execution(lensId: string, executionId: string, offset: number): Promise<ExecutionContent>;
@ -57,45 +59,73 @@ export interface LensApi {
watchAll(): Promise<components["schemas"]["WatchAllResult"]>;
cancelRun(lensId: string): Promise<void>;
reviewFinding(lensId: string, findingId: string, status: FindingStatus, reason: string): Promise<void>;
registerWorker(analysisKeyId: string | null): Promise<WorkerCreated>;
setWorkerBillingKey(workerId: string, analysisKeyId: string | null): Promise<void>;
registerWorker(analysisKeyId: string): Promise<WorkerCreated>;
setWorkerBillingKey(workerId: string, analysisKeyId: string): Promise<void>;
revokeWorker(workerId: string): Promise<void>;
generateAnalysisKey(request: AnalysisKeyRequest): Promise<{ token_id?: string }>;
deleteKeys(keys: readonly string[]): Promise<void>;
}
export function liveLensApi(apiClient: ApiClient, accessToken: string): LensApi {
const encode = encodeURIComponent;
type LensClient = Client<paths>;
async function required<T>(request: Promise<{ data?: T }>): Promise<T> {
const { data } = await request;
if (data === undefined) throw new Error("The proxy returned an empty response");
return data;
}
async function sent(request: Promise<unknown>): Promise<void> {
await request;
}
export function liveLensApi(client: LensClient, apiClient: ApiClient, accessToken: string): LensApi {
const headers = { [getAuthHeaderName()]: `Bearer ${accessToken}` };
const lens = (lens_id: string) => ({ headers, params: { path: { lens_id } } });
const worker = (worker_id: string) => ({ headers, params: { path: { worker_id } } });
return {
scope: accessToken,
lenses: () => apiClient.get<LensList>("/lens", { accessToken }),
activity: () => apiClient.get("/lens/activity/available", { accessToken }),
runs: (lensId, offset) => apiClient.get<Job[]>(`/lens/${lensId}/runs`, { accessToken, query: { offset } }),
run: (lensId, jobId) => apiClient.get<Job>(`/lens/${lensId}/runs/${jobId}`, { accessToken }),
lenses: () => required(client.GET("/lens", { headers })),
activity: () => required(client.GET("/lens/activity/available", { headers })),
runs: (lensId, offset) =>
required(
client.GET("/lens/{lens_id}/runs", { headers, params: { path: { lens_id: lensId }, query: { offset } } }),
),
run: (lensId, jobId) =>
required(
client.GET("/lens/{lens_id}/runs/{job_id}", {
headers,
params: { path: { lens_id: lensId, job_id: jobId } },
}),
),
execution: (lensId, executionId, offset) =>
apiClient.get<ExecutionContent>(`/lens/${lensId}/executions/${encode(executionId)}`, {
accessToken,
query: { offset },
}),
sample: (selection, offset, asOf) => {
const { lookback_hours, ...selectionSettings } = selection;
return apiClient.post<Sample>("/lens/preview/sample", {
accessToken,
body: {
offset,
as_of: asOf,
settings: {
...selectionSettings,
execution_ids: [],
name: "Preview",
model: "preview",
checks: [{ id: "preview", instruction: "Preview recorded activity" }],
required(
client.GET("/lens/{lens_id}/executions/{execution_id}", {
headers,
params: { path: { lens_id: lensId, execution_id: executionId }, query: { offset } },
}),
),
sample: (selection, offset, asOf) =>
required(
client.POST("/lens/preview/sample", {
headers,
body: {
offset,
as_of: asOf,
selection: {
source: selection.source,
service: selection.service ?? "",
agent_name: selection.agent_name ?? "",
filters: selection.filters ?? [],
sample_size: selection.sample_size,
sample_percent: selection.sample_percent ?? 100,
team_id: selection.team_id ?? "",
execution_ids: [],
},
lookback_hours: selection.lookback_hours ?? 24,
},
lookback_hours: lookback_hours ?? 24,
},
});
},
agents: () => apiClient.get<string[]>("/lens/agents", { accessToken }),
}),
),
agents: () => required(client.GET("/lens/agents", { headers })),
models: () => apiClient.get("/models", { accessToken }),
modelDetails: () => apiClient.get("/model_group/info", { accessToken }),
keys: async (alias, page, signal) =>
@ -118,21 +148,37 @@ export function liveLensApi(apiClient: ApiClient, accessToken: string): LensApi
keyInfo: async (keyId) =>
keyInfoSchema.parse(await apiClient.get("/key/info", { accessToken, query: { key: keyId } })).info,
saveLens: (id, settings) =>
apiClient.request<Lens>(id ? "PUT" : "POST", id ? `/lens/${id}` : "/lens", { accessToken, body: settings }),
startRun: (lensId, request = {}) => apiClient.post(`/lens/${lensId}/runs`, { accessToken, body: request }),
watchAll: () =>
apiClient.post<components["schemas"]["WatchAllResult"]>("/lens/watch-all", { accessToken, body: {} }),
cancelRun: (lensId) => apiClient.post(`/lens/${lensId}/cancel`, { accessToken, body: {} }),
required(
id
? client.PUT("/lens/{lens_id}", { ...lens(id), body: settings })
: client.POST("/lens", { headers, body: settings }),
),
startRun: (lensId, request = {}) => sent(client.POST("/lens/{lens_id}/runs", { ...lens(lensId), body: request })),
watchAll: () => required(client.POST("/lens/watch-all", { headers })),
cancelRun: (lensId) => sent(client.POST("/lens/{lens_id}/cancel", lens(lensId))),
reviewFinding: (lensId, findingId, status, reason) =>
apiClient.patch(`/lens/${lensId}/findings/${findingId}`, { accessToken, body: { status, reason } }),
sent(
client.PATCH("/lens/{lens_id}/findings/{finding_id}", {
headers,
params: { path: { lens_id: lensId, finding_id: findingId } },
body: { status, reason },
}),
),
registerWorker: (analysisKeyId) =>
apiClient.post<WorkerCreated>("/lens/workers/register", {
accessToken,
body: { name: "Lens worker", analysis_key_id: analysisKeyId },
}),
required(
client.POST("/lens/workers/register", {
headers,
body: { name: "Lens worker", analysis_key_id: analysisKeyId },
}),
),
setWorkerBillingKey: (workerId, analysisKeyId) =>
apiClient.put(`/lens/workers/${workerId}/billing-key`, { accessToken, body: { analysis_key_id: analysisKeyId } }),
revokeWorker: (workerId) => apiClient.delete(`/lens/workers/${workerId}`, { accessToken }),
sent(
client.PUT("/lens/workers/{worker_id}/billing-key", {
...worker(workerId),
body: { analysis_key_id: analysisKeyId },
}),
),
revokeWorker: (workerId) => sent(client.DELETE("/lens/workers/{worker_id}", worker(workerId))),
generateAnalysisKey: (request) =>
apiClient.post<{ token_id?: string }>("/key/generate", {
accessToken,

View file

@ -2,9 +2,8 @@ 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 { testQueryClient } from "@/../tests/test-utils";
import { renderWithLens } from "@/../tests/lens-test-utils";
import { renderWithLens, stubGateway } from "@/../tests/lens-test-utils";
import { ApiError } from "@/lib/http/client";
import { apiClient } from "@/components/networking";
import { lensKeys } from "../data/queries";
import { InvestigationsView } from "./InvestigationsView";
import { LensPreviewContext } from "@/components/lens/ui/LensPreviewButton";
@ -21,16 +20,20 @@ function renderWithProviders(ui: React.ReactElement, options?: Parameters<typeof
return renderWithLens(ui, { searchParams: window.location.search, ...options });
}
vi.mock("@/components/networking", () => ({
apiClient: { get: vi.fn(), post: vi.fn(), patch: vi.fn(), request: vi.fn() },
vi.mock("@/components/networking", async (importOriginal) => ({
...(await importOriginal<typeof import("@/components/networking")>()),
proxyBaseUrl: "",
getProxyBaseUrl: () => "",
}));
let proxy = stubGateway();
const sentBody = (handler: typeof proxy.post, path: string) =>
handler.mock.calls.filter(([called]) => called === path).map(([, request]) => request.body);
beforeEach(() => {
window.history.replaceState({}, "", "/lens/?lens=lens");
vi.mocked(apiClient.post).mockReset();
vi.mocked(apiClient.post).mockResolvedValue({ eligible: 0, selected: 0, executions: [] });
proxy = stubGateway();
proxy.post.mockResolvedValue({ eligible: 0, selected: 0, executions: [] });
});
const executionId = btoa(JSON.stringify(["traces", "", "trace-42"]));
@ -161,8 +164,8 @@ const lens: Lens = {
describe("Lens findings and runs", () => {
beforeEach(() => {
vi.mocked(apiClient.get).mockReset();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockReset();
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return lens.jobs;
return { data: [] };
@ -197,7 +200,7 @@ describe("Lens findings and runs", () => {
async function openIssue(finding: Finding) {
testQueryClient.clear();
const jobs = lens.jobs.map((job) => ({ ...job, findings: [finding] }));
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens")
return { lenses: [{ ...lens, findings: [finding], jobs }], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return jobs;
@ -245,7 +248,7 @@ describe("Lens findings and runs", () => {
it("closes the open run when the keyboard switches to another investigation run", async () => {
testQueryClient.clear();
const older = { ...lens.jobs[0], id: "older", created_at: "2026-09-29T10:00:00Z" };
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return [...lens.jobs, older];
return { data: [] };
@ -269,7 +272,7 @@ describe("Lens findings and runs", () => {
it("runs with saved settings from Run now without opening setup, then accepts an agent and window", async () => {
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens")
return {
lenses: [lens],
@ -289,32 +292,29 @@ it("runs with saved settings from Run now without opening setup, then accepts an
if (path === "/lens/activity/available") return { traces: true, requests: false };
return { data: [] };
});
vi.mocked(apiClient.post).mockResolvedValue(lens);
proxy.post.mockResolvedValue(lens);
const user = userEvent.setup();
renderWithProviders(<InvestigationsView />);
await user.click(await screen.findByRole("button", { name: "Run now" }));
const choices = await screen.findByRole("dialog", { name: "Run now" });
expect(within(choices).getByRole("button", { name: "Since last run" })).toHaveAttribute("aria-pressed", "true");
await user.click(within(choices).getByRole("button", { name: "Run now" }));
expect(apiClient.post).toHaveBeenCalledWith("/lens/lens/runs", { accessToken: "test", body: {} });
expect(sentBody(proxy.post, "/lens/lens/runs")).toEqual([{}]);
await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument());
vi.mocked(apiClient.post).mockClear();
proxy.post.mockClear();
await user.click(screen.getByRole("button", { name: "Run now" }));
const custom = await screen.findByRole("dialog", { name: "Run now" });
fireEvent.change(within(custom).getByRole("combobox", { name: "Agent" }), { target: { value: "billing" } });
await user.click(within(custom).getByRole("button", { name: "Last 24h" }));
await user.click(within(custom).getByRole("button", { name: "Run now" }));
expect(apiClient.post).toHaveBeenCalledWith("/lens/lens/runs", {
accessToken: "test",
body: { agent_name: "billing", lookback_hours: 24 },
});
expect(sentBody(proxy.post, "/lens/lens/runs")).toEqual([{ agent_name: "billing", lookback_hours: 24 }]);
});
it("offers the interactive demo without starting an investigation", async () => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: false };
return { traces: false, requests: false };
});
@ -323,13 +323,13 @@ it("offers the interactive demo without starting an investigation", async () =>
renderWithProviders(withPreview(<InvestigationsView />, onPreview));
await user.click(await screen.findByRole("button", { name: "Preview sample" }));
expect(onPreview).toHaveBeenCalledOnce();
expect(apiClient.post).not.toHaveBeenCalled();
expect(proxy.post).not.toHaveBeenCalled();
});
it("guides a first-time administrator into worker connection and lens setup", async () => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: true };
if (path === "/lens/agents") return [];
return { traces: true, requests: false, data: [] };
@ -341,7 +341,10 @@ it("guides a first-time administrator into worker connection and lens setup", as
onboarding: { connect, create },
});
const guide = within(await screen.findByRole("region", { name: "Get Lens running" }));
expect(apiClient.get).toHaveBeenCalledWith("/lens/activity/available", { accessToken: "test" });
expect(proxy.get).toHaveBeenCalledWith(
"/lens/activity/available",
expect.objectContaining({ authorization: "Bearer test" }),
);
expect(guide.getByRole("button", { name: /Send your first trace/ })).toContainElement(
guide.getByLabelText("Step 2 complete"),
);
@ -382,7 +385,7 @@ it("opens the saved results of an older batch", async () => {
finished_at: "2026-09-29T10:02:13Z",
findings: [{ ...issue, title: "Earlier batch finding" }],
};
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return [lens.jobs[0], older];
if (path === "/lens/lens/runs/older") return older;
@ -413,11 +416,11 @@ it("reads request content from the beginning after its abbreviated preview", asy
executions: [{ ...lens.jobs[0].sample!.executions[0], id: requestId, source: "requests" as const }],
},
};
vi.mocked(apiClient.get).mockImplementation(async (path, options) => {
proxy.get.mockImplementation(async (path, options) => {
if (path === "/lens") return { lenses: [{ ...lens, jobs: [job] }], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return [job];
if (!path.includes("/executions/")) return { data: [] };
const offset = options?.query?.offset ?? 0;
const offset = Number(options.query.offset ?? 0);
return {
parts: [
{
@ -453,7 +456,7 @@ it.each([false, true])(
async (enabled) => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) =>
proxy.get.mockImplementation(async (path) =>
path === "/lens" ? { lenses: [], workers: [], tracing_enabled: enabled } : { data: [] },
);
const user = userEvent.setup();
@ -472,7 +475,7 @@ it("enables first-lens setup when a trace arrives without leaving Investigations
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
const traceCheck = vi.fn().mockResolvedValue({ traces: false, requests: false });
vi.mocked(apiClient.get).mockImplementation(async (path) =>
proxy.get.mockImplementation(async (path) =>
path === "/lens" ? { lenses: [], workers: [], tracing_enabled: true } : traceCheck(),
);
vi.useFakeTimers();
@ -502,7 +505,7 @@ it("allows retrying a failed trace readiness check without treating it as an emp
.fn()
.mockRejectedValueOnce(new ApiError("Trace storage unavailable", 503, {}))
.mockResolvedValue({ data: [] });
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: true };
if (path === "/lens/activity/available") return traceCheck();
return { data: [] };
@ -523,7 +526,7 @@ it("shows a centered failure with a retry when investigations cannot load, then
.fn()
.mockRejectedValueOnce(new ApiError("Proxy timed out", 504, {}))
.mockResolvedValue({ lenses: [], workers: [], tracing_enabled: true });
vi.mocked(apiClient.get).mockImplementation(async (path) => (path === "/lens" ? list() : { data: [] }));
proxy.get.mockImplementation(async (path) => (path === "/lens" ? list() : { data: [] }));
const user = userEvent.setup();
renderWithProviders(<InvestigationsView />);
const alert = await screen.findByRole("alert");
@ -536,7 +539,7 @@ it("shows a centered failure with a retry when investigations cannot load, then
it("keeps saved investigations accessible when tracing is disabled", async () => {
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: false };
if (path === "/lens/lens/runs") return lens.jobs;
return { data: [] };
@ -549,7 +552,7 @@ it("keeps saved investigations accessible when tracing is disabled", async () =>
it("allows request-only accounts to connect a worker without requiring agent traces", async () => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: true };
if (path === "/lens/activity/available") return { traces: false, requests: true };
return { data: [] };
@ -565,7 +568,7 @@ it("allows request-only accounts to connect a worker without requiring agent tra
it("reopens the inline editor from a shared link and drops it from the URL on cancel", async () => {
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return lens.jobs;
if (path === "/lens/agents") return [];
@ -587,12 +590,13 @@ it("reopens the inline editor from a shared link and drops it from the URL on ca
const url = new URLSearchParams(String(onUrlUpdate.mock.lastCall?.[0].queryString ?? ""));
expect(url.has("dialog")).toBe(false);
expect(url.get("lens")).toBe(lens.id);
expect(apiClient.request).not.toHaveBeenCalled();
expect(proxy.put).not.toHaveBeenCalled();
expect(sentBody(proxy.post, "/lens")).toEqual([]);
});
it("reopens a finding and a results section from shared links", async () => {
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return lens.jobs;
return { data: [] };
@ -628,7 +632,7 @@ it("steps across findings and investigations with J and K, skipping hidden findi
findings: [twinIssue],
jobs: lens.jobs.map((job) => ({ ...job, findings: [twinIssue] })),
};
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens, twin], tracing_enabled: true, workers: [] };
if (path === "/lens/activity/available") return { traces: true, requests: false };
if (path.endsWith("/runs")) return [];
@ -665,7 +669,7 @@ it("steps across findings and investigations with J and K, skipping hidden findi
it("opens an investigation beside the list and walks from it into its findings with J and K", async () => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], tracing_enabled: true, workers: [] };
if (path === "/lens/activity/available") return { traces: true, requests: false };
if (path.endsWith("/runs")) return [];
@ -698,13 +702,13 @@ it("lists each finding under the investigation that owns it and resolves only th
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
const twin: Lens = { ...lens, id: "twin", settings: { ...lens.settings, name: "Twin reviews" } };
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens, twin], tracing_enabled: true, workers: [] };
if (path === "/lens/activity/available") return { traces: true, requests: false };
if (path.endsWith("/runs")) return [];
return { data: [] };
});
vi.mocked(apiClient.patch).mockResolvedValue(undefined);
proxy.patch.mockResolvedValue(undefined);
const user = userEvent.setup();
renderWithProviders(<InvestigationsView />);
const rows = await screen.findAllByRole("row", { name: issue.title });
@ -714,14 +718,14 @@ it("lists each finding under the investigation that owns it and resolves only th
expect(remaining.previousElementSibling).toBe(screen.getByRole("row", { name: twin.settings.name }));
await user.click(remaining);
await user.click(await screen.findByRole("button", { name: "Mark resolved" }));
await waitFor(() => expect(apiClient.patch).toHaveBeenCalledTimes(1));
expect(vi.mocked(apiClient.patch).mock.calls[0][0]).toBe("/lens/twin/findings/issue");
await waitFor(() => expect(proxy.patch).toHaveBeenCalledTimes(1));
expect(proxy.patch.mock.calls[0][0]).toBe("/lens/twin/findings/issue");
});
it("lists investigations without edit or run controls for read-only viewers", async () => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], tracing_enabled: true, workers: [] };
if (path === "/lens/activity/available") return { traces: true, requests: false };
if (path === "/lens/lens/runs") return [];
@ -740,7 +744,7 @@ it("lists investigations without edit or run controls for read-only viewers", as
it("opens investigations from the keyboard without treating nested edit keys as row activation", async () => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], tracing_enabled: true, workers: [] };
if (path === "/lens/activity/available") return { traces: true, requests: false };
if (path === "/lens/lens/runs") return [];
@ -779,7 +783,7 @@ it("opens a failed investigation's details from its row and edits only from the
error: "boom",
findings: [],
};
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [{ ...lens, jobs: [job] }], workers: [], tracing_enabled: true };
if (path === "/lens/activity/available") return { traces: true, requests: false };
if (path === "/lens/lens/runs") return [job];
@ -804,7 +808,7 @@ it("shows the actual saved failure and run context without opening backend logs"
"Grouping observations failed: Clusters response invalid after 2 attempts.\n" +
"candidates.0.check_id: Field required [missing]";
const job = { ...lens.jobs[0], id: "failed-run", status: "failed" as const, stage: "Failed", error, findings: [] };
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [{ ...lens, jobs: [job] }], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return [job];
if (path === "/lens/lens/runs/failed-run") return job;
@ -822,14 +826,14 @@ it("keeps a finding open to retry when its update fails", async () => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
const twin: Lens = { ...lens, id: "twin", settings: { ...lens.settings, name: "Twin reviews" } };
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens, twin], tracing_enabled: true, workers: [] };
if (path === "/lens/activity/available") return { traces: true, requests: false };
if (path.endsWith("/runs")) return [];
return { data: [] };
});
vi.mocked(apiClient.patch).mockReset();
vi.mocked(apiClient.patch).mockImplementation(async (path) => {
proxy.patch.mockReset();
proxy.patch.mockImplementation(async (path) => {
if (String(path).startsWith("/lens/twin/")) throw new Error("Twin reviews could not be updated");
});
const user = userEvent.setup();
@ -844,28 +848,24 @@ it("keeps a finding open to retry when its update fails", async () => {
it("pauses monitoring from the detail menu by saving the investigation with monitoring off", async () => {
testQueryClient.clear();
const watching: Lens = { ...lens, settings: { ...lens.settings, enabled: true } };
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [watching], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return watching.jobs;
return { data: [] };
});
vi.mocked(apiClient.request).mockReset();
vi.mocked(apiClient.request).mockResolvedValue({ ...watching, settings: lens.settings });
proxy.put.mockResolvedValue({ ...watching, settings: lens.settings });
const user = userEvent.setup();
renderWithProviders(<InvestigationsView />);
await user.click(await screen.findByRole("button", { name: "Investigation actions" }));
await user.click(await screen.findByRole("menuitem", { name: "Pause monitoring" }));
await waitFor(() => expect(apiClient.request).toHaveBeenCalledTimes(1));
expect(apiClient.request).toHaveBeenCalledWith("PUT", "/lens/lens", {
accessToken: "test",
body: { ...watching.settings, enabled: false },
});
await waitFor(() => expect(proxy.put).toHaveBeenCalledTimes(1));
expect(sentBody(proxy.put, "/lens/lens")).toEqual([{ ...watching.settings, enabled: false }]);
});
it("cancels the running job from the progress banner", async () => {
testQueryClient.clear();
const running = { ...lens.jobs[0], id: "live", status: "running" as const, stage: "Reading executions" };
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens")
return { lenses: [{ ...lens, jobs: [running, lens.jobs[0]] }], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return [running, lens.jobs[0]];
@ -874,15 +874,13 @@ it("cancels the running job from the progress banner", async () => {
const user = userEvent.setup();
renderWithProviders(<InvestigationsView />);
await user.click(await screen.findByRole("button", { name: "Cancel" }));
await waitFor(() =>
expect(apiClient.post).toHaveBeenCalledWith("/lens/lens/cancel", { accessToken: "test", body: {} }),
);
await waitFor(() => expect(proxy.post).toHaveBeenCalledWith("/lens/lens/cancel", expect.anything()));
});
it("refreshes run history as soon as the list reports a job the scheduler started", async () => {
testQueryClient.clear();
const runs = vi.fn().mockResolvedValue(lens.jobs);
vi.mocked(apiClient.get).mockImplementation(async (path) => {
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true };
if (path === "/lens/lens/runs") return runs();
return { data: [] };

View file

@ -3,18 +3,32 @@ import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { useQuery } from "@tanstack/react-query";
import { testQueryClient } from "@/../tests/test-utils";
import { renderWithLens } from "@/../tests/lens-test-utils";
import { apiClient } from "@/components/networking";
import { readRequest, renderWithLens, type StubbedRequest } from "@/../tests/lens-test-utils";
import { lensQueries } from "../../data/queries";
import { useLensApi } from "../../data/LensServices";
import type { LensList } from "../../model/types";
import { WorkerSettings } from "./WorkerSettings";
vi.mock("@/components/networking", () => ({
apiClient: { get: vi.fn(), post: vi.fn(), put: vi.fn(), delete: vi.fn() },
vi.mock("@/components/networking", async (importOriginal) => ({
...(await importOriginal<typeof import("@/components/networking")>()),
proxyBaseUrl: "https://gateway.example/proxy",
}));
type Route = (request: StubbedRequest) => unknown;
const sent: StubbedRequest[] = [];
const network = vi.fn<typeof fetch>();
function serve(route: Route) {
network.mockImplementation(async (input, init) => {
const request = await readRequest(input, init);
sent.push(request);
const result = route(request);
return result instanceof Response ? result : Response.json(result ?? null);
});
}
const calls = (method: string, path: string) =>
sent.filter((request) => request.method === method && request.path === path);
const writes = () => sent.filter((request) => request.method !== "GET");
const created = {
token: "lens-test-token",
image: "ghcr.io/berriai/litellm-lens-worker:v1.2.3",
@ -35,23 +49,22 @@ function WorkerSettingsHost() {
return list.data ? <WorkerSettings workers={list.data.workers} /> : null;
}
const listCalls = () => vi.mocked(apiClient.get).mock.calls.filter(([path]) => path === "/lens").length;
const listCalls = () => calls("GET", "/lens").length;
const keyRoute: Route = ({ path }) =>
path === "/key/info"
? { info: { models: ["analysis-model"], max_budget: 15, budget_duration: "1mo" } }
: { keys: [{ token: "b".repeat(64), key_alias: "Analysis" }], total_pages: 1 };
describe("Worker setup", () => {
beforeEach(() => {
testQueryClient.clear();
vi.clearAllMocks();
vi.mocked(apiClient.get).mockImplementation(async (path) =>
path === "/key/info"
? { info: { models: ["analysis-model"], max_budget: 15, budget_duration: "1mo" } }
: {
keys: [{ token: "b".repeat(64), key_alias: "Analysis" }],
total_pages: 1,
},
);
sent.length = 0;
network.mockReset();
vi.stubGlobal("fetch", network);
serve(keyRoute);
});
it("generates a complete command using one worker credential and the configured proxy address", async () => {
vi.mocked(apiClient.post).mockResolvedValue(created);
serve((request) => (request.path === "/lens/workers/register" ? created : keyRoute(request)));
const user = userEvent.setup();
const { rerender } = renderWithLens(<WorkerSettings workers={[]} />, { accessToken: "admin" });
await user.click(screen.getByText("Advanced options"));
@ -61,10 +74,9 @@ describe("Worker setup", () => {
await user.click(screen.getByRole("combobox", { name: "Charge analysis to" }));
await user.click(await screen.findByRole("option", { name: "Analysis" }));
await user.click(screen.getByRole("button", { name: "Get install command" }));
expect(apiClient.post).toHaveBeenCalledWith("/lens/workers/register", {
accessToken: "admin",
body: { name: "Lens worker", analysis_key_id: "b".repeat(64) },
});
expect(calls("POST", "/lens/workers/register").map(({ body }) => body)).toEqual([
{ name: "Lens worker", analysis_key_id: "b".repeat(64) },
]);
expect(screen.getByRole("status")).toHaveTextContent("Waiting for your worker to connect");
expect(screen.getByLabelText("Docker command preview")).not.toBeVisible();
await user.click(screen.getByRole("button", { name: "Copy Docker command" }));
@ -87,14 +99,13 @@ describe("Worker setup", () => {
it("assigns billing to an existing worker without replacing its access token", async () => {
const user = userEvent.setup();
const workers = vi.fn((): LensList["workers"] => [{ ...created.worker, analysis_key_id: null }]);
vi.mocked(apiClient.get).mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [], workers: workers(), tracing_enabled: true };
if (path === "/key/info") return { info: { models: ["analysis-model"], max_budget: 15, budget_duration: "1mo" } };
return { keys: [{ token: "b".repeat(64), key_alias: "Analysis" }], total_pages: 1 };
});
vi.mocked(apiClient.put).mockImplementation(async () => {
workers.mockReturnValue([created.worker]);
return created.worker;
serve((request) => {
if (request.method === "PUT") {
workers.mockReturnValue([created.worker]);
return created.worker;
}
if (request.path === "/lens") return { lenses: [], workers: workers(), tracing_enabled: true };
return keyRoute(request);
});
renderWithLens(<WorkerSettingsHost />, { accessToken: "admin" });
expect(await screen.findByText("Billing key required")).toBeInTheDocument();
@ -103,48 +114,49 @@ describe("Worker setup", () => {
await user.click(screen.getByRole("combobox", { name: "Charge analysis to" }));
await user.click(await screen.findByRole("option", { name: "Analysis" }));
await user.click(screen.getByRole("button", { name: "Save analysis access" }));
expect(apiClient.put).toHaveBeenCalledWith("/lens/workers/worker/billing-key", {
accessToken: "admin",
body: { analysis_key_id: "b".repeat(64) },
});
expect(calls("PUT", "/lens/workers/worker/billing-key").map(({ body }) => body)).toEqual([
{ analysis_key_id: "b".repeat(64) },
]);
expect(await screen.findByText(/Not connected/)).toBeVisible();
expect(screen.queryByText("Billing key required")).not.toBeInTheDocument();
expect(listCalls()).toBe(listedBefore + 1);
expect(apiClient.post).not.toHaveBeenCalled();
expect(writes().filter((request) => request.method === "POST")).toEqual([]);
});
it("requires revoking the current worker before setting up a replacement", async () => {
const user = userEvent.setup();
const workers = vi.fn((): LensList["workers"] => [created.worker]);
vi.mocked(apiClient.get).mockImplementation(async (path) =>
path === "/lens" ? { lenses: [], workers: workers(), tracing_enabled: true } : { data: [] },
);
vi.mocked(apiClient.delete).mockImplementation(async () => {
workers.mockReturnValue([{ ...created.worker, revoked: true }]);
return true;
serve(({ method, path }) => {
if (method === "DELETE") {
workers.mockReturnValue([{ ...created.worker, revoked: true }]);
return true;
}
return path === "/lens" ? { lenses: [], workers: workers(), tracing_enabled: true } : { data: [] };
});
renderWithLens(<WorkerSettingsHost />, { accessToken: "admin" });
const revoke = await screen.findByRole("button", { name: "Revoke access" });
expect(screen.queryByRole("button", { name: "Add worker" })).not.toBeInTheDocument();
expect(screen.queryByRole("button", { name: "Get install command" })).not.toBeInTheDocument();
await user.click(revoke);
expect(apiClient.delete).toHaveBeenCalledWith("/lens/workers/worker", { accessToken: "admin" });
expect(calls("DELETE", "/lens/workers/worker")).toHaveLength(1);
expect(await screen.findByRole("button", { name: "Get install command" })).toBeDisabled();
expect(screen.getByRole("combobox", { name: "Analysis model" })).toBeVisible();
expect(listCalls()).toBe(2);
});
it("cleans up a newly created key when registration fails before retrying", async () => {
const user = userEvent.setup();
vi.mocked(apiClient.get).mockImplementation(async (path) => {
const posts = [
Response.json({ token_id: "limited-key-id", key: "sk-secret-not-displayed" }),
Response.json({ detail: "Registration unavailable" }, { status: 503 }),
Response.json({}),
Response.json({ token_id: "retry-key-id" }),
Response.json(created),
];
serve(({ method, path }) => {
if (method === "POST") return posts.shift();
if (path === "/models") return { data: [{ id: "analysis-model" }] };
if (path === "/lens") return { workers: [] };
return { keys: [], total_pages: 0 };
});
vi.mocked(apiClient.post)
.mockResolvedValueOnce({ token_id: "limited-key-id", key: "sk-secret-not-displayed" })
.mockRejectedValueOnce(new Error("Registration unavailable"))
.mockResolvedValueOnce({})
.mockResolvedValueOnce({ token_id: "retry-key-id" })
.mockResolvedValueOnce(created);
renderWithLens(<WorkerSettings workers={[]} />, { accessToken: "admin" });
expect(screen.getByRole("button", { name: "Get install command" })).toBeDisabled();
expect(screen.getByRole("textbox", { name: "LiteLLM proxy URL", hidden: true })).not.toBeVisible();
@ -154,8 +166,8 @@ describe("Worker setup", () => {
await user.type(screen.getByLabelText("Monthly limit (USD)"), "12");
await user.click(screen.getByRole("button", { name: "Get install command" }));
expect(await screen.findByRole("alert")).toHaveTextContent("Registration unavailable");
expect(apiClient.post).toHaveBeenNthCalledWith(1, "/key/generate", {
accessToken: "admin",
expect(writes()[0]).toMatchObject({
path: "/key/generate",
body: {
key_alias: "Lens analysis",
models: ["analysis-model"],
@ -166,15 +178,15 @@ describe("Worker setup", () => {
});
await user.click(screen.getByRole("button", { name: "Get install command" }));
expect(await screen.findByRole("status")).toHaveTextContent("Waiting for your worker");
expect(apiClient.post).toHaveBeenCalledWith("/key/delete", {
accessToken: "admin",
body: { keys: ["limited-key-id"] },
});
expect(apiClient.post).toHaveBeenCalledTimes(5);
expect(apiClient.post).toHaveBeenLastCalledWith("/lens/workers/register", {
accessToken: "admin",
body: { name: "Lens worker", analysis_key_id: "retry-key-id" },
});
expect(calls("POST", "/key/delete").map(({ body }) => body)).toEqual([{ keys: ["limited-key-id"] }]);
expect(writes().map(({ path }) => path)).toEqual([
"/key/generate",
"/lens/workers/register",
"/key/delete",
"/key/generate",
"/lens/workers/register",
]);
expect(writes().at(-1)?.body).toEqual({ name: "Lens worker", analysis_key_id: "retry-key-id" });
expect(screen.queryByText("sk-secret-not-displayed")).not.toBeInTheDocument();
});
});

View file

@ -32,7 +32,8 @@ async function releaseUnusedKey(api: LensApi, keyId: string): Promise<void> {
async function prepareWorker(api: LensApi, registration: WorkerRegistration): Promise<WorkerCreated | null> {
const { address, useExisting, analysisKey, access, editingWorker } = registration;
validateWorkerAddress(address);
const keyId = useExisting ? analysisKey : await createAnalysisKey(api, access);
if (useExisting && !analysisKey) throw new Error("Choose an existing key");
const keyId = useExisting && analysisKey ? analysisKey : await createAnalysisKey(api, access);
const newKey = useExisting ? null : keyId;
try {
if (editingWorker) {

View file

@ -3,14 +3,13 @@ import userEvent from "@testing-library/user-event";
import { mockAllIsIntersecting, setupIntersectionMocking } from "react-intersection-observer/test-utils";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { testQueryClient } from "@/../tests/test-utils";
import { renderWithLens as renderWithProviders } from "@/../tests/lens-test-utils";
import { renderWithLens as renderWithProviders, stubGateway } from "@/../tests/lens-test-utils";
import { MonitoringDialog } from "./MonitoringDialog";
import { InvestigationSetup } from "./InvestigationSetup";
import { apiClient } from "@/components/networking";
import { initialWatches, watchChecks } from "../model/watches";
import { type AnalysisModelInfo, type Settings } from "../model/types";
vi.mock("@/components/networking", () => ({ apiClient: { post: vi.fn(), get: vi.fn() } }));
let proxy = stubGateway();
const settings: Settings = {
lookback_hours: 24,
@ -65,21 +64,20 @@ function gatewayResponse(path: string, gateway: Gateway): unknown {
}
function mockGateway(gateway: Gateway = {}) {
vi.mocked(apiClient.get).mockImplementation(async (path) => gatewayResponse(path, gateway));
proxy.get.mockImplementation(async (path) => gatewayResponse(path, gateway));
}
beforeEach(() => {
testQueryClient.clear();
setupIntersectionMocking(vi.fn);
vi.mocked(apiClient.get).mockReset();
proxy = stubGateway();
mockGateway();
vi.mocked(apiClient.post).mockReset();
vi.mocked(apiClient.post).mockResolvedValue({ eligible: 1, selected: 1, executions: [] });
proxy.post.mockResolvedValue({ eligible: 1, selected: 1, executions: [] });
});
describe("Investigation setup", () => {
it("preserves saved manual run selections when editing and lets the preview footer clear them", async () => {
const user = userEvent.setup();
vi.mocked(apiClient.post).mockResolvedValue({
proxy.post.mockResolvedValue({
eligible: 2,
selected: 2,
executions: [
@ -132,9 +130,9 @@ describe("Investigation setup", () => {
it("previews identifiable matching runs and saves the same filter selection", async () => {
const save = vi.fn().mockResolvedValue(undefined);
const user = userEvent.setup();
vi.mocked(apiClient.post).mockImplementation(async (_path, options) => {
const body = options?.body as { settings: Settings };
return body.settings.filters?.some((f) => f.key === "swarm" && f.value === "research")
proxy.post.mockImplementation(async (_path, options) => {
const body = options?.body as { selection: Settings };
return body.selection.filters?.some((f) => f.key === "swarm" && f.value === "research")
? {
eligible: 1,
selected: 1,
@ -308,8 +306,8 @@ it("keeps the draft when readiness changes and blocks a run until the worker rec
});
it.each(["empty", "error"])("allows editing saved settings when the preview is %s", async (state) => {
if (state === "error") vi.mocked(apiClient.post).mockRejectedValue(new Error("Storage unavailable"));
else vi.mocked(apiClient.post).mockResolvedValue({ eligible: 0, selected: 0, executions: [] });
if (state === "error") proxy.post.mockRejectedValue(new Error("Storage unavailable"));
else proxy.post.mockResolvedValue({ eligible: 0, selected: 0, executions: [] });
const user = userEvent.setup();
renderWithProviders(<InvestigationSetup mode="edit" initial={settings} onClose={vi.fn()} onSave={vi.fn()} />);
await user.click(screen.getByRole("button", { name: "Continue" }));
@ -322,7 +320,7 @@ it.each(["empty", "error"])("allows editing saved settings when the preview is %
it.each(["loading", "error"])("saves edits with the existing model while models are %s", async (state) => {
const user = userEvent.setup();
const save = vi.fn().mockResolvedValue(undefined);
vi.mocked(apiClient.get).mockImplementation((path) => {
proxy.get.mockImplementation((path) => {
if (path !== "/models") return Promise.resolve(gatewayResponse(path, {}));
return state === "loading" ? new Promise(() => {}) : Promise.reject(new Error("Temporarily unavailable"));
});
@ -341,7 +339,7 @@ it.each(["loading", "error"])("saves edits with the existing model while models
it.each(["new", "duplicate"] as const)("blocks a %s investigation until its model is verified", async (mode) => {
const user = userEvent.setup();
const save = vi.fn();
vi.mocked(apiClient.get).mockImplementation((path) =>
proxy.get.mockImplementation((path) =>
path === "/models"
? Promise.reject(new Error("Temporarily unavailable"))
: Promise.resolve(gatewayResponse(path, {})),
@ -397,7 +395,7 @@ it("appends the next preview page as the list scrolls near its end, then stops a
span_count: 2,
});
let finishSecondPage = (): void => {};
vi.mocked(apiClient.post).mockImplementation((_path, options) => {
proxy.post.mockImplementation((_path, options) => {
const { offset } = options?.body as { offset: number };
const firstPage = { eligible: 2, selected: 2, executions: [run("one")], next_offset: 1 };
const secondPage = { eligible: 2, selected: 2, executions: [run("two")], next_offset: null };
@ -411,7 +409,7 @@ it("appends the next preview page as the list scrolls near its end, then stops a
expect(screen.getByText(/Showing 1 of 2/)).toBeVisible();
expect(screen.getByRole("status")).toHaveTextContent("2 matching runs");
const nextPageCalls = () =>
vi.mocked(apiClient.post).mock.calls.filter(([, options]) => (options?.body as { offset: number }).offset === 1);
proxy.post.mock.calls.filter(([, options]) => (options?.body as { offset: number }).offset === 1);
expect(nextPageCalls()).toHaveLength(0);
act(() => mockAllIsIntersecting(true));
await waitFor(() => expect(nextPageCalls()).toHaveLength(1));
@ -428,7 +426,7 @@ it("appends the next preview page as the list scrolls near its end, then stops a
it("does not silently analyze everything after individual selection is enabled", async () => {
const user = userEvent.setup();
vi.mocked(apiClient.post).mockResolvedValue({
proxy.post.mockResolvedValue({
eligible: 1,
selected: 1,
executions: [
@ -456,13 +454,13 @@ it("refreshes agent suggestions when the first activity arrives", async () => {
vi.useFakeTimers({ shouldAdvanceTime: true });
try {
const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
vi.mocked(apiClient.post).mockResolvedValue({ eligible: 0, selected: 0, executions: [] });
proxy.post.mockResolvedValue({ eligible: 0, selected: 0, executions: [] });
renderWithProviders(<InvestigationSetup mode="new" onClose={vi.fn()} onSave={vi.fn()} />);
await vi.advanceTimersByTimeAsync(400);
await user.click(screen.getByRole("combobox", { name: "Agent (optional)" }));
expect(await screen.findByText(/No matches. You can enter/)).toBeVisible();
await user.keyboard("{Escape}");
vi.mocked(apiClient.post).mockResolvedValue({
proxy.post.mockResolvedValue({
eligible: 1,
selected: 1,
executions: [
@ -493,9 +491,9 @@ it("fetches one preview for two keystrokes inside the debounce window", async ()
renderWithProviders(<InvestigationSetup mode="new" onClose={vi.fn()} onSave={vi.fn()} />);
await user.click(screen.getByText("Advanced filters"));
const previewsFor = (teamId: string) =>
vi
.mocked(apiClient.post)
.mock.calls.filter(([, options]) => (options?.body as { settings: Settings }).settings.team_id === teamId);
proxy.post.mock.calls.filter(
([, options]) => (options?.body as { selection: Settings }).selection.team_id === teamId,
);
await user.type(screen.getByRole("textbox", { name: "Team ID (optional)" }), "ab");
expect(previewsFor("a")).toHaveLength(0);
expect(previewsFor("ab")).toHaveLength(0);
@ -510,8 +508,8 @@ it("fetches one preview for two keystrokes inside the debounce window", async ()
});
it.each(["empty", "error"])("blocks a new investigation when its preview is %s", async (state) => {
if (state === "error") vi.mocked(apiClient.post).mockRejectedValue(new Error("Storage unavailable"));
else vi.mocked(apiClient.post).mockResolvedValue({ eligible: 0, selected: 0, executions: [] });
if (state === "error") proxy.post.mockRejectedValue(new Error("Storage unavailable"));
else proxy.post.mockResolvedValue({ eligible: 0, selected: 0, executions: [] });
const user = userEvent.setup();
renderWithProviders(<InvestigationSetup mode="duplicate" initial={settings} onClose={vi.fn()} onSave={vi.fn()} />);
await user.click(screen.getByRole("button", { name: "Continue" }));

View file

@ -25541,6 +25541,47 @@ export interface components {
*/
traces: boolean;
};
/** ActivitySelection */
ActivitySelection: {
/**
* Agent Name
* @default
*/
agent_name: string;
/**
* Execution Ids
* @default []
*/
execution_ids: string[];
/**
* Filters
* @default []
*/
filters: components["schemas"]["MetadataFilter"][];
/**
* Sample Percent
* @default 100
*/
sample_percent: number;
/** Sample Size */
sample_size?: number | null;
/**
* Service
* @default
*/
service: string;
/**
* Source
* @default traces
* @enum {string}
*/
source: "traces" | "requests" | "both";
/**
* Team Id
* @default
*/
team_id: string;
};
/** AdaptiveRouterWeights */
AdaptiveRouterWeights: {
/**
@ -41454,7 +41495,7 @@ export interface components {
* @default 0
*/
offset: number;
settings: components["schemas"]["LensSettings"];
selection: components["schemas"]["ActivitySelection"];
};
/** Progress */
Progress: {

View file

@ -1,10 +1,72 @@
import type { ReactElement } from "react";
import { vi } from "vitest";
import { LensServicesProvider, liveLensServices } from "@/components/lens/data/LensServices";
import { OnboardingProvider, type Onboarding } from "@/components/lens/onboarding/OnboardingContext";
import { LENS_INTRO_DISMISSED } from "@/components/lens/storage";
import { writeStorage } from "@/lib/storage";
import { renderWithProviders } from "./test-utils";
export interface StubbedRequest {
readonly path: string;
readonly method: string;
readonly query: URLSearchParams;
readonly body: unknown;
readonly authorization: string | null;
}
/** The path of a stubbed fetch call, whether a client passed a URL or a Request. */
export const requestPath = (input: RequestInfo | URL): string =>
new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname;
/** Reads a stubbed fetch call the same way whether a client passed a URL and init or a Request. */
export async function readRequest(input: RequestInfo | URL, init?: RequestInit): Promise<StubbedRequest> {
const request = input instanceof Request ? input : new Request(new URL(String(input), "http://localhost"), init);
const url = new URL(request.url);
const text = await request.clone().text();
return {
path: url.pathname,
method: request.method,
query: url.searchParams,
body: text ? JSON.parse(text) : undefined,
authorization: request.headers.get("Authorization"),
};
}
export interface GatewayRequest {
readonly query: Readonly<Record<string, string>>;
readonly body: unknown;
readonly authorization: string | null;
}
type GatewayHandler = (path: string, request: GatewayRequest) => unknown;
/**
* Fakes the proxy at fetch, so Lens code runs its real HTTP clients. Each method's handler gets the path and the
* parsed query and body; a rejection becomes a 500 carrying its message, and a pending promise stays pending.
*/
export function stubGateway() {
const gateway = {
get: vi.fn<GatewayHandler>(),
post: vi.fn<GatewayHandler>(),
put: vi.fn<GatewayHandler>(),
patch: vi.fn<GatewayHandler>(),
delete: vi.fn<GatewayHandler>(),
};
vi.stubGlobal(
"fetch",
vi.fn<typeof fetch>(async (input, init) => {
const { path, method, query, body, authorization } = await readRequest(input, init);
const handler = gateway[method.toLowerCase() as keyof typeof gateway];
try {
return Response.json((await handler(path, { query: Object.fromEntries(query), body, authorization })) ?? null);
} catch (error) {
return Response.json({ detail: error instanceof Error ? error.message : String(error) }, { status: 500 });
}
}),
);
return gateway;
}
/** The Lens introduction opens on a first visit; tests about anything else start with it dismissed. */
export function dismissLensIntro() {
writeStorage(LENS_INTRO_DISMISSED, true);