From 6d819de6aea36af47ab7365c351272f9594d8346 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 7 Oct 2026 13:23:04 -0700 Subject: [PATCH] fix(lens): address CI failures and review findings --- .github/workflows/lens-worker.yml | 1 + .github/workflows/test-rust.yml | 8 + deploy/lens/README.md | 4 + helm/litellm-helm/templates/_helpers.tpl | 6 + .../templates/lens/deployment.yaml | 8 +- .../tests/lens_service_tests.yaml | 167 +++++++++++++++ helm/litellm-helm/values.yaml | 2 + helm/litellm/templates/_helpers.tpl | 6 + helm/litellm/templates/lens/deployment.yaml | 8 +- helm/litellm/tests/lens_service_tests.yaml | 196 ++++++++++++++++++ helm/litellm/values.yaml | 2 + .../migration.sql | 2 +- litellm-rust/crates/lens/src/agent.rs | 13 +- litellm-rust/crates/lens/tests/worker.rs | 106 ++++++++++ litellm/integrations/otel/model/spans.py | 3 +- litellm/proxy/db/prisma_query_span.py | 1 + litellm/proxy/lens/endpoints.py | 2 +- litellm/proxy/tracing_runtime.py | 4 +- litellm/tracing/exporter.py | 7 +- litellm/tracing/remote.py | 65 ++++-- tests/test_litellm_rust/test_traces.py | 11 +- .../common_utils/test_http_parsing_utils.py | 17 +- tests/unit/proxy/db/test_prisma_query_span.py | 3 +- tests/unit/proxy/lens/test_endpoints.py | 154 +++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 61 +++--- tests/unit/proxy/test_tracing_endpoints.py | 176 +++------------- tests/unit/tracing/test_config.py | 18 ++ tests/unit/tracing/test_exporter.py | 155 ++++++++++++++ tests/unit/tracing/test_remote.py | 187 +++++++++++++++++ .../InvestigationsView.integration.test.tsx | 15 +- .../onboarding/tracing/TracingSetupCard.tsx | 2 +- .../AgentTracesSection.integration.test.tsx | 20 +- 32 files changed, 1191 insertions(+), 239 deletions(-) create mode 100644 helm/litellm-helm/tests/lens_service_tests.yaml create mode 100644 helm/litellm/tests/lens_service_tests.yaml create mode 100644 tests/unit/tracing/test_remote.py diff --git a/.github/workflows/lens-worker.yml b/.github/workflows/lens-worker.yml index d440dd6e5a7..46fb292060d 100644 --- a/.github/workflows/lens-worker.yml +++ b/.github/workflows/lens-worker.yml @@ -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 diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 740cfc222a8..c07487d5b49 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -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 diff --git a/deploy/lens/README.md b/deploy/lens/README.md index 5ad40e2d078..7fe964f2f8b 100644 --- a/deploy/lens/README.md +++ b/deploy/lens/README.md @@ -82,9 +82,13 @@ lensWorker: clickhouseSecret: name: litellm-lens-clickhouse key: url + clickhouseDatabase: litellm + retentionDays: 14 publicUrl: https:///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: diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl index 7ba8331bdcd..5e5b47b586a 100644 --- a/helm/litellm-helm/templates/_helpers.tpl +++ b/helm/litellm-helm/templates/_helpers.tpl @@ -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 -}} diff --git a/helm/litellm-helm/templates/lens/deployment.yaml b/helm/litellm-helm/templates/lens/deployment.yaml index aa7686ad16b..dee141acde8 100644 --- a/helm/litellm-helm/templates/lens/deployment.yaml +++ b/helm/litellm-helm/templates/lens/deployment.yaml @@ -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: diff --git a/helm/litellm-helm/tests/lens_service_tests.yaml b/helm/litellm-helm/tests/lens_service_tests.yaml new file mode 100644 index 00000000000..197f447f5f3 --- /dev/null +++ b/helm/litellm-helm/tests/lens_service_tests.yaml @@ -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 diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 2e2b73ddc9a..821557fd116 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -667,6 +667,8 @@ lensWorker: serviceTokenSecret: name: "" key: service-token + clickhouseDatabase: litellm + retentionDays: 14 clickhouseSecret: name: "" key: url diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 0b5ccc14ddb..9c76e2da748 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -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 -}} diff --git a/helm/litellm/templates/lens/deployment.yaml b/helm/litellm/templates/lens/deployment.yaml index f3ba3780231..93772a7a650 100644 --- a/helm/litellm/templates/lens/deployment.yaml +++ b/helm/litellm/templates/lens/deployment.yaml @@ -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: diff --git a/helm/litellm/tests/lens_service_tests.yaml b/helm/litellm/tests/lens_service_tests.yaml new file mode 100644 index 00000000000..c5504025572 --- /dev/null +++ b/helm/litellm/tests/lens_service_tests.yaml @@ -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 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index e8820102098..3fd3245d166 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -644,6 +644,8 @@ lensWorker: serviceTokenSecret: name: "" key: service-token + clickhouseDatabase: litellm + retentionDays: 14 clickhouseSecret: name: "" key: url diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261007000100_lens_ingestion_keys/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261007000100_lens_ingestion_keys/migration.sql index b68370a01d0..a05c066da43 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261007000100_lens_ingestion_keys/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261007000100_lens_ingestion_keys/migration.sql @@ -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") diff --git a/litellm-rust/crates/lens/src/agent.rs b/litellm-rust/crates/lens/src/agent.rs index 7f6022089f4..4669ab750dd 100644 --- a/litellm-rust/crates/lens/src/agent.rs +++ b/litellm-rust/crates/lens/src/agent.rs @@ -269,17 +269,16 @@ pub async fn run( 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 diff --git a/litellm-rust/crates/lens/tests/worker.rs b/litellm-rust/crates/lens/tests/worker.rs index f2d2d525163..b6e2cb71a5b 100644 --- a/litellm-rust/crates/lens/tests/worker.rs +++ b/litellm-rust/crates/lens/tests/worker.rs @@ -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::()); + 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); +} diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 3b17aea3cd8..7cdb375f9d3 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -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_.]*") diff --git a/litellm/proxy/db/prisma_query_span.py b/litellm/proxy/db/prisma_query_span.py index d99f4dcc768..75625c6c0f2 100644 --- a/litellm/proxy/db/prisma_query_span.py +++ b/litellm/proxy/db/prisma_query_span.py @@ -76,6 +76,7 @@ _VERB_BY_KEYWORD: Final[Mapping[str, str]] = MappingProxyType( "REFRESH": "ddl", "TRUNCATE": "delete", "SET": "set", + "LOCK": "lock", } ) diff --git a/litellm/proxy/lens/endpoints.py b/litellm/proxy/lens/endpoints.py index e070324f004..9e655e83645 100644 --- a/litellm/proxy/lens/endpoints.py +++ b/litellm/proxy/lens/endpoints.py @@ -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()) diff --git a/litellm/proxy/tracing_runtime.py b/litellm/proxy/tracing_runtime.py index 437d9fb6e8e..dd739a6803c 100644 --- a/litellm/proxy/tracing_runtime.py +++ b/litellm/proxy/tracing_runtime.py @@ -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 diff --git a/litellm/tracing/exporter.py b/litellm/tracing/exporter.py index ad9b6a2eb4e..dae78dd0d7a 100644 --- a/litellm/tracing/exporter.py +++ b/litellm/tracing/exporter.py @@ -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 diff --git a/litellm/tracing/remote.py b/litellm/tracing/remote.py index 8612dbc82a1..2be5a7bd061 100644 --- a/litellm/tracing/remote.py +++ b/litellm/tracing/remote.py @@ -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": diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index 50982f9774e..f07bc776bce 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -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 diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index fdc24708ff6..222a2cb329b 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -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() diff --git a/tests/unit/proxy/db/test_prisma_query_span.py b/tests/unit/proxy/db/test_prisma_query_span.py index abd95dec517..9f2ba28ea7e 100644 --- a/tests/unit/proxy/db/test_prisma_query_span.py +++ b/tests/unit/proxy/db/test_prisma_query_span.py @@ -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)), ], ) diff --git a/tests/unit/proxy/lens/test_endpoints.py b/tests/unit/proxy/lens/test_endpoints.py index 74eb649aad0..d51e9d1b739 100644 --- a/tests/unit/proxy/lens/test_endpoints.py +++ b/tests/unit/proxy/lens/test_endpoints.py @@ -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 diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index c1f2b9c876f..903557194ec 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -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() + # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index a3c8c8d9181..123538e49a4 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -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=[]) diff --git a/tests/unit/tracing/test_config.py b/tests/unit/tracing/test_config.py index 030c4247c62..a442e098083 100644 --- a/tests/unit/tracing/test_config.py +++ b/tests/unit/tracing/test_config.py @@ -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 diff --git a/tests/unit/tracing/test_exporter.py b/tests/unit/tracing/test_exporter.py index e0e96e1bf79..da40e46ee3c 100644 --- a/tests/unit/tracing/test_exporter.py +++ b/tests/unit/tracing/test_exporter.py @@ -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")}) diff --git a/tests/unit/tracing/test_remote.py b/tests/unit/tracing/test_remote.py new file mode 100644 index 00000000000..2572c2a2044 --- /dev/null +++ b/tests/unit/tracing/test_remote.py @@ -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"}] diff --git a/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx index 86e1bb9f316..f856faa6d8e 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx @@ -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(); 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/ })); diff --git a/ui/litellm-dashboard/src/components/lens/onboarding/tracing/TracingSetupCard.tsx b/ui/litellm-dashboard/src/components/lens/onboarding/tracing/TracingSetupCard.tsx index 33009ea77f7..299954ada99 100644 --- a/ui/litellm-dashboard/src/components/lens/onboarding/tracing/TracingSetupCard.tsx +++ b/ui/litellm-dashboard/src/components/lens/onboarding/tracing/TracingSetupCard.tsx @@ -760,7 +760,7 @@ export function TracingSetupCard(props: TracingSetupProps) {

{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."}

diff --git a/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesSection.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesSection.integration.test.tsx index db43b3610b9..37937dd58d1 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesSection.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesSection.integration.test.tsx @@ -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 () => {