scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57

Merge branch 'main' into litellm_tokenlab_provider

Browse source
This commit is contained in:
hedging8563 2026-10-03 05:48:08 +08:00 • committed by GitHub
commit 46b7b1a617
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
711 changed files with 63854 additions and 19272 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.circleci/scripts/run_integration.sh
View file

@ -169,11 +169,11 @@ start_proxy() {
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True LITELLM_ENABLE_MCP_STDIO=true "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
--host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
--use_prisma_db_push --enforce_prisma_migration_check \
--use_prisma_db_push \
> "$results/$log_name" 2>&1 &
launched_pid=$!
}

2
.circleci/scripts/unit_selection.sh
View file

@ -77,6 +77,7 @@ legacy_paths() {
echo tests/unit/embeddings
echo tests/unit/endpoints
echo tests/unit/files
echo tests/unit/harness
echo tests/unit/images
echo tests/unit/interactions
echo tests/unit/messages
@ -145,6 +146,7 @@ legacy_paths() {
echo tests/unit/proxy/test_proxy_token_counter.py
echo tests/unit/proxy/test_server_root_path.py ;;
proxy-db-proxy-server-core)
echo tests/unit/proxy/test__lazy_features.py
echo tests/unit/proxy/test_aproxy_startup.py
echo tests/unit/proxy/test_proxy_server.py ;;
proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;

8
.github/ci-coverage-allowlist.yml vendored
View file

@ -4,6 +4,14 @@ description: >-
by a job nor listed here, so every entry below is a decision on the record.
test_paths:
- reason: >-
litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
before changing litellm/harness; the mocked coverage runs in tests/unit/harness and
tests/unit/llms/*/harness
paths:
- tests/harness_e2e
- reason: >-
The Rust/Python parity harness is run manually through its local CLI. Recorded replay,
fixture generation, and harness checks are intentionally outside pull request CI

1
.github/workflows/test-e2e-changed.yml vendored
View file

@ -176,6 +176,7 @@ jobs:
TESTS: ${{ needs.detect.outputs.tests }}
E2E_FIXTURE_MODE: live
E2E_PROVIDER_EDGE_HOST_REACHABLE: '1'
E2E_OWNED_GATEWAY: '1'
COLUMNS: '400'
run: |
umask 077

25
README.md
View file

@ -268,6 +268,31 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
</details>
<details>
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
### Python SDK - Agents
```python
import litellm
from litellm import Harness, sandbox
result = litellm.agent(
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
"Find why tests/test_router.py is flaky and fix it.",
sandbox=sandbox.local("./repo"),
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
)
print(result.text, result.cost, [f.path for f in result.files])
```
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)
</details>
### Supported Providers ([Website Supported Models](https://models.litellm.ai/) | [Docs](https://docs.litellm.ai/docs/providers))
| Provider | `/chat/completions` | `/messages` | `/responses` | `/embeddings` | `/image/generations` | `/audio/transcriptions` | `/audio/speech` | `/moderations` | `/batches` | `/rerank` |

16
backend/main.py
View file

@ -8,9 +8,13 @@ Run with:
uvicorn backend.main:app --host 0.0.0.0 --port 4001
"""
from collections.abc import AsyncGenerator, Mapping
from contextlib import asynccontextmanager
from typing import Final
from fastapi.routing import Mount
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.types import Lifespan
# See gateway/main.py for why we assemble DATABASE_URL(s) here before
# importing proxy_server.
@ -43,14 +47,16 @@ def _is_backend_route(route) -> bool:
# See gateway/main.py for why the trim runs inside the lifespan instead of at
# module scope.
_proxy_lifespan = app.router.lifespan_context
_proxy_lifespan: Final = app.router.lifespan_context
@asynccontextmanager
async def _backend_lifespan(app_):
async with _proxy_lifespan(app_):
async def _backend_lifespan(
app_: Starlette, lifespan: Lifespan[Starlette] = _proxy_lifespan
) -> AsyncGenerator[Mapping[str, object], None]:
async with lifespan(app_) as state:
app_.router.routes = [r for r in app_.router.routes if _is_backend_route(r)]
yield
yield state if state is not None else {}
app.router.lifespan_context = _backend_lifespan

1
backend/routes/allowlist.py
View file

@ -60,6 +60,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
# Tools / agents (registry & policy admin)
"/v1/tool/",
"/v1/agents",
"/agent/daily/activity/",
# Guardrails admin
"/v2/guardrails/",
# MCP server admin + BYOK OAuth flow (UI-initiated) + dynamic per-server endpoints

15
deploy/lens/README.md
View file

@ -4,7 +4,20 @@ Lens reviews recorded activity and saves evidence-linked findings in the LiteLLM
## Start a worker
Upgrade your existing LiteLLM proxy to a release that includes Lens with PostgreSQL, agent tracing (`general_settings.tracing: {store: clickhouse}`), and ClickHouse configured through `CLICKHOUSE_URL` and a separate SELECT-only `CLICKHOUSE_READER_URL`. Enable the ClickHouse callback and request/response logging to analyze LLM requests. Lens can only inspect content you actually retain
Upgrade your existing LiteLLM proxy to a release that includes Lens with PostgreSQL and agent tracing. Configure one ClickHouse URL for trace writes, bounded reads, and Lens queries:
```yaml
general_settings:
tracing:
store:
type: clickhouse
url: os.environ/CLICKHOUSE_URL
retention_days: 14
```
The URL, database, and retention settings can also come from `CLICKHOUSE_URL`, `CLICKHOUSE_DATABASE`, and `AGENT_TRACING_RETENTION_DAYS` when omitted from YAML. A YAML value wins when both are set. The database defaults to `litellm`. `retention_days` defaults to 14 and applies to both traces and spend logs
Retention changes require a proxy restart. ClickHouse removes expired rows during background merges, not immediately at startup. Enable request/response logging to analyze LLM requests. Lens can only inspect content you actually retain
In Lens, click **Set up analysis**, choose an existing virtual key or **Create worker key**, then **Generate setup command**. The LiteLLM address is filled in for you; change it only if the server running Docker needs a different network address. Copy the command and run it on your server. The dialog changes to **Analyzer connected** when the container checks in

2
deploy/lens/compose.yaml
View file

@ -1,6 +1,6 @@
services:
lens-worker:
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:a8e8731d954916594eea462969946b9292fb771681ff515a9fd296b53f856c77}
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:67eba741c1b97c749975c5c38e2370a603e1105babc908d613c1b79d7b995393}
environment:
LITELLM_URL: ${LITELLM_URL:?Set the URL reachable from this container}
LENS_WORKER_TOKEN: ${LENS_WORKER_TOKEN:?Create a worker credential in the Lens UI}

2
docker-compose.hardened.yml
View file

@ -6,8 +6,6 @@ services:
context: .
dockerfile: docker/Dockerfile.non_root
target: runtime
args:
PROXY_EXTRAS_SOURCE: "local"
depends_on:
- squid
user: "101:101"

31
docker/Dockerfile.non_root
View file

@ -3,7 +3,6 @@
# Base images
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG PROXY_EXTRAS_SOURCE=published
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
@ -44,7 +43,6 @@ COPY ui/litellm-dashboard/ ./
RUN npm run build
FROM $LITELLM_BUILD_IMAGE AS builder
ARG PROXY_EXTRAS_SOURCE
WORKDIR /app
USER root
@ -107,26 +105,14 @@ RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
touch /var/lib/litellm/ui/.litellm_ui_ready
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13 \
--no-sources-package litellm-proxy-extras; \
else \
uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13; \
fi
uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
@ -136,7 +122,6 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
FROM $LITELLM_RUNTIME_IMAGE AS runtime
ARG PROXY_EXTRAS_SOURCE
WORKDIR /app
USER root

8
docker/docker-compose.quickstart.yml
View file

@ -13,11 +13,13 @@ services:
litellm:
image: docker.litellm.ai/berriai/litellm:main-stable
ports:
- "4000:4000"
# LITELLM_BIND is empty by default, so this stays "4000:4000". The quickstart
# script sets it to "127.0.0.1:" so new installs listen on this machine only.
- "${LITELLM_BIND:-}${LITELLM_PORT:-4000}:4000"
environment:
LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY:?set it in .env - see the header of this file}
LITELLM_SALT_KEY: ${LITELLM_SALT_KEY:?set it in .env - see the header of this file}
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
DATABASE_URL: postgresql://litellm:${POSTGRES_PASSWORD:-litellm}@db:5432/litellm
STORE_MODEL_IN_DB: "True"
depends_on:
db:
@ -27,7 +29,7 @@ services:
image: postgres:16
environment:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: litellm
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-litellm}
POSTGRES_DB: litellm
healthcheck:
test: ["CMD-SHELL", "pg_isready -U litellm"]

1
docker/docker-compose.tracing.yml
View file

@ -12,7 +12,6 @@ services:
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
STORE_MODEL_IN_DB: "True"
CLICKHOUSE_URL: http://default:local-tracing@clickhouse:8123
CLICKHOUSE_READER_URL: http://default:local-tracing@clickhouse:8123
CLICKHOUSE_DATABASE: litellm
OPENAI_API_KEY: ${OPENAI_API_KEY:-}
volumes:

5
docker/tracing-config.yaml
View file

@ -7,4 +7,7 @@ model_list:
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
tracing:
store: clickhouse
store:
type: clickhouse
url: os.environ/CLICKHOUSE_URL
retention_days: 14

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.72"
version = "0.1.73"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.72"
version = "0.1.73"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

16
gateway/main.py
View file

@ -9,9 +9,13 @@ Run with:
uvicorn gateway.main:app --host 0.0.0.0 --port 4000
"""
from collections.abc import AsyncGenerator, Mapping
from contextlib import asynccontextmanager
from typing import Final
from fastapi.routing import Mount
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.types import Lifespan
# Assemble DATABASE_URL (+ DATABASE_URL_READ_REPLICA) from the discrete
# DATABASE_* env vars before proxy_server imports spin up Prisma. Handles
@ -54,14 +58,16 @@ def _is_gateway_route(route) -> bool:
# register routes. A module-load filter would miss routes added during
# startup; running inside the lifespan, after the inner __aenter__, catches
# them while still completing before uvicorn opens the listener.
_proxy_lifespan = app.router.lifespan_context
_proxy_lifespan: Final = app.router.lifespan_context
@asynccontextmanager
async def _gateway_lifespan(app_):
async with _proxy_lifespan(app_):
async def _gateway_lifespan(
app_: Starlette, lifespan: Lifespan[Starlette] = _proxy_lifespan
) -> AsyncGenerator[Mapping[str, object], None]:
async with lifespan(app_) as state:
app_.router.routes = [r for r in app_.router.routes if _is_gateway_route(r)]
yield
yield state if state is not None else {}
app.router.lifespan_context = _gateway_lifespan

18
helm/litellm-helm/tests/migrations-job_tests.yaml
View file

@ -112,6 +112,24 @@ tests:
name: CUSTOM_VAR
value: "custom_value"
- it: should override a user-supplied DISABLE_SCHEMA_UPDATE so the Job always migrates
template: migrations-job.yaml
set:
envVars:
DISABLE_SCHEMA_UPDATE: "true"
migrationJob:
enabled: true
asserts:
# The Job is what owns the schema, so it renders its own
# DISABLE_SCHEMA_UPDATE=false after envVars and extraEnvVars. Kubernetes
# takes the last value for a duplicated name, so the user's "true" cannot
# leave the schema unmigrated. Skipping migrations is migrationJob.enabled.
- equal:
path: spec.template.spec.containers[0].env[-1]
value:
name: DISABLE_SCHEMA_UPDATE
value: "false"
- it: should not include DATABASE_URL when deployStandalone is false
template: migrations-job.yaml
set:

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

@ -545,7 +545,6 @@ redis:
# Prisma migration job settings
migrationJob:
enabled: true # Enable or disable the schema migration Job
retries: 3 # Number of retries for the Job in case of failure
backoffLimit: 4 # Backoff limit for Job restarts
# Wall-clock budget for the whole Job, shared across every `backoffLimit`
# retry rather than granted per attempt. Without it a migration that blocks
@ -554,7 +553,6 @@ migrationJob:
# stop reconciling the whole chart until someone deletes the Job by hand.
# Set to null to opt out and restore the unbounded behaviour.
activeDeadlineSeconds: 1800
disableSchemaUpdate: false # Skip schema migrations for specific environments. When True, the job will exit with code 0.
# Optional service account for the migration job.
# Only used when migrationJob.hooks.helm.enabled=true and serviceAccount.create=true.
# In that case, pre-install/pre-upgrade hooks run before normal resources, so this defaults to "default".

17
litellm-proxy-extras/litellm_proxy_extras/migrations/20261001200000_add_autorouter_daily_spend/migration.sql Normal file
View file

@ -0,0 +1,17 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterDailySpend" (
"date" TEXT NOT NULL,
"api_key" TEXT NOT NULL,
"user_id" TEXT NOT NULL,
"router_name" TEXT NOT NULL,
"router_type" TEXT NOT NULL,
"turns" INTEGER NOT NULL DEFAULT 0,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"savings_estimated_turns" INTEGER NOT NULL DEFAULT 0,
"savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
"classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0,
CONSTRAINT "LiteLLM_AutoRouterDailySpend_pkey" PRIMARY KEY ("date", "api_key", "user_id", "router_name", "router_type")
);

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20261002220000_lens_worker_scope_index/migration.sql Normal file
View file

@ -0,0 +1,3 @@
CREATE INDEX IF NOT EXISTS "LiteLLM_LensWorker_active_scope_idx"
ON "LiteLLM_LensWorker" USING GIN ((data->'scope') jsonb_path_ops)
WHERE data @> '{"revoked": false}'::jsonb;

57
litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py
View file

@ -56,8 +56,10 @@ REQUEST_LOG_INDEXES: Final = (
)
_IDENTIFIER_MAX_BYTES: Final = 63
_PARENT_LOCK_TIMEOUT: Final = "2s"
_PARENT_LOCK_ATTEMPTS: Final = 30
_DDL_LOCK_TIMEOUT: Final = "200ms"
_DDL_LOCK_ATTEMPTS: Final = 10
_DDL_RETRY_BASE_SECONDS: Final = 0.25
_DDL_RETRY_MAX_SECONDS: Final = 8.0
_LOCK_HANDOVER_SECONDS: Final = 2.0
_DIGEST_LENGTH: Final = 8
_CREATE_INDEX_STATEMENT: Final = re.compile(
@ -181,6 +183,32 @@ def _under_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]",
return step()
def _with_bounded_lock(
connection: "psycopg.Connection[tuple[object, ...]]", step: Callable[[], bool], what: str
) -> bool:
"""Run `step` under the migration lock with a short lock_timeout, so a DDL statement that has to wait for open
transactions holds new writes back for at most that long; retry with capped exponential backoff, holding the
migration lock per attempt only and releasing it while sleeping. False when another process holds the migration
lock or every attempt timed out."""
import psycopg
from psycopg import sql
for attempt in range(_DDL_LOCK_ATTEMPTS):
if attempt:
time.sleep(min(_DDL_RETRY_MAX_SECONDS, _DDL_RETRY_BASE_SECONDS * 2.0**attempt) * random.uniform(0.5, 1.0))
connection.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(_DDL_LOCK_TIMEOUT)))
try:
return _under_migration_lock(connection, step)
except psycopg.errors.LockNotAvailable:
logger.info("Waiting for open transactions before %s", what)
finally:
connection.execute("SET lock_timeout = 0")
logger.warning(
"Could not get the lock for %s without holding writes back, leaving it for the next index build", what
)
return False
def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex) -> bool:
from psycopg.rows import class_row
@ -368,12 +396,13 @@ def build_index_on_partitioned_table(
"Index %s already exists on %s rather than %s, leaving it alone", parent_index, existing.table, parent_table
)
return False
if existing is None and not _under_migration_lock(
if existing is None and not _with_bounded_lock(
connection,
lambda: (
_adopt_equivalent_index(connection, schema, parent_table, parent_index, index)
or _create_parent_index(connection, schema, parent_index, parent_table, index)
),
f"creating the parent index {parent_index}",
):
return False
children: Final = _children_without_the_index(connection, schema, parent_table, parent_index)
@ -393,29 +422,15 @@ def _create_parent_index(
table: str,
index: RequestLogIndex,
) -> bool:
"""Create the metadata-only parent index. Postgres takes a SHARE lock on the
parent for that statement, so it waits for in-flight writes and queues new ones
behind it; a short lock_timeout with retries keeps every such pause bounded."""
import psycopg
"""Create the metadata-only parent index. The caller bounds Postgres's SHARE lock wait on the parent."""
from psycopg import sql
prefix: Final = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON ONLY {} ").format(
sql.Identifier(name), sql.Identifier(schema, table)
)
statement: Final = _create_index_statement(connection, prefix, index.definition)
connection.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(_PARENT_LOCK_TIMEOUT)))
try:
for _ in range(_PARENT_LOCK_ATTEMPTS):
try:
connection.execute(statement)
return True
except psycopg.errors.LockNotAvailable:
logger.info("Waiting for in-flight writes to %s before creating the parent index %s", table, name)
time.sleep(random.uniform(0.1, 0.5))
finally:
connection.execute("SET lock_timeout = 0")
logger.warning("Could not get the parent lock on %s to create %s, leaving it for the next index build", table, name)
return False
connection.execute(statement)
return True
def _attach_child_index(
@ -445,4 +460,4 @@ def _attach_child_index(
logger.info("Attached index %s on partition %s to %s", child_index, child.name, parent_index)
return True
return _under_migration_lock(connection, attach)
return _with_bounded_lock(connection, attach, f"attaching {child_index}")

21
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1744,6 +1744,27 @@ model LiteLLM_AutoRouterUserSession {
@@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn")
}
// Auto-routed requests per UTC request day and router: the selected-day money behind the
// auto-router usage view. Written in the same statement as the session rollup, so a day row
// and its session row never disagree; corrected in the same transaction as late baselines.
model LiteLLM_AutoRouterDailySpend {
date String
api_key String
user_id String
router_name String
router_type String
turns Int @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
@@id([date, api_key, user_id, router_name, router_type])
}
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
// either direction. forward duplicates the requests the keys did not route through the
// router through it, answering whether they should adopt it; reverse duplicates the

70
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -78,6 +78,23 @@ class _InvalidIndex:
table_size: str
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
LIBPQ_URL_PARAMS: Final = frozenset(
{
"sslmode",
"sslcert",
"sslkey",
"sslrootcert",
"sslpassword",
"application_name",
"connect_timeout",
"client_encoding",
"options",
"service",
"gssencmode",
"krbsrvname",
"target_session_attrs",
}
)
@dataclass(frozen=True)
@ -689,30 +706,43 @@ class ProxyExtrasDBManager:
@staticmethod
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
schema, etc.) from DATABASE_URL so psycopg can parse it."""
"""Rewrite a Prisma-dialect URL for libpq: drop the Prisma-only params
(connection_limit, pool_timeout, schema, pgbouncer, sslaccept, ...) and
translate Prisma's TLS params back, since libpq reads ``sslcert`` as a
client certificate where Prisma reads it as the CA."""
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
parsed = urlparse(url)
parsed: Final = urlparse(url)
if not parsed.query:
return url
libpq_params = {
"sslmode",
"sslcert",
"sslkey",
"sslrootcert",
"sslpassword",
"application_name",
"connect_timeout",
"client_encoding",
"options",
"service",
"gssencmode",
"krbsrvname",
"target_session_attrs",
}
kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
pairs: Final = tuple(parse_qsl(parsed.query))
kept: Final = tuple((k, v) for k, v in pairs if k in LIBPQ_URL_PARAMS)
sslaccept: Final = next((v for k, v in pairs if k == "sslaccept"), None)
libpq_pairs: Final = ProxyExtrasDBManager._libpq_tls_params(kept, sslaccept)
return urlunparse(parsed._replace(query=urlencode(libpq_pairs, quote_via=quote)))
@staticmethod
def _libpq_tls_params(
pairs: "tuple[tuple[str, str], ...]", sslaccept: "str | None"
) -> "tuple[tuple[str, str], ...]":
"""Undo ``translate_libpq_ssl_params``. Prisma's ``sslcert`` is the CA and
``sslaccept=strict`` checks chain and hostname, which libpq only does in
``sslmode=verify-full``, so strict becomes ``sslrootcert`` plus
``verify-full`` whatever ``sslmode`` said (``disable`` stays off). Prisma
defaults an absent ``sslaccept`` to ``accept_invalid_certs`` and anything
else to strict. Without strict it checks nothing, so the CA is dropped and
``sslmode`` is kept as is: libpq only verifies when a root cert is present.
A URL that also carries ``sslkey`` is libpq's own client-certificate form
and is kept."""
keys: Final = frozenset(k for k, _ in pairs)
if "sslcert" not in keys or "sslkey" in keys:
return pairs
sslmode: Final = next((v for k, v in pairs if k == "sslmode"), None)
rest: Final = tuple((k, v) for k, v in pairs if k not in ("sslcert", "sslmode"))
if sslaccept in (None, "accept_invalid_certs") or sslmode == "disable":
return rest if sslmode is None else rest + (("sslmode", sslmode),)
root_cert: Final = tuple(("sslrootcert", v) for k, v in pairs if k == "sslcert" and "sslrootcert" not in keys)
return rest + root_cert + (("sslmode", "verify-full"),)
@staticmethod
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.103"
version = "0.4.105"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.103"
version = "0.4.105"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

74
litellm-rust/Cargo.lock generated
View file

@ -97,6 +97,53 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
[[package]]
name = "askama"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6024d73179f43f15ccd2b881bfea6fee7f3a46ec53f33b52210dea749ebebaa4"
dependencies = [
"askama_macros",
"itoa",
"percent-encoding",
"serde",
"serde_json",
]
[[package]]
name = "askama_derive"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "071ee5ebf2138e3ad180e0aacf6940c2cab5e6d8333741d9925c7bee2b153f39"
dependencies = [
"askama_parser",
"memchr",
"proc-macro2",
"quote",
"rustc-hash",
"syn 3.0.6",
]
[[package]]
name = "askama_macros"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "643e1c7cbb6aec1d920332fe51a7c0d8219e273dcb8602db03f5263e4d16487b"
dependencies = [
"askama_derive",
]
[[package]]
name = "askama_parser"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2c5ae75772275d268b03ab8bdccdd12117b6169ee23256942b34e46c9f476583"
dependencies = [
"rustc-hash",
"unicode-ident",
"winnow 1.0.4",
]
[[package]]
name = "asn1-rs"
version = "0.7.2"
@ -4038,6 +4085,26 @@ dependencies = [
"strum",
]
[[package]]
name = "litellm-migrate"
version = "0.1.0"
dependencies = [
"litellm-migrate-macros",
"rstest",
]
[[package]]
name = "litellm-migrate-macros"
version = "0.1.0"
dependencies = [
"proc-macro2",
"quote",
"rstest",
"syn 2.0.119",
"tempfile",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-model-catalog"
version = "0.1.0"
@ -4384,11 +4451,17 @@ dependencies = [
name = "litellm-traces"
version = "0.1.0"
dependencies = [
"askama",
"base64 0.22.1",
"criterion",
"flate2",
"futures-util",
"hmac 0.12.1",
"indexmap 2.14.0",
"litellm-http",
"litellm-migrate",
"litellm-storage-clickhouse",
"moka",
"opentelemetry-proto",
"prost",
"rstest",
@ -4400,6 +4473,7 @@ dependencies = [
"thiserror 2.0.19",
"time",
"tokio",
"url",
"wiremock",
]

6
litellm-rust/Cargo.toml
View file

@ -14,6 +14,8 @@ litellm-router = { path = "crates/router" }
litellm-tracing = { path = "crates/tracing" }
litellm-traces = { path = "crates/traces" }
litellm-storage-clickhouse = { path = "crates/storage-clickhouse" }
litellm-migrate = { path = "crates/migrate" }
litellm-migrate-macros = { path = "crates/migrate-macros" }
litellm-core = { path = "crates/core" }
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
litellm-gateway = { path = "crates/gateway" }
@ -63,6 +65,7 @@ litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
litellm-host-python = { path = "crates/host-python" }
litellm-python-compat = { path = "crates/python-compat" }
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
tracing = "0.1"
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }
axum-login = "0.18.0"
@ -93,7 +96,10 @@ serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
serde_with = { version = "=3.16.1", default-features = false, features = ["std", "macros"] }
sha2 = "0.10"
syn = { version = "2", default-features = false }
sqlx = { version = "0.9.0", default-features = false, features = ["json", "macros", "postgres", "runtime-tokio", "chrono", "tls-rustls-ring-native-roots"] }
proc-macro2 = "1"
quote = "1"
subtle = "2"
thiserror = "2.0"
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }

5
litellm-rust/crates/config/src/lib.rs
View file

@ -12,7 +12,10 @@ use serde::Deserialize;
pub use error::Error;
pub use mcp::{McpAuth, McpServer, McpTransport};
pub use model::{LiteLlmParams, Model};
pub use settings::{GeneralSettings, LiteLlmSettings, RouterSettings};
pub use settings::{
ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings,
TracingStoreSettings,
};
pub use value::{AdditionalFields, Flag, NumberOrString, Object, OneOrMany, Value};
#[derive(Clone, Default, Deserialize)]

44
litellm-rust/crates/config/src/settings.rs
View file

@ -5,6 +5,47 @@ use serde::Deserialize;
use crate::{AdditionalFields, Flag, NumberOrString, Object, OneOrMany, Value};
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum TracingStoreKind {
Clickhouse,
}
#[derive(Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ClickHouseStoreSettings {
#[serde(rename = "type")]
pub kind: TracingStoreKind,
pub url: Option<SecretValue>,
pub database: Option<String>,
pub retention_days: Option<NumberOrString>,
}
impl fmt::Debug for ClickHouseStoreSettings {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ClickHouseStoreSettings")
.field("kind", &self.kind)
.field("database", &self.database)
.field("retention_days", &self.retention_days)
.finish()
}
}
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
pub enum TracingStoreSettings {
ClickHouse(ClickHouseStoreSettings),
}
#[derive(Clone, Default, Debug, Deserialize)]
#[serde(default)]
pub struct TracingSettings {
pub store: Option<TracingStoreSettings>,
#[serde(flatten)]
pub additional_fields: AdditionalFields,
}
#[derive(Clone, Deserialize)]
#[serde(default)]
pub struct GeneralSettings {
@ -14,6 +55,7 @@ pub struct GeneralSettings {
pub admission_queue_timeout_seconds: f64,
pub master_key: Option<SecretValue>,
pub database_url: Option<SecretValue>,
pub tracing: Option<TracingSettings>,
pub database_connection_pool_limit: Option<u64>,
pub database_connection_timeout: Option<f64>,
pub database_connect_timeout: Option<f64>,
@ -50,6 +92,7 @@ impl Default for GeneralSettings {
admission_queue_timeout_seconds: 1.0,
master_key: None,
database_url: None,
tracing: None,
database_connection_pool_limit: Some(10),
database_connection_timeout: Some(60.0),
database_connect_timeout: None,
@ -97,6 +140,7 @@ impl fmt::Debug for GeneralSettings {
)
.field("master_key", &self.master_key)
.field("database_url", &self.database_url)
.field("tracing", &self.tracing)
.field("store_model_in_db", &self.store_model_in_db)
.field("additional_fields", &self.additional_fields.keys())
.finish_non_exhaustive()

50
litellm-rust/crates/config/tests/config.rs
View file

@ -1,4 +1,4 @@
use litellm_config::{Config, Error, Flag, NumberOrString};
use litellm_config::{Config, Error, Flag, NumberOrString, TracingStoreSettings};
use rstest::{fixture, rstest};
use tempfile::TempDir;
@ -113,6 +113,54 @@ fn missing_general_settings_has_no_master_key() {
assert!(config.general_settings.master_key.is_none());
}
#[test]
fn tracing_settings_are_typed_and_redact_the_url() {
let config = Config::from_yaml(
"general_settings:\n tracing:\n store:\n type: clickhouse\n url: https://writer:password@example.com\n database: analytics\n retention_days: 7\n",
)
.unwrap();
let tracing = config.general_settings.tracing.as_ref().unwrap();
let Some(TracingStoreSettings::ClickHouse(store)) = tracing.store.as_ref() else {
panic!("expected ClickHouse tracing store")
};
assert_eq!(
store.url.as_ref().unwrap().expose(),
"https://writer:password@example.com"
);
assert_eq!(store.database.as_deref(), Some("analytics"));
assert_eq!(store.retention_days, Some(NumberOrString::Number(7.0)));
assert!(!format!("{config:?}").contains("password"));
}
#[test]
fn tracing_settings_accept_environment_references() {
let config = Config::from_yaml(
"general_settings:\n tracing:\n store:\n type: clickhouse\n url: os.environ/CLICKHOUSE_URL\n retention_days: os.environ/RETENTION_DAYS\n",
)
.unwrap();
let Some(TracingStoreSettings::ClickHouse(store)) =
config.general_settings.tracing.unwrap().store
else {
panic!("expected ClickHouse tracing store")
};
assert_eq!(
store.retention_days,
Some(NumberOrString::String(
"os.environ/RETENTION_DAYS".to_owned()
))
);
}
#[test]
fn tracing_settings_reject_string_store() {
assert!(Config::from_yaml("general_settings:\n tracing:\n store: clickhouse\n").is_err());
}
#[test]
fn tracing_settings_reject_removed_reader_configuration() {
assert!(Config::from_yaml("general_settings:\n tracing:\n store:\n type: clickhouse\n reader_url: http://localhost:8123\n").is_err());
}
#[rstest]
fn empty_config_matches_python_defaults() {
let config = Config::from_yaml("{}").unwrap();

19
litellm-rust/crates/migrate-macros/Cargo.toml Normal file
View file

@ -0,0 +1,19 @@
[package]
name = "litellm-migrate-macros"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[lib]
proc-macro = true
[dependencies]
proc-macro2.workspace = true
quote.workspace = true
syn = { workspace = true, features = ["parsing", "printing", "proc-macro"] }
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
tempfile.workspace = true

21
litellm-rust/crates/migrate-macros/src/error.rs Normal file
View file

@ -0,0 +1,21 @@
use std::io;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("could not read migrations directory `{path}`")]
ReadDirectory {
path: String,
#[source]
source: io::Error,
},
#[error(
"migration name `{name}` must be `<digits>_<description>.sql` with a `[a-z0-9_]` description"
)]
InvalidName { name: String },
#[error("migration version `{version}` is declared more than once")]
DuplicateVersion { version: u64 },
#[error("migrations directory `{path}` contains no migrations")]
Empty { path: String },
#[error("migration path `{path}` is not valid UTF-8")]
NonUtf8Path { path: String },
}

199
litellm-rust/crates/migrate-macros/src/lib.rs Normal file
View file

@ -0,0 +1,199 @@
mod error;
use std::path::{Path, PathBuf};
use error::Error;
use proc_macro::TokenStream;
use quote::quote;
use syn::LitStr;
struct Entry {
version: u64,
description: String,
path: PathBuf,
}
fn resolve(dir: &Path) -> Result<Vec<Entry>, Error> {
let mut entries = Vec::new();
let files = std::fs::read_dir(dir).map_err(|source| Error::ReadDirectory {
path: dir.display().to_string(),
source,
})?;
for file in files {
let file = file.map_err(|source| Error::ReadDirectory {
path: dir.display().to_string(),
source,
})?;
let path = file.path();
let name = path
.file_name()
.and_then(|name| name.to_str())
.ok_or_else(|| Error::NonUtf8Path {
path: path.display().to_string(),
})?
.to_owned();
let invalid = || Error::InvalidName { name: name.clone() };
let stem = name
.strip_suffix(".sql")
.filter(|_| file.file_type().is_ok_and(|kind| kind.is_file()))
.and_then(|stem| stem.split_once('_'))
.filter(|(version, description)| {
!version.is_empty()
&& version.bytes().all(|b| b.is_ascii_digit())
&& !description.is_empty()
&& description
.bytes()
.all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'_')
})
.ok_or_else(invalid)?;
let version = stem.0.parse::<u64>().map_err(|_| invalid())?;
entries.push(Entry {
version,
description: stem.1.to_owned(),
path,
});
}
if entries.is_empty() {
return Err(Error::Empty {
path: dir.display().to_string(),
});
}
entries.sort_by_key(|entry| entry.version);
for pair in entries.windows(2) {
if pair[0].version == pair[1].version {
return Err(Error::DuplicateVersion {
version: pair[0].version,
});
}
}
Ok(entries)
}
fn resolve_input(lit: &LitStr) -> Result<Vec<Entry>, Error> {
let root = std::env::var("CARGO_MANIFEST_DIR")
.map(PathBuf::from)
.unwrap_or_default();
let dir = root.join(lit.value());
let dir = dir.canonicalize().map_err(|source| Error::ReadDirectory {
path: dir.display().to_string(),
source,
})?;
if dir.to_str().is_none() {
return Err(Error::NonUtf8Path {
path: dir.display().to_string(),
});
}
resolve(&dir)
}
#[proc_macro]
pub fn migrate(input: TokenStream) -> TokenStream {
let lit = syn::parse_macro_input!(input as LitStr);
match resolve_input(&lit) {
Ok(entries) => {
let migrations = entries.iter().map(|entry| {
let version = entry.version;
let description = &entry.description;
let path = entry
.path
.to_str()
.expect("canonical migration path is UTF-8");
quote! {
::litellm_migrate::Migration {
version: #version,
description: #description,
sql: ::core::include_str!(#path),
}
}
});
quote! { &[#(#migrations),*] }.into()
}
Err(err) => syn::Error::new(lit.span(), err).to_compile_error().into(),
}
}
#[cfg(test)]
mod tests {
use std::fs;
use rstest::rstest;
use tempfile::TempDir;
use super::{Error, resolve};
fn migrations_dir(files: &[&str]) -> TempDir {
let dir = TempDir::new().expect("tempdir");
for file in files {
fs::write(dir.path().join(file), "SELECT 1").expect("write fixture");
}
dir
}
#[rstest]
fn orders_versions_numerically() {
let dir = migrations_dir(&["10_tenth.sql", "2_second.sql", "1_first.sql"]);
let entries = resolve(dir.path()).expect("resolves");
let versions: Vec<u64> = entries.iter().map(|entry| entry.version).collect();
let descriptions: Vec<&str> = entries
.iter()
.map(|entry| entry.description.as_str())
.collect();
assert_eq!(versions, [1, 2, 10]);
assert_eq!(descriptions, ["first", "second", "tenth"]);
}
#[rstest]
#[case::dash_in_version(&["0001-dash.sql"])]
#[case::not_sql(&["notes.txt"])]
#[case::empty_description(&["0001_.sql"])]
#[case::non_digit_version(&["x_name.sql"])]
#[case::uppercase_description(&["0001_Upper.sql"])]
#[case::no_underscore(&["0001.sql"])]
#[case::plus_sign_version(&["+10_add.sql"])]
fn rejects_invalid_names(#[case] files: &[&str]) {
let dir = migrations_dir(files);
assert!(matches!(
resolve(dir.path()),
Err(Error::InvalidName { .. })
));
}
#[rstest]
fn rejects_subdirectories() {
let dir = migrations_dir(&["0001_a.sql"]);
fs::create_dir(dir.path().join("0002_b.sql")).expect("subdir");
assert!(matches!(
resolve(dir.path()),
Err(Error::InvalidName { .. })
));
}
#[cfg(unix)]
#[rstest]
fn rejects_symlinks() {
let dir = migrations_dir(&["0001_a.sql"]);
let target = TempDir::new().expect("tempdir");
let target_file = target.path().join("real.sql");
fs::write(&target_file, "SELECT 2").expect("write fixture");
std::os::unix::fs::symlink(&target_file, dir.path().join("0002_b.sql")).expect("symlink");
assert!(matches!(
resolve(dir.path()),
Err(Error::InvalidName { .. })
));
}
#[rstest]
fn rejects_duplicate_versions() {
let dir = migrations_dir(&["0001_a.sql", "1_b.sql"]);
assert!(matches!(
resolve(dir.path()),
Err(Error::DuplicateVersion { version: 1 })
));
}
#[rstest]
fn rejects_empty_directory() {
let dir = migrations_dir(&[]);
assert!(matches!(resolve(dir.path()), Err(Error::Empty { .. })));
}
}

12
litellm-rust/crates/migrate/Cargo.toml Normal file
View file

@ -0,0 +1,12 @@
[package]
name = "litellm-migrate"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-migrate-macros.workspace = true
[dev-dependencies]
rstest.workspace = true

5
litellm-rust/crates/migrate/README.md Normal file
View file

@ -0,0 +1,5 @@
# Migrations
`litellm-migrate` exports the `Migration` struct and the `migrate!` macro that embeds a directory of `<digits>_<description>.sql` files at compile time, sorted by numeric version
The crate does not apply or track migrations; callers decide how and when the embedded SQL runs

8
litellm-rust/crates/migrate/src/lib.rs Normal file
View file

@ -0,0 +1,8 @@
pub use litellm_migrate_macros::migrate;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Migration {
pub version: u64,
pub description: &'static str,
pub sql: &'static str,
}

1
litellm-rust/crates/migrate/tests/fixtures/migrations/10_tenth.sql vendored Normal file
View file

@ -0,0 +1 @@
SELECT 10;

1
litellm-rust/crates/migrate/tests/fixtures/migrations/1_first.sql vendored Normal file
View file

@ -0,0 +1 @@
SELECT 1;

1
litellm-rust/crates/migrate/tests/fixtures/migrations/2_second.sql vendored Normal file
View file

@ -0,0 +1 @@
SELECT 2;

21
litellm-rust/crates/migrate/tests/migrate.rs Normal file
View file

@ -0,0 +1,21 @@
use litellm_migrate::Migration;
use rstest::rstest;
const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("tests/fixtures/migrations");
#[rstest]
#[case::first(0, 1, "first", include_str!("fixtures/migrations/1_first.sql"))]
#[case::second(1, 2, "second", include_str!("fixtures/migrations/2_second.sql"))]
#[case::tenth(2, 10, "tenth", include_str!("fixtures/migrations/10_tenth.sql"))]
fn embeds_every_file_sorted_by_numeric_version(
#[case] index: usize,
#[case] version: u64,
#[case] description: &str,
#[case] sql: &str,
) {
assert_eq!(MIGRATIONS.len(), 3);
let migration = &MIGRATIONS[index];
assert_eq!(migration.version, version);
assert_eq!(migration.description, description);
assert_eq!(migration.sql, sql);
}

7
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -44,7 +44,10 @@ mod _native {
#[pymodule_export]
use crate::routes::token_counter::TokenCounter;
#[pymodule_export]
use crate::routes::traces::{NativeTraceStorage, trace_decode_otlp, trace_encode_error};
use crate::routes::traces::{
NativeTraceConfig, NativeTraceStorage, trace_decode_otlp, trace_encode_error,
trace_normalized_field_definitions,
};
#[cfg(feature = "huggingface")]
#[pymodule_export]
use crate::tokenizer::HuggingFaceEncoding;
@ -109,9 +112,11 @@ mod tests {
"aresponses",
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"NativeTraceConfig",
"NativeTraceStorage",
"trace_decode_otlp",
"trace_encode_error",
"trace_normalized_field_definitions",
"TokenCounter",
"Tokenizer",
"gil_stats",

141
litellm-rust/crates/python-bridge/src/routes/traces.rs
View file

@ -2,8 +2,10 @@ use std::collections::BTreeMap;
use litellm_host_python::{FromPythonCache, ToPythonCache};
use litellm_http::ClientVariant;
use litellm_storage_clickhouse::Storage;
use litellm_traces::{Error, InsertTable, Parameter, ReadQuery, Shared};
use litellm_traces::{
Config, Error, InsertTable, Parameter, QueryAccessError, QueryReaders, QueryScope, ReadQuery,
Shared,
};
use prost::Message;
use pyo3::{
exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
@ -46,42 +48,64 @@ fn map_error(error: Error) -> PyErr {
}
}
fn map_sql_error(error: Error) -> PyErr {
match error {
Error::QueryFailed(400 | 404) => PyValueError::new_err(error.to_string()),
error => map_error(error),
}
}
fn map_query_access_error(error: QueryAccessError) -> PyErr {
match error {
QueryAccessError::Storage(error) => map_sql_error(error),
QueryAccessError::InvalidScope => PyValueError::new_err(error.to_string()),
error => PyRuntimeError::new_err(error.to_string()),
}
}
#[pyclass(frozen)]
pub struct NativeTraceConfig {
inner: Config,
}
#[pymethods]
impl NativeTraceConfig {
#[new]
fn new(database: String, url: &str, retention_days: u32) -> PyResult<Self> {
Ok(Self {
inner: Config::new(database, url, retention_days).map_err(map_error)?,
})
}
}
#[pyclass]
pub struct NativeTraceStorage {
storage: Storage,
config: Config,
query_readers: QueryReaders,
}
#[pymethods]
impl NativeTraceStorage {
#[new]
#[pyo3(signature = (database, url, reader_url = None))]
fn new(database: String, url: &str, reader_url: Option<&str>) -> PyResult<Self> {
litellm_traces::schema_statements(&database, 1, 1).map_err(map_error)?;
fn new(config: PyRef<'_, NativeTraceConfig>) -> PyResult<Self> {
Ok(Self {
storage: Storage::new(database, url, reader_url).map_err(map_error)?,
query_readers: QueryReaders::new(
config.inner.storage().writer().clone(),
config.inner.storage().database().to_owned(),
),
config: config.inner.clone(),
})
}
fn ensure_schema<'py>(
&self,
py: Python<'py>,
trace_retention_days: u32,
spend_log_retention_days: u32,
) -> PyResult<Bound<'py, PyAny>> {
fn ensure_schema<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
let connection = self.storage.writer().clone();
let database = self.storage.database().to_owned();
let connection = self.config.storage().writer().clone();
let database = self.config.storage().database().to_owned();
let retention_days = self.config.retention_days();
crate::execution::run_async(
py,
async move {
litellm_traces::ensure_schema(
&client,
&connection,
&database,
trace_retention_days,
spend_log_retention_days,
)
.await
litellm_traces::ensure_schema(&client, &connection, &database, retention_days).await
},
map_error,
)
@ -95,8 +119,8 @@ impl NativeTraceStorage {
) -> PyResult<Bound<'py, PyAny>> {
let table = InsertTable::parse(table).map_err(map_error)?;
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
let connection = self.storage.writer().clone();
let database = self.storage.database().to_owned();
let connection = self.config.storage().writer().clone();
let database = self.config.storage().database().to_owned();
crate::execution::run_async(
py,
async move {
@ -107,6 +131,52 @@ impl NativeTraceStorage {
)
}
fn query_sql<'py>(
&self,
py: Python<'py>,
sql: String,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
secret: String,
) -> PyResult<Bound<'py, PyAny>> {
if sql.trim().is_empty() {
return Err(map_error(Error::EmptySql));
}
let readers = self.query_readers.clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
async move {
let _permit = readers.acquire()?;
let connection = readers.connection(&client, &scope, &secret).await?;
litellm_traces::query_sql(&client, &connection, &sql)
.await
.map_err(QueryAccessError::Storage)
},
map_query_access_error,
)
}
fn query_help<'py>(
&self,
py: Python<'py>,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
secret: String,
) -> PyResult<Bound<'py, PyAny>> {
let readers = self.query_readers.clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
async move {
let _permit = readers.acquire()?;
let connection = readers.connection(&client, &scope, &secret).await?;
litellm_traces::query_help(&client, &connection)
.await
.map_err(QueryAccessError::Storage)
},
map_query_access_error,
)
}
fn lens_query<'py>(
&self,
py: Python<'py>,
@ -117,9 +187,7 @@ impl NativeTraceStorage {
>,
) -> PyResult<Bound<'py, PyAny>> {
let query = litellm_traces::LensQuery::parse(name).map_err(map_error)?;
let connection = self.storage.reader().cloned().ok_or_else(|| {
PyRuntimeError::new_err("Trace reads require a separate ClickHouse reader URL")
})?;
let connection = self.config.storage().reader().clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
@ -140,9 +208,7 @@ impl NativeTraceStorage {
>,
) -> PyResult<Bound<'py, PyAny>> {
let query = ReadQuery::parse(query).map_err(map_error)?;
let connection = self.storage.reader().cloned().ok_or_else(|| {
PyRuntimeError::new_err("Trace reads require a separate ClickHouse reader URL")
})?;
let connection = self.config.storage().reader().clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
@ -236,6 +302,14 @@ fn spans_to_py<'py>(
"events",
litellm_host_python::Pythonized(&span.events).into_pyobject(py)?,
)?;
row.set_item(
"normalized",
litellm_host_python::Pythonized(&span.normalized).into_pyobject(py)?,
)?;
row.set_item(
"consumed_attributes",
litellm_host_python::Pythonized(&span.consumed_attributes).into_pyobject(py)?,
)?;
result.append(row)?;
}
Ok(result)
@ -288,3 +362,8 @@ mod tests {
});
}
}
#[pyfunction]
pub fn trace_normalized_field_definitions<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
litellm_host_python::Pythonized(litellm_traces::NORMALIZED_FIELD_DEFINITIONS).into_pyobject(py)
}

2
litellm-rust/crates/storage-clickhouse/README.md
View file

@ -1,5 +1,5 @@
# ClickHouse storage
`litellm-storage-clickhouse` exports `Storage`, a shared writer connection and optional reader connection for one ClickHouse database. It also exports bounded HTTP read and insert execution
`litellm-storage-clickhouse` exports `Storage`, a writer and bounded reader derived from one ClickHouse URL and database. It also exports bounded HTTP read and insert execution
The crate has no trace tables, OTLP types, or named trace queries. `litellm-traces` supplies those rules and uses this storage for both trace rows and spend rows

12
litellm-rust/crates/storage-clickhouse/src/lib.rs
View file

@ -89,19 +89,17 @@ impl Connection {
pub struct Storage {
database: String,
writer: Connection,
reader: Option<Connection>,
reader: Connection,
}
impl Storage {
pub fn new(database: String, url: &str, reader_url: Option<&str>) -> Result<Self, Error> {
pub fn new(database: String, url: &str) -> Result<Self, Error> {
if !valid_identifier(&database) {
return Err(Error::InvalidSchema);
}
Ok(Self {
writer: Connection::writer(url)?,
reader: reader_url
.map(|value| Connection::reader(value, &database))
.transpose()?,
reader: Connection::reader(url, &database)?,
database,
})
}
@ -114,8 +112,8 @@ impl Storage {
&self.writer
}
pub fn reader(&self) -> Option<&Connection> {
self.reader.as_ref()
pub fn reader(&self) -> &Connection {
&self.reader
}
}

27
litellm-rust/crates/storage-clickhouse/tests/connection.rs
View file

@ -10,25 +10,30 @@ fn accepts_only_clickhouse_http_urls(#[case] value: &str, #[case] expected: bool
assert_eq!(Connection::parse(value).is_ok(), expected);
}
#[rstest]
#[case::writer_only(None, false)]
#[case::separate_reader(Some("http://localhost:8124"), true)]
fn storage_exports_writer_and_optional_reader(
#[case] reader_url: Option<&str>,
#[case] has_reader: bool,
) {
let storage = Storage::new("litellm".to_owned(), "http://localhost:8123", reader_url)
.expect("valid ClickHouse URLs");
#[test]
fn storage_uses_one_url_for_writes_and_bounded_reads() {
let storage =
Storage::new("litellm".to_owned(), "http://localhost:8123").expect("valid ClickHouse URLs");
assert_eq!(storage.database(), "litellm");
assert_eq!(storage.writer().url().host_str(), Some("localhost"));
assert_eq!(storage.writer().url().port(), Some(8123));
assert_eq!(storage.reader().is_some(), has_reader);
assert_eq!(storage.reader().url().port(), Some(8123));
assert_eq!(
storage
.reader()
.url()
.query_pairs()
.find(|(key, _)| key == "database")
.unwrap()
.1,
"litellm"
);
}
#[rstest]
#[case::empty("")]
#[case::injection("db; DROP DATABASE default")]
fn storage_rejects_invalid_database(#[case] database: &str) {
assert!(Storage::new(database.to_owned(), "http://localhost:8123", None).is_err());
assert!(Storage::new(database.to_owned(), "http://localhost:8123").is_err());
}

2
litellm-rust/crates/traces/AGENTS.md
View file

@ -1,6 +1,6 @@
- Keep OTLP decoding, trace schema, row encoding and named query selection here. Generic ClickHouse connections and HTTP execution belong in `litellm-storage-clickhouse`
- Keep this crate independent of Python; PyO3 conversion and public Python exceptions belong in `python-bridge`
- Keep the SQL migrations here as the only ClickHouse schema definition
- Keep the SQL migrations here as the only ClickHouse schema definition, as `migrations/NNNN_description.sql` files embedded by `litellm_migrate::migrate!`; adding a file is the only step
- Use typed query parameters and a dedicated SELECT-only reader with server-side limits
- Keep `config/reader.xml` grants on the database the schema is created in (CLICKHOUSE_DATABASE, default `litellm`)
- Bound insert time and encoded bytes; make retry deduplication behavior explicit for supported ClickHouse versions

8
litellm-rust/crates/traces/Cargo.toml
View file

@ -6,18 +6,26 @@ license.workspace = true
repository.workspace = true
[dependencies]
askama.workspace = true
base64.workspace = true
flate2.workspace = true
futures-util.workspace = true
hmac = "0.12.1"
indexmap = { version = "2", features = ["serde"] }
moka.workspace = true
opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] }
prost.workspace = true
time = { workspace = true, features = ["formatting"] }
litellm-http.workspace = true
litellm-migrate.workspace = true
litellm-storage-clickhouse.workspace = true
sha2.workspace = true
serde = { workspace = true, features = ["rc"] }
serde_json.workspace = true
strum.workspace = true
thiserror.workspace = true
tokio.workspace = true
url.workspace = true
[dev-dependencies]
criterion.workspace = true

3
litellm-rust/crates/traces/build.rs Normal file
View file

@ -0,0 +1,3 @@
fn main() {
println!("cargo:rerun-if-changed=migrations");
}

3
litellm-rust/crates/traces/migrations/0001_otel_traces.sql
View file

@ -38,10 +38,11 @@ CREATE TABLE IF NOT EXISTS {database}.otel_traces
Input String CODEC(ZSTD(3)),
Output String CODEC(ZSTD(3)),
InputPreview String DEFAULT substring(Input, 1, 240),
EngineReceivedMs UInt64 DEFAULT 0,
INDEX idx_trace_id TraceId TYPE bloom_filter(0.001) GRANULARITY 1,
INDEX idx_req_id LiteLLMRequestId TYPE bloom_filter(0.01) GRANULARITY 1
)
ENGINE = MergeTree
PARTITION BY toDate(Timestamp)
ORDER BY (TeamId, ServiceName, toDateTime(Timestamp), TraceId)
SETTINGS ttl_only_drop_parts = 1, non_replicated_deduplication_window = 1000
SETTINGS ttl_only_drop_parts = 1, materialize_ttl_recalculate_only = 1, non_replicated_deduplication_window = 1000

2
litellm-rust/crates/traces/migrations/0005_otel_traces_ttl.sql → litellm-rust/crates/traces/migrations/0002_otel_traces_ttl.sql
View file

@ -1 +1 @@
ALTER TABLE {database}.otel_traces MODIFY TTL toDateTime(Timestamp) + INTERVAL {trace_retention_days} DAY
ALTER TABLE {database}.otel_traces MODIFY TTL toDateTime(Timestamp) + INTERVAL {retention_days} DAY

2
litellm-rust/crates/traces/migrations/0002_agent_traces.sql → litellm-rust/crates/traces/migrations/0003_agent_traces.sql
View file

@ -22,4 +22,4 @@ CREATE TABLE IF NOT EXISTS {database}.agent_traces_by_key
)
ENGINE = AggregatingMergeTree
ORDER BY (TeamId, ApiKeyHash, TraceId)
SETTINGS non_replicated_deduplication_window = 1000
SETTINGS materialize_ttl_recalculate_only = 1, non_replicated_deduplication_window = 1000

2
litellm-rust/crates/traces/migrations/0006_agent_traces_ttl.sql → litellm-rust/crates/traces/migrations/0004_agent_traces_ttl.sql
View file

@ -1 +1 @@
ALTER TABLE {database}.agent_traces_by_key MODIFY TTL toDateTime(StartTs) + INTERVAL {trace_retention_days} DAY
ALTER TABLE {database}.agent_traces_by_key MODIFY TTL toDateTime(StartTs) + INTERVAL {retention_days} DAY

0
litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql → litellm-rust/crates/traces/migrations/0005_agent_traces_mv.sql
View file

2
litellm-rust/crates/traces/migrations/0004_spend_logs.sql → litellm-rust/crates/traces/migrations/0006_spend_logs.sql
View file

@ -34,9 +34,11 @@ CREATE TABLE IF NOT EXISTS {database}.spend_logs
metadata String CODEC(ZSTD(3)),
messages String CODEC(ZSTD(3)),
response String CODEC(ZSTD(3)),
EngineReceivedMs UInt64 DEFAULT 0,
INDEX idx_response_id response_id TYPE bloom_filter(0.001) GRANULARITY 1,
INDEX idx_trace_id trace_id TYPE bloom_filter(0.001) GRANULARITY 1
)
ENGINE = ReplacingMergeTree(end_time)
PARTITION BY toYYYYMM(start_time)
ORDER BY (team_id, start_time, request_id)
SETTINGS materialize_ttl_recalculate_only = 1

2
litellm-rust/crates/traces/migrations/0007_spend_logs_ttl.sql
View file

@ -1 +1 @@
ALTER TABLE {database}.spend_logs MODIFY TTL toDateTime(start_time) + INTERVAL {spend_log_retention_days} DAY
ALTER TABLE {database}.spend_logs MODIFY TTL toDateTime(start_time) + INTERVAL {retention_days} DAY

1
litellm-rust/crates/traces/migrations/0008_otel_traces_framework.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE {database}.otel_traces ADD COLUMN IF NOT EXISTS Framework LowCardinality(String) AFTER AgentName

1
litellm-rust/crates/traces/migrations/0008_trace_received.sql
View file

@ -1 +0,0 @@
ALTER TABLE {database}.otel_traces ADD COLUMN IF NOT EXISTS EngineReceivedMs UInt64 DEFAULT 0

1
litellm-rust/crates/traces/migrations/0009_spend_received.sql
View file

@ -1 +0,0 @@
ALTER TABLE {database}.spend_logs ADD COLUMN IF NOT EXISTS EngineReceivedMs UInt64 DEFAULT 0

6
litellm-rust/crates/traces/query/lens_agents.sql Normal file
View file

@ -0,0 +1,6 @@
SELECT DISTINCT AgentName AS agent_name
FROM otel_traces
WHERE AgentName != ''
AND ({all_teams:UInt8}=1 OR TeamId={team:String})
AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String})
ORDER BY agent_name

8
litellm-rust/crates/traces/query/lens_availability.sql Normal file
View file

@ -0,0 +1,8 @@
SELECT
EXISTS(SELECT 1 FROM otel_traces
WHERE ({all_teams:UInt8}=1 OR TeamId={team:String})
AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String})) AS traces,
EXISTS(SELECT 1 FROM spend_logs
WHERE ({all_teams:UInt8}=1 OR team_id={team:String})
AND ({key_hash:String}='' OR api_key={key_hash:String})
AND NOT JSONExtractBool(metadata,'litellm_lens_internal')) AS requests

2
litellm-rust/crates/traces/query/lens_sample.sql
View file

@ -28,6 +28,7 @@ SELECT *, selection_key FROM (
GROUP BY TeamId,ApiKeyHash,TraceId
HAVING max(EngineReceivedMs) < {end:UInt64}
AND max(toUnixTimestamp64Milli(Timestamp)+toInt64(intDiv(Duration,1000000))) < {end:UInt64}
AND ({agent_name:String}='' OR countIf(AgentName={agent_name:String}) > 0)
AND countIf(arrayAll((k,v) -> ResourceAttributes[k]=v OR SpanAttributes[k]=v,
{filter_keys:Array(String)},{filter_values:Array(String)})
AND ({service:String}='' OR ServiceName={service:String})) > 0
@ -48,6 +49,7 @@ SELECT *, selection_key FROM (
OR JSONExtractString(metadata,'requester_metadata',k)=v OR (k='tag' AND has(request_tags,v)),
{filter_keys:Array(String)},{filter_values:Array(String)})
AND ({service:String}='' OR model_group={service:String})
AND {agent_name:String}=''
AND NOT JSONExtractBool(metadata,'litellm_lens_internal')
AND ({source:String}!='both' OR (team_id,api_key,response_id) NOT IN (
SELECT TeamId,ApiKeyHash,LiteLLMRequestId FROM otel_traces

24
litellm-rust/crates/traces/query/list_traces.sql
View file

@ -1,11 +1,13 @@
WITH page AS (
SELECT TraceId AS trace_id,
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref,
TeamId AS team_id, ApiKeyHash AS api_key_hash,
ifNull(any(RootName), '') AS name, any(ServiceName) AS service,
ifNull(any(RootInput), '') AS input_preview, ifNull(any(RootStatus), '') AS status,
toUnixTimestamp64Milli(min(StartTs)) AS start_ms,
min(StartTs) AS trace_start, max(EndTs) AS trace_end,
dateDiff('millisecond', min(StartTs), max(EndTs)) AS duration_ms,
sum(SpanCount) AS span_count, length(groupUniqArrayArray(AgentNames)) AS agent_count,
sum(SpanCount) AS span_count,
sum(AgentCount) AS agent_invocations,
sum(LlmCount) AS llm_calls, sum(ToolCount) AS tool_calls,
sum(InputTokens) AS input_tokens, sum(OutputTokens) AS output_tokens,
@ -21,3 +23,23 @@ HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64})
< ({cursor_ms:Int64}, {cursor_trace_id:String}))
ORDER BY start_ms DESC, trace_ref DESC
LIMIT {limit:UInt32}
)
SELECT page.* EXCEPT (trace_start, trace_end),
identities.agent_names AS agent_names, identities.agent_count AS agent_count,
identities.frameworks AS frameworks
FROM page
LEFT JOIN (
SELECT TeamId, ApiKeyHash, TraceId,
arraySort(groupUniqArrayIf(AgentName, AgentName != '')) AS agent_names,
arraySort(groupUniqArrayIf(toString(Framework), Framework != '')) AS frameworks,
uniqExactIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS agent_count
FROM otel_traces
WHERE Timestamp >= (SELECT min(trace_start) FROM page)
AND Timestamp <= (SELECT max(trace_end) FROM page)
AND TraceId IN (SELECT trace_id FROM page)
AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page)
GROUP BY TeamId, ApiKeyHash, TraceId
) AS identities
ON page.team_id = identities.TeamId AND page.api_key_hash = identities.ApiKeyHash
AND page.trace_id = identities.TraceId
ORDER BY page.start_ms DESC, page.trace_ref DESC

23
litellm-rust/crates/traces/query/span_detail.sql
View file

@ -1,8 +1,19 @@
SELECT SpanId AS span_id, Input AS input, Output AS output, SpanAttributes AS attributes
FROM otel_traces
WHERE TraceId = {trace_id:String} AND SpanId = {span_id:String}
AND (empty({team_ids:Array(String)}) OR TeamId IN {team_ids:Array(String)})
AND ({api_key_hash:String} = '' OR ApiKeyHash = {api_key_hash:String})
SELECT o.SpanId AS span_id, o.Input AS input,
if(o.Output = '' AND o.ObservationType = 'agent', answer.output, o.Output) AS output,
o.SpanAttributes AS attributes
FROM otel_traces AS o
LEFT JOIN (
SELECT ParentSpanId AS parent_span_id, argMax(Output, Timestamp) AS output
FROM otel_traces
WHERE TraceId = {trace_id:String} AND ParentSpanId = {span_id:String}
AND ObservationType = 'llm' AND Output != ''
AND (empty({team_ids:Array(String)}) OR TeamId IN {team_ids:Array(String)})
AND ({api_key_hash:String} = '' OR ApiKeyHash = {api_key_hash:String})
GROUP BY ParentSpanId
) AS answer ON answer.parent_span_id = o.SpanId
WHERE o.TraceId = {trace_id:String} AND o.SpanId = {span_id:String}
AND (empty({team_ids:Array(String)}) OR o.TeamId IN {team_ids:Array(String)})
AND ({api_key_hash:String} = '' OR o.ApiKeyHash = {api_key_hash:String})
AND ({trace_ref:String} = '' OR
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
LIMIT 1

3
litellm-rust/crates/traces/query/trace_spans.sql
View file

@ -1,5 +1,6 @@
SELECT o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name,
o.ObservationType AS type, o.AgentName AS agent, o.StatusCode AS status,
o.ObservationType AS type, o.AgentName AS agent,
o.Framework AS framework, o.StatusCode AS status,
substringUTF8(o.StatusMessage, 1, 128) AS status_message,
lengthUTF8(o.StatusMessage) > 128 AS error_truncated,
toUnixTimestamp64Nano(o.Timestamp) AS start_ns, o.Duration AS duration_ns,

25
litellm-rust/crates/traces/src/config.rs Normal file
View file

@ -0,0 +1,25 @@
use litellm_storage_clickhouse::{Error, Storage};
#[derive(Clone)]
pub struct Config {
storage: Storage,
retention_days: u32,
}
impl Config {
pub fn new(database: String, url: &str, retention_days: u32) -> Result<Self, Error> {
crate::schema_statements(&database, retention_days)?;
Ok(Self {
storage: Storage::new(database, url)?,
retention_days,
})
}
pub fn storage(&self) -> &Storage {
&self.storage
}
pub fn retention_days(&self) -> u32 {
self.retention_days
}
}

22
litellm-rust/crates/traces/src/error.rs
View file

@ -4,4 +4,26 @@ pub enum DecodeError {
InvalidPayload,
#[error("OTLP trace payload exceeds the decoding budget")]
TooLarge,
#[error("OTLP token count is outside the storage range")]
TokenCountOutOfRange,
}
#[derive(Debug, thiserror::Error)]
pub enum QueryAccessError {
#[error("trace SQL queries require a configured proxy master key")]
MissingSecret,
#[error("invalid trace query scope")]
InvalidScope,
#[error("trace SQL query concurrency limit exceeded")]
Busy,
#[error(
"ClickHouse reader provisioning failed with HTTP status {0}; the configured connection must be allowed to manage users, row policies, and SELECT grants on the trace tables"
)]
ProvisionFailed(u16),
#[error("ClickHouse reader provisioning transport failed")]
ProvisionTransport,
#[error(transparent)]
Storage(#[from] litellm_storage_clickhouse::Error),
#[error(transparent)]
Cached(#[from] std::sync::Arc<QueryAccessError>),
}

13
litellm-rust/crates/traces/src/lib.rs
View file

@ -1,14 +1,25 @@
mod config;
mod error;
mod insert;
mod normalize;
mod otlp;
mod query;
mod query_access;
mod schema;
mod shared;
mod sql;
pub use error::DecodeError;
pub use config::Config;
pub use error::{DecodeError, QueryAccessError};
pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows};
pub use litellm_storage_clickhouse::{Connection, Error, Parameter, execute_read};
pub use normalize::{
NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, NormalizedSpan, ObservationType,
};
pub use otlp::{DecodedSpan, decode_otlp};
pub use query_access::{QueryReaders, QueryScope};
pub use schema::{ensure_schema, schema_statements};
pub use shared::{Shared, SharedIdentity};
pub use sql::{LensQuery, ReadQuery, execute_named_read};
pub use query::{query_help, query_sql};

367
litellm-rust/crates/traces/src/normalize/claude_code.rs Normal file
View file

@ -0,0 +1,367 @@
use std::collections::BTreeMap;
use serde_json::{Map, Value, json};
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, first, tokens};
use crate::{DecodeError, otlp::DecodedEvent};
pub(crate) const CLAUDE_CODE_SCOPE: &str = "com.anthropic.claude_code.tracing";
pub(crate) const CLAUDE_CODE_AGENT: &str = "claude-code";
const AGENT_SDK_FRAMEWORK: &str = "claude-agent-sdk";
pub(super) struct ClaudeCodeNormalizer;
enum SpanType {
Interaction,
LlmRequest,
Tool,
Other,
}
fn span_type(name: &str, attributes: &BTreeMap<String, String>) -> SpanType {
let kind = attr(attributes, "span.type");
let kind = if kind.is_empty() {
name.strip_prefix("claude_code.").unwrap_or(name)
} else {
kind
};
match kind {
"interaction" => SpanType::Interaction,
"llm_request" => SpanType::LlmRequest,
"tool" => SpanType::Tool,
_ => SpanType::Other,
}
}
fn framework(attributes: &BTreeMap<String, String>) -> &'static str {
if attr(attributes, "query_source_safe") == "sdk"
|| attr(attributes, "system_prompt_preview").contains("cc_entrypoint=sdk")
{
AGENT_SDK_FRAMEWORK
} else {
CLAUDE_CODE_AGENT
}
}
fn split_header(text: &str) -> Option<(&str, &str)> {
let (header, body) = text.strip_prefix('[')?.split_once("]\n")?;
Some((header, body))
}
fn without_header<'a>(text: &'a str, prefix: &str) -> &'a str {
split_header(text)
.filter(|(header, _)| header.starts_with(prefix))
.map_or(text, |(_, body)| body)
}
fn tool_arguments(attributes: &BTreeMap<String, String>) -> Option<&str> {
let arguments = without_header(attr(attributes, "tool_input"), "TOOL INPUT");
serde_json::from_str::<Map<String, Value>>(arguments)
.is_ok()
.then_some(arguments)
}
fn tool_input(attributes: &BTreeMap<String, String>) -> String {
if let Some(arguments) = tool_arguments(attributes) {
return arguments.to_owned();
}
let fields: Map<String, Value> = [
("command", "full_command"),
("file_path", "file_path"),
("bash_argv0", "bash_argv0"),
]
.into_iter()
.filter_map(|(key, source)| {
let value = attr(attributes, source);
(!value.is_empty()).then(|| (key.to_owned(), Value::String(value.to_owned())))
})
.collect();
if fields.is_empty() {
String::new()
} else {
Value::Object(fields).to_string()
}
}
fn tool_output(attributes: &BTreeMap<String, String>, events: &[DecodedEvent]) -> String {
events
.iter()
.filter(|event| event.name == "tool.output")
.flat_map(|event| {
["output", "content", "diff"]
.into_iter()
.map(|key| attr(&event.attributes, key))
})
.find(|value| !value.is_empty())
.unwrap_or_else(|| without_header(attr(attributes, "new_context"), "TOOL RESULT"))
.to_owned()
}
fn context_message(context: &str) -> Value {
let (role, content) = match split_header(context) {
Some(("USER" | "USER PROMPT", body)) => ("user", body),
Some(("ASSISTANT", body)) => ("assistant", body),
Some((header, body)) if header.starts_with("TOOL RESULT") => ("tool", body),
_ => ("user", context),
};
json!({"role": role, "content": content})
}
fn user_prompt(attributes: &BTreeMap<String, String>) -> String {
let prompt = attr(attributes, "user_prompt");
if prompt.is_empty() {
String::new()
} else {
json!([{"role": "user", "content": prompt}]).to_string()
}
}
fn llm_input(attributes: &BTreeMap<String, String>) -> String {
let messages: Vec<Value> = [
Some(attr(attributes, "system_prompt_preview"))
.filter(|system| !system.is_empty())
.map(|system| json!({"role": "system", "content": system})),
Some(attr(attributes, "new_context"))
.filter(|context| !context.is_empty())
.map(context_message),
]
.into_iter()
.flatten()
.collect();
if messages.is_empty() {
String::new()
} else {
Value::Array(messages).to_string()
}
}
fn llm_output(attributes: &BTreeMap<String, String>) -> String {
let output = attr(attributes, "response.model_output");
if output.is_empty() {
String::new()
} else {
json!({"role": "assistant", "content": output}).to_string()
}
}
fn input_tokens(attributes: &BTreeMap<String, String>) -> Result<u32, DecodeError> {
["input_tokens", "cache_read_tokens", "cache_creation_tokens"]
.into_iter()
.try_fold(0u32, |total, key| {
total
.checked_add(tokens(attributes, key)?)
.ok_or(DecodeError::TokenCountOutOfRange)
})
}
impl SpanNormalizer for ClaudeCodeNormalizer {
fn matches(&self, scope_name: &str, _attributes: &BTreeMap<String, String>) -> bool {
scope_name == CLAUDE_CODE_SCOPE
}
fn consumed_attributes(&self, attributes: &BTreeMap<String, String>) -> [&'static str; 2] {
match span_type("", attributes) {
SpanType::Interaction => ["user_prompt", ""],
SpanType::LlmRequest => ["new_context", "response.model_output"],
SpanType::Tool if tool_arguments(attributes).is_some() => ["tool_input", ""],
SpanType::Tool | SpanType::Other => ["", ""],
}
}
fn display_name(&self, attributes: &BTreeMap<String, String>) -> Option<String> {
let tool_name = attr(attributes, "tool_name");
(matches!(span_type("", attributes), SpanType::Tool) && !tool_name.is_empty())
.then(|| tool_name.to_owned())
}
fn normalize(
&self,
name: &str,
_parent_span_id: &str,
attributes: &BTreeMap<String, String>,
events: &[DecodedEvent],
) -> Result<NormalizedSpan, DecodeError> {
let base = NormalizedSpan {
observation_type: ObservationType::Framework,
agent_name: CLAUDE_CODE_AGENT.to_owned(),
framework: framework(attributes).to_owned(),
litellm_request_id: String::new(),
model: String::new(),
input_tokens: 0,
output_tokens: 0,
input: String::new(),
output: String::new(),
};
Ok(match span_type(name, attributes) {
SpanType::Interaction => NormalizedSpan {
observation_type: ObservationType::Agent,
input: user_prompt(attributes),
..base
},
SpanType::LlmRequest => NormalizedSpan {
observation_type: ObservationType::Llm,
litellm_request_id: first(attributes, "gen_ai.response.id", "request_id")
.to_owned(),
model: first(attributes, "model", "gen_ai.request.model").to_owned(),
input_tokens: input_tokens(attributes)?,
output_tokens: tokens(attributes, "output_tokens")?,
input: llm_input(attributes),
output: llm_output(attributes),
..base
},
SpanType::Tool => NormalizedSpan {
observation_type: ObservationType::Tool,
input: tool_input(attributes),
output: tool_output(attributes, events),
..base
},
SpanType::Other => base,
})
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use rstest::rstest;
use serde_json::Value;
use super::{CLAUDE_CODE_SCOPE, ClaudeCodeNormalizer, SpanNormalizer};
use crate::{DecodeError, normalize::ObservationType, otlp::DecodedEvent};
fn attributes(pairs: &[(&str, &str)]) -> BTreeMap<String, String> {
pairs
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect()
}
#[rstest]
fn tool_without_detailed_input_lists_known_arguments() {
let span = ClaudeCodeNormalizer
.normalize(
"claude_code.tool",
"parent",
&attributes(&[
("span.type", "tool"),
("tool_name", "Bash"),
("full_command", "git status"),
("bash_argv0", "git"),
]),
&[],
)
.expect("valid span");
let input: Value = serde_json::from_str(&span.input).expect("argument object");
assert_eq!(input["command"], "git status");
assert_eq!(input["bash_argv0"], "git");
assert!(input.get("file_path").is_none());
assert!(input.get("role").is_none());
}
#[rstest]
fn malformed_tool_input_falls_back_and_stays_in_attributes() {
let attrs = attributes(&[
("span.type", "tool"),
("tool_input", "[TOOL INPUT: Read]\nnot json"),
("file_path", "/workspace/a.py"),
]);
let span = ClaudeCodeNormalizer
.normalize("claude_code.tool", "parent", &attrs, &[])
.expect("valid span");
let input: Value = serde_json::from_str(&span.input).expect("argument object");
assert_eq!(input["file_path"], "/workspace/a.py");
assert!(
!ClaudeCodeNormalizer
.consumed_attributes(&attrs)
.contains(&"tool_input")
);
}
#[rstest]
#[case::event_output(
vec![DecodedEvent { name: "tool.output".to_owned(), attributes: attributes(&[("output", "stdout text")]) }],
"stdout text"
)]
#[case::event_diff(
vec![DecodedEvent { name: "tool.output".to_owned(), attributes: attributes(&[("diff", "+line")]) }],
"+line"
)]
#[case::other_event_ignored(
vec![DecodedEvent { name: "other".to_owned(), attributes: attributes(&[("output", "nope")]) }],
"{\"stdout\":\"ctx\"}"
)]
fn tool_output_prefers_event_then_context(
#[case] events: Vec<DecodedEvent>,
#[case] expected: &str,
) {
let span = ClaudeCodeNormalizer
.normalize(
"claude_code.tool",
"parent",
&attributes(&[
("span.type", "tool"),
("new_context", "[TOOL RESULT: Bash]\n{\"stdout\":\"ctx\"}"),
]),
&events,
)
.expect("valid span");
assert_eq!(span.output, expected);
}
#[rstest]
fn llm_tool_result_context_becomes_tool_message() {
let span = ClaudeCodeNormalizer
.normalize(
"claude_code.llm_request",
"parent",
&attributes(&[
("span.type", "llm_request"),
("new_context", "[TOOL RESULT: toolu_1]\n1\timport os"),
]),
&[],
)
.expect("valid span");
let input: Value = serde_json::from_str(&span.input).expect("messages");
assert_eq!(input[0]["role"], "tool");
assert_eq!(input[0]["content"], "1\timport os");
assert_eq!(span.output, "");
assert_eq!(span.framework, "claude-code");
}
#[rstest]
fn llm_token_sum_overflow_is_rejected() {
let result = ClaudeCodeNormalizer.normalize(
"claude_code.llm_request",
"parent",
&attributes(&[
("span.type", "llm_request"),
("input_tokens", "4294967295"),
("cache_read_tokens", "1"),
]),
&[],
);
assert!(matches!(result, Err(DecodeError::TokenCountOutOfRange)));
}
#[rstest]
#[case::span_type_wins("claude_code.tool", "hook", ObservationType::Framework)]
#[case::name_fallback("claude_code.interaction", "", ObservationType::Agent)]
#[case::unknown("claude_code.something_new", "", ObservationType::Framework)]
fn span_type_attribute_then_name_select_the_observation(
#[case] name: &str,
#[case] kind: &str,
#[case] expected: ObservationType,
) {
let attrs = if kind.is_empty() {
BTreeMap::new()
} else {
attributes(&[("span.type", kind)])
};
let span = ClaudeCodeNormalizer
.normalize(name, "parent", &attrs, &[])
.expect("valid span");
assert_eq!(span.observation_type, expected);
assert!(ClaudeCodeNormalizer.matches(CLAUDE_CODE_SCOPE, &attrs));
}
}

65
litellm-rust/crates/traces/src/normalize/genai.rs Normal file
View file

@ -0,0 +1,65 @@
use std::collections::BTreeMap;
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, first, usage_tokens};
use crate::{DecodeError, otlp::DecodedEvent};
pub(super) struct GenAiNormalizer;
impl SpanNormalizer for GenAiNormalizer {
fn matches(&self, _scope_name: &str, _attributes: &BTreeMap<String, String>) -> bool {
true
}
fn consumed_attributes(&self, attributes: &BTreeMap<String, String>) -> [&'static str; 2] {
[
if attr(attributes, "gen_ai.input.messages").is_empty() {
"gen_ai.tool.call.arguments"
} else {
"gen_ai.input.messages"
},
if attr(attributes, "gen_ai.output.messages").is_empty() {
"gen_ai.tool.call.result"
} else {
"gen_ai.output.messages"
},
]
}
fn normalize(
&self,
_name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
_events: &[DecodedEvent],
) -> Result<NormalizedSpan, DecodeError> {
let (input_tokens, output_tokens) = usage_tokens(attributes)?;
let observation_type = match attr(attributes, "gen_ai.operation.name") {
"invoke_agent" => ObservationType::Agent,
"chat" | "text_completion" | "generate_content" => ObservationType::Llm,
"execute_tool" => ObservationType::Tool,
_ if parent_span_id.is_empty() => ObservationType::Agent,
_ => ObservationType::Chain,
};
Ok(NormalizedSpan {
observation_type,
agent_name: attr(attributes, "gen_ai.agent.name").to_owned(),
framework: String::new(),
litellm_request_id: attr(attributes, "gen_ai.response.id").to_owned(),
model: first(attributes, "gen_ai.request.model", "gen_ai.response.model").to_owned(),
input_tokens,
output_tokens,
input: first(
attributes,
"gen_ai.input.messages",
"gen_ai.tool.call.arguments",
)
.to_owned(),
output: first(
attributes,
"gen_ai.output.messages",
"gen_ai.tool.call.result",
)
.to_owned(),
})
}
}

470
litellm-rust/crates/traces/src/normalize/langsmith.rs Normal file
View file

@ -0,0 +1,470 @@
use std::{collections::BTreeMap, io};
use indexmap::IndexMap;
use serde::{Deserialize, Deserializer, Serialize, de::DeserializeOwned};
use serde_json::{Value, ser::Formatter};
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, usage_tokens};
use crate::{DecodeError, otlp::DecodedEvent};
pub(super) struct LangSmithNormalizer;
#[derive(Deserialize)]
#[serde(untagged)]
enum MessageContent {
Text(String),
Blocks(Vec<ContentBlock>),
Other(Value),
}
impl MessageContent {
fn display_text(&self) -> String {
match self {
Self::Text(text) => text.clone(),
Self::Blocks(blocks) => blocks
.iter()
.filter_map(|block| match block {
ContentBlock::Text { text } => Some(text.as_str()),
ContentBlock::Hidden(kind) => match kind {
HiddenBlock::Reasoning
| HiddenBlock::Thinking
| HiddenBlock::RedactedThinking
| HiddenBlock::FunctionCall
| HiddenBlock::ToolUse
| HiddenBlock::ToolCall => None,
},
})
.collect::<Vec<_>>()
.join("\n\n"),
Self::Other(value) => encode(value),
}
}
}
#[derive(Deserialize)]
#[serde(untagged)]
enum ContentBlock {
Text { text: String },
Hidden(HiddenBlock),
}
#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum HiddenBlock {
Reasoning,
Thinking,
RedactedThinking,
FunctionCall,
ToolUse,
ToolCall,
}
#[derive(Deserialize, Serialize)]
#[serde(transparent)]
struct RawToolCall(IndexMap<String, Value>);
#[derive(Deserialize)]
struct ResponseMetadata {
id: Option<String>,
}
#[derive(Deserialize)]
struct RawMessage {
kwargs: Option<Box<RawMessage>>,
#[serde(rename = "type")]
kind: Option<String>,
role: Option<String>,
content: Option<MessageContent>,
tool_calls: Option<Vec<RawToolCall>>,
name: Option<Value>,
response_metadata: Option<ResponseMetadata>,
}
impl RawMessage {
fn unwrapped(&self) -> &Self {
self.kwargs.as_deref().unwrap_or(self)
}
fn normalized(&self) -> NormalizedMessage<'_> {
let fields = self.unwrapped();
let raw_role = fields
.kind
.as_deref()
.filter(|role| !role.is_empty())
.or_else(|| fields.role.as_deref().filter(|role| !role.is_empty()))
.unwrap_or_default();
let role = match raw_role {
"human" => "user",
"ai" => "assistant",
other => other,
};
NormalizedMessage {
role,
content: fields
.content
.as_ref()
.map_or_else(String::new, MessageContent::display_text),
tool_calls: fields
.tool_calls
.as_deref()
.filter(|calls| !calls.is_empty()),
name: (role == "tool")
.then_some(fields.name.as_ref())
.flatten()
.filter(|name| !name.is_null() && name != &&Value::String(String::new())),
}
}
}
#[derive(Serialize)]
struct NormalizedMessage<'a> {
role: &'a str,
content: String,
#[serde(skip_serializing_if = "Option::is_none")]
tool_calls: Option<&'a [RawToolCall]>,
#[serde(skip_serializing_if = "Option::is_none")]
name: Option<&'a Value>,
}
enum MessageBatch {
Flat(Vec<RawMessage>),
Nested(Vec<Vec<RawMessage>>),
}
impl<'de> Deserialize<'de> for MessageBatch {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let value = Value::deserialize(deserializer)?;
let Value::Array(items) = value else {
return Err(serde::de::Error::custom("messages must be an array"));
};
let parse = |items: Vec<Value>| {
items
.into_iter()
.filter_map(|item| serde_json::from_value(item).ok())
.collect()
};
Ok(if items.first().is_some_and(Value::is_array) {
Self::Nested(
items
.into_iter()
.filter_map(|item| item.as_array().cloned())
.map(parse)
.collect(),
)
} else {
Self::Flat(parse(items))
})
}
}
fn lenient<'de, D: Deserializer<'de>, T: DeserializeOwned>(
deserializer: D,
) -> Result<Option<T>, D::Error> {
let value = Value::deserialize(deserializer)?;
Ok(serde_json::from_value(value).ok())
}
impl MessageBatch {
fn first_batch(&self) -> &[RawMessage] {
match self {
Self::Flat(messages) => messages,
Self::Nested(batches) => batches.first().map(Vec::as_slice).unwrap_or_default(),
}
}
fn agent_messages(&self) -> &[RawMessage] {
match self {
Self::Flat(messages) => messages,
Self::Nested(_) => &[],
}
}
}
#[derive(Deserialize)]
struct GenerationMessage {
kwargs: Option<RawMessage>,
}
#[derive(Deserialize)]
struct Generation {
message: Option<GenerationMessage>,
}
#[derive(Default, Deserialize)]
struct Payload {
#[serde(default, deserialize_with = "lenient")]
messages: Option<MessageBatch>,
#[serde(default, deserialize_with = "lenient")]
generations: Option<Vec<Vec<Generation>>>,
}
#[derive(Deserialize)]
struct Command {
update: CommandUpdate,
}
#[derive(Deserialize)]
struct CommandUpdate {
messages: Vec<Value>,
}
#[derive(Deserialize)]
struct ContentValue {
content: Value,
}
struct SpanIo {
input: String,
output: String,
request_id: String,
}
struct PythonJsonFormatter;
impl Formatter for PythonJsonFormatter {
fn begin_array_value<W: ?Sized + io::Write>(
&mut self,
writer: &mut W,
first: bool,
) -> io::Result<()> {
if first {
Ok(())
} else {
writer.write_all(b", ")
}
}
fn begin_object_key<W: ?Sized + io::Write>(
&mut self,
writer: &mut W,
first: bool,
) -> io::Result<()> {
if first {
Ok(())
} else {
writer.write_all(b", ")
}
}
fn begin_object_value<W: ?Sized + io::Write>(&mut self, writer: &mut W) -> io::Result<()> {
writer.write_all(b": ")
}
}
fn encode<T: Serialize>(value: &T) -> String {
let mut output = Vec::new();
let mut serializer = serde_json::Serializer::with_formatter(&mut output, PythonJsonFormatter);
if value.serialize(&mut serializer).is_err() {
return String::new();
}
String::from_utf8(output).unwrap_or_default()
}
fn normalized_messages(messages: &[RawMessage]) -> String {
encode(
&messages
.iter()
.map(RawMessage::normalized)
.collect::<Vec<_>>(),
)
}
fn span_type(
name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
) -> ObservationType {
match attr(attributes, "langsmith.span.kind") {
"llm" => ObservationType::Llm,
"tool" => ObservationType::Tool,
_ if parent_span_id.is_empty()
|| name == attr(attributes, "langsmith.metadata.lc_agent_name") =>
{
ObservationType::Agent
}
_ if [
".wrap_model_call",
".wrap_tool_call",
".before_agent",
".after_agent",
".before_model",
".after_model",
]
.iter()
.any(|suffix| name.ends_with(suffix)) =>
{
ObservationType::Framework
}
_ => ObservationType::Chain,
}
}
fn tool_output(raw_completion: &str) -> String {
let completion = serde_json::from_str::<Value>(raw_completion).unwrap_or(Value::Null);
let raw = completion.get("output").cloned().unwrap_or(completion);
let selected = serde_json::from_value::<Command>(raw.clone())
.ok()
.and_then(|command| command.update.messages.into_iter().last())
.unwrap_or(raw);
let output = serde_json::from_value::<ContentValue>(selected.clone())
.map(|message| message.content)
.unwrap_or(selected);
output
.as_str()
.map(str::to_owned)
.unwrap_or_else(|| encode(&output))
}
fn span_io(kind: ObservationType, attributes: &BTreeMap<String, String>) -> SpanIo {
let raw_prompt = attr(attributes, "gen_ai.prompt");
let raw_completion = attr(attributes, "gen_ai.completion");
let prompt = serde_json::from_str::<Payload>(raw_prompt).unwrap_or_default();
let completion = serde_json::from_str::<Payload>(raw_completion).unwrap_or_default();
if kind == ObservationType::Llm
&& serde_json::from_str::<Value>(raw_completion).is_ok_and(|value| value.is_object())
{
let input = prompt.messages.as_ref().map_or_else(
|| "[]".to_owned(),
|messages| normalized_messages(messages.first_batch()),
);
let generation = completion
.generations
.as_ref()
.and_then(|batches| batches.first())
.and_then(|batch| batch.first())
.and_then(|generation| generation.message.as_ref())
.and_then(|message| message.kwargs.as_ref());
if let Some(generation) = generation {
let id = generation
.response_metadata
.as_ref()
.and_then(|metadata| metadata.id.as_deref())
.unwrap_or_default()
.to_owned();
return SpanIo {
input,
output: encode(&generation.normalized()),
request_id: id,
};
}
return SpanIo {
input,
output: raw_completion.to_owned(),
request_id: String::new(),
};
}
if kind == ObservationType::Tool {
return SpanIo {
input: raw_prompt.to_owned(),
output: tool_output(raw_completion),
request_id: String::new(),
};
}
if kind == ObservationType::Agent {
let input = prompt
.messages
.as_ref()
.filter(|messages| !messages.agent_messages().is_empty())
.map_or_else(
|| raw_prompt.to_owned(),
|messages| normalized_messages(messages.agent_messages()),
);
let output = completion
.messages
.as_ref()
.and_then(|messages| messages.agent_messages().last())
.map_or_else(
|| raw_completion.to_owned(),
|message| encode(&message.normalized()),
);
return SpanIo {
input,
output,
request_id: String::new(),
};
}
SpanIo {
input: raw_prompt.to_owned(),
output: raw_completion.to_owned(),
request_id: String::new(),
}
}
impl SpanNormalizer for LangSmithNormalizer {
fn matches(&self, scope_name: &str, attributes: &BTreeMap<String, String>) -> bool {
scope_name == "langsmith" || attributes.contains_key("langsmith.span.kind")
}
fn consumed_attributes(&self, _attributes: &BTreeMap<String, String>) -> [&'static str; 2] {
["gen_ai.prompt", "gen_ai.completion"]
}
fn normalize(
&self,
name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
_events: &[DecodedEvent],
) -> Result<NormalizedSpan, DecodeError> {
let (input_tokens, output_tokens) = usage_tokens(attributes)?;
let observation_type = span_type(name, parent_span_id, attributes);
let io = span_io(observation_type, attributes);
Ok(NormalizedSpan {
observation_type,
agent_name: attr(attributes, "langsmith.metadata.lc_agent_name").to_owned(),
framework: String::new(),
litellm_request_id: io.request_id,
model: attr(attributes, "gen_ai.request.model").to_owned(),
input_tokens,
output_tokens,
input: io.input,
output: io.output,
})
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use rstest::rstest;
use serde_json::Value;
use super::{ObservationType, span_io};
#[rstest]
fn malformed_messages_preserve_valid_input_and_response_id() {
let attributes = BTreeMap::from([
(
"gen_ai.prompt".to_owned(),
r#"{"messages":[[{"kwargs":{"type":"human","content":"hello"}},null]]}"#.to_owned(),
),
(
"gen_ai.completion".to_owned(),
r#"{"messages":"unexpected","generations":[[{"message":{"kwargs":{"type":"ai","content":"hi","response_metadata":{"id":"response-1"}}}}]]}"#.to_owned(),
),
]);
let io = span_io(ObservationType::Llm, &attributes);
let input: Value = serde_json::from_str(&io.input).expect("normalized input");
assert_eq!(input.as_array().expect("messages").len(), 1);
assert_eq!(input[0]["content"], "hello");
assert_eq!(io.request_id, "response-1");
}
#[rstest]
fn explicit_null_tool_output_is_preserved() {
let attributes = BTreeMap::from([(
"gen_ai.completion".to_owned(),
r#"{"output":null}"#.to_owned(),
)]);
let io = span_io(ObservationType::Tool, &attributes);
assert_eq!(io.output, "null");
}
#[rstest]
fn absent_llm_messages_render_as_an_empty_list() {
let attributes = BTreeMap::from([("gen_ai.completion".to_owned(), "{}".to_owned())]);
let io = span_io(ObservationType::Llm, &attributes);
assert_eq!(io.input, "[]");
}
}

326
litellm-rust/crates/traces/src/normalize/mod.rs Normal file
View file

@ -0,0 +1,326 @@
use std::collections::BTreeMap;
use crate::{DecodeError, otlp::DecodedEvent};
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum ObservationType {
Agent,
Llm,
Tool,
Chain,
Framework,
}
#[derive(Debug, Serialize)]
pub struct NormalizedSpan {
pub observation_type: ObservationType,
pub agent_name: String,
pub framework: String,
pub litellm_request_id: String,
pub model: String,
pub input_tokens: u32,
pub output_tokens: u32,
pub input: String,
pub output: String,
}
pub(crate) struct Normalization {
pub span: NormalizedSpan,
pub display_name: Option<String>,
pub consumed_attributes: [&'static str; 2],
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
pub struct NormalizedFieldDefinition {
pub name: &'static str,
pub clickhouse_column: &'static str,
pub clickhouse_type: &'static str,
pub meaning: &'static str,
}
pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 9] = [
NormalizedFieldDefinition {
name: "observation_type",
clickhouse_column: "ObservationType",
clickhouse_type: "LowCardinality(String)",
meaning: "Agent, LLM, tool, chain, or framework span",
},
NormalizedFieldDefinition {
name: "agent_name",
clickhouse_column: "AgentName",
clickhouse_type: "LowCardinality(String)",
meaning: "Agent associated with this span",
},
NormalizedFieldDefinition {
name: "framework",
clickhouse_column: "Framework",
clickhouse_type: "LowCardinality(String)",
meaning: "Agent framework or SDK that emitted this span, e.g. claude-agent-sdk",
},
NormalizedFieldDefinition {
name: "litellm_request_id",
clickhouse_column: "LiteLLMRequestId",
clickhouse_type: "String",
meaning: "LiteLLM response ID used to link a span to a spend log",
},
NormalizedFieldDefinition {
name: "model",
clickhouse_column: "Model",
clickhouse_type: "LowCardinality(String)",
meaning: "Model used by this span",
},
NormalizedFieldDefinition {
name: "input_tokens",
clickhouse_column: "InputTokens",
clickhouse_type: "UInt32",
meaning: "Input token count",
},
NormalizedFieldDefinition {
name: "output_tokens",
clickhouse_column: "OutputTokens",
clickhouse_type: "UInt32",
meaning: "Output token count",
},
NormalizedFieldDefinition {
name: "input",
clickhouse_column: "Input",
clickhouse_type: "String",
meaning: "Normalized input payload",
},
NormalizedFieldDefinition {
name: "output",
clickhouse_column: "Output",
clickhouse_type: "String",
meaning: "Normalized output payload",
},
];
trait SpanNormalizer {
fn matches(&self, scope_name: &str, attributes: &BTreeMap<String, String>) -> bool;
fn consumed_attributes(&self, attributes: &BTreeMap<String, String>) -> [&'static str; 2];
fn normalize(
&self,
name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
events: &[DecodedEvent],
) -> Result<NormalizedSpan, DecodeError>;
fn display_name(&self, _attributes: &BTreeMap<String, String>) -> Option<String> {
None
}
}
mod claude_code;
mod genai;
mod langsmith;
mod openinference;
use claude_code::ClaudeCodeNormalizer;
pub(crate) use claude_code::{CLAUDE_CODE_AGENT, CLAUDE_CODE_SCOPE};
use genai::GenAiNormalizer;
use langsmith::LangSmithNormalizer;
use openinference::OpenInferenceNormalizer;
fn attr<'a>(attributes: &'a BTreeMap<String, String>, key: &str) -> &'a str {
attributes.get(key).map(String::as_str).unwrap_or_default()
}
fn first<'a>(attributes: &'a BTreeMap<String, String>, left: &str, right: &str) -> &'a str {
let value = attr(attributes, left);
if value.is_empty() {
attr(attributes, right)
} else {
value
}
}
fn tokens(attributes: &BTreeMap<String, String>, key: &str) -> Result<u32, DecodeError> {
let value = attr(attributes, key).trim();
if value.is_empty() {
return Ok(0);
}
match value.parse::<i128>() {
Ok(number) if (0..=u32::MAX as i128).contains(&number) => Ok(number as u32),
Ok(_) => Err(DecodeError::TokenCountOutOfRange),
Err(_)
if value
.trim_start_matches(['+', '-'])
.bytes()
.all(|byte| byte.is_ascii_digit()) =>
{
Err(DecodeError::TokenCountOutOfRange)
}
Err(_) => Ok(0),
}
}
fn usage_tokens(attributes: &BTreeMap<String, String>) -> Result<(u32, u32), DecodeError> {
Ok((
tokens(attributes, "gen_ai.usage.input_tokens")?,
tokens(attributes, "gen_ai.usage.output_tokens")?,
))
}
#[derive(Default, Deserialize)]
struct AgentMetadata {
#[serde(default)]
lc_agent_name: String,
#[serde(default)]
ls_integration: String,
}
fn recorded_agent_name(
name: &str,
attributes: &BTreeMap<String, String>,
span: &NormalizedSpan,
) -> String {
let explicit = [
span.agent_name.as_str(),
attr(attributes, "gen_ai.agent.name"),
attr(attributes, "agent.name"),
attr(attributes, "openclaw.agent"),
]
.into_iter()
.find(|value| !value.is_empty());
if let Some(value) = explicit {
return value.to_owned();
}
let metadata =
serde_json::from_str::<AgentMetadata>(attr(attributes, "metadata")).unwrap_or_default();
if !metadata.lc_agent_name.is_empty() {
return metadata.lc_agent_name;
}
if span.observation_type == ObservationType::Agent {
let node = attr(attributes, "graph.node.id");
if !node.is_empty() {
return node.to_owned();
}
if metadata.ls_integration == "langgraph" && name != "LangGraph" && !is_middleware(name) {
return name.to_owned();
}
}
String::new()
}
fn is_middleware(name: &str) -> bool {
[
".wrap_model_call",
".wrap_tool_call",
".before_agent",
".after_agent",
".before_model",
".after_model",
]
.iter()
.any(|suffix| name.ends_with(suffix))
}
pub fn normalize(
scope_name: &str,
name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
events: &[DecodedEvent],
) -> Result<Normalization, DecodeError> {
let normalizers: [&dyn SpanNormalizer; 4] = [
&ClaudeCodeNormalizer,
&LangSmithNormalizer,
&OpenInferenceNormalizer,
&GenAiNormalizer,
];
let normalizer = normalizers
.into_iter()
.find(|normalizer| normalizer.matches(scope_name, attributes))
.expect("GenAI fallback always matches");
let span = normalizer.normalize(name, parent_span_id, attributes, events)?;
let agent_name = recorded_agent_name(name, attributes, &span);
let observation_type = if !parent_span_id.is_empty()
&& scope_name == "openinference.instrumentation.langchain"
&& is_middleware(name)
{
ObservationType::Framework
} else {
span.observation_type
};
Ok(Normalization {
span: NormalizedSpan {
agent_name,
observation_type,
..span
},
display_name: normalizer.display_name(attributes),
consumed_attributes: normalizer.consumed_attributes(attributes),
})
}
#[cfg(test)]
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use rstest::rstest;
use super::{NORMALIZED_FIELD_DEFINITIONS, ObservationType, normalize};
#[rstest]
#[case::langsmith("langsmith", [("langsmith.span.kind", "llm"), ("openinference.span.kind", "TOOL")], ObservationType::Llm)]
#[case::openinference("other", [("openinference.span.kind", "LLM"), ("gen_ai.operation.name", "execute_tool")], ObservationType::Llm)]
#[case::genai("other", [("gen_ai.operation.name", "execute_tool"), ("gen_ai.usage.input_tokens", "7")], ObservationType::Tool)]
#[case::claude_code("com.anthropic.claude_code.tracing", [("span.type", "llm_request"), ("openinference.span.kind", "TOOL")], ObservationType::Llm)]
fn convention_dispatch_preserves_precedence(
#[case] scope: &str,
#[case] attributes: [(&str, &str); 2],
#[case] expected: ObservationType,
) {
let attributes = attributes
.into_iter()
.map(|(key, value)| (key.to_owned(), value.to_owned()))
.collect();
let fields = normalize(scope, "step", "parent", &attributes, &[])
.expect("valid tokens")
.span;
assert_eq!(fields.observation_type, expected);
if expected == ObservationType::Tool {
assert_eq!(fields.input_tokens, 7);
}
}
#[rstest]
fn field_definitions_match_serialized_normalized_span() {
let fields = normalize("", "root", "", &BTreeMap::new(), &[])
.expect("valid tokens")
.span;
let serialized = serde_json::to_value(fields).expect("serializable fields");
let keys: BTreeSet<_> = serialized
.as_object()
.expect("field object")
.keys()
.map(String::as_str)
.collect();
let mapped: BTreeSet<_> = NORMALIZED_FIELD_DEFINITIONS
.iter()
.map(|field| field.name)
.collect();
assert_eq!(keys, mapped);
}
#[rstest]
fn token_counts_accept_surrounding_whitespace() {
let attributes =
BTreeMap::from([("gen_ai.usage.input_tokens".to_owned(), " 7 ".to_owned())]);
let fields = normalize("", "root", "", &attributes, &[])
.expect("valid tokens")
.span;
assert_eq!(fields.input_tokens, 7);
}
#[rstest]
#[case::negative("-1")]
#[case::overflow("4294967296")]
fn token_counts_outside_storage_range_are_rejected(#[case] value: &str) {
let attributes =
BTreeMap::from([("gen_ai.usage.input_tokens".to_owned(), value.to_owned())]);
assert!(normalize("", "root", "", &attributes, &[]).is_err());
}
}

55
litellm-rust/crates/traces/src/normalize/openinference.rs Normal file
View file

@ -0,0 +1,55 @@
use std::collections::BTreeMap;
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, tokens, usage_tokens};
use crate::{DecodeError, otlp::DecodedEvent};
pub(super) struct OpenInferenceNormalizer;
impl SpanNormalizer for OpenInferenceNormalizer {
fn matches(&self, _scope_name: &str, attributes: &BTreeMap<String, String>) -> bool {
attributes.contains_key("openinference.span.kind")
}
fn consumed_attributes(&self, _attributes: &BTreeMap<String, String>) -> [&'static str; 2] {
["input.value", "output.value"]
}
fn normalize(
&self,
_name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
_events: &[DecodedEvent],
) -> Result<NormalizedSpan, DecodeError> {
let (usage_input, usage_output) = usage_tokens(attributes)?;
let observation_type = match attr(attributes, "openinference.span.kind")
.to_ascii_uppercase()
.as_str()
{
"AGENT" => ObservationType::Agent,
"LLM" => ObservationType::Llm,
"TOOL" => ObservationType::Tool,
_ if parent_span_id.is_empty() => ObservationType::Agent,
_ => ObservationType::Chain,
};
Ok(NormalizedSpan {
observation_type,
agent_name: attr(attributes, "agent.name").to_owned(),
framework: String::new(),
litellm_request_id: String::new(),
model: attr(attributes, "llm.model_name").to_owned(),
input_tokens: if attributes.contains_key("llm.token_count.prompt") {
tokens(attributes, "llm.token_count.prompt")?
} else {
usage_input
},
output_tokens: if attributes.contains_key("llm.token_count.completion") {
tokens(attributes, "llm.token_count.completion")?
} else {
usage_output
},
input: attr(attributes, "input.value").to_owned(),
output: attr(attributes, "output.value").to_owned(),
})
}
}

4
litellm-rust/crates/traces/src/otlp/mod.rs
View file

@ -6,7 +6,7 @@ mod wire;
use serde::Serialize;
use std::collections::BTreeMap;
use crate::{DecodeError, Shared};
use crate::{DecodeError, NormalizedSpan, Shared};
#[derive(Serialize)]
pub struct DecodedEvent {
@ -31,6 +31,8 @@ pub struct DecodedSpan {
pub status_code: String,
pub status_message: String,
pub events: Vec<DecodedEvent>,
pub normalized: NormalizedSpan,
pub consumed_attributes: [&'static str; 2],
}
pub fn decode_otlp(

73
litellm-rust/crates/traces/src/otlp/span.rs
View file

@ -10,7 +10,10 @@ use super::{
attributes::attributes,
limits::{Budget, MAX_ATTRIBUTES, MAX_DECODED_SPAN_BYTES, MAX_EVENTS, MAX_SPANS},
};
use crate::{DecodeError, Shared};
use crate::{
DecodeError, Shared,
normalize::{CLAUDE_CODE_AGENT, CLAUDE_CODE_SCOPE, normalize},
};
pub(super) fn flatten(request: ExportTraceServiceRequest) -> Result<Vec<DecodedSpan>, DecodeError> {
let mut budget = Budget::new(MAX_DECODED_SPAN_BYTES);
@ -125,12 +128,60 @@ fn decoded_span(
budget: &mut Budget,
) -> Result<DecodedSpan, DecodeError> {
let status = span.status.unwrap_or_default();
let parent_span_id = hex_bytes(&span.parent_span_id);
let span_attributes = attributes(span.attributes, budget)?;
let events = span
.events
.into_iter()
.map(|event| {
budget.consume(event.name.len() + 96)?;
Ok(DecodedEvent {
name: event.name,
attributes: attributes(event.attributes, budget)?,
})
})
.collect::<Result<Vec<_>, DecodeError>>()?;
let normalization = normalize(
scope_name.as_ref(),
&span.name,
&parent_span_id,
&span_attributes,
&events,
)?;
let resource_agent_name = resource_attributes
.get("gen_ai.agent.name")
.filter(|name| !name.is_empty());
let agent_name = match (resource_agent_name, normalization.span.agent_name.as_str()) {
(Some(name), "") => name.clone(),
(Some(name), "hermes-agent") if scope_name.as_ref() == "hermes-otel-plugin" => name.clone(),
(Some(name), CLAUDE_CODE_AGENT) if scope_name.as_ref() == CLAUDE_CODE_SCOPE => name.clone(),
(None, CLAUDE_CODE_AGENT) if scope_name.as_ref() == CLAUDE_CODE_SCOPE => {
resource_attributes
.get("service.name")
.filter(|name| !name.is_empty())
.map_or_else(|| CLAUDE_CODE_AGENT.to_owned(), Clone::clone)
}
(_, name) => name.to_owned(),
};
let normalized = crate::normalize::NormalizedSpan {
agent_name,
..normalization.span
};
budget.consume(
normalized.input.len()
+ normalized.output.len()
+ normalized.agent_name.len()
+ normalized.framework.len()
+ normalized.litellm_request_id.len()
+ normalized.model.len()
+ normalization.display_name.as_ref().map_or(0, String::len),
)?;
Ok(DecodedSpan {
trace_id: hex_bytes(&span.trace_id),
span_id: hex_bytes(&span.span_id),
parent_span_id: hex_bytes(&span.parent_span_id),
parent_span_id,
trace_state: span.trace_state,
name: span.name,
name: normalization.display_name.unwrap_or(span.name),
kind: SpanKind::try_from(span.kind)
.unwrap_or(SpanKind::Unspecified)
.as_str_name()
@ -143,7 +194,7 @@ fn decoded_span(
})?,
scope_name: budget.clone_shared(scope_name, String::len)?,
scope_version: budget.clone_shared(scope_version, String::len)?,
attributes: attributes(span.attributes, budget)?,
attributes: span_attributes,
start_ns: span.start_time_unix_nano,
end_ns: span.end_time_unix_nano,
status_code: StatusCode::try_from(status.code)
@ -151,16 +202,8 @@ fn decoded_span(
.as_str_name()
.to_owned(),
status_message: status.message,
events: span
.events
.into_iter()
.map(|event| {
budget.consume(event.name.len() + 96)?;
Ok(DecodedEvent {
name: event.name,
attributes: attributes(event.attributes, budget)?,
})
})
.collect::<Result<Vec<_>, DecodeError>>()?,
events,
normalized,
consumed_attributes: normalization.consumed_attributes,
})
}

346
litellm-rust/crates/traces/src/query.rs Normal file
View file

@ -0,0 +1,346 @@
use std::collections::{BTreeMap, BTreeSet};
use futures_util::{
StreamExt,
stream::{self, TryStreamExt},
};
use litellm_http::Client;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use crate::{Connection, Error, NORMALIZED_FIELD_DEFINITIONS, execute_read};
mod guide;
const SAMPLE_ROWS: usize = 200;
const MAX_FIELDS: usize = 200;
const MAX_DEPTH: usize = 16;
const METADATA_SQL: &str = "SELECT metadata FROM spend_logs FINAL \
WHERE start_time >= now() - INTERVAL 7 DAY AND length(metadata) <= 8192 \
LIMIT 201";
const METADATA_SCOPE: &str = "Up to 200 unordered rows from the last 7 days, excluding metadata larger than 8192 bytes; up to 200 paths and 16 levels. Missing paths may exist outside this sample. Array indexes are 1-based and describe sampled positions, not a fixed schema";
const ATTRIBUTE_SCOPE: &str = "Distinct keys from up to 200 unordered spans in the last 7 days; up to 200 keys per map. Missing keys may exist outside this sample";
#[derive(Deserialize)]
struct Rows<T> {
data: Vec<T>,
}
#[derive(Deserialize)]
struct MetadataRow {
metadata: String,
}
#[derive(Deserialize)]
struct AttributeRow {
key: String,
}
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)]
#[serde(untagged)]
enum PathPart {
Key(String),
Index(usize),
}
#[derive(Serialize)]
struct MetadataField {
path: Vec<PathPart>,
types: BTreeSet<&'static str>,
expression: String,
}
#[derive(Deserialize, Serialize)]
struct ColumnSchema {
name: String,
#[serde(rename = "type")]
kind: String,
#[serde(flatten)]
details: BTreeMap<String, Value>,
}
#[derive(Serialize)]
struct TableSchema {
name: &'static str,
columns: Vec<ColumnSchema>,
}
#[derive(Serialize)]
struct MetadataCatalog {
table: &'static str,
column: &'static str,
fields: Vec<MetadataField>,
sampled_rows: usize,
invalid_json_rows: usize,
truncated: bool,
sample_sql: &'static str,
scope: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
#[derive(Serialize)]
struct AttributeField {
key: String,
#[serde(rename = "type")]
kind: &'static str,
expression: String,
}
#[derive(Serialize)]
struct AttributeCatalog {
table: &'static str,
column: &'static str,
fields: Vec<AttributeField>,
truncated: bool,
discovery_sql: String,
scope: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
pub async fn query_sql(
client: &Client,
connection: &Connection,
sql: &str,
) -> Result<String, Error> {
execute_read(client, connection, sql, &BTreeMap::new()).await
}
async fn rows<T: serde::de::DeserializeOwned>(
client: &Client,
connection: &Connection,
sql: &str,
) -> Result<Vec<T>, Error> {
let body = query_sql(client, connection, sql).await?;
serde_json::from_str::<Rows<T>>(&body)
.map(|result| result.data)
.map_err(|_| Error::InvalidResponse)
}
fn literal(value: &str) -> String {
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
}
fn metadata_expression(path: &[PathPart]) -> String {
let arguments = path
.iter()
.map(|part| match part {
PathPart::Key(key) => literal(key),
PathPart::Index(index) => index.to_string(),
})
.collect::<Vec<_>>()
.join(", ");
format!("JSONExtractRaw(metadata, {arguments})")
}
fn discover(
value: &Value,
path: Vec<PathPart>,
fields: &mut BTreeMap<Vec<PathPart>, BTreeSet<&'static str>>,
) -> bool {
if path.len() > MAX_DEPTH || (fields.len() >= MAX_FIELDS && !fields.contains_key(&path)) {
return true;
}
if !path.is_empty() {
let kind = match value {
Value::Null => "null",
Value::Bool(_) => "boolean",
Value::Number(number) if number.is_i64() || number.is_u64() => "integer",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
};
fields.entry(path.clone()).or_default().insert(kind);
}
match value {
Value::Object(object) => object.iter().fold(false, |limited, (key, value)| {
let child = path
.iter()
.cloned()
.chain([PathPart::Key(key.clone())])
.collect();
discover(value, child, fields) | limited
}),
Value::Array(array) => array
.iter()
.enumerate()
.fold(false, |limited, (index, value)| {
let child = path
.iter()
.cloned()
.chain([PathPart::Index(index + 1)])
.collect();
discover(value, child, fields) | limited
}),
_ => false,
}
}
fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog {
let (fields, limited, invalid_rows) = sample.iter().take(SAMPLE_ROWS).fold(
(BTreeMap::new(), sample.len() > SAMPLE_ROWS, 0),
|(fields, limited, invalid_rows), row| match serde_json::from_str::<Value>(&row.metadata) {
Ok(value) => {
let mut fields = fields;
let limited = limited | discover(&value, Vec::new(), &mut fields);
(fields, limited, invalid_rows)
}
Err(_) => (fields, limited, invalid_rows + 1),
},
);
let fields: Vec<_> = fields
.into_iter()
.map(|(path, types)| MetadataField {
expression: metadata_expression(&path),
path,
types,
})
.collect();
MetadataCatalog {
table: "spend_logs",
column: "metadata",
fields,
sampled_rows: sample.len().min(SAMPLE_ROWS),
invalid_json_rows: invalid_rows,
truncated: limited,
sample_sql: METADATA_SQL,
error: None,
scope: METADATA_SCOPE,
}
}
pub async fn query_help(client: &Client, connection: &Connection) -> Result<String, Error> {
let tables = stream::iter(["otel_traces", "agent_traces_by_key", "spend_logs"])
.then(|table| async move {
Ok::<_, Error>(TableSchema {
name: table,
columns: rows::<ColumnSchema>(
client,
connection,
&format!("DESCRIBE TABLE {table}"),
)
.await?,
})
})
.try_collect::<Vec<_>>()
.await?;
let metadata = match rows::<MetadataRow>(client, connection, METADATA_SQL).await {
Ok(sample) => metadata_catalog(&sample),
Err(error) => MetadataCatalog {
error: Some(error.to_string()),
truncated: true,
..metadata_catalog(&[])
},
};
let attributes = stream::iter(["SpanAttributes", "ResourceAttributes"])
.then(|column| async move {
let sql = format!(
"SELECT DISTINCT arrayJoin(mapKeys({column})) AS key FROM \
(SELECT {column} FROM otel_traces WHERE Timestamp >= now() - INTERVAL 7 DAY \
LIMIT 200) ORDER BY key LIMIT 201"
);
let (keys, error) = match rows::<AttributeRow>(client, connection, &sql).await {
Ok(keys) => (keys, None),
Err(error) => (Vec::new(), Some(error.to_string())),
};
let fields = keys
.iter()
.take(MAX_FIELDS)
.map(|row| AttributeField {
key: row.key.clone(),
kind: "String",
expression: format!("{column}[{}]", literal(&row.key)),
})
.collect();
AttributeCatalog {
table: "otel_traces",
column,
fields,
truncated: error.is_some() || keys.len() > MAX_FIELDS,
discovery_sql: sql,
scope: ATTRIBUTE_SCOPE,
error,
}
})
.collect::<Vec<_>>()
.await;
let guide = guide::QueryGuide {
tables: &tables,
normalized_fields: &NORMALIZED_FIELD_DEFINITIONS,
metadata: &metadata,
attributes: &attributes,
};
Ok(json!({
"dialect": "ClickHouse SQL",
"access": "Authenticated team scope enforced by ClickHouse row policies; proxy admins can read all teams, while project-bound and teamless keys can read only their own rows",
"response": "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings",
"tables": tables,
"normalized_fields": NORMALIZED_FIELD_DEFINITIONS.iter().map(|field| json!({
"table": "otel_traces", "name": field.name, "column": field.clickhouse_column,
"type": field.clickhouse_type, "meaning": field.meaning
})).collect::<Vec<_>>(),
"metadata": metadata,
"attributes": attributes,
"relationships": [{
"left": "otel_traces.LiteLLMRequestId", "right": "spend_logs.response_id",
"additional_predicates": "otel_traces.TeamId = spend_logs.team_id AND otel_traces.ApiKeyHash = spend_logs.api_key",
"meaning": "The normalized ID is the response ID, not request_id. Cached requests can share response_id; joins may return multiple spend rows"
}],
"examples": guide.examples()?,
"gotchas": guide.gotchas()?,
"guide": guide::render(&guide)?,
}).to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[rstest]
fn metadata_discovery_preserves_mixed_types_and_reports_invalid_rows() {
let sample = [
MetadataRow {
metadata: r#"{"x": 1}"#.into(),
},
MetadataRow {
metadata: r#"{"x": "one"}"#.into(),
},
MetadataRow {
metadata: "invalid".into(),
},
];
let catalog = json!(metadata_catalog(&sample));
assert_eq!(
catalog["fields"],
json!([{
"path": ["x"], "types": ["integer", "string"], "expression": "JSONExtractRaw(metadata, 'x')"
}])
);
assert_eq!(catalog["invalid_json_rows"], 1);
assert_eq!(catalog["sampled_rows"], sample.len());
}
#[rstest]
#[case::rows(SAMPLE_ROWS + 1, 1)]
#[case::paths(1, MAX_FIELDS + 1)]
fn metadata_discovery_reports_truncation(#[case] row_count: usize, #[case] field_count: usize) {
let metadata: BTreeMap<_, _> = (0..field_count)
.map(|index| (format!("field{index}"), index))
.collect();
let sample: Vec<_> = (0..row_count)
.map(|_| MetadataRow {
metadata: json!(metadata).to_string(),
})
.collect();
let catalog = json!(metadata_catalog(&sample));
assert_eq!(catalog["truncated"], true);
assert_eq!(catalog["sampled_rows"], row_count.min(SAMPLE_ROWS));
assert_eq!(
catalog["fields"].as_array().unwrap().len(),
field_count.min(MAX_FIELDS)
);
}
}

89
litellm-rust/crates/traces/src/query/guide.rs Normal file
View file

@ -0,0 +1,89 @@
use askama::Template;
use serde::Serialize;
use super::{AttributeCatalog, MetadataCatalog, TableSchema};
use crate::{Error, NormalizedFieldDefinition};
#[derive(Template)]
#[template(path = "query_help.jinja", escape = "none", blocks = [
"recent_spans_name",
"recent_spans_sql",
"custom_metadata_name",
"custom_metadata_sql",
"nested_metadata_name",
"nested_metadata_sql",
"correlated_calls_name",
"correlated_calls_sql",
"discover_keys_name",
"discover_keys_sql",
"time_window",
"reader_limits",
"reader_profile",
"output_format",
"json_values",
"map_values",
"literal_keys",
"time_units",
"spend_totals",
"trace_rollups",
"sampling",
])]
pub(super) struct QueryGuide<'a> {
pub tables: &'a [TableSchema],
pub normalized_fields: &'a [NormalizedFieldDefinition],
pub metadata: &'a MetadataCatalog,
pub attributes: &'a [AttributeCatalog],
}
#[derive(Serialize)]
pub(super) struct Example {
name: String,
sql: String,
}
impl QueryGuide<'_> {
pub fn examples(&self) -> Result<[Example; 5], Error> {
Ok([
Example {
name: render(&self.as_recent_spans_name())?,
sql: render(&self.as_recent_spans_sql())?,
},
Example {
name: render(&self.as_custom_metadata_name())?,
sql: render(&self.as_custom_metadata_sql())?,
},
Example {
name: render(&self.as_nested_metadata_name())?,
sql: render(&self.as_nested_metadata_sql())?,
},
Example {
name: render(&self.as_correlated_calls_name())?,
sql: render(&self.as_correlated_calls_sql())?,
},
Example {
name: render(&self.as_discover_keys_name())?,
sql: render(&self.as_discover_keys_sql())?,
},
])
}
pub fn gotchas(&self) -> Result<[String; 11], Error> {
Ok([
render(&self.as_time_window())?,
render(&self.as_reader_limits())?,
render(&self.as_reader_profile())?,
render(&self.as_output_format())?,
render(&self.as_json_values())?,
render(&self.as_map_values())?,
render(&self.as_literal_keys())?,
render(&self.as_time_units())?,
render(&self.as_spend_totals())?,
render(&self.as_trace_rollups())?,
render(&self.as_sampling())?,
])
}
}
pub(super) fn render(template: &impl Template) -> Result<String, Error> {
template.render().map_err(|_| Error::InvalidResponse)
}

200
litellm-rust/crates/traces/src/query_access.rs Normal file
View file

@ -0,0 +1,200 @@
use std::{sync::Arc, time::Duration};
use hmac::{Hmac, Mac};
use litellm_http::Client;
use moka::future::Cache;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use crate::{Connection, QueryAccessError};
const TABLES: [&str; 3] = ["otel_traces", "agent_traces_by_key", "spend_logs"];
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum QueryScope {
Admin,
Team {
team_id: String,
},
Key {
team_id: String,
api_key_hash: String,
},
}
impl QueryScope {
fn validate(&self) -> Result<(), QueryAccessError> {
match self {
Self::Admin => Ok(()),
Self::Team { team_id } if !team_id.is_empty() => Ok(()),
Self::Key { api_key_hash, .. } if !api_key_hash.is_empty() => Ok(()),
_ => Err(QueryAccessError::InvalidScope),
}
}
fn predicate(&self, table: &str) -> String {
let (team, key) = if table == "spend_logs" {
("team_id", "api_key")
} else {
("TeamId", "ApiKeyHash")
};
match self {
Self::Admin => "1".to_owned(),
Self::Team { team_id } => format!("{team} = {}", literal(team_id)),
Self::Key {
team_id,
api_key_hash,
} => format!(
"{team} = {} AND {key} = {}",
literal(team_id),
literal(api_key_hash)
),
}
}
}
#[derive(Clone)]
pub struct QueryReaders {
writer: Connection,
database: String,
readers: Cache<String, Connection>,
slots: Arc<Semaphore>,
}
impl QueryReaders {
pub fn new(writer: Connection, database: String) -> Self {
Self {
writer,
database,
readers: Cache::builder().max_capacity(1024).build(),
slots: Arc::new(Semaphore::new(8)),
}
}
pub fn acquire(&self) -> Result<OwnedSemaphorePermit, QueryAccessError> {
self.slots
.clone()
.try_acquire_owned()
.map_err(|_| QueryAccessError::Busy)
}
pub async fn connection(
&self,
client: &Client,
scope: &QueryScope,
secret: &str,
) -> Result<Connection, QueryAccessError> {
scope.validate()?;
if secret.is_empty() {
return Err(QueryAccessError::MissingSecret);
}
let identity = serde_json::to_vec(&("litellm_trace_reader_v1", &self.database, scope))
.map_err(|_| QueryAccessError::InvalidScope)?;
let user = format!("litellm_traces_{:x}", Sha256::digest(&identity));
let password = credential(secret, b"password", &identity)?;
self.readers
.try_get_with(
user.clone(),
self.provision(client, scope, &user, &password),
)
.await
.map_err(QueryAccessError::Cached)
}
async fn provision(
&self,
client: &Client,
scope: &QueryScope,
user: &str,
password: &str,
) -> Result<Connection, QueryAccessError> {
let database = &self.database;
if database.is_empty()
|| !database
.bytes()
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
{
return Err(QueryAccessError::InvalidScope);
}
let password_hash = format!("{:x}", Sha256::digest(password));
self.execute(
client,
format!(
"CREATE USER IF NOT EXISTS {user} IDENTIFIED WITH sha256_hash BY '{password_hash}' \
SETTINGS readonly = 1 CONST, max_execution_time = 10 CONST, \
max_result_rows = 1000 CONST, max_result_bytes = 4194304 CONST, \
result_overflow_mode = 'throw' CONST, max_memory_usage = 268435456 CONST, \
max_threads = 2 CONST, max_concurrent_queries_for_user = 8 CONST"
),
)
.await?;
self.execute(
client,
format!("ALTER USER {user} IDENTIFIED WITH sha256_hash BY '{password_hash}'"),
)
.await?;
for table in TABLES {
let predicate = scope.predicate(table);
self.execute(
client,
format!(
"CREATE ROW POLICY IF NOT EXISTS {user}_allow ON `{database}`.{table} \
USING 1 TO {user}"
),
)
.await?;
self.execute(
client,
format!(
"CREATE ROW POLICY IF NOT EXISTS {user}_scope ON `{database}`.{table} \
AS RESTRICTIVE USING {predicate} TO {user}"
),
)
.await?;
}
for table in TABLES {
self.execute(
client,
format!("GRANT SELECT ON `{database}`.{table} TO {user}"),
)
.await?;
}
Connection::configured(
&self.writer.url()[..url::Position::AfterPath],
database,
user,
password,
)
.map_err(QueryAccessError::Storage)
}
async fn execute(&self, client: &Client, sql: String) -> Result<(), QueryAccessError> {
let response = client
.post(self.writer.url().clone())
.timeout(Duration::from_secs(15))
.body(sql)
.send()
.await
.map_err(|_| QueryAccessError::ProvisionTransport)?;
if !response.status().is_success() {
return Err(QueryAccessError::ProvisionFailed(
response.status().as_u16(),
));
}
Ok(())
}
}
fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result<String, QueryAccessError> {
let mut mac = Hmac::<Sha256>::new_from_slice(secret.as_bytes())
.map_err(|_| QueryAccessError::MissingSecret)?;
mac.update(purpose);
mac.update(identity);
Ok(format!("{:x}", mac.finalize().into_bytes()))
}
fn literal(value: &str) -> String {
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
}

45
litellm-rust/crates/traces/src/schema.rs
View file

@ -1,4 +1,5 @@
use litellm_http::Client;
use litellm_migrate::Migration;
use std::time::Duration;
use crate::Connection;
@ -6,42 +7,25 @@ use crate::Error;
const SCHEMA_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const MIGRATIONS: [&str; 9] = [
include_str!("../migrations/0001_otel_traces.sql"),
include_str!("../migrations/0002_agent_traces.sql"),
include_str!("../migrations/0003_agent_traces_mv.sql"),
include_str!("../migrations/0004_spend_logs.sql"),
include_str!("../migrations/0005_otel_traces_ttl.sql"),
include_str!("../migrations/0006_agent_traces_ttl.sql"),
include_str!("../migrations/0007_spend_logs_ttl.sql"),
include_str!("../migrations/0008_trace_received.sql"),
include_str!("../migrations/0009_spend_received.sql"),
];
const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("migrations");
pub fn schema_statements(
database: &str,
trace_retention_days: u32,
spend_log_retention_days: u32,
) -> Result<Vec<String>, Error> {
pub fn schema_statements(database: &str, retention_days: u32) -> Result<Vec<String>, Error> {
if database.is_empty()
|| !database
.bytes()
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
|| trace_retention_days == 0
|| spend_log_retention_days == 0
|| retention_days == 0
{
return Err(Error::InvalidSchema);
}
let database = format!("`{database}`");
Ok(
std::iter::once(format!("CREATE DATABASE IF NOT EXISTS {database}"))
.chain(MIGRATIONS.iter().map(|sql| {
sql.replace("{database}", &database)
.replace("{trace_retention_days}", &trace_retention_days.to_string())
.replace(
"{spend_log_retention_days}",
&spend_log_retention_days.to_string(),
)
.chain(MIGRATIONS.iter().map(|migration| {
migration
.sql
.replace("{database}", &database)
.replace("{retention_days}", &retention_days.to_string())
}))
.collect(),
)
@ -51,15 +35,13 @@ pub async fn ensure_schema(
client: &Client,
connection: &Connection,
database: &str,
trace_retention_days: u32,
spend_log_retention_days: u32,
retention_days: u32,
) -> Result<(), Error> {
ensure_schema_with_timeout(
client,
connection,
database,
trace_retention_days,
spend_log_retention_days,
retention_days,
SCHEMA_REQUEST_TIMEOUT,
)
.await
@ -69,11 +51,10 @@ async fn ensure_schema_with_timeout(
client: &Client,
connection: &Connection,
database: &str,
trace_retention_days: u32,
spend_log_retention_days: u32,
retention_days: u32,
request_timeout: Duration,
) -> Result<(), Error> {
for statement in schema_statements(database, trace_retention_days, spend_log_retention_days)? {
for statement in schema_statements(database, retention_days)? {
let response = client
.post(connection.url().clone())
.timeout(request_timeout)

6
litellm-rust/crates/traces/src/sql.rs
View file

@ -37,6 +37,8 @@ impl ReadQuery {
#[derive(Clone, Copy)]
pub enum LensQuery {
Availability,
Agents,
Sample,
Content,
Evidence,
@ -45,6 +47,8 @@ pub enum LensQuery {
impl LensQuery {
pub fn parse(name: &str) -> Result<Self, Error> {
match name {
"availability" => Ok(Self::Availability),
"agents" => Ok(Self::Agents),
"sample" => Ok(Self::Sample),
"content" => Ok(Self::Content),
"evidence" => Ok(Self::Evidence),
@ -53,6 +57,8 @@ impl LensQuery {
}
pub fn sql(self) -> &'static str {
match self {
Self::Availability => include_str!("../query/lens_availability.sql"),
Self::Agents => include_str!("../query/lens_agents.sql"),
Self::Sample => include_str!("../query/lens_sample.sql"),
Self::Content => include_str!("../query/lens_content.sql"),
Self::Evidence => include_str!("../query/lens_evidence.sql"),

65
litellm-rust/crates/traces/templates/query_help.jinja Normal file
View file

@ -0,0 +1,65 @@
Trace SQL query guide
Live ClickHouse schema
{% for table in tables %}
{{ table.name }}
{% for column in table.columns %}{{ column.name }}: {{ column.kind }}
{% endfor %}{% endfor %}
Normalized span fields
{% for field in normalized_fields %}{{ field.name }}: otel_traces.{{ field.clickhouse_column }} ({{ field.clickhouse_type }})
{{ field.meaning }}
{% endfor %}
Observed LLM call metadata
{{ metadata.scope }}
Sampled rows: {{ metadata.sampled_rows }}; invalid JSON rows: {{ metadata.invalid_json_rows }}; truncated: {{ metadata.truncated }}
{% if let Some(error) = metadata.error %}Metadata discovery unavailable: {{ error }}
{% else if metadata.fields.is_empty() %}No metadata paths found in the sampled rows
{% else %}{% for field in metadata.fields %}{{ field.expression }}: {% for kind in field.types %}{{ kind }} {% endfor %}
{% endfor %}{% endif %}
Observed span and resource attributes
{% for catalog in attributes %}{{ catalog.table }}.{{ catalog.column }}
{{ catalog.scope }}
{% if let Some(error) = catalog.error %}Attribute discovery unavailable: {{ error }}
{% else if catalog.fields.is_empty() %}No attribute keys found in the sampled spans
{% else %}{% for field in catalog.fields %}{{ field.expression }}: {{ field.kind }}
{% endfor %}{% endif %}{% endfor %}
Examples
{% block recent_spans_name %}Recent normalized LLM spans{% endblock %}
{% block recent_spans_sql %}SELECT TraceId, SpanId, Model, InputTokens, OutputTokens, Duration / 1000000 AS duration_ms FROM otel_traces WHERE Timestamp >= now() - INTERVAL 1 DAY AND ObservationType = 'llm' ORDER BY Timestamp DESC LIMIT 100{% endblock %}
{% block custom_metadata_name %}Find calls by custom metadata{% endblock %}
{% block custom_metadata_sql %}SELECT request_id, response_id, model, spend, JSONExtractString(metadata, 'project') AS project FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'project') AND JSONExtractString(metadata, 'project') = 'example' ORDER BY start_time DESC LIMIT 100{% endblock %}
{% block nested_metadata_name %}Nested metadata with unknown types{% endblock %}
{% block nested_metadata_sql %}SELECT request_id, JSONType(metadata, 'labels', 'priority') AS type, JSONExtractRaw(metadata, 'labels', 'priority') AS value FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'labels', 'priority') LIMIT 100{% endblock %}
{% block correlated_calls_name %}Traces correlated with LLM call metadata{% endblock %}
{% block correlated_calls_sql %}SELECT t.TraceId, t.SpanId, s.request_id, s.spend, s.metadata FROM otel_traces AS t INNER JOIN (SELECT * FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY) AS s ON t.LiteLLMRequestId = s.response_id AND t.TeamId = s.team_id AND t.ApiKeyHash = s.api_key WHERE t.Timestamp >= now() - INTERVAL 1 DAY AND t.LiteLLMRequestId != '' AND JSONExtractString(s.metadata, 'project') = 'example' LIMIT 100{% endblock %}
{% block discover_keys_name %}Discover metadata keys over a different window{% endblock %}
{% block discover_keys_sql %}SELECT DISTINCT arrayJoin(JSONExtractKeys(metadata)) AS key FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 30 DAY ORDER BY key LIMIT 200{% endblock %}
Gotchas
{% block time_window %}Always bound Timestamp or start_time and use LIMIT; add TeamId/ApiKeyHash or team_id/api_key filters when investigating one tenant{% endblock %}
{% block reader_limits %}The reader enforces 1000 result rows, 4 MiB response bytes, 256 MiB memory and a 10 second query limit; exceeding limits fails instead of returning partial results{% endblock %}
{% block reader_profile %}LiteLLM provisions SELECT-only readers from the configured ClickHouse connection and enforces authenticated team scope through row policies. Project-bound and teamless keys see only their own rows. Provisioning requires CREATE USER, ALTER USER, CREATE ROW POLICY, and GRANT SELECT permissions{% endblock %}
{% block output_format %}Do not add FORMAT clauses; the endpoint requires ClickHouse JSON output{% endblock %}
{% block json_values %}metadata is a JSON-encoded String; use JSONHas before typed extraction to distinguish missing values from empty strings, zero and false{% endblock %}
{% block map_values %}SpanAttributes and ResourceAttributes are Map(String, String); missing map keys return an empty string, so use mapContains for existence checks{% endblock %}
{% block literal_keys %}Use the discovered path components as separate JSONExtract arguments; a dot inside a key is literal, not a path separator{% endblock %}
{% block time_units %}Duration is nanoseconds; Timestamp has nanosecond precision, spend start_time has millisecond precision{% endblock %}
{% block spend_totals %}Use spend_logs FINAL to collapse replacement rows before totals. Shared response IDs and multiple spans can multiply costs in joins; aggregate spend separately{% endblock %}
{% block trace_rollups %}agent_traces_by_key uses SimpleAggregateFunction columns; group by TeamId, ApiKeyHash and TraceId, using min(StartTs), max(EndTs), sum(SpanCount) and groupUniqArrayArray(Models). Do not use Merge combinators{% endblock %}
{% block sampling %}Discovery is sampled, contains no metadata values, and is not an exhaustive schema. Edit the supplied discovery SQL for older data or nested JSONExtractKeys(metadata, 'parent'){% endblock %}

720
litellm-rust/crates/traces/tests/migrations.rs
View file

@ -2,8 +2,8 @@ use std::{collections::BTreeMap, time::Duration};
use litellm_http::Client;
use litellm_traces::{
Connection, Error, InsertTable, Parameter, ReadQuery, encode_rows, ensure_schema,
execute_named_read, execute_read, schema_statements,
Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, Parameter, ReadQuery,
encode_rows, ensure_schema, execute_named_read, execute_read, schema_statements,
};
use rstest::{fixture, rstest};
use testcontainers_modules::{
@ -109,8 +109,8 @@ async fn schema_supports_span_rollups_and_spend_joins(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let span = serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1", "ParentSpanId": "",
@ -210,6 +210,42 @@ async fn schema_supports_span_rollups_and_spend_joins(
Ok(())
}
#[rstest]
#[tokio::test]
async fn normalized_fields_match_clickhouse_catalog(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
ensure_schema(
&database.client,
&Connection::writer(&database.url)?,
"trace_test",
7,
)
.await?;
let catalog = read_json(&database, "SELECT name, type FROM system.columns WHERE database = 'trace_test' AND table = 'otel_traces'").await?;
let columns: BTreeMap<&str, &str> = catalog["data"]
.as_array()
.expect("catalog rows")
.iter()
.map(|row| {
(
row["name"].as_str().expect("column name"),
row["type"].as_str().expect("column type"),
)
})
.collect();
for field in NORMALIZED_FIELD_DEFINITIONS {
assert_eq!(
columns.get(field.clickhouse_column).copied(),
Some(field.clickhouse_type),
"{}",
field.name
);
}
Ok(())
}
#[rstest]
#[tokio::test]
async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
@ -220,7 +256,7 @@ async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
"{}?input_format_skip_unknown_fields=1",
database.url
))?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let row = BTreeMap::from([
(
"Timestamp".to_owned(),
@ -254,7 +290,7 @@ async fn retried_trace_insert_does_not_inflate_rollup(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let row: BTreeMap<String, serde_json::Value> = serde_json::from_value(serde_json::json!({
"Timestamp": time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64,
"TraceId": "retried-trace", "SpanId": "span-1", "ParentSpanId": "",
@ -289,7 +325,7 @@ async fn keyed_rollup_keeps_same_trace_ids_separate_by_api_key(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let rows = vec![
serde_json::from_value(serde_json::json!({
@ -326,6 +362,192 @@ async fn keyed_rollup_keeps_same_trace_ids_separate_by_api_key(
Ok(())
}
#[rstest]
#[tokio::test]
async fn listed_agent_names_preserve_scope_and_cursor(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
for (team, key, trace, agent, span, parent, framework) in [
(
"alpha",
"one",
"shared",
"research_agent",
"root",
"",
"claude-code",
),
(
"alpha",
"one",
"shared",
"reviewer",
"child",
"root",
"claude-agent-sdk",
),
(
"alpha",
"one",
"shared",
"reviewer",
"repeated",
"root",
"claude-agent-sdk",
),
("alpha", "one", "shared", "", "unnamed", "root", ""),
("alpha", "one", "second", "support_agent", "root", "", ""),
(
"alpha",
"two",
"shared",
"private_agent",
"root",
"",
"private-sdk",
),
(
"beta",
"one",
"shared",
"other_agent",
"root",
"",
"other-sdk",
),
] {
insert_rows(
&database,
"otel_traces",
vec![serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": trace, "SpanId": span, "ParentSpanId": parent,
"ServiceName": "shared-app", "SpanName": span, "AgentName": agent,
"Framework": framework, "ObservationType": "agent",
"ResourceAttributes": {"litellm.team_id": team, "litellm.api_key_hash": key}
}))?],
)
.await?;
}
let historical_rows = (0..5000)
.map(|index| {
serde_json::from_value(serde_json::json!({
"Timestamp": timestamp - 86_400_000_000_000_i64,
"TraceId": "shared", "SpanId": format!("historical-{index}"),
"ParentSpanId": "", "SpanName": "historical", "AgentName": "private_agent",
"ObservationType": "agent", "ServiceName": "shared-app",
"ResourceAttributes": {"litellm.team_id": "alpha", "litellm.api_key_hash": "history"}
}))
})
.collect::<Result<Vec<_>, _>>()?;
insert_rows(&database, "otel_traces", historical_rows).await?;
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
let parameters = BTreeMap::from([
("team_ids".into(), Parameter::Strings(vec!["alpha".into()])),
("api_key_hash".into(), Parameter::Text("one".into())),
(
"start_ms".into(),
Parameter::Integer(timestamp / 1_000_000 - 1000),
),
(
"end_ms".into(),
Parameter::Integer(timestamp / 1_000_000 + 1000),
),
("cursor_ms".into(), Parameter::Integer(0)),
("cursor_trace_id".into(), Parameter::Text(String::new())),
("limit".into(), Parameter::Integer(1)),
]);
let first: serde_json::Value = serde_json::from_str(
&execute_named_read(
&database.client,
&connection,
ReadQuery::ListTraces,
&parameters,
)
.await?,
)?;
let cursor = first["data"][0]["trace_ref"]
.as_str()
.ok_or("missing cursor")?;
let next_parameters = parameters
.into_iter()
.chain([
(
"cursor_ms".into(),
Parameter::Integer(timestamp / 1_000_000),
),
("cursor_trace_id".into(), Parameter::Text(cursor.into())),
])
.collect();
let second: serde_json::Value = serde_json::from_str(
&execute_named_read(
&database.client,
&connection,
ReadQuery::ListTraces,
&next_parameters,
)
.await?,
)?;
assert_eq!(
first["data"].as_array().ok_or("missing first page")?.len(),
1
);
assert_eq!(
second["data"]
.as_array()
.ok_or("missing second page")?
.len(),
1
);
assert_ne!(first["data"][0]["trace_id"], second["data"][0]["trace_id"]);
let names = [&first["data"][0], &second["data"][0]]
.into_iter()
.map(|row| {
(
row["trace_id"].as_str().unwrap(),
row["agent_names"].clone(),
)
})
.collect::<BTreeMap<_, _>>();
assert_eq!(
names["shared"],
serde_json::json!(["research_agent", "reviewer"])
);
assert_eq!(names["second"], serde_json::json!(["support_agent"]));
let frameworks = [&first["data"][0], &second["data"][0]]
.into_iter()
.map(|row| (row["trace_id"].as_str().unwrap(), row["frameworks"].clone()))
.collect::<BTreeMap<_, _>>();
assert_eq!(
frameworks["shared"],
serde_json::json!(["claude-agent-sdk", "claude-code"])
);
assert_eq!(frameworks["second"], serde_json::json!([]));
let counts = [&first["data"][0], &second["data"][0]]
.into_iter()
.map(|row| {
(
row["trace_id"].as_str().unwrap(),
row["agent_count"].as_u64(),
)
})
.collect::<BTreeMap<_, _>>();
assert_eq!(counts["shared"], Some(3));
assert_eq!(counts["second"], Some(1));
for page in [&first, &second] {
assert!(
page["statistics"]["rows_read"]
.as_u64()
.ok_or("missing read statistics")?
< 5000
);
}
Ok(())
}
#[rstest]
#[tokio::test]
async fn rollup_merges_spans_across_days_without_losing_root_fields(
@ -333,13 +555,14 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let day_start = time::OffsetDateTime::now_utc()
.replace_time(time::Time::MIDNIGHT)
.unix_timestamp_nanos() as i64;
let root = serde_json::from_value(serde_json::json!({
"Timestamp": day_start - 1_000_000_000, "TraceId": "cross-day", "SpanId": "span-root",
"ParentSpanId": "", "ServiceName": "proxy", "SpanName": "root", "Input": "root input",
"AgentName": "lead", "ObservationType": "agent",
"StatusCode": "STATUS_CODE_ERROR",
"ResourceAttributes": {"litellm.team_id": "team-1"}
}))?;
@ -347,6 +570,7 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
let child = serde_json::from_value(serde_json::json!({
"Timestamp": day_start + 1_000_000_000, "TraceId": "cross-day", "SpanId": "span-child",
"ParentSpanId": "span-root", "ServiceName": "proxy", "SpanName": "child",
"AgentName": "researcher", "ObservationType": "agent",
"StatusCode": "STATUS_CODE_UNSET",
"ResourceAttributes": {"litellm.team_id": "team-1"}
}))?;
@ -370,6 +594,33 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
"RootStatus": "STATUS_CODE_ERROR", "SpanCount": 2
}])
);
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
let parameters = BTreeMap::from([
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
("api_key_hash".into(), Parameter::Text(String::new())),
(
"start_ms".into(),
Parameter::Integer(day_start / 1_000_000 - 2000),
),
("end_ms".into(), Parameter::Integer(day_start / 1_000_000)),
("cursor_ms".into(), Parameter::Integer(0)),
("cursor_trace_id".into(), Parameter::Text(String::new())),
("limit".into(), Parameter::Integer(10)),
]);
let listed: serde_json::Value = serde_json::from_str(
&execute_named_read(
&database.client,
&connection,
ReadQuery::ListTraces,
&parameters,
)
.await?,
)?;
assert_eq!(
listed["data"][0]["agent_names"],
serde_json::json!(["lead", "researcher"])
);
assert_eq!(listed["data"][0]["agent_count"], 2);
Ok(())
}
@ -380,7 +631,7 @@ async fn spend_deduplication_preserves_subsecond_requests_and_retries(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let now_ms = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64 / 1_000_000;
let base_start_time = now_ms / 1000 * 1000;
let first_start_time = base_start_time + 100;
@ -432,7 +683,21 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 30, 30).await?;
ensure_schema(&database.client, &writer, "trace_test", 30).await?;
let tables = read_json(
&database,
"SELECT name FROM system.tables WHERE database = 'trace_test' \
AND match(engine_full, 'materialize_ttl_recalculate_only = 1') ORDER BY name",
)
.await?;
assert_eq!(
tables["data"],
serde_json::json!([
{"name": "agent_traces_by_key"},
{"name": "otel_traces"},
{"name": "spend_logs"}
])
);
let old_time = time::OffsetDateTime::now_utc() - time::Duration::days(20);
let old_timestamp_ns = old_time.unix_timestamp_nanos() as i64;
let old_timestamp_ms = old_timestamp_ns / 1_000_000;
@ -448,7 +713,7 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
insert_rows(&database, "otel_traces", vec![span]).await?;
insert_rows(&database, "spend_logs", vec![spend]).await?;
assert_eq!(table_rows(&database, "agent_traces_by_key").await?, 1);
ensure_schema(&database.client, &writer, "trace_test", 14, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 14).await?;
let deadline = tokio::time::Instant::now() + Duration::from_secs(60);
loop {
let response = read_json(
@ -480,7 +745,7 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
assert_eq!(table_rows(&database, "agent_traces_by_key").await?, 0);
assert_eq!(table_rows(&database, "spend_logs").await?, 0);
let mutation_count = mutation_rows(&database).await?;
ensure_schema(&database.client, &writer, "trace_test", 14, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 14).await?;
assert_eq!(mutation_rows(&database).await?, mutation_count);
Ok(())
}
@ -499,7 +764,7 @@ async fn schema_statement_timeout_maps_to_transport_error() -> TestResult {
let writer = Connection::writer(&url)?;
let result = tokio::time::timeout(
Duration::from_secs(35),
ensure_schema(&client, &writer, "trace_test", 7, 14),
ensure_schema(&client, &writer, "trace_test", 7),
)
.await;
server.abort();
@ -508,16 +773,11 @@ async fn schema_statement_timeout_maps_to_transport_error() -> TestResult {
}
#[rstest]
#[case::empty("", 7, 14)]
#[case::sql("db; DROP DATABASE default", 7, 14)]
#[case::trace_retention("traces", 0, 14)]
#[case::spend_retention("traces", 7, 0)]
fn schema_rejects_invalid_configuration(
#[case] database: &str,
#[case] traces: u32,
#[case] spend: u32,
) {
assert!(schema_statements(database, traces, spend).is_err());
#[case::empty("", 7)]
#[case::sql("db; DROP DATABASE default", 7)]
#[case::retention("traces", 0)]
fn schema_rejects_invalid_configuration(#[case] database: &str, #[case] retention_days: u32) {
assert!(schema_statements(database, retention_days).is_err());
}
#[rstest]
@ -528,7 +788,7 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
use litellm_traces::{LensQuery, Parameter};
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
for (key, text) in [("one", "timeout"), ("two", "success")] {
insert_rows(&database, "otel_traces", vec![serde_json::from_value(serde_json::json!({
@ -551,6 +811,7 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
"end".into(),
Parameter::Integer(timestamp / 1_000_000 + 1000),
),
("agent_name".into(), Parameter::Text(String::new())),
("service".into(), Parameter::Text("review".into())),
(
"filter_keys".into(),
@ -634,7 +895,7 @@ async fn lens_request_sample_does_not_trust_caller_tags(
use litellm_traces::{LensQuery, Parameter};
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64 / 1_000_000;
for (id, internal) in [("external", false), ("internal", true)] {
let row = serde_json::from_value(serde_json::json!({
@ -652,6 +913,7 @@ async fn lens_request_sample_does_not_trust_caller_tags(
("key_hash".into(), Parameter::Text(String::new())),
("start".into(), Parameter::Integer(timestamp - 1000)),
("end".into(), Parameter::Integer(timestamp + 60000)),
("agent_name".into(), Parameter::Text(String::new())),
("service".into(), Parameter::Text(String::new())),
("filter_keys".into(), Parameter::Strings(vec![])),
("filter_values".into(), Parameter::Strings(vec![])),
@ -702,7 +964,6 @@ async fn lens_selection_pages_without_losing_or_repeating_runs(
&Connection::writer(&database.url)?,
"trace_test",
7,
14,
)
.await?;
execute_write(&database, "INSERT INTO trace_test.spend_logs (request_id,team_id,start_time,end_time) SELECT toString(number),'team',now64(3)-INTERVAL 5 MINUTE,now64(3)-INTERVAL 5 MINUTE FROM numbers(1001)").await?;
@ -719,6 +980,7 @@ async fn lens_selection_pages_without_losing_or_repeating_runs(
("key_hash".into(), Parameter::Text(String::new())),
("start".into(), Parameter::Integer(0)),
("end".into(), Parameter::Integer(end)),
("agent_name".into(), Parameter::Text(String::new())),
("service".into(), Parameter::Text(String::new())),
("filter_keys".into(), Parameter::Strings(vec![])),
("filter_values".into(), Parameter::Strings(vec![])),
@ -782,7 +1044,6 @@ async fn lens_content_keeps_output_visible_after_long_input(
&Connection::writer(&database.url)?,
"trace_test",
7,
14,
)
.await?;
insert_rows(&database, "spend_logs", vec![serde_json::from_value(serde_json::json!({
@ -848,7 +1109,7 @@ async fn trace_error_previews_preserve_paginated_diagnostics(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let rows = (0..span_count)
.map(|index| {
@ -938,7 +1199,7 @@ async fn duplicate_span_preview_matches_diagnostic(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let message = "a".repeat(200);
let rows = [
@ -980,3 +1241,404 @@ async fn duplicate_span_preview_matches_diagnostic(
assert_eq!(diagnostic["data"][0]["message"], message);
Ok(())
}
#[rstest]
fn schema_includes_every_migration_file() -> TestResult {
let files = std::fs::read_dir(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations"))?
.filter_map(|entry| entry.ok())
.filter(|entry| entry.path().extension().is_some_and(|ext| ext == "sql"))
.count();
assert_eq!(schema_statements("trace_test", 7)?.len(), 1 + files);
Ok(())
}
#[rstest]
#[tokio::test]
async fn lens_agent_discovery_and_selection_preserve_scope(
#[future] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
use litellm_traces::LensQuery;
let database = database.await?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
for (team, key, trace, agent, span, parent) in [
("alpha", "one", "research", "research_agent", "root", ""),
("alpha", "one", "research", "", "tool", "root"),
("alpha", "one", "support", "support_agent", "root", ""),
("alpha", "two", "hidden-key", "private_agent", "root", ""),
("beta", "one", "hidden-team", "other_agent", "root", ""),
] {
insert_rows(
&database,
"otel_traces",
vec![serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": trace, "SpanId": span, "ParentSpanId": parent,
"ServiceName": "shared-app", "SpanName": "run", "Input": "test",
"SpanAttributes": {"gen_ai.agent.name": agent},
"ResourceAttributes": {"litellm.team_id": team, "litellm.api_key_hash": key}
}))?],
)
.await?;
}
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
let scope_parameters = BTreeMap::from([
("all_teams".into(), Parameter::Integer(0)),
("team".into(), Parameter::Text("alpha".into())),
("key_hash".into(), Parameter::Text("one".into())),
]);
let agents: serde_json::Value = serde_json::from_str(
&execute_read(
&database.client,
&connection,
LensQuery::Agents.sql(),
&scope_parameters,
)
.await?,
)?;
assert_eq!(
agents["data"],
serde_json::json!([
{"agent_name": "research_agent"}, {"agent_name": "support_agent"}
])
);
let parameters = scope_parameters
.into_iter()
.chain([
("source".into(), Parameter::Text("traces".into())),
(
"start".into(),
Parameter::Integer(timestamp / 1_000_000 - 1000),
),
(
"end".into(),
Parameter::Integer(timestamp / 1_000_000 + 1000),
),
("service".into(), Parameter::Text("shared-app".into())),
(
"agent_name".into(),
Parameter::Text("research_agent".into()),
),
("filter_keys".into(), Parameter::Strings(vec![])),
("filter_values".into(), Parameter::Strings(vec![])),
("limit".into(), Parameter::Integer(100)),
("offset".into(), Parameter::Integer(0)),
("after".into(), Parameter::Text(String::new())),
("sample_percent".into(), Parameter::Text("100".into())),
("sample_cap".into(), Parameter::Integer(0)),
("preview".into(), Parameter::Integer(1)),
("selected_team".into(), Parameter::Text(String::new())),
("execution_ids".into(), Parameter::Strings(vec![])),
])
.collect::<BTreeMap<_, _>>();
let sample: serde_json::Value = serde_json::from_str(
&execute_read(
&database.client,
&connection,
LensQuery::Sample.sql(),
&parameters,
)
.await?,
)?;
assert_eq!(sample["data"].as_array().expect("rows").len(), 1);
assert_eq!(sample["data"][0]["trace_id"], "research");
assert_eq!(sample["data"][0]["span_count"], 2);
let available: serde_json::Value = serde_json::from_str(
&execute_read(
&database.client,
&connection,
LensQuery::Availability.sql(),
&parameters,
)
.await?,
)?;
assert_eq!(available["data"][0]["traces"], 1);
assert_eq!(available["data"][0]["requests"], 0);
Ok(())
}
#[rstest]
#[case::empty(false)]
#[case::custom_metadata(true)]
#[tokio::test]
async fn query_help_discovers_live_schema_and_runs_its_examples(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
#[case] populated: bool,
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
execute_write(&database, "CREATE USER help_reader").await?;
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
execute_write(
&database,
&format!("GRANT SELECT ON trace_test.{table} TO help_reader"),
)
.await?;
}
let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
if populated {
execute_write(&database, "SYSTEM STOP MERGES trace_test.spend_logs").await?;
insert_rows(
&database,
"spend_logs",
vec![serde_json::from_value(serde_json::json!({
"request_id": "request-1", "response_id": "response-1", "team_id": "team-1",
"api_key": "key-1", "metadata": r#"{"obsolete":true,"labels":{"priority":"old"}}"#,
"start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000
}))?],
)
.await?;
let metadata = serde_json::json!({
"project": "example", "labels": {"priority": 3, "enabled": true},
"dotted.key": "literal", "quote'\\key": null, "items": [{"name": "first"}],
"<custom>&{{key}}": {"nested.key": true}
});
insert_rows(
&database,
"spend_logs",
vec![serde_json::from_value(serde_json::json!({
"request_id": "request-1", "response_id": "response-1", "team_id": "team-1",
"api_key": "key-1", "metadata": metadata.to_string(), "spend": 0.25,
"start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000 + 100
}))?],
)
.await?;
insert_rows(
&database,
"otel_traces",
vec![serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1",
"TeamId": "team-1", "ApiKeyHash": "key-1", "ObservationType": "llm",
"LiteLLMRequestId": "response-1", "SpanAttributes": {"custom.tag": "value"},
"ResourceAttributes": {"custom.resource": "value"}
}))?],
)
.await?;
execute_write(
&database,
"ALTER TABLE trace_test.otel_traces ADD COLUMN CustomColumn String",
)
.await?;
}
let help: serde_json::Value =
serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?;
let keys: std::collections::BTreeSet<_> = help
.as_object()
.ok_or("missing help object")?
.keys()
.map(String::as_str)
.collect();
assert_eq!(
keys,
std::collections::BTreeSet::from([
"access",
"attributes",
"dialect",
"examples",
"gotchas",
"guide",
"metadata",
"normalized_fields",
"relationships",
"response",
"tables",
])
);
let guide = help["guide"].as_str().ok_or("missing rendered guide")?;
assert!(guide.starts_with("Trace SQL query guide"));
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
let described = read_json(&database, &format!("DESCRIBE TABLE {table}")).await?;
let schema = help["tables"]
.as_array()
.ok_or("missing tables")?
.iter()
.find(|schema| schema["name"] == table)
.ok_or("missing table")?;
assert_eq!(schema["columns"], described["data"]);
for column in described["data"].as_array().ok_or("missing live columns")? {
assert!(guide.contains(&format!(
"{}: {}",
column["name"].as_str().ok_or("column name")?,
column["type"].as_str().ok_or("column type")?
)));
}
}
for gotcha in help["gotchas"].as_array().ok_or("missing gotchas")? {
assert!(guide.contains(gotcha.as_str().ok_or("gotcha text")?));
}
let tables = help["tables"].as_array().ok_or("missing tables")?;
assert_eq!(tables.len(), 3);
let columns = tables[0]["columns"].as_array().ok_or("missing columns")?;
for field in NORMALIZED_FIELD_DEFINITIONS {
assert!(
columns
.iter()
.any(|column| column["name"] == field.clickhouse_column
&& column["type"] == field.clickhouse_type)
);
assert!(
help["normalized_fields"]
.as_array()
.ok_or("missing mappings")?
.iter()
.any(|mapped| {
mapped["name"] == field.name && mapped["column"] == field.clickhouse_column
})
);
}
let fields = help["metadata"]["fields"]
.as_array()
.ok_or("missing metadata fields")?;
assert_eq!(fields.is_empty(), !populated);
assert_eq!(help["metadata"]["truncated"], false);
assert!(guide.contains(help["metadata"]["scope"].as_str().ok_or("missing scope")?));
assert_eq!(
guide.contains("No metadata paths found in the sampled rows"),
!populated
);
if populated {
let versions = read_json(&database, "SELECT count() AS count FROM spend_logs").await?;
assert_eq!(versions["data"][0]["count"], 2);
assert_eq!(help["metadata"]["sampled_rows"], 1);
assert!(
!fields
.iter()
.any(|field| field["path"] == serde_json::json!(["obsolete"]))
);
assert!(
columns
.iter()
.any(|column| column["name"] == "CustomColumn")
);
assert!(fields.iter().any(|field| field["path"]
== serde_json::json!(["labels", "priority"])
&& field["types"] == serde_json::json!(["integer"])));
assert!(
fields
.iter()
.any(|field| field["path"] == serde_json::json!(["items", 1, "name"]))
);
assert!(guide.contains("CustomColumn: String"));
assert!(guide.contains("JSONExtractRaw(metadata, '<custom>&{{key}}', 'nested.key')"));
assert!(guide.contains("SpanAttributes['custom.tag']"));
assert!(guide.contains("ResourceAttributes['custom.resource']"));
assert_eq!(help["attributes"][0]["fields"][0]["key"], "custom.tag");
assert_eq!(help["attributes"][1]["fields"][0]["key"], "custom.resource");
for field in fields {
let expression = field["expression"].as_str().ok_or("missing expression")?;
assert!(
guide.contains(expression),
"missing plain-text expression: {expression}"
);
let sql = format!("SELECT {expression} AS value FROM spend_logs FINAL");
let body = litellm_traces::query_sql(&database.client, &reader, &sql).await?;
let values: serde_json::Value = serde_json::from_str(&body)?;
assert_ne!(values["data"][0]["value"], "");
}
}
for example in help["examples"].as_array().ok_or("missing examples")? {
let sql = example["sql"].as_str().ok_or("missing example SQL")?;
assert!(guide.contains(example["name"].as_str().ok_or("missing example name")?));
assert!(guide.contains(sql));
assert_eq!(
example
.as_object()
.ok_or("example object")?
.keys()
.map(String::as_str)
.collect::<std::collections::BTreeSet<_>>(),
std::collections::BTreeSet::from(["name", "sql"])
);
let body = litellm_traces::query_sql(&database.client, &reader, sql).await?;
let values: serde_json::Value = serde_json::from_str(&body)?;
assert_eq!(
values["data"].as_array().ok_or("missing data")?.is_empty(),
!populated,
"{sql}"
);
if populated && example["name"] == "Traces correlated with LLM call metadata" {
assert_eq!(values["data"][0]["TraceId"], "trace-1");
assert_eq!(values["data"][0]["spend"], 0.25);
}
}
Ok(())
}
#[rstest]
#[case::metadata(2, 1)]
#[case::attributes(1, 2)]
#[case::all(2, 2)]
#[tokio::test]
async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
#[case] spend_rows: usize,
#[case] span_rows: usize,
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
execute_write(
&database,
"CREATE USER help_reader SETTINGS max_rows_to_read = 1",
)
.await?;
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
execute_write(
&database,
&format!("GRANT SELECT ON trace_test.{table} TO help_reader"),
)
.await?;
}
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let spend = (0..spend_rows)
.map(|index| {
serde_json::from_value(serde_json::json!({
"request_id": format!("request-{index}"), "start_time": timestamp / 1_000_000,
"end_time": timestamp / 1_000_000, "metadata": r#"{"custom":{"enabled":true}}"#
}))
})
.collect::<Result<Vec<_>, _>>()?;
insert_rows(&database, "spend_logs", spend).await?;
let spans = (0..span_rows).map(|index| serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": "trace", "SpanId": format!("span-{index}"),
"SpanAttributes": {"custom.span": "value"}, "ResourceAttributes": {"custom.resource": "value"}
}))).collect::<Result<Vec<_>, _>>()?;
insert_rows(&database, "otel_traces", spans).await?;
let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?;
let help: serde_json::Value =
serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?;
assert_eq!(help["tables"].as_array().ok_or("tables")?.len(), 3);
assert!(!help["examples"].as_array().ok_or("examples")?.is_empty());
assert_eq!(
help["normalized_fields"]
.as_array()
.ok_or("normalized fields")?
.len(),
NORMALIZED_FIELD_DEFINITIONS.len()
);
let guide = help["guide"].as_str().ok_or("guide")?;
assert!(guide.contains("TraceId: String"));
assert_eq!(
guide.contains("Metadata discovery unavailable:"),
spend_rows > 1
);
assert_eq!(
guide.contains("Attribute discovery unavailable:"),
span_rows > 1
);
for (catalog, unavailable) in [
(&help["metadata"], spend_rows > 1),
(&help["attributes"][0], span_rows > 1),
(&help["attributes"][1], span_rows > 1),
] {
assert_eq!(catalog.get("error").is_some(), unavailable);
assert_eq!(catalog["truncated"], unavailable);
assert_eq!(
catalog["fields"].as_array().ok_or("fields")?.is_empty(),
unavailable
);
}
Ok(())
}

267
litellm-rust/crates/traces/tests/otlp.rs
View file

@ -1,5 +1,5 @@
use litellm_traces::Shared;
use litellm_traces::decode_otlp;
use litellm_traces::{ObservationType, Shared};
use rstest::rstest;
const FIXTURE: &[u8] = include_bytes!(
@ -341,3 +341,268 @@ fn escaped_attribute_expansion_is_bounded_below_four_mib(
Err(litellm_traces::DecodeError::TooLarge)
));
}
#[rstest]
fn normalizes_langsmith_fixture() {
let spans = decode_otlp(FIXTURE, Some("application/json")).expect("valid OTLP export");
let llm = spans
.iter()
.find(|span| span.name == "ChatOpenAI")
.expect("LLM span");
assert_eq!(llm.normalized.observation_type, ObservationType::Llm);
assert_eq!(llm.normalized.agent_name, "deep_research_agent");
assert_eq!(llm.normalized.model, "claude-sonnet-4-5");
assert_eq!(
(llm.normalized.input_tokens, llm.normalized.output_tokens),
(3332, 467)
);
assert_eq!(
llm.normalized.litellm_request_id,
"chatcmpl-4077bb36-9380-4a3b-9481-245700cef09a"
);
let input: serde_json::Value =
serde_json::from_str(&llm.normalized.input).expect("message input");
assert_eq!(input[0]["role"], "system");
assert_eq!(input[1]["role"], "user");
let output: serde_json::Value =
serde_json::from_str(&llm.normalized.output).expect("message output");
assert_eq!(output["role"], "assistant");
assert!(output["tool_calls"][0]["name"].is_string());
assert!(output["tool_calls"][0]["id"].is_string());
assert_eq!(output["tool_calls"][0]["type"], "tool_call");
let root = spans
.iter()
.find(|span| span.name == "deep_research_agent")
.expect("root span");
assert_eq!(root.normalized.observation_type, ObservationType::Agent);
assert_eq!(
root.normalized.input,
"[{\"role\": \"user\", \"content\": \"Should we store OTEL agent spans in ClickHouse or Postgres at 50k spans/sec?\"}]"
);
let tool = spans
.iter()
.find(|span| span.name == "task")
.expect("tool span");
assert_eq!(tool.normalized.observation_type, ObservationType::Tool);
assert!(tool.normalized.output.starts_with("Based on my research"));
}
const CLAUDE_AGENT_SDK_FIXTURE: &[u8] =
include_bytes!("../../../../tests/test_litellm/tracing/fixtures/claude_agent_sdk_export.json");
const CLAUDE_AGENT_SDK_DETAILED_FIXTURE: &[u8] = include_bytes!(
"../../../../tests/test_litellm/tracing/fixtures/claude_agent_sdk_detailed_export.json"
);
fn raw_spans(fixture: &[u8]) -> Vec<serde_json::Value> {
let export: serde_json::Value = serde_json::from_slice(fixture).expect("fixture JSON");
export["resourceSpans"][0]["scopeSpans"][0]["spans"]
.as_array()
.expect("spans")
.clone()
}
fn raw_attribute(span: &serde_json::Value, key: &str) -> Option<serde_json::Value> {
span["attributes"]
.as_array()
.expect("attributes")
.iter()
.find(|attribute| attribute["key"] == key)
.map(|attribute| attribute["value"].clone())
}
fn raw_string(span: &serde_json::Value, key: &str) -> String {
raw_attribute(span, key)
.and_then(|value| value["stringValue"].as_str().map(str::to_owned))
.unwrap_or_default()
}
fn raw_int(span: &serde_json::Value, key: &str) -> u64 {
raw_attribute(span, key).map_or(0, |value| match &value["intValue"] {
serde_json::Value::String(text) => text.parse().expect("integer"),
number => number.as_u64().expect("integer"),
})
}
fn raw_span<'a>(raw: &'a [serde_json::Value], span_id: &str) -> &'a serde_json::Value {
raw.iter()
.find(|span| {
span["spanId"]
.as_str()
.is_some_and(|id| id.eq_ignore_ascii_case(span_id))
})
.expect("raw span")
}
#[rstest]
#[case::default_telemetry(CLAUDE_AGENT_SDK_FIXTURE)]
#[case::detailed_telemetry(CLAUDE_AGENT_SDK_DETAILED_FIXTURE)]
fn normalizes_claude_agent_sdk_fixture(#[case] fixture: &[u8]) {
let spans = decode_otlp(fixture, Some("application/json")).expect("valid OTLP export");
let raw = raw_spans(fixture);
let types: std::collections::BTreeSet<_> = spans
.iter()
.map(|span| format!("{:?}", span.normalized.observation_type))
.collect();
assert_eq!(
types,
["Agent", "Framework", "Llm", "Tool"]
.into_iter()
.map(str::to_owned)
.collect()
);
let root = spans
.iter()
.find(|span| span.normalized.observation_type == ObservationType::Agent)
.expect("interaction root");
assert!(root.parent_span_id.is_empty());
let root_input: serde_json::Value =
serde_json::from_str(&root.normalized.input).expect("root input messages");
assert_eq!(root_input[0]["role"], "user");
assert_eq!(
root_input[0]["content"],
raw_string(raw_span(&raw, &root.span_id), "user_prompt")
);
assert!(root.consumed_attributes.contains(&"user_prompt"));
let tools: Vec<_> = spans
.iter()
.filter(|span| span.normalized.observation_type == ObservationType::Tool)
.collect();
assert_eq!(tools.len(), 2);
for tool in &tools {
assert_eq!(
tool.name,
raw_string(raw_span(&raw, &tool.span_id), "tool_name")
);
let input: serde_json::Value =
serde_json::from_str(&tool.normalized.input).expect("tool argument object");
assert!(input.is_object());
assert!(input.get("role").is_none());
let event = tool
.events
.iter()
.find(|event| event.name == "tool.output")
.expect("tool output event");
let expected_output = ["output", "content", "diff"]
.into_iter()
.filter_map(|key| event.attributes.get(key))
.find(|value| !value.is_empty())
.expect("event output");
assert_eq!(&tool.normalized.output, expected_output);
}
let bash = tools
.iter()
.find(|tool| tool.name == "Bash")
.expect("Bash tool");
assert_eq!(
serde_json::from_str::<serde_json::Value>(&bash.normalized.input).unwrap()["command"],
raw_string(raw_span(&raw, &bash.span_id), "full_command")
);
let llms: Vec<_> = spans
.iter()
.filter(|span| span.normalized.observation_type == ObservationType::Llm)
.collect();
assert!(!llms.is_empty());
for llm in &llms {
let raw_llm = raw_span(&raw, &llm.span_id);
let expected = raw_int(raw_llm, "input_tokens")
+ raw_int(raw_llm, "cache_read_tokens")
+ raw_int(raw_llm, "cache_creation_tokens");
assert_eq!(u64::from(llm.normalized.input_tokens), expected);
assert_eq!(
u64::from(llm.normalized.output_tokens),
raw_int(raw_llm, "output_tokens")
);
assert_eq!(llm.normalized.model, raw_string(raw_llm, "model"));
if raw_string(raw_llm, "query_source_safe") == "sdk" {
assert_eq!(llm.normalized.framework, "claude-agent-sdk");
}
}
assert!(spans.iter().all(|span| {
span.normalized.agent_name == span.resource_attributes["service.name"].as_str()
}));
}
#[rstest]
fn claude_agent_sdk_detailed_fixture_keeps_full_tool_arguments_and_llm_messages() {
let spans = decode_otlp(CLAUDE_AGENT_SDK_DETAILED_FIXTURE, Some("application/json"))
.expect("valid OTLP export");
let raw = raw_spans(CLAUDE_AGENT_SDK_DETAILED_FIXTURE);
let bash = spans
.iter()
.find(|span| span.name == "Bash")
.expect("Bash tool");
let tool_input = raw_string(raw_span(&raw, &bash.span_id), "tool_input");
let (_, arguments) = tool_input.split_once('\n').expect("tool input header");
assert_eq!(
serde_json::from_str::<serde_json::Value>(&bash.normalized.input).unwrap(),
serde_json::from_str::<serde_json::Value>(arguments).unwrap()
);
assert!(bash.consumed_attributes.contains(&"tool_input"));
let answer = spans
.iter()
.find(|span| {
span.normalized.observation_type == ObservationType::Llm
&& span.attributes.get("query_source_safe").map(String::as_str) == Some("sdk")
&& !span.normalized.output.is_empty()
})
.expect("final SDK answer");
let raw_answer = raw_span(&raw, &answer.span_id);
let input: serde_json::Value =
serde_json::from_str(&answer.normalized.input).expect("llm input messages");
assert_eq!(input[0]["role"], "system");
assert_eq!(
input[0]["content"],
raw_string(raw_answer, "system_prompt_preview")
);
let output: serde_json::Value =
serde_json::from_str(&answer.normalized.output).expect("llm output message");
assert_eq!(output["role"], "assistant");
assert_eq!(
output["content"],
raw_string(raw_answer, "response.model_output")
);
let title = spans
.iter()
.find(|span| {
span.attributes.get("query_source_safe").map(String::as_str)
== Some("generate_session_title")
})
.expect("side query");
assert_eq!(title.normalized.framework, "claude-agent-sdk");
}
#[rstest]
fn claude_code_scope_takes_precedence_over_openinference_attributes(
mut span: opentelemetry_proto::tonic::trace::v1::Span,
) {
use opentelemetry_proto::tonic::common::v1::{
AnyValue, InstrumentationScope, KeyValue, any_value::Value,
};
let string = |key: &str, value: &str| KeyValue {
key: key.to_owned(),
value: Some(AnyValue {
value: Some(Value::StringValue(value.to_owned())),
}),
..Default::default()
};
span.attributes = vec![
string("span.type", "tool"),
string("tool_name", "Grep"),
string("openinference.span.kind", "LLM"),
];
let mut request = request_with(span);
request.resource_spans[0].scope_spans[0].scope = Some(InstrumentationScope {
name: "com.anthropic.claude_code.tracing".to_owned(),
..Default::default()
});
let spans = decode_otlp(&prost::Message::encode_to_vec(&request), None).expect("valid span");
assert_eq!(spans[0].normalized.observation_type, ObservationType::Tool);
assert_eq!(spans[0].name, "Grep");
assert_eq!(spans[0].normalized.framework, "claude-code");
assert_eq!(spans[0].normalized.agent_name, "claude-code");
}

268
litellm-rust/crates/traces/tests/query_access.rs Normal file
View file

@ -0,0 +1,268 @@
use std::collections::BTreeMap;
use litellm_http::Client;
use litellm_traces::{
Connection, Error, QueryReaders, QueryScope, ensure_schema, query_help, query_sql,
};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use testcontainers_modules::{
clickhouse::ClickHouse,
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
};
struct Database {
_container: ContainerAsync<ClickHouse>,
client: Client,
writer: Connection,
readers: QueryReaders,
}
#[fixture]
async fn database() -> Result<Database, Box<dyn std::error::Error>> {
let container = ClickHouse::default()
.with_tag(
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e",
)
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
.start()
.await?;
let writer = Connection::parse(&format!(
"http://{}:{}",
container.get_host().await?,
container.get_host_port_ipv4(8123).await?
))?;
let client = Client::no_redirect_for_test();
ensure_schema(&client, &writer, "trace_test", 7).await?;
for sql in [
"INSERT INTO trace_test.otel_traces (TeamId, ApiKeyHash, TraceId, SpanId, Timestamp, SpanAttributes) VALUES ('team-a', 'key-a1', 'shared-trace', 'a1', now(), map('visible', 'a')), ('team-a', 'key-a2', 'shared-trace', 'a2', now(), map('visible', 'a')), ('team-b', 'key-b', 'shared-trace', 'b', now(), map('secret-b', 'b'))",
"INSERT INTO trace_test.spend_logs (team_id, api_key, request_id, start_time, end_time, metadata) VALUES ('team-a', 'key-a1', 'a1', now(), now(), '{\"visible\":1}'), ('team-a', 'key-a2', 'a2', now(), now(), '{\"visible\":1}'), ('team-b', 'key-b', 'b', now(), now(), '{\"secret_b\":1}')",
"CREATE TABLE trace_test.private_data (secret String) ENGINE = Memory",
"INSERT INTO trace_test.private_data VALUES ('hidden')",
] {
let response = client.post(writer.url().clone()).body(sql).send().await?;
assert!(response.status().is_success(), "{}", response.text().await?);
}
let readers = QueryReaders::new(writer.clone(), "trace_test".to_owned());
Ok(Database {
_container: container,
client,
writer,
readers,
})
}
#[rstest]
#[case::team(QueryScope::Team { team_id: "team-a".to_owned() }, vec!["a1", "a2"])]
#[case::project_key(QueryScope::Key { team_id: "team-a".to_owned(), api_key_hash: "key-a1".to_owned() }, vec!["a1"])]
#[case::admin(QueryScope::Admin, vec!["a1", "a2", "b"])]
#[case::quoted_team(QueryScope::Team { team_id: "team-a' OR 1=1 --\\".to_owned() }, vec![])]
#[tokio::test]
async fn queries_and_help_are_scoped_by_the_database(
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
#[case] scope: QueryScope,
#[case] expected: Vec<&str>,
) -> Result<(), Box<dyn std::error::Error>> {
let database = database?;
let reader = database
.readers
.connection(&database.client, &scope, "test-master-secret")
.await?;
let queries = [
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
"SELECT SpanId AS id FROM trace_test.otel_traces WHERE 1 = 1 ORDER BY id",
"SELECT SpanId AS id FROM merge('trace_test', '^otel_traces$') ORDER BY id",
"WITH source AS (SELECT * FROM trace_test.otel_traces) SELECT SpanId AS id FROM source ORDER BY id",
"SELECT SpanId AS id FROM otel_traces UNION DISTINCT SELECT SpanId AS id FROM trace_test.otel_traces ORDER BY id",
"SELECT t.SpanId AS id FROM otel_traces t INNER JOIN spend_logs s ON t.SpanId = s.request_id ORDER BY id",
"SELECT request_id AS id FROM spend_logs FINAL ORDER BY id",
];
for sql in queries {
let body: Value = serde_json::from_str(&query_sql(&database.client, &reader, sql).await?)?;
assert_eq!(
body["data"],
json!(
expected
.iter()
.map(|id| json!({"id": id}))
.collect::<Vec<_>>()
),
"{sql}"
);
}
let summary: Value = serde_json::from_str(
&query_sql(
&database.client,
&reader,
"SELECT sum(SpanCount) AS count FROM agent_traces_by_key",
)
.await?,
)?;
assert_eq!(summary["data"][0]["count"], json!(expected.len()));
let help = query_help(&database.client, &reader).await?;
assert_eq!(help.contains("secret_b"), expected.contains(&"b"));
assert_eq!(help.contains("secret-b"), expected.contains(&"b"));
let recreated = QueryReaders::new(database.writer.clone(), "trace_test".to_owned());
let repeated = recreated
.connection(&database.client, &scope, "test-master-secret")
.await?;
assert_eq!(reader.url(), repeated.url());
Ok(())
}
#[rstest]
#[tokio::test]
async fn rotating_master_secret_revokes_previous_reader_credentials(
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
) -> Result<(), Box<dyn std::error::Error>> {
let database = database?;
let scope = QueryScope::Team {
team_id: "team-a".to_owned(),
};
let old_reader = database
.readers
.connection(&database.client, &scope, "old-master-secret")
.await?;
let old_result = query_sql(
&database.client,
&old_reader,
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
)
.await?;
let old_rows: Value = serde_json::from_str(&old_result)?;
assert_eq!(old_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }]));
let rotated_readers = QueryReaders::new(database.writer.clone(), "trace_test".into());
let new_reader = rotated_readers
.connection(&database.client, &scope, "new-master-secret")
.await?;
assert!(
query_sql(
&database.client,
&old_reader,
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
)
.await
.is_err()
);
let new_result = query_sql(
&database.client,
&new_reader,
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
)
.await?;
let new_rows: Value = serde_json::from_str(&new_result)?;
assert_eq!(new_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }]));
assert_eq!(old_reader.url().username(), new_reader.url().username());
assert_ne!(old_reader.url().password(), new_reader.url().password());
Ok(())
}
#[rstest]
#[tokio::test]
async fn managed_reader_rejects_privilege_and_scope_bypasses(
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
) -> Result<(), Box<dyn std::error::Error>> {
let database = database?;
let scope = QueryScope::Team {
team_id: "team-a".to_owned(),
};
let reader = database
.readers
.connection(&database.client, &scope, "test-master-secret")
.await?;
for sql in [
"INSERT INTO otel_traces (TraceId) VALUES ('injected')",
"DROP TABLE otel_traces",
"SELECT * FROM private_data",
"SELECT * FROM otel_traces SETTINGS readonly = 0",
"SELECT * FROM otel_traces SETTINGS max_memory_usage = 0",
"SELECT * FROM otel_traces SETTINGS max_execution_time = 0",
"CREATE USER scope_bypass",
"CREATE NAMED COLLECTION scope_bypass AS host = 'localhost'",
"BACKUP TABLE otel_traces TO Disk('default', 'scope-bypass')",
"SELECT * FROM url('http://127.0.0.1:1/', 'LineAsString', 'line String')",
"SELECT * FROM remote('127.0.0.1', 'trace_test', 'otel_traces')",
] {
assert!(
matches!(
query_sql(&database.client, &reader, sql).await,
Err(Error::QueryFailed(_))
),
"{sql}"
);
}
let roles: Value = serde_json::from_str(
&query_sql(&database.client, &reader, "SELECT enabledRoles() AS roles").await?,
)?;
assert_eq!(roles["data"], json!([{ "roles": [] }]));
let rows: Value = serde_json::from_str(
&query_sql(
&database.client,
&reader,
"SELECT DISTINCT TeamId FROM otel_traces",
)
.await?,
)?;
assert_eq!(rows["data"], json!([{ "TeamId": "team-a" }]));
Ok(())
}
#[rstest]
#[tokio::test]
async fn provisioning_failure_never_returns_a_writer_connection(
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
) -> Result<(), Box<dyn std::error::Error>> {
let database = database?;
let reader = database
.readers
.connection(&database.client, &QueryScope::Admin, "test-master-secret")
.await?;
let no_provision_privileges = QueryReaders::new(reader, "trace_test".to_owned());
let result = no_provision_privileges
.connection(
&database.client,
&QueryScope::Team {
team_id: "team-a".to_owned(),
},
"other-secret",
)
.await;
assert!(result.is_err());
assert!(
database
.readers
.connection(&database.client, &QueryScope::Admin, "")
.await
.is_err()
);
assert!(
database
.readers
.connection(
&database.client,
&QueryScope::Team {
team_id: String::new()
},
"test-master-secret"
)
.await
.is_err()
);
let permits = (0..8)
.map(|_| database.readers.acquire())
.collect::<Result<Vec<_>, _>>()?;
assert!(database.readers.acquire().is_err());
drop(permits);
assert!(database.readers.acquire().is_ok());
let rows = litellm_traces::execute_read(
&database.client,
&database.writer,
"SELECT count() AS count FROM trace_test.otel_traces",
&BTreeMap::new(),
)
.await?;
let rows: Value = serde_json::from_str(&rows)?;
assert_eq!(rows["data"][0]["count"], 3);
Ok(())
}

23
litellm/__init__.py
View file

@ -2287,6 +2287,23 @@ if TYPE_CHECKING:
# Track if async client cleanup has been registered (for lazy loading)
_async_client_cleanup_registered = False
_AGENT_EXPORTS: Final = frozenset(
{
"agent",
"aagent",
"agent_session",
"aagent_session",
"agent_resume",
"aagent_resume",
"agent_capabilities",
"Harness",
"ClaudeCodeOptions",
"CodexOptions",
"OpenCodeOptions",
"DeepAgentsOptions",
}
)
# Eager loading for backwards compatibility with VCR and other HTTP recording tools
# When LITELLM_DISABLE_LAZY_LOADING is set, lazy-loaded attributes are loaded at import time
# For now, this only affects encoding (tiktoken) as it was the only reported issue
@ -2320,6 +2337,12 @@ def __getattr__(name: str) -> Any:
handler_func: Final = registry[name]
return handler_func(name)
if name == "harness" or name in _AGENT_EXPORTS:
import importlib
harness_module = importlib.import_module("litellm.harness")
return harness_module if name == "harness" else getattr(harness_module, name)
# Lazy load encoding from main.py to avoid heavy tiktoken import
if name == "encoding":
from ._lazy_imports import get_litellm_globals

4
litellm/anthropic_beta_headers_config.json
View file

@ -67,6 +67,7 @@
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": null,
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": null,
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
"web-search-2025-03-05": "web-search-2025-03-05"
@ -100,6 +101,7 @@
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": null,
"tool-search-tool-2025-10-19": null,
"web-fetch-2025-09-10": null,
@ -172,6 +174,7 @@
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"tool-examples-2025-10-29": "tool-examples-2025-10-29",
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
@ -207,6 +210,7 @@
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": null,
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
"web-fetch-2025-09-10": null,

38
litellm/constants.py
View file

@ -50,8 +50,8 @@ CLICKHOUSE_BATCH_SIZE: Final = get_env_int("CLICKHOUSE_BATCH_SIZE", 10_000)
CLICKHOUSE_FLUSH_INTERVAL_SECONDS: Final = float(os.getenv("CLICKHOUSE_FLUSH_INTERVAL_SECONDS", "1.0"))
CLICKHOUSE_MAX_BUFFERED_ROWS: Final = get_env_int("CLICKHOUSE_MAX_BUFFERED_ROWS", 200_000)
CLICKHOUSE_MAX_RETRIES: Final = get_env_int("CLICKHOUSE_MAX_RETRIES", 3)
AGENT_TRACING_RETENTION_DAYS: Final = get_env_int("AGENT_TRACING_RETENTION_DAYS", 30)
AGENT_TRACING_SPEND_LOG_RETENTION_DAYS: Final = get_env_int("AGENT_TRACING_SPEND_LOG_RETENTION_DAYS", 90)
DEFAULT_CLICKHOUSE_DATABASE: Final = "litellm"
DEFAULT_AGENT_TRACING_RETENTION_DAYS: Final = 14
OTLP_MAX_BODY_BYTES: Final = get_env_int("OTLP_MAX_BODY_BYTES", 16 * 1024 * 1024)
OTLP_MAX_ATTRIBUTE_VALUE_BYTES: Final = get_env_int("OTLP_MAX_ATTRIBUTE_VALUE_BYTES", 64 * 1024)
OTLP_RETRY_AFTER_SECONDS: Final = get_env_int("OTLP_RETRY_AFTER_SECONDS", 2)
@ -2164,6 +2164,17 @@ MCP_SPEND_LOG_MODEL_PREFIX: Final[str] = "MCP: "
PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__"
PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job"
PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
USAGE_TOP_API_KEYS_DEFAULT: Final[int] = 100
USAGE_TOP_API_KEYS_MAX: Final[int] = 1000
USAGE_KEY_PAGE_DEFAULT: Final[int] = 50
USAGE_KEY_PAGE_MAX: Final[int] = 100
USAGE_KEY_SEARCH_DEFAULT: Final[int] = 100
USAGE_KEY_SEARCH_MAX: Final[int] = 100
USAGE_MODEL_TOP_KEYS_DEFAULT: Final[int] = 5
USAGE_MODEL_TOP_KEYS_MAX: Final[int] = 100
USAGE_CACHE_LEAKAGE_KEYS_DEFAULT: Final[int] = 20
USAGE_CACHE_LEAKAGE_KEYS_MAX: Final[int] = 100
USAGE_EXPORT_BATCH_SIZE: Final[int] = 1000
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
# declares no ptu_effective_from, bounding the scan for an open-ended window.
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90
@ -2201,3 +2212,26 @@ EMPTY_MAPPING: Final = MappingProxyType({})
# API endpoint for breached password k-anonymity search
HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range"
# litellm.harness defaults
HARNESS_ENDPOINT_HOST: Final = "127.0.0.1"
HARNESS_ENDPOINT_STARTUP_TIMEOUT_SECONDS: Final = 10.0
HARNESS_ENDPOINT_REQUEST_TIMEOUT_SECONDS: Final = 600.0
HARNESS_SESSION_TOKEN_BYTES: Final = 32
HARNESS_MAX_DIFF_BYTES: Final = 256 * 1024
HARNESS_STDERR_TAIL_LINES: Final = 40
HARNESS_STREAM_READ_CHUNK_BYTES: Final = 64 * 1024
HARNESS_EVENT_QUEUE_MAX_SIZE: Final = 1024
HARNESS_PROCESS_KILL_GRACE_SECONDS: Final = 5.0
HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset(
{
".git",
"node_modules",
".venv",
"venv",
"__pycache__",
".mypy_cache",
".pytest_cache",
".ruff_cache",
}
)

98
litellm/harness/__init__.py Normal file
View file

@ -0,0 +1,98 @@
"""Agent harnesses: run Claude Code, Codex, OpenCode or Deep Agents on any LiteLLM model.
The entrypoints live on the top-level package:
import litellm
from litellm import Harness, sandbox
result = litellm.agent(
Harness.CLAUDE_CODE,
"fix the failing test",
sandbox=sandbox.local("."),
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
)
This module holds the types you get back: events, Result, State, errors.
"""
from litellm.harness.errors import (
CapabilityUnsupported,
HarnessError,
HarnessInstallFailed,
OptionsMismatch,
OutputInvalid,
SandboxError,
SessionClosed,
StateIncompatible,
)
from litellm.harness.options import (
ClaudeCodeOptions,
CodexOptions,
DeepAgentsOptions,
OpenCodeOptions,
)
from litellm.harness.runtime import (
AsyncEventStream,
AsyncSession,
aagent,
aagent_resume,
aagent_session,
agent_capabilities,
)
from litellm.harness.sync import EventStream, Session, agent, agent_resume, agent_session
from litellm.harness.types import (
Approval,
Capabilities,
Compaction,
Done,
Event,
FileChange,
Harness,
Reasoning,
Result,
State,
Text,
ToolCall,
ToolResult,
Usage,
)
__all__ = (
"Approval",
"AsyncEventStream",
"AsyncSession",
"Capabilities",
"CapabilityUnsupported",
"ClaudeCodeOptions",
"CodexOptions",
"Compaction",
"DeepAgentsOptions",
"Done",
"Event",
"EventStream",
"FileChange",
"Harness",
"HarnessError",
"HarnessInstallFailed",
"OpenCodeOptions",
"OptionsMismatch",
"OutputInvalid",
"Reasoning",
"Result",
"SandboxError",
"Session",
"SessionClosed",
"State",
"StateIncompatible",
"Text",
"ToolCall",
"ToolResult",
"Usage",
"aagent",
"aagent_resume",
"aagent_session",
"agent",
"agent_capabilities",
"agent_resume",
"agent_session",
)

62
litellm/harness/context.py Normal file
View file

@ -0,0 +1,62 @@
"""Per-session state shared by the runtime, handlers and harness configs."""
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, TypeAlias
from pydantic import BaseModel
from litellm.harness.options import HarnessOptions
from litellm.harness.sandbox.base import Sandbox
from litellm.harness.types import Approval, Harness, PermissionMode
if TYPE_CHECKING:
from litellm.harness.endpoint import ModelEndpoint
ApprovalHandler: TypeAlias = Callable[
[Approval], bool | Awaitable[bool] # mutable-ok: Callable parameter list in a type alias, not a runtime collection
]
@dataclass(frozen=True)
class GatewayTarget:
"""Resolved LiteLLM AI Gateway for `litellm_proxy/` models. Internal, not exported."""
api_base: str
api_key: str
@dataclass
class SessionContext:
"""Everything a handler and config need for a session. Owned by the runtime."""
harness: Harness
sandbox: Sandbox
session_id: str
# Model name as sent to the runtime (litellm_proxy/ prefix already stripped).
model: str | None = None
gateway: GatewayTarget | None = None
api_key: str | None = None
api_base: str | None = None
endpoint: ModelEndpoint | None = None
instructions: str | None = None
tools: Sequence[Callable[..., Any]] = ()
skills: Sequence[str] = ()
disable_tools: Sequence[str] = ()
permissions: PermissionMode = "full"
on_approval: ApprovalHandler | None = None
output: type[BaseModel] | None = None
max_turns: int | None = None
timeout: float | None = None
metadata: Mapping[str, Any] = field(default_factory=dict)
options: HarnessOptions | None = None
# Set by the handler after each turn.
final_text: str = ""
output_json: str | None = None
# Usage for in-process harnesses that call LiteLLM directly (no model endpoint).
input_tokens: int = 0
output_tokens: int = 0
cost: float = 0.0
calls: int = 0

665
litellm/harness/endpoint.py Normal file
View file

@ -0,0 +1,665 @@
"""Per-session local model endpoint every CLI harness talks to.
The runtime inside the sandbox points its Anthropic / OpenAI base URL at this endpoint and
authenticates with a random per-session token. The endpoint either reverse-proxies to a LiteLLM
AI Gateway (gateway mode) or calls the LiteLLM SDK directly (SDK mode), and counts usage + cost.
starlette and uvicorn are optional: they are imported only when an endpoint starts.
"""
from __future__ import annotations
import asyncio
import contextlib
import itertools
import json
import logging
import secrets
from collections.abc import AsyncIterable, AsyncIterator, Mapping
from dataclasses import dataclass
from types import MappingProxyType, ModuleType
from typing import TYPE_CHECKING, Any, Final
import httpx
import openai
import litellm
from litellm.constants import (
DEFAULT_POLLING_INTERVAL,
HARNESS_ENDPOINT_HOST,
HARNESS_ENDPOINT_REQUEST_TIMEOUT_SECONDS,
HARNESS_ENDPOINT_STARTUP_TIMEOUT_SECONDS,
HARNESS_PROCESS_KILL_GRACE_SECONDS,
HARNESS_SESSION_TOKEN_BYTES,
)
from litellm.harness.context import GatewayTarget
from litellm.harness.errors import HarnessError, HarnessInstallFailed
from litellm.harness.types import Harness, Usage
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import Response
from uvicorn import Server
verbose_logger: Final = logging.getLogger("LiteLLM")
MISSING_DEPS_MESSAGE = "litellm.harness needs starlette and uvicorn: pip install starlette uvicorn"
ROUTE_MESSAGES = "messages"
ROUTE_CHAT = "chat/completions"
ROUTE_RESPONSES = "responses"
POST_ROUTES = (ROUTE_MESSAGES, ROUTE_CHAT, ROUTE_RESPONSES)
ROUTE_PREFIXES: Final = ("", "/v1")
# What an SDK call or its stream can raise: LiteLLM maps provider failures onto openai's
# exception hierarchy; transport errors, bad request kwargs and unserializable chunks remain.
SDK_ERRORS: Final = (openai.OpenAIError, httpx.HTTPError, HarnessError, ValueError, TypeError)
HOP_BY_HOP_HEADERS = frozenset(
{
"connection",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"te",
"trailer",
"trailers",
"transfer-encoding",
"upgrade",
"host",
"content-length",
}
)
DROPPED_REQUEST_HEADERS = HOP_BY_HOP_HEADERS | frozenset(
(
"authorization",
"x-api-key",
"accept-encoding",
)
)
DROPPED_RESPONSE_HEADERS = HOP_BY_HOP_HEADERS | frozenset(("content-encoding",))
COST_HEADER = "x-litellm-response-cost"
SSE_MEDIA_TYPE = "text/event-stream"
@dataclass(frozen=True)
class _ServerDeps:
uvicorn: ModuleType
applications: ModuleType
routing: ModuleType
responses: ModuleType
def _load_server_deps() -> _ServerDeps:
"""Import starlette + uvicorn on demand; they are not litellm dependencies."""
try:
import uvicorn
from starlette import applications, responses, routing
except ImportError as e:
raise HarnessInstallFailed(MISSING_DEPS_MESSAGE) from e
return _ServerDeps(
uvicorn=uvicorn,
applications=applications,
routing=routing,
responses=responses,
)
@dataclass
class UsageTracker:
"""Running token + cost totals for one session."""
input_tokens: int = 0
output_tokens: int = 0
cost: float = 0.0
calls: int = 0
def add(self, input_tokens: int = 0, output_tokens: int = 0, cost: float = 0.0) -> None:
self.input_tokens += input_tokens
self.output_tokens += output_tokens
self.cost += cost
self.calls += 1
def snapshot(self) -> Usage:
return Usage(
input_tokens=self.input_tokens,
output_tokens=self.output_tokens,
calls=self.calls,
)
def _as_int(value: object) -> int:
if isinstance(value, bool):
return 0
if isinstance(value, (int, float)):
return int(value)
return 0
def usage_from_mapping(usage: object) -> tuple[int, int]:
"""(input, output) from a usage dict using OpenAI or Anthropic/Responses field names."""
if not isinstance(usage, Mapping):
return 0, 0
input_tokens = usage.get("input_tokens", usage.get("prompt_tokens"))
output_tokens = usage.get("output_tokens", usage.get("completion_tokens"))
return _as_int(input_tokens), _as_int(output_tokens)
def usage_from_body(body: object) -> tuple[int, int]:
"""Usage from a non-streaming JSON response body."""
if not isinstance(body, Mapping):
return 0, 0
if isinstance(body.get("usage"), Mapping):
return usage_from_mapping(body["usage"])
response = body.get("response")
if isinstance(response, Mapping):
return usage_from_mapping(response.get("usage"))
return 0, 0
class SSEUsageParser:
"""Collects token usage from an SSE byte stream as it passes through."""
def __init__(self) -> None:
self.input_tokens = 0
self.output_tokens = 0
self._buffer = b""
def feed(self, chunk: bytes) -> None:
self._buffer += chunk
*lines, self._buffer = self._buffer.split(b"\n")
for line in lines:
self._feed_line(line)
def close(self) -> None:
if self._buffer:
self._feed_line(self._buffer)
self._buffer = b""
def _feed_line(self, line: bytes) -> None:
text = line.strip()
if not text.startswith(b"data:"):
return
payload = text[len(b"data:") :].strip()
if not payload or payload == b"[DONE]":
return
try:
event = json.loads(payload)
except ValueError:
return
if isinstance(event, Mapping):
self.absorb(event)
def absorb(self, event: Mapping[str, Any]) -> None:
event_type = event.get("type")
if event_type == "message_start":
self._absorb_message_start(event)
elif event_type == "message_delta":
self._absorb_message_delta(event)
elif event_type == "response.completed":
self._absorb_response_completed(event)
elif isinstance(event.get("usage"), Mapping):
self._set(*usage_from_mapping(event["usage"]))
def _absorb_message_start(self, event: Mapping[str, Any]) -> None:
message = event.get("message")
if isinstance(message, Mapping):
self._set(*usage_from_mapping(message.get("usage")))
def _absorb_message_delta(self, event: Mapping[str, Any]) -> None:
# message_delta output_tokens is cumulative for the whole message.
self._set(*usage_from_mapping(event.get("usage")))
def _absorb_response_completed(self, event: Mapping[str, Any]) -> None:
response = event.get("response")
if isinstance(response, Mapping):
self._set(*usage_from_mapping(response.get("usage")))
def _set(self, input_tokens: int, output_tokens: int) -> None:
if input_tokens:
self.input_tokens = input_tokens
if output_tokens:
self.output_tokens = output_tokens
def compute_cost(model: str | None, input_tokens: int, output_tokens: int) -> float:
"""Cost from LiteLLM's price map. Never raises; unknown models cost 0.0."""
if not model or not (input_tokens or output_tokens):
return 0.0
try:
prompt_cost, completion_cost = litellm.cost_per_token(
model=model, prompt_tokens=input_tokens, completion_tokens=output_tokens
)
return float(prompt_cost) + float(completion_cost)
except Exception: # accounting must never break a call; the price-map lookup raises bare Exception
verbose_logger.debug("harness endpoint: cost lookup failed for %s", model, exc_info=True)
return 0.0
def header_cost(headers: Mapping[str, str]) -> float | None:
raw = headers.get(COST_HEADER)
if raw is None:
return None
try:
return float(raw)
except (TypeError, ValueError):
return None
def hidden_cost(response: object) -> float | None:
hidden = getattr(response, "_hidden_params", None)
if not isinstance(hidden, Mapping):
return None
try:
cost = hidden.get("response_cost")
return None if cost is None else float(cost)
except (TypeError, ValueError):
return None
def extract_token(headers: Mapping[str, str]) -> str | None:
auth = headers.get("authorization") or ""
if auth.lower().startswith("bearer "):
return auth[len("bearer ") :].strip()
return headers.get("x-api-key")
def gateway_headers(
incoming: Mapping[str, str],
gateway: GatewayTarget,
harness: Harness,
metadata: Mapping[str, Any] | None,
) -> Mapping[str, str]:
"""Incoming headers minus hop-by-hop/auth/x-litellm-*, plus gateway auth, tags, metadata."""
kept = (
(name, value)
for name, value in incoming.items()
if name.lower() not in DROPPED_REQUEST_HEADERS and not name.lower().startswith("x-litellm-")
)
metadata_json = json.dumps(dict(metadata), default=str) if metadata else None # mutable-ok: for json.dumps
metadata_header = (("x-litellm-spend-logs-metadata", metadata_json),) if metadata_json is not None else ()
added = (
("authorization", f"Bearer {gateway.api_key}"),
("x-litellm-tags", f"harness,{harness.value}"),
*metadata_header,
)
return MappingProxyType(dict(itertools.chain(kept, added)))
def response_headers(upstream: Mapping[str, str]) -> Mapping[str, str]:
return MappingProxyType(
{name: value for name, value in upstream.items() if name.lower() not in DROPPED_RESPONSE_HEADERS}
)
def sanitize(message: str, secret_values: tuple[str | None, ...]) -> str:
for value in secret_values:
if value:
message = message.replace(value, "***")
return message
def error_status(exc: BaseException) -> int:
status = getattr(exc, "status_code", None)
if isinstance(status, int) and 400 <= status <= 599:
return status
return 500
def error_body(exc: BaseException, message: str) -> dict[str, Any]: # mutable-ok: JSONResponse body
return {"error": {"type": type(exc).__name__, "message": message}} # mutable-ok: JSONResponse body
def to_jsonable(obj: object) -> object:
if hasattr(obj, "model_dump"):
return obj.model_dump(mode="json", exclude_none=True)
if isinstance(obj, Mapping):
return dict(obj) # mutable-ok: plain-dict copy so json.dumps can serialize any Mapping
return obj
def encode_anthropic_chunk(chunk: object) -> bytes:
if isinstance(chunk, bytes):
return chunk
if isinstance(chunk, str):
return chunk.encode()
data = to_jsonable(chunk)
event_type = data.get("type", "message") if isinstance(data, Mapping) else "message"
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
def encode_chat_chunk(chunk: object) -> bytes:
if hasattr(chunk, "model_dump_json"):
return f"data: {chunk.model_dump_json()}\n\n".encode()
return f"data: {json.dumps(to_jsonable(chunk))}\n\n".encode()
def encode_responses_chunk(chunk: object) -> bytes:
data = to_jsonable(chunk)
event_type = data.get("type", "message") if isinstance(data, Mapping) else "message"
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
STREAM_ENCODERS = MappingProxyType(
{
ROUTE_MESSAGES: encode_anthropic_chunk,
ROUTE_CHAT: encode_chat_chunk,
ROUTE_RESPONSES: encode_responses_chunk,
}
)
STREAM_TRAILERS = MappingProxyType({ROUTE_CHAT: b"data: [DONE]\n\n"})
def route_of(path: str) -> str:
stripped = path.strip("/")
stripped = stripped.removeprefix("v1/")
return stripped
def _noop() -> None:
return None
class ModelEndpoint:
"""Local HTTP endpoint for one harness session. Use as an async context manager."""
def __init__(
self,
harness: Harness,
model: str | None,
gateway: GatewayTarget | None,
api_key: str | None = None,
api_base: str | None = None,
metadata: Mapping[str, Any] | None = None,
*,
client: httpx.AsyncClient | None = None,
) -> None:
self.harness = harness
self.model = model
self.gateway = gateway
self.api_key = api_key
self.api_base = api_base
self.metadata: Mapping[str, Any] = MappingProxyType(dict(metadata or ()))
self.token = secrets.token_urlsafe(HARNESS_SESSION_TOKEN_BYTES)
self.usage = UsageTracker()
self.port = 0
self._injected_client = client
self._deps: _ServerDeps | None = None
self._client: httpx.AsyncClient | None = None
self._server: Any = None
self._task: asyncio.Task[None] | None = None
@property
def url(self) -> str:
return f"http://{HARNESS_ENDPOINT_HOST}:{self.port}"
async def __aenter__(self) -> ModelEndpoint:
await self.start()
return self
async def __aexit__(self, *exc_info: object) -> None:
await self.stop()
async def start(self) -> None:
self._deps = _load_server_deps()
if self.gateway is not None:
self._client = self._gateway_client()
self._server = self._build_server(self._deps)
self._task = asyncio.create_task(self._server.serve())
try:
await asyncio.wait_for(self._wait_started(), HARNESS_ENDPOINT_STARTUP_TIMEOUT_SECONDS)
except BaseException:
await self.stop()
raise
self.port = self._server.servers[0].sockets[0].getsockname()[1]
async def stop(self) -> None:
if self._server is not None:
self._server.should_exit = True
if self._task is not None:
with contextlib.suppress(BaseException):
await self._task
self._task = None
# Never close the client: the shared cached one may still serve other requests,
# and an injected one belongs to its caller.
self._client = None
def _gateway_client(self) -> httpx.AsyncClient:
"""LiteLLM's shared cached async client, unless one was injected."""
if self._injected_client is not None:
return self._injected_client
handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.AgentHarness,
params={ # mutable-ok: get_async_httpx_client takes a dict params argument
"timeout": HARNESS_ENDPOINT_REQUEST_TIMEOUT_SECONDS
},
)
return handler.client
async def _wait_started(self) -> None:
while not self._server.started:
if self._task is not None and self._task.done():
raise HarnessError("harness model endpoint failed to start")
await asyncio.sleep(DEFAULT_POLLING_INTERVAL)
def _build_server(self, deps: _ServerDeps) -> Server:
config = deps.uvicorn.Config(
self._build_app(deps),
host=HARNESS_ENDPOINT_HOST,
port=0,
log_config=None,
log_level="warning",
access_log=False,
lifespan="off",
timeout_graceful_shutdown=HARNESS_PROCESS_KILL_GRACE_SECONDS,
)
server = deps.uvicorn.Server(config)
# Never touch the host process's signal handlers.
if hasattr(server, "capture_signals"):
server.capture_signals = contextlib.nullcontext
if hasattr(server, "install_signal_handlers"):
server.install_signal_handlers = _noop
return server
def _build_app(self, deps: _ServerDeps) -> Starlette:
Route = deps.routing.Route
post_routes = tuple(
Route(
f"{prefix}/{route}",
self._handle,
methods=["POST"], # mutable-ok: Starlette Route takes a methods list
)
for prefix, route in itertools.product(ROUTE_PREFIXES, POST_ROUTES)
)
get_routes = tuple(
Route(f"{prefix}/models", self._models, methods=["GET"]) # mutable-ok: Starlette Route takes a methods list
for prefix in ROUTE_PREFIXES
)
return deps.applications.Starlette(
routes=[*post_routes, *get_routes] # mutable-ok: Starlette takes a routes list
)
@property
def _responses(self) -> ModuleType:
if self._deps is None:
raise HarnessError("harness model endpoint is not started")
return self._deps.responses
def _authorized(self, request: Request) -> bool:
token = extract_token(request.headers)
return token is not None and secrets.compare_digest(token.encode(), self.token.encode())
def _json(self, body: object, status_code: int = 200) -> Response:
return self._responses.JSONResponse(body, status_code=status_code)
def _unauthorized(self) -> Response:
return self._json(
{"error": {"type": "authentication_error", "message": "invalid token"}}, # mutable-ok: JSONResponse body
401,
)
def _error(self, exc: BaseException, status_code: int | None = None) -> Response:
message = sanitize(str(exc), self._secrets())
return self._json(error_body(exc, message), status_code or error_status(exc))
def _secrets(self) -> tuple[str | None, ...]:
gateway_key = self.gateway.api_key if self.gateway else None
return (gateway_key, self.api_key, self.token)
async def _models(self, request: Request) -> Response:
if not self._authorized(request):
return self._unauthorized()
entry = {"id": self.model, "object": "model", "created": 0, "owned_by": "litellm"} # mutable-ok: JSON body
data = (entry,) if self.model else ()
return self._json({"object": "list", "data": data}) # mutable-ok: JSON response body for Starlette JSONResponse
async def _handle(self, request: Request) -> Response:
if not self._authorized(request):
return self._unauthorized()
try:
body = json.loads(await request.body())
except ValueError as e:
return self._error(e, 400)
if not isinstance(body, dict):
return self._error(ValueError("request body must be a JSON object"), 400)
route = route_of(request.url.path)
if self.gateway is not None:
return await self._forward(request, route, body)
return await self._call_sdk(route, body)
def _cost_model(self, body: Mapping[str, Any]) -> str | None:
model = self.model or body.get("model")
return model if isinstance(model, str) else None
def _record(
self,
model: str | None,
input_tokens: int,
output_tokens: int,
cost: float | None,
) -> None:
if cost is None:
cost = compute_cost(model, input_tokens, output_tokens)
self.usage.add(input_tokens, output_tokens, cost)
async def _forward(self, request: Request, route: str, body: Mapping[str, Any]) -> Response:
if self._client is None or self.gateway is None:
raise HarnessError("gateway client is not started")
if self.model:
body = {**body, "model": self.model} # mutable-ok: JSON request body re-sent upstream via httpx json=
upstream_request = self._client.build_request(
"POST",
f"{self.gateway.api_base}/v1/{route}",
json=body,
headers=gateway_headers(request.headers, self.gateway, self.harness, self.metadata),
)
try:
upstream = await self._client.send(upstream_request, stream=True)
except httpx.HTTPError as e:
return self._error(e, 502)
return self._responses.StreamingResponse(
self._relay(upstream, self._cost_model(body)),
status_code=upstream.status_code,
headers=response_headers(upstream.headers),
)
async def _relay(self, upstream: httpx.Response, model: str | None) -> AsyncIterator[bytes]:
is_sse = SSE_MEDIA_TYPE in upstream.headers.get("content-type", "")
parser = SSEUsageParser()
collected = bytearray()
try:
async for chunk in upstream.aiter_bytes():
if is_sse:
parser.feed(chunk)
else:
collected.extend(chunk)
yield chunk
finally:
await upstream.aclose()
if upstream.status_code < 400:
self._record_relayed(upstream, model, parser, is_sse, bytes(collected))
def _record_relayed(
self,
upstream: httpx.Response,
model: str | None,
parser: SSEUsageParser,
is_sse: bool,
collected: bytes,
) -> None:
if is_sse:
parser.close()
tokens = (parser.input_tokens, parser.output_tokens)
else:
try:
tokens = usage_from_body(json.loads(collected))
except ValueError:
tokens = (0, 0)
self._record(model, tokens[0], tokens[1], header_cost(upstream.headers))
def _sdk_kwargs(
self, body: Mapping[str, Any]
) -> dict[str, Any]: # mutable-ok: SDK call kwargs, mutated by _invoke_sdk then splatted
kwargs: dict[str, Any] = {**body} # mutable-ok: SDK call kwargs built from the JSON body, then overridden
if self.model:
kwargs["model"] = self.model
if self.api_key:
kwargs["api_key"] = self.api_key
if self.api_base:
kwargs["api_base"] = self.api_base
return kwargs
async def _invoke_sdk(
self,
route: str,
kwargs: dict[str, Any], # mutable-ok: injects stream_options into the SDK kwargs
) -> object:
if route == ROUTE_MESSAGES:
return await litellm.anthropic.messages.acreate(**kwargs)
if route == ROUTE_CHAT:
if kwargs.get("stream"):
stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field
kwargs["stream_options"] = { # mutable-ok: JSON field sent to litellm.acompletion
"include_usage": True,
**stream_options,
}
return await litellm.acompletion(**kwargs)
return await litellm.aresponses(**kwargs)
async def _call_sdk(self, route: str, body: Mapping[str, Any]) -> Response:
kwargs = self._sdk_kwargs(body)
model = self._cost_model(kwargs)
try:
response = await self._invoke_sdk(route, kwargs)
except SDK_ERRORS as e:
verbose_logger.debug("harness endpoint: SDK call failed: %s", type(e).__name__)
return self._error(e)
if kwargs.get("stream") and isinstance(response, AsyncIterable):
return self._responses.StreamingResponse(
self._sdk_stream(route, response, model), media_type=SSE_MEDIA_TYPE
)
data = to_jsonable(response)
input_tokens, output_tokens = usage_from_body(data)
self._record(model, input_tokens, output_tokens, hidden_cost(response))
return self._json(data)
async def _sdk_stream(self, route: str, iterator: AsyncIterable[object], model: str | None) -> AsyncIterator[bytes]:
encode = STREAM_ENCODERS[route]
parser = SSEUsageParser()
try:
async for chunk in iterator:
encoded = encode(chunk)
parser.feed(encoded)
yield encoded
trailer = STREAM_TRAILERS.get(route)
if trailer:
yield trailer
except SDK_ERRORS as e:
message = sanitize(str(e), self._secrets())
yield f"event: error\ndata: {json.dumps(error_body(e, message))}\n\n".encode()
finally:
parser.close()
self._record(model, parser.input_tokens, parser.output_tokens, None)

45
litellm/harness/errors.py Normal file
View file

@ -0,0 +1,45 @@
"""Exceptions raised by litellm.harness."""
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from litellm.harness.types import Result
class HarnessError(Exception):
"""Base class for every litellm.harness error."""
class CapabilityUnsupported(HarnessError):
"""The harness cannot do what was asked. Raised before the runtime starts."""
class OptionsMismatch(HarnessError):
"""Options for a different harness, or a native option LiteLLM manages itself."""
class HarnessInstallFailed(HarnessError):
"""The runtime is missing from the sandbox or failed to start."""
class SandboxError(HarnessError):
"""The sandbox failed to start, run a command, or reach the host."""
class SessionClosed(HarnessError):
"""A turn was started on a session that is closed or detached."""
class StateIncompatible(HarnessError):
"""resume() was given a State from another harness or an unreadable version."""
class OutputInvalid(HarnessError):
"""The final answer did not validate against output=."""
def __init__(self, message: str, raw: str, result: Result | None = None) -> None:
super().__init__(message)
self.raw = raw
self.result = result

35
litellm/harness/handlers/__init__.py Normal file
View file

@ -0,0 +1,35 @@
"""Handlers run a harness config: CLI runtimes as subprocesses, Deep Agents in-process."""
from __future__ import annotations
from litellm.harness.errors import HarnessError
from litellm.harness.handlers.base import BaseHarnessHandler
from litellm.harness.types import Harness, require_harness
from litellm.llms.base_llm.harness.transformation import (
BaseCLIHarnessConfig,
BaseHarnessConfig,
)
from litellm.utils import ProviderConfigManager
def get_harness_config(harness: Harness) -> BaseHarnessConfig:
config = ProviderConfigManager.get_provider_harness_config(require_harness(harness))
if config is None:
raise HarnessError(f"No harness config registered for Harness.{harness.name}")
return config
def get_harness_handler(config: BaseHarnessConfig) -> BaseHarnessHandler:
"""The handler that knows how to run this kind of config."""
if isinstance(config, BaseCLIHarnessConfig):
from litellm.harness.handlers.cli_handler import CLIHarnessHandler
return CLIHarnessHandler(config)
if config.harness is Harness.DEEPAGENTS:
from litellm.harness.handlers.deepagents_handler import DeepAgentsHandler
return DeepAgentsHandler(config)
raise HarnessError(f"No handler for Harness.{config.harness.name}")
__all__ = ("BaseHarnessHandler", "get_harness_config", "get_harness_handler")

42
litellm/harness/handlers/base.py Normal file
View file

@ -0,0 +1,42 @@
"""The handler interface the runtime drives. A handler owns I/O for one session."""
from __future__ import annotations
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from typing import Any
from litellm.harness.context import SessionContext
from litellm.harness.errors import CapabilityUnsupported
from litellm.harness.types import Event
from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
class BaseHarnessHandler(ABC):
"""Runs one harness session. The config decides what to run; the handler runs it."""
def __init__(self, config: BaseHarnessConfig) -> None:
self.config = config
@abstractmethod
async def start(self, ctx: SessionContext) -> None:
"""Prepare the runtime (config files, skills, agent build). Called again after an interrupt."""
@abstractmethod
def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
"""Run one turn and yield events (never Done). Sets ctx.final_text / ctx.output_json."""
@abstractmethod
async def stop(self, ctx: SessionContext) -> None:
"""Stop anything this handler started. Safe to call twice."""
@abstractmethod
def native_session_id(self) -> str | None:
"""The runtime's own session id, for State / resume."""
@abstractmethod
async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
"""Continue the runtime's own session on the next turn."""
async def history(self, ctx: SessionContext) -> list[dict[str, Any]]: # mutable-ok: public history() API shape
raise CapabilityUnsupported(f"Harness.{self.config.harness.name} does not expose history")

161
litellm/harness/handlers/cli_handler.py Normal file
View file

@ -0,0 +1,161 @@
"""
Generic handler for CLI harnesses (Claude Code, Codex, OpenCode).
The config (`litellm/llms/<harness>/harness/transformation.py`) says what to run and how to
read it; this handler does every sandbox and process operation: binary check, private dir,
config files, persisted dirs, skills, spawning the turn, streaming stdout lines into the
config's parser, collecting stderr, and killing the process on early exit.
"""
from __future__ import annotations
import asyncio
import os
from collections import deque
from collections.abc import AsyncIterator, Sequence
from typing import Any, Final
from litellm._logging import verbose_logger
from litellm.constants import HARNESS_STDERR_TAIL_LINES, HARNESS_STREAM_READ_CHUNK_BYTES
from litellm.harness.context import SessionContext
from litellm.harness.errors import HarnessInstallFailed, SandboxError
from litellm.harness.handlers.base import BaseHarnessHandler
from litellm.harness.sandbox.base import Process, Sandbox
from litellm.harness.types import Event
from litellm.llms.base_llm.harness.transformation import BaseCLIHarnessConfig, HarnessSessionSetup
from litellm.llms.base_llm.harness.utils import decode_json_line, read_skill_files
# Link <private_dir>/<dir> to a LiteLLM-owned cache dir so a later session can resume.
PERSIST_DIR_SCRIPT: Final = (
'd="${HOME:-/tmp}/.cache/litellm-harness/$2"; mkdir -p "$d" && mkdir -p "$(dirname "$1")" && ln -sfn "$d" "$1"'
)
async def iter_stream_lines(stream: asyncio.StreamReader) -> AsyncIterator[bytes]:
"""Newline-delimited lines without StreamReader's 64KiB readline limit."""
buffer = b""
while True:
chunk = await stream.read(HARNESS_STREAM_READ_CHUNK_BYTES)
if not chunk:
break
buffer += chunk
*lines, buffer = buffer.split(b"\n")
for line in lines:
yield line
if buffer:
yield buffer
async def drain_stderr(stream: asyncio.StreamReader, tail: deque[str]) -> None: # mutable-ok: stderr ring
async for line in iter_stream_lines(stream):
tail.append(line.decode("utf-8", errors="replace"))
async def send_stdin(proc: Process, data: str) -> None:
if proc.stdin is None:
raise SandboxError("harness process has no stdin")
proc.stdin.write(data.encode("utf-8"))
await proc.stdin.drain()
proc.stdin.close()
async def private_dir_for(sandbox: Sandbox) -> str:
tempdir = getattr(sandbox, "tempdir", None)
if tempdir is None:
raise SandboxError(f"{type(sandbox).__name__} has no tempdir(); CLI harnesses need a private config dir")
path: str = await tempdir()
return path
def sandbox_path(private_dir: str, path: str) -> str:
return path if path.startswith("/") else f"{private_dir}/{path}"
async def persist_dir(sandbox: Sandbox, link_path: str, cache_subpath: str) -> None:
script_args: Final = ("-c", PERSIST_DIR_SCRIPT, "sh", link_path, cache_subpath)
cmd: Final = ["sh", *script_args] # mutable-ok: Sandbox.run takes list[str]
run = await sandbox.run(cmd)
if run.exit_code != 0:
verbose_logger.debug(
"harness: could not persist %s, resume across sessions disabled: %s", cache_subpath, run.stderr.strip()
)
async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str) -> None:
for skill in skills:
name = os.path.basename(os.path.realpath(os.fspath(skill)))
for rel, data in await asyncio.to_thread(read_skill_files, skill):
await sandbox.write(f"{skills_root}/{name}/{rel.replace(os.sep, '/')}", data)
class CLIHarnessHandler(BaseHarnessHandler):
config: BaseCLIHarnessConfig
def __init__(self, config: BaseCLIHarnessConfig) -> None:
super().__init__(config)
self._private_dir: str | None = None
self._setup: HarnessSessionSetup | None = None
self._native_id: str | None = None
self._proc: Process | None = None
async def start(self, ctx: SessionContext) -> None:
self.config.validate_environment(ctx)
binary = self.config.get_binary()
if not await ctx.sandbox.which(binary):
raise HarnessInstallFailed(
f"`{binary}` was not found on PATH in the sandbox. Install it with: {self.config.get_install_hint()}"
)
private_dir = await private_dir_for(ctx.sandbox)
setup = self.config.transform_session_setup(ctx, private_dir)
for link, cache_subpath in setup.persisted_dirs:
await persist_dir(ctx.sandbox, sandbox_path(private_dir, link), cache_subpath)
for rel_path, data in setup.files.items():
await ctx.sandbox.write(sandbox_path(private_dir, rel_path), data)
if ctx.skills and setup.skills_dir:
await copy_skills(ctx.sandbox, tuple(ctx.skills), sandbox_path(private_dir, setup.skills_dir))
self._private_dir = private_dir
self._setup = setup
async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
if self._setup is None or self._private_dir is None:
raise RuntimeError("CLIHarnessHandler.turn() called before start()")
request = self.config.transform_turn_request(ctx, self._setup, self._private_dir, prompt, self._native_id)
argv: Final = list(request.argv) # mutable-ok: Sandbox.exec takes list[str]
proc = await ctx.sandbox.exec(argv, env=request.env, cwd=request.cwd)
self._proc = proc
tail: Final[deque[str]] = deque(maxlen=HARNESS_STDERR_TAIL_LINES) # mutable-ok: bounded stderr ring buffer
stderr_task = asyncio.ensure_future(drain_stderr(proc.stderr, tail))
state: Any = self.config.create_stream_state()
exit_code: int | None = None
try:
await send_stdin(proc, request.stdin)
async for raw in iter_stream_lines(proc.stdout):
line = decode_json_line(raw)
if line is None:
continue
for event in self.config.transform_stream_line(line, state):
yield event
self._native_id = self.config.get_native_session_id(state) or self._native_id
exit_code = await proc.wait()
await stderr_task
finally:
self._proc = None
if exit_code is None:
# Consumer stopped early, timed out or errored: don't leave the runtime running.
await proc.kill()
if not stderr_task.done():
stderr_task.cancel()
response = self.config.transform_turn_response(ctx, state, exit_code, tuple(tail))
ctx.final_text = response.final_text
ctx.output_json = response.output_json
async def stop(self, ctx: SessionContext) -> None:
proc, self._proc = self._proc, None
if proc is not None:
await proc.kill()
def native_session_id(self) -> str | None:
return self._native_id
async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
self._native_id = native_session_id

261
litellm/harness/handlers/deepagents_handler.py Normal file
View file

@ -0,0 +1,261 @@
"""
In-process handler for Deep Agents.
Deep Agents is a Python library, so there is no process or model endpoint: the handler
builds the agent with a LiteLLM chat model, streams the LangGraph run, turns interrupts into
Approval events and counts usage. Translation lives in
`litellm/llms/deepagents/harness/transformation.py`.
"""
from __future__ import annotations
import asyncio
import importlib
from collections.abc import AsyncIterator, Mapping
from dataclasses import dataclass
from types import MappingProxyType, ModuleType
from typing import TYPE_CHECKING, Any
from litellm.harness.context import SessionContext
from litellm.harness.errors import HarnessError, HarnessInstallFailed
from litellm.harness.handlers.base import BaseHarnessHandler
from litellm.harness.handlers.cli_handler import copy_skills
from litellm.harness.options import DeepAgentsOptions
from litellm.harness.types import Approval, Event
from litellm.llms.deepagents.harness.transformation import (
EXECUTE_TOOLS,
INSTALL_HINT,
SKILLS_DIR,
WRITE_TOOLS,
TurnState,
approval_requests,
blocked_tools,
chat_model_kwargs,
decision,
final_ai_text,
interrupt_config,
interrupts_in,
normalized_tool_name,
recursion_limit,
stream_events,
structured_json,
update_events,
)
if TYPE_CHECKING:
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.language_models import BaseChatModel
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph.state import CompiledStateGraph
from langgraph.types import Command
from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
_MODEL_NODE = "model"
@dataclass(frozen=True)
class DeepAgentsDeps:
"""The optional-dependency entrypoints this handler uses."""
create_deep_agent: Any
chat_litellm: Any
checkpointer_cls: Any
command_cls: Any
subagent_defaults: Mapping[str, Any]
convert_to_openai_messages: Any
backend: ModuleType
def load_deps() -> DeepAgentsDeps:
"""Import deepagents + langchain-litellm, or raise HarnessInstallFailed."""
try:
deepagents = importlib.import_module("deepagents")
subagents = importlib.import_module("deepagents.middleware.subagents")
chat = importlib.import_module("langchain_litellm")
memory = importlib.import_module("langgraph.checkpoint.memory")
lg_types = importlib.import_module("langgraph.types")
messages = importlib.import_module("langchain_core.messages")
backend = importlib.import_module("litellm.llms.deepagents.harness.sandbox_backend")
except ImportError as e:
raise HarnessInstallFailed(f"{INSTALL_HINT} ({e})") from e
return DeepAgentsDeps(
create_deep_agent=deepagents.create_deep_agent,
chat_litellm=chat.ChatLiteLLM,
checkpointer_cls=memory.InMemorySaver,
command_cls=lg_types.Command,
subagent_defaults=subagents.GENERAL_PURPOSE_SUBAGENT,
convert_to_openai_messages=messages.convert_to_openai_messages,
backend=backend,
)
_SHARED_CHECKPOINTER: dict[str, Any] = {} # mutable-ok: process-wide lazy singleton slot for the in-memory checkpointer
def shared_checkpointer(deps: DeepAgentsDeps) -> BaseCheckpointSaver:
"""One in-memory checkpointer per process, so resume() works across sessions in-process."""
saver = _SHARED_CHECKPOINTER.get("saver")
if saver is None:
saver = deps.checkpointer_cls()
_SHARED_CHECKPOINTER["saver"] = saver
return saver
def build_chat_model(ctx: SessionContext, deps: DeepAgentsDeps) -> BaseChatModel:
"""The LangChain chat model for this session. Tests monkeypatch this."""
return deps.chat_litellm(**chat_model_kwargs(ctx))
class DeepAgentsHandler(BaseHarnessHandler):
def __init__(self, config: BaseHarnessConfig) -> None:
super().__init__(config)
self._deps: DeepAgentsDeps | None = None
self._agent: Any = None
self._thread_id: str | None = None
self._skip_tools: frozenset[str] = frozenset()
async def start(self, ctx: SessionContext) -> None:
self.config.validate_environment(ctx)
deps = load_deps()
self._deps = deps
blocked = blocked_tools(ctx.permissions, ctx.disable_tools)
backend = deps.backend.SandboxBackend(
ctx.sandbox,
loop=asyncio.get_running_loop(),
writable=WRITE_TOOLS.isdisjoint(blocked),
allow_execute=EXECUTE_TOOLS.isdisjoint(blocked),
)
self._agent = deps.create_deep_agent(
model=build_chat_model(ctx, deps),
tools=list(ctx.tools), # mutable-ok: deepagents create_deep_agent(tools=) takes a list
system_prompt=ctx.instructions,
middleware=self._middleware(deps, blocked),
subagents=self._subagents(ctx, deps, blocked),
skills=await self._install_skills(ctx),
backend=backend,
interrupt_on=interrupt_config(ctx.permissions, blocked),
response_format=ctx.output,
checkpointer=shared_checkpointer(deps),
)
self._skip_tools = frozenset({ctx.output.__name__}) if ctx.output is not None else frozenset()
if self._thread_id is None:
self._thread_id = ctx.session_id
async def stop(self, ctx: SessionContext) -> None:
self._agent = None
def native_session_id(self) -> str | None:
return self._thread_id
async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
self._thread_id = native_session_id
async def history(
self, ctx: SessionContext
) -> list[dict[str, Any]]: # mutable-ok: BaseHarnessHandler.history API returns OpenAI message dicts
agent, deps = self._require_agent()
snapshot = await agent.aget_state(self._run_config(ctx, None))
messages = (snapshot.values or MappingProxyType({})).get("messages") or ()
converted: list[dict[str, Any]] = deps.convert_to_openai_messages( # mutable-ok: LangChain returns a list
messages
)
return converted
async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
agent, deps = self._require_agent()
run_config = self._run_config(ctx, deps.backend.UsageCallback(ctx, ctx.model))
user_message = {"role": "user", "content": prompt} # mutable-ok: LangGraph input message dict
payload: dict[str, object] | Command = {"messages": [user_message]} # mutable-ok: LangGraph input state
while True:
state = TurnState()
async for event in self._stream_pass(agent, payload, run_config, state):
yield event
if not state.interrupts:
break
resume: dict[str, Any] = {} # mutable-ok: Command(resume=) payload, filled per answered approval
for interrupt in state.interrupts:
decisions: list[dict[str, Any]] = [] # mutable-ok: HITL decisions collected across awaited approvals
for request in approval_requests(getattr(interrupt, "value", None)):
approval = Approval(
tool=normalized_tool_name(str(request.get("name") or "")),
input=dict(request.get("args") or ()), # mutable-ok: Approval.input is a public dict field
)
yield approval
decisions.append(decision(*await approval.wait()))
resume[interrupt.id] = {"decisions": decisions} # mutable-ok: LangGraph HITL resume payload
payload = deps.command_cls(resume=resume)
await self._finish_turn(ctx, agent, run_config)
async def _stream_pass(
self,
agent: CompiledStateGraph,
payload: dict[str, object] | Command, # mutable-ok: LangGraph astream input type
run_config: RunnableConfig,
state: TurnState,
) -> AsyncIterator[Event]:
async for part in agent.astream(
payload,
run_config,
stream_mode=["messages", "updates"], # mutable-ok: LangGraph stream_mode takes a list
):
# A list stream_mode yields (mode, chunk) tuples; LangGraph's overloads don't say so.
if not isinstance(part, tuple) or len(part) != 2:
continue
mode, chunk = part
if mode == "messages":
message, meta = chunk
if isinstance(meta, Mapping) and meta.get("langgraph_node") == _MODEL_NODE:
for event in stream_events(message):
yield event
elif mode == "updates":
state.interrupts = (*state.interrupts, *interrupts_in(chunk))
for event in update_events(chunk, self._skip_tools):
yield event
async def _finish_turn(self, ctx: SessionContext, agent: CompiledStateGraph, run_config: RunnableConfig) -> None:
snapshot = await agent.aget_state(run_config)
values = snapshot.values or MappingProxyType({})
ctx.final_text = final_ai_text(values.get("messages") or ())
if ctx.output is not None:
ctx.output_json = structured_json(values.get("structured_response"))
def _require_agent(self) -> tuple[Any, DeepAgentsDeps]:
if self._agent is None or self._deps is None:
raise HarnessError("Deep Agents session is not started")
return self._agent, self._deps
def _run_config(self, ctx: SessionContext, usage_callback: BaseCallbackHandler | None) -> RunnableConfig:
run_config: RunnableConfig = {
"configurable": {"thread_id": self._thread_id or ctx.session_id},
"recursion_limit": recursion_limit(ctx),
}
if usage_callback is not None:
run_config["callbacks"] = [usage_callback] # mutable-ok: LangChain RunnableConfig.callbacks is a list
return run_config
@staticmethod
def _middleware(deps: DeepAgentsDeps, blocked: frozenset[str]) -> list[Any]: # mutable-ok: deepagents API
filters = (deps.backend.ToolFilterMiddleware(blocked),) if blocked else ()
return list(filters) # mutable-ok: deepagents create_deep_agent(middleware=) takes a list
def _subagents(
self, ctx: SessionContext, deps: DeepAgentsDeps, blocked: frozenset[str]
) -> list[Any]: # mutable-ok: deepagents create_deep_agent(subagents=) takes a list
"""User subagents, plus a general-purpose one that honours disable_tools when set."""
options = ctx.options if isinstance(ctx.options, DeepAgentsOptions) else None
user_subagents = tuple(options.subagents) if options is not None else ()
has_general = any(
isinstance(s, Mapping) and s.get("name") == deps.subagent_defaults["name"] for s in user_subagents
)
spec = {**deps.subagent_defaults, "middleware": self._middleware(deps, blocked)} # mutable-ok: SubAgent dict
general = (spec,) if blocked and not has_general else ()
return [*general, *user_subagents] # mutable-ok: deepagents create_deep_agent(subagents=) takes a list
@staticmethod
async def _install_skills(ctx: SessionContext) -> list[str] | None: # mutable-ok: deepagents skills= takes a list
if not ctx.skills:
return None
await copy_skills(ctx.sandbox, ctx.skills, f"{ctx.sandbox.workdir}/{SKILLS_DIR}")
return [f"/{SKILLS_DIR}/"] # mutable-ok: deepagents create_deep_agent(skills=) takes a list

37
litellm/harness/options.py Normal file
View file

@ -0,0 +1,37 @@
"""Typed per-harness options. Settings that only make sense for one runtime live here."""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import Any, Literal
@dataclass(frozen=True)
class ClaudeCodeOptions:
config: Mapping[str, Any] = field(default_factory=dict)
env: Mapping[str, str] = field(default_factory=dict)
@dataclass(frozen=True)
class CodexOptions:
reasoning_effort: Literal["low", "medium", "high", "xhigh"] | None = None
web_search: bool = False
config: Mapping[str, Any] = field(default_factory=dict)
env: Mapping[str, str] = field(default_factory=dict)
@dataclass(frozen=True)
class OpenCodeOptions:
agent: str = "build"
config: Mapping[str, Any] = field(default_factory=dict)
env: Mapping[str, str] = field(default_factory=dict)
@dataclass(frozen=True)
class DeepAgentsOptions:
subagents: Sequence[Any] = ()
recursion_limit: int | None = None
HarnessOptions = ClaudeCodeOptions | CodexOptions | OpenCodeOptions | DeepAgentsOptions

1028
litellm/harness/runtime.py Normal file
View file

File diff suppressed because it is too large Load diff

25
litellm/harness/sandbox/__init__.py Normal file
View file

@ -0,0 +1,25 @@
"""Sandboxes for litellm.harness: where the runtime runs and which files it can touch."""
from litellm.harness.sandbox.base import CompletedRun, Process, Sandbox
from litellm.harness.sandbox.docker import DockerSandbox, docker
from litellm.harness.sandbox.local import LocalSandbox, local
from litellm.harness.sandbox.snapshot import (
build_file_changes,
capture_text_contents,
diff_snapshots,
snapshot_local,
)
__all__ = (
"CompletedRun",
"DockerSandbox",
"LocalSandbox",
"Process",
"Sandbox",
"build_file_changes",
"capture_text_contents",
"diff_snapshots",
"docker",
"local",
"snapshot_local",
)

62
litellm/harness/sandbox/base.py Normal file
View file

@ -0,0 +1,62 @@
"""The Sandbox protocol: where a harness runtime runs and which files it can touch."""
from __future__ import annotations
import asyncio
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Protocol, runtime_checkable
@dataclass(frozen=True)
class CompletedRun:
stdout: str
stderr: str
exit_code: int
@runtime_checkable
class Process(Protocol):
stdin: asyncio.StreamWriter | None
stdout: asyncio.StreamReader
stderr: asyncio.StreamReader
async def wait(self) -> int: ...
async def kill(self) -> None: ...
@runtime_checkable
class Sandbox(Protocol):
workdir: str
async def exec(
self,
cmd: Sequence[str],
*,
env: Mapping[str, str] | None = None,
cwd: str | None = None,
) -> Process: ...
async def run(
self,
cmd: Sequence[str],
*,
env: Mapping[str, str] | None = None,
cwd: str | None = None,
timeout: float | None = None,
) -> CompletedRun: ...
async def read(self, path: str) -> bytes: ...
async def write(self, path: str, data: bytes) -> None: ...
def host_url(self, port: int) -> str: ...
async def which(self, binary: str) -> str | None: ...
async def snapshot(self) -> Mapping[str, str]: ...
async def tempdir(self) -> str: ...
async def close(self) -> None: ...

290
litellm/harness/sandbox/docker.py Normal file
View file

@ -0,0 +1,290 @@
"""DockerSandbox: run the harness runtime inside a container via the docker CLI."""
from __future__ import annotations
import asyncio
import os
import posixpath
import shutil
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from litellm.constants import HARNESS_SNAPSHOT_SKIP_DIRS
from litellm.harness.errors import SandboxError
from litellm.harness.sandbox.base import CompletedRun
from litellm.harness.sandbox.local import SubprocessHandle, collect_output
from litellm.harness.sandbox.snapshot import HARNESS_SNAPSHOT_MAX_FILE_BYTES
DOCKER_HOST_ALIAS: Final = "host.docker.internal"
_SHA256_HEX_LEN: Final = 64
_WRITE_SCRIPT: Final = 'mkdir -p "$(dirname "$1")" && cat > "$1"'
_WHICH_SCRIPT: Final = 'command -v "$1"'
def _snapshot_script() -> str:
prune = " -o ".join(f"-name '{name}'" for name in sorted(HARNESS_SNAPSHOT_SKIP_DIRS))
return (
'cd "$1" && find . -type d \\( '
+ prune
+ " \\) -prune -o -type f -size -"
+ f"{HARNESS_SNAPSHOT_MAX_FILE_BYTES + 1}c"
+ " -exec sha256sum {} +"
)
def parse_sha256sum(output: str) -> Mapping[str, str]:
"""Parse `sha256sum` lines ("<hex> ./rel/path") into {rel/path: hex}."""
return MappingProxyType(
{
line[_SHA256_HEX_LEN + 2 :].removeprefix("./"): line[:_SHA256_HEX_LEN]
for line in output.splitlines()
if len(line) > _SHA256_HEX_LEN + 2
}
)
class DockerSandbox:
"""Sandbox backed by a long-lived `sleep infinity` container."""
# Harness configs read this to skip a runtime's own nested OS sandbox.
is_container = True
def __init__(
self,
image: str,
mounts: Mapping[str | os.PathLike[str], str] | None = None,
workdir: str = "/workspace",
env: Mapping[str, str] | None = None,
name: str | None = None,
) -> None:
if not image:
raise SandboxError("docker sandbox needs an image")
if not posixpath.isabs(workdir):
raise SandboxError(f"docker workdir must be absolute: {workdir}")
self.image = image
self.workdir: str = posixpath.normpath(workdir)
self.mounts: Mapping[str, str] = MappingProxyType(
{os.path.abspath(os.fspath(host)): container for host, container in (mounts.items() if mounts else ())}
)
self.env: Mapping[str, str] = MappingProxyType(dict(env or ()))
self.name = name
self.container_id: str | None = None
self._start_lock = asyncio.Lock()
self._processes: set[SubprocessHandle] = set() # mutable-ok: live-process registry (add/discard)
self._closed = False
def __repr__(self) -> str:
return f"DockerSandbox({self.image!r}, workdir={self.workdir!r})"
def _docker_binary(self) -> str:
binary = shutil.which("docker")
if binary is None:
raise SandboxError(
"docker sandbox requires the `docker` CLI on PATH; install Docker or use sandbox.local(path)"
)
return binary
async def _spawn(self, args: Sequence[str]) -> SubprocessHandle:
"""Start `docker <args>` with stdin/stdout/stderr pipes."""
try:
proc = await asyncio.create_subprocess_exec(
self._docker_binary(),
*args,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
except (FileNotFoundError, PermissionError) as exc:
raise SandboxError(f"could not run docker: {exc}") from exc
return SubprocessHandle(proc)
async def _docker(
self,
args: Sequence[str],
*,
input: bytes | None = None,
timeout: float | None = None,
) -> tuple[int, bytes, bytes]:
"""Run `docker <args>` to completion; returns (exit_code, stdout, stderr)."""
handle = await self._spawn(args)
try:
return await asyncio.wait_for(_communicate(handle, input), timeout)
except asyncio.TimeoutError:
await handle.kill()
raise SandboxError(f"docker {args[0]} timed out after {timeout}s")
def run_args(
self,
) -> list[str]: # mutable-ok: argv is returned as a list, the shape callers and tests compare against
name_args = ("--name", self.name) if self.name else ()
mount_args = tuple(
arg for host, container in self.mounts.items() for arg in ("-v", f"{host}:{container}")
) # comprehension-ok: flattens (flag, value) pairs into argv
env_args = tuple(
arg for key, value in self.env.items() for arg in ("-e", f"{key}={value}")
) # comprehension-ok: flattens (flag, value) pairs into argv
return [ # mutable-ok: argv is returned as a list, the shape callers and tests compare against
"run",
"-d",
"--rm",
f"--add-host={DOCKER_HOST_ALIAS}:host-gateway",
*name_args,
*mount_args,
*env_args,
"-w",
self.workdir,
self.image,
"sleep",
"infinity",
]
def exec_args(
self,
container_id: str,
cmd: Sequence[str],
*,
env: Mapping[str, str] | None = None,
cwd: str | None = None,
) -> list[str]: # mutable-ok: argv is returned as a list, the shape callers and tests compare against
env_args = tuple(
arg for key, value in (env.items() if env else ()) for arg in ("-e", f"{key}={value}")
) # comprehension-ok: flattens (flag, value) pairs into argv
return [ # mutable-ok: argv is returned as a list, the shape callers and tests compare against
"exec",
"-i",
"-w",
self.container_path(cwd or self.workdir),
*env_args,
container_id,
*cmd,
]
def container_path(self, path: str) -> str:
"""Absolute container path; relative paths resolve against workdir."""
joined = path if posixpath.isabs(path) else posixpath.join(self.workdir, path)
return posixpath.normpath(joined)
async def start(self) -> str:
"""Start the container if needed and return its id."""
if self._closed:
raise SandboxError("sandbox is closed")
async with self._start_lock:
if self.container_id is not None:
return self.container_id
code, out, err = await self._docker(self.run_args())
if code != 0:
raise SandboxError(f"docker run {self.image} failed ({code}): {err.decode(errors='replace').strip()}")
container_id = out.decode().strip()
if not container_id:
raise SandboxError("docker run returned no container id")
self.container_id = container_id
return container_id
async def _exec_capture(self, cmd: Sequence[str], *, input: bytes | None = None) -> tuple[int, bytes, bytes]:
container_id = await self.start()
return await self._docker(self.exec_args(container_id, cmd), input=input)
async def exec(
self,
cmd: Sequence[str],
*,
env: Mapping[str, str] | None = None,
cwd: str | None = None,
) -> SubprocessHandle:
if not cmd:
raise SandboxError("exec() needs a non-empty command")
container_id = await self.start()
handle = await self._spawn(self.exec_args(container_id, cmd, env=env, cwd=cwd))
self._processes.add(handle)
return handle
async def run(
self,
cmd: Sequence[str],
*,
env: Mapping[str, str] | None = None,
cwd: str | None = None,
timeout: float | None = None,
) -> CompletedRun:
handle = await self.exec(cmd, env=env, cwd=cwd)
try:
return await collect_output(handle, cmd, timeout)
finally:
self._processes.discard(handle)
async def read(self, path: str) -> bytes:
target = self.container_path(path)
code, out, err = await self._exec_capture(("cat", target))
if code != 0:
raise SandboxError(f"could not read {target}: {err.decode(errors='replace').strip()}")
return out
async def write(self, path: str, data: bytes) -> None:
target = self.container_path(path)
code, _, err = await self._exec_capture(("sh", "-c", _WRITE_SCRIPT, "sh", target), input=data)
if code != 0:
raise SandboxError(f"could not write {target}: {err.decode(errors='replace').strip()}")
def host_url(self, port: int) -> str:
return f"http://{DOCKER_HOST_ALIAS}:{port}"
async def which(self, binary: str) -> str | None:
code, out, _ = await self._exec_capture(("sh", "-lc", _WHICH_SCRIPT, "sh", binary))
found = out.decode(errors="replace").strip()
return found if code == 0 and found else None
async def tempdir(self) -> str:
"""A fresh `mktemp -d` directory inside the container."""
code, out, err = await self._exec_capture(("mktemp", "-d"))
path = out.decode(errors="replace").strip()
if code != 0 or not path:
raise SandboxError(f"mktemp -d failed: {err.decode(errors='replace').strip()}")
return path
async def snapshot(self) -> Mapping[str, str]:
code, out, err = await self._exec_capture(("sh", "-c", _snapshot_script(), "sh", self.workdir))
if code != 0:
raise SandboxError(f"snapshot failed: {err.decode(errors='replace').strip()}")
return parse_sha256sum(out.decode("utf-8", errors="replace"))
async def close(self) -> None:
if self._closed:
return
self._closed = True
live = tuple(h for h in self._processes if h.returncode is None)
await asyncio.gather(*(h.kill() for h in live), return_exceptions=True)
self._processes.clear()
if self.container_id is not None:
container_id, self.container_id = self.container_id, None
await self._docker(
["rm", "-f", container_id] # mutable-ok: argv list, the shape _spawn records and tests assert on
)
async def __aenter__(self) -> DockerSandbox:
await self.start()
return self
async def __aexit__(self, *exc_info: object) -> None:
await self.close()
async def _communicate(handle: SubprocessHandle, data: bytes | None) -> tuple[int, bytes, bytes]:
if handle.stdin is not None:
if data:
handle.stdin.write(data)
await handle.stdin.drain()
handle.stdin.close()
stdout, stderr = await asyncio.gather(handle.stdout.read(), handle.stderr.read())
return await handle.wait(), stdout, stderr
def docker(
image: str,
mounts: Mapping[str | os.PathLike[str], str] | None = None,
workdir: str = "/workspace",
env: Mapping[str, str] | None = None,
name: str | None = None,
) -> DockerSandbox:
"""Sandbox in a new container of `image`, started lazily on first use."""
return DockerSandbox(image, mounts=mounts, workdir=workdir, env=env, name=name)

277
litellm/harness/sandbox/local.py Normal file
View file

@ -0,0 +1,277 @@
"""LocalSandbox: run the harness runtime as a subprocess on this machine."""
from __future__ import annotations
import asyncio
import itertools
import os
import shutil
import signal
import tempfile
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from litellm.constants import HARNESS_PROCESS_KILL_GRACE_SECONDS
from litellm.harness.errors import SandboxError
from litellm.harness.sandbox.base import CompletedRun
from litellm.harness.sandbox.snapshot import snapshot_local
_SECRET_PREFIXES: Final = (
"ANTHROPIC_",
"OPENAI_",
"LITELLM_",
"AZURE_",
"AWS_",
"GEMINI_",
"CODEX_",
"CURSOR_",
"VERTEX",
# A parent Claude Code session's socket/session vars make a child `claude` attach to
# the parent's login instead of the harness token.
"CLAUDE_CODE_",
"CLAUDE_PID",
"CLAUDECODE",
)
_SECRET_NAMES: Final = frozenset({"GOOGLE_API_KEY", "GOOGLE_APPLICATION_CREDENTIALS"})
_SECRET_SUBSTRINGS: Final = ("API_KEY", "TOKEN", "SECRET")
_TEMPDIR_PREFIX: Final = "litellm-harness-"
def is_secret_env_name(name: str) -> bool:
"""True if an env var name looks like a provider credential."""
upper = name.upper()
if upper in _SECRET_NAMES or upper.startswith(_SECRET_PREFIXES):
return True
return any(part in upper for part in _SECRET_SUBSTRINGS)
def filtered_environ(
base: Mapping[str, str] | None = None,
extra: Mapping[str, str] | None = None,
) -> Mapping[str, str]:
"""base (default os.environ) without provider secrets, then extra on top."""
source = os.environ if base is None else base
kept = ((k, v) for k, v in source.items() if not is_secret_env_name(k))
overlay = extra.items() if extra else ()
return MappingProxyType(dict(itertools.chain(kept, overlay)))
def _signal_process(proc: asyncio.subprocess.Process, sig: int) -> None:
try:
os.killpg(proc.pid, sig)
except (ProcessLookupError, PermissionError, OSError):
try:
proc.send_signal(sig)
except ProcessLookupError:
pass
class SubprocessHandle:
"""Process-protocol wrapper around an asyncio subprocess."""
def __init__(self, proc: asyncio.subprocess.Process) -> None:
if proc.stdout is None or proc.stderr is None:
raise SandboxError("subprocess was started without stdout/stderr pipes")
self._proc = proc
self.stdin: asyncio.StreamWriter | None = proc.stdin
self.stdout: asyncio.StreamReader = proc.stdout
self.stderr: asyncio.StreamReader = proc.stderr
@property
def pid(self) -> int:
return self._proc.pid
@property
def returncode(self) -> int | None:
return self._proc.returncode
async def wait(self) -> int:
return await self._proc.wait()
async def kill(self) -> None:
"""SIGTERM, wait HARNESS_PROCESS_KILL_GRACE_SECONDS, then SIGKILL."""
if self._proc.returncode is not None:
return
_signal_process(self._proc, signal.SIGTERM)
try:
await asyncio.wait_for(self._proc.wait(), timeout=HARNESS_PROCESS_KILL_GRACE_SECONDS)
return
except asyncio.TimeoutError:
pass
_signal_process(self._proc, signal.SIGKILL)
await self._proc.wait()
async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]:
if handle.stdin is not None:
handle.stdin.close()
stdout, stderr = await asyncio.gather(handle.stdout.read(), handle.stderr.read())
exit_code = await handle.wait()
return stdout, stderr, exit_code
async def collect_output(handle: SubprocessHandle, cmd: Sequence[str], timeout: float | None) -> CompletedRun:
"""Close stdin, read stdout/stderr to EOF; kill and raise SandboxError on timeout."""
try:
stdout, stderr, code = await asyncio.wait_for(_read_all(handle), timeout)
except asyncio.TimeoutError:
await handle.kill()
raise SandboxError(f"command timed out after {timeout}s: {cmd[0]}")
return CompletedRun(
stdout=stdout.decode("utf-8", errors="replace"),
stderr=stderr.decode("utf-8", errors="replace"),
exit_code=code,
)
def _is_within(path: str, root: str) -> bool:
return path == root or path.startswith(root.rstrip(os.sep) + os.sep)
class LocalSandbox:
"""Sandbox backed by the local filesystem and asyncio subprocesses."""
def __init__(self, path: str | os.PathLike[str]) -> None:
resolved = os.path.realpath(os.path.abspath(os.fspath(path)))
if not os.path.isdir(resolved):
raise SandboxError(f"sandbox path does not exist or is not a directory: {resolved}")
self.workdir: str = resolved
self._processes: set[SubprocessHandle] = set() # mutable-ok: live-process registry (add/discard)
self._tempdirs: list[str] = [] # mutable-ok: tempdirs created on demand by tempdir(), removed on close()
self._closed = False
def __repr__(self) -> str:
return f"LocalSandbox({self.workdir!r})"
def _check_open(self) -> None:
if self._closed:
raise SandboxError("sandbox is closed")
def _allowed_roots(self) -> tuple[str, ...]:
return (self.workdir, *self._tempdirs)
def resolve_path(self, path: str) -> str:
"""Absolute real path for path; SandboxError if it escapes the sandbox."""
joined = path if os.path.isabs(path) else os.path.join(self.workdir, path)
real = os.path.realpath(joined)
if not any(_is_within(real, root) for root in self._allowed_roots()):
raise SandboxError(f"path escapes the sandbox: {path}")
return real
def _resolve_cwd(self, cwd: str | None) -> str:
if cwd is None:
return self.workdir
resolved = self.resolve_path(cwd)
if not os.path.isdir(resolved):
raise SandboxError(f"cwd is not a directory: {cwd}")
return resolved
def child_env(self, env: Mapping[str, str] | None = None) -> Mapping[str, str]:
return filtered_environ(extra=env)
async def exec(
self,
cmd: Sequence[str],
*,
env: Mapping[str, str] | None = None,
cwd: str | None = None,
) -> SubprocessHandle:
self._check_open()
if not cmd:
raise SandboxError("exec() needs a non-empty command")
try:
proc = await asyncio.create_subprocess_exec(
*cmd,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=self._resolve_cwd(cwd),
env=self.child_env(env),
start_new_session=True,
)
except (FileNotFoundError, PermissionError) as exc:
raise SandboxError(f"could not start {cmd[0]}: {exc}") from exc
handle = SubprocessHandle(proc)
self._processes.add(handle)
return handle
async def run(
self,
cmd: Sequence[str],
*,
env: Mapping[str, str] | None = None,
cwd: str | None = None,
timeout: float | None = None,
) -> CompletedRun:
handle = await self.exec(cmd, env=env, cwd=cwd)
try:
return await collect_output(handle, cmd, timeout)
finally:
self._processes.discard(handle)
async def read(self, path: str) -> bytes:
self._check_open()
resolved = self.resolve_path(path)
try:
return await asyncio.to_thread(_read_bytes, resolved)
except OSError as exc:
raise SandboxError(f"could not read {path}: {exc}") from exc
async def write(self, path: str, data: bytes) -> None:
self._check_open()
resolved = self.resolve_path(path)
try:
await asyncio.to_thread(_write_bytes, resolved, data)
except OSError as exc:
raise SandboxError(f"could not write {path}: {exc}") from exc
def host_url(self, port: int) -> str:
return f"http://127.0.0.1:{port}"
async def which(self, binary: str) -> str | None:
return shutil.which(binary, path=self.child_env().get("PATH"))
async def tempdir(self) -> str:
"""A private temp dir (e.g. for CODEX_HOME), removed on close()."""
self._check_open()
path = os.path.realpath(tempfile.mkdtemp(prefix=_TEMPDIR_PREFIX))
self._tempdirs.append(path)
return path
async def snapshot(self) -> Mapping[str, str]:
self._check_open()
return await snapshot_local(self.workdir)
async def close(self) -> None:
if self._closed:
return
self._closed = True
live = tuple(h for h in self._processes if h.returncode is None)
await asyncio.gather(*(h.kill() for h in live), return_exceptions=True)
self._processes.clear()
for path in self._tempdirs:
shutil.rmtree(path, ignore_errors=True)
self._tempdirs.clear()
async def __aenter__(self) -> LocalSandbox:
return self
async def __aexit__(self, *exc_info: object) -> None:
await self.close()
def _read_bytes(path: str) -> bytes:
with open(path, "rb") as fh:
return fh.read()
def _write_bytes(path: str, data: bytes) -> None:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "wb") as fh:
fh.write(data)
def local(path: str | os.PathLike[str]) -> LocalSandbox:
"""Sandbox rooted at an existing local directory."""
return LocalSandbox(path)

183
litellm/harness/sandbox/snapshot.py Normal file
View file

@ -0,0 +1,183 @@
"""Workspace snapshots and FileChange construction.
A snapshot maps a workspace-relative POSIX path to the sha256 of its contents.
Diffing two snapshots tells us which files a turn created, modified or deleted;
`build_file_changes` turns that into `FileChange` events with unified diffs for
small text files.
"""
from __future__ import annotations
import asyncio
import difflib
import functools
import hashlib
import os
from collections.abc import Iterator, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from litellm.constants import HARNESS_MAX_DIFF_BYTES, HARNESS_SNAPSHOT_SKIP_DIRS
from litellm.harness.errors import HarnessError
from litellm.harness.types import FileChange, FileChangeKind
if TYPE_CHECKING:
from litellm.harness.sandbox.base import Sandbox
# Files larger than this are left out of snapshots entirely.
HARNESS_SNAPSHOT_MAX_FILE_BYTES: Final = 50 * 1024 * 1024
# Upper bound on bytes read by capture_text_contents() for one turn.
HARNESS_SNAPSHOT_MAX_TOTAL_BYTES: Final = 16 * 1024 * 1024
_HASH_CHUNK_BYTES: Final = 1024 * 1024
_NO_NEWLINE_MARKER: Final = "\\ No newline at end of file\n"
def _hash_file(path: str) -> str:
digest = hashlib.sha256()
with open(path, "rb") as fh:
for chunk in iter(functools.partial(fh.read, _HASH_CHUNK_BYTES), b""):
digest.update(chunk)
return digest.hexdigest()
def _hash_entry(root: str, dirpath: str, filename: str) -> tuple[str, str] | None:
full = os.path.join(dirpath, filename)
try:
info = os.lstat(full)
except OSError:
return None
if not os.path.isfile(full) or os.path.islink(full):
return None
if info.st_size > HARNESS_SNAPSHOT_MAX_FILE_BYTES:
return None
try:
digest = _hash_file(full)
except OSError:
return None
rel = os.path.relpath(full, root).replace(os.sep, "/")
return rel, digest
def _walk_entries(root: str) -> Iterator[tuple[str, str]]:
for dirpath, dirnames, filenames in os.walk(root, followlinks=False):
dirnames[:] = [ # mutable-ok: os.walk prunes only via in-place mutation of its dirnames list
d for d in dirnames if d not in HARNESS_SNAPSHOT_SKIP_DIRS
]
for filename in filenames:
entry = _hash_entry(root, dirpath, filename)
if entry is not None:
yield entry
def snapshot_local_sync(root: str) -> Mapping[str, str]:
"""Hash every regular file under root. Symlinks are never followed."""
return MappingProxyType(dict(_walk_entries(root)))
async def snapshot_local(root: str) -> Mapping[str, str]:
"""Async wrapper around snapshot_local_sync (runs in a worker thread)."""
return await asyncio.to_thread(snapshot_local_sync, root)
def _change_kind(path: str, before: Mapping[str, str], after: Mapping[str, str]) -> FileChangeKind | None:
if path not in before:
return "created"
if path not in after:
return "deleted"
if before[path] != after[path]:
return "modified"
return None
def diff_snapshots(
before: Mapping[str, str], after: Mapping[str, str]
) -> list[tuple[str, FileChangeKind]]: # mutable-ok: public sandbox helper; callers compare against a list
"""Return (path, kind) for every changed file, sorted by path."""
kinds = ((path, _change_kind(path, before, after)) for path in sorted(frozenset(before) | frozenset(after)))
return [ # mutable-ok: public sandbox helper returns a list
(path, kind) for path, kind in kinds if kind is not None
]
def _as_text(data: bytes) -> str | None:
if len(data) > HARNESS_MAX_DIFF_BYTES or b"\0" in data:
return None
try:
return data.decode("utf-8")
except UnicodeDecodeError:
return None
def unified_diff(path: str, old: str | None, new: str | None) -> str:
"""Unified diff between two versions of path; None means the file is absent."""
from_file = "/dev/null" if old is None else f"a/{path}"
to_file = "/dev/null" if new is None else f"b/{path}"
lines = difflib.unified_diff(
(old or "").splitlines(keepends=True),
(new or "").splitlines(keepends=True),
fromfile=from_file,
tofile=to_file,
)
return "".join(line if line.endswith("\n") else line + "\n" + _NO_NEWLINE_MARKER for line in lines)
async def _read_or_none(sandbox: Sandbox, path: str) -> bytes | None:
try:
return await sandbox.read(path)
except (HarnessError, OSError):
return None
async def capture_text_contents(sandbox: Sandbox, paths_hashes: Mapping[str, str]) -> Mapping[str, bytes]:
"""Read small text files before a turn so "modified"/"deleted" diffs can be built.
Each kept file is <= HARNESS_MAX_DIFF_BYTES; every byte read (kept or not) counts
toward HARNESS_SNAPSHOT_MAX_TOTAL_BYTES, after which capture stops.
"""
captured: dict[str, bytes] = {} # mutable-ok: async accumulator (awaits per read), frozen on return
total = 0
for path in sorted(paths_hashes):
if total >= HARNESS_SNAPSHOT_MAX_TOTAL_BYTES:
break
data = await _read_or_none(sandbox, path)
if data is None:
continue
total += len(data)
if _as_text(data) is not None:
captured[path] = data
return MappingProxyType(captured)
async def _change_for(
sandbox: Sandbox,
path: str,
kind: FileChangeKind,
before_contents: Mapping[str, bytes],
) -> FileChange:
old_bytes = before_contents.get(path)
old = _as_text(old_bytes) if old_bytes is not None else None
if kind == "deleted":
diff = unified_diff(path, old, None) if old is not None else None
return FileChange(path=path, kind=kind, diff=diff)
new_bytes = await _read_or_none(sandbox, path)
new = _as_text(new_bytes) if new_bytes is not None else None
if new is None or (kind == "modified" and old is None):
return FileChange(path=path, kind=kind, diff=None)
return FileChange(
path=path,
kind=kind,
diff=unified_diff(path, old if kind == "modified" else None, new),
)
async def build_file_changes(
sandbox: Sandbox,
before: Mapping[str, str],
after: Mapping[str, str],
before_contents: Mapping[str, bytes] | None = None,
) -> list[FileChange]: # mutable-ok: feeds the public Result.files list
"""FileChange per changed path. diff is None when it cannot be built as text."""
contents: Mapping[str, bytes] = before_contents or MappingProxyType({})
return [ # mutable-ok: feeds the public Result.files list
await _change_for(sandbox, path, kind, contents) for path, kind in diff_snapshots(before, after)
]

Some files were not shown because too many files have changed in this diff Show more