diff --git a/.gitignore b/.gitignore index 399458eced5..763f8db2940 100644 --- a/.gitignore +++ b/.gitignore @@ -151,3 +151,6 @@ litellm.log .coverage-rust coverage-rust.xml + +# make lens-dev worker token, generated config and logs +.lens-dev/ diff --git a/Makefile b/Makefile index e512960949c..40573fb68a3 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ .PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc test-unit-proxy-root \ test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \ test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \ - test-rust-extension rust-sqlx-prepare \ + test-rust-extension rust-sqlx-prepare lens-dev \ info lint lint-inner lint-dev lint-checks format \ lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \ lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \ @@ -58,6 +58,7 @@ help: @echo " make test-unit-helm - Run helm unit tests" @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" @echo " make rust-sqlx-prepare - Refresh litellm-rust/crates/db/.sqlx against a migrated Postgres container" + @echo " make lens-dev - Run proxy + Lens worker + hot-reload dashboard (LENS_DEV_PROXY_PORT, LENS_DEV_UI_PORT)" @echo "" @echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide" @echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine." @@ -311,6 +312,9 @@ test-rust-extension: rust-sqlx-prepare: cd litellm-rust && cargo run -p litellm-db-testing --bin sqlx-prepare +lens-dev: + ./scripts/lens_dev.sh + test: install-test-deps $(UV_RUN) pytest tests/ diff --git a/deploy/lens/README.md b/deploy/lens/README.md index 4b65c88afcd..4f78b7bfb59 100644 --- a/deploy/lens/README.md +++ b/deploy/lens/README.md @@ -33,7 +33,7 @@ For deployments managed with Compose, download `compose.yaml` and provide `LITEL docker compose --env-file /path/to/lens.env -f compose.yaml up -d ``` -Developers can build locally with `LENS_WORKER_IMAGE=litellm-lens-worker:local docker compose -f deploy/lens/compose.yaml -f deploy/lens/compose.build.yaml up -d --build` +Developers can build locally with `LENS_WORKER_IMAGE=litellm-lens-worker:local docker compose -f deploy/lens/compose.yaml -f deploy/lens/compose.build.yaml up -d --build`. To work on Lens itself, `make lens-dev` runs the proxy, a worker from source and the hot-reload dashboard together; set `LENS_DEV_PROXY_PORT` / `LENS_DEV_UI_PORT` to move them off 4000/3000 The generated command gives the worker 1 GiB of temporary memory-backed storage, shared across parallel reviews. Change `size=1g` in the Docker command or set `LENS_WORKER_TMP_SIZE` with Compose to fit your server and workload. A storage failure marks the scan as failed, cleans up temporary traces, and leaves the worker available for other scans; it does not silently truncate the review. Existing workers must be recreated with the new image and mount options diff --git a/scripts/lens_dev.sh b/scripts/lens_dev.sh new file mode 100755 index 00000000000..7ab4b9bb433 --- /dev/null +++ b/scripts/lens_dev.sh @@ -0,0 +1,268 @@ +#!/usr/bin/env bash +# One-command Lens local dev loop: proxy + Lens worker + hot-reload dashboard. +# +# LENS_DEV_PROXY_PORT proxy port (default 4000) +# LENS_DEV_UI_PORT next dev port (default 3000) +# LENS_DEV_MASTER_KEY master key, also the admin UI password +# (default: random, generated once into .lens-dev/master_key) +# LENS_DEV_CONFIG proxy config to use instead of the generated one +# LENS_DEV_DATABASE_URL Postgres URL (default: the tracing stack's litellm DB on :15432) +# LENS_DEV_REBUILD_RUST=1 rebuild the Rust bridge even if it imports +# +# State (master key, worker token, generated config, logs) lives in .lens-dev/ (gitignored). +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +proxy_port="${LENS_DEV_PROXY_PORT:-4000}" +ui_port="${LENS_DEV_UI_PORT:-3000}" +state_dir="${LENS_DEV_STATE_DIR:-$repo_root/.lens-dev}" +log_dir="$state_dir/logs" +token_file="$state_dir/worker_token" +key_file="$state_dir/master_key" +proxy_url="http://localhost:$proxy_port" +py="${LENS_DEV_PYTHON:-$repo_root/.venv/bin/python}" +database_url="${LENS_DEV_DATABASE_URL:-postgresql://litellm:litellm@127.0.0.1:15432/litellm}" +clickhouse_url=http://default:local-tracing@127.0.0.1:18123 +master_key="" +pids=() + +die() { echo "lens-dev: $*" >&2; exit 1; } +listening() { lsof -nP -iTCP:"$1" -sTCP:LISTEN >/dev/null 2>&1; } + +# A fixed key would let anyone who can reach the proxy sign in as admin, so default to a +# random key generated once per checkout and kept next to the worker token. +load_master_key() { + if [ -n "${LENS_DEV_MASTER_KEY:-}" ]; then + master_key="$LENS_DEV_MASTER_KEY" + return + fi + if [ ! -s "$key_file" ]; then + (umask 077 && printf 'sk-%s\n' "$(openssl rand -hex 24)" > "$key_file") + fi + master_key="$(cat "$key_file")" +} + +# Only reuse a listener on 15432/18123 if it accepts the tracing stack's credentials; +# start the compose service when nothing is listening; fail if something else is. +# A LENS_DEV_DATABASE_URL is left to the proxy, which may use Prisma-only URL params. +ensure_services() { + local services=() + if [ -z "${LENS_DEV_DATABASE_URL:-}" ] && listening 15432; then + "$py" -c 'import sys, psycopg; psycopg.connect(sys.argv[1], connect_timeout=5).close()' "$database_url" 2>/dev/null \ + || die "port 15432 is taken by something that isn't the tracing Postgres (litellm/litellm)" + elif [ -z "${LENS_DEV_DATABASE_URL:-}" ]; then + services+=(db) + fi + if listening 18123; then + [ "$(curl -fsS --max-time 5 "$clickhouse_url/?query=SELECT%201" 2>/dev/null)" = 1 ] \ + || die "port 18123 is taken by something that isn't the tracing ClickHouse (default/local-tracing)" + else + services+=(clickhouse) + fi + if [ "${#services[@]}" -gt 0 ]; then + docker compose -f docker/docker-compose.tracing.yml up -d --wait "${services[@]}" + else + echo "lens-dev: reusing running Postgres and ClickHouse" + fi +} + +write_default_config() { + cat > "$1" <<'EOF' +model_list: + - model_name: gpt-4.1-mini + litellm_params: + model: openai/gpt-4.1-mini + api_key: os.environ/OPENAI_API_KEY +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + store_prompts_in_spend_logs: true + tracing: + store: + type: clickhouse + url: os.environ/CLICKHOUSE_URL + retention_days: 14 +EOF +} + +# litellm's implicit load_dotenv() walks up from a worktree into the parent checkout's +# .env and picks up REDIS_* / UI_* from there. LITELLM_MODE=PRODUCTION turns that off; +# this prints export lines for the same .env minus those vars, so provider keys still load. +dotenv_exports() { + "$py" - <<'PY' +import os, re, shlex +from dotenv import dotenv_values, find_dotenv + +path = find_dotenv(usecwd=True) +skip = re.compile(r"REDIS_.*|UI_USERNAME|UI_PASSWORD|LITELLM_MODE|ANTHROPIC_BASE_URL|ANTHROPIC_AUTH_TOKEN|ANTHROPIC_CUSTOM_HEADERS|OPENAI_BASE_URL|OPENAI_API_BASE") +for key, value in (dotenv_values(path) if path else {}).items(): + if value is not None and key not in os.environ and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key) and not skip.fullmatch(key): + print(f"export {key}={shlex.quote(value)}") +PY +} + +# Run in the proxy's subshell: drop inherited settings that would point it at someone +# else's services, then set the local stack's. +proxy_env() { + local var + # Claude Code and similar tools export these; provider calls would go to them. + unset ANTHROPIC_BASE_URL ANTHROPIC_AUTH_TOKEN ANTHROPIC_CUSTOM_HEADERS OPENAI_BASE_URL OPENAI_API_BASE + for var in $(compgen -e | grep '^REDIS_' || true); do unset "$var"; done + eval "$1" + export LITELLM_MODE=PRODUCTION + export LITELLM_MASTER_KEY="$master_key" + if [ "$master_key" = sk-1234 ]; then export LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true; fi + export LITELLM_SALT_KEY=sk-local-tracing-salt-key + export DATABASE_URL="$database_url" + export STORE_MODEL_IN_DB=True + export CLICKHOUSE_URL="$clickhouse_url" + export CLICKHOUSE_DATABASE=litellm + export LITELLM_LOCAL_MODEL_COST_MAP=True + export PROXY_BASE_URL="$proxy_url" + export UI_USERNAME=admin + export UI_PASSWORD="$master_key" +} + +# POST JSON as the admin and print one field of the response; dies with the body on failure. +admin_post() { + local body + body="$(curl -sS --fail-with-body "$proxy_url$1" -H "Authorization: Bearer $master_key" \ + -H "Content-Type: application/json" -d "$2")" || die "POST $1 failed: $body" + "$py" -c 'import json, sys; print(json.loads(sys.argv[1])[sys.argv[2]])' "$body" "$3" +} + +register_worker() { + local key_hash worker_token + key_hash="$(admin_post /key/generate "{\"key_alias\": \"lens-dev-$(date +%s)\"}" token)" + worker_token="$(admin_post /lens/workers/register "{\"name\": \"lens-dev\", \"analysis_key_id\": \"$key_hash\"}" token)" + (umask 077 && printf '%s\n' "$worker_token" > "$token_file") + echo "lens-dev: registered a new Lens worker (token in $token_file)" +} + +# Auth runs before the handler, so protocol_version=1 answers 401 for a bad token and +# 409 for a good one without claiming a job. +ensure_worker_token() { + local status + if [ ! -s "$token_file" ]; then + register_worker + return + fi + status="$(curl -sS -o /dev/null -w '%{http_code}' -X POST "$proxy_url/lens/worker/claim?protocol_version=1" \ + -H "Authorization: Bearer $(cat "$token_file")")" + case "$status" in + 409) echo "lens-dev: reusing worker token from $token_file" ;; + 401) echo "lens-dev: stored worker token was rejected"; register_worker ;; + *) die "unexpected HTTP $status checking the worker token" ;; + esac +} + +wait_for_proxy() { + local proxy_pid="$1" + echo "lens-dev: waiting for the proxy (log: $log_dir/proxy.log)" + for _ in $(seq 1 300); do + kill -0 "$proxy_pid" 2>/dev/null || die "proxy exited; see $log_dir/proxy.log" + curl -fsS "$proxy_url/health/readiness" -H "Authorization: Bearer $master_key" >/dev/null 2>&1 && return + sleep 1 + done + die "proxy not ready after 300s; see $log_dir/proxy.log" +} + +# Children run in their own process groups (set -m), so killing -pid takes their trees too. +cleanup() { + local alive pid + trap - EXIT INT TERM + [ "${#pids[@]}" -gt 0 ] || return 0 + echo "lens-dev: stopping" + for pid in "${pids[@]}"; do kill -TERM -- "-$pid" 2>/dev/null || true; done + for _ in $(seq 1 20); do + alive=0 + for pid in "${pids[@]}"; do kill -0 "$pid" 2>/dev/null && alive=1; done + [ "$alive" = 0 ] && break + sleep 0.5 + done + for pid in "${pids[@]}"; do kill -KILL -- "-$pid" 2>/dev/null || true; done +} + +main() { + local config_file exports proxy_pid pid key_hint + if [ -n "${LENS_DEV_CONFIG:-}" ]; then + [ -f "$LENS_DEV_CONFIG" ] || die "LENS_DEV_CONFIG not found: $LENS_DEV_CONFIG" + config_file="$(cd "$(dirname "$LENS_DEV_CONFIG")" && pwd)/$(basename "$LENS_DEV_CONFIG")" + fi + cd "$repo_root" + + listening "$proxy_port" && die "port $proxy_port is in use; set LENS_DEV_PROXY_PORT" + listening "$ui_port" && die "port $ui_port is in use; set LENS_DEV_UI_PORT" + [ "$proxy_port" != "$ui_port" ] || die "proxy and UI ports must differ" + mkdir -p "$log_dir" + load_master_key + + uv sync --inexact --frozen --extra proxy --group proxy-dev --no-install-project + ensure_services + "$py" scripts/prisma_generate_if_needed.py + + if [ "${LENS_DEV_REBUILD_RUST:-0}" = "1" ] || ! "$py" -c "import litellm.rust_bridge._native" >/dev/null 2>&1; then + echo "lens-dev: building the Rust bridge (litellm.rust_bridge._native); the ClickHouse trace store uses it" + VIRTUAL_ENV="$repo_root/.venv" uvx --from maturin==1.15.0 maturin develop \ + --release --manifest-path litellm-rust/crates/python-bridge/Cargo.toml --features extension-module + fi + + if [ ! -x ui/litellm-dashboard/node_modules/.bin/next ]; then + (cd ui/litellm-dashboard && "$repo_root/scripts/with_dashboard_node.sh" npm ci) + fi + + if [ -z "${config_file:-}" ]; then + config_file="$state_dir/config.yaml" + write_default_config "$config_file" + fi + exports="$(dotenv_exports)" + + trap cleanup EXIT + trap 'exit 130' INT TERM + set -m + + ( + proxy_env "$exports" + exec "$py" litellm/proxy/proxy_cli.py --config "$config_file" --host 127.0.0.1 --port "$proxy_port" + ) < /dev/null > "$log_dir/proxy.log" 2>&1 & + proxy_pid=$! + pids+=("$proxy_pid") + + ( + cd ui/litellm-dashboard + NEXT_PUBLIC_BASE_URL="$proxy_url" exec "$repo_root/scripts/with_dashboard_node.sh" npx next dev -p "$ui_port" + ) < /dev/null > "$log_dir/ui.log" 2>&1 & + pids+=("$!") + + wait_for_proxy "$proxy_pid" + ensure_worker_token + + LITELLM_MODE=PRODUCTION LITELLM_URL="$proxy_url" LENS_WORKER_TOKEN="$(cat "$token_file")" \ + "$py" -c "import asyncio, logging; from litellm.proxy.lens.worker import main; logging.basicConfig(level=logging.INFO); asyncio.run(main())" \ + < /dev/null > "$log_dir/worker.log" 2>&1 & + pids+=("$!") + + key_hint="password in $key_file" + [ -z "${LENS_DEV_MASTER_KEY:-}" ] || key_hint="password from LENS_DEV_MASTER_KEY" + cat </dev/null || die "a child process (pid $pid) exited; check the logs above" + done + sleep 2 + done +} + +# Sourcing (tests) only defines the functions. +if [ "${BASH_SOURCE[0]}" = "$0" ]; then + main "$@" +fi diff --git a/tests/unit/test_lens_dev.py b/tests/unit/test_lens_dev.py new file mode 100644 index 00000000000..e14b27fd27e --- /dev/null +++ b/tests/unit/test_lens_dev.py @@ -0,0 +1,151 @@ +import os +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +SCRIPT = ROOT / "scripts" / "lens_dev.sh" + +# Fake curl: answers the worker-token check with $CLAIM_STATUS, and key/generate and +# workers/register with a JSON "token". Every call is appended to $CURL_LOG. +FAKE_CURL = """#!/bin/sh +echo "$@" >> "$CURL_LOG" +case "$*" in + *worker/claim*) printf '%s' "$CLAIM_STATUS" ;; + */key/generate*) printf '{"token": "%064d"}' 0 ;; + */lens/workers/register*) printf '{"token": "lens-fresh"}' ;; +esac +""" + + +def _run(tmp_path: Path, snippet: str, **env: str) -> subprocess.CompletedProcess[str]: + bin_dir = tmp_path / "bin" + bin_dir.mkdir(exist_ok=True) + curl = bin_dir / "curl" + curl.write_text(FAKE_CURL) + curl.chmod(0o755) + state = tmp_path / "state" + state.mkdir(exist_ok=True) + return subprocess.run( + ["bash", "-c", f'source "{SCRIPT}"\n{snippet}'], + capture_output=True, + text=True, + env={ + "PATH": f"{bin_dir}{os.pathsep}/usr/bin{os.pathsep}/bin", + "HOME": str(tmp_path), + "LENS_DEV_STATE_DIR": str(state), + "LENS_DEV_PYTHON": sys.executable, + "CURL_LOG": str(tmp_path / "curl.log"), + "CLAIM_STATUS": "409", + **env, + }, + ) + + +def _curl_calls(tmp_path: Path) -> str: + log = tmp_path / "curl.log" + return log.read_text() if log.exists() else "" + + +def test_missing_worker_token_registers_a_worker(tmp_path): + proc = _run(tmp_path, "ensure_worker_token") + assert proc.returncode == 0, proc.stderr + assert (tmp_path / "state" / "worker_token").read_text().strip() == "lens-fresh" + assert oct((tmp_path / "state" / "worker_token").stat().st_mode & 0o777) == "0o600" + assert f'"analysis_key_id": "{0:064d}"' in _curl_calls(tmp_path) + + +def test_accepted_worker_token_is_reused(tmp_path): + (tmp_path / "state").mkdir() + (tmp_path / "state" / "worker_token").write_text("lens-saved\n") + proc = _run(tmp_path, "ensure_worker_token", CLAIM_STATUS="409") + assert proc.returncode == 0, proc.stderr + assert "reusing worker token" in proc.stdout + assert (tmp_path / "state" / "worker_token").read_text().strip() == "lens-saved" + assert "/lens/workers/register" not in _curl_calls(tmp_path) + + +def test_rejected_worker_token_is_replaced(tmp_path): + (tmp_path / "state").mkdir() + (tmp_path / "state" / "worker_token").write_text("lens-revoked\n") + proc = _run(tmp_path, "ensure_worker_token", CLAIM_STATUS="401") + assert proc.returncode == 0, proc.stderr + assert "was rejected" in proc.stdout + assert (tmp_path / "state" / "worker_token").read_text().strip() == "lens-fresh" + + +def test_unexpected_token_check_status_fails(tmp_path): + (tmp_path / "state").mkdir() + (tmp_path / "state" / "worker_token").write_text("lens-saved\n") + proc = _run(tmp_path, "ensure_worker_token", CLAIM_STATUS="500") + assert proc.returncode == 1 + assert "unexpected HTTP 500" in proc.stderr + + +def test_default_master_key_is_random_and_stable(tmp_path): + first = _run(tmp_path, 'load_master_key; echo "$master_key"') + second = _run(tmp_path, 'load_master_key; echo "$master_key"') + assert first.returncode == 0, first.stderr + key = first.stdout.strip() + assert key.startswith("sk-") and len(key) == 51 and key != "sk-1234" + assert second.stdout.strip() == key + assert oct((tmp_path / "state" / "master_key").stat().st_mode & 0o777) == "0o600" + + +def test_master_key_override_wins(tmp_path): + proc = _run(tmp_path, 'load_master_key; echo "$master_key"', LENS_DEV_MASTER_KEY="sk-mine") + assert proc.stdout.strip() == "sk-mine" + assert not (tmp_path / "state" / "master_key").exists() + + +def test_proxy_env_drops_inherited_redis_and_base_urls(tmp_path): + proc = _run( + tmp_path, + 'master_key=sk-strong; proxy_env "export OPENAI_API_KEY=from-dotenv"; env', + REDIS_HOST="redis.example", + REDIS_PORT="6379", + REDIS_PASSWORD="secret", + ANTHROPIC_BASE_URL="http://elsewhere", + OPENAI_BASE_URL="http://elsewhere", + ) + assert proc.returncode == 0, proc.stderr + names = {line.split("=", 1)[0] for line in proc.stdout.splitlines()} + assert not {n for n in names if n.startswith("REDIS_")} + assert not names & {"ANTHROPIC_BASE_URL", "OPENAI_BASE_URL"} + assert "OPENAI_API_KEY=from-dotenv" in proc.stdout + assert "LITELLM_MODE=PRODUCTION" in proc.stdout + assert "UI_PASSWORD=sk-strong" in proc.stdout + assert "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY" not in names + + +def test_proxy_env_permits_the_weak_key_only_when_chosen(tmp_path): + proc = _run(tmp_path, 'master_key=sk-1234; proxy_env ""; env') + assert "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true" in proc.stdout + + +def test_external_database_url_never_starts_compose_postgres(tmp_path): + docker = tmp_path / "bin" / "docker" + proc = _run( + tmp_path, + f"listening() {{ return 1; }}\n" + f"printf '#!/bin/sh\\necho \"$@\" > {tmp_path}/docker.log\\n' > {docker}; chmod +x {docker}\n" + "ensure_services", + LENS_DEV_DATABASE_URL="postgresql://elsewhere/db?schema=public", + ) + assert proc.returncode == 0, proc.stderr + assert (tmp_path / "docker.log").read_text().split()[-2:] == ["--wait", "clickhouse"] + + +def test_cleanup_kills_child_process_trees(tmp_path): + proc = _run( + tmp_path, + "set -m\n" + "(sleep 300 & wait) & pids+=($!)\n" + "child=$!; sleep 0.3\n" + "cleanup\n" + "sleep 0.3\n" + 'pgrep -g "$child" >/dev/null && echo LEFTOVER || echo CLEAN', + ) + assert proc.returncode == 0, proc.stderr + assert "lens-dev: stopping" in proc.stdout + assert proc.stdout.strip().endswith("CLEAN")