scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(lens): address CI failures and review findings

Browse source
This commit is contained in:
moe-berri 2026-10-07 13:23:04 -07:00
parent 2fe8d56ca0
commit 6d819de6ae
32 changed files with 1191 additions and 239 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/lens-worker.yml vendored
View file

@ -88,6 +88,7 @@ jobs:
retention-days: 1
publish:
name: Publish Lens development index
needs: lens-worker-image
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm' && github.ref == 'refs/heads/main'
runs-on: ubuntu-latest

8
.github/workflows/test-rust.yml vendored
View file

@ -6,6 +6,8 @@ on:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "scripts/generate_trace_types.py"
- "scripts/generate_lens_contract.py"
- "litellm/proxy/lens/**"
- "scripts/trace_codegen/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/custom_logger.py"
@ -35,6 +37,8 @@ on:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "scripts/generate_trace_types.py"
- "scripts/generate_lens_contract.py"
- "litellm/proxy/lens/**"
- "scripts/trace_codegen/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/custom_logger.py"
@ -132,6 +136,10 @@ jobs:
working-directory: .
run: uv run scripts/generate_trace_types.py --check
- name: Check generated Lens contracts
working-directory: .
run: uv run scripts/generate_lens_contract.py --check
- run: cargo nextest run --workspace --locked --features litellm-traces/schema,litellm-traces-clickhouse/schema
- run: cargo test --workspace --doc --locked

4
deploy/lens/README.md
View file

@ -82,9 +82,13 @@ lensWorker:
clickhouseSecret:
name: litellm-lens-clickhouse
key: url
clickhouseDatabase: litellm
retentionDays: 14
publicUrl: https://<your-litellm-host>/lens-ingest
```
Set `clickhouseDatabase` and `retentionDays` to your existing database and retention before upgrading
When the chart's main ingress is enabled, it routes `/lens-ingest` directly to Lens. With a custom ingress, add that route yourself. For a dedicated hostname, use `lensWorker.ingress.enabled`, `host`, `className`, and `tls`, and set `publicUrl` to that hostname. The chart connects LiteLLM to Lens internally and gives both services the shared secret
Update your existing component image overrides to matching builds, then use the chart from that checkout:

6
helm/litellm-helm/templates/_helpers.tpl
View file

@ -365,3 +365,9 @@ shutdown drain window.
value: {{ .drainTimeoutSeconds | quote }}
{{- end }}
{{- end -}}
{{- define "litellm.lensWorker.labels" -}}
{{- $labels := include "litellm.labels" . | fromYaml -}}
{{- $_ := set $labels "app.kubernetes.io/name" (printf "%s-lens-worker" (include "litellm.name" . | trunc 51 | trimSuffix "-")) -}}
{{- toYaml $labels -}}
{{- end -}}

8
helm/litellm-helm/templates/lens/deployment.yaml
View file

@ -4,7 +4,7 @@ kind: Deployment
metadata:
name: {{ include "litellm.fullname" . }}-lens-worker
labels:
{{- include "litellm.labels" . | nindent 4 }}
{{- include "litellm.lensWorker.labels" . | nindent 4 }}
app.kubernetes.io/component: lens-worker
spec:
replicas: {{ .Values.lensWorker.replicaCount }}
@ -15,7 +15,7 @@ spec:
template:
metadata:
labels:
{{- include "litellm.labels" . | nindent 8 }}
{{- include "litellm.lensWorker.labels" . | nindent 8 }}
app.kubernetes.io/component: lens-worker
spec:
automountServiceAccountToken: false
@ -52,6 +52,10 @@ spec:
secretKeyRef:
name: {{ required "lensWorker.clickhouseSecret.name is required" .Values.lensWorker.clickhouseSecret.name | quote }}
key: {{ .Values.lensWorker.clickhouseSecret.key | quote }}
- name: CLICKHOUSE_DATABASE
value: {{ .Values.lensWorker.clickhouseDatabase | quote }}
- name: AGENT_TRACING_RETENTION_DAYS
value: {{ .Values.lensWorker.retentionDays | quote }}
{{- if .Values.lensWorker.tokenSecret.name }}
- name: LENS_WORKER_TOKEN
valueFrom:

167
helm/litellm-helm/tests/lens_service_tests.yaml Normal file
View file

@ -0,0 +1,167 @@
suite: Lens service isolation and ingestion routing
templates:
- configmap-litellm.yaml
- deployment.yaml
- ingress.yaml
- lens/ingress.yaml
- lens/service.yaml
- lens/deployment.yaml
tests:
- it: connects deployment.yaml to the shared Lens service
template: deployment.yaml
set: &id001
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: https://gateway.example/lens-ingest
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_URL
value: http://lens-test-lens-worker:4318
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://gateway.example/lens-ingest
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_SERVICE_TOKEN
valueFrom:
secretKeyRef:
name: lens-service
key: service-token
- it: routes uploads directly to Lens instead of the gateway
template: ingress.yaml
set:
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: https://gateway.example/lens-ingest
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
ingress.enabled: true
ingress.hosts:
- host: gateway.example
paths:
- path: /
pathType: Prefix
asserts:
- contains:
path: spec.rules[0].http.paths
content:
path: /lens-ingest
pathType: Prefix
backend:
service:
name: lens-test-lens-worker
port:
number: 4318
- it: keeps internal routes out of a dedicated ingestion hostname
template: lens/ingress.yaml
set:
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: https://gateway.example/lens-ingest
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
lensWorker.ingress.enabled: true
lensWorker.ingress.host: traces.example
asserts:
- equal:
path: spec.rules[0].http.paths
value:
- path: /v1/
pathType: Prefix
backend:
service:
name: lens-test-lens-worker
port:
name: otlp
- it: maps the Lens service to the ingestion listener
template: lens/service.yaml
set: *id001
asserts:
- equal:
path: spec.selector
value:
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: lens-worker
- equal:
path: spec.ports
value:
- name: otlp
port: 4318
targetPort: otlp
- it: gives only Lens the ClickHouse secret
template: lens/deployment.yaml
set: *id001
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_URL
valueFrom:
secretKeyRef:
name: lens-storage
key: url
- it: requires an agent reachable ingestion URL
template: deployment.yaml
set:
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: ''
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
asserts:
- failedTemplate:
errorMessage: lensWorker.publicUrl is required
- it: omits Lens connection settings when disabled in deployment.yaml
template: deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_URL
any: true
- it: preserves an existing ClickHouse database and retention
template: lens/deployment.yaml
set:
lensWorker.enabled: true
lensWorker.publicUrl: https://traces.example
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
lensWorker.clickhouseDatabase: existing_traces
lensWorker.retentionDays: 45
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_DATABASE
value: existing_traces
- contains:
path: spec.template.spec.containers[0].env
content:
name: AGENT_TRACING_RETENTION_DAYS
value: '45'
- it: keeps Lens pods outside the gateway autoscaling selector
template: lens/deployment.yaml
set: &id002
nameOverride: inference
lensWorker.enabled: true
lensWorker.publicUrl: https://traces.example
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
asserts:
- equal:
path: spec.template.metadata.labels["app.kubernetes.io/name"]
value: inference-lens-worker
- it: preserves the existing gateway deployment selector
template: deployment.yaml
set: *id002
asserts:
- equal:
path: spec.selector.matchLabels["app.kubernetes.io/name"]
value: inference

2
helm/litellm-helm/values.yaml
View file

@ -667,6 +667,8 @@ lensWorker:
serviceTokenSecret:
name: ""
key: service-token
clickhouseDatabase: litellm
retentionDays: 14
clickhouseSecret:
name: ""
key: url

6
helm/litellm/templates/_helpers.tpl
View file

@ -528,3 +528,9 @@ shutdown drain window.
key: {{ .Values.lensWorker.serviceTokenSecret.key | quote }}
{{- end }}
{{- end -}}
{{- define "litellm.lensWorker.labels" -}}
{{- $labels := include "litellm.commonLabels" . | fromYaml -}}
{{- $_ := set $labels "app.kubernetes.io/name" (printf "%s-lens-worker" (include "litellm.name" . | trunc 51 | trimSuffix "-")) -}}
{{- toYaml $labels -}}
{{- end -}}

8
helm/litellm/templates/lens/deployment.yaml
View file

@ -4,7 +4,7 @@ kind: Deployment
metadata:
name: {{ include "litellm.fullname" . }}-lens-worker
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
{{- include "litellm.lensWorker.labels" . | nindent 4 }}
app.kubernetes.io/component: lens-worker
spec:
replicas: {{ .Values.lensWorker.replicaCount }}
@ -15,7 +15,7 @@ spec:
template:
metadata:
labels:
{{- include "litellm.commonLabels" . | nindent 8 }}
{{- include "litellm.lensWorker.labels" . | nindent 8 }}
app.kubernetes.io/component: lens-worker
spec:
automountServiceAccountToken: false
@ -52,6 +52,10 @@ spec:
secretKeyRef:
name: {{ required "lensWorker.clickhouseSecret.name is required" .Values.lensWorker.clickhouseSecret.name | quote }}
key: {{ .Values.lensWorker.clickhouseSecret.key | quote }}
- name: CLICKHOUSE_DATABASE
value: {{ .Values.lensWorker.clickhouseDatabase | quote }}
- name: AGENT_TRACING_RETENTION_DAYS
value: {{ .Values.lensWorker.retentionDays | quote }}
{{- if .Values.lensWorker.tokenSecret.name }}
- name: LENS_WORKER_TOKEN
valueFrom:

196
helm/litellm/tests/lens_service_tests.yaml Normal file
View file

@ -0,0 +1,196 @@
suite: Lens service isolation and ingestion routing
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- backend/deployment.yaml
- ingress.yaml
- lens/ingress.yaml
- lens/service.yaml
- lens/deployment.yaml
tests:
- it: connects gateway/deployment.yaml to the shared Lens service
template: gateway/deployment.yaml
set: &id001
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: https://gateway.example/lens-ingest
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_URL
value: http://lens-test-lens-worker:4318
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://gateway.example/lens-ingest
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_SERVICE_TOKEN
valueFrom:
secretKeyRef:
name: lens-service
key: service-token
- it: connects backend/deployment.yaml to the shared Lens service
template: backend/deployment.yaml
set: *id001
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_URL
value: http://lens-test-lens-worker:4318
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://gateway.example/lens-ingest
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_SERVICE_TOKEN
valueFrom:
secretKeyRef:
name: lens-service
key: service-token
- it: routes uploads directly to Lens instead of the gateway
template: ingress.yaml
set:
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: https://gateway.example/lens-ingest
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
ingress.enabled: true
ingress.host: gateway.example
asserts:
- contains:
path: spec.rules[0].http.paths
content:
path: /lens-ingest
pathType: Prefix
backend:
service:
name: lens-test-lens-worker
port:
number: 4318
- it: keeps internal routes out of a dedicated ingestion hostname
template: lens/ingress.yaml
set:
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: https://gateway.example/lens-ingest
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
lensWorker.ingress.enabled: true
lensWorker.ingress.host: traces.example
asserts:
- equal:
path: spec.rules[0].http.paths
value:
- path: /v1/
pathType: Prefix
backend:
service:
name: lens-test-lens-worker
port:
name: otlp
- it: maps the Lens service to the ingestion listener
template: lens/service.yaml
set: *id001
asserts:
- equal:
path: spec.selector
value:
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: lens-worker
- equal:
path: spec.ports
value:
- name: otlp
port: 4318
targetPort: otlp
- it: gives only Lens the ClickHouse secret
template: lens/deployment.yaml
set: *id001
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_URL
valueFrom:
secretKeyRef:
name: lens-storage
key: url
- it: requires an agent reachable ingestion URL
template: gateway/deployment.yaml
set:
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: ''
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
asserts:
- failedTemplate:
errorMessage: lensWorker.publicUrl is required
- it: omits Lens connection settings when disabled in gateway/deployment.yaml
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_URL
any: true
- it: omits Lens connection settings when disabled in backend/deployment.yaml
template: backend/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_URL
any: true
- it: preserves an existing ClickHouse database and retention
template: lens/deployment.yaml
set:
lensWorker.enabled: true
lensWorker.publicUrl: https://traces.example
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
lensWorker.clickhouseDatabase: existing_traces
lensWorker.retentionDays: 45
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_DATABASE
value: existing_traces
- contains:
path: spec.template.spec.containers[0].env
content:
name: AGENT_TRACING_RETENTION_DAYS
value: '45'
- it: keeps Lens pods outside the gateway autoscaling selector
template: lens/deployment.yaml
set: &id002
nameOverride: inference
lensWorker.enabled: true
lensWorker.publicUrl: https://traces.example
lensWorker.serviceTokenSecret.name: lens-service
lensWorker.clickhouseSecret.name: lens-storage
asserts:
- equal:
path: spec.template.metadata.labels["app.kubernetes.io/name"]
value: inference-lens-worker
- it: preserves the existing gateway deployment selector
template: gateway/deployment.yaml
set: *id002
asserts:
- equal:
path: spec.selector.matchLabels["app.kubernetes.io/name"]
value: inference
values:
- ./values/required.yaml

2
helm/litellm/values.yaml
View file

@ -644,6 +644,8 @@ lensWorker:
serviceTokenSecret:
name: ""
key: service-token
clickhouseDatabase: litellm
retentionDays: 14
clickhouseSecret:
name: ""
key: url

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20261007000100_lens_ingestion_keys/migration.sql
View file

@ -1,4 +1,4 @@
CREATE TABLE "LiteLLM_LensIngestionKey" (
CREATE TABLE IF NOT EXISTS "LiteLLM_LensIngestionKey" (
"id" TEXT NOT NULL,
"data" JSONB NOT NULL,
CONSTRAINT "LiteLLM_LensIngestionKey_pkey" PRIMARY KEY ("id")

13
litellm-rust/crates/lens/src/agent.rs
View file

@ -269,17 +269,16 @@ pub async fn run<T: Output>(
Ok(value) => value.to_string(),
Err(error) => json!({"request": tool, "error": error.to_string()}).to_string(),
};
bytes += result.len();
if bytes > MAX_TOOL_BYTES {
let error = json!({"request": tool, "error": "Combined tool output exceeds 8 MiB. Request smaller ranges or fewer tools per turn."}).to_string();
results.push(error.clone());
archived.push(error);
continue;
}
archived.push(match &tool {
Tool::Evidence(r) => journal.reference(r).unwrap_or_else(|| result.clone()),
_ => result.clone(),
});
bytes += result.len();
if bytes > MAX_TOOL_BYTES {
let error = json!({"request": tool, "error": "Combined tool output exceeds 8 MiB. Request smaller ranges or fewer tools per turn."}).to_string();
results.push(error);
continue;
}
results.push(result);
}
journal

106
litellm-rust/crates/lens/tests/worker.rs
View file

@ -259,3 +259,109 @@ async fn model_failures_expose_only_bounded_sanitized_gateway_diagnostics(
litellm_lens::Error::Control { status: 400, .. }
));
}
#[rstest]
#[tokio::test]
async fn configured_private_dns_names_are_reachable_without_following_redirects() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/private-service"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/redirect"))
.respond_with(ResponseTemplate::new(302).insert_header("location", "/private-service"))
.mount(&server)
.await;
let base = server.uri().replace("127.0.0.1", "localhost");
let client = http_client().unwrap();
let response = client
.get(format!("{base}/private-service"))
.send()
.await
.unwrap();
assert_eq!(response.status(), 200);
let redirected = client.get(format!("{base}/redirect")).send().await.unwrap();
assert_eq!(redirected.status(), 302);
}
#[tokio::test]
async fn oversized_combined_tool_replies_remain_readable_after_a_checkpoint() {
use litellm_lens::{
activity::Tracker,
agent,
evidence::{MAX_TOOL_BYTES, Workspace},
};
let server = MockServer::start().await;
let claim: wire::Claim = serde_json::from_value(fixture()).unwrap();
let sample: wire::Sample = serde_json::from_str(include_str!("fixtures/sample.json")).unwrap();
let filler_size = MAX_TOOL_BYTES * 3 / 5;
let page_calls = Arc::new(AtomicUsize::new(0));
let page_count = page_calls.clone();
let execution = sample.executions[0].clone();
Mock::given(method("GET"))
.and(path("/lens/worker/lens-test/job-test/content"))
.respond_with(move |_: &Request| {
let marker = if page_count.fetch_add(1, Ordering::SeqCst) == 0 { "FIRST_REPLY" } else { "ARCHIVED_SECOND_REPLY" };
ResponseTemplate::new(200).set_body_json(json!({"execution":execution,"parts":[{
"execution_id":"run-test","span_id":"span-test","name":format!("{marker}{}", "x".repeat(filler_size)),"kind":"tool","content":"evidence","truncated":false
}]}))
}).expect(2).mount(&server).await;
Mock::given(method("POST"))
.and(path("/lens/worker/lens-test/job-test/progress"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({})))
.mount(&server)
.await;
let model_calls = Arc::new(AtomicUsize::new(0));
let model_count = model_calls.clone();
Mock::given(method("POST"))
.and(path("/lens/worker/lens-test/job-test/model"))
.respond_with(move |request: &Request| {
let model: wire::ModelRequest = request.body_json().unwrap();
let turn = match model_count.fetch_add(1, Ordering::SeqCst) {
0 => json!({"tools":[{"action":"catalog","execution_id":"run-test"},{"action":"catalog","execution_id":"run-test"}],"checkpoint":"Inspect the archived second reply"}),
1 => {
let reply: Value = serde_json::from_str(&model.messages.last().unwrap().content).unwrap();
assert!(reply["tool_results"][1].as_str().unwrap().contains("Combined tool output exceeds"), "{}", reply["tool_results"][1].as_str().unwrap().chars().take(600).collect::<String>());
json!({"tools":[{"action":"history","turn_start":0,"turn_end":1,"char_start":filler_size,"char_end":filler_size+6000}]})
},
2 => {
let reply: Value = serde_json::from_str(&model.messages.last().unwrap().content).unwrap();
let history: Value = serde_json::from_str(reply["tool_results"][0].as_str().unwrap()).unwrap();
assert!(history["excerpt"].as_str().unwrap().contains("ARCHIVED_SECOND_REPLY"));
json!({"result":{"observations":[]}})
},
_ => panic!("Unexpected model retry"),
};
ResponseTemplate::new(200).set_body_json(json!({"content":turn.to_string(),"cost":0}))
}).expect(3).mount(&server).await;
let client = client(&server);
let workspace = Workspace::new(sample.executions, client.clone());
let tracker = Tracker::start(
&client,
"test".into(),
wire::ActivityPhase::Review,
"Archive".into(),
vec![],
)
.await
.unwrap();
let output: wire::Extraction = agent::run(
&claim,
&workspace,
agent::Assignment {
stage: "test",
task: "Read two tools and recover the second from history".into(),
purpose: wire::ModelRequestPurpose::Extract,
supplied: json!({}),
},
&tracker,
)
.await
.unwrap();
assert!(output.observations.is_empty());
assert_eq!(page_calls.load(Ordering::SeqCst), 2);
assert_eq!(model_calls.load(Ordering::SeqCst), 3);
}

3
litellm/integrations/otel/model/spans.py
View file

@ -352,6 +352,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_LensDataset",
"LiteLLM_LensRun",
"LiteLLM_LensReview",
"LiteLLM_LensIngestionKey",
"LiteLLM_LensWorker",
)
)
@ -472,7 +473,7 @@ _POSTGRES_OPERATION_BY_CALL_TYPE: Final[Mapping[str, PostgresOperation]] = Mappi
_RAW_PRISMA_CALL_TYPES: Final[frozenset[str]] = frozenset(("query_raw", "execute_raw"))
_DB_OPERATION_METADATA_KEY: Final = "db_operation"
_POSTGRES_VERBS: Final[frozenset[str]] = frozenset(
("select", "insert", "update", "delete", "upsert", "ddl", "set", "ping")
("select", "insert", "update", "delete", "upsert", "ddl", "set", "ping", "lock")
)
_TARGETLESS_VERBS: Final[frozenset[str]] = frozenset(("ping",))
_SETTING_NAME: Final = re.compile(r"[a-z_][a-z0-9_.]*")

1
litellm/proxy/db/prisma_query_span.py
View file

@ -76,6 +76,7 @@ _VERB_BY_KEYWORD: Final[Mapping[str, str]] = MappingProxyType(
"REFRESH": "ddl",
"TRUNCATE": "delete",
"SET": "set",
"LOCK": "lock",
}
)

2
litellm/proxy/lens/endpoints.py
View file

@ -630,7 +630,7 @@ async def claim(worker: WorkerAuth, protocol_version: int = 1, worker_release: s
if protocol_version != PROTOCOL_VERSION or worker_release != expected:
raise HTTPException(409, f"Upgrade the Lens worker to {image} and retry")
if worker.analysis_key_id is None:
raise HTTPException(409, "Assign an analysis key to this worker in Lens setup")
return None
now: Final = datetime.now(timezone.utc)
lens_repository: Final = repository()
await lens_repository.heartbeat(worker.id, now.isoformat())

4
litellm/proxy/tracing_runtime.py
View file

@ -2,6 +2,7 @@ from collections.abc import AsyncGenerator, Callable, Mapping
from contextlib import asynccontextmanager
from typing import Final
import httpx
from fastapi import HTTPException, Request
from pydantic import ConfigDict, TypeAdapter
@ -38,6 +39,7 @@ async def manage_tracing(
enabled: bool,
receiver_factory: Callable[[], TraceReceiver] | None = None,
settings: Mapping[str, object] | None = None,
client_factory: Callable[[LensConnection], httpx.AsyncClient] = LensConnection.client,
) -> AsyncGenerator[TraceReceiver | None, None]:
if not enabled:
yield None
@ -50,7 +52,7 @@ async def manage_tracing(
)
yield None
return
async with connection.client() as client:
async with client_factory(connection) as client:
tracing: Final = (
receiver_factory()
if receiver_factory

7
litellm/tracing/exporter.py
View file

@ -1,7 +1,7 @@
import asyncio
import json
from collections import deque
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from contextlib import suppress
from io import BytesIO
from typing import Final
@ -72,9 +72,10 @@ def encode_record(value: Mapping[str, object]) -> bytes:
class LensExporter(CustomLogger):
def __init__(self, client: httpx.AsyncClient) -> None:
def __init__(self, client: httpx.AsyncClient, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep) -> None:
super().__init__()
self.client: Final = client
self.sleep: Final = sleep
self.queue: Final[deque[bytes]] = deque() # mutable-ok: bounded producer-consumer queue
self.wake: Final = asyncio.Event()
self.closed = False
@ -173,7 +174,7 @@ class LensExporter(CustomLogger):
except httpx.HTTPError:
pass
if attempt < 2:
await asyncio.sleep(float(1 << attempt))
await self.sleep(float(1 << attempt))
self._warn("retry limit reached")
return False

65
litellm/tracing/remote.py
View file

@ -2,11 +2,13 @@ import json
import os
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Final
from enum import Enum
from typing import Final, NoReturn
from urllib.parse import urlsplit
import httpx
from pydantic import JsonValue, TypeAdapter
from typing_extensions import assert_never
from litellm.rust_bridge.trace.errors import TraceChanged
from litellm.rust_bridge.trace.generated.types import QueryScope, ReadQueryName, TraceScope
@ -47,6 +49,33 @@ class LensConnection:
)
class _ReadFailure(Enum):
INVALID_QUERY = "invalid_query"
CHANGED = "changed"
QUERY_TOO_LARGE = "query_too_large"
UNAVAILABLE = "unavailable"
RESPONSE_TOO_LARGE = "response_too_large"
INVALID_RESPONSE = "invalid_response"
def _raise_read_failure(failure: _ReadFailure) -> NoReturn:
match failure:
case _ReadFailure.INVALID_QUERY:
raise ValueError("Invalid trace query")
case _ReadFailure.CHANGED:
raise TraceChanged("Trace changed while paging; refresh the trace to continue")
case _ReadFailure.QUERY_TOO_LARGE:
raise OverflowError("Trace exceeds the interactive read budget")
case _ReadFailure.UNAVAILABLE:
raise RuntimeError("Lens trace storage is unavailable")
case _ReadFailure.RESPONSE_TOO_LARGE:
raise RuntimeError("Lens response exceeds the size limit")
case _ReadFailure.INVALID_RESPONSE:
raise ValueError("Invalid Lens response")
case _:
assert_never(failure)
class RemoteTraceStore:
def __init__(self, client: httpx.AsyncClient) -> None:
self.client: Final = client
@ -55,19 +84,31 @@ class RemoteTraceStore:
return
async def _read(self, request: Mapping[str, object]) -> JsonValue:
result: Final = await self._read_result(request)
if isinstance(result, _ReadFailure):
_raise_read_failure(result)
return result
async def _read_result(self, request: Mapping[str, object]) -> JsonValue | _ReadFailure:
try:
async with self.client.stream("POST", "/internal/read", json=dict(request)) as response:
if response.status_code == 400:
raise ValueError("Invalid trace query")
if response.status_code == 409:
raise TraceChanged("Trace changed while paging; refresh the trace to continue")
if response.status_code == 413:
raise OverflowError("Trace exceeds the interactive read budget")
if response.status_code != 200:
raise RuntimeError("Lens trace storage is unavailable")
return _JSON.validate_json(await bounded_response(response, MAX_RESPONSE_BYTES))
except httpx.HTTPError as error:
raise RuntimeError("Lens trace storage is unavailable") from error
match response.status_code:
case 400:
return _ReadFailure.INVALID_QUERY
case 409:
return _ReadFailure.CHANGED
case 413:
return _ReadFailure.QUERY_TOO_LARGE
case 200:
return _JSON.validate_json(await bounded_response(response, MAX_RESPONSE_BYTES))
case _:
return _ReadFailure.UNAVAILABLE
except httpx.HTTPError:
return _ReadFailure.UNAVAILABLE
except RuntimeError:
return _ReadFailure.RESPONSE_TOO_LARGE
except ValueError:
return _ReadFailure.INVALID_RESPONSE
async def insert_rows(self, table: str, rows: Sequence[Mapping[str, object]]) -> None:
if table != "spend_logs":

11
tests/test_litellm_rust/test_traces.py
View file

@ -533,9 +533,14 @@ def _fixture_trace_api(
with TestClient(app) as client:
assert client.portal is not None
client.portal.call(storage.ensure_schema)
ingested: Final = tuple(client.post("/v1/traces", json=replay.export) for replay in replays)
for result in ingested:
assert result.status_code == 200, result.text
for replay in replays:
client.portal.call(
TraceReceiver(storage).ingest,
json.dumps(replay.export).encode(),
"application/json",
None,
Tenant(team_id="team-a", api_key_hash="fixture-key", user_id="fixture-user"),
)
client.portal.call(storage.insert_rows, "spend_logs", stamped)
response: Final = client.get("/v1/traces/query/help")
assert response.status_code == 200, response.text

17
tests/unit/proxy/common_utils/test_http_parsing_utils.py
View file

@ -1394,12 +1394,10 @@ async def test_otlp_auth_does_not_consume_chunked_bodies_before_the_receiver_lim
@pytest.mark.asyncio
async def test_auth_body_read_and_trace_handler_leave_stream_for_receiver_limit() -> None:
async def test_auth_and_retired_trace_handler_never_consume_upload_body() -> None:
from litellm.constants import OTLP_MAX_BODY_BYTES
from litellm.proxy import tracing_endpoints
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import _read_request_body_deferring_parse_failure
from litellm.tracing import TraceReceiver
chunk: Final = b"x" * (OTLP_MAX_BODY_BYTES // 2 + 1)
receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": chunk, "more_body": True}] * 2)
@ -1407,21 +1405,14 @@ async def test_auth_body_read_and_trace_handler_leave_stream_for_receiver_limit(
{"type": "http", "method": "POST", "path": "/v1/traces", "headers": [(b"content-type", b"application/json")]},
receive,
)
storage: Final = MagicMock()
storage.ingest = AsyncMock()
context: Final = await tracing_endpoints.provide_trace_access(
auth=UserAPIKeyAuth(token="key", team_id="team"), tracing=TraceReceiver(storage), log_team_lookup=AsyncMock()
)
parsed, parse_error = await _read_request_body_deferring_parse_failure(request)
assert parsed == {}
assert parse_error is None
receive.assert_not_awaited()
response: Final = await tracing_endpoints.ingest_otlp_traces(request, context)
assert response.status_code == 413
assert receive.await_count == 2
storage.ingest.assert_not_awaited()
response: Final = await tracing_endpoints.ingest_otlp_traces(request)
assert response.status_code == 410
receive.assert_not_awaited()
@pytest.fixture()

3
tests/unit/proxy/db/test_prisma_query_span.py
View file

@ -41,7 +41,7 @@ _MODEL_BY_ACCESSOR: Final[Mapping[str, str]] = {relation.lower(): relation for r
_GENERIC_CRUD_HELPERS: Final = frozenset({"get_data", "get_generic_data", "insert_data", "update_data", "delete_data"})
_TRANSACTION_BODIES: Final[Mapping[str, str]] = {"litellm/proxy/db/baseline_accounting.py": "baseline_accounting"}
_RENDERED_NAME: Final = re.compile(
r"postgres\.(select|insert|update|delete|upsert|ddl|set|transaction) .+|postgres\.ping"
r"postgres\.(select|insert|update|delete|upsert|ddl|set|transaction|lock) .+|postgres\.ping"
)
@ -115,6 +115,7 @@ def test_a_payload_the_parser_does_not_know_stays_the_legacy_function_named_span
'WITH team_rows AS (UPDATE "LiteLLM_TeamTable" SET models = $1 RETURNING team_id) SELECT team_id FROM team_rows',
("update", "LiteLLM_TeamTable"),
),
('LOCK TABLE "LiteLLM_LensIngestionKey" IN EXCLUSIVE MODE', ("lock", "LiteLLM_LensIngestionKey")),
("BEGIN", (None, None)),
],
)

154
tests/unit/proxy/lens/test_endpoints.py
View file

@ -45,9 +45,9 @@ from litellm.proxy.lens.models import (
)
from litellm.proxy.lens.repository import DueLens, Row
from litellm.proxy.lens.state import claim_job, queue_job, replace_job
from litellm.rust_bridge.trace.generated.models import ExecutionRow, LensSampleParams
from litellm.rust_bridge.trace.storage import ClickHouseStorage
from litellm.tracing.remote import RemoteTraceStore
from litellm.rust_bridge.trace.generated.models import ExecutionRow, LensSampleParams
from tests.unit.proxy.lens.test_agent_workspace import execution
from tests.unit.proxy.lens.test_state import NOW, lens, worker
@ -800,3 +800,155 @@ async def test_claim_due_pages_through_more_than_a_thousand_full_pages() -> None
assert claim is None
assert len(repository.after_calls) == 1_201
assert repository.after_calls == expected_after
@pytest.mark.asyncio
async def test_compatible_worker_without_an_analysis_key_waits_without_claiming_jobs(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.lens.endpoints import claim
from litellm.proxy.lens.release import PROTOCOL_VERSION
monkeypatch.setenv("LITELLM_RELEASE_TAG", "v1.2.3")
monkeypatch.setattr(proxy_server, "prisma_client", None)
unassigned: Final = worker().model_copy(update={"analysis_key_id": None})
assert await claim(unassigned, protocol_version=PROTOCOL_VERSION, worker_release="v1.2.3") is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
"configured,credential,expected", ((False, "x" * 32, 503), (True, "wrong", 401), (True, "x" * 32, None))
)
async def test_internal_service_authentication_is_separate_from_gateway_keys(
monkeypatch: pytest.MonkeyPatch, configured: bool, credential: str, expected: int | None
) -> None:
from fastapi.security import HTTPAuthorizationCredentials
from litellm.proxy.lens.endpoints import service_auth
monkeypatch.setenv("LITELLM_LENS_URL", "http://lens" if configured else "")
monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", "x" * 32)
credentials: Final = HTTPAuthorizationCredentials(scheme="Bearer", credentials=credential)
if expected is None:
assert await service_auth(credentials) is None
else:
with pytest.raises(HTTPException) as failure:
await service_auth(credentials)
assert failure.value.status_code == expected
@pytest.mark.asyncio
@pytest.mark.parametrize(
"status,content,connected",
(
(200, b'{"storage_ready":true,"credentials_ready":true,"release":"v1.2.3","protocol_version":2}', True),
(503, b"private storage details", False),
(200, b"invalid JSON", False),
(200, b"x" * 17000, False),
),
)
async def test_service_status_uses_internal_auth_and_only_advertises_the_public_url(
monkeypatch: pytest.MonkeyPatch, status: int, content: bytes, connected: bool
) -> None:
import respx
from litellm.proxy.lens.endpoints import service_connection
monkeypatch.setenv("LITELLM_LENS_URL", "http://lens/private-prefix")
monkeypatch.setenv("LITELLM_LENS_PUBLIC_URL", "https://traces.example/lens-ingest/")
monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", "x" * 32)
with respx.mock as network:
route: Final = network.get("http://lens/private-prefix/internal/status").respond(status, content=content)
result: Final = await service_connection(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN))
assert result.url == "https://traces.example/lens-ingest"
assert result.connected is connected
assert result.status.storage_ready is connected
assert route.calls[0].request.headers["Authorization"] == "Bearer " + "x" * 32
assert "private storage details" not in result.model_dump_json()
@pytest.mark.asyncio
async def test_credential_snapshot_excludes_expired_keys_and_disables_caching(monkeypatch: pytest.MonkeyPatch) -> None:
from unittest.mock import AsyncMock
from fastapi import Response
from litellm.proxy import proxy_server
from litellm.proxy.lens.endpoints import ingestion_credentials
from litellm.proxy.lens.ingestion import IngestionCredential, IngestionKeyRequest, new_key
current: Final = new_key(IngestionKeyRequest(team_id="team"), "owner").record
expired: Final = current.model_copy(update={"id": "expired", "expires_at": 1})
db: Final = SimpleNamespace(
query_raw=AsyncMock(return_value=tuple(Row(data=key.model_dump(mode="json")) for key in (current, expired)))
)
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db))
response: Final = Response()
snapshot: Final = await ingestion_credentials(None, response)
assert snapshot.keys == (
IngestionCredential(token_hash=current.tenant.api_key_hash, tenant=current.tenant, expires_at=None),
)
assert response.headers["Cache-Control"] == "no-store"
assert snapshot.issued_at >= int(current.created_at.timestamp())
@pytest.mark.asyncio
@pytest.mark.parametrize("accepted", (True, False))
async def test_created_ingestion_keys_report_activation_only_after_the_service_acknowledges(
monkeypatch: pytest.MonkeyPatch, accepted: bool
) -> None:
import hashlib
import json
from unittest.mock import AsyncMock, MagicMock
import respx
from litellm.proxy import proxy_server
from litellm.proxy.lens.endpoints import create_ingestion_key, list_ingestion_keys, revoke_ingestion_key
from litellm.proxy.lens.ingestion import IngestionKey, IngestionKeyRequest
db: Final = SimpleNamespace(query_raw=AsyncMock(return_value=()), execute_raw=AsyncMock(return_value=1))
context: Final = AsyncMock()
context.__aenter__.return_value = db
db.tx = MagicMock(return_value=context)
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db))
monkeypatch.setenv("LITELLM_LENS_URL", "http://lens")
monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", "x" * 32)
auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="owner")
with respx.mock as network:
route: Final = network.post("http://lens/internal/credentials").respond(204 if accepted else 503)
created: Final = await create_ingestion_key(IngestionKeyRequest(name="Agent", team_id="team"), auth)
assert created.active is accepted
persisted: Final = IngestionKey.model_validate_json(db.execute_raw.call_args.args[2])
assert persisted == created.record
assert persisted.tenant.api_key_hash == hashlib.sha256(created.key.encode()).hexdigest()
assert persisted.tenant.user_id == "owner"
assert persisted.tenant.team_id == "team"
assert created.key not in persisted.model_dump_json()
db.query_raw.return_value = (Row(data=persisted.model_dump(mode="json")),)
assert await list_ingestion_keys(auth) == (persisted,)
db.query_raw.return_value = ()
assert await revoke_ingestion_key(persisted.id, auth)
assert db.execute_raw.call_args.args == ('DELETE FROM "LiteLLM_LensIngestionKey" WHERE id=$1', persisted.id)
assert json.loads(route.calls[-1].request.content)["keys"] == []
@pytest.mark.asyncio
async def test_ingestion_keys_reject_expired_requests_and_read_only_admins(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.lens.endpoints import create_ingestion_key
from litellm.proxy.lens.ingestion import IngestionKeyRequest
monkeypatch.setattr(proxy_server, "prisma_client", None)
with pytest.raises(HTTPException) as expired:
await create_ingestion_key(
IngestionKeyRequest(expires_at=datetime(2000, 1, 1, tzinfo=timezone.utc)),
UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert expired.value.status_code == 422
with pytest.raises(HTTPException) as forbidden:
await create_ingestion_key(
IngestionKeyRequest(), UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
)
assert forbidden.value.status_code == 403

61
tests/unit/proxy/proxy_server/test_proxy_config.py
View file

@ -71,49 +71,44 @@ async def test_proxy_config_loads_tracing_url_and_retention_from_yaml(tmp_path,
@pytest.mark.asyncio
@pytest.mark.parametrize("shutdown_error", [False, True])
async def test_tracing_config_automatically_logs_spend_without_callback_setting(shutdown_error: bool) -> None:
from litellm.integrations.clickhouse.clickhouse_spend_logger import ClickHouseSpendLogger
from litellm.proxy.tracing_runtime import manage_tracing
from litellm.tracing import TraceReceiver
async def test_tracing_config_automatically_exports_spend_without_a_storage_dependency(
shutdown_error: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
import httpx
storage: Final = MagicMock()
storage.ensure_schema = AsyncMock()
storage.insert_rows = AsyncMock()
receiver: Final = TraceReceiver(storage)
from litellm.proxy.tracing_runtime import manage_tracing
from litellm.tracing.exporter import LensExporter
from litellm.tracing.remote import LensConnection
monkeypatch.setenv("LITELLM_LENS_URL", "http://lens.test")
monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", "test-service-token-with-32-characters")
received: Final = asyncio.Future[httpx.Request]()
def accept(request: httpx.Request) -> httpx.Response:
received.set_result(request)
return httpx.Response(204)
def client(connection: LensConnection) -> httpx.AsyncClient:
return httpx.AsyncClient(base_url=connection.url, transport=httpx.MockTransport(accept))
outcome: Final = pytest.raises(RuntimeError, match="shutdown failure") if shutdown_error else nullcontext()
with outcome:
async with manage_tracing(enabled=True, receiver_factory=lambda: receiver):
storage.ensure_schema.assert_awaited_once()
logger: Final = next(
callback
for callback in litellm._async_success_callback
if isinstance(callback, ClickHouseSpendLogger) and callback.storage is storage
)
now: Final = datetime.now()
async with manage_tracing(enabled=True, client_factory=client):
logger: Final = next(callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter))
await logger.async_log_success_event(
{
"standard_logging_object": {
"id": "response-1",
"startTime": now.timestamp(),
"endTime": now.timestamp(),
"response_cost": 0.25,
}
},
None,
now,
now,
{"standard_logging_object": {"id": "response-1", "response_cost": 0.25}}, None, None, None
)
storage.insert_rows.assert_not_awaited()
if shutdown_error:
raise RuntimeError("shutdown failure")
assert storage.insert_rows.await_args.args[0] == "spend_logs"
assert storage.insert_rows.await_args.args[1][0]["spend"] == 0.25
request: Final = received.result()
rows: Final = json.loads(request.content)
assert request.url.path == "/internal/spend"
assert rows[0]["spend"] == 0.25
assert rows[0]["response_id"] == "response-1"
assert logger not in litellm._async_success_callback
assert logger._flush_task is not None and logger._flush_task.done()
assert not logger._flush_task.cancelled()
assert logger.task is not None and logger.task.done() and not logger.task.cancelled()
# ---------------------------------------------------------------------------

176
tests/unit/proxy/test_tracing_endpoints.py
View file

@ -9,7 +9,7 @@ from typing import Final, Literal, TypedDict
from unittest.mock import AsyncMock, MagicMock, call
import pytest
from fastapi import FastAPI, HTTPException
from fastapi import FastAPI
from fastapi.testclient import TestClient
from httpx import Response
from pydantic import JsonValue, TypeAdapter
@ -28,7 +28,7 @@ from litellm.rust_bridge.trace.generated.models import TraceQueryHelp
from litellm.rust_bridge.trace.generated.responses import TraceSQLResponse
from litellm.rust_bridge.trace.generated.types import AllQueryScope, TraceScope
from litellm.rust_bridge.trace.storage import ClickHouseStorage, TraceStorageConfig
from litellm.tracing import Tenant, TraceReceiver, TracingPayloadTooLargeError
from litellm.tracing import TraceReceiver
SQL_ROWS: Final[tuple[Mapping[str, JsonValue], ...]] = (
{
@ -131,42 +131,37 @@ def _assert_validation_error(response: Response, error_type: str, location: tupl
@pytest.mark.parametrize(
("auth", "scope", "can_write"),
("auth", "scope"),
(
pytest.param(
UserAPIKeyAuth(token="admin-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN),
TraceScope(all_teams=1, user_id="", team_ids=()),
True,
id="admin",
),
pytest.param(
UserAPIKeyAuth(token="view-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
TraceScope(all_teams=1, user_id="", team_ids=()),
False,
id="view-only-admin",
),
pytest.param(
TEAM_KEY,
TraceScope(all_teams=0, user_id="user", team_ids=()),
True,
id="team-key",
),
pytest.param(
UserAPIKeyAuth(user_id="user", token="hashed-key", user_role=LitellmUserRoles.INTERNAL_USER),
TraceScope(all_teams=0, user_id="user", team_ids=()),
True,
id="teamless-key",
),
pytest.param(
UserAPIKeyAuth(token="hashed-key", user_role=LitellmUserRoles.INTERNAL_USER),
None,
True,
id="key-without-user-can-only-write",
),
),
)
def test_trace_read_and_write_permissions(
client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, scope: TraceScope | None, can_write: bool
def test_trace_read_permissions_with_retired_uploads(
client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, scope: TraceScope | None
) -> None:
client.app.dependency_overrides[user_api_key_auth] = lambda: auth
@ -178,17 +173,8 @@ def test_trace_read_and_write_permissions(
receiver.list_traces.assert_awaited_once_with(scope=scope, start_ms=1, end_ms=2, cursor=None)
write: Final = client.post("/v1/traces", json={})
assert write.status_code == (200 if can_write else 403), write.text
if not can_write:
receiver.ingest.assert_not_awaited()
return
receiver.ingest.assert_awaited_once()
tenant: Final = receiver.ingest.await_args.kwargs["tenant"]
assert (tenant.team_id, tenant.api_key_hash, tenant.org_id) == (
auth.team_id or "",
auth.token or "",
auth.org_id or "",
)
assert write.status_code == 410
receiver.ingest.assert_not_called()
@pytest.fixture
@ -215,68 +201,21 @@ def client() -> TestClient:
return TestClient(app)
@pytest.mark.parametrize("native_available", [True, False])
def test_501_when_tracing_not_enabled(
client: TestClient, native_available: bool, monkeypatch: pytest.MonkeyPatch
@pytest.mark.parametrize("endpoint", ("/v1/traces", "/v1/logs"))
@pytest.mark.parametrize("media_type", ("application/json", "application/x-protobuf"))
def test_gateway_uploads_return_setup_guidance_without_reading_the_body(
client: TestClient, receiver: MagicMock, endpoint: str, media_type: str
) -> None:
from google.rpc.status_pb2 import Status
from litellm.rust_bridge import loader
if not native_available:
monkeypatch.setattr(loader, "_cached_bridge", None)
response: Final = client.post("/v1/traces", content=b"")
assert response.status_code == 501
assert response.headers["content-type"] == "application/x-protobuf"
assert Status.FromString(response.content).message == (
"Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL."
if native_available
else ""
response: Final = client.post(endpoint, content=b"invalid payload", headers={"content-type": media_type})
assert response.status_code == 410
assert response.headers["content-type"] == media_type
message: Final = (
response.json()["message"] if media_type == "application/json" else Status.FromString(response.content).message
)
assert client.get("/v1/traces").status_code == 501
@pytest.mark.parametrize("endpoint", ("/v1/traces", "/v1/logs"))
def test_post_protobuf_returns_empty_protobuf(client, receiver, endpoint):
response = client.post(
endpoint,
content=b"\x0a\x00",
headers={"content-type": "application/x-protobuf", "content-encoding": "gzip"},
)
assert response.status_code == 200
assert response.content == b""
assert response.headers["content-type"] == "application/x-protobuf"
kwargs = receiver.ingest.call_args.kwargs
assert kwargs["body"] is not None
assert kwargs["logs"] is (endpoint == "/v1/logs")
assert kwargs["content_type"] == "application/x-protobuf"
assert kwargs["content_encoding"] == "gzip"
assert kwargs["tenant"].team_id == "team-research"
@pytest.mark.parametrize("endpoint", ("/v1/traces", "/v1/logs"))
def test_post_json_returns_empty_json(client, receiver, endpoint):
response = client.post(endpoint, content=b"{}", headers={"content-type": "application/json"})
assert response.status_code == 200
assert response.json() == {}
@pytest.mark.parametrize("endpoint", ("/v1/traces", "/v1/logs"))
def test_post_clickhouse_failure_is_503_with_retry_after(client, receiver, endpoint):
receiver.ingest.side_effect = RuntimeError("ClickHouse unavailable")
response = client.post(endpoint, content=b"", headers={"content-type": "application/x-protobuf"})
assert response.status_code == 503
assert response.headers["retry-after"] == str(tracing_endpoints.OTLP_RETRY_AFTER_SECONDS)
@pytest.mark.parametrize("endpoint", ("/v1/traces", "/v1/logs"))
def test_post_too_large_is_413(client, receiver, endpoint):
receiver.ingest.side_effect = TracingPayloadTooLargeError("OTLP body exceeds 10 bytes")
response = client.post(endpoint, content=b"x" * 20)
assert response.status_code == 413
from google.rpc.status_pb2 import Status
assert "exceeds" in Status.FromString(response.content).message
assert message == "Send traces and logs directly to the Lens endpoint shown in Lens setup."
receiver.ingest.assert_not_called()
def test_list_traces_passes_scope_window_and_cursor(client, receiver):
@ -531,11 +470,7 @@ def test_trace_page_rejects_unbounded_parameters(client: TestClient, receiver: M
receiver.get_trace.assert_not_awaited()
def test_invalid_export_and_cursor_are_client_errors(client, receiver):
from litellm.tracing.otlp_http import InvalidOTLPPayloadError
receiver.ingest.side_effect = InvalidOTLPPayloadError("invalid OTLP trace payload")
assert client.post("/v1/traces", content=b"broken").status_code == 400
def test_invalid_cursor_is_a_client_error(client: TestClient, receiver: MagicMock) -> None:
receiver.list_traces.side_effect = ValueError("Invalid trace cursor")
assert client.get("/v1/traces?cursor=broken").status_code == 400
@ -571,38 +506,6 @@ def test_key_without_user_cannot_read_traces(client: TestClient, auth: UserAPIKe
storage.query_help.assert_not_called()
@pytest.mark.parametrize("endpoint", ("/v1/traces", "/v1/logs"))
def test_view_only_admin_cannot_ingest_traces(client, receiver, endpoint):
client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
token="admin-key", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
)
response = client.post(endpoint, content=b"{}")
assert response.status_code == 403
receiver.ingest.assert_not_called()
@pytest.mark.parametrize(
"status_code, field, message",
[(401, "detail", "Invalid API key"), (403, "message", "Not allowed to ingest agent traces")],
)
def test_auth_failure_precedes_disabled_receiver(
client: TestClient, status_code: int, field: str, message: str
) -> None:
def unavailable() -> None:
return None
def authenticate() -> UserAPIKeyAuth:
if status_code == 401:
raise HTTPException(status_code=401, detail="Invalid API key")
return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
client.app.dependency_overrides[user_api_key_auth] = authenticate
client.app.dependency_overrides[tracing_endpoints.provide_receiver] = unavailable
response: Final = client.post("/v1/traces", content=b"{}", headers={"content-type": "application/json"})
assert response.status_code == status_code
assert response.json() == {field: message}
def test_disabled_receiver_precedes_read_scope_rejection(client: TestClient) -> None:
client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER
@ -610,33 +513,13 @@ def test_disabled_receiver_precedes_read_scope_rejection(client: TestClient) ->
response: Final = client.get("/v1/traces")
assert response.status_code == 501
assert response.json() == {
"detail": "Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL."
"detail": "Agent tracing is not enabled. Configure the Lens service and LITELLM_LENS_URL."
}
def test_injected_receiver_ingests_with_the_authenticated_tenant(client: TestClient) -> None:
storage: Final = MagicMock(spec=ClickHouseStorage)
storage.ingest = AsyncMock(return_value=1)
client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: TraceReceiver(storage)
response: Final = client.post(
"/v1/traces", content=b'{"resourceSpans": []}', headers={"content-type": "application/json"}
)
assert response.status_code == 200, response.text
assert response.json() == {}
storage.ingest.assert_awaited_once_with(
b'{"resourceSpans": []}',
"application/json",
Tenant(
team_id=TEAM_KEY.team_id or "",
api_key_hash=TEAM_KEY.token or "",
org_id=TEAM_KEY.org_id or "",
user_id=TEAM_KEY.user_id or "",
),
False,
)
def test_lifespan_receivers_are_app_local() -> None:
def test_lifespan_receivers_are_app_local(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LENS_URL", "http://lens.test")
monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", "test-service-token-with-32-characters")
first_storage: Final = MagicMock(spec=ClickHouseStorage)
first_storage.get_span = AsyncMock(return_value={**SPAN_DETAIL_RESPONSE, "span_id": "first-span"})
second_storage: Final = MagicMock(spec=ClickHouseStorage)
@ -671,8 +554,8 @@ def test_lifespan_receivers_are_app_local() -> None:
simultaneous: Final = first_client.get("/v1/traces/t1/spans/first-span?trace_ref=first-run")
first_response: Final = first_client.get("/v1/traces/t1/spans/first-span?trace_ref=first-run")
assert simultaneous.json() == first_response.json()
first_storage.ensure_schema.assert_awaited_once()
second_storage.ensure_schema.assert_awaited_once()
first_storage.ensure_schema.assert_not_awaited()
second_storage.ensure_schema.assert_not_awaited()
assert first_response.status_code == second_response.status_code == 200
assert first_response.json()["span_id"] == "first-span"
@ -692,7 +575,8 @@ def test_query_validation_precedes_trace_access_checks(client: TestClient, auth:
@pytest.mark.parametrize("enabled", [True, False])
def test_unavailable_lifespan_receiver_returns_501(enabled: bool) -> None:
def test_unconfigured_lifespan_receiver_returns_501(enabled: bool, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("LITELLM_LENS_URL", raising=False)
storage: Final = MagicMock(spec=ClickHouseStorage)
storage.ensure_schema = AsyncMock(side_effect=RuntimeError("storage unavailable"))
tracing: Final = TraceReceiver(storage)
@ -709,13 +593,15 @@ def test_unavailable_lifespan_receiver_returns_501(enabled: bool) -> None:
with TestClient(app) as client:
response: Final = client.get("/v1/traces")
assert response.status_code == 501
assert storage.ensure_schema.await_count == int(enabled)
storage.ensure_schema.assert_not_awaited()
storage.list_traces.assert_not_called()
def test_lens_reads_from_the_lifespan_storage() -> None:
def test_lens_reads_from_the_lifespan_storage(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy.lens.endpoints import router as lens_router
monkeypatch.setenv("LITELLM_LENS_URL", "http://lens.test")
monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", "test-service-token-with-32-characters")
storage: Final = MagicMock(spec=ClickHouseStorage)
storage.ensure_schema = AsyncMock()
storage.lens_sample = AsyncMock(return_value=[])

18
tests/unit/tracing/test_config.py
View file

@ -114,3 +114,21 @@ def test_legacy_reader_and_split_retention_fields_are_rejected() -> None:
},
{},
)
@pytest.mark.parametrize(
"settings,environ,enabled",
(
(None, {}, False),
({"store": {"type": "clickhouse"}}, {}, False),
({"store": {"type": "lens"}}, {}, True),
({"store": "lens"}, {}, False),
(None, {"LITELLM_LENS_URL": "http://lens"}, True),
),
)
def test_lens_enablement_requires_its_service_or_an_explicit_lens_store(
settings: object, environ: dict[str, str], enabled: bool
) -> None:
from litellm.tracing.config import is_lens_tracing_enabled
assert is_lens_tracing_enabled(settings, environ) is enabled

155
tests/unit/tracing/test_exporter.py
View file

@ -83,3 +83,158 @@ async def test_inflight_records_count_toward_the_queue_limit() -> None:
def test_oversized_event_is_rejected_before_queueing(value: str) -> None:
with pytest.raises(OverflowError):
encode_record({"messages": value})
@pytest.mark.asyncio
@pytest.mark.parametrize("status", (429, 502, 503, 504, 0))
async def test_transient_failures_retry_the_same_batch_and_recover(status: int) -> None:
requests: Final = asyncio.Queue[bytes]()
waits: Final = asyncio.Queue[float]()
async def retry_delay(seconds: float) -> None:
waits.put_nowait(seconds)
def respond(request: httpx.Request) -> httpx.Response:
requests.put_nowait(request.content)
if requests.qsize() == 3:
return httpx.Response(204)
if status == 0:
raise httpx.ConnectError("private storage host", request=request)
return httpx.Response(status)
async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(respond)) as client:
exporter: Final = LensExporter(client, sleep=retry_delay)
assert exporter.enqueue(b'{"id":1}')
exporter.start()
await exporter.aclose()
assert tuple(requests.get_nowait() for _ in range(3)) == (b'[{"id":1}]',) * 3
assert tuple(waits.get_nowait() for _ in range(2)) == (1.0, 2.0)
assert (exporter.rows_written, exporter.rows_dropped, exporter.buffered_bytes, exporter.buffered_events) == (
1,
0,
0,
0,
)
assert exporter.last_error == ""
@pytest.mark.asyncio
@pytest.mark.parametrize(
"status,attempts,reason", ((401, 1, "HTTP 401"), (500, 1, "HTTP 500"), (503, 3, "retry limit reached"))
)
async def test_failed_exports_are_counted_and_release_all_buffer_capacity(
status: int, attempts: int, reason: str
) -> None:
requests: Final = asyncio.Queue[bytes]()
async def no_wait(_: float) -> None:
return None
def reject(request: httpx.Request) -> httpx.Response:
requests.put_nowait(request.content)
return httpx.Response(status, text="private credentials")
async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(reject)) as client:
exporter: Final = LensExporter(client, sleep=no_wait)
assert exporter.enqueue(b"{}")
exporter.start()
exporter.start()
await exporter.aclose()
assert requests.qsize() == attempts
assert (exporter.rows_written, exporter.rows_dropped, exporter.buffered_bytes, exporter.buffered_events) == (
0,
1,
0,
0,
)
assert exporter.last_error == reason
assert not exporter.enqueue(b"{}")
assert exporter.rows_dropped == 2
@pytest.mark.asyncio
async def test_cancelled_inflight_export_drops_the_batch_and_pending_records() -> None:
started: Final = asyncio.Event()
async def block(_: httpx.Request) -> httpx.Response:
started.set()
await asyncio.Future[None]()
return httpx.Response(204)
async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(block)) as client:
exporter: Final = LensExporter(client)
assert exporter.enqueue(b"{}")
exporter.start()
await started.wait()
assert exporter.enqueue(b"{}")
assert exporter.task is not None
exporter.task.cancel()
await exporter.aclose()
assert (exporter.rows_written, exporter.rows_dropped, exporter.buffered_bytes, exporter.buffered_events) == (
0,
2,
0,
0,
)
@pytest.mark.asyncio
async def test_byte_budget_rejects_large_queue_and_shutdown_without_start_discards_it() -> None:
from litellm.tracing.exporter import MAX_BUFFER_BYTES
async with httpx.AsyncClient(base_url="http://lens") as client:
exporter: Final = LensExporter(client)
assert not exporter.enqueue(b"x" * (MAX_EVENT_BYTES + 1))
for _ in range(MAX_BUFFER_BYTES // MAX_EVENT_BYTES):
assert exporter.enqueue(b"x" * MAX_EVENT_BYTES)
assert not exporter.enqueue(b"x")
assert exporter.buffered_bytes == MAX_BUFFER_BYTES
await exporter.aclose()
assert exporter.rows_dropped == MAX_BUFFER_BYTES // MAX_EVENT_BYTES + 2
assert (exporter.buffered_bytes, exporter.buffered_events) == (0, 0)
@pytest.mark.asyncio
async def test_batches_stay_bounded_without_losing_or_reordering_records() -> None:
from litellm.tracing.exporter import MAX_BATCH_BYTES
bodies: Final = asyncio.Queue[bytes]()
def accept(request: httpx.Request) -> httpx.Response:
bodies.put_nowait(request.content)
return httpx.Response(204)
records: Final = tuple(encode_record({"id": index, "text": "x" * (MAX_EVENT_BYTES // 2)}) for index in range(10))
async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client:
exporter: Final = LensExporter(client)
for record in records:
assert exporter.enqueue(record)
exporter.start()
await exporter.aclose()
sent: Final = tuple(bodies.get_nowait() for _ in range(bodies.qsize()))
assert len(sent) == 2
assert all(len(body) <= MAX_BATCH_BYTES for body in sent)
assert [row["id"] for body in sent for row in json.loads(body)] == list(range(10))
assert exporter.rows_written == 10
assert exporter.rows_dropped == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("payload", (None, {"id": []}, {"id": "x", "messages": [{"content": "x" * MAX_EVENT_BYTES}]}))
async def test_invalid_callback_data_never_interrupts_model_requests(payload: object) -> None:
async with httpx.AsyncClient(base_url="http://lens") as client:
exporter: Final = LensExporter(client)
await exporter.async_log_failure_event({"standard_logging_object": payload}, None, None, None)
await exporter.aclose()
assert exporter.rows_written == 0
assert exporter.buffered_bytes == 0
assert exporter.rows_dropped == (0 if payload is None else 1)
def test_serialization_rejects_recursive_and_non_finite_payloads() -> None:
cyclic: Final[dict[str, object]] = {} # mutable-ok: deliberately constructs a cyclic callback payload
cyclic["self"] = cyclic
with pytest.raises(OverflowError):
encode_record(cyclic)
with pytest.raises(ValueError, match="Out of range float values"):
encode_record({"cost": float("nan")})

187
tests/unit/tracing/test_remote.py Normal file
View file

@ -0,0 +1,187 @@
import asyncio
import json
from collections.abc import Mapping
from typing import Final
import httpx
import pytest
from pydantic import JsonValue
from litellm.rust_bridge.trace.errors import TraceChanged
from litellm.rust_bridge.trace.generated.types import AllQueryScope, TraceScope
from litellm.tracing.remote import LensConnection, RemoteTraceStore, bounded_response
@pytest.mark.parametrize(
"url", ("", "ftp://lens", "http://", "http://user:secret@lens", "https://lens?q=1", "https://lens/#x")
)
def test_service_url_rejects_unsupported_or_credential_bearing_destinations(url: str) -> None:
with pytest.raises(ValueError, match="URL"):
LensConnection.from_env({"LITELLM_LENS_URL": url, "LITELLM_LENS_SERVICE_TOKEN": "x" * 32})
def test_connection_requires_a_strong_secret_and_preserves_the_configured_prefix() -> None:
with pytest.raises(ValueError, match="secret"):
LensConnection.from_env({"LITELLM_LENS_URL": "https://lens", "LITELLM_LENS_SERVICE_TOKEN": "short"})
connection: Final = LensConnection.from_env(
{"LITELLM_LENS_URL": "https://lens/prefix/", "LITELLM_LENS_SERVICE_TOKEN": "x" * 32}
)
assert connection.url == "https://lens/prefix"
assert "x" * 32 not in repr(connection)
async def _read_case(
store: RemoteTraceStore, operation: str, scope: TraceScope, query_scope: AllQueryScope
) -> tuple[JsonValue, Mapping[str, object]]:
match operation:
case "list":
return (
await store.list_traces(scope, 10, 20, "next", 17),
{
"operation": operation,
"scope": scope,
"start_ms": 10,
"end_ms": 20,
"cursor": "next",
"limit": 17,
},
)
case "trace":
return (
await store.get_trace("trace", scope, "ref", "next", 17),
{
"operation": operation,
"scope": scope,
"trace_id": "trace",
"trace_ref": "ref",
"cursor": "next",
"page_size": 17,
},
)
case "span":
return (
await store.get_span("trace", "span", scope, "ref"),
{
"operation": operation,
"scope": scope,
"trace_id": "trace",
"trace_ref": "ref",
"span_id": "span",
},
)
case "span_error":
return (
await store.get_span_error("trace", "span", scope, "ref", "next"),
{
"operation": operation,
"scope": scope,
"trace_id": "trace",
"trace_ref": "ref",
"span_id": "span",
"cursor": "next",
},
)
case "sql":
return (
json.loads(await store.query_sql("SELECT 1", query_scope, "unused-local-secret")),
{"operation": operation, "scope": query_scope, "sql": "SELECT 1"},
)
case "help":
return (
await store.query_help(query_scope, "unused-local-secret"),
{"operation": operation, "scope": query_scope},
)
case _:
return (
json.loads(await store.query("lens_sample", {"source": "traces"})),
{"operation": operation, "name": "lens_sample", "parameters": {"source": "traces"}},
)
@pytest.mark.asyncio
@pytest.mark.parametrize("operation", ("list", "trace", "span", "span_error", "sql", "help", "query"))
async def test_remote_reads_preserve_scope_and_pagination(operation: str) -> None:
scope: Final = TraceScope(all_teams=0, user_id="owner", team_ids=("team",))
query_scope: Final = AllQueryScope(kind="all")
requests: Final = asyncio.Queue[httpx.Request]()
def accept(request: httpx.Request) -> httpx.Response:
requests.put_nowait(request)
return httpx.Response(200, json={"data": [{"value": "safe"}]})
async with httpx.AsyncClient(base_url="http://lens/prefix/", transport=httpx.MockTransport(accept)) as client:
store: Final = RemoteTraceStore(client)
result, expected = await _read_case(store, operation, scope, query_scope)
request: Final = requests.get_nowait()
assert request.url.path == "/prefix/internal/read"
assert json.loads(request.content) == json.loads(json.dumps(expected))
assert result == {"data": [{"value": "safe"}]}
assert b"unused-local-secret" not in request.content
@pytest.mark.asyncio
@pytest.mark.parametrize(
"status,error",
((400, ValueError), (409, TraceChanged), (413, OverflowError), (503, RuntimeError), (302, RuntimeError)),
)
async def test_remote_failures_preserve_public_error_categories_without_leaking_storage_details(
status: int, error: type[Exception]
) -> None:
async with httpx.AsyncClient(
base_url="http://lens",
transport=httpx.MockTransport(lambda request: httpx.Response(status, text="private storage credentials")),
) as client:
with pytest.raises(error) as failure:
await RemoteTraceStore(client).get_trace("trace", TraceScope(all_teams=1, user_id="", team_ids=()), "ref")
assert "private storage credentials" not in str(failure.value)
@pytest.mark.asyncio
async def test_network_failure_is_retryable_without_exposing_the_remote_url() -> None:
def fail(request: httpx.Request) -> httpx.Response:
raise httpx.ConnectError("private storage credentials", request=request)
async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(fail)) as client:
with pytest.raises(RuntimeError, match="Lens trace storage is unavailable"):
await RemoteTraceStore(client).query_help(AllQueryScope(kind="all"), "secret")
@pytest.mark.asyncio
async def test_reads_reject_oversized_responses_and_invalid_json() -> None:
with pytest.raises(RuntimeError, match="size limit"):
await bounded_response(httpx.Response(200, content=b"abcd"), 3)
assert await bounded_response(httpx.Response(200, content=b"abcd"), 4) == b"abcd"
async with httpx.AsyncClient(
base_url="http://lens", transport=httpx.MockTransport(lambda request: httpx.Response(200, content=b"{"))
) as client:
with pytest.raises(ValueError, match="Invalid Lens response"):
await RemoteTraceStore(client).query_help(AllQueryScope(kind="all"), "secret")
@pytest.mark.asyncio
async def test_gateway_cannot_relay_otlp_or_write_arbitrary_tables() -> None:
def fail(request: httpx.Request) -> httpx.Response:
raise AssertionError("No network access is allowed for schema setup or refused uploads")
async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(fail)) as client:
store: Final = RemoteTraceStore(client)
await store.ensure_schema()
with pytest.raises(RuntimeError, match="directly"):
await store.ingest(b"{}", "application/json", {})
with pytest.raises(ValueError, match="request records"):
await store.insert_rows("otel_traces", ())
@pytest.mark.asyncio
async def test_request_records_use_the_internal_service_endpoint() -> None:
requests: Final = asyncio.Queue[httpx.Request]()
def accept(request: httpx.Request) -> httpx.Response:
requests.put_nowait(request)
return httpx.Response(204)
async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client:
await RemoteTraceStore(client).insert_rows("spend_logs", ({"request_id": "r"},))
request: Final = requests.get_nowait()
assert request.url.path == "/internal/spend"
assert json.loads(request.content) == [{"request_id": "r"}]

15
ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx
View file

@ -589,14 +589,21 @@ it.each([false, true])(
async (enabled) => {
window.history.replaceState({}, "", "/lens/");
testQueryClient.clear();
proxy.get.mockImplementation(async (path) =>
path === "/lens" ? { lenses: [], workers: [], tracing_enabled: enabled } : { data: [] },
);
proxy.get.mockImplementation(async (path) => {
if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: enabled };
if (path === "/lens/service")
return {
url: "https://traces.test",
connected: true,
status: { storage_ready: true, credentials_ready: true },
};
return { data: [] };
});
const user = userEvent.setup();
renderWithProviders(<InvestigationsView />);
const guide = within(await screen.findByRole("region", { name: "Get Lens running" }));
expect(guide.getByRole("button", { name: /Send your first trace/ })).toHaveAttribute("aria-expanded", "true");
expect(guide.getByRole("button", { name: "Check for traces" })).toBeVisible();
expect(await guide.findByRole("button", { name: "Check for traces" })).toBeVisible();
await user.click(guide.getByRole("button", { name: /Connect a worker/ }));
expect(guide.getByRole("button", { name: "Connect worker" })).toBeDisabled();
await user.click(guide.getByRole("button", { name: /Run your first investigation/ }));

2
ui/litellm-dashboard/src/components/lens/onboarding/tracing/TracingSetupCard.tsx
View file

@ -760,7 +760,7 @@ export function TracingSetupCard(props: TracingSetupProps) {
<p className="mt-2 text-sm leading-6 text-muted-foreground">
{enabled
? "Send your agent’s runs to LiteLLM to see its inputs, outputs, and tool calls."
: "Tracing needs ClickHouse and a small update to your LiteLLM proxy configuration."}
: "Tracing needs a Lens service with ClickHouse access and a connection from LiteLLM."}
</p>
<TracingSetupFields {...props} />
</div>

20
ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesSection.integration.test.tsx
View file

@ -77,7 +77,11 @@ describe("AgentTracesSection", () => {
setupIntersectionMocking(vi.fn);
testQueryClient.clear();
vi.mocked(agentTraceListCall).mockReset();
vi.mocked(apiClient.get).mockResolvedValue({ data: [] });
vi.mocked(apiClient.get).mockImplementation(async (path) =>
path === "/lens/service"
? { url: "https://traces.test", connected: true, status: { storage_ready: true, credentials_ready: true } }
: { data: [] },
);
vi.mocked(apiClient.post).mockImplementation(async (_path, options) => {
const body = options?.body as { traces: { trace_id: string; trace_ref?: string }[] };
return body.traces.map((trace) => ({ ...trace, finding_count: null }));
@ -163,12 +167,12 @@ describe("AgentTracesSection", () => {
const card = await screen.findByTestId("tracing-setup-card");
expect(card).toHaveTextContent("Tracing is not enabled");
expect(screen.queryByRole("button", { name: "Preview sample" })).not.toBeInTheDocument();
expect(card).toHaveTextContent("type: clickhouse");
expect(card).toHaveTextContent("url: os.environ/CLICKHOUSE_URL");
expect(card).toHaveTextContent("LITELLM_LENS_URL");
expect(card).toHaveTextContent("LITELLM_LENS_SERVICE_TOKEN");
expect(screen.getByRole("button", { name: "Check setup" })).toBeEnabled();
expect(card).not.toHaveTextContent(/langsmith/i);
expect(card).toHaveTextContent("ClickHouse and proxy setup");
expect(card).toHaveTextContent("Ask your proxy administrator");
expect(card).toHaveTextContent("Lens service setup");
expect(card).toHaveTextContent("Run the Lens service");
});
it("shows the waiting guide when tracing is on but no runs have arrived", async () => {
@ -177,7 +181,7 @@ describe("AgentTracesSection", () => {
const card = await screen.findByTestId("tracing-setup-card");
expect(card).toHaveTextContent("Connect your agent");
expect(card).toHaveTextContent("Waiting for your first trace");
expect(await screen.findByText("Waiting for your first trace")).toBeInTheDocument();
expect(screen.queryByRole("button", { name: "Preview sample" })).not.toBeInTheDocument();
expect(card).not.toHaveTextContent("store: clickhouse");
});
@ -201,7 +205,7 @@ describe("AgentTracesSection", () => {
vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), data: [] });
fireEvent.click(checkSetup);
expect(await screen.findByRole("heading", { name: "Connect your agent" })).toBeVisible();
expect(screen.getByText("Waiting for your first trace")).toBeVisible();
expect(await screen.findByText("Waiting for your first trace")).toBeVisible();
vi.mocked(agentTraceListCall).mockResolvedValue(traceList as TracePage);
fireEvent.click(screen.getByRole("button", { name: "Check for traces" }));
expect(await screen.findAllByTestId("agent-trace-row")).toHaveLength(runs.length);
@ -278,7 +282,7 @@ describe("AgentTracesSection", () => {
const card = await screen.findByTestId("tracing-setup-card");
expect(card).toHaveTextContent("Tracing is not enabled");
expect(card).toHaveTextContent("url: os.environ/CLICKHOUSE_URL");
expect(card).toHaveTextContent("LITELLM_LENS_SERVICE_TOKEN");
});
it("lists uninvestigated runs without presenting tool errors as failures", async () => {