From e4b893a7b8560870cd10f7ac7f1f2f1556254a89 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 29 Sep 2026 23:34:25 -0700 Subject: [PATCH 01/12] wip --- litellm-rust/Cargo.lock | 17 ++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 5 + .../crates/python-bridge/src/routes/mod.rs | 1 + .../crates/python-bridge/src/routes/traces.rs | 87 +++++++ litellm-rust/crates/traces/Cargo.toml | 18 ++ litellm-rust/crates/traces/README.md | 13 ++ litellm-rust/crates/traces/config/reader.xml | 30 +++ .../traces/migrations/0001_otel_traces.sql | 49 ++++ litellm-rust/crates/traces/src/error.rs | 21 ++ litellm-rust/crates/traces/src/lib.rs | 67 ++++++ litellm-rust/crates/traces/src/sql.rs | 70 ++++++ litellm-rust/crates/traces/tests/admin_sql.rs | 215 ++++++++++++++++++ .../crates/traces/tests/migrations.rs | 61 +++++ litellm-rust/crates/traces/tests/queries.rs | 41 ++++ litellm/proxy/_types.py | 1 + litellm/proxy/logging_endpoints/traces.py | 95 ++++++++ litellm/proxy/proxy_server.py | 3 + litellm/rust_bridge/_native.pyi | 12 + .../proxy/logging_endpoints/test_traces.py | 45 ++++ 21 files changed, 853 insertions(+) create mode 100644 litellm-rust/crates/python-bridge/src/routes/traces.rs create mode 100644 litellm-rust/crates/traces/Cargo.toml create mode 100644 litellm-rust/crates/traces/README.md create mode 100644 litellm-rust/crates/traces/config/reader.xml create mode 100644 litellm-rust/crates/traces/migrations/0001_otel_traces.sql create mode 100644 litellm-rust/crates/traces/src/error.rs create mode 100644 litellm-rust/crates/traces/src/lib.rs create mode 100644 litellm-rust/crates/traces/src/sql.rs create mode 100644 litellm-rust/crates/traces/tests/admin_sql.rs create mode 100644 litellm-rust/crates/traces/tests/migrations.rs create mode 100644 litellm-rust/crates/traces/tests/queries.rs create mode 100644 litellm/proxy/logging_endpoints/traces.py create mode 100644 tests/unit/proxy/logging_endpoints/test_traces.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8d189c8c515a..2d823692e87c 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4075,6 +4075,7 @@ dependencies = [ "litellm-secrets-aws", "litellm-secrets-types", "litellm-token-counter", + "litellm-traces", "litellm-tracing", "pyo3", "pyo3-async-runtimes", @@ -4351,6 +4352,19 @@ dependencies = [ "tiktoken-rs", ] +[[package]] +name = "litellm-traces" +version = "0.1.0" +dependencies = [ + "litellm-http", + "rstest", + "serde_json", + "testcontainers-modules", + "thiserror 2.0.19", + "tokio", + "url", +] + [[package]] name = "litellm-tracing" version = "0.1.0" @@ -5704,6 +5718,7 @@ checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029" dependencies = [ "base64 0.23.1", "bytes", + "encoding_rs", "futures-core", "futures-util", "h2 0.4.15", @@ -5715,6 +5730,7 @@ dependencies = [ "hyper-util", "js-sys", "log", + "mime", "percent-encoding", "pin-project-lite", "quinn", @@ -6945,6 +6961,7 @@ dependencies = [ "memchr", "parse-display", "pin-project-lite", + "reqwest 0.13.5", "serde", "serde_json", "serde_with", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 53aaf7a4d523..257a47268e44 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -12,6 +12,7 @@ repository = "https://github.com/BerriAI/litellm" litellm-config = { path = "crates/config" } litellm-router = { path = "crates/router" } litellm-tracing = { path = "crates/tracing" } +litellm-traces = { path = "crates/traces" } litellm-core = { path = "crates/core" } litellm-gateway-mcp = { path = "crates/gateway-mcp" } litellm-gateway = { path = "crates/gateway" } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index f8ed125f2296..99c95632bb38 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -21,6 +21,7 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] fancy-regex.workspace = true litellm-tracing.workspace = true +litellm-traces.workspace = true litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 37c21cec2de6..dc43475c2af1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -43,6 +43,8 @@ mod _native { use crate::routes::responses::{ResponsesWebSocketConnection, aresponses, responses}; #[pymodule_export] use crate::routes::token_counter::TokenCounter; + #[pymodule_export] + use crate::routes::traces::{get_span, get_trace, list_traces}; #[cfg(feature = "huggingface")] #[pymodule_export] use crate::tokenizer::HuggingFaceEncoding; @@ -107,6 +109,9 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", + "get_span", + "get_trace", + "list_traces", "TokenCounter", "Tokenizer", "gil_stats", diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 0ea10c52c08f..2380274001e6 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -6,6 +6,7 @@ pub(crate) mod messages; pub(crate) mod ocr; pub(crate) mod responses; pub(crate) mod token_counter; +pub(crate) mod traces; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_callbacks_legacy_python::{LegacyLogging, PublicCall}; diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs new file mode 100644 index 000000000000..ef70aa6e01d6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -0,0 +1,87 @@ +use litellm_traces::{Connection, Error, ListQuery, Query}; +use pyo3::{ + exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, + prelude::*, + types::PyAny, +}; + +fn map_error(error: Error) -> PyErr { + match error { + Error::SchemaPending => PyNotImplementedError::new_err(error.to_string()), + Error::InvalidUrl => PyRuntimeError::new_err(error.to_string()), + Error::InvalidListQuery | Error::InvalidIdentifier | Error::EmptySql => { + PyValueError::new_err(error.to_string()) + } + Error::QueryFailed(_) + | Error::ResponseTooLarge + | Error::InvalidResponse + | Error::Transport => PyRuntimeError::new_err(error.to_string()), + } +} + +fn dispatch<'py>( + py: Python<'py>, + clickhouse_url: &str, + request: Query, +) -> PyResult> { + let clickhouse_url = clickhouse_url.to_owned(); + litellm_host_python::run_async( + py, + async move { + let connection = Connection::parse(&clickhouse_url)?; + litellm_traces::query(connection, request).await + }, + map_error, + ) +} + +#[pyfunction] +#[pyo3(signature = (clickhouse_url, start_ms, end_ms, service=None, status=None, search=None, cursor=None, limit=50))] +#[expect( + clippy::too_many_arguments, + reason = "the list API has seven typed query fields" +)] +pub fn list_traces<'py>( + py: Python<'py>, + clickhouse_url: &str, + start_ms: i64, + end_ms: i64, + service: Option, + status: Option, + search: Option, + cursor: Option, + limit: u8, +) -> PyResult> { + dispatch( + py, + clickhouse_url, + Query::List(ListQuery { + start_ms, + end_ms, + service, + status, + search, + cursor, + limit, + }), + ) +} + +#[pyfunction] +pub fn get_trace<'py>( + py: Python<'py>, + clickhouse_url: &str, + trace_id: String, +) -> PyResult> { + dispatch(py, clickhouse_url, Query::Trace { trace_id }) +} + +#[pyfunction] +pub fn get_span<'py>( + py: Python<'py>, + clickhouse_url: &str, + trace_id: String, + span_id: String, +) -> PyResult> { + dispatch(py, clickhouse_url, Query::Span { trace_id, span_id }) +} diff --git a/litellm-rust/crates/traces/Cargo.toml b/litellm-rust/crates/traces/Cargo.toml new file mode 100644 index 000000000000..3c375d2ebe6e --- /dev/null +++ b/litellm-rust/crates/traces/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "litellm-traces" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-http.workspace = true +serde_json.workspace = true +thiserror.workspace = true +url.workspace = true + +[dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } +rstest.workspace = true +testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] } +tokio.workspace = true diff --git a/litellm-rust/crates/traces/README.md b/litellm-rust/crates/traces/README.md new file mode 100644 index 000000000000..e3d4e2bc50f6 --- /dev/null +++ b/litellm-rust/crates/traces/README.md @@ -0,0 +1,13 @@ +# Trace queries + +Use a dedicated ClickHouse reader account for `general_settings.clickhouse_url`. The SQL helper sends the supplied SQL to ClickHouse without parsing it. Its `readonly=1` request setting does not make an account with write or administrative grants safe + +Pass `execute_admin_sql` a shared HTTP client obtained from `HttpClientPool` with `ClientVariant::NoRedirect` and the host's HTTP configuration. The helper applies a 15-second request timeout and validates the bounded JSON response, including errors returned with HTTP 200 + +For self-hosted ClickHouse, install [config/reader.xml](config/reader.xml) as `/etc/clickhouse-server/users.d/litellm-traces-reader.xml` and provide `LITELLM_TRACES_READER_PASSWORD` in the ClickHouse server environment. Restrict the configured network range to your proxy network. The configuration grants SELECT only on `default.otel_traces` and locks read-only mode, execution time, result rows and bytes, overflow behavior and memory limits. Adjust the database name if traces live elsewhere + +Set `clickhouse_url` to `https://litellm_traces_reader:@:8443/?database=default`. Keep ingestion and migration credentials separate. For managed ClickHouse, provision the equivalent SELECT grant and settings constraints through its access-management interface + +Do not grant write, backup, named-collection management or grant-option privileges to the reader, directly or through roles. Do not make `readonly` changeable in read-only mode. ClickHouse documents exceptions to `readonly` in [query permissions](https://clickhouse.com/docs/concepts/features/configuration/settings/permissions-for-queries) and explains locked settings in [settings constraints](https://clickhouse.com/docs/concepts/features/configuration/settings/constraints-on-settings) + +Run the reader permission, authentication and response-limit tests with `cargo test -p litellm-traces --test admin_sql -- --test-threads=2` from `litellm-rust`. They start disposable ClickHouse instances through Testcontainers and load the same reader configuration diff --git a/litellm-rust/crates/traces/config/reader.xml b/litellm-rust/crates/traces/config/reader.xml new file mode 100644 index 000000000000..704f56255bbc --- /dev/null +++ b/litellm-rust/crates/traces/config/reader.xml @@ -0,0 +1,30 @@ + + + + 1 + 10 + 1000 + 4194304 + throw + 268435456 + + + + + + + + + + + + + + ::/0 + litellm_traces_reader + + GRANT SELECT ON default.otel_traces + + + + diff --git a/litellm-rust/crates/traces/migrations/0001_otel_traces.sql b/litellm-rust/crates/traces/migrations/0001_otel_traces.sql new file mode 100644 index 000000000000..0dc54ac2942c --- /dev/null +++ b/litellm-rust/crates/traces/migrations/0001_otel_traces.sql @@ -0,0 +1,49 @@ +CREATE TABLE IF NOT EXISTS otel_traces +( + Timestamp DateTime64(9) CODEC(Delta, ZSTD(1)), + TraceId String CODEC(ZSTD(1)), + SpanId String CODEC(ZSTD(1)), + ParentSpanId String CODEC(ZSTD(1)), + TraceState String CODEC(ZSTD(1)), + SpanName LowCardinality(String) CODEC(ZSTD(1)), + SpanKind LowCardinality(String) CODEC(ZSTD(1)), + ServiceName LowCardinality(String) CODEC(ZSTD(1)), + ResourceAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), + ScopeName String CODEC(ZSTD(1)), + ScopeVersion String CODEC(ZSTD(1)), + SpanAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), + Duration UInt64 CODEC(ZSTD(1)), + StatusCode LowCardinality(String) CODEC(ZSTD(1)), + StatusMessage String CODEC(ZSTD(1)), + `Events.Timestamp` Array(DateTime64(9)) CODEC(ZSTD(1)), + `Events.Name` Array(LowCardinality(String)) CODEC(ZSTD(1)), + `Events.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), + `Links.TraceId` Array(String) CODEC(ZSTD(1)), + `Links.SpanId` Array(String) CODEC(ZSTD(1)), + `Links.TraceState` Array(String) CODEC(ZSTD(1)), + `Links.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), + TeamId LowCardinality(String) DEFAULT SpanAttributes['litellm.team_id'], + ProjectId String DEFAULT SpanAttributes['litellm.project_id'], + UserId String DEFAULT SpanAttributes['litellm.user_id'], + ApiKeyHash String DEFAULT SpanAttributes['litellm.api_key_hash'], + ObservationType LowCardinality(String) DEFAULT multiIf( + ParentSpanId = '', 'agent', + SpanAttributes['gen_ai.operation.name'] = 'invoke_agent', 'agent', + SpanAttributes['gen_ai.operation.name'] IN ('chat', 'text_completion', 'generate_content'), 'llm', + SpanAttributes['gen_ai.operation.name'] = 'execute_tool', 'tool', + 'chain' + ), + AgentName LowCardinality(String) DEFAULT SpanAttributes['gen_ai.agent.name'], + LiteLLMRequestId String DEFAULT SpanAttributes['gen_ai.response.id'], + Model LowCardinality(String) DEFAULT SpanAttributes['gen_ai.request.model'], + InputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.input_tokens']), + OutputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.output_tokens']), + Input String CODEC(ZSTD(3)), + Output String CODEC(ZSTD(3)), + InputPreview String DEFAULT substring(Input, 1, 240), + 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) diff --git a/litellm-rust/crates/traces/src/error.rs b/litellm-rust/crates/traces/src/error.rs new file mode 100644 index 000000000000..c0e70b09aa8b --- /dev/null +++ b/litellm-rust/crates/traces/src/error.rs @@ -0,0 +1,21 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid ClickHouse HTTP URL")] + InvalidUrl, + #[error("trace query requires start before end and a limit from 1 to 100")] + InvalidListQuery, + #[error("trace query requires nonempty identifiers")] + InvalidIdentifier, + #[error("trace queries require the ClickHouse table schema")] + SchemaPending, + #[error("SQL query must not be empty")] + EmptySql, + #[error("ClickHouse query failed with HTTP status {0}")] + QueryFailed(u16), + #[error("ClickHouse query exceeded the response size limit")] + ResponseTooLarge, + #[error("ClickHouse returned an invalid or failed JSON query response")] + InvalidResponse, + #[error("ClickHouse query transport failed")] + Transport, +} diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs new file mode 100644 index 000000000000..c5764c7c5695 --- /dev/null +++ b/litellm-rust/crates/traces/src/lib.rs @@ -0,0 +1,67 @@ +mod error; +mod sql; + +pub use error::Error; +pub use sql::execute_admin_sql; +use url::Url; + +#[derive(Clone)] +pub struct Connection { + url: Url, +} + +impl Connection { + pub fn parse(value: &str) -> Result { + let url = Url::parse(value).map_err(|_| Error::InvalidUrl)?; + if !matches!(url.scheme(), "http" | "https") || url.host().is_none() { + return Err(Error::InvalidUrl); + } + Ok(Self { url }) + } + + pub fn url(&self) -> &Url { + &self.url + } +} + +#[derive(Debug, Clone)] +pub struct ListQuery { + pub start_ms: i64, + pub end_ms: i64, + pub service: Option, + pub status: Option, + pub search: Option, + pub cursor: Option, + pub limit: u8, +} + +#[derive(Debug, Clone)] +pub enum Query { + List(ListQuery), + Trace { trace_id: String }, + Span { trace_id: String, span_id: String }, +} + +impl Query { + pub fn validate(&self) -> Result<(), Error> { + match self { + Self::List(query) + if query.start_ms >= query.end_ms || !(1..=100).contains(&query.limit) => + { + Err(Error::InvalidListQuery) + } + Self::Trace { trace_id } + | Self::Span { + trace_id, + span_id: _, + } if trace_id.is_empty() => Err(Error::InvalidIdentifier), + Self::Span { span_id, .. } if span_id.is_empty() => Err(Error::InvalidIdentifier), + _ => Ok(()), + } + } +} + +pub async fn query(_connection: Connection, request: Query) -> Result<(), Error> { + request.validate()?; + Err(Error::SchemaPending) +} diff --git a/litellm-rust/crates/traces/src/sql.rs b/litellm-rust/crates/traces/src/sql.rs new file mode 100644 index 000000000000..408c57914091 --- /dev/null +++ b/litellm-rust/crates/traces/src/sql.rs @@ -0,0 +1,70 @@ +use std::time::Duration; + +use litellm_http::Client; + +use crate::{Connection, Error}; + +const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024; + +pub async fn execute_admin_sql( + client: &Client, + connection: &Connection, + sql: &str, +) -> Result { + if sql.trim().is_empty() { + return Err(Error::EmptySql); + } + + let mut url = connection.url().clone(); + + let existing_pairs: Vec<(String, String)> = url + .query_pairs() + .filter(|(key, _)| { + !matches!( + key.as_ref(), + "query" + | "readonly" + | "default_format" + | "max_result_rows" + | "result_overflow_mode" + | "max_execution_time" + | "wait_end_of_query" + ) + }) + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect(); + url.query_pairs_mut() + .clear() + .extend_pairs(existing_pairs) + .append_pair("readonly", "1") + .append_pair("max_result_rows", "1000") + .append_pair("result_overflow_mode", "throw") + .append_pair("max_execution_time", "10") + .append_pair("wait_end_of_query", "1") + .append_pair("default_format", "JSON"); + + let request = client + .post(url) + .timeout(Duration::from_secs(15)) + .body(sql.to_owned()); + let mut response = request.send().await.map_err(|_| Error::Transport)?; + if !response.status().is_success() { + return Err(Error::QueryFailed(response.status().as_u16())); + } + + let mut body = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? { + if body.len() + chunk.len() > MAX_RESPONSE_BYTES { + return Err(Error::ResponseTooLarge); + } + body.extend_from_slice(&chunk); + } + + let json: serde_json::Value = + serde_json::from_slice(&body).map_err(|_| Error::InvalidResponse)?; + if json.get("exception").is_some() || !json.get("data").is_some_and(serde_json::Value::is_array) + { + return Err(Error::InvalidResponse); + } + String::from_utf8(body).map_err(|_| Error::InvalidResponse) +} diff --git a/litellm-rust/crates/traces/tests/admin_sql.rs b/litellm-rust/crates/traces/tests/admin_sql.rs new file mode 100644 index 000000000000..44326e15278b --- /dev/null +++ b/litellm-rust/crates/traces/tests/admin_sql.rs @@ -0,0 +1,215 @@ +use litellm_http::Client; +use litellm_traces::{Connection, Error, execute_admin_sql}; +use rstest::{fixture, rstest}; +use serde_json::Value; +use testcontainers_modules::{ + clickhouse::ClickHouse, + testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner}, +}; + +const CLICKHOUSE_TAG: &str = + "26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e"; + +struct Database { + _container: ContainerAsync, + url: String, + admin_url: String, + client: Client, +} + +#[fixture] +async fn database() -> Result> { + let container = ClickHouse::default() + .with_tag(CLICKHOUSE_TAG) + .with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1") + .with_env_var("LITELLM_TRACES_READER_PASSWORD", "test_password") + .with_copy_to( + "/etc/clickhouse-server/users.d/litellm-traces-reader.xml", + include_bytes!("../config/reader.xml").to_vec(), + ) + .start() + .await?; + let admin_url = format!( + "http://{}:{}", + container.get_host().await?, + container.get_host_port_ipv4(8123).await?, + ); + let client = Client::no_redirect_for_test(); + for sql in [ + "CREATE TABLE otel_traces (n UInt8) ENGINE = Memory", + "INSERT INTO otel_traces VALUES (1)", + "CREATE TABLE private_traces (n UInt8) ENGINE = Memory", + ] { + client + .post(&admin_url) + .body(sql) + .send() + .await? + .error_for_status()?; + } + let url = admin_url.replacen("http://", "http://litellm_traces_reader:test_password@", 1); + Ok(Database { + _container: container, + url, + admin_url, + client, + }) +} + +#[rstest] +#[tokio::test] +async fn admin_sql_reads_rows_with_enforced_settings( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!( + "{}?readonly=0&default_format=TabSeparated&query=SELECT+2", + database.url, + ))?; + + let result = execute_admin_sql( + &database.client, + &connection, + "SELECT n AS answer FROM otel_traces", + ) + .await?; + let json: Value = serde_json::from_str(&result)?; + assert_eq!(json["data"][0]["answer"], 1); + + Ok(()) +} + +#[rstest] +#[case::table("CREATE TABLE admin_sql_test (n UInt8) ENGINE = Memory")] +#[case::insert("INSERT INTO otel_traces VALUES (2)")] +#[case::drop("DROP TABLE otel_traces")] +#[case::named_collection("CREATE NAMED COLLECTION admin_sql_test AS host = 'localhost'")] +#[case::settings("SET readonly = 0")] +#[case::inline_settings("SELECT n FROM otel_traces SETTINGS readonly = 0")] +#[case::time_limit("SELECT n FROM otel_traces SETTINGS max_execution_time = 0")] +#[case::row_limit("SELECT n FROM otel_traces SETTINGS max_result_rows = 0")] +#[case::byte_limit("SELECT n FROM otel_traces SETTINGS max_result_bytes = 0")] +#[case::memory_limit("SELECT n FROM otel_traces SETTINGS max_memory_usage = 0")] +#[case::other_table("SELECT * FROM private_traces")] +#[tokio::test] +async fn reader_rejects_writes_and_privilege_escalation( + #[future(awt)] database: Result>, + #[case] sql: &str, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!("{}?readonly=0", database.url))?; + + let result = execute_admin_sql(&database.client, &connection, sql).await; + + assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}"); + let rows = + execute_admin_sql(&database.client, &connection, "SELECT n FROM otel_traces").await?; + let json: Value = serde_json::from_str(&rows)?; + assert_eq!(json["data"], serde_json::json!([{ "n": 1 }])); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn admin_sql_rejects_errors_after_output_starts( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!( + "{}?max_block_size=1&buffer_size=1&http_write_exception_in_output_format=1\ + &send_progress_in_http_headers=1&http_headers_progress_interval_ms=0", + database.admin_url, + ))?; + + let result = execute_admin_sql( + &database.client, + &connection, + "SELECT sleepEachRow(0.2), throwIf(number = 2) FROM numbers(5)", + ) + .await; + + assert!( + matches!(result, Err(Error::InvalidResponse)), + "expected an error embedded in a successful HTTP response: {result:?}" + ); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn admin_sql_enforces_result_row_limit( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!( + "{}?max_result_rows=0&result_overflow_mode=throw&wait_end_of_query=1", + database.url, + ))?; + + let result = execute_admin_sql( + &database.client, + &connection, + "SELECT number FROM numbers(1001)", + ) + .await; + + assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}"); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn admin_sql_enforces_response_byte_limit( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&database.admin_url)?; + + let result = execute_admin_sql( + &database.client, + &connection, + "SELECT repeat('x', 5 * 1024 * 1024)", + ) + .await; + + assert!(matches!(result, Err(Error::ResponseTooLarge)), "{result:?}"); + Ok(()) +} + +#[rstest] +#[case::plain("test_password", "test_password")] +#[case::encoded("p@ss/word%", "p%40ss%2Fword%25")] +#[tokio::test] +async fn admin_sql_authenticates_url_credentials( + #[future(awt)] database: Result>, + #[case] password: &str, + #[case] encoded_password: &str, +) -> Result<(), Box> { + let database = database?; + database + .client + .post(&database.admin_url) + .body(format!( + "CREATE USER sql_reader IDENTIFIED WITH plaintext_password BY '{password}'" + )) + .send() + .await? + .error_for_status()?; + let connection = Connection::parse(&database.admin_url.replacen( + "http://", + &format!("http://sql_reader:{encoded_password}@"), + 1, + ))?; + + let result = execute_admin_sql( + &database.client, + &connection, + "SELECT currentUser() AS username", + ) + .await?; + let json: Value = serde_json::from_str(&result)?; + + assert_eq!(json["data"][0]["username"], "sql_reader"); + + Ok(()) +} diff --git a/litellm-rust/crates/traces/tests/migrations.rs b/litellm-rust/crates/traces/tests/migrations.rs new file mode 100644 index 000000000000..b10c508d275a --- /dev/null +++ b/litellm-rust/crates/traces/tests/migrations.rs @@ -0,0 +1,61 @@ +use litellm_http::Client; +use testcontainers_modules::{ + clickhouse::ClickHouse, + testcontainers::{ImageExt, runners::AsyncRunner}, +}; + +const CLICKHOUSE_TAG: &str = + "26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e"; +const INIT_SQL: &str = include_str!("../migrations/0001_otel_traces.sql"); + +#[rstest::rstest] +#[tokio::test] +async fn migration_creates_a_queryable_trace_table() -> Result<(), Box> { + let container = ClickHouse::default() + .with_tag(CLICKHOUSE_TAG) + .with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1") + .start() + .await?; + let url = format!( + "http://{}:{}", + container.get_host().await?, + container.get_host_port_ipv4(8123).await?, + ); + let client = Client::no_redirect_for_test(); + + for _ in 0..2 { + client + .post(&url) + .body(INIT_SQL) + .send() + .await? + .error_for_status()?; + } + + client + .post(&url) + .body( + "INSERT INTO otel_traces \ + (Timestamp, TraceId, SpanId, ParentSpanId, ServiceName, SpanName, SpanAttributes, Input) \ + VALUES (now64(9), 'trace-1', 'span-1', '', 'proxy', 'request', \ + map('litellm.team_id', 'team-1', 'litellm.api_key_hash', 'hash-1'), 'hello world')", + ) + .send() + .await? + .error_for_status()?; + + let row = client + .post(&url) + .body( + "SELECT TeamId, ApiKeyHash, ObservationType, InputPreview \ + FROM otel_traces WHERE TraceId = 'trace-1' FORMAT TSV", + ) + .send() + .await? + .error_for_status()? + .text() + .await?; + assert_eq!(row, "team-1\thash-1\tagent\thello world\n"); + + Ok(()) +} diff --git a/litellm-rust/crates/traces/tests/queries.rs b/litellm-rust/crates/traces/tests/queries.rs new file mode 100644 index 000000000000..eaa9b7578f48 --- /dev/null +++ b/litellm-rust/crates/traces/tests/queries.rs @@ -0,0 +1,41 @@ +use litellm_traces::{Connection, Error, ListQuery, Query}; +use rstest::rstest; + +fn list(start_ms: i64, end_ms: i64, limit: u8) -> Query { + Query::List(ListQuery { + start_ms, + end_ms, + service: None, + status: None, + search: None, + cursor: None, + limit, + }) +} + +#[rstest] +#[case::reversed_window(list(2, 1, 50))] +#[case::empty_window(list(1, 1, 50))] +#[case::zero_limit(list(1, 2, 0))] +#[case::large_limit(list(1, 2, 101))] +#[case::empty_trace(Query::Trace { trace_id: String::new() })] +#[case::empty_span(Query::Span { trace_id: "trace".into(), span_id: String::new() })] +fn rejects_invalid_queries(#[case] request: Query) { + assert!(request.validate().is_err()); +} + +#[rstest] +#[case::http("http://localhost:8123", true)] +#[case::https("https://localhost:8443", true)] +#[case::tcp("tcp://localhost:9000", false)] +#[case::missing_host("http://", false)] +fn accepts_only_clickhouse_http_urls(#[case] value: &str, #[case] expected: bool) { + assert_eq!(Connection::parse(value).is_ok(), expected); +} + +#[tokio::test] +async fn schema_pending_is_explicit_after_valid_request() { + let connection = Connection::parse("http://localhost:8123").expect("valid URL"); + let result = litellm_traces::query(connection, list(1, 2, 50)).await; + assert!(matches!(result, Err(Error::SchemaPending))); +} diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index be06d2e73217..776783206e78 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2624,6 +2624,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ + clickhouse_url: str | None = Field(None, description="ClickHouse HTTP URL for trace queries") completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls") max_in_flight_requests_per_worker: int | None = Field( None, gt=0, description="maximum concurrent requests handled by each worker" diff --git a/litellm/proxy/logging_endpoints/traces.py b/litellm/proxy/logging_endpoints/traces.py new file mode 100644 index 000000000000..ac82f8cb0227 --- /dev/null +++ b/litellm/proxy/logging_endpoints/traces.py @@ -0,0 +1,95 @@ +from collections.abc import Awaitable +from datetime import datetime, timezone +from typing import Final, Protocol, cast + +from fastapi import APIRouter, Depends, HTTPException, Query + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.rust_bridge.loader import get_native_bridge + +router: Final = APIRouter(prefix="/v1/traces", tags=["Traces"]) + + +class NativeTraces(Protocol): + def list_traces( + self, + clickhouse_url: str, + start_ms: int, + end_ms: int, + service: str | None, + status: str | None, + search: str | None, + cursor: str | None, + limit: int, + ) -> Awaitable[object]: ... + + def get_trace(self, clickhouse_url: str, trace_id: str) -> Awaitable[object]: ... + + def get_span(self, clickhouse_url: str, trace_id: str, span_id: str) -> Awaitable[object]: ... + + +def _connection(user_api_key_dict: UserAPIKeyAuth) -> tuple[str, NativeTraces]: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException(status_code=403, detail="Trace queries require proxy admin access") + + from litellm.proxy.proxy_server import general_settings_view + + clickhouse_url: Final = general_settings_view().get("clickhouse_url") + if not isinstance(clickhouse_url, str) or not clickhouse_url: + raise HTTPException(status_code=503, detail="ClickHouse trace queries are not configured") + + native: Final = get_native_bridge() + if native is None or not all(hasattr(native, name) for name in ("list_traces", "get_trace", "get_span")): + raise HTTPException(status_code=503, detail="Rust trace query bridge is unavailable") + return clickhouse_url, cast(NativeTraces, native) + + +async def _result(awaitable: Awaitable[object]) -> object: + try: + return await awaitable + except NotImplementedError as exc: + raise HTTPException(status_code=501, detail=str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except RuntimeError as exc: + raise HTTPException(status_code=503, detail=str(exc)) from exc + + +def _utc_ms(value: datetime) -> int: + utc_value: Final = value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) + return int(utc_value.timestamp() * 1000) + + +@router.get("") +async def list_traces( + start: datetime, + end: datetime, + service: str | None = None, + status: str | None = None, + q: str | None = None, + cursor: str | None = None, + limit: int = Query(default=50, ge=1, le=100), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> object: + clickhouse_url, native = _connection(user_api_key_dict) + return await _result(native.list_traces(clickhouse_url, _utc_ms(start), _utc_ms(end), service, status, q, cursor, limit)) + + +@router.get("/{trace_id}") +async def get_trace( + trace_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> object: + clickhouse_url, native = _connection(user_api_key_dict) + return await _result(native.get_trace(clickhouse_url, trace_id)) + + +@router.get("/{trace_id}/spans/{span_id}") +async def get_span( + trace_id: str, + span_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> object: + clickhouse_url, native = _connection(user_api_key_dict) + return await _result(native.get_span(clickhouse_url, trace_id, span_id)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a994293a4d89..d777d3dc9066 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -568,6 +568,7 @@ def generate_feedback_box(): from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( rust_control_plane_router, ) +from litellm.proxy.logging_endpoints.traces import router as traces_router from litellm.proxy.management_endpoints.auto_router_endpoints import ( router as auto_router_management_router, ) @@ -972,6 +973,7 @@ def generate_feedback_box(): _EXTRA_SECRET_GENERAL_SETTINGS_FIELDS: Final = frozenset( { "database_url", + "clickhouse_url", "database_extra_connection_params", "pass_through_endpoints", "alert_to_webhook_url", @@ -19884,6 +19886,7 @@ async def get_routes(): app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) app.include_router(rust_control_plane_router) +app.include_router(traces_router) app.include_router(ui_crud_endpoints_router) app.include_router(user_banner_endpoints_router) app.include_router(latest_release_endpoints_router) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 6a5798898698..173ee08279a3 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -19,6 +19,18 @@ class RustBridgeDeclined(Exception): ... class RustUpstreamError(Exception): ... class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... +def list_traces( + clickhouse_url: str, + start_ms: int, + end_ms: int, + service: str | None = None, + status: str | None = None, + search: str | None = None, + cursor: str | None = None, + limit: int = 50, +) -> Future[None]: ... +def get_trace(clickhouse_url: str, trace_id: str) -> Future[None]: ... +def get_span(clickhouse_url: str, trace_id: str, span_id: str) -> Future[None]: ... @final class NativeDiagnosticProcessor: diff --git a/tests/unit/proxy/logging_endpoints/test_traces.py b/tests/unit/proxy/logging_endpoints/test_traces.py new file mode 100644 index 000000000000..2c92e47d78bb --- /dev/null +++ b/tests/unit/proxy/logging_endpoints/test_traces.py @@ -0,0 +1,45 @@ +import asyncio +from collections.abc import Callable, Coroutine +from datetime import datetime, timezone + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.logging_endpoints.traces import get_span, get_trace, list_traces + + +def _list(auth: UserAPIKeyAuth) -> Coroutine[object, object, object]: + return list_traces( + start=datetime(2026, 1, 1, tzinfo=timezone.utc), + end=datetime(2026, 1, 2, tzinfo=timezone.utc), + limit=50, + user_api_key_dict=auth, + ) + + +def _trace(auth: UserAPIKeyAuth) -> Coroutine[object, object, object]: + return get_trace(trace_id="trace", user_api_key_dict=auth) + + +def _span(auth: UserAPIKeyAuth) -> Coroutine[object, object, object]: + return get_span(trace_id="trace", span_id="span", user_api_key_dict=auth) + + +@pytest.mark.parametrize( + "run_route", + ( + _list, + _trace, + _span, + ), +) +def test_trace_routes_reject_non_admin_callers_before_accessing_storage( + run_route: Callable[[UserAPIKeyAuth], Coroutine[object, object, object]], +) -> None: + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) + + with pytest.raises(HTTPException) as error: + asyncio.run(run_route(auth)) + + assert error.value.status_code == 403 From 28e24077ff3f25aaf713db07a859aeb1e7962348 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 07:11:55 -0700 Subject: [PATCH 02/12] feat(traces): establish shared Rust storage foundation --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 7 +- .../crates/python-bridge/src/routes/traces.rs | 99 +++++++------------ litellm-rust/crates/traces/Cargo.toml | 1 + litellm-rust/crates/traces/README.md | 16 +-- litellm-rust/crates/traces/config/reader.xml | 2 + .../traces/migrations/0001_otel_traces.sql | 83 ++++++++-------- .../traces/migrations/0002_agent_traces.sql | 25 +++++ .../migrations/0003_agent_traces_mv.sql | 21 ++++ .../traces/migrations/0004_spend_logs.sql | 43 ++++++++ litellm-rust/crates/traces/src/error.rs | 8 +- litellm-rust/crates/traces/src/lib.rs | 76 +++++++------- litellm-rust/crates/traces/src/schema.rs | 37 +++++++ litellm-rust/crates/traces/src/sql.rs | 61 +++++++++--- litellm-rust/crates/traces/tests/admin_sql.rs | 58 +++++++++-- .../crates/traces/tests/migrations.rs | 97 ++++++++++++------ litellm-rust/crates/traces/tests/queries.rs | 32 +----- litellm/proxy/_types.py | 1 - litellm/proxy/logging_endpoints/traces.py | 95 ------------------ litellm/proxy/proxy_server.py | 3 - litellm/rust_bridge/_native.pyi | 22 ++--- litellm/rust_bridge/traces.py | 50 ++++++++++ tests/test_litellm_rust/test_traces.py | 46 +++++++++ .../proxy/logging_endpoints/test_traces.py | 45 --------- 24 files changed, 523 insertions(+), 406 deletions(-) create mode 100644 litellm-rust/crates/traces/migrations/0002_agent_traces.sql create mode 100644 litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql create mode 100644 litellm-rust/crates/traces/migrations/0004_spend_logs.sql create mode 100644 litellm-rust/crates/traces/src/schema.rs delete mode 100644 litellm/proxy/logging_endpoints/traces.py create mode 100644 litellm/rust_bridge/traces.py create mode 100644 tests/test_litellm_rust/test_traces.py delete mode 100644 tests/unit/proxy/logging_endpoints/test_traces.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2d823692e87c..236e3bb0c008 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4358,6 +4358,7 @@ version = "0.1.0" dependencies = [ "litellm-http", "rstest", + "serde", "serde_json", "testcontainers-modules", "thiserror 2.0.19", diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index dc43475c2af1..8b55a40fa226 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -44,7 +44,7 @@ mod _native { #[pymodule_export] use crate::routes::token_counter::TokenCounter; #[pymodule_export] - use crate::routes::traces::{get_span, get_trace, list_traces}; + use crate::routes::traces::{trace_query, trace_schema_statements}; #[cfg(feature = "huggingface")] #[pymodule_export] use crate::tokenizer::HuggingFaceEncoding; @@ -109,9 +109,8 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", - "get_span", - "get_trace", - "list_traces", + "trace_query", + "trace_schema_statements", "TokenCounter", "Tokenizer", "gil_stats", diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index ef70aa6e01d6..557bde793ee4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -1,87 +1,58 @@ -use litellm_traces::{Connection, Error, ListQuery, Query}; +use std::collections::BTreeMap; + +use litellm_http::ClientVariant; +use litellm_traces::{Connection, Error, Parameter}; use pyo3::{ - exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, + exceptions::{PyRuntimeError, PyValueError}, prelude::*, - types::PyAny, }; fn map_error(error: Error) -> PyErr { match error { - Error::SchemaPending => PyNotImplementedError::new_err(error.to_string()), - Error::InvalidUrl => PyRuntimeError::new_err(error.to_string()), - Error::InvalidListQuery | Error::InvalidIdentifier | Error::EmptySql => { - PyValueError::new_err(error.to_string()) - } - Error::QueryFailed(_) + Error::InvalidSchema | Error::EmptySql => PyValueError::new_err(error.to_string()), + Error::InvalidUrl + | Error::QueryFailed(_) | Error::ResponseTooLarge | Error::InvalidResponse | Error::Transport => PyRuntimeError::new_err(error.to_string()), } } -fn dispatch<'py>( - py: Python<'py>, - clickhouse_url: &str, - request: Query, -) -> PyResult> { - let clickhouse_url = clickhouse_url.to_owned(); - litellm_host_python::run_async( - py, - async move { - let connection = Connection::parse(&clickhouse_url)?; - litellm_traces::query(connection, request).await - }, - map_error, - ) +#[pyfunction] +pub fn trace_schema_statements( + py: Python<'_>, + database: String, + trace_retention_days: u32, + spend_log_retention_days: u32, +) -> PyResult> { + py.detach(|| { + litellm_traces::schema_statements(&database, trace_retention_days, spend_log_retention_days) + }) + .map_err(map_error) } #[pyfunction] -#[pyo3(signature = (clickhouse_url, start_ms, end_ms, service=None, status=None, search=None, cursor=None, limit=50))] #[expect( clippy::too_many_arguments, - reason = "the list API has seven typed query fields" + reason = "the internal adapter supplies connection fields and a parameterized query" )] -pub fn list_traces<'py>( +pub fn trace_query<'py>( py: Python<'py>, - clickhouse_url: &str, - start_ms: i64, - end_ms: i64, - service: Option, - status: Option, - search: Option, - cursor: Option, - limit: u8, + url: &str, + database: &str, + user: &str, + password: &str, + sql: String, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] parameters: BTreeMap< + String, + Parameter, + >, ) -> PyResult> { - dispatch( + let connection = Connection::configured(url, database, user, password).map_err(map_error)?; + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + litellm_host_python::run_async( py, - clickhouse_url, - Query::List(ListQuery { - start_ms, - end_ms, - service, - status, - search, - cursor, - limit, - }), + async move { litellm_traces::execute_read(&client, &connection, &sql, ¶meters).await }, + map_error, ) } - -#[pyfunction] -pub fn get_trace<'py>( - py: Python<'py>, - clickhouse_url: &str, - trace_id: String, -) -> PyResult> { - dispatch(py, clickhouse_url, Query::Trace { trace_id }) -} - -#[pyfunction] -pub fn get_span<'py>( - py: Python<'py>, - clickhouse_url: &str, - trace_id: String, - span_id: String, -) -> PyResult> { - dispatch(py, clickhouse_url, Query::Span { trace_id, span_id }) -} diff --git a/litellm-rust/crates/traces/Cargo.toml b/litellm-rust/crates/traces/Cargo.toml index 3c375d2ebe6e..87d2709c9063 100644 --- a/litellm-rust/crates/traces/Cargo.toml +++ b/litellm-rust/crates/traces/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-http.workspace = true +serde.workspace = true serde_json.workspace = true thiserror.workspace = true url.workspace = true diff --git a/litellm-rust/crates/traces/README.md b/litellm-rust/crates/traces/README.md index e3d4e2bc50f6..657bdc58d6ce 100644 --- a/litellm-rust/crates/traces/README.md +++ b/litellm-rust/crates/traces/README.md @@ -1,13 +1,15 @@ -# Trace queries +# Trace storage foundation -Use a dedicated ClickHouse reader account for `general_settings.clickhouse_url`. The SQL helper sends the supplied SQL to ClickHouse without parsing it. Its `readonly=1` request setting does not make an account with write or administrative grants safe +This crate owns the initial ClickHouse schema and parameterized read transport. Python adapters use `trace_schema_statements` and `trace_query` from the native bridge. HTTP endpoints, ingestion and trace response assembly belong to the integration built on this foundation -Pass `execute_admin_sql` a shared HTTP client obtained from `HttpClientPool` with `ClientVariant::NoRedirect` and the host's HTTP configuration. The helper applies a 15-second request timeout and validates the bounded JSON response, including errors returned with HTTP 200 +The SQL files under `migrations/` are the canonical definitions for `otel_traces`, `agent_traces`, its materialized view, and `spend_logs`. `schema_statements` supplies the database and retention settings. Apply the returned statements in order using migration credentials. Setup is repeatable for a new installation; these initial `CREATE IF NOT EXISTS` statements do not upgrade an existing incompatible table. Future schema changes need explicit upgrade statements -For self-hosted ClickHouse, install [config/reader.xml](config/reader.xml) as `/etc/clickhouse-server/users.d/litellm-traces-reader.xml` and provide `LITELLM_TRACES_READER_PASSWORD` in the ClickHouse server environment. Restrict the configured network range to your proxy network. The configuration grants SELECT only on `default.otel_traces` and locks read-only mode, execution time, result rows and bytes, overflow behavior and memory limits. Adjust the database name if traces live elsewhere +`execute_read` uses ClickHouse typed placeholders and binds text, integer and string-array parameters separately from SQL. It obtains no privileges itself: callers must use a dedicated SELECT-only reader. The helper enforces a 15-second HTTP timeout, a 10-second query setting, 1,000 result rows and a 4 MiB response cap, and rejects errors embedded in HTTP 200 JSON responses. Large traces fail explicitly at those limits; span pagination is future work -Set `clickhouse_url` to `https://litellm_traces_reader:@:8443/?database=default`. Keep ingestion and migration credentials separate. For managed ClickHouse, provision the equivalent SELECT grant and settings constraints through its access-management interface +Pass a shared `HttpClientPool` client using `ClientVariant::NoRedirect` and the host HTTP settings. The native bridge already does this. The SQL helper is internal infrastructure, not a public arbitrary-SQL endpoint -Do not grant write, backup, named-collection management or grant-option privileges to the reader, directly or through roles. Do not make `readonly` changeable in read-only mode. ClickHouse documents exceptions to `readonly` in [query permissions](https://clickhouse.com/docs/concepts/features/configuration/settings/permissions-for-queries) and explains locked settings in [settings constraints](https://clickhouse.com/docs/concepts/features/configuration/settings/constraints-on-settings) +For self-hosted ClickHouse, install [config/reader.xml](config/reader.xml) as `/etc/clickhouse-server/users.d/litellm-traces-reader.xml` and set `LITELLM_TRACES_READER_PASSWORD` on the server. Change `default` in the three SELECT grants to the configured database and restrict the network range to the proxy network. The profile locks read-only mode, execution time, result rows and bytes, overflow behavior and memory limits. Managed ClickHouse deployments need equivalent grants and settings constraints -Run the reader permission, authentication and response-limit tests with `cargo test -p litellm-traces --test admin_sql -- --test-threads=2` from `litellm-rust`. They start disposable ClickHouse instances through Testcontainers and load the same reader configuration +Keep reader, ingestion and migration credentials separate. Do not grant the reader write, backup, named-collection management or grant-option privileges, including through roles. `readonly=1` alone is insufficient for privileged accounts; see [query permissions](https://clickhouse.com/docs/concepts/features/configuration/settings/permissions-for-queries) and [settings constraints](https://clickhouse.com/docs/concepts/features/configuration/settings/constraints-on-settings) + +Run `cargo test -p litellm-traces -- --test-threads=2` from `litellm-rust`. Testcontainers applies the canonical schema, checks span rollups and spend joins, and exercises reader permissions, parameter binding, credentials and response limits diff --git a/litellm-rust/crates/traces/config/reader.xml b/litellm-rust/crates/traces/config/reader.xml index 704f56255bbc..4ff2e9d9d848 100644 --- a/litellm-rust/crates/traces/config/reader.xml +++ b/litellm-rust/crates/traces/config/reader.xml @@ -24,6 +24,8 @@ litellm_traces_reader GRANT SELECT ON default.otel_traces + GRANT SELECT ON default.agent_traces + GRANT SELECT ON default.spend_logs diff --git a/litellm-rust/crates/traces/migrations/0001_otel_traces.sql b/litellm-rust/crates/traces/migrations/0001_otel_traces.sql index 0dc54ac2942c..f228ee8144c9 100644 --- a/litellm-rust/crates/traces/migrations/0001_otel_traces.sql +++ b/litellm-rust/crates/traces/migrations/0001_otel_traces.sql @@ -1,49 +1,48 @@ -CREATE TABLE IF NOT EXISTS otel_traces +CREATE TABLE IF NOT EXISTS {database}.otel_traces ( - Timestamp DateTime64(9) CODEC(Delta, ZSTD(1)), - TraceId String CODEC(ZSTD(1)), - SpanId String CODEC(ZSTD(1)), - ParentSpanId String CODEC(ZSTD(1)), - TraceState String CODEC(ZSTD(1)), - SpanName LowCardinality(String) CODEC(ZSTD(1)), - SpanKind LowCardinality(String) CODEC(ZSTD(1)), - ServiceName LowCardinality(String) CODEC(ZSTD(1)), + Timestamp DateTime64(9) CODEC(Delta, ZSTD(1)), + TraceId String CODEC(ZSTD(1)), + SpanId String CODEC(ZSTD(1)), + ParentSpanId String CODEC(ZSTD(1)), + TraceState String CODEC(ZSTD(1)), + SpanName LowCardinality(String) CODEC(ZSTD(1)), + SpanKind LowCardinality(String) CODEC(ZSTD(1)), + ServiceName LowCardinality(String) CODEC(ZSTD(1)), ResourceAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), - ScopeName String CODEC(ZSTD(1)), - ScopeVersion String CODEC(ZSTD(1)), - SpanAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), - Duration UInt64 CODEC(ZSTD(1)), - StatusCode LowCardinality(String) CODEC(ZSTD(1)), - StatusMessage String CODEC(ZSTD(1)), - `Events.Timestamp` Array(DateTime64(9)) CODEC(ZSTD(1)), - `Events.Name` Array(LowCardinality(String)) CODEC(ZSTD(1)), + ScopeName String CODEC(ZSTD(1)), + ScopeVersion String CODEC(ZSTD(1)), + SpanAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), + Duration UInt64 CODEC(ZSTD(1)), + StatusCode LowCardinality(String) CODEC(ZSTD(1)), + StatusMessage String CODEC(ZSTD(1)), + `Events.Timestamp` Array(DateTime64(9)) CODEC(ZSTD(1)), + `Events.Name` Array(LowCardinality(String)) CODEC(ZSTD(1)), `Events.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), - `Links.TraceId` Array(String) CODEC(ZSTD(1)), - `Links.SpanId` Array(String) CODEC(ZSTD(1)), - `Links.TraceState` Array(String) CODEC(ZSTD(1)), - `Links.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), - TeamId LowCardinality(String) DEFAULT SpanAttributes['litellm.team_id'], - ProjectId String DEFAULT SpanAttributes['litellm.project_id'], - UserId String DEFAULT SpanAttributes['litellm.user_id'], - ApiKeyHash String DEFAULT SpanAttributes['litellm.api_key_hash'], - ObservationType LowCardinality(String) DEFAULT multiIf( - ParentSpanId = '', 'agent', - SpanAttributes['gen_ai.operation.name'] = 'invoke_agent', 'agent', - SpanAttributes['gen_ai.operation.name'] IN ('chat', 'text_completion', 'generate_content'), 'llm', - SpanAttributes['gen_ai.operation.name'] = 'execute_tool', 'tool', - 'chain' - ), - AgentName LowCardinality(String) DEFAULT SpanAttributes['gen_ai.agent.name'], - LiteLLMRequestId String DEFAULT SpanAttributes['gen_ai.response.id'], - Model LowCardinality(String) DEFAULT SpanAttributes['gen_ai.request.model'], - InputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.input_tokens']), - OutputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.output_tokens']), - Input String CODEC(ZSTD(3)), - Output String CODEC(ZSTD(3)), - InputPreview String DEFAULT substring(Input, 1, 240), - INDEX idx_trace_id TraceId TYPE bloom_filter(0.001) GRANULARITY 1, - INDEX idx_req_id LiteLLMRequestId TYPE bloom_filter(0.01) GRANULARITY 1 + `Links.TraceId` Array(String) CODEC(ZSTD(1)), + `Links.SpanId` Array(String) CODEC(ZSTD(1)), + `Links.TraceState` Array(String) CODEC(ZSTD(1)), + `Links.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), + TeamId LowCardinality(String) DEFAULT ResourceAttributes['litellm.team_id'], + ApiKeyHash String DEFAULT ResourceAttributes['litellm.api_key_hash'], + ObservationType LowCardinality(String) DEFAULT multiIf( + ParentSpanId = '', 'agent', + SpanAttributes['gen_ai.operation.name'] = 'invoke_agent', 'agent', + SpanAttributes['gen_ai.operation.name'] IN ('chat', 'text_completion', 'generate_content'), 'llm', + SpanAttributes['gen_ai.operation.name'] = 'execute_tool', 'tool', + 'chain'), + AgentName LowCardinality(String) DEFAULT SpanAttributes['gen_ai.agent.name'], + LiteLLMRequestId String DEFAULT SpanAttributes['gen_ai.response.id'], + Model LowCardinality(String) DEFAULT SpanAttributes['gen_ai.request.model'], + InputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.input_tokens']), + OutputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.output_tokens']), + Input String CODEC(ZSTD(3)), + Output String CODEC(ZSTD(3)), + InputPreview String DEFAULT substring(Input, 1, 240), + 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) +TTL toDateTime(Timestamp) + INTERVAL {trace_retention_days} DAY +SETTINGS ttl_only_drop_parts = 1 diff --git a/litellm-rust/crates/traces/migrations/0002_agent_traces.sql b/litellm-rust/crates/traces/migrations/0002_agent_traces.sql new file mode 100644 index 000000000000..ea8177fa12fa --- /dev/null +++ b/litellm-rust/crates/traces/migrations/0002_agent_traces.sql @@ -0,0 +1,25 @@ +CREATE TABLE IF NOT EXISTS {database}.agent_traces +( + TeamId LowCardinality(String), + TraceId String, + StartTs SimpleAggregateFunction(min, DateTime64(9)), + EndTs SimpleAggregateFunction(max, DateTime64(9)), + ServiceName SimpleAggregateFunction(any, LowCardinality(String)), + RootName SimpleAggregateFunction(anyLast, String), + RootInput SimpleAggregateFunction(anyLast, String), + RootStatus SimpleAggregateFunction(anyLast, String), + SpanCount SimpleAggregateFunction(sum, UInt64), + AgentCount SimpleAggregateFunction(sum, UInt64), + LlmCount SimpleAggregateFunction(sum, UInt64), + ToolCount SimpleAggregateFunction(sum, UInt64), + ErrorCount SimpleAggregateFunction(sum, UInt64), + InputTokens SimpleAggregateFunction(sum, UInt64), + OutputTokens SimpleAggregateFunction(sum, UInt64), + Models SimpleAggregateFunction(groupUniqArrayArray, Array(String)), + AgentNames SimpleAggregateFunction(groupUniqArrayArray, Array(String)), + RequestIds SimpleAggregateFunction(groupArrayArray, Array(String)) +) +ENGINE = AggregatingMergeTree +PARTITION BY toDate(StartTs) +ORDER BY (TeamId, TraceId) +TTL toDateTime(StartTs) + INTERVAL {trace_retention_days} DAY diff --git a/litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql b/litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql new file mode 100644 index 000000000000..3b3c5ed3adce --- /dev/null +++ b/litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql @@ -0,0 +1,21 @@ +CREATE MATERIALIZED VIEW IF NOT EXISTS {database}.agent_traces_mv TO {database}.agent_traces AS +SELECT + TeamId, TraceId, + min(Timestamp) AS StartTs, + max(Timestamp + toIntervalNanosecond(Duration)) AS EndTs, + any(ServiceName) AS ServiceName, + anyLastIf(SpanName, ParentSpanId = '') AS RootName, + anyLastIf(InputPreview, ParentSpanId = '') AS RootInput, + anyLastIf(StatusCode, ParentSpanId = '') AS RootStatus, + count() AS SpanCount, + countIf(ObservationType = 'agent') AS AgentCount, + countIf(ObservationType = 'llm') AS LlmCount, + countIf(ObservationType = 'tool') AS ToolCount, + countIf(StatusCode = 'STATUS_CODE_ERROR') AS ErrorCount, + sum(InputTokens) AS InputTokens, + sum(OutputTokens) AS OutputTokens, + groupUniqArrayIf(toString(Model), Model != '') AS Models, + groupUniqArrayIf(SpanName, ObservationType = 'agent') AS AgentNames, + groupArrayIf(LiteLLMRequestId, LiteLLMRequestId != '') AS RequestIds +FROM {database}.otel_traces +GROUP BY TeamId, TraceId diff --git a/litellm-rust/crates/traces/migrations/0004_spend_logs.sql b/litellm-rust/crates/traces/migrations/0004_spend_logs.sql new file mode 100644 index 000000000000..c20c182ba3c8 --- /dev/null +++ b/litellm-rust/crates/traces/migrations/0004_spend_logs.sql @@ -0,0 +1,43 @@ +CREATE TABLE IF NOT EXISTS {database}.spend_logs +( + request_id String, + response_id String, + call_type LowCardinality(String), + api_key String, + key_alias String, + team_id LowCardinality(String), + team_alias String, + organization_id String, + user String, + end_user String, + model LowCardinality(String), + model_group LowCardinality(String), + model_id String, + custom_llm_provider LowCardinality(String), + api_base String, + spend Float64, + prompt_tokens UInt32, + completion_tokens UInt32, + total_tokens UInt32, + cache_read_tokens UInt32, + cache_write_tokens UInt32, + start_time DateTime64(3), + end_time DateTime64(3), + completion_start_time Nullable(DateTime64(3)), + status LowCardinality(String), + error_str String, + cache_hit Bool, + session_id String, + trace_id String, + span_id String, + request_tags Array(String), + metadata String CODEC(ZSTD(3)), + messages String CODEC(ZSTD(3)), + response String CODEC(ZSTD(3)), + 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, toDateTime(start_time), request_id) +TTL toDateTime(start_time) + INTERVAL {spend_log_retention_days} DAY diff --git a/litellm-rust/crates/traces/src/error.rs b/litellm-rust/crates/traces/src/error.rs index c0e70b09aa8b..1219bd9700cb 100644 --- a/litellm-rust/crates/traces/src/error.rs +++ b/litellm-rust/crates/traces/src/error.rs @@ -2,12 +2,8 @@ pub enum Error { #[error("invalid ClickHouse HTTP URL")] InvalidUrl, - #[error("trace query requires start before end and a limit from 1 to 100")] - InvalidListQuery, - #[error("trace query requires nonempty identifiers")] - InvalidIdentifier, - #[error("trace queries require the ClickHouse table schema")] - SchemaPending, + #[error("database must be a nonempty SQL identifier and retention must be positive")] + InvalidSchema, #[error("SQL query must not be empty")] EmptySql, #[error("ClickHouse query failed with HTTP status {0}")] diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index c5764c7c5695..b97a2b889014 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -1,8 +1,10 @@ mod error; +mod schema; mod sql; pub use error::Error; -pub use sql::execute_admin_sql; +pub use schema::schema_statements; +pub use sql::{Parameter, execute_read}; use url::Url; #[derive(Clone)] @@ -19,49 +21,37 @@ impl Connection { Ok(Self { url }) } - pub fn url(&self) -> &Url { - &self.url + pub fn configured( + url: &str, + database: &str, + user: &str, + password: &str, + ) -> Result { + let mut connection = Self::parse(url)?; + connection + .url + .set_username(user) + .map_err(|_| Error::InvalidUrl)?; + connection + .url + .set_password(Some(password)) + .map_err(|_| Error::InvalidUrl)?; + let pairs: Vec<_> = connection + .url + .query_pairs() + .filter(|(key, _)| !matches!(key.as_ref(), "database" | "user" | "password")) + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect(); + connection + .url + .query_pairs_mut() + .clear() + .extend_pairs(pairs) + .append_pair("database", database); + Ok(connection) } -} - -#[derive(Debug, Clone)] -pub struct ListQuery { - pub start_ms: i64, - pub end_ms: i64, - pub service: Option, - pub status: Option, - pub search: Option, - pub cursor: Option, - pub limit: u8, -} - -#[derive(Debug, Clone)] -pub enum Query { - List(ListQuery), - Trace { trace_id: String }, - Span { trace_id: String, span_id: String }, -} -impl Query { - pub fn validate(&self) -> Result<(), Error> { - match self { - Self::List(query) - if query.start_ms >= query.end_ms || !(1..=100).contains(&query.limit) => - { - Err(Error::InvalidListQuery) - } - Self::Trace { trace_id } - | Self::Span { - trace_id, - span_id: _, - } if trace_id.is_empty() => Err(Error::InvalidIdentifier), - Self::Span { span_id, .. } if span_id.is_empty() => Err(Error::InvalidIdentifier), - _ => Ok(()), - } + pub fn url(&self) -> &Url { + &self.url } } - -pub async fn query(_connection: Connection, request: Query) -> Result<(), Error> { - request.validate()?; - Err(Error::SchemaPending) -} diff --git a/litellm-rust/crates/traces/src/schema.rs b/litellm-rust/crates/traces/src/schema.rs new file mode 100644 index 000000000000..341672b3bc4c --- /dev/null +++ b/litellm-rust/crates/traces/src/schema.rs @@ -0,0 +1,37 @@ +use crate::Error; + +const MIGRATIONS: [&str; 4] = [ + 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"), +]; + +pub fn schema_statements( + database: &str, + trace_retention_days: u32, + spend_log_retention_days: u32, +) -> Result, Error> { + if database.is_empty() + || !database + .bytes() + .all(|c| c.is_ascii_alphanumeric() || c == b'_') + || trace_retention_days == 0 + || spend_log_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(), + ) + })) + .collect(), + ) +} diff --git a/litellm-rust/crates/traces/src/sql.rs b/litellm-rust/crates/traces/src/sql.rs index 408c57914091..0989987eca8e 100644 --- a/litellm-rust/crates/traces/src/sql.rs +++ b/litellm-rust/crates/traces/src/sql.rs @@ -1,4 +1,6 @@ -use std::time::Duration; +use std::{collections::BTreeMap, time::Duration}; + +use serde::Deserialize; use litellm_http::Client; @@ -6,10 +8,38 @@ use crate::{Connection, Error}; const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024; -pub async fn execute_admin_sql( +#[derive(Debug, Deserialize)] +#[serde(untagged)] +pub enum Parameter { + Text(String), + Integer(i64), + Strings(Vec), +} + +impl Parameter { + fn encoded(&self) -> String { + match self { + Self::Text(value) => value.clone(), + Self::Integer(value) => value.to_string(), + Self::Strings(values) => format!( + "[{}]", + values + .iter() + .map(|value| { + format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'")) + }) + .collect::>() + .join(",") + ), + } + } +} + +pub async fn execute_read( client: &Client, connection: &Connection, sql: &str, + parameters: &BTreeMap, ) -> Result { if sql.trim().is_empty() { return Err(Error::EmptySql); @@ -20,16 +50,17 @@ pub async fn execute_admin_sql( let existing_pairs: Vec<(String, String)> = url .query_pairs() .filter(|(key, _)| { - !matches!( - key.as_ref(), - "query" - | "readonly" - | "default_format" - | "max_result_rows" - | "result_overflow_mode" - | "max_execution_time" - | "wait_end_of_query" - ) + !key.starts_with("param_") + && !matches!( + key.as_ref(), + "query" + | "readonly" + | "default_format" + | "max_result_rows" + | "result_overflow_mode" + | "max_execution_time" + | "wait_end_of_query" + ) }) .map(|(key, value)| (key.into_owned(), value.into_owned())) .collect(); @@ -43,6 +74,12 @@ pub async fn execute_admin_sql( .append_pair("wait_end_of_query", "1") .append_pair("default_format", "JSON"); + url.query_pairs_mut().extend_pairs( + parameters + .iter() + .map(|(name, value)| (format!("param_{name}"), value.encoded())), + ); + let request = client .post(url) .timeout(Duration::from_secs(15)) diff --git a/litellm-rust/crates/traces/tests/admin_sql.rs b/litellm-rust/crates/traces/tests/admin_sql.rs index 44326e15278b..ed30a36b1641 100644 --- a/litellm-rust/crates/traces/tests/admin_sql.rs +++ b/litellm-rust/crates/traces/tests/admin_sql.rs @@ -1,7 +1,8 @@ use litellm_http::Client; -use litellm_traces::{Connection, Error, execute_admin_sql}; +use litellm_traces::{Connection, Error, Parameter, execute_read}; use rstest::{fixture, rstest}; use serde_json::Value; +use std::collections::BTreeMap; use testcontainers_modules::{ clickhouse::ClickHouse, testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner}, @@ -67,7 +68,7 @@ async fn admin_sql_reads_rows_with_enforced_settings( database.url, ))?; - let result = execute_admin_sql( + let result = read( &database.client, &connection, "SELECT n AS answer FROM otel_traces", @@ -99,11 +100,10 @@ async fn reader_rejects_writes_and_privilege_escalation( let database = database?; let connection = Connection::parse(&format!("{}?readonly=0", database.url))?; - let result = execute_admin_sql(&database.client, &connection, sql).await; + let result = read(&database.client, &connection, sql).await; assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}"); - let rows = - execute_admin_sql(&database.client, &connection, "SELECT n FROM otel_traces").await?; + let rows = read(&database.client, &connection, "SELECT n FROM otel_traces").await?; let json: Value = serde_json::from_str(&rows)?; assert_eq!(json["data"], serde_json::json!([{ "n": 1 }])); Ok(()) @@ -121,7 +121,7 @@ async fn admin_sql_rejects_errors_after_output_starts( database.admin_url, ))?; - let result = execute_admin_sql( + let result = read( &database.client, &connection, "SELECT sleepEachRow(0.2), throwIf(number = 2) FROM numbers(5)", @@ -146,7 +146,7 @@ async fn admin_sql_enforces_result_row_limit( database.url, ))?; - let result = execute_admin_sql( + let result = read( &database.client, &connection, "SELECT number FROM numbers(1001)", @@ -165,7 +165,7 @@ async fn admin_sql_enforces_response_byte_limit( let database = database?; let connection = Connection::parse(&database.admin_url)?; - let result = execute_admin_sql( + let result = read( &database.client, &connection, "SELECT repeat('x', 5 * 1024 * 1024)", @@ -201,7 +201,7 @@ async fn admin_sql_authenticates_url_credentials( 1, ))?; - let result = execute_admin_sql( + let result = read( &database.client, &connection, "SELECT currentUser() AS username", @@ -213,3 +213,43 @@ async fn admin_sql_authenticates_url_credentials( Ok(()) } + +async fn read(client: &Client, connection: &Connection, sql: &str) -> Result { + execute_read(client, connection, sql, &BTreeMap::new()).await +} + +#[rstest] +#[case::sql("'; DROP TABLE otel_traces; --")] +#[case::escapes("back\\slash\ttab\nline\0null")] +#[tokio::test] +async fn query_parameters_preserve_values_and_replace_url_parameters( + #[case] value: &str, + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!("{}?param_value=wrong", database.url))?; + let values = vec![ + "a'b".to_owned(), + "back\\slash".to_owned(), + "line\nbreak".to_owned(), + "雪".to_owned(), + ]; + let parameters = BTreeMap::from([ + ("value".to_owned(), Parameter::Text(value.into())), + ("teams".to_owned(), Parameter::Strings(values.clone())), + ("number".to_owned(), Parameter::Integer(-42)), + ]); + let body = execute_read(&database.client, &connection, + "SELECT {value:String} AS value, {teams:Array(String)} AS teams, toInt32({number:Int64}) AS number", + ¶meters).await?; + let json: Value = serde_json::from_str(&body)?; + assert_eq!(json["data"][0]["value"], value); + assert_eq!(json["data"][0]["teams"], serde_json::json!(values)); + assert_eq!(json["data"][0]["number"], -42); + assert!( + read(&database.client, &connection, "SELECT n FROM otel_traces") + .await + .is_ok() + ); + Ok(()) +} diff --git a/litellm-rust/crates/traces/tests/migrations.rs b/litellm-rust/crates/traces/tests/migrations.rs index b10c508d275a..6c03ae230ffc 100644 --- a/litellm-rust/crates/traces/tests/migrations.rs +++ b/litellm-rust/crates/traces/tests/migrations.rs @@ -1,4 +1,8 @@ +use std::collections::BTreeMap; + use litellm_http::Client; +use litellm_traces::{Connection, execute_read, schema_statements}; +use rstest::rstest; use testcontainers_modules::{ clickhouse::ClickHouse, testcontainers::{ImageExt, runners::AsyncRunner}, @@ -6,11 +10,10 @@ use testcontainers_modules::{ const CLICKHOUSE_TAG: &str = "26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e"; -const INIT_SQL: &str = include_str!("../migrations/0001_otel_traces.sql"); -#[rstest::rstest] +#[rstest] #[tokio::test] -async fn migration_creates_a_queryable_trace_table() -> Result<(), Box> { +async fn schema_supports_span_rollups_and_spend_joins() -> Result<(), Box> { let container = ClickHouse::default() .with_tag(CLICKHOUSE_TAG) .with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1") @@ -19,43 +22,73 @@ async fn migration_creates_a_queryable_trace_table() -> Result<(), Box Query { - Query::List(ListQuery { - start_ms, - end_ms, - service: None, - status: None, - search: None, - cursor: None, - limit, - }) -} - -#[rstest] -#[case::reversed_window(list(2, 1, 50))] -#[case::empty_window(list(1, 1, 50))] -#[case::zero_limit(list(1, 2, 0))] -#[case::large_limit(list(1, 2, 101))] -#[case::empty_trace(Query::Trace { trace_id: String::new() })] -#[case::empty_span(Query::Span { trace_id: "trace".into(), span_id: String::new() })] -fn rejects_invalid_queries(#[case] request: Query) { - assert!(request.validate().is_err()); -} - #[rstest] #[case::http("http://localhost:8123", true)] #[case::https("https://localhost:8443", true)] @@ -32,10 +9,3 @@ fn rejects_invalid_queries(#[case] request: Query) { fn accepts_only_clickhouse_http_urls(#[case] value: &str, #[case] expected: bool) { assert_eq!(Connection::parse(value).is_ok(), expected); } - -#[tokio::test] -async fn schema_pending_is_explicit_after_valid_request() { - let connection = Connection::parse("http://localhost:8123").expect("valid URL"); - let result = litellm_traces::query(connection, list(1, 2, 50)).await; - assert!(matches!(result, Err(Error::SchemaPending))); -} diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 776783206e78..be06d2e73217 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2624,7 +2624,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ - clickhouse_url: str | None = Field(None, description="ClickHouse HTTP URL for trace queries") completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls") max_in_flight_requests_per_worker: int | None = Field( None, gt=0, description="maximum concurrent requests handled by each worker" diff --git a/litellm/proxy/logging_endpoints/traces.py b/litellm/proxy/logging_endpoints/traces.py deleted file mode 100644 index ac82f8cb0227..000000000000 --- a/litellm/proxy/logging_endpoints/traces.py +++ /dev/null @@ -1,95 +0,0 @@ -from collections.abc import Awaitable -from datetime import datetime, timezone -from typing import Final, Protocol, cast - -from fastapi import APIRouter, Depends, HTTPException, Query - -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.rust_bridge.loader import get_native_bridge - -router: Final = APIRouter(prefix="/v1/traces", tags=["Traces"]) - - -class NativeTraces(Protocol): - def list_traces( - self, - clickhouse_url: str, - start_ms: int, - end_ms: int, - service: str | None, - status: str | None, - search: str | None, - cursor: str | None, - limit: int, - ) -> Awaitable[object]: ... - - def get_trace(self, clickhouse_url: str, trace_id: str) -> Awaitable[object]: ... - - def get_span(self, clickhouse_url: str, trace_id: str, span_id: str) -> Awaitable[object]: ... - - -def _connection(user_api_key_dict: UserAPIKeyAuth) -> tuple[str, NativeTraces]: - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException(status_code=403, detail="Trace queries require proxy admin access") - - from litellm.proxy.proxy_server import general_settings_view - - clickhouse_url: Final = general_settings_view().get("clickhouse_url") - if not isinstance(clickhouse_url, str) or not clickhouse_url: - raise HTTPException(status_code=503, detail="ClickHouse trace queries are not configured") - - native: Final = get_native_bridge() - if native is None or not all(hasattr(native, name) for name in ("list_traces", "get_trace", "get_span")): - raise HTTPException(status_code=503, detail="Rust trace query bridge is unavailable") - return clickhouse_url, cast(NativeTraces, native) - - -async def _result(awaitable: Awaitable[object]) -> object: - try: - return await awaitable - except NotImplementedError as exc: - raise HTTPException(status_code=501, detail=str(exc)) from exc - except ValueError as exc: - raise HTTPException(status_code=400, detail=str(exc)) from exc - except RuntimeError as exc: - raise HTTPException(status_code=503, detail=str(exc)) from exc - - -def _utc_ms(value: datetime) -> int: - utc_value: Final = value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) - return int(utc_value.timestamp() * 1000) - - -@router.get("") -async def list_traces( - start: datetime, - end: datetime, - service: str | None = None, - status: str | None = None, - q: str | None = None, - cursor: str | None = None, - limit: int = Query(default=50, ge=1, le=100), - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> object: - clickhouse_url, native = _connection(user_api_key_dict) - return await _result(native.list_traces(clickhouse_url, _utc_ms(start), _utc_ms(end), service, status, q, cursor, limit)) - - -@router.get("/{trace_id}") -async def get_trace( - trace_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> object: - clickhouse_url, native = _connection(user_api_key_dict) - return await _result(native.get_trace(clickhouse_url, trace_id)) - - -@router.get("/{trace_id}/spans/{span_id}") -async def get_span( - trace_id: str, - span_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> object: - clickhouse_url, native = _connection(user_api_key_dict) - return await _result(native.get_span(clickhouse_url, trace_id, span_id)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d777d3dc9066..a994293a4d89 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -568,7 +568,6 @@ def generate_feedback_box(): from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( rust_control_plane_router, ) -from litellm.proxy.logging_endpoints.traces import router as traces_router from litellm.proxy.management_endpoints.auto_router_endpoints import ( router as auto_router_management_router, ) @@ -973,7 +972,6 @@ def generate_feedback_box(): _EXTRA_SECRET_GENERAL_SETTINGS_FIELDS: Final = frozenset( { "database_url", - "clickhouse_url", "database_extra_connection_params", "pass_through_endpoints", "alert_to_webhook_url", @@ -19886,7 +19884,6 @@ async def get_routes(): app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) app.include_router(rust_control_plane_router) -app.include_router(traces_router) app.include_router(ui_crud_endpoints_router) app.include_router(user_banner_endpoints_router) app.include_router(latest_release_endpoints_router) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 173ee08279a3..978f0069923f 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -19,18 +19,16 @@ class RustBridgeDeclined(Exception): ... class RustUpstreamError(Exception): ... class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... -def list_traces( - clickhouse_url: str, - start_ms: int, - end_ms: int, - service: str | None = None, - status: str | None = None, - search: str | None = None, - cursor: str | None = None, - limit: int = 50, -) -> Future[None]: ... -def get_trace(clickhouse_url: str, trace_id: str) -> Future[None]: ... -def get_span(clickhouse_url: str, trace_id: str, span_id: str) -> Future[None]: ... + +def trace_schema_statements(database: str, trace_retention_days: int, spend_log_retention_days: int) -> list[str]: ... +def trace_query( + url: str, + database: str, + user: str, + password: str, + sql: str, + parameters: Mapping[str, str | int | Sequence[str]], +) -> Future[str]: ... @final class NativeDiagnosticProcessor: diff --git a/litellm/rust_bridge/traces.py b/litellm/rust_bridge/traces.py new file mode 100644 index 000000000000..06321423b52c --- /dev/null +++ b/litellm/rust_bridge/traces.py @@ -0,0 +1,50 @@ +from collections.abc import Awaitable, Mapping, Sequence +from typing import Final, Protocol, cast + +from pydantic import BaseModel, ConfigDict, JsonValue + +from litellm.rust_bridge.loader import get_native_bridge + + +class NativeTraces(Protocol): + def trace_schema_statements( + self, database: str, trace_retention_days: int, spend_log_retention_days: int + ) -> list[str]: ... + + def trace_query( + self, + url: str, + database: str, + user: str, + password: str, + sql: str, + parameters: Mapping[str, str | int | Sequence[str]], + ) -> Awaitable[str]: ... + + +class QueryResponse(BaseModel): + model_config = ConfigDict(frozen=True) + data: list[dict[str, JsonValue]] + + +def _native() -> NativeTraces: + native: Final = get_native_bridge() + if native is None: + raise RuntimeError("Agent tracing requires the Rust extension") + return cast(NativeTraces, native) + + +def schema_statements(database: str, trace_retention_days: int, spend_log_retention_days: int) -> list[str]: + return _native().trace_schema_statements(database, trace_retention_days, spend_log_retention_days) + + +async def query( + url: str, + database: str, + user: str, + password: str, + sql: str, + parameters: Mapping[str, str | int | Sequence[str]], +) -> list[dict[str, JsonValue]]: + result: Final = await _native().trace_query(url, database, user, password, sql, parameters) + return QueryResponse.model_validate_json(result).data diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py new file mode 100644 index 000000000000..df7f38af665d --- /dev/null +++ b/tests/test_litellm_rust/test_traces.py @@ -0,0 +1,46 @@ +import base64 +from typing import Final +from urllib.parse import parse_qs, urlsplit + +import pytest + +from litellm.rust_bridge.traces import query, schema_statements +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +async def test_trace_reader_projects_connection_and_parameters(recording_server: RecordingServer) -> None: + recording_server.enqueue(ResponseSpec(body={"data": [{"trace_id": "trace-1"}]})) + rows: Final = await query( + recording_server.base_url + "?database=wrong&user=wrong&password=wrong", + "trace_test", + "reader", + "p@ss/word%", + "SELECT {trace_id:String} AS trace_id", + {"trace_id": "trace-1"}, + ) + request: Final = recording_server.requests[0] + parameters: Final = parse_qs(urlsplit(request.path).query) + assert rows == [{"trace_id": "trace-1"}] + assert request.raw_body == b"SELECT {trace_id:String} AS trace_id" + assert parameters["database"] == ["trace_test"] + assert parameters["param_trace_id"] == ["trace-1"] + assert parameters["readonly"] == ["1"] + assert "user" not in parameters + assert "password" not in parameters + assert request.headers["authorization"] == "Basic " + base64.b64encode(b"reader:p@ss/word%").decode() + + +@pytest.mark.asyncio +async def test_trace_reader_rejects_success_status_with_embedded_error(recording_server: RecordingServer) -> None: + recording_server.enqueue(ResponseSpec(body={"data": [], "exception": "query failed"})) + with pytest.raises(RuntimeError, match="invalid or failed JSON"): + await query(recording_server.base_url, "trace_test", "reader", "password", "SELECT 1", {}) + + +@pytest.mark.parametrize("database,retention", [("db; DROP DATABASE default", 7), ("traces", 0)]) +def test_schema_binding_preserves_configuration_validation(database: str, retention: int) -> None: + with pytest.raises(ValueError, match="database.*retention"): + schema_statements(database, retention, 14) diff --git a/tests/unit/proxy/logging_endpoints/test_traces.py b/tests/unit/proxy/logging_endpoints/test_traces.py deleted file mode 100644 index 2c92e47d78bb..000000000000 --- a/tests/unit/proxy/logging_endpoints/test_traces.py +++ /dev/null @@ -1,45 +0,0 @@ -import asyncio -from collections.abc import Callable, Coroutine -from datetime import datetime, timezone - -import pytest -from fastapi import HTTPException - -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.logging_endpoints.traces import get_span, get_trace, list_traces - - -def _list(auth: UserAPIKeyAuth) -> Coroutine[object, object, object]: - return list_traces( - start=datetime(2026, 1, 1, tzinfo=timezone.utc), - end=datetime(2026, 1, 2, tzinfo=timezone.utc), - limit=50, - user_api_key_dict=auth, - ) - - -def _trace(auth: UserAPIKeyAuth) -> Coroutine[object, object, object]: - return get_trace(trace_id="trace", user_api_key_dict=auth) - - -def _span(auth: UserAPIKeyAuth) -> Coroutine[object, object, object]: - return get_span(trace_id="trace", span_id="span", user_api_key_dict=auth) - - -@pytest.mark.parametrize( - "run_route", - ( - _list, - _trace, - _span, - ), -) -def test_trace_routes_reject_non_admin_callers_before_accessing_storage( - run_route: Callable[[UserAPIKeyAuth], Coroutine[object, object, object]], -) -> None: - auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) - - with pytest.raises(HTTPException) as error: - asyncio.run(run_route(auth)) - - assert error.value.status_code == 403 From 50dc0950924ccc857a9b82ac5bd02892760dc988 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 07:16:21 -0700 Subject: [PATCH 03/12] fix(traces): escape ClickHouse text parameters --- litellm-rust/crates/traces/src/sql.rs | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/litellm-rust/crates/traces/src/sql.rs b/litellm-rust/crates/traces/src/sql.rs index 0989987eca8e..1aa21a59caad 100644 --- a/litellm-rust/crates/traces/src/sql.rs +++ b/litellm-rust/crates/traces/src/sql.rs @@ -19,15 +19,13 @@ pub enum Parameter { impl Parameter { fn encoded(&self) -> String { match self { - Self::Text(value) => value.clone(), + Self::Text(value) => escaped(value), Self::Integer(value) => value.to_string(), Self::Strings(values) => format!( "[{}]", values .iter() - .map(|value| { - format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'")) - }) + .map(|value| format!("'{}'", escaped(value).replace('\'', "\\'"))) .collect::>() .join(",") ), @@ -35,6 +33,15 @@ impl Parameter { } } +fn escaped(value: &str) -> String { + value + .replace('\\', "\\\\") + .replace('\t', "\\t") + .replace('\n', "\\n") + .replace('\r', "\\r") + .replace('\0', "\\0") +} + pub async fn execute_read( client: &Client, connection: &Connection, From 835118f4fc7b3fc0385595a36c27953e4389f893 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 07:18:40 -0700 Subject: [PATCH 04/12] test(traces): exercise response cap with bounded strings --- litellm-rust/crates/traces/tests/admin_sql.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/crates/traces/tests/admin_sql.rs b/litellm-rust/crates/traces/tests/admin_sql.rs index ed30a36b1641..b2ef6113a15e 100644 --- a/litellm-rust/crates/traces/tests/admin_sql.rs +++ b/litellm-rust/crates/traces/tests/admin_sql.rs @@ -168,7 +168,7 @@ async fn admin_sql_enforces_response_byte_limit( let result = read( &database.client, &connection, - "SELECT repeat('x', 5 * 1024 * 1024)", + "SELECT repeat('x', 512 * 1024) AS payload FROM numbers(9)", ) .await; From c120adb6a296d1ed8451f40270a969da042e560a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 07:22:11 -0700 Subject: [PATCH 05/12] fix(traces): remove unnecessary lint expectation --- litellm-rust/crates/python-bridge/src/routes/traces.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index 557bde793ee4..9da1224ff82e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -32,10 +32,6 @@ pub fn trace_schema_statements( } #[pyfunction] -#[expect( - clippy::too_many_arguments, - reason = "the internal adapter supplies connection fields and a parameterized query" -)] pub fn trace_query<'py>( py: Python<'py>, url: &str, From 350d06a7e80fa21187267846a0216f78b80a2922 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 07:29:08 -0700 Subject: [PATCH 06/12] fix(traces): encode ClickHouse timestamp units in Rust --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 3 +- .../crates/python-bridge/src/routes/traces.rs | 15 ++++++- litellm-rust/crates/traces/Cargo.toml | 1 + litellm-rust/crates/traces/README.md | 4 +- litellm-rust/crates/traces/src/error.rs | 2 + litellm-rust/crates/traces/src/insert.rs | 37 +++++++++++++++++ litellm-rust/crates/traces/src/lib.rs | 2 + litellm-rust/crates/traces/tests/insert.rs | 40 +++++++++++++++++++ .../crates/traces/tests/migrations.rs | 40 +++++++++++++------ litellm/rust_bridge/_native.pyi | 1 + litellm/rust_bridge/traces.py | 6 +++ tests/test_litellm_rust/test_traces.py | 10 ++++- 13 files changed, 145 insertions(+), 17 deletions(-) create mode 100644 litellm-rust/crates/traces/src/insert.rs create mode 100644 litellm-rust/crates/traces/tests/insert.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 236e3bb0c008..5fdf0266e8c8 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4362,6 +4362,7 @@ dependencies = [ "serde_json", "testcontainers-modules", "thiserror 2.0.19", + "time", "tokio", "url", ] diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 8b55a40fa226..4dd73e3654dc 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -44,7 +44,7 @@ mod _native { #[pymodule_export] use crate::routes::token_counter::TokenCounter; #[pymodule_export] - use crate::routes::traces::{trace_query, trace_schema_statements}; + use crate::routes::traces::{trace_encode_rows, trace_query, trace_schema_statements}; #[cfg(feature = "huggingface")] #[pymodule_export] use crate::tokenizer::HuggingFaceEncoding; @@ -109,6 +109,7 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", + "trace_encode_rows", "trace_query", "trace_schema_statements", "TokenCounter", diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index 9da1224ff82e..3a2585f8baea 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -9,7 +9,9 @@ use pyo3::{ fn map_error(error: Error) -> PyErr { match error { - Error::InvalidSchema | Error::EmptySql => PyValueError::new_err(error.to_string()), + Error::InvalidRow | Error::InvalidSchema | Error::EmptySql => { + PyValueError::new_err(error.to_string()) + } Error::InvalidUrl | Error::QueryFailed(_) | Error::ResponseTooLarge @@ -52,3 +54,14 @@ pub fn trace_query<'py>( map_error, ) } + +#[pyfunction] +pub fn trace_encode_rows( + py: Python<'_>, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] rows: Vec< + BTreeMap, + >, +) -> PyResult { + py.detach(|| litellm_traces::encode_rows(rows)) + .map_err(map_error) +} diff --git a/litellm-rust/crates/traces/Cargo.toml b/litellm-rust/crates/traces/Cargo.toml index 87d2709c9063..a8091cf9ec80 100644 --- a/litellm-rust/crates/traces/Cargo.toml +++ b/litellm-rust/crates/traces/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +time = { workspace = true, features = ["formatting"] } litellm-http.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/traces/README.md b/litellm-rust/crates/traces/README.md index 657bdc58d6ce..7a2cb0870bbd 100644 --- a/litellm-rust/crates/traces/README.md +++ b/litellm-rust/crates/traces/README.md @@ -1,6 +1,6 @@ # Trace storage foundation -This crate owns the initial ClickHouse schema and parameterized read transport. Python adapters use `trace_schema_statements` and `trace_query` from the native bridge. HTTP endpoints, ingestion and trace response assembly belong to the integration built on this foundation +This crate owns the initial ClickHouse schema, insert row encoding and parameterized read transport. Python adapters use `trace_schema_statements`, `trace_encode_rows` and `trace_query` from the native bridge. HTTP endpoints, ingestion and trace response assembly belong to the integration built on this foundation The SQL files under `migrations/` are the canonical definitions for `otel_traces`, `agent_traces`, its materialized view, and `spend_logs`. `schema_statements` supplies the database and retention settings. Apply the returned statements in order using migration credentials. Setup is repeatable for a new installation; these initial `CREATE IF NOT EXISTS` statements do not upgrade an existing incompatible table. Future schema changes need explicit upgrade statements @@ -13,3 +13,5 @@ For self-hosted ClickHouse, install [config/reader.xml](config/reader.xml) as `/ Keep reader, ingestion and migration credentials separate. Do not grant the reader write, backup, named-collection management or grant-option privileges, including through roles. `readonly=1` alone is insufficient for privileged accounts; see [query permissions](https://clickhouse.com/docs/concepts/features/configuration/settings/permissions-for-queries) and [settings constraints](https://clickhouse.com/docs/concepts/features/configuration/settings/constraints-on-settings) Run `cargo test -p litellm-traces -- --test-threads=2` from `litellm-rust`. Testcontainers applies the canonical schema, checks span rollups and spend joins, and exercises reader permissions, parameter binding, credentials and response limits + +`encode_rows` produces JSONEachRow and formats span nanoseconds and spend-log milliseconds as UTC DateTime64 strings without losing precision. The integration owns compression, batching and writes diff --git a/litellm-rust/crates/traces/src/error.rs b/litellm-rust/crates/traces/src/error.rs index 1219bd9700cb..6832582272aa 100644 --- a/litellm-rust/crates/traces/src/error.rs +++ b/litellm-rust/crates/traces/src/error.rs @@ -1,5 +1,7 @@ #[derive(Debug, thiserror::Error)] pub enum Error { + #[error("invalid ClickHouse insert row")] + InvalidRow, #[error("invalid ClickHouse HTTP URL")] InvalidUrl, #[error("database must be a nonempty SQL identifier and retention must be positive")] diff --git a/litellm-rust/crates/traces/src/insert.rs b/litellm-rust/crates/traces/src/insert.rs new file mode 100644 index 000000000000..05b14651ce07 --- /dev/null +++ b/litellm-rust/crates/traces/src/insert.rs @@ -0,0 +1,37 @@ +use std::collections::BTreeMap; + +use serde_json::Value; +use time::{OffsetDateTime, format_description::well_known::Rfc3339}; + +use crate::Error; + +pub fn encode_rows(rows: Vec>) -> Result { + rows.into_iter() + .map(|row| { + let encoded = row + .into_iter() + .map(|(name, value)| insert_value(&name, value).map(|value| (name, value))) + .collect::, _>>()?; + serde_json::to_string(&encoded).map_err(|_| Error::InvalidRow) + }) + .collect::, _>>() + .map(|rows| rows.join("\n")) +} + +fn insert_value(name: &str, value: Value) -> Result { + let multiplier = match name { + "Timestamp" => 1, + "start_time" | "end_time" | "completion_start_time" => 1_000_000, + _ => return Ok(value), + }; + if name == "completion_start_time" && value.is_null() { + return Ok(value); + } + let timestamp = value.as_i64().ok_or(Error::InvalidRow)?; + let datetime = OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp) * multiplier) + .map_err(|_| Error::InvalidRow)?; + datetime + .format(&Rfc3339) + .map(Value::String) + .map_err(|_| Error::InvalidRow) +} diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index b97a2b889014..f2e9e411b965 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -1,8 +1,10 @@ mod error; +mod insert; mod schema; mod sql; pub use error::Error; +pub use insert::encode_rows; pub use schema::schema_statements; pub use sql::{Parameter, execute_read}; use url::Url; diff --git a/litellm-rust/crates/traces/tests/insert.rs b/litellm-rust/crates/traces/tests/insert.rs new file mode 100644 index 000000000000..cba678152b9e --- /dev/null +++ b/litellm-rust/crates/traces/tests/insert.rs @@ -0,0 +1,40 @@ +use std::collections::BTreeMap; + +use litellm_traces::encode_rows; +use rstest::rstest; +use serde_json::{Value, json}; + +#[rstest] +#[case::span("Timestamp", json!(1_234_567_890), json!("1970-01-01T00:00:01.23456789Z"))] +#[case::start("start_time", json!(1_234), json!("1970-01-01T00:00:01.234Z"))] +#[case::end("end_time", json!(2_345), json!("1970-01-01T00:00:02.345Z"))] +#[case::completion("completion_start_time", json!(1_345), json!("1970-01-01T00:00:01.345Z"))] +#[case::absent_completion("completion_start_time", Value::Null, Value::Null)] +#[case::before_epoch("Timestamp", json!(-1), json!("1969-12-31T23:59:59.999999999Z"))] +fn insert_encoding_preserves_timestamp_precision_and_other_fields( + #[case] field: &str, + #[case] value: Value, + #[case] expected: Value, +) { + let rows = vec![BTreeMap::from([ + (field.to_owned(), value), + ("SpanAttributes".into(), json!({"message": "a\nb\\c\"雪"})), + ("InputTokens".into(), json!(42)), + ])]; + let encoded = encode_rows(rows).expect("valid row"); + let actual: Value = serde_json::from_str(&encoded).expect("JSONEachRow record"); + assert_eq!( + actual, + json!({ + field: expected, "SpanAttributes": {"message": "a\nb\\c\"雪"}, "InputTokens": 42 + }) + ); +} + +#[rstest] +#[case::fractional(json!(1.25))] +#[case::out_of_range(json!(u64::MAX))] +#[case::null(Value::Null)] +fn insert_encoding_rejects_invalid_span_timestamps(#[case] timestamp: Value) { + assert!(encode_rows(vec![BTreeMap::from([("Timestamp".into(), timestamp)])]).is_err()); +} diff --git a/litellm-rust/crates/traces/tests/migrations.rs b/litellm-rust/crates/traces/tests/migrations.rs index 6c03ae230ffc..50b9ee01d216 100644 --- a/litellm-rust/crates/traces/tests/migrations.rs +++ b/litellm-rust/crates/traces/tests/migrations.rs @@ -1,7 +1,7 @@ use std::collections::BTreeMap; use litellm_http::Client; -use litellm_traces::{Connection, execute_read, schema_statements}; +use litellm_traces::{Connection, encode_rows, execute_read, schema_statements}; use rstest::rstest; use testcontainers_modules::{ clickhouse::ClickHouse, @@ -35,25 +35,38 @@ async fn schema_supports_span_rollups_and_spend_joins() -> Result<(), Box Result<(), Box str: ... def trace_schema_statements(database: str, trace_retention_days: int, spend_log_retention_days: int) -> list[str]: ... def trace_query( url: str, diff --git a/litellm/rust_bridge/traces.py b/litellm/rust_bridge/traces.py index 06321423b52c..3beee285d4db 100644 --- a/litellm/rust_bridge/traces.py +++ b/litellm/rust_bridge/traces.py @@ -7,6 +7,8 @@ class NativeTraces(Protocol): + def trace_encode_rows(self, rows: Sequence[Mapping[str, JsonValue]]) -> str: ... + def trace_schema_statements( self, database: str, trace_retention_days: int, spend_log_retention_days: int ) -> list[str]: ... @@ -48,3 +50,7 @@ async def query( ) -> list[dict[str, JsonValue]]: result: Final = await _native().trace_query(url, database, user, password, sql, parameters) return QueryResponse.model_validate_json(result).data + + +def encode_rows(rows: Sequence[Mapping[str, JsonValue]]) -> bytes: + return _native().trace_encode_rows(rows).encode("utf-8") diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index df7f38af665d..961071619be7 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -1,10 +1,11 @@ import base64 +import json from typing import Final from urllib.parse import parse_qs, urlsplit import pytest -from litellm.rust_bridge.traces import query, schema_statements +from litellm.rust_bridge.traces import encode_rows, query, schema_statements from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension @@ -44,3 +45,10 @@ async def test_trace_reader_rejects_success_status_with_embedded_error(recording def test_schema_binding_preserves_configuration_validation(database: str, retention: int) -> None: with pytest.raises(ValueError, match="database.*retention"): schema_statements(database, retention, 14) + + +def test_insert_encoding_preserves_nanoseconds_through_bridge() -> None: + assert json.loads(encode_rows([{"Timestamp": 1_234_567_890, "Input": "hello"}])) == { + "Input": "hello", + "Timestamp": "1970-01-01T00:00:01.23456789Z", + } From 1da14377183332ebc15641fa00c6b5d98ceb6d9e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 07:36:22 -0700 Subject: [PATCH 07/12] test(traces): mark exception match as a regex --- tests/test_litellm_rust/test_traces.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index 961071619be7..5ee14d569904 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -43,7 +43,7 @@ async def test_trace_reader_rejects_success_status_with_embedded_error(recording @pytest.mark.parametrize("database,retention", [("db; DROP DATABASE default", 7), ("traces", 0)]) def test_schema_binding_preserves_configuration_validation(database: str, retention: int) -> None: - with pytest.raises(ValueError, match="database.*retention"): + with pytest.raises(ValueError, match=r"database.*retention"): schema_statements(database, retention, 14) From 829e9a726a66635ee4fdc5e5052e86d3195f645e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 08:28:38 -0700 Subject: [PATCH 08/12] refactor(traces): execute schema setup in Rust --- litellm-rust/crates/python-bridge/src/lib.rs | 4 +-- .../crates/python-bridge/src/routes/traces.rs | 30 ++++++++++++++----- litellm-rust/crates/traces/README.md | 4 +-- litellm-rust/crates/traces/src/error.rs | 2 ++ litellm-rust/crates/traces/src/lib.rs | 16 +++++++++- litellm-rust/crates/traces/src/schema.rs | 24 +++++++++++++++ .../crates/traces/tests/migrations.rs | 15 +++------- litellm/rust_bridge/_native.pyi | 7 ++++- litellm/rust_bridge/traces.py | 23 ++++++++++---- tests/test_litellm_rust/test_traces.py | 30 +++++++++++++++++-- 10 files changed, 123 insertions(+), 32 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 4dd73e3654dc..f65cb350cfe6 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -44,7 +44,7 @@ mod _native { #[pymodule_export] use crate::routes::token_counter::TokenCounter; #[pymodule_export] - use crate::routes::traces::{trace_encode_rows, trace_query, trace_schema_statements}; + use crate::routes::traces::{trace_encode_rows, trace_ensure_schema, trace_query}; #[cfg(feature = "huggingface")] #[pymodule_export] use crate::tokenizer::HuggingFaceEncoding; @@ -110,8 +110,8 @@ mod tests { "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", "trace_encode_rows", + "trace_ensure_schema", "trace_query", - "trace_schema_statements", "TokenCounter", "Tokenizer", "gil_stats", diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index 3a2585f8baea..2c4fc8573fec 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -14,6 +14,7 @@ fn map_error(error: Error) -> PyErr { } Error::InvalidUrl | Error::QueryFailed(_) + | Error::SchemaFailed(_) | Error::ResponseTooLarge | Error::InvalidResponse | Error::Transport => PyRuntimeError::new_err(error.to_string()), @@ -21,16 +22,31 @@ fn map_error(error: Error) -> PyErr { } #[pyfunction] -pub fn trace_schema_statements( - py: Python<'_>, +pub fn trace_ensure_schema<'py>( + py: Python<'py>, + url: &str, database: String, + user: &str, + password: &str, trace_retention_days: u32, spend_log_retention_days: u32, -) -> PyResult> { - py.detach(|| { - litellm_traces::schema_statements(&database, trace_retention_days, spend_log_retention_days) - }) - .map_err(map_error) +) -> PyResult> { + let connection = Connection::writer(url, user, password).map_err(map_error)?; + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + litellm_host_python::run_async( + py, + async move { + litellm_traces::ensure_schema( + &client, + &connection, + &database, + trace_retention_days, + spend_log_retention_days, + ) + .await + }, + map_error, + ) } #[pyfunction] diff --git a/litellm-rust/crates/traces/README.md b/litellm-rust/crates/traces/README.md index 7a2cb0870bbd..7f7af15a59e2 100644 --- a/litellm-rust/crates/traces/README.md +++ b/litellm-rust/crates/traces/README.md @@ -1,8 +1,8 @@ # Trace storage foundation -This crate owns the initial ClickHouse schema, insert row encoding and parameterized read transport. Python adapters use `trace_schema_statements`, `trace_encode_rows` and `trace_query` from the native bridge. HTTP endpoints, ingestion and trace response assembly belong to the integration built on this foundation +This crate owns the initial ClickHouse schema, insert row encoding and parameterized read transport. Python adapters use `trace_ensure_schema`, `trace_encode_rows` and `trace_query` from the native bridge. HTTP endpoints, ingestion and trace response assembly belong to the integration built on this foundation -The SQL files under `migrations/` are the canonical definitions for `otel_traces`, `agent_traces`, its materialized view, and `spend_logs`. `schema_statements` supplies the database and retention settings. Apply the returned statements in order using migration credentials. Setup is repeatable for a new installation; these initial `CREATE IF NOT EXISTS` statements do not upgrade an existing incompatible table. Future schema changes need explicit upgrade statements +The SQL files under `migrations/` are the canonical definitions for `otel_traces`, `agent_traces`, its materialized view, and `spend_logs`. `ensure_schema` applies them in order using writer credentials and checks each response. Setup is repeatable for a new installation; these initial `CREATE IF NOT EXISTS` statements do not upgrade an existing incompatible table. Future schema changes need explicit upgrade statements `execute_read` uses ClickHouse typed placeholders and binds text, integer and string-array parameters separately from SQL. It obtains no privileges itself: callers must use a dedicated SELECT-only reader. The helper enforces a 15-second HTTP timeout, a 10-second query setting, 1,000 result rows and a 4 MiB response cap, and rejects errors embedded in HTTP 200 JSON responses. Large traces fail explicitly at those limits; span pagination is future work diff --git a/litellm-rust/crates/traces/src/error.rs b/litellm-rust/crates/traces/src/error.rs index 6832582272aa..a06eb27397a9 100644 --- a/litellm-rust/crates/traces/src/error.rs +++ b/litellm-rust/crates/traces/src/error.rs @@ -10,6 +10,8 @@ pub enum Error { EmptySql, #[error("ClickHouse query failed with HTTP status {0}")] QueryFailed(u16), + #[error("ClickHouse schema setup failed with HTTP status {0}")] + SchemaFailed(u16), #[error("ClickHouse query exceeded the response size limit")] ResponseTooLarge, #[error("ClickHouse returned an invalid or failed JSON query response")] diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index f2e9e411b965..1f56bb5b3a6a 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -5,7 +5,7 @@ mod sql; pub use error::Error; pub use insert::encode_rows; -pub use schema::schema_statements; +pub use schema::{ensure_schema, schema_statements}; pub use sql::{Parameter, execute_read}; use url::Url; @@ -53,6 +53,20 @@ impl Connection { Ok(connection) } + pub fn writer(url: &str, user: &str, password: &str) -> Result { + let mut connection = Self::parse(url)?; + connection + .url + .set_username(user) + .map_err(|_| Error::InvalidUrl)?; + connection + .url + .set_password(Some(password)) + .map_err(|_| Error::InvalidUrl)?; + connection.url.set_query(None); + Ok(connection) + } + pub fn url(&self) -> &Url { &self.url } diff --git a/litellm-rust/crates/traces/src/schema.rs b/litellm-rust/crates/traces/src/schema.rs index 341672b3bc4c..3959d78b4c69 100644 --- a/litellm-rust/crates/traces/src/schema.rs +++ b/litellm-rust/crates/traces/src/schema.rs @@ -1,3 +1,6 @@ +use litellm_http::Client; + +use crate::Connection; use crate::Error; const MIGRATIONS: [&str; 4] = [ @@ -35,3 +38,24 @@ pub fn schema_statements( .collect(), ) } + +pub async fn ensure_schema( + client: &Client, + connection: &Connection, + database: &str, + trace_retention_days: u32, + spend_log_retention_days: u32, +) -> Result<(), Error> { + for statement in schema_statements(database, trace_retention_days, spend_log_retention_days)? { + let response = client + .post(connection.url().clone()) + .body(statement) + .send() + .await + .map_err(|_| Error::Transport)?; + if !response.status().is_success() { + return Err(Error::SchemaFailed(response.status().as_u16())); + } + } + Ok(()) +} diff --git a/litellm-rust/crates/traces/tests/migrations.rs b/litellm-rust/crates/traces/tests/migrations.rs index 50b9ee01d216..a77d57cfb0fe 100644 --- a/litellm-rust/crates/traces/tests/migrations.rs +++ b/litellm-rust/crates/traces/tests/migrations.rs @@ -1,7 +1,7 @@ use std::collections::BTreeMap; use litellm_http::Client; -use litellm_traces::{Connection, encode_rows, execute_read, schema_statements}; +use litellm_traces::{Connection, encode_rows, ensure_schema, execute_read, schema_statements}; use rstest::rstest; use testcontainers_modules::{ clickhouse::ClickHouse, @@ -25,16 +25,9 @@ async fn schema_supports_span_rollups_and_spend_joins() -> Result<(), Box str: ... -def trace_schema_statements(database: str, trace_retention_days: int, spend_log_retention_days: int) -> list[str]: ... +def trace_ensure_schema( + url: str, database: str, user: str, password: str, trace_retention_days: int, spend_log_retention_days: int +) -> Future[None]: ... def trace_query( url: str, database: str, @@ -349,6 +351,9 @@ __all__ = [ "process_state_started", "reserve_process_for_forking", "responses", + "trace_encode_rows", + "trace_ensure_schema", + "trace_query", "transcription", ] diff --git a/litellm/rust_bridge/traces.py b/litellm/rust_bridge/traces.py index 3beee285d4db..782d63a3aa78 100644 --- a/litellm/rust_bridge/traces.py +++ b/litellm/rust_bridge/traces.py @@ -9,9 +9,15 @@ class NativeTraces(Protocol): def trace_encode_rows(self, rows: Sequence[Mapping[str, JsonValue]]) -> str: ... - def trace_schema_statements( - self, database: str, trace_retention_days: int, spend_log_retention_days: int - ) -> list[str]: ... + def trace_ensure_schema( + self, + url: str, + database: str, + user: str, + password: str, + trace_retention_days: int, + spend_log_retention_days: int, + ) -> Awaitable[None]: ... def trace_query( self, @@ -36,8 +42,15 @@ def _native() -> NativeTraces: return cast(NativeTraces, native) -def schema_statements(database: str, trace_retention_days: int, spend_log_retention_days: int) -> list[str]: - return _native().trace_schema_statements(database, trace_retention_days, spend_log_retention_days) +async def ensure_schema( + url: str, + database: str, + user: str, + password: str, + trace_retention_days: int, + spend_log_retention_days: int, +) -> None: + await _native().trace_ensure_schema(url, database, user, password, trace_retention_days, spend_log_retention_days) async def query( diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index 5ee14d569904..21f61ad24bce 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -5,7 +5,7 @@ import pytest -from litellm.rust_bridge.traces import encode_rows, query, schema_statements +from litellm.rust_bridge.traces import encode_rows, ensure_schema, query from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension @@ -41,10 +41,34 @@ async def test_trace_reader_rejects_success_status_with_embedded_error(recording await query(recording_server.base_url, "trace_test", "reader", "password", "SELECT 1", {}) +@pytest.mark.asyncio @pytest.mark.parametrize("database,retention", [("db; DROP DATABASE default", 7), ("traces", 0)]) -def test_schema_binding_preserves_configuration_validation(database: str, retention: int) -> None: +async def test_schema_binding_preserves_configuration_validation(database: str, retention: int) -> None: with pytest.raises(ValueError, match=r"database.*retention"): - schema_statements(database, retention, 14) + await ensure_schema("http://localhost:8123", database, "writer", "password", retention, 14) + + +@pytest.mark.asyncio +async def test_schema_setup_uses_writer_credentials_and_rejects_failed_statement(recording_server: RecordingServer) -> None: + recording_server.expected_requests = 2 + recording_server.enqueue(ResponseSpec(body="")) + recording_server.enqueue(ResponseSpec(status=403, body="denied")) + with pytest.raises(RuntimeError, match="schema setup failed with HTTP status 403"): + await ensure_schema( + recording_server.base_url + "?database=wrong&readonly=1", + "trace_test", + "writer", + "p@ss/word%", + 7, + 14, + ) + assert len(recording_server.requests) == 2 + assert recording_server.requests[0].raw_body.startswith(b"CREATE DATABASE IF NOT EXISTS") + assert recording_server.requests[1].raw_body.startswith(b"CREATE TABLE IF NOT EXISTS") + assert "readonly" not in parse_qs(urlsplit(recording_server.requests[0].path).query) + assert recording_server.requests[0].headers["authorization"] == "Basic " + base64.b64encode( + b"writer:p@ss/word%" + ).decode() def test_insert_encoding_preserves_nanoseconds_through_bridge() -> None: From 1f1c9856eab94d71dd1a24561a03b713e8fc4bcb Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 08:36:13 -0700 Subject: [PATCH 09/12] refactor(traces): use shared logging execution wrapper --- litellm-rust/crates/python-bridge/src/routes/traces.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index 2c4fc8573fec..eb7ff16af28e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -33,7 +33,7 @@ pub fn trace_ensure_schema<'py>( ) -> PyResult> { let connection = Connection::writer(url, user, password).map_err(map_error)?; let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; - litellm_host_python::run_async( + crate::logger::run_async( py, async move { litellm_traces::ensure_schema( @@ -64,7 +64,7 @@ pub fn trace_query<'py>( ) -> PyResult> { let connection = Connection::configured(url, database, user, password).map_err(map_error)?; let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; - litellm_host_python::run_async( + crate::logger::run_async( py, async move { litellm_traces::execute_read(&client, &connection, &sql, ¶meters).await }, map_error, From c3f3b1436f3e4f198a1f6a76072b2ca2eeec6996 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 09:16:55 -0700 Subject: [PATCH 10/12] docs(traces): replace foundation README with boundary rules --- litellm-rust/crates/traces/AGENTS.md | 6 ++++++ litellm-rust/crates/traces/README.md | 17 ----------------- 2 files changed, 6 insertions(+), 17 deletions(-) create mode 100644 litellm-rust/crates/traces/AGENTS.md delete mode 100644 litellm-rust/crates/traces/README.md diff --git a/litellm-rust/crates/traces/AGENTS.md b/litellm-rust/crates/traces/AGENTS.md new file mode 100644 index 000000000000..d330429b3157 --- /dev/null +++ b/litellm-rust/crates/traces/AGENTS.md @@ -0,0 +1,6 @@ +- Rust owns OTLP wire decoding, ClickHouse schema, row encoding, named reads, connection validation and transport +- 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 +- Use typed query parameters and a dedicated SELECT-only reader with server-side limits +- Bound insert time and encoded bytes; make retry deduplication behavior explicit for supported ClickHouse versions +- Test storage behavior through the crate's public API against ClickHouse diff --git a/litellm-rust/crates/traces/README.md b/litellm-rust/crates/traces/README.md deleted file mode 100644 index 7f7af15a59e2..000000000000 --- a/litellm-rust/crates/traces/README.md +++ /dev/null @@ -1,17 +0,0 @@ -# Trace storage foundation - -This crate owns the initial ClickHouse schema, insert row encoding and parameterized read transport. Python adapters use `trace_ensure_schema`, `trace_encode_rows` and `trace_query` from the native bridge. HTTP endpoints, ingestion and trace response assembly belong to the integration built on this foundation - -The SQL files under `migrations/` are the canonical definitions for `otel_traces`, `agent_traces`, its materialized view, and `spend_logs`. `ensure_schema` applies them in order using writer credentials and checks each response. Setup is repeatable for a new installation; these initial `CREATE IF NOT EXISTS` statements do not upgrade an existing incompatible table. Future schema changes need explicit upgrade statements - -`execute_read` uses ClickHouse typed placeholders and binds text, integer and string-array parameters separately from SQL. It obtains no privileges itself: callers must use a dedicated SELECT-only reader. The helper enforces a 15-second HTTP timeout, a 10-second query setting, 1,000 result rows and a 4 MiB response cap, and rejects errors embedded in HTTP 200 JSON responses. Large traces fail explicitly at those limits; span pagination is future work - -Pass a shared `HttpClientPool` client using `ClientVariant::NoRedirect` and the host HTTP settings. The native bridge already does this. The SQL helper is internal infrastructure, not a public arbitrary-SQL endpoint - -For self-hosted ClickHouse, install [config/reader.xml](config/reader.xml) as `/etc/clickhouse-server/users.d/litellm-traces-reader.xml` and set `LITELLM_TRACES_READER_PASSWORD` on the server. Change `default` in the three SELECT grants to the configured database and restrict the network range to the proxy network. The profile locks read-only mode, execution time, result rows and bytes, overflow behavior and memory limits. Managed ClickHouse deployments need equivalent grants and settings constraints - -Keep reader, ingestion and migration credentials separate. Do not grant the reader write, backup, named-collection management or grant-option privileges, including through roles. `readonly=1` alone is insufficient for privileged accounts; see [query permissions](https://clickhouse.com/docs/concepts/features/configuration/settings/permissions-for-queries) and [settings constraints](https://clickhouse.com/docs/concepts/features/configuration/settings/constraints-on-settings) - -Run `cargo test -p litellm-traces -- --test-threads=2` from `litellm-rust`. Testcontainers applies the canonical schema, checks span rollups and spend joins, and exercises reader permissions, parameter binding, credentials and response limits - -`encode_rows` produces JSONEachRow and formats span nanoseconds and spend-log milliseconds as UTC DateTime64 strings without losing precision. The integration owns compression, batching and writes From 1f3eb17a3d4e682f5674b47de541f76278abbc11 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 09:40:20 -0700 Subject: [PATCH 11/12] fix(traces): use current bridge execution facade --- litellm-rust/crates/python-bridge/src/routes/traces.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index eb7ff16af28e..fdd2a7eb93d8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -33,7 +33,7 @@ pub fn trace_ensure_schema<'py>( ) -> PyResult> { let connection = Connection::writer(url, user, password).map_err(map_error)?; let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; - crate::logger::run_async( + crate::execution::run_async( py, async move { litellm_traces::ensure_schema( @@ -64,7 +64,7 @@ pub fn trace_query<'py>( ) -> PyResult> { let connection = Connection::configured(url, database, user, password).map_err(map_error)?; let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; - crate::logger::run_async( + crate::execution::run_async( py, async move { litellm_traces::execute_read(&client, &connection, &sql, ¶meters).await }, map_error, From b0518f2c196025e319ed906d104979fab5bc5841 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 30 Sep 2026 18:49:11 +0000 Subject: [PATCH 12/12] fix(traces): account for protocol cast in lint budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/rust_bridge/traces.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/rust_bridge/traces.py b/litellm/rust_bridge/traces.py index 782d63a3aa78..cb4fd4a5bd3e 100644 --- a/litellm/rust_bridge/traces.py +++ b/litellm/rust_bridge/traces.py @@ -39,7 +39,7 @@ def _native() -> NativeTraces: native: Final = get_native_bridge() if native is None: raise RuntimeError("Agent tracing requires the Rust extension") - return cast(NativeTraces, native) + return cast(NativeTraces, native) # cast-ok: the native extension is validated against this protocol at call sites async def ensure_schema(