diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 348beeb3813..f7b667c8ab2 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4156,6 +4156,7 @@ dependencies = [ "litellm-storage-clickhouse", "litellm-token-counter", "litellm-traces", + "litellm-traces-cache", "litellm-traces-clickhouse", "litellm-tracing", "prost", @@ -4474,13 +4475,17 @@ dependencies = [ name = "litellm-traces-cache" version = "0.1.0" dependencies = [ + "base64 0.22.1", "litellm-traces", "moka", "rstest", + "serde", "serde_json", "sha2 0.10.9", "thiserror 2.0.19", + "time", "tokio", + "tracing", ] [[package]] @@ -4488,11 +4493,9 @@ name = "litellm-traces-clickhouse" version = "0.1.0" dependencies = [ "askama", - "base64 0.22.1", "flate2", "futures-util", "hmac 0.12.1", - "itertools 0.14.0", "jsonschema", "litellm-http", "litellm-migrate", @@ -4511,7 +4514,6 @@ dependencies = [ "thiserror 2.0.19", "time", "tokio", - "tracing", "url", "wiremock", ] diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index c3a86009111..226cecd5b55 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -22,6 +22,7 @@ tiktoken = ["litellm-token-counter/tiktoken"] fancy-regex.workspace = true litellm-tracing.workspace = true litellm-traces.workspace = true +litellm-traces-cache.workspace = true litellm-traces-clickhouse.workspace = true litellm-storage-clickhouse.workspace = true litellm-host.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index dd5c6b9860f..f442d6c31de 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -1,8 +1,11 @@ -use std::collections::BTreeMap; +use std::{collections::BTreeMap, sync::Arc}; use litellm_http::ClientVariant; use litellm_traces::{QueryScope, ReadQuery, Tenant, query::named::ReadAccessParams}; -use litellm_traces_clickhouse::{Config, Error, InsertTable, Parameter, QueryReaders}; +use litellm_traces_cache::{ReadError, TraceReader}; +use litellm_traces_clickhouse::{ + ClickHouseTraces, Config, Error, InsertTable, Parameter, QueryReaders, +}; use prost::Message; use pyo3::{ exceptions::{PyOverflowError, PyRuntimeError, PyValueError}, @@ -10,6 +13,8 @@ use pyo3::{ types::PyBytes, }; +pyo3::import_exception!(litellm.rust_bridge.trace.errors, TraceChanged); + #[derive(Message)] struct OtlpErrorStatus { #[prost(int32, tag = "1")] @@ -35,15 +40,12 @@ fn map_error_ref(error: &Error) -> PyErr { use litellm_storage_clickhouse::Error as StorageError; match error { - Error::Decode(litellm_traces::Error::TooLarge) - | Error::InsertTooLarge - | Error::ReadTooLarge => PyOverflowError::new_err(error.to_string()), + Error::Decode(litellm_traces::Error::TooLarge) | Error::InsertTooLarge => { + PyOverflowError::new_err(error.to_string()) + } Error::InvalidRow | Error::InvalidLimit(_) | Error::InvalidTable - | Error::InvalidCursor(_) - | Error::AmbiguousTrace - | Error::TraceChanged | Error::Decode(_) | Error::InvalidSchema | Error::InvalidQuery @@ -78,6 +80,18 @@ fn map_error_ref(error: &Error) -> PyErr { } } +fn map_read_error(error: ReadError) -> PyErr { + match error { + error @ (ReadError::InvalidParameters + | ReadError::InvalidCursor(_) + | ReadError::AmbiguousTrace) => PyValueError::new_err(error.to_string()), + error @ ReadError::TraceChanged => TraceChanged::new_err(error.to_string()), + error @ ReadError::TooLarge => PyOverflowError::new_err(error.to_string()), + error @ ReadError::Encode(_) => PyRuntimeError::new_err(error.to_string()), + ReadError::Store(error) => map_error_ref(&error), + } +} + fn map_sql_error(error: Error) -> PyErr { match error { Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(400 | 404)) => { @@ -112,6 +126,7 @@ impl NativeTraceConfig { pub struct NativeTraceStorage { config: Config, query_readers: QueryReaders, + reader: Arc, } #[pymethods] @@ -123,6 +138,9 @@ impl NativeTraceStorage { config.inner.storage().writer().clone(), config.inner.storage().database().to_owned(), ), + reader: Arc::new(TraceReader::new( + litellm_storage_clickhouse::READ_LIMITS.response_bytes, + )), config: config.inner.clone(), }) } @@ -218,21 +236,16 @@ impl NativeTraceStorage { ) -> PyResult> { let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); + let reader = Arc::clone(&self.reader); crate::execution::run_async( py, async move { - litellm_traces_clickhouse::list_traces( - &client, - &connection, - &scope, - start_ms, - end_ms, - cursor.as_deref(), - limit, - ) - .await + let store = ClickHouseTraces::new(client, connection); + reader + .list_traces(&store, &scope, start_ms, end_ms, cursor.as_deref(), limit) + .await }, - map_error, + map_read_error, ) } @@ -248,34 +261,31 @@ impl NativeTraceStorage { ) -> PyResult> { let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); + let reader = Arc::clone(&self.reader); crate::execution::run_async( py, async move { + let store = ClickHouseTraces::new(client, connection); if let Some(page_size) = page_size { - litellm_traces_clickhouse::get_trace_page( - &client, - &connection, - &scope, - &trace_id, - &trace_ref, - cursor.as_deref(), - page_size, - ) - .await + reader + .get_trace_page( + &store, + &scope, + &trace_id, + &trace_ref, + cursor.as_deref(), + page_size, + ) + .await } else if cursor.is_some() { - Err(Error::InvalidParameters) + Err(ReadError::InvalidParameters) } else { - litellm_traces_clickhouse::get_trace( - &client, - &connection, - &scope, - &trace_id, - &trace_ref, - ) - .await + reader + .get_trace(&store, &scope, &trace_id, &trace_ref) + .await } }, - map_error, + map_read_error, ) } @@ -289,20 +299,16 @@ impl NativeTraceStorage { ) -> PyResult> { let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); + let reader = Arc::clone(&self.reader); crate::execution::run_async( py, async move { - litellm_traces_clickhouse::get_span( - &client, - &connection, - &scope, - &trace_id, - &span_id, - &trace_ref, - ) - .await + let store = ClickHouseTraces::new(client, connection); + reader + .get_span(&store, &scope, &trace_id, &span_id, &trace_ref) + .await }, - map_error, + map_read_error, ) } @@ -318,21 +324,23 @@ impl NativeTraceStorage { ) -> PyResult> { let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); + let reader = Arc::clone(&self.reader); crate::execution::run_async( py, async move { - litellm_traces_clickhouse::get_span_error( - &client, - &connection, - &scope, - &trace_id, - &span_id, - &trace_ref, - cursor.as_deref(), - ) - .await + let store = ClickHouseTraces::new(client, connection); + reader + .get_span_error( + &store, + &scope, + &trace_id, + &span_id, + &trace_ref, + cursor.as_deref(), + ) + .await }, - map_error, + map_read_error, ) } @@ -495,11 +503,7 @@ mod tests { Error::Decode(litellm_traces::Error::InvalidLimit("OTLP_MAX_SPANS")), "ValueError" )] - #[case::cursor(Error::InvalidCursor("trace"), "ValueError")] - #[case::ambiguous(Error::AmbiguousTrace, "ValueError")] - #[case::changed_snapshot(Error::TraceChanged, "ValueError")] - #[case::read_budget(Error::ReadTooLarge, "OverflowError")] - fn trace_read_and_ingest_failures_preserve_public_exception_types( + fn trace_ingest_failures_preserve_public_exception_types( #[case] error: Error, #[case] exception_name: &str, ) { @@ -511,4 +515,43 @@ mod tests { ); }); } + + #[rstest] + #[case::invalid_parameters(ReadError::InvalidParameters, "ValueError")] + #[case::invalid_cursor(ReadError::InvalidCursor("trace"), "ValueError")] + #[case::ambiguous(ReadError::AmbiguousTrace, "ValueError")] + #[case::changed_snapshot(ReadError::TraceChanged, "TraceChanged")] + #[case::read_budget(ReadError::TooLarge, "OverflowError")] + #[case::encode( + ReadError::Encode(Arc::new(serde_json::Error::io(std::io::Error::other("invalid")))), + "RuntimeError" + )] + #[case::store(ReadError::Store(Arc::new(Error::InvalidScope)), "ValueError")] + fn trace_read_failures_preserve_public_exception_types( + #[case] error: ReadError, + #[case] exception_name: &str, + ) { + Python::initialize(); + Python::attach(|py| { + let repository = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .ancestors() + .nth(3) + .unwrap() + .to_str() + .unwrap(); + pyo3::types::PyModule::import(py, "sys") + .unwrap() + .getattr("path") + .unwrap() + .call_method1("insert", (0, repository)) + .unwrap(); + let message = error.to_string(); + let exception = map_read_error(error); + assert_eq!(exception.get_type(py).name().unwrap(), exception_name); + assert_eq!( + exception.value(py).str().unwrap().to_str().unwrap(), + message + ); + }); + } } diff --git a/litellm-rust/crates/traces-cache/AGENTS.md b/litellm-rust/crates/traces-cache/AGENTS.md index bfaea42d901..7e4caacd15e 100644 --- a/litellm-rust/crates/traces-cache/AGENTS.md +++ b/litellm-rust/crates/traces-cache/AGENTS.md @@ -1,5 +1,5 @@ -Own resolved-trace snapshot storage, cache identity, weighting, and expiry -Depend on trace domain types, never storage, HTTP, or Python +Own storage-independent trace reads over `TraceStore`: the in-process read cache (identity, single-flight, freshness expiry, weighting), cursor formats, paging, response splitting, spend windows and run batching +Depend on trace domain types, never storage, HTTP or Python Preserve the full source and authorization scope in every cache key Keep snapshots immutable and expose borrowed data -Keep cursor formats and database reads in their existing owners +Storage adapters implement `TraceStore`; keep SQL and row encoding there diff --git a/litellm-rust/crates/traces-cache/Cargo.toml b/litellm-rust/crates/traces-cache/Cargo.toml index 3e58ea4c965..16eb01c6c8a 100644 --- a/litellm-rust/crates/traces-cache/Cargo.toml +++ b/litellm-rust/crates/traces-cache/Cargo.toml @@ -6,11 +6,15 @@ license.workspace = true repository.workspace = true [dependencies] +base64.workspace = true litellm-traces.workspace = true moka.workspace = true +serde.workspace = true serde_json.workspace = true sha2.workspace = true +time.workspace = true thiserror.workspace = true +tracing.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/traces-cache/src/cache.rs b/litellm-rust/crates/traces-cache/src/cache.rs new file mode 100644 index 00000000000..f308714146e --- /dev/null +++ b/litellm-rust/crates/traces-cache/src/cache.rs @@ -0,0 +1,395 @@ +use std::{future::Future, sync::Arc, time::Duration}; + +use litellm_traces::{ + Trace, TraceSummary, + query::named::{ReadAccessParams, TraceSpansRow}, +}; +use moka::{Expiry, future::Cache}; +use serde::Serialize; +use sha2::{Digest, Sha256}; + +use crate::Error; + +pub const LIVE_TTL: Duration = Duration::from_secs(5); +pub const SETTLED_TTL: Duration = Duration::from_secs(10 * 60); +const SETTLED_AFTER_MS: u64 = 5 * 60 * 1000; +const MAX_INDEX_ENTRIES: u64 = 100_000; + +#[derive(Clone, Eq, Hash, PartialEq)] +pub struct SnapshotKey(String); + +impl SnapshotKey { + fn digest(fields: &impl Serialize) -> Result { + Ok(Self(format!( + "{:x}", + Sha256::digest(serde_json::to_vec(fields)?) + ))) + } + + pub fn new( + source: &str, + access: &ReadAccessParams, + trace_id: &str, + trace_ref: &str, + snapshot_ms: u64, + ) -> Result { + Self::digest(&(source, access, trace_id, trace_ref, snapshot_ms)) + } + + pub fn latest( + source: &str, + access: &ReadAccessParams, + trace_id: &str, + trace_ref: &str, + ) -> Result { + Self::digest(&(source, access, trace_id, trace_ref)) + } + + pub(crate) fn run( + source: &str, + access: &ReadAccessParams, + run: (&str, &str, &str, &str), + ) -> Result { + Self::digest(&("run", source, access, run)) + } + + pub(crate) fn scope(source: &str, access: &ReadAccessParams) -> Result { + Self::digest(&("scope", source, access)) + } +} + +/// How long a read result stays reusable: traces still receiving spans, or read with spend +/// unavailable, are re-read after `LIVE_TTL`; traces quiet for `SETTLED_AFTER_MS` are kept for +/// `SETTLED_TTL`. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Freshness { + Live, + Settled, +} + +impl Freshness { + pub fn of(rows: &[TraceSpansRow], spend_known: bool, snapshot_ms: u64) -> Self { + let last_end_ms = rows + .iter() + .map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns) / 1_000_000) + .max() + .unwrap_or(i64::MAX); + let quiet_ms = i64::try_from(snapshot_ms) + .unwrap_or(i64::MAX) + .saturating_sub(last_end_ms); + if spend_known && quiet_ms >= SETTLED_AFTER_MS as i64 { + Self::Settled + } else { + Self::Live + } + } + + fn ttl(self) -> Duration { + match self { + Self::Live => LIVE_TTL, + Self::Settled => SETTLED_TTL, + } + } +} + +trait Fresh { + fn freshness(&self) -> Freshness; +} + +struct ByFreshness; + +impl Expiry for ByFreshness { + fn expire_after_create(&self, _: &K, value: &V, _: std::time::Instant) -> Option { + Some(value.freshness().ttl()) + } +} + +pub struct Snapshot { + trace: Trace, + version: String, + snapshot_ms: u64, + freshness: Freshness, + weight: u32, +} + +impl Snapshot { + pub fn trace(&self) -> &Trace { + &self.trace + } + + pub fn version(&self) -> &str { + &self.version + } + + pub fn snapshot_ms(&self) -> u64 { + self.snapshot_ms + } + + pub fn freshness(&self) -> Freshness { + self.freshness + } +} + +#[derive(Clone, Copy)] +struct Latest { + snapshot_ms: u64, + freshness: Freshness, +} + +impl Fresh for Latest { + fn freshness(&self) -> Freshness { + self.freshness + } +} + +/// Resolved trace snapshots pinned by `snapshot_ms` for paging, plus which snapshot each trace +/// currently serves so repeated opens reuse one read until its freshness expires. +pub struct SnapshotCache { + pinned: Cache>, + latest: Cache, + max_graph_bytes: usize, +} + +impl SnapshotCache { + pub fn new(max_graph_bytes: usize, idle: Duration) -> Self { + Self { + pinned: Cache::builder() + .max_capacity((max_graph_bytes as u64).saturating_mul(2)) + .weigher(|_: &SnapshotKey, snapshot: &Arc| snapshot.weight) + .time_to_idle(idle) + .build(), + latest: Cache::builder() + .max_capacity(MAX_INDEX_ENTRIES) + .expire_after(ByFreshness) + .build(), + max_graph_bytes, + } + } + + pub async fn get(&self, key: &SnapshotKey) -> Option> { + self.pinned.get(key).await + } + + /// Returns the snapshot pinned at `key`, running `load` once for all concurrent callers on a + /// miss. A failed load is not cached. + pub async fn pinned_or_load( + &self, + key: SnapshotKey, + snapshot_ms: u64, + load: F, + ) -> Result, Arc> + where + E: From + Send + Sync + 'static, + F: Future>, + { + self.pinned + .try_get_with(key, async { + let (trace, freshness) = load.await?; + Ok(Arc::new(self.snapshot(trace, snapshot_ms, freshness)?)) + }) + .await + } + + /// Returns the snapshot `latest` currently serves. On a miss, `load_at(now_ms)` runs once for + /// all concurrent callers and its snapshot is served until its freshness expires. + pub async fn latest_or_load( + &self, + latest: SnapshotKey, + now_ms: u64, + load_at: F, + ) -> Result, Arc> + where + E: Send + Sync + 'static, + F: Fn(u64) -> Fut, + Fut: Future, Arc>>, + { + let entry = self + .latest + .try_get_with(latest, async { + let snapshot = load_at(now_ms).await?; + Ok::<_, Arc>(Latest { + snapshot_ms: snapshot.snapshot_ms, + freshness: snapshot.freshness, + }) + }) + .await + .map_err(|error| Arc::clone(&*error))?; + load_at(entry.snapshot_ms).await + } + + fn snapshot( + &self, + trace: Trace, + snapshot_ms: u64, + freshness: Freshness, + ) -> Result { + let encoded = serde_json::to_vec(&trace)?; + if encoded.len() > self.max_graph_bytes { + return Err(Error::ReadTooLarge); + } + let span_ids: Vec<&str> = trace + .spans + .iter() + .map(|span| span.span_id.as_str()) + .collect(); + let version = format!("{:x}", Sha256::digest(serde_json::to_vec(&span_ids)?)); + Ok(Snapshot { + trace, + version, + snapshot_ms, + freshness, + weight: u32::try_from(encoded.len().saturating_mul(2)).unwrap_or(u32::MAX), + }) + } + + #[cfg(test)] + async fn weighted_size(&self) -> u64 { + self.pinned.run_pending_tasks().await; + self.pinned.weighted_size() + } +} + +#[derive(Clone)] +pub(crate) enum ListedRun { + Resolved(Box, Freshness), + Limited, +} + +impl Fresh for ListedRun { + fn freshness(&self) -> Freshness { + match self { + Self::Resolved(_, freshness) => *freshness, + Self::Limited => Freshness::Settled, + } + } +} + +pub(crate) struct ListCache { + pub(crate) runs: Cache, + pub(crate) limits: Cache, +} + +impl ListCache { + pub(crate) fn new() -> Self { + Self { + runs: Cache::builder() + .max_capacity(MAX_INDEX_ENTRIES) + .expire_after(ByFreshness) + .build(), + limits: Cache::builder() + .max_capacity(MAX_INDEX_ENTRIES) + .time_to_live(SETTLED_TTL) + .build(), + } + } +} + +#[cfg(test)] +mod tests { + use litellm_traces::{ + SpanStatus, + query::named::{SpendByResponseIdsRow, TraceSpansRow}, + resolve_trace, + }; + use rstest::rstest; + + use super::*; + + fn row(span_id: &str) -> TraceSpansRow { + TraceSpansRow { + trace_id: String::new(), + span_id: span_id.into(), + parent_span_id: String::new(), + name: "run".into(), + kind: litellm_traces::ObservationType::Agent, + wrapper_candidate: false, + agent: "agent".into(), + framework: String::new(), + status: SpanStatus::Ok, + status_message: String::new(), + error_truncated: false, + start_ns: 1_790_742_989_000_000_000, + duration_ns: 10_000_000, + service: "agent-demo".into(), + input_preview: format!("input of {span_id}"), + model: String::new(), + input_tokens: 0, + output_tokens: 0, + litellm_request_id: String::new(), + call_keys: Vec::new(), + call_evidence: None, + tool_call_id: String::new(), + team_id: String::new(), + api_key_hash: String::new(), + user_id: String::new(), + } + } + + fn trace(span_id: &str) -> Trace { + resolve_trace( + "trace", + "ref", + &[row(span_id)], + &[] as &[SpendByResponseIdsRow], + ) + .expect("fixture should resolve") + } + + fn key(suffix: &str) -> SnapshotKey { + SnapshotKey::new( + "source", + &ReadAccessParams { + all_teams: false, + user_id: String::new(), + team_ids: vec!["team".into()], + }, + suffix, + "ref", + 100, + ) + .unwrap() + } + + #[tokio::test] + async fn weighted_capacity_bounds_retained_snapshots() { + let limit = ["first", "second", "third"] + .iter() + .map(|span_id| serde_json::to_vec(&trace(span_id)).unwrap().len()) + .max() + .unwrap(); + let cache = SnapshotCache::new(limit, Duration::from_secs(120)); + + for (key, span_id) in [ + (key("a"), "first"), + (key("b"), "second"), + (key("c"), "third"), + ] { + cache + .pinned_or_load(key, 100, async { + Ok::<_, Error>((trace(span_id), Freshness::Settled)) + }) + .await + .unwrap(); + } + + assert!(cache.weighted_size().await <= (limit as u64) * 2); + } + + const LAST_END_MS: u64 = 1_790_742_989_010; + + #[rstest] + #[case::just_ended(LAST_END_MS, true, Freshness::Live)] + #[case::quiet_just_under(LAST_END_MS + SETTLED_AFTER_MS - 1, true, Freshness::Live)] + #[case::quiet_long_enough(LAST_END_MS + SETTLED_AFTER_MS, true, Freshness::Settled)] + #[case::spend_unknown(LAST_END_MS + SETTLED_AFTER_MS, false, Freshness::Live)] + fn freshness_settles_once_spans_stop_and_spend_is_known( + #[case] snapshot_ms: u64, + #[case] spend_known: bool, + #[case] expected: Freshness, + ) { + assert_eq!( + Freshness::of(&[row("root")], spend_known, snapshot_ms), + expected + ); + } +} diff --git a/litellm-rust/crates/traces-cache/src/cursor.rs b/litellm-rust/crates/traces-cache/src/cursor.rs new file mode 100644 index 00000000000..8858f243f09 --- /dev/null +++ b/litellm-rust/crates/traces-cache/src/cursor.rs @@ -0,0 +1,124 @@ +use base64::{Engine, engine::general_purpose::URL_SAFE}; +use serde::{Deserialize, Serialize}; + +use crate::ReadError; + +pub(super) fn encode_cursor(position: &T) -> String { + URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default()) +} + +pub(super) fn decode_cursor Deserialize<'de>, E>( + cursor: &str, + kind: &'static str, +) -> Result> { + URL_SAFE + .decode(cursor) + .ok() + .and_then(|json| serde_json::from_slice(&json).ok()) + .ok_or(ReadError::InvalidCursor(kind)) +} + +pub(super) fn trace_position(cursor: Option<&str>) -> Result<(i64, String), ReadError> { + let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else { + return Ok((0, String::new())); + }; + match decode_cursor::<(i64, String), E>(cursor, "trace")? { + (start_ms, trace_ref) if start_ms > 0 && !trace_ref.is_empty() => Ok((start_ms, trace_ref)), + _ => Err(ReadError::InvalidCursor("trace")), + } +} + +#[derive(Deserialize, Serialize)] +pub(super) struct ErrorPosition { + pub(super) offset: u64, + pub(super) version: String, +} + +pub(super) fn error_position( + cursor: Option<&str>, +) -> Result, ReadError> { + let Some(cursor) = cursor else { + return Ok(None); + }; + let position: ErrorPosition = decode_cursor(cursor, "diagnostic")?; + let valid_version = position.version.len() == 64 + && position + .version + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte)); + if i64::try_from(position.offset).is_err() || !valid_version { + return Err(ReadError::InvalidCursor("diagnostic")); + } + Ok(Some(position)) +} + +#[derive(Deserialize, Serialize)] +pub(super) struct SpanPosition { + pub(super) trace_ref: String, + pub(super) snapshot_ms: u64, + pub(super) offset: usize, + pub(super) version: String, +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + fn trace_cursor_round_trips_the_last_listed_run() { + let cursor = encode_cursor(&(1_790_742_989_377_i64, "4bad42b84e9de3ba46fc870185f8f023")); + assert_eq!( + trace_position::(Some(&cursor)).unwrap(), + ( + 1_790_742_989_377, + "4bad42b84e9de3ba46fc870185f8f023".to_owned() + ) + ); + assert_eq!( + trace_position::(None).unwrap(), + (0, String::new()) + ); + assert_eq!( + trace_position::(Some("")).unwrap(), + (0, String::new()) + ); + } + + #[rstest] + #[case::not_base64("abc")] + #[case::not_json("bm90LWpzb24=")] + #[case::numeric_reference("WzEsIDJd")] + #[case::zero_start("WzAsICJ0Il0=")] + fn malformed_trace_cursors_are_rejected(#[case] cursor: &str) { + let result: Result<(i64, String), ReadError> = trace_position(Some(cursor)); + assert!(matches!(result, Err(ReadError::InvalidCursor("trace")))); + } + + #[rstest] + #[case::not_base64("garbage")] + #[case::missing_fields("e30=")] + #[case::not_an_object("WzEsMl0=")] + fn malformed_diagnostic_cursors_are_rejected(#[case] cursor: &str) { + let result: Result, ReadError> = + error_position(Some(cursor)); + assert!(matches!( + result, + Err(ReadError::InvalidCursor("diagnostic")) + )); + } + + #[rstest] + #[case::lowercase_version("a".repeat(64))] + #[case::short_version("A".repeat(63))] + fn diagnostic_cursor_requires_a_content_version(#[case] version: String) { + let cursor = encode_cursor(&ErrorPosition { offset: 1, version }); + let result: Result, ReadError> = + error_position(Some(&cursor)); + assert!(matches!( + result, + Err(ReadError::InvalidCursor("diagnostic")) + )); + } +} diff --git a/litellm-rust/crates/traces-cache/src/error.rs b/litellm-rust/crates/traces-cache/src/error.rs index ac5f19369fe..c7aaa3fd332 100644 --- a/litellm-rust/crates/traces-cache/src/error.rs +++ b/litellm-rust/crates/traces-cache/src/error.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + #[derive(Debug, thiserror::Error)] pub enum Error { #[error("trace snapshot serialization failed")] @@ -5,3 +7,45 @@ pub enum Error { #[error("trace snapshot exceeds the size limit")] ReadTooLarge, } + +/// Cheap to clone so one failed single-flight read can be returned to every waiting caller. +#[derive(Debug, thiserror::Error)] +pub enum ReadError { + #[error("invalid trace read parameters")] + InvalidParameters, + #[error("Invalid {0} cursor")] + InvalidCursor(&'static str), + #[error("Multiple traces have this ID; provide trace_ref")] + AmbiguousTrace, + #[error("Trace changed while paging; refresh the trace to continue")] + TraceChanged, + #[error("Trace exceeds the interactive read budget; use a filtered trace query")] + TooLarge, + #[error("trace could not be encoded")] + Encode(#[source] Arc), + #[error(transparent)] + Store(Arc), +} + +impl Clone for ReadError { + fn clone(&self) -> Self { + match self { + Self::InvalidParameters => Self::InvalidParameters, + Self::InvalidCursor(kind) => Self::InvalidCursor(kind), + Self::AmbiguousTrace => Self::AmbiguousTrace, + Self::TraceChanged => Self::TraceChanged, + Self::TooLarge => Self::TooLarge, + Self::Encode(error) => Self::Encode(Arc::clone(error)), + Self::Store(error) => Self::Store(Arc::clone(error)), + } + } +} + +impl From for ReadError { + fn from(error: Error) -> Self { + match error { + Error::ReadTooLarge => Self::TooLarge, + Error::Serialization(error) => Self::Encode(Arc::new(error)), + } + } +} diff --git a/litellm-rust/crates/traces-cache/src/lib.rs b/litellm-rust/crates/traces-cache/src/lib.rs index 7dfd3bf32a1..8ed79c1a547 100644 --- a/litellm-rust/crates/traces-cache/src/lib.rs +++ b/litellm-rust/crates/traces-cache/src/lib.rs @@ -1,166 +1,12 @@ -use std::{sync::Arc, time::Duration}; - -use litellm_traces::{Trace, query::named::ReadAccessParams}; -use moka::future::Cache; -use sha2::{Digest, Sha256}; - +mod cache; +mod cursor; mod error; +mod list; +mod reader; +mod spend; +mod store; -pub use error::Error; - -#[derive(Clone, Eq, Hash, PartialEq)] -pub struct SnapshotKey(String); - -impl SnapshotKey { - pub fn new( - source: &str, - access: &ReadAccessParams, - trace_id: &str, - trace_ref: &str, - snapshot_ms: u64, - ) -> Result { - let encoded = serde_json::to_vec(&(source, access, trace_id, trace_ref, snapshot_ms))?; - Ok(Self(format!("{:x}", Sha256::digest(encoded)))) - } -} - -pub struct Snapshot { - trace: Trace, - version: String, - weight: u32, -} - -impl Snapshot { - pub fn trace(&self) -> &Trace { - &self.trace - } - - pub fn version(&self) -> &str { - &self.version - } -} - -pub struct SnapshotCache { - entries: Cache>, - max_graph_bytes: usize, -} - -impl SnapshotCache { - pub fn new(max_graph_bytes: usize, ttl: Duration) -> Self { - Self { - entries: Cache::builder() - .max_capacity((max_graph_bytes as u64).saturating_mul(2)) - .weigher(|_: &SnapshotKey, snapshot: &Arc| snapshot.weight) - .time_to_live(ttl) - .build(), - max_graph_bytes, - } - } - - pub async fn get(&self, key: &SnapshotKey) -> Option> { - self.entries.get(key).await - } - - pub async fn insert(&self, key: SnapshotKey, trace: Trace) -> Result, Error> { - let encoded = serde_json::to_vec(&trace)?; - if encoded.len() > self.max_graph_bytes { - return Err(Error::ReadTooLarge); - } - - let span_ids: Vec<&str> = trace - .spans - .iter() - .map(|span| span.span_id.as_str()) - .collect(); - - let version = format!("{:x}", Sha256::digest(serde_json::to_vec(&span_ids)?)); - - let snapshot = Arc::new(Snapshot { - trace, - version, - weight: u32::try_from(encoded.len().saturating_mul(2)).unwrap_or(u32::MAX), - }); - - self.entries.insert(key, Arc::clone(&snapshot)).await; - Ok(snapshot) - } -} - -#[cfg(test)] -mod tests { - use litellm_traces::{ - SpanStatus, - query::named::{SpendByResponseIdsRow, TraceSpansRow}, - resolve_trace, - }; - - use super::*; - - fn trace(span_id: &str) -> Trace { - let rows = [TraceSpansRow { - trace_id: String::new(), - span_id: span_id.into(), - parent_span_id: String::new(), - name: "run".into(), - kind: litellm_traces::ObservationType::Agent, - wrapper_candidate: false, - agent: "agent".into(), - framework: String::new(), - status: SpanStatus::Ok, - status_message: String::new(), - error_truncated: false, - start_ns: 1_790_742_989_000_000_000, - duration_ns: 10_000_000, - service: "agent-demo".into(), - input_preview: format!("input of {span_id}"), - model: String::new(), - input_tokens: 0, - output_tokens: 0, - litellm_request_id: String::new(), - call_keys: Vec::new(), - call_evidence: None, - tool_call_id: String::new(), - team_id: String::new(), - api_key_hash: String::new(), - user_id: String::new(), - }]; - resolve_trace("trace", "ref", &rows, &[] as &[SpendByResponseIdsRow]) - .expect("fixture should resolve") - } - - fn key(suffix: &str) -> SnapshotKey { - SnapshotKey::new( - "source", - &ReadAccessParams { - all_teams: false, - user_id: String::new(), - team_ids: vec!["team".into()], - }, - suffix, - "ref", - 100, - ) - .unwrap() - } - - #[tokio::test] - async fn weighted_capacity_bounds_retained_snapshots() { - let limit = ["first", "second", "third"] - .iter() - .map(|span_id| serde_json::to_vec(&trace(span_id)).unwrap().len()) - .max() - .unwrap(); - let cache = SnapshotCache::new(limit, Duration::from_secs(120)); - - for (key, span_id) in [ - (key("a"), "first"), - (key("b"), "second"), - (key("c"), "third"), - ] { - cache.insert(key, trace(span_id)).await.unwrap(); - } - - cache.entries.run_pending_tasks().await; - assert!(cache.entries.weighted_size() <= (limit as u64) * 2); - } -} +pub use cache::{Freshness, LIVE_TTL, SETTLED_TTL, Snapshot, SnapshotCache, SnapshotKey}; +pub use error::{Error, ReadError}; +pub use reader::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, TraceReader}; +pub use store::{StoreError, TraceStore}; diff --git a/litellm-rust/crates/traces-cache/src/list.rs b/litellm-rust/crates/traces-cache/src/list.rs new file mode 100644 index 00000000000..7093411bc8c --- /dev/null +++ b/litellm-rust/crates/traces-cache/src/list.rs @@ -0,0 +1,191 @@ +use std::collections::HashMap; + +use crate::{ + ReadError, SnapshotKey, TraceReader, TraceStore, + cache::{Freshness, ListedRun}, + reader::{map_store_error, now_ms}, + spend::{spend, spend_window, spend_within}, + store::StoreError, +}; +use litellm_traces::{ + TraceSummary, listed_summary, + query::named::{ListTracesRow, ReadAccessParams, TracePageSpansParams, TraceSpansRow}, + resolve_trace, +}; + +const RUNS_PER_SPAN_READ: usize = 16; + +fn run_key(team_id: &str, api_key_hash: &str, trace_id: &str) -> (String, String, String) { + ( + team_id.to_owned(), + api_key_hash.to_owned(), + trace_id.to_owned(), + ) +} + +fn cache_key( + source: &str, + access: &ReadAccessParams, + row: &ListTracesRow, +) -> Result> { + Ok(SnapshotKey::run( + source, + access, + ( + &row.team_id, + &row.api_key_hash, + &row.trace_id, + &row.trace_ref, + ), + )?) +} + +fn summary(row: &ListTracesRow, listed: Option<&ListedRun>) -> TraceSummary { + match listed { + Some(ListedRun::Resolved(summary, _)) => (**summary).clone(), + Some(ListedRun::Limited) | None => listed_summary(row), + } +} + +/// Summaries for one batch of listed runs. Runs resolved within their freshness window come +/// from the cache; only the rest are read from storage, with one span and one spend read. +pub(super) async fn list_summaries( + reader: &TraceReader, + store: &S, + access: &ReadAccessParams, + runs: &[ListTracesRow], +) -> Result, ReadError> { + let mut keys = Vec::with_capacity(runs.len()); + let mut listed = Vec::with_capacity(runs.len()); + for row in runs { + let key = cache_key(store.source(), access, row)?; + listed.push(reader.lists.runs.get(&key).await); + keys.push(key); + } + let misses: Vec<&ListTracesRow> = runs + .iter() + .zip(&listed) + .filter_map(|(row, listed)| listed.is_none().then_some(row)) + .collect(); + let mut resolved = resolve_runs(reader, store, access, &misses) + .await? + .into_iter(); + let mut summaries = Vec::with_capacity(runs.len()); + for ((row, key), cached) in runs.iter().zip(keys).zip(listed) { + let listed = match cached { + Some(listed) => Some(listed), + None => { + let listed = resolved.next().flatten(); + if let Some(listed) = &listed { + reader.lists.runs.insert(key, listed.clone()).await; + } + listed + } + }; + summaries.push(summary(row, listed.as_ref())); + } + Ok(summaries) +} + +/// One entry per run; `None` means the run could not be resolved and keeps its listed summary +/// without being cached. +async fn resolve_runs( + reader: &TraceReader, + store: &S, + access: &ReadAccessParams, + runs: &[&ListTracesRow], +) -> Result>, ReadError> { + let (Some(start_ms), Some(end_ms)) = ( + runs.iter().map(|row| row.start_ms).min(), + runs.iter() + .map(|row| row.start_ms.saturating_add(row.duration_ms)) + .max(), + ) else { + return Ok(Vec::new()); + }; + let params = TracePageSpansParams { + access: access.clone(), + trace_refs: runs.iter().map(|row| row.trace_ref.clone()).collect(), + start_ms, + end_ms: end_ms.saturating_add(1), + }; + let snapshot_ms = now_ms(); + let spans = match store.run_spans(¶ms, snapshot_ms).await { + Ok(spans) => spans, + Err(StoreError::TooLarge) => { + let mut resolved = Vec::with_capacity(runs.len()); + for row in runs { + resolved.push(resolve_run(reader, store, access, row).await?); + } + return Ok(resolved); + } + Err(error) => return Err(map_store_error(error)), + }; + let spend_rows = spend(store, access, &spans).await; + let spend_known = spend_rows.is_some(); + let spend_rows = spend_rows.unwrap_or_default(); + let mut spans = spans; + spans.sort_by(|left, right| { + run_key(&left.team_id, &left.api_key_hash, &left.trace_id) + .cmp(&run_key( + &right.team_id, + &right.api_key_hash, + &right.trace_id, + )) + .then(left.start_ns.cmp(&right.start_ns)) + }); + let by_run: HashMap<_, &[TraceSpansRow]> = spans + .chunk_by(|left, right| { + (&left.team_id, &left.api_key_hash, &left.trace_id) + == (&right.team_id, &right.api_key_hash, &right.trace_id) + }) + .map(|run| { + ( + run_key(&run[0].team_id, &run[0].api_key_hash, &run[0].trace_id), + run, + ) + }) + .collect(); + Ok(runs + .iter() + .map(|row| { + let spans = by_run + .get(&run_key(&row.team_id, &row.api_key_hash, &row.trace_id)) + .copied() + .unwrap_or_default(); + let spend = + spend_window(spans).map_or(&[][..], |window| spend_within(&spend_rows, window)); + resolve_trace(&row.trace_id, &row.trace_ref, spans, spend).map(|trace| { + ListedRun::Resolved( + Box::new(trace.summary), + Freshness::of(spans, spend_known, snapshot_ms), + ) + }) + }) + .collect()) +} + +async fn resolve_run( + reader: &TraceReader, + store: &S, + access: &ReadAccessParams, + row: &ListTracesRow, +) -> Result, ReadError> { + match reader + .current(store, access, &row.trace_id, &row.trace_ref) + .await + { + Ok(snapshot) => Ok(snapshot.map(|snapshot| { + ListedRun::Resolved( + Box::new(snapshot.trace().summary.clone()), + snapshot.freshness(), + ) + })), + Err(ReadError::TooLarge) => Ok(Some(ListedRun::Limited)), + Err(error) => Err(error), + } +} + +pub(super) fn run_batches(runs: &[T]) -> impl Iterator + '_ { + runs.chunks(RUNS_PER_SPAN_READ) +} diff --git a/litellm-rust/crates/traces-cache/src/reader.rs b/litellm-rust/crates/traces-cache/src/reader.rs new file mode 100644 index 00000000000..5c68d6e58bd --- /dev/null +++ b/litellm-rust/crates/traces-cache/src/reader.rs @@ -0,0 +1,365 @@ +use std::{sync::Arc, time::Duration}; + +use crate::{ + ReadError, Snapshot, SnapshotCache, SnapshotKey, StoreError, TraceStore, + cache::{Freshness, ListCache}, + cursor::{ + ErrorPosition, SpanPosition, decode_cursor, encode_cursor, error_position, trace_position, + }, + list::{list_summaries, run_batches}, + spend::spend, +}; +use litellm_traces::{ + SpanDetail, SpanErrorPage, Trace, TracePage, + query::named::{ + ListTracesParams, ReadAccessParams, SpanDetailParams, SpanErrorParams, TraceIdentityParams, + TraceSpansParams, + }, + resolve_trace, to_ui_content, +}; + +pub const MAX_GRAPH_BYTES: usize = 64 * 1024 * 1024; +pub const MAX_GRAPH_SPANS: usize = 100_000; + +const SNAPSHOT_IDLE: Duration = Duration::from_secs(120); + +/// A read that found no trace, kept apart from failures so single-flight waiters share it +/// without it being cached. +pub(super) enum Miss { + Absent, + Read(ReadError), +} + +impl From for Miss { + fn from(error: crate::Error) -> Self { + Self::Read(error.into()) + } +} + +fn settle(result: Result>>) -> Result, ReadError> { + match result { + Ok(value) => Ok(Some(value)), + Err(miss) => match &*miss { + Miss::Absent => Ok(None), + Miss::Read(error) => Err(error.clone()), + }, + } +} + +pub struct TraceReader { + snapshots: SnapshotCache, + pub(super) lists: ListCache, + response_bytes: usize, +} + +impl TraceReader { + pub fn new(response_bytes: usize) -> Self { + Self { + snapshots: SnapshotCache::new(MAX_GRAPH_BYTES, SNAPSHOT_IDLE), + lists: ListCache::new(), + response_bytes, + } + } + + pub async fn list_traces( + &self, + store: &S, + access: &ReadAccessParams, + start_ms: i64, + end_ms: i64, + cursor: Option<&str>, + limit: u32, + ) -> Result> { + if limit == 0 { + return Err(ReadError::InvalidParameters); + } + let (cursor_ms, cursor_trace_id) = trace_position(cursor)?; + let scope = SnapshotKey::scope(store.source(), access)?; + let accepted = self.lists.limits.get(&scope).await.unwrap_or(u32::MAX); + let mut params = ListTracesParams { + access: access.clone(), + start_ms, + end_ms, + cursor_ms, + cursor_trace_id, + limit: limit.min(500).min(accepted), + }; + let page = loop { + match store.list_runs(¶ms).await { + Err(StoreError::TooLarge) if params.limit > 1 => { + params.limit /= 2; + self.lists.limits.insert(scope.clone(), params.limit).await; + } + Err(StoreError::TooLarge) => return Err(ReadError::TooLarge), + result => break result.map_err(map_store_error)?, + } + }; + let next_cursor = page + .last() + .filter(|_| page.len() == params.limit as usize) + .map(|last| encode_cursor(&(last.start_ms, &last.trace_ref))); + let data = { + let mut summaries = Vec::with_capacity(page.len()); + for batch in run_batches(&page) { + summaries.extend(list_summaries(self, store, access, batch).await?); + } + summaries + }; + Ok(TracePage { data, next_cursor }) + } + + pub async fn get_trace( + &self, + store: &S, + access: &ReadAccessParams, + trace_id: &str, + trace_ref: &str, + ) -> Result, ReadError> { + let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else { + return Ok(None); + }; + Ok(self + .current(store, access, trace_id, &trace_ref) + .await? + .map(|snapshot| snapshot.trace().clone())) + } + + pub async fn get_trace_page( + &self, + store: &S, + access: &ReadAccessParams, + trace_id: &str, + trace_ref: &str, + cursor: Option<&str>, + page_size: u32, + ) -> Result, ReadError> { + if !(1..=500).contains(&page_size) { + return Err(ReadError::InvalidParameters); + } + let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else { + return Ok(None); + }; + let Some(cursor) = cursor else { + let Some(snapshot) = self.current(store, access, trace_id, &trace_ref).await? else { + return Ok(None); + }; + let position = SpanPosition { + trace_ref, + snapshot_ms: snapshot.snapshot_ms(), + offset: 0, + version: snapshot.version().to_owned(), + }; + return page(&snapshot, &position, page_size, self.response_bytes).map(Some); + }; + let position: SpanPosition = decode_cursor(cursor, "span")?; + if position.trace_ref != trace_ref || position.snapshot_ms == 0 { + return Err(ReadError::InvalidCursor("span")); + } + let Some(snapshot) = settle( + self.pinned(store, access, trace_id, &trace_ref, position.snapshot_ms) + .await, + )? + else { + return Ok(None); + }; + if position.version != snapshot.version() { + return Err(ReadError::TraceChanged); + } + if position.offset > snapshot.trace().spans.len() { + return Err(ReadError::InvalidCursor("span")); + } + page(&snapshot, &position, page_size, self.response_bytes).map(Some) + } + + pub(super) async fn current( + &self, + store: &S, + access: &ReadAccessParams, + trace_id: &str, + trace_ref: &str, + ) -> Result>, ReadError> { + let latest = SnapshotKey::latest(store.source(), access, trace_id, trace_ref)?; + settle( + self.snapshots + .latest_or_load(latest, now_ms(), |snapshot_ms| { + self.pinned(store, access, trace_id, trace_ref, snapshot_ms) + }) + .await, + ) + } + + async fn pinned( + &self, + store: &S, + access: &ReadAccessParams, + trace_id: &str, + trace_ref: &str, + snapshot_ms: u64, + ) -> Result, Arc>> { + let key = SnapshotKey::new(store.source(), access, trace_id, trace_ref, snapshot_ms) + .map_err(|error| Arc::new(error.into()))?; + self.snapshots + .pinned_or_load(key, snapshot_ms, async { + let params = TraceSpansParams { + access: access.clone(), + trace_id: trace_id.to_owned(), + trace_ref: trace_ref.to_owned(), + }; + let rows = store + .trace_spans(¶ms, snapshot_ms) + .await + .map_err(|error| Miss::Read(map_store_error(error)))?; + let spend_rows = spend(store, access, &rows).await; + let freshness = Freshness::of(&rows, spend_rows.is_some(), snapshot_ms); + resolve_trace( + trace_id, + trace_ref, + &rows, + spend_rows.as_deref().unwrap_or_default(), + ) + .map(|trace| (trace, freshness)) + .ok_or(Miss::Absent) + }) + .await + } + + pub async fn get_span( + &self, + store: &S, + access: &ReadAccessParams, + trace_id: &str, + span_id: &str, + trace_ref: &str, + ) -> Result, ReadError> { + let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else { + return Ok(None); + }; + let params = SpanDetailParams { + access: access.clone(), + trace_id: trace_id.to_owned(), + trace_ref, + span_id: span_id.to_owned(), + }; + let row = store.span_detail(¶ms).await.map_err(map_store_error)?; + Ok(row.map(|row| SpanDetail { + input_ui: to_ui_content(&row.input), + output_ui: to_ui_content(&row.output), + span_id: row.span_id, + input: row.input, + output: row.output, + attributes: row.attributes, + })) + } + + pub async fn get_span_error( + &self, + store: &S, + access: &ReadAccessParams, + trace_id: &str, + span_id: &str, + trace_ref: &str, + cursor: Option<&str>, + ) -> Result, ReadError> { + let position = error_position(cursor)?; + let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else { + return Ok(None); + }; + let offset = position.as_ref().map_or(0, |position| position.offset); + let params = SpanErrorParams { + access: access.clone(), + trace_id: trace_id.to_owned(), + trace_ref, + span_id: span_id.to_owned(), + error_offset: offset, + error_version: position + .map(|position| position.version) + .unwrap_or_default(), + }; + let Some(row) = store.span_error(¶ms).await.map_err(map_store_error)? else { + return Ok(None); + }; + let next_offset = offset + row.message.chars().count() as u64; + let next_cursor = (next_offset < row.total_chars).then(|| { + encode_cursor(&ErrorPosition { + offset: next_offset, + version: row.version, + }) + }); + Ok(Some(SpanErrorPage { + span_id: row.span_id, + message: row.message, + total_chars: row.total_chars, + next_cursor, + })) + } +} + +async fn reference( + store: &S, + access: &ReadAccessParams, + trace_id: &str, + trace_ref: &str, +) -> Result, ReadError> { + if !trace_ref.is_empty() { + return Ok(Some(trace_ref.to_owned())); + } + let params = TraceIdentityParams { + access: access.clone(), + trace_id: trace_id.to_owned(), + }; + let identities = store.trace_refs(¶ms).await.map_err(map_store_error)?; + if identities.len() > 1 { + return Err(ReadError::AmbiguousTrace); + } + Ok(identities.into_iter().next()) +} + +fn page( + snapshot: &Snapshot, + position: &SpanPosition, + page_size: u32, + response_bytes: usize, +) -> Result> { + let spans = &snapshot.trace().spans; + let create_page = |count: usize| { + let end = position.offset.saturating_add(count).min(spans.len()); + Trace { + summary: snapshot.trace().summary.clone(), + agents: snapshot.trace().agents.clone(), + spans: spans[position.offset..end].to_vec(), + next_cursor: (end < spans.len()).then(|| { + encode_cursor(&SpanPosition { + trace_ref: position.trace_ref.clone(), + snapshot_ms: position.snapshot_ms, + offset: end, + version: snapshot.version().to_owned(), + }) + }), + } + }; + let mut trace = create_page(page_size as usize); + loop { + if serde_json::to_vec(&trace) + .map_err(|error| ReadError::Encode(Arc::new(error)))? + .len() + <= response_bytes + { + return Ok(trace); + } + if trace.spans.len() <= 1 { + return Err(ReadError::TooLarge); + } + trace = create_page(trace.spans.len() / 2); + } +} + +pub(super) fn now_ms() -> u64 { + (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as u64 +} + +pub(super) fn map_store_error(error: StoreError) -> ReadError { + match error { + StoreError::TooLarge => ReadError::TooLarge, + StoreError::Failed(error) => ReadError::Store(Arc::new(error)), + } +} diff --git a/litellm-rust/crates/traces-cache/src/spend.rs b/litellm-rust/crates/traces-cache/src/spend.rs new file mode 100644 index 00000000000..1cfcde841d0 --- /dev/null +++ b/litellm-rust/crates/traces-cache/src/spend.rs @@ -0,0 +1,67 @@ +use std::ops::Range; + +use crate::TraceStore; +use litellm_traces::{ + SpendLookup, + query::named::{ + ReadAccessParams, SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow, + }, +}; + +const NANOS_PER_MS: i64 = 1_000_000; +const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000; + +pub(super) fn spend_window(rows: &[TraceSpansRow]) -> Option> { + let start_ns = rows.iter().map(|row| row.start_ns).min()?; + let end_ns = rows + .iter() + .map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns)) + .max()?; + Some( + start_ns.div_euclid(NANOS_PER_MS) - SPEND_WINDOW_MS + ..end_ns.div_euclid(NANOS_PER_MS) + SPEND_WINDOW_MS, + ) +} + +pub(super) fn spend_within( + spend: &[SpendByResponseIdsRow], + window: Range, +) -> &[SpendByResponseIdsRow] { + let first = spend.partition_point(|row| row.start_ms < window.start); + let end = spend.partition_point(|row| row.start_ms < window.end); + &spend[first..end.max(first)] +} + +/// Spend rows sorted by `start_ms`, or `None` when the lookup failed and spend is unknown. +pub(super) async fn spend( + store: &S, + access: &ReadAccessParams, + rows: &[TraceSpansRow], +) -> Option> { + let lookup = SpendLookup::new(rows); + let Some(window) = spend_window(rows) else { + return Some(Vec::new()); + }; + if lookup.is_empty() { + return Some(Vec::new()); + } + let params = SpendByResponseIdsParams { + access: access.clone(), + response_ids: lookup.response_ids, + request_ids: lookup.request_ids, + trace_ids: lookup.trace_ids, + start_ms: window.start, + end_ms: window.end, + }; + match store.spend(¶ms).await { + Ok(rows) => { + let mut rows = rows; + rows.sort_by_key(|row| row.start_ms); + Some(rows) + } + Err(error) => { + tracing::warn!(%error, "trace spend lookup unavailable"); + None + } + } +} diff --git a/litellm-rust/crates/traces-cache/src/store.rs b/litellm-rust/crates/traces-cache/src/store.rs new file mode 100644 index 00000000000..3f2a536d4cf --- /dev/null +++ b/litellm-rust/crates/traces-cache/src/store.rs @@ -0,0 +1,62 @@ +use std::future::Future; + +use litellm_traces::query::named::{ + ListTracesParams, ListTracesRow, SpanDetailParams, SpanDetailRow, SpanErrorParams, + SpanErrorRow, SpendByResponseIdsParams, SpendByResponseIdsRow, TraceIdentityParams, + TracePageSpansParams, TraceSpansParams, TraceSpansRow, +}; + +#[derive(Debug, thiserror::Error)] +pub enum StoreError { + #[error("trace read exceeds the storage read budget")] + TooLarge, + #[error(transparent)] + Failed(E), +} + +pub trait TraceStore: Sync { + type Error: std::error::Error + Send + Sync + 'static; + + /// Identifies the backing storage for snapshot cache keys. + fn source(&self) -> &str; + + fn trace_refs( + &self, + params: &TraceIdentityParams, + ) -> impl Future, StoreError>> + Send; + + /// Returns `TooLarge` when the response exceeds the storage limit so the reader can halve `limit`. + fn list_runs( + &self, + params: &ListTracesParams, + ) -> impl Future, StoreError>> + Send; + + /// Returns spans visible at `snapshot_ms`, sorted by `start_ns`, or `TooLarge` past `MAX_GRAPH_BYTES`/`MAX_GRAPH_SPANS`. + fn trace_spans( + &self, + params: &TraceSpansParams, + snapshot_ms: u64, + ) -> impl Future, StoreError>> + Send; + + /// Returns spans visible at `snapshot_ms`, sorted by `start_ns`, or `TooLarge` past `MAX_GRAPH_BYTES`/`MAX_GRAPH_SPANS`. + fn run_spans( + &self, + params: &TracePageSpansParams, + snapshot_ms: u64, + ) -> impl Future, StoreError>> + Send; + + fn spend( + &self, + params: &SpendByResponseIdsParams, + ) -> impl Future, StoreError>> + Send; + + fn span_detail( + &self, + params: &SpanDetailParams, + ) -> impl Future, StoreError>> + Send; + + fn span_error( + &self, + params: &SpanErrorParams, + ) -> impl Future, StoreError>> + Send; +} diff --git a/litellm-rust/crates/traces-cache/tests/read.rs b/litellm-rust/crates/traces-cache/tests/read.rs new file mode 100644 index 00000000000..b5d2442582c --- /dev/null +++ b/litellm-rust/crates/traces-cache/tests/read.rs @@ -0,0 +1,687 @@ +use std::{ + collections::{HashMap, HashSet}, + sync::{ + Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use litellm_traces::{ + CallEvidenceKind, CallKey, ObservationType, SpanStatus, + query::named::{ + ListTracesParams, ListTracesRow, ReadAccessParams, SpanDetailParams, SpanDetailRow, + SpanErrorParams, SpanErrorRow, SpendByResponseIdsParams, SpendByResponseIdsRow, + TraceIdentityParams, TracePageSpansParams, TraceSpansParams, TraceSpansRow, + }, +}; +use litellm_traces_cache::{LIVE_TTL, ReadError, StoreError, TraceReader, TraceStore}; +use rstest::rstest; + +const START_NS: i64 = 1_790_742_989_000_000_000; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +enum Operation { + TraceRefs, + ListRuns, + TraceSpans, + RunSpans, + Spend, + SpanDetail, + SpanError, +} + +#[derive(Clone, Copy)] +enum Failure { + TooLarge, + Failed, +} + +#[derive(Debug, thiserror::Error)] +#[error("fake trace store failed")] +struct FakeError; + +#[derive(Default)] +struct State { + failures: HashMap, + trace_refs: Vec, + list_runs: Vec, + trace_spans: HashMap>, + run_spans: Vec, + spend: Vec, + span_detail: Option, + span_error: Option, + list_runs_too_large_above: Option, + trace_too_large_refs: HashSet, +} + +#[derive(Default)] +struct Calls { + trace_refs: AtomicUsize, + list_runs: AtomicUsize, + trace_spans: AtomicUsize, + run_spans: AtomicUsize, + spend: AtomicUsize, + span_detail: AtomicUsize, + span_error: AtomicUsize, +} + +#[derive(Default)] +struct FakeStore { + state: Mutex, + calls: Calls, +} + +impl FakeStore { + fn with_spans(trace_ref: &str, spans: Vec) -> Self { + Self { + state: Mutex::new(State { + trace_spans: HashMap::from([(trace_ref.to_owned(), spans)]), + ..State::default() + }), + calls: Calls::default(), + } + } + + fn set_failure(&self, operation: Operation, failure: Failure) { + self.state + .lock() + .unwrap() + .failures + .insert(operation, failure); + } + + fn set_trace_refs(&self, trace_refs: Vec) { + self.state.lock().unwrap().trace_refs = trace_refs; + } + + fn set_list_runs(&self, rows: Vec) { + self.state.lock().unwrap().list_runs = rows; + } + + fn set_list_runs_too_large_above(&self, limit: u32) { + self.state.lock().unwrap().list_runs_too_large_above = Some(limit); + } + + fn set_run_spans(&self, rows: Vec) { + self.state.lock().unwrap().run_spans = rows; + } + + fn set_trace_spans_too_large(&self, trace_ref: &str) { + self.state + .lock() + .unwrap() + .trace_too_large_refs + .insert(trace_ref.to_owned()); + } + + fn calls(&self, operation: Operation) -> usize { + match operation { + Operation::TraceRefs => self.calls.trace_refs.load(Ordering::SeqCst), + Operation::ListRuns => self.calls.list_runs.load(Ordering::SeqCst), + Operation::TraceSpans => self.calls.trace_spans.load(Ordering::SeqCst), + Operation::RunSpans => self.calls.run_spans.load(Ordering::SeqCst), + Operation::Spend => self.calls.spend.load(Ordering::SeqCst), + Operation::SpanDetail => self.calls.span_detail.load(Ordering::SeqCst), + Operation::SpanError => self.calls.span_error.load(Ordering::SeqCst), + } + } + + fn failure(state: &State, operation: Operation) -> Result<(), StoreError> { + match state.failures.get(&operation) { + Some(Failure::TooLarge) => Err(StoreError::TooLarge), + Some(Failure::Failed) => Err(StoreError::Failed(FakeError)), + None => Ok(()), + } + } +} + +impl TraceStore for FakeStore { + type Error = FakeError; + + fn source(&self) -> &str { + "fake" + } + + async fn trace_refs( + &self, + _: &TraceIdentityParams, + ) -> Result, StoreError> { + self.calls.trace_refs.fetch_add(1, Ordering::SeqCst); + let state = self.state.lock().unwrap(); + Self::failure(&state, Operation::TraceRefs)?; + Ok(state.trace_refs.clone()) + } + + async fn list_runs( + &self, + params: &ListTracesParams, + ) -> Result, StoreError> { + self.calls.list_runs.fetch_add(1, Ordering::SeqCst); + let state = self.state.lock().unwrap(); + Self::failure(&state, Operation::ListRuns)?; + if state + .list_runs_too_large_above + .is_some_and(|limit| params.limit > limit) + { + return Err(StoreError::TooLarge); + } + Ok(state + .list_runs + .iter() + .take(params.limit as usize) + .cloned() + .collect()) + } + + async fn trace_spans( + &self, + params: &TraceSpansParams, + _: u64, + ) -> Result, StoreError> { + self.calls.trace_spans.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + let state = self.state.lock().unwrap(); + Self::failure(&state, Operation::TraceSpans)?; + if state.trace_too_large_refs.contains(¶ms.trace_ref) { + return Err(StoreError::TooLarge); + } + Ok(state + .trace_spans + .get(¶ms.trace_ref) + .cloned() + .unwrap_or_default()) + } + + async fn run_spans( + &self, + _: &TracePageSpansParams, + _: u64, + ) -> Result, StoreError> { + self.calls.run_spans.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + let state = self.state.lock().unwrap(); + Self::failure(&state, Operation::RunSpans)?; + Ok(state.run_spans.clone()) + } + + async fn spend( + &self, + _: &SpendByResponseIdsParams, + ) -> Result, StoreError> { + self.calls.spend.fetch_add(1, Ordering::SeqCst); + let state = self.state.lock().unwrap(); + Self::failure(&state, Operation::Spend)?; + Ok(state.spend.clone()) + } + + async fn span_detail( + &self, + _: &SpanDetailParams, + ) -> Result, StoreError> { + self.calls.span_detail.fetch_add(1, Ordering::SeqCst); + let state = self.state.lock().unwrap(); + Self::failure(&state, Operation::SpanDetail)?; + Ok(state.span_detail.clone()) + } + + async fn span_error( + &self, + _: &SpanErrorParams, + ) -> Result, StoreError> { + self.calls.span_error.fetch_add(1, Ordering::SeqCst); + let state = self.state.lock().unwrap(); + Self::failure(&state, Operation::SpanError)?; + Ok(state.span_error.clone()) + } +} + +fn access() -> ReadAccessParams { + ReadAccessParams { + all_teams: true, + user_id: String::new(), + team_ids: Vec::new(), + } +} + +fn span(index: usize) -> TraceSpansRow { + TraceSpansRow { + trace_id: "trace".into(), + span_id: format!("span-{index}"), + parent_span_id: if index == 0 { + String::new() + } else { + "span-0".into() + }, + name: "agent".into(), + kind: ObservationType::Agent, + wrapper_candidate: false, + agent: "agent".into(), + framework: String::new(), + status: SpanStatus::Ok, + status_message: String::new(), + error_truncated: false, + start_ns: START_NS + index as i64 * 1_000_000, + duration_ns: 10_000_000, + service: "test".into(), + input_preview: format!("span input {index}"), + model: String::new(), + input_tokens: 0, + output_tokens: 0, + litellm_request_id: String::new(), + call_keys: Vec::new(), + call_evidence: None, + tool_call_id: String::new(), + team_id: "team".into(), + api_key_hash: "key".into(), + user_id: "user".into(), + } +} + +fn run(trace_id: &str, trace_ref: &str) -> ListTracesRow { + ListTracesRow { + trace_id: trace_id.into(), + trace_ref: trace_ref.into(), + team_id: "team".into(), + api_key_hash: "key".into(), + user_id: "user".into(), + name: "listed".into(), + service: "test".into(), + input_preview: String::new(), + status: SpanStatus::Ok, + start_ms: 1_790_742_989_000, + duration_ms: 10, + span_count: 1, + agent_count: 1, + agent_invocations: 1, + agent_names: vec!["agent".into()], + frameworks: Vec::new(), + llm_calls: 0, + tool_calls: 0, + input_tokens: 0, + output_tokens: 0, + models: Vec::new(), + error_count: 0, + request_ids: Vec::new(), + } +} + +#[rstest] +#[tokio::test] +async fn pages_reuse_one_trace_snapshot_and_concatenate_in_order() { + let store = FakeStore::with_spans("ref", (0..5).map(span).collect()); + let reader = TraceReader::new(usize::MAX); + let access = access(); + let first = reader + .get_trace_page(&store, &access, "trace", "ref", None, 2) + .await + .unwrap() + .unwrap(); + let second = reader + .get_trace_page( + &store, + &access, + "trace", + "ref", + first.next_cursor.as_deref(), + 2, + ) + .await + .unwrap() + .unwrap(); + let third = reader + .get_trace_page( + &store, + &access, + "trace", + "ref", + second.next_cursor.as_deref(), + 2, + ) + .await + .unwrap() + .unwrap(); + let ids: Vec<_> = first + .spans + .iter() + .chain(&second.spans) + .chain(&third.spans) + .map(|span| span.span_id.as_str()) + .collect(); + assert_eq!(ids, ["span-0", "span-1", "span-2", "span-3", "span-4"]); + assert!(third.next_cursor.is_none()); + assert_eq!(store.calls(Operation::TraceSpans), 1); +} + +#[rstest] +#[tokio::test] +async fn snapshot_versions_are_stable_across_readers_and_detect_changes() { + let access = access(); + let original = FakeStore::with_spans("ref", vec![span(0), span(1)]); + let reader_a = TraceReader::new(usize::MAX); + let first = reader_a + .get_trace_page(&original, &access, "trace", "ref", None, 1) + .await + .unwrap() + .unwrap(); + let cursor = first.next_cursor.unwrap(); + + let changed = FakeStore::with_spans("ref", vec![span(0), span(1), span(2)]); + let reader_b = TraceReader::new(usize::MAX); + let result = reader_b + .get_trace_page(&changed, &access, "trace", "ref", Some(&cursor), 1) + .await; + assert!(matches!(result, Err(ReadError::TraceChanged))); + + let unchanged = FakeStore::with_spans("ref", vec![span(0), span(1)]); + let reader_c = TraceReader::new(usize::MAX); + let next = reader_c + .get_trace_page(&unchanged, &access, "trace", "ref", Some(&cursor), 1) + .await + .unwrap() + .unwrap(); + assert_eq!(next.spans[0].span_id, "span-1"); +} + +#[rstest] +#[tokio::test] +async fn response_size_splits_pages_and_rejects_a_single_oversized_span() { + let spans: Vec<_> = (0..4) + .map(|index| { + let mut row = span(index); + row.input_preview = "x".repeat(256); + row + }) + .collect(); + let access = access(); + let full_budget_reader = TraceReader::new(usize::MAX); + let one_span = full_budget_reader + .get_trace_page( + &FakeStore::with_spans("ref", spans.clone()), + &access, + "trace", + "ref", + None, + 1, + ) + .await + .unwrap() + .unwrap(); + let response_bytes = serde_json::to_vec(&one_span).unwrap().len() + 128; + let reader = TraceReader::new(response_bytes); + let store = FakeStore::with_spans("ref", spans.clone()); + let page = reader + .get_trace_page(&store, &access, "trace", "ref", None, 4) + .await + .unwrap() + .unwrap(); + assert!(!page.spans.is_empty()); + assert!(page.spans.len() < 4); + assert!(page.next_cursor.is_some()); + let continued = reader + .get_trace_page( + &store, + &access, + "trace", + "ref", + page.next_cursor.as_deref(), + 4, + ) + .await + .unwrap() + .unwrap(); + assert!(!continued.spans.is_empty()); + assert!(matches!( + TraceReader::new(1) + .get_trace_page(&store, &access, "trace", "ref", None, 1) + .await, + Err(ReadError::TooLarge) + )); +} + +#[rstest] +#[tokio::test] +async fn list_run_budget_halves_the_limit_and_cursor_requires_a_full_page() { + let store = FakeStore::default(); + store.set_list_runs( + (0..3) + .map(|index| run(&format!("trace-{index}"), &format!("ref-{index}"))) + .collect(), + ); + store.set_list_runs_too_large_above(2); + let reader = TraceReader::new(usize::MAX); + let access = access(); + let page = reader + .list_traces(&store, &access, 0, i64::MAX, None, 8) + .await + .unwrap(); + assert_eq!(page.data.len(), 2); + assert!(page.next_cursor.is_some()); + assert_eq!(store.calls(Operation::ListRuns), 3); + + let shorter = FakeStore::default(); + shorter.set_list_runs(vec![run("only", "ref-only")]); + shorter.set_list_runs_too_large_above(2); + let page = reader + .list_traces(&shorter, &access, 0, i64::MAX, None, 8) + .await + .unwrap(); + assert_eq!(page.data.len(), 1); + assert!(page.next_cursor.is_none()); + assert_eq!(shorter.calls(Operation::ListRuns), 1); +} + +#[rstest] +#[tokio::test] +async fn oversized_run_batch_falls_back_to_each_run_and_keeps_listed_summaries() { + let store = FakeStore::with_spans("ref-good", vec![span(0)]); + store.set_list_runs(vec![ + run("trace-large", "ref-large"), + run("trace-good", "ref-good"), + ]); + store.set_trace_spans_too_large("ref-large"); + store.set_run_spans(Vec::new()); + store.set_failure(Operation::RunSpans, Failure::TooLarge); + let reader = TraceReader::new(usize::MAX); + let page = reader + .list_traces(&store, &access(), 0, i64::MAX, None, 2) + .await + .unwrap(); + assert_eq!(page.data.len(), 2); + assert!(page.data[0].resolution_limited); + assert_eq!(page.data[0].trace_ref, "ref-large"); + assert!(!page.data[1].resolution_limited); + assert_eq!(page.data[1].trace_ref, "ref-good"); + assert_eq!(store.calls(Operation::RunSpans), 1); + assert_eq!(store.calls(Operation::TraceSpans), 2); + + let again = reader + .list_traces(&store, &access(), 0, i64::MAX, None, 2) + .await + .unwrap(); + assert_eq!(again.data, page.data); + assert_eq!(store.calls(Operation::RunSpans), 1); + assert_eq!(store.calls(Operation::TraceSpans), 2); +} + +#[rstest] +#[tokio::test] +async fn failed_spend_lookup_preserves_the_trace_with_unknown_spend() { + let mut row = span(0); + row.litellm_request_id = "response".into(); + row.call_keys = vec![CallKey::ProviderResponse("response".into())]; + row.call_evidence = Some(CallEvidenceKind::Complete); + let store = FakeStore::with_spans("ref", vec![row]); + store.set_failure(Operation::Spend, Failure::Failed); + let trace = TraceReader::new(usize::MAX) + .get_trace(&store, &access(), "trace", "ref") + .await + .unwrap() + .unwrap(); + assert_eq!(trace.summary.spend, None); + assert_eq!(trace.spans[0].spend, None); + assert_eq!(store.calls(Operation::Spend), 1); +} + +#[rstest] +#[tokio::test] +async fn ambiguous_trace_references_fail_and_a_single_reference_is_resolved() { + let reader = TraceReader::new(usize::MAX); + let access = access(); + let ambiguous = FakeStore::default(); + ambiguous.set_trace_refs(vec!["ref-a".into(), "ref-b".into()]); + assert!(matches!( + reader.get_trace(&ambiguous, &access, "trace", "").await, + Err(ReadError::AmbiguousTrace) + )); + + let unique = FakeStore::with_spans("ref-only", vec![span(0)]); + unique.set_trace_refs(vec!["ref-only".into()]); + let trace = reader + .get_trace(&unique, &access, "trace", "") + .await + .unwrap() + .unwrap(); + assert_eq!(trace.summary.trace_ref, "ref-only"); +} + +#[rstest] +#[case::zero(0)] +#[case::above_max(501)] +#[tokio::test] +async fn invalid_page_sizes_are_rejected(#[case] page_size: u32) { + let store = FakeStore::default(); + let reader = TraceReader::new(usize::MAX); + let access = access(); + assert!(matches!( + reader + .get_trace_page(&store, &access, "trace", "ref", None, page_size) + .await, + Err(ReadError::InvalidParameters) + )); +} + +#[rstest] +#[tokio::test] +async fn zero_list_limit_is_rejected() { + let store = FakeStore::default(); + let reader = TraceReader::new(usize::MAX); + let access = access(); + assert!(matches!( + reader + .list_traces(&store, &access, 0, i64::MAX, None, 0) + .await, + Err(ReadError::InvalidParameters) + )); +} + +fn now_ns() -> i64 { + time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64 +} + +#[rstest] +#[tokio::test] +async fn concurrent_and_repeated_opens_share_one_storage_read() { + let store = FakeStore::with_spans("ref", (0..3).map(span).collect()); + let reader = TraceReader::new(usize::MAX); + let access = access(); + let (first, second) = tokio::join!( + reader.get_trace(&store, &access, "trace", "ref"), + reader.get_trace_page(&store, &access, "trace", "ref", None, 2), + ); + let first = first.unwrap().unwrap(); + let second = second.unwrap().unwrap(); + let reopened = reader + .get_trace_page(&store, &access, "trace", "ref", None, 2) + .await + .unwrap() + .unwrap(); + assert_eq!(first.spans.len(), 3); + assert_eq!(second.spans, first.spans[..2]); + assert_eq!(reopened.next_cursor, second.next_cursor); + assert_eq!(store.calls(Operation::TraceSpans), 1); +} + +#[rstest] +#[tokio::test] +async fn failed_reads_are_not_cached() { + let store = FakeStore::with_spans("ref", vec![span(0)]); + store.set_failure(Operation::TraceSpans, Failure::Failed); + let reader = TraceReader::new(usize::MAX); + let access = access(); + assert!(matches!( + reader.get_trace(&store, &access, "trace", "ref").await, + Err(ReadError::Store(_)) + )); + store.state.lock().unwrap().failures.clear(); + let trace = reader + .get_trace(&store, &access, "trace", "ref") + .await + .unwrap() + .unwrap(); + assert_eq!(trace.spans.len(), 1); + assert_eq!(store.calls(Operation::TraceSpans), 2); +} + +#[rstest] +#[tokio::test] +async fn listed_runs_are_read_once_until_a_live_run_expires() { + let live = TraceSpansRow { + trace_id: "trace-live".into(), + start_ns: now_ns(), + ..span(0) + }; + let settled = TraceSpansRow { + trace_id: "trace-settled".into(), + ..span(0) + }; + let store = FakeStore::default(); + store.set_list_runs(vec![ + run("trace-live", "ref-live"), + run("trace-settled", "ref-settled"), + ]); + store.set_run_spans(vec![live, settled]); + let reader = TraceReader::new(usize::MAX); + let access = access(); + let list = || reader.list_traces(&store, &access, 0, i64::MAX, None, 2); + + let first = list().await.unwrap(); + assert!(first.data.iter().all(|summary| summary.name == "agent")); + list().await.unwrap(); + assert_eq!(store.calls(Operation::RunSpans), 1); + + store.set_run_spans(Vec::new()); + tokio::time::sleep(LIVE_TTL + Duration::from_millis(200)).await; + let after = list().await.unwrap(); + assert_eq!(store.calls(Operation::RunSpans), 2); + assert_eq!(after.data[0].name, "listed"); + assert_eq!(after.data[1], first.data[1]); +} + +#[rstest] +#[tokio::test] +async fn concurrent_pages_of_an_evicted_snapshot_share_one_storage_read() { + let access = access(); + let first = TraceReader::new(usize::MAX) + .get_trace_page( + &FakeStore::with_spans("ref", (0..3).map(span).collect()), + &access, + "trace", + "ref", + None, + 1, + ) + .await + .unwrap() + .unwrap(); + let cursor = first.next_cursor.as_deref(); + let store = FakeStore::with_spans("ref", (0..3).map(span).collect()); + let reader = TraceReader::new(usize::MAX); + let (left, right) = tokio::join!( + reader.get_trace_page(&store, &access, "trace", "ref", cursor, 1), + reader.get_trace_page(&store, &access, "trace", "ref", cursor, 1), + ); + assert_eq!(left.unwrap().unwrap().spans[0].span_id, "span-1"); + assert_eq!(right.unwrap().unwrap().spans[0].span_id, "span-1"); + assert_eq!(store.calls(Operation::TraceSpans), 1); +} diff --git a/litellm-rust/crates/traces-cache/tests/snapshots.rs b/litellm-rust/crates/traces-cache/tests/snapshots.rs index de199ae0301..42085d3e95b 100644 --- a/litellm-rust/crates/traces-cache/tests/snapshots.rs +++ b/litellm-rust/crates/traces-cache/tests/snapshots.rs @@ -5,7 +5,9 @@ use litellm_traces::{ query::named::{ReadAccessParams, SpendByResponseIdsRow, TraceSpansRow}, resolve_trace, }; -use litellm_traces_cache::{Error, SnapshotCache, SnapshotKey}; +use std::sync::Arc; + +use litellm_traces_cache::{Error, Freshness, Snapshot, SnapshotCache, SnapshotKey}; use rstest::{fixture, rstest}; const T0: i64 = 1_790_742_989_000_000_000; @@ -60,6 +62,18 @@ fn key( SnapshotKey::new(source, access, trace_id, trace_ref, ms).unwrap() } +async fn insert( + cache: &SnapshotCache, + key: SnapshotKey, + trace: Trace, +) -> Result, Arc> { + cache + .pinned_or_load(key, 100, async { + Ok::<_, Error>((trace, Freshness::Settled)) + }) + .await +} + #[fixture] fn trace() -> Trace { resolve_trace( @@ -85,7 +99,7 @@ async fn cached_trace_is_isolated_by_access_scope( let cache = SnapshotCache::new(1024 * 1024, TTL); let stored = key("source", &access(), "trace", "ref", 100); - cache.insert(stored.clone(), trace.clone()).await.unwrap(); + insert(&cache, stored.clone(), trace.clone()).await.unwrap(); let other_access = ReadAccessParams { all_teams, @@ -115,7 +129,7 @@ async fn cached_trace_is_isolated_by_key_fields( let cache = SnapshotCache::new(1024 * 1024, TTL); let stored = key("source", &access(), "trace", "ref", 100); - cache.insert(stored.clone(), trace.clone()).await.unwrap(); + insert(&cache, stored.clone(), trace.clone()).await.unwrap(); let other = key(source, &access(), trace_id, trace_ref, snapshot_ms); assert!(cache.get(&other).await.is_none()); @@ -129,7 +143,7 @@ async fn snapshot_at_the_size_limit_is_accepted(trace: Trace) { let cache = SnapshotCache::new(size, TTL); let stored = key("source", &access(), "trace", "ref", 100); - cache.insert(stored.clone(), trace).await.unwrap(); + insert(&cache, stored.clone(), trace).await.unwrap(); assert!(cache.get(&stored).await.is_some()); } @@ -141,8 +155,8 @@ async fn snapshot_one_byte_over_the_size_limit_is_rejected(trace: Trace) { let stored = key("source", &access(), "trace", "ref", 100); assert!(matches!( - cache.insert(stored.clone(), trace).await, - Err(Error::ReadTooLarge) + insert(&cache, stored.clone(), trace).await, + Err(error) if matches!(*error, Error::ReadTooLarge) )); assert!(cache.get(&stored).await.is_none()); } @@ -166,25 +180,31 @@ async fn snapshot_version_tracks_the_ordered_span_ids( }; let cache = SnapshotCache::new(1024 * 1024, TTL); - let first = cache - .insert(key("source", &access(), "a", "ref", 100), build(first_ids)) - .await - .unwrap(); - let second = cache - .insert(key("source", &access(), "b", "ref", 100), build(second_ids)) - .await - .unwrap(); + let first = insert( + &cache, + key("source", &access(), "a", "ref", 100), + build(first_ids), + ) + .await + .unwrap(); + let second = insert( + &cache, + key("source", &access(), "b", "ref", 100), + build(second_ids), + ) + .await + .unwrap(); assert_eq!(first.version() == second.version(), equal); } #[rstest] #[tokio::test] -async fn snapshots_expire_after_the_ttl(trace: Trace) { +async fn snapshots_expire_when_idle(trace: Trace) { let cache = SnapshotCache::new(1024 * 1024, Duration::from_millis(50)); let stored = key("source", &access(), "trace", "ref", 100); - cache.insert(stored.clone(), trace).await.unwrap(); + insert(&cache, stored.clone(), trace).await.unwrap(); tokio::time::sleep(Duration::from_millis(200)).await; assert!(cache.get(&stored).await.is_none()); diff --git a/litellm-rust/crates/traces-clickhouse/Cargo.toml b/litellm-rust/crates/traces-clickhouse/Cargo.toml index 6b95eb149d3..b90ad8a7bf8 100644 --- a/litellm-rust/crates/traces-clickhouse/Cargo.toml +++ b/litellm-rust/crates/traces-clickhouse/Cargo.toml @@ -12,11 +12,9 @@ schema = ["dep:schemars", "litellm-traces/schema"] macro_rules_attribute.workspace = true schemars = { workspace = true, optional = true } askama.workspace = true -base64.workspace = true flate2.workspace = true futures-util.workspace = true hmac = "0.12.1" -itertools = "0.14.0" litellm-http.workspace = true litellm-migrate.workspace = true litellm-storage-clickhouse.workspace = true @@ -30,7 +28,6 @@ strum.workspace = true thiserror.workspace = true time = { workspace = true, features = ["formatting"] } tokio.workspace = true -tracing.workspace = true url.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/traces-clickhouse/src/error.rs b/litellm-rust/crates/traces-clickhouse/src/error.rs index 78df9a2121c..1d15c556316 100644 --- a/litellm-rust/crates/traces-clickhouse/src/error.rs +++ b/litellm-rust/crates/traces-clickhouse/src/error.rs @@ -16,8 +16,6 @@ pub enum Error { InvalidResponse, #[error("ClickHouse insert exceeds the encoded size limit")] InsertTooLarge, - #[error("Trace exceeds the interactive read budget; use a filtered trace query")] - ReadTooLarge, #[error("ClickHouse schema setup failed with HTTP status {0}")] SchemaFailed(u16), #[error("ClickHouse schema setup transport failed")] @@ -34,12 +32,6 @@ pub enum Error { ProvisionFailed(u16), #[error("ClickHouse reader provisioning transport failed")] ProvisionTransport, - #[error("Invalid {0} cursor")] - InvalidCursor(&'static str), - #[error("Multiple traces have this ID; provide trace_ref")] - AmbiguousTrace, - #[error("Trace changed while paging; refresh the trace to continue")] - TraceChanged, #[error(transparent)] Decode(#[from] litellm_traces::Error), #[error("trace ingestion task failed")] @@ -49,12 +41,3 @@ pub enum Error { #[error(transparent)] Cached(#[from] std::sync::Arc), } - -impl From for Error { - fn from(error: litellm_traces_cache::Error) -> Self { - match error { - litellm_traces_cache::Error::Serialization(_) => Self::InvalidResponse, - litellm_traces_cache::Error::ReadTooLarge => Self::ReadTooLarge, - } - } -} diff --git a/litellm-rust/crates/traces-clickhouse/src/lib.rs b/litellm-rust/crates/traces-clickhouse/src/lib.rs index d83708d27f1..43ff0b8bf33 100644 --- a/litellm-rust/crates/traces-clickhouse/src/lib.rs +++ b/litellm-rust/crates/traces-clickhouse/src/lib.rs @@ -31,7 +31,7 @@ pub use litellm_storage_clickhouse::{Connection, Parameter}; pub use litellm_traces::{QueryScope, ReadQuery}; pub use query::{QueryHelp, execute_read, query_help, query_sql}; pub use query_access::QueryReaders; -pub use reads::{get_span, get_span_error, get_trace, get_trace_page, list_traces}; +pub use reads::ClickHouseTraces; pub use schema::{ NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements, }; diff --git a/litellm-rust/crates/traces-clickhouse/src/reads.rs b/litellm-rust/crates/traces-clickhouse/src/reads.rs index 6938aab06b3..7a5005b8cb0 100644 --- a/litellm-rust/crates/traces-clickhouse/src/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/src/reads.rs @@ -1,26 +1,14 @@ -//! Scoped trace reads: the trace list, one trace resolved with its spend, and span payloads. - -use std::sync::LazyLock; -use std::time::Duration; - -use base64::{Engine, engine::general_purpose::URL_SAFE}; -use futures_util::{StreamExt, TryStreamExt, stream}; -use itertools::Itertools; use litellm_http::Client; +use litellm_storage_clickhouse::Error as StorageError; use litellm_storage_clickhouse::{Query, fetch}; -use litellm_traces::{ - SpanDetail, SpanErrorPage, SpendLookup, Trace, TracePage, listed_summary, - query::named as contracts, resolve_trace, to_ui_content, -}; -use litellm_traces_cache::{SnapshotCache, SnapshotKey}; -use serde::{Deserialize, Serialize}; +use litellm_traces::query::named as contracts; +use litellm_traces_cache::{StoreError, TraceStore}; use crate::{ Connection, Error, query::named::{ - ListTracesParams, ListTracesRow, ReadAccessParams, SpanDetail as SpanDetailQuery, - SpanDetailParams, SpanError, SpanErrorParams, SpendByResponseIdsParams, TraceIdentity, - TraceIdentityParams, TracePageSpansParams, TraceSpansParams, + ListTracesParams, ListTracesRow, SpanDetail as SpanDetailQuery, SpanError, SpanErrorParams, + SpendByResponseIdsParams, TraceIdentity, TracePageSpansParams, }, }; @@ -36,501 +24,103 @@ impl Query for RunCandidates { ); } -// Cursor pages share a bounded snapshot so advancing does not resolve the whole graph again. -static TRACE_SNAPSHOTS: LazyLock = LazyLock::new(|| { - SnapshotCache::new( - crate::span_batches::MAX_GRAPH_BYTES, - Duration::from_secs(120), - ) -}); - -const NANOS_PER_MS: i64 = 1_000_000; -const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000; -const SPEND_CONCURRENCY: usize = 4; - -fn encode_cursor(position: &T) -> String { - URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default()) +pub struct ClickHouseTraces { + client: Client, + connection: Connection, } -fn decode_cursor Deserialize<'de>>( - cursor: &str, - kind: &'static str, -) -> Result { - URL_SAFE - .decode(cursor) - .ok() - .and_then(|json| serde_json::from_slice(&json).ok()) - .ok_or(Error::InvalidCursor(kind)) -} - -fn trace_position(cursor: Option<&str>) -> Result<(i64, String), Error> { - let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else { - return Ok((0, String::new())); - }; - match decode_cursor::<(i64, String)>(cursor, "trace")? { - (start_ms, trace_ref) if start_ms > 0 && !trace_ref.is_empty() => Ok((start_ms, trace_ref)), - _ => Err(Error::InvalidCursor("trace")), +impl ClickHouseTraces { + pub fn new(client: Client, connection: Connection) -> Self { + Self { client, connection } } } -#[derive(Deserialize, Serialize)] -struct ErrorPosition { - offset: u64, - version: String, -} +impl TraceStore for ClickHouseTraces { + type Error = Error; -fn error_position(cursor: Option<&str>) -> Result, Error> { - let Some(cursor) = cursor else { - return Ok(None); - }; - let position = decode_cursor::(cursor, "diagnostic")?; - let valid_version = position.version.len() == 64 - && position - .version - .bytes() - .all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte)); - if i64::try_from(position.offset).is_err() || !valid_version { - return Err(Error::InvalidCursor("diagnostic")); + fn source(&self) -> &str { + self.connection.url().as_str() } - Ok(Some(position)) -} -/// The stored run a trace id names for this caller; ids can repeat across tenants and runs. -async fn reference( - client: &Client, - connection: &Connection, - access: &ReadAccessParams, - trace_id: &str, - trace_ref: &str, -) -> Result, Error> { - if !trace_ref.is_empty() { - return Ok(Some(trace_ref.to_owned())); + async fn trace_refs( + &self, + params: &contracts::TraceIdentityParams, + ) -> Result, StoreError> { + fetch::(&self.client, &self.connection, params) + .await + .map(|rows| rows.into_iter().map(|row| row.trace_ref).collect()) + .map_err(failed) } - let params = TraceIdentityParams { - access: access.clone(), - trace_id: trace_id.to_owned(), - }; - let mut identities = fetch::(client, connection, ¶ms).await?; - if identities.len() > 1 { - return Err(Error::AmbiguousTrace); - } - Ok(identities.pop().map(|identity| identity.trace_ref)) -} -/// Spend records behind the spans' calls. A failed lookup leaves cost unknown instead of failing -/// the read. -async fn spend( - client: &Client, - connection: &Connection, - access: &ReadAccessParams, - rows: &[contracts::TraceSpansRow], -) -> Vec { - let lookup = SpendLookup::new(rows); - let (Some(start_ns), Some(end_ns)) = ( - rows.iter().map(|row| row.start_ns).min(), - rows.iter() - .map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns)) - .max(), - ) else { - return Vec::new(); - }; - if lookup.is_empty() { - return Vec::new(); - } - let params = SpendByResponseIdsParams::from(contracts::SpendByResponseIdsParams { - access: access.clone(), - response_ids: lookup.response_ids, - request_ids: lookup.request_ids, - trace_ids: lookup.trace_ids, - start_ms: start_ns.div_euclid(NANOS_PER_MS) - SPEND_WINDOW_MS, - end_ms: end_ns.div_euclid(NANOS_PER_MS) + SPEND_WINDOW_MS, - }); - match crate::span_batches::read_spend(client, connection, params).await { - Ok(rows) => rows, - Err(error) => { - tracing::warn!(%error, "trace spend lookup unavailable"); - Vec::new() + async fn list_runs( + &self, + params: &contracts::ListTracesParams, + ) -> Result, StoreError> { + let storage_params = ListTracesParams::from(params.clone()); + match fetch::(&self.client, &self.connection, &storage_params).await { + Ok(rows) => Ok(rows.into_iter().map(|row| row.0).collect()), + Err(StorageError::ResponseTooLarge) => Err(StoreError::TooLarge), + Err(error) => Err(failed(error)), } } -} -pub async fn list_traces( - client: &Client, - connection: &Connection, - access: &ReadAccessParams, - start_ms: i64, - end_ms: i64, - cursor: Option<&str>, - limit: u32, -) -> Result { - if limit == 0 { - return Err(Error::InvalidParameters); + async fn trace_spans( + &self, + params: &contracts::TraceSpansParams, + snapshot_ms: u64, + ) -> Result, StoreError> { + crate::span_batches::read_spans(&self.client, &self.connection, params.clone(), snapshot_ms) + .await } - let (cursor_ms, cursor_trace_id) = trace_position(cursor)?; - let mut params = ListTracesParams::from(contracts::ListTracesParams { - access: access.clone(), - start_ms, - end_ms, - cursor_ms, - cursor_trace_id, - limit: limit.min(500), - }); - let page: Vec = loop { - match fetch::(client, connection, ¶ms).await { - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) if params.0.limit > 1 => { - params.0.limit /= 2; - } - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => { - return Err(Error::ReadTooLarge); - } - result => break result?.into_iter().map(|row| row.0).collect(), - } - }; - let next_cursor = page - .last() - .filter(|_| page.len() == params.0.limit as usize) - .map(|last| encode_cursor(&(last.start_ms, &last.trace_ref))); - let data = stream::iter(page.chunks(16)) - .then(|batch| list_summaries(client, connection, access, batch)) - .try_collect::>() - .await? - .into_iter() - .flatten() - .collect(); - Ok(TracePage { data, next_cursor }) -} -async fn list_summaries( - client: &Client, - connection: &Connection, - access: &ReadAccessParams, - runs: &[contracts::ListTracesRow], -) -> Result, Error> { - let (Some(start_ms), Some(end_ms)) = ( - runs.iter().map(|row| row.start_ms).min(), - runs.iter() - .map(|row| row.start_ms.saturating_add(row.duration_ms)) - .max(), - ) else { - return Ok(Vec::new()); - }; - let params = TracePageSpansParams::from(contracts::TracePageSpansParams { - access: access.clone(), - trace_refs: runs.iter().map(|row| row.trace_ref.clone()).collect(), - start_ms, - end_ms: end_ms.saturating_add(1), - }); - let spans = match crate::span_batches::read_list_spans(client, connection, params).await { - Ok(spans) => spans, - Err(Error::ReadTooLarge) => { - return stream::iter(runs) - .then(|row| async move { - match get_trace(client, connection, access, &row.trace_id, &row.trace_ref).await - { - Ok(trace) => { - Ok(trace.map_or_else(|| listed_summary(row), |trace| trace.summary)) - } - Err(Error::ReadTooLarge) => Ok(listed_summary(row)), - Err(error) => Err(error), - } - }) - .try_collect() - .await; - } - Err(error) => return Err(error), - }; - let by_trace = spans.into_iter().into_group_map_by(|span| { - ( - span.team_id.clone(), - span.api_key_hash.clone(), - span.trace_id.clone(), + async fn run_spans( + &self, + params: &contracts::TracePageSpansParams, + snapshot_ms: u64, + ) -> Result, StoreError> { + crate::span_batches::read_list_spans( + &self.client, + &self.connection, + TracePageSpansParams::from(params.clone()), + snapshot_ms, ) - }); - let summaries = runs - .iter() - .map(|row| { - let spans = by_trace - .get(&( - row.team_id.clone(), - row.api_key_hash.clone(), - row.trace_id.clone(), - )) - .map(Vec::as_slice) - .unwrap_or_default(); - async move { - let spend_rows = spend(client, connection, access, spans).await; - resolve_trace(&row.trace_id, &row.trace_ref, spans, &spend_rows) - .map_or_else(|| listed_summary(row), |trace| trace.summary) - } - }) - .collect::>(); - Ok(stream::iter(summaries) - .buffered(SPEND_CONCURRENCY) - .collect() - .await) -} - -pub async fn get_trace( - client: &Client, - connection: &Connection, - access: &ReadAccessParams, - trace_id: &str, - trace_ref: &str, -) -> Result, Error> { - let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else { - return Ok(None); - }; - let params = TraceSpansParams { - access: access.clone(), - trace_id: trace_id.to_owned(), - trace_ref: trace_ref.clone(), - }; - let rows = crate::span_batches::read_spans(client, connection, params, u64::MAX).await?; - if rows.is_empty() { - return Ok(None); + .await } - let spend_rows = spend(client, connection, access, &rows).await; - Ok(resolve_trace(trace_id, &trace_ref, &rows, &spend_rows)) -} -#[derive(Deserialize, Serialize)] -struct SpanPosition { - trace_ref: String, - snapshot_ms: u64, - offset: usize, - version: String, -} - -pub async fn get_trace_page( - client: &Client, - connection: &Connection, - access: &ReadAccessParams, - trace_id: &str, - trace_ref: &str, - cursor: Option<&str>, - page_size: u32, -) -> Result, Error> { - if !(1..=500).contains(&page_size) { - return Err(Error::InvalidParameters); + async fn spend( + &self, + params: &contracts::SpendByResponseIdsParams, + ) -> Result, StoreError> { + crate::span_batches::read_spend( + &self.client, + &self.connection, + SpendByResponseIdsParams::from(params.clone()), + ) + .await } - let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else { - return Ok(None); - }; - let position = match cursor { - Some(cursor) => { - let position: SpanPosition = decode_cursor(cursor, "span")?; - if position.trace_ref != trace_ref || position.snapshot_ms == 0 { - return Err(Error::InvalidCursor("span")); - } - position + + async fn span_detail( + &self, + params: &contracts::SpanDetailParams, + ) -> Result, StoreError> { + match fetch::(&self.client, &self.connection, params).await { + Ok(rows) => Ok(rows.into_iter().next()), + Err(error) => Err(failed(error)), } - None => SpanPosition { - trace_ref: trace_ref.clone(), - snapshot_ms: (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) - as u64, - offset: 0, - version: String::new(), - }, - }; - let key = SnapshotKey::new( - connection.url().as_str(), - access, - trace_id, - &trace_ref, - position.snapshot_ms, - ) - .map_err(|_| Error::InvalidParameters)?; - let snapshot = match TRACE_SNAPSHOTS.get(&key).await { - Some(snapshot) => snapshot, - None => { - let params = TraceSpansParams { - access: access.clone(), - trace_id: trace_id.to_owned(), - trace_ref: trace_ref.clone(), - }; - let rows = - crate::span_batches::read_spans(client, connection, params, position.snapshot_ms) - .await?; - let spend_rows = spend(client, connection, access, &rows).await; - let Some(trace) = resolve_trace(trace_id, &trace_ref, &rows, &spend_rows) else { - return Ok(None); - }; - TRACE_SNAPSHOTS.insert(key, trace).await? + } + + async fn span_error( + &self, + params: &contracts::SpanErrorParams, + ) -> Result, StoreError> { + let storage_params = SpanErrorParams::from(params.clone()); + match fetch::(&self.client, &self.connection, &storage_params).await { + Ok(rows) => Ok(rows.into_iter().next().map(|row| row.0)), + Err(error) => Err(failed(error)), } - }; - let spans = &snapshot.trace().spans; - if cursor.is_some() && position.version != snapshot.version() { - return Err(Error::TraceChanged); - } - let mut trace = Trace { - summary: snapshot.trace().summary.clone(), - agents: snapshot.trace().agents.clone(), - spans: Vec::new(), - next_cursor: None, - }; - if position.offset > spans.len() { - return Err(Error::InvalidCursor("span")); - } - let end = position - .offset - .saturating_add(page_size as usize) - .min(spans.len()); - trace.next_cursor = (end < spans.len()).then(|| { - encode_cursor(&SpanPosition { - offset: end, - version: snapshot.version().to_owned(), - ..position - }) - }); - trace.spans = spans[position.offset..end].to_vec(); - while serde_json::to_vec(&trace) - .map_err(|_| Error::InvalidResponse)? - .len() - > litellm_storage_clickhouse::READ_LIMITS.response_bytes - { - if trace.spans.len() <= 1 { - return Err(Error::ReadTooLarge); - } - trace.spans.truncate(trace.spans.len() / 2); - trace.next_cursor = Some(encode_cursor(&SpanPosition { - trace_ref: trace_ref.clone(), - snapshot_ms: position.snapshot_ms, - offset: position.offset + trace.spans.len(), - version: snapshot.version().to_owned(), - })); - } - Ok(Some(trace)) -} - -pub async fn get_span( - client: &Client, - connection: &Connection, - access: &ReadAccessParams, - trace_id: &str, - span_id: &str, - trace_ref: &str, -) -> Result, Error> { - let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else { - return Ok(None); - }; - let params = SpanDetailParams { - access: access.clone(), - trace_id: trace_id.to_owned(), - trace_ref, - span_id: span_id.to_owned(), - }; - let row = fetch::(client, connection, ¶ms) - .await? - .into_iter() - .next(); - Ok(row.map(|row| SpanDetail { - input_ui: to_ui_content(&row.input), - output_ui: to_ui_content(&row.output), - span_id: row.span_id, - input: row.input, - output: row.output, - attributes: row.attributes, - })) -} - -pub async fn get_span_error( - client: &Client, - connection: &Connection, - access: &ReadAccessParams, - trace_id: &str, - span_id: &str, - trace_ref: &str, - cursor: Option<&str>, -) -> Result, Error> { - let position = error_position(cursor)?; - let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else { - return Ok(None); - }; - let offset = position.as_ref().map_or(0, |position| position.offset); - let params = SpanErrorParams::from(contracts::SpanErrorParams { - access: access.clone(), - trace_id: trace_id.to_owned(), - trace_ref, - span_id: span_id.to_owned(), - error_offset: offset, - error_version: position - .map(|position| position.version) - .unwrap_or_default(), - }); - let Some(row) = fetch::(client, connection, ¶ms) - .await? - .into_iter() - .next() - else { - return Ok(None); - }; - let row = row.0; - let next_offset = offset + row.message.chars().count() as u64; - let next_cursor = (next_offset < row.total_chars).then(|| { - encode_cursor(&ErrorPosition { - offset: next_offset, - version: row.version, - }) - }); - Ok(Some(SpanErrorPage { - span_id: row.span_id, - message: row.message, - total_chars: row.total_chars, - next_cursor, - })) -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - - use super::*; - - #[rstest] - fn trace_cursor_round_trips_the_last_listed_run() { - let cursor = encode_cursor(&(1_790_742_989_377_i64, "4bad42b84e9de3ba46fc870185f8f023")); - assert_eq!( - trace_position(Some(&cursor)).unwrap(), - ( - 1_790_742_989_377, - "4bad42b84e9de3ba46fc870185f8f023".to_owned() - ) - ); - assert_eq!(trace_position(None).unwrap(), (0, String::new())); - assert_eq!(trace_position(Some("")).unwrap(), (0, String::new())); - } - - #[rstest] - #[case::not_base64("abc")] - #[case::not_json("bm90LWpzb24=")] - #[case::numeric_reference("WzEsIDJd")] - #[case::zero_start("WzAsICJ0Il0=")] - fn malformed_trace_cursors_are_rejected(#[case] cursor: &str) { - assert!(matches!( - trace_position(Some(cursor)), - Err(Error::InvalidCursor("trace")) - )); - } - - #[rstest] - #[case::not_base64("garbage")] - #[case::missing_fields("e30=")] - #[case::not_an_object("WzEsMl0=")] - fn malformed_diagnostic_cursors_are_rejected(#[case] cursor: &str) { - assert!(matches!( - error_position(Some(cursor)), - Err(Error::InvalidCursor("diagnostic")) - )); - } - - #[rstest] - #[case::lowercase_version("a".repeat(64))] - #[case::short_version("A".repeat(63))] - fn diagnostic_cursor_requires_a_content_version(#[case] version: String) { - let cursor = encode_cursor(&ErrorPosition { offset: 1, version }); - assert!(matches!( - error_position(Some(&cursor)), - Err(Error::InvalidCursor("diagnostic")) - )); } } + +fn failed(error: StorageError) -> StoreError { + StoreError::Failed(Error::Storage(error)) +} diff --git a/litellm-rust/crates/traces-clickhouse/src/span_batches.rs b/litellm-rust/crates/traces-clickhouse/src/span_batches.rs index d66ad9506da..7d77feadeee 100644 --- a/litellm-rust/crates/traces-clickhouse/src/span_batches.rs +++ b/litellm-rust/crates/traces-clickhouse/src/span_batches.rs @@ -1,15 +1,20 @@ -use futures_util::{TryStreamExt, stream}; -use itertools::Itertools; +//! Keyset-paged reads that shrink their page when ClickHouse rejects a response as too large and +//! stop accumulating once a graph exceeds the interactive budget. + +use std::{future::Future, marker::PhantomData}; + use litellm_http::Client; use litellm_storage_clickhouse::{Query, fetch}; use litellm_traces::query::named as contracts; -use serde::Serialize; +use litellm_traces_cache::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, StoreError}; +use serde::{Serialize, de::DeserializeOwned}; -use crate::{Connection, Error, query::named::TraceSpansRow}; +use crate::{ + Connection, Error, + query::named::{SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow}, +}; const PAGE_SIZE: u32 = 256; -pub(crate) const MAX_GRAPH_BYTES: usize = 64 * 1024 * 1024; -const MAX_GRAPH_SPANS: usize = 100_000; #[derive(Default)] struct ReadBudget { @@ -18,40 +23,138 @@ struct ReadBudget { } impl ReadBudget { - fn checked_add(&self, bytes: usize, rows: usize) -> Result { - let next = Self { - bytes: self.bytes.saturating_add(bytes), - rows: self.rows.saturating_add(rows), - }; - if next.bytes > MAX_GRAPH_BYTES || next.rows > MAX_GRAPH_SPANS { - return Err(Error::ReadTooLarge); + fn reserve(&mut self, bytes: usize) -> Result<(), StoreError> { + self.bytes = self.bytes.saturating_add(bytes); + if self.bytes > MAX_GRAPH_BYTES || self.rows == MAX_GRAPH_SPANS { + return Err(StoreError::TooLarge); } - Ok(next) - } - - fn record(&mut self, row: &impl Serialize) -> Result<(), Error> { - let bytes = serde_json::to_vec(row).map_err(|_| Error::InvalidResponse)?; - *self = self.checked_add(bytes.len(), 1)?; + self.rows += 1; Ok(()) } + + fn record(&mut self, row: &impl Serialize) -> Result<(), StoreError> { + let bytes = + serde_json::to_vec(row).map_err(|_| StoreError::Failed(Error::InvalidResponse))?; + self.reserve(bytes.len()) + } +} + +/// One keyset position in a paged query: the SQL reads the cursor fields of `Self` plus the +/// `page_size` that [`Batch`] adds. +trait Keyset: Serialize + Sized + Send + Sync { + type Row: Serialize + DeserializeOwned + Send; + const SQL: &'static str; + + /// The position just after `last`. + fn after(self, last: &Self::Row) -> Self; } #[derive(Serialize)] -struct Parameters { +struct Batch { + #[serde(flatten)] + keyset: K, + page_size: u32, +} + +trait PageSource { + fn page( + &self, + batch: &Batch, + ) -> impl Future, litellm_storage_clickhouse::Error>> + Send; +} + +struct Paged(PhantomData); + +impl Query for Paged { + type Params = Batch; + type Row = K::Row; + const SQL: &'static str = K::SQL; +} + +/// Reads every row after `keyset`. A page ClickHouse rejects as too large is retried at half the +/// size, and the smaller page is kept for the rest of the read because row sizes within one graph +/// rarely shrink again. Halving a one-row page means a single row exceeds the response limit. +async fn read_all>( + source: &S, + keyset: K, +) -> Result, StoreError> { + let mut batch = Batch { + keyset, + page_size: PAGE_SIZE, + }; + let mut rows = Vec::new(); + let mut budget = ReadBudget::default(); + loop { + let page = match source.page(&batch).await { + Err(litellm_storage_clickhouse::Error::ResponseTooLarge) if batch.page_size > 1 => { + batch.page_size /= 2; + continue; + } + Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => { + return Err(StoreError::TooLarge); + } + result => result.map_err(|error| StoreError::Failed(Error::Storage(error)))?, + }; + let complete = page.len() < batch.page_size as usize; + for row in &page { + budget.record(row)?; + } + if let Some(last) = page.last() { + batch.keyset = batch.keyset.after(last); + } + rows.extend(page); + if complete { + return Ok(rows); + } + } +} + +struct ClickHouse<'a> { + client: &'a Client, + connection: &'a Connection, +} + +impl PageSource for ClickHouse<'_> { + fn page( + &self, + batch: &Batch, + ) -> impl Future, litellm_storage_clickhouse::Error>> + Send { + fetch::>(self.client, self.connection, batch) + } +} + +async fn read_paged( + client: &Client, + connection: &Connection, + keyset: K, +) -> Result, StoreError> { + let source = ClickHouse { client, connection }; + read_all(&source, keyset).await +} + +fn by_start(mut rows: Vec) -> Vec { + rows.sort_by_key(|row| row.start_ns); + rows +} + +#[derive(Serialize)] +struct SpanKeyset { #[serde(flatten)] trace: contracts::TraceSpansParams, after_span_id: String, - page_size: u32, snapshot_ms: u64, } -struct SpanBatch; - -impl Query for SpanBatch { - type Params = Parameters; +impl Keyset for SpanKeyset { type Row = TraceSpansRow; - const SQL: &'static str = include_str!("../query/trace_span_batch.sql"); + + fn after(self, last: &TraceSpansRow) -> Self { + Self { + after_span_id: last.0.span_id.clone(), + ..self + } + } } pub(crate) async fn read_spans( @@ -59,199 +162,109 @@ pub(crate) async fn read_spans( connection: &Connection, trace: contracts::TraceSpansParams, snapshot_ms: u64, -) -> Result, Error> { - let mut parameters = Parameters { +) -> Result, StoreError> { + let keyset = SpanKeyset { trace, after_span_id: String::new(), - page_size: PAGE_SIZE, snapshot_ms, }; - let mut spans = Vec::new(); - let mut budget = ReadBudget::default(); - loop { - let page = match fetch::(client, connection, ¶meters).await { - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) - if parameters.page_size > 1 => - { - parameters.page_size /= 2; - continue; - } - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => { - return Err(Error::ReadTooLarge); - } - result => result?, - }; - let complete = page.len() < parameters.page_size as usize; - if let Some(last) = page.last() { - parameters.after_span_id.clone_from(&last.0.span_id); - } - for row in page { - budget.record(&row)?; - spans.push(row.0); - } - if complete { - spans.sort_by_key(|row| row.start_ns); - return Ok(spans); - } - parameters.page_size = (parameters.page_size * 2).min(PAGE_SIZE); - } + let rows = read_paged(client, connection, keyset).await?; + Ok(by_start(rows.into_iter().map(|row| row.0).collect())) } #[derive(Serialize)] -struct ListParameters { +struct ListSpanKeyset { #[serde(flatten)] runs: crate::query::named::TracePageSpansParams, after_team: String, after_key: String, after_trace: String, after_span: String, - page_size: u32, snapshot_ms: u64, } -struct ListSpanBatch; - -impl Query for ListSpanBatch { - type Params = ListParameters; +impl Keyset for ListSpanKeyset { type Row = TraceSpansRow; - const SQL: &'static str = include_str!("../query/trace_list_span_batch.sql"); + + fn after(self, last: &TraceSpansRow) -> Self { + Self { + after_team: last.0.team_id.clone(), + after_key: last.0.api_key_hash.clone(), + after_trace: last.0.trace_id.clone(), + after_span: last.0.span_id.clone(), + ..self + } + } } pub(crate) async fn read_list_spans( client: &Client, connection: &Connection, runs: crate::query::named::TracePageSpansParams, -) -> Result, Error> { - let parameters = ListParameters { + snapshot_ms: u64, +) -> Result, StoreError> { + let keyset = ListSpanKeyset { runs, after_team: String::new(), after_key: String::new(), after_trace: String::new(), after_span: String::new(), - page_size: PAGE_SIZE, - snapshot_ms: (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as u64, + snapshot_ms, }; - let pages = stream::try_unfold( - (Some(parameters), ReadBudget::default()), - |(parameters, budget)| async move { - let Some(parameters) = parameters else { - return Ok(None); - }; - let page = match fetch::(client, connection, ¶meters).await { - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) - if parameters.page_size > 1 => - { - let retry = ListParameters { - page_size: parameters.page_size / 2, - ..parameters - }; - return Ok(Some((Vec::new(), (Some(retry), budget)))); - } - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => { - return Err(Error::ReadTooLarge); - } - result => result?, - }; - let next = page - .last() - .filter(|_| page.len() == parameters.page_size as usize) - .map(|last| ListParameters { - after_team: last.0.team_id.clone(), - after_key: last.0.api_key_hash.clone(), - after_trace: last.0.trace_id.clone(), - after_span: last.0.span_id.clone(), - page_size: (parameters.page_size * 2).min(PAGE_SIZE), - ..parameters - }); - let next_budget = page.iter().try_fold(budget, |budget, row| { - let bytes = serde_json::to_vec(row).map_err(|_| Error::InvalidResponse)?; - budget.checked_add(bytes.len(), 1) - })?; - Ok(Some((page, (next, next_budget)))) - }, - ) - .try_collect::>() - .await?; - Ok(pages - .into_iter() - .flatten() - .map(|row| row.0) - .sorted_by_key(|row| row.start_ns) - .collect()) + let rows = read_paged(client, connection, keyset).await?; + Ok(by_start(rows.into_iter().map(|row| row.0).collect())) } #[derive(Serialize)] -struct SpendParameters { +struct SpendKeyset { #[serde(flatten)] - lookup: crate::query::named::SpendByResponseIdsParams, + lookup: SpendByResponseIdsParams, has_cursor: u8, after_team: String, after_ms: i64, after_id: String, - page_size: u32, } -struct SpendBatch; - -impl Query for SpendBatch { - type Params = SpendParameters; - type Row = crate::query::named::SpendByResponseIdsRow; - +impl Keyset for SpendKeyset { + type Row = SpendByResponseIdsRow; const SQL: &'static str = include_str!("../query/spend_batch.sql"); + + fn after(self, last: &SpendByResponseIdsRow) -> Self { + Self { + has_cursor: 1, + after_team: last.0.team_id.clone(), + after_ms: last.0.start_ms, + after_id: last.0.request_id.clone(), + ..self + } + } } pub(crate) async fn read_spend( client: &Client, connection: &Connection, - lookup: crate::query::named::SpendByResponseIdsParams, -) -> Result, Error> { - let mut parameters = SpendParameters { + lookup: SpendByResponseIdsParams, +) -> Result, StoreError> { + let keyset = SpendKeyset { lookup, has_cursor: 0, after_team: String::new(), after_ms: 0, after_id: String::new(), - page_size: PAGE_SIZE, }; - let mut rows = Vec::new(); - let mut budget = ReadBudget::default(); - loop { - let page = match fetch::(client, connection, ¶meters).await { - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) - if parameters.page_size > 1 => - { - parameters.page_size /= 2; - continue; - } - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => { - return Err(Error::ReadTooLarge); - } - result => result?, - }; - let complete = page.len() < parameters.page_size as usize; - if let Some(last) = page.last() { - parameters.has_cursor = 1; - parameters.after_team.clone_from(&last.0.team_id); - parameters.after_ms = last.0.start_ms; - parameters.after_id.clone_from(&last.0.request_id); - } - for row in page { - budget.record(&row)?; - rows.push(row.0); - } - if complete { - return Ok(rows); - } - parameters.page_size = (parameters.page_size * 2).min(PAGE_SIZE); - } + let rows = read_paged(client, connection, keyset).await?; + Ok(rows.into_iter().map(|row| row.0).collect()) } #[cfg(test)] mod tests { - use super::*; + use std::sync::Mutex; + use rstest::rstest; + use super::*; + #[rstest] #[case::byte_boundary(MAX_GRAPH_BYTES - 1, 0, 1, false)] #[case::byte_overflow(MAX_GRAPH_BYTES - 1, 0, 2, true)] @@ -264,7 +277,77 @@ mod tests { #[case] next: usize, #[case] rejected: bool, ) { - let budget = ReadBudget { bytes, rows }; - assert_eq!(budget.checked_add(next, 1).is_err(), rejected); + let mut budget = ReadBudget { bytes, rows }; + assert_eq!(budget.reserve(next).is_err(), rejected); + } + + #[derive(Serialize)] + struct Numbers { + after: u32, + } + + impl Keyset for Numbers { + type Row = u32; + const SQL: &'static str = ""; + + fn after(self, last: &u32) -> Self { + Self { after: *last } + } + } + + /// A table of `total` rows whose transport rejects any page larger than `largest_page`. + struct Table { + total: u32, + largest_page: u32, + requests: Mutex>, + } + + impl PageSource for Table { + async fn page( + &self, + batch: &Batch, + ) -> Result, litellm_storage_clickhouse::Error> { + self.requests.lock().unwrap().push(batch.page_size); + if batch.page_size > self.largest_page { + return Err(litellm_storage_clickhouse::Error::ResponseTooLarge); + } + let end = (batch.keyset.after + batch.page_size).min(self.total); + Ok((batch.keyset.after + 1..=end).collect()) + } + } + + #[rstest] + #[case::fits(1000, PAGE_SIZE, &[256, 256, 256, 256])] + #[case::uniform_large_rows(1000, 100, &[256, 128, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64])] + #[tokio::test] + async fn a_rejected_page_size_is_not_retried( + #[case] total: u32, + #[case] largest_page: u32, + #[case] requests: &[u32], + ) { + let table = Table { + total, + largest_page, + requests: Mutex::new(Vec::new()), + }; + let rows = read_all(&table, Numbers { after: 0 }).await.unwrap(); + assert_eq!(rows, (1..=total).collect::>()); + assert_eq!(table.requests.lock().unwrap().as_slice(), requests); + } + + #[rstest] + #[tokio::test] + async fn a_single_oversized_row_fails_the_read() { + let table = Table { + total: 10, + largest_page: 0, + requests: Mutex::new(Vec::new()), + }; + let result = read_all(&table, Numbers { after: 0 }).await; + assert!(matches!(result, Err(StoreError::TooLarge)), "{result:?}"); + assert_eq!( + table.requests.lock().unwrap().as_slice(), + &[256, 128, 64, 32, 16, 8, 4, 2, 1] + ); } } diff --git a/litellm-rust/crates/traces-clickhouse/tests/reads.rs b/litellm-rust/crates/traces-clickhouse/tests/reads.rs index 2abb308f9c1..e1996d6af04 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/reads.rs @@ -1,8 +1,10 @@ use std::collections::BTreeMap; +use litellm_http::Client; use litellm_traces::query::named::ReadAccessParams; +use litellm_traces_cache::{ReadError, TraceReader}; use litellm_traces_clickhouse::{ - Connection, InsertTable, QueryScope, get_trace, get_trace_page, insert_rows, list_traces, + ClickHouseTraces, Connection, InsertTable, QueryScope, insert_rows, }; use rstest::rstest; use serde_json::json; @@ -14,6 +16,13 @@ mod support; use fixtures::{DATABASE, SeededDatabase, migrated_database, seeded_database}; use support::TestResult; +fn make_reader(client: &Client, connection: Connection) -> (TraceReader, ClickHouseTraces) { + ( + TraceReader::new(litellm_storage_clickhouse::READ_LIMITS.response_bytes), + ClickHouseTraces::new(client.clone(), connection), + ) +} + #[rstest] #[case::api_key("key-a", "")] #[case::user("", "user-a")] @@ -74,16 +83,19 @@ async fn list_costs_match_each_run_when_response_ids_are_reused( .collect(), ) .await?; - let reader = fixture + let connection = fixture .readers .connection(client, &QueryScope::All, "fixture-secret") .await?; + let (reader, store) = make_reader(client, connection); let access = ReadAccessParams { all_teams: false, user_id: user_id.into(), team_ids: vec!["team-a".into()], }; - let page = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?; + let page = reader + .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .await?; assert_eq!(page.data.len(), runs.len()); for (trace_id, _, cost) in runs { let summary = page @@ -91,7 +103,8 @@ async fn list_costs_match_each_run_when_response_ids_are_reused( .iter() .find(|summary| summary.trace_id == trace_id) .ok_or("missing run")?; - let detail = get_trace(client, &reader, &access, trace_id, &summary.trace_ref) + let detail = reader + .get_trace(&store, &access, trace_id, &summary.trace_ref) .await? .ok_or("missing trace")?; assert_eq!(detail.summary.spend, Some(cost)); @@ -199,16 +212,19 @@ async fn large_runs_remain_complete_under_default_reader_limits( .collect::>(); insert_rows(client, &writer, DATABASE, InsertTable::SpendLogs, costs).await?; } - let reader = fixture + let connection = fixture .readers .connection(client, &QueryScope::All, "fixture-secret") .await?; + let (reader, store) = make_reader(client, connection); let access = ReadAccessParams { all_teams: false, user_id: String::new(), team_ids: vec!["team-a".into()], }; - let page = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 500).await?; + let page = reader + .list_traces(&store, &access, 0, 2_000_000_000_000, None, 500) + .await?; assert_eq!(page.data.len(), runs); assert!( page.data @@ -222,42 +238,17 @@ async fn large_runs_remain_complete_under_default_reader_limits( .send() .await? .error_for_status()?; - let read_queries = client.post(writer.url().clone()).body(format!( - "SELECT count() FROM system.query_log WHERE type = 'QueryFinish' AND current_database = '{DATABASE}' AND query LIKE '%FROM otel_traces AS o%' AND query NOT LIKE '%system.query_log%'" - )).send().await?.error_for_status()?.text().await?; - let read_queries = read_queries.trim().parse::()?; - assert!( - read_queries > 0 && read_queries < runs, - "{read_queries} span queries for {runs} runs" - ); - if costed { - let overlapping = client - .post(writer.url().clone()) - .body(format!( - "WITH spend_reads AS ( - SELECT query_start_time_microseconds AS started, event_time_microseconds AS finished - FROM system.query_log - WHERE type = 'QueryFinish' AND current_database = '{DATABASE}' - AND query LIKE '%FROM spend_logs FINAL%' AND query NOT LIKE '%system.query_log%' - ), events AS ( - SELECT started AS at, 1 AS delta FROM spend_reads - UNION ALL SELECT finished AS at, -1 AS delta FROM spend_reads - ) - SELECT max(active) FROM ( - SELECT sum(delta) OVER (ORDER BY at, delta ROWS UNBOUNDED PRECEDING) AS active - FROM events - )" - )) - .send() - .await? - .error_for_status()? - .text() - .await? - .trim() - .parse::()?; + for table in ["otel_traces AS o", "spend_logs FINAL"] + .into_iter() + .take(if costed { 2 } else { 1 }) + { + let read_queries = client.post(writer.url().clone()).body(format!( + "SELECT count() FROM system.query_log WHERE type = 'QueryFinish' AND current_database = '{DATABASE}' AND query LIKE '%FROM {table}%' AND query NOT LIKE '%system.query_log%'" + )).send().await?.error_for_status()?.text().await?; + let read_queries = read_queries.trim().parse::()?; assert!( - (2..=4).contains(&overlapping), - "{overlapping} simultaneous spend reads for {runs} runs" + read_queries > 0 && read_queries < runs, + "{read_queries} {table} queries for {runs} runs" ); } } @@ -281,7 +272,8 @@ async fn large_runs_remain_complete_under_default_reader_limits( .find(|run| run.trace_id == "trace-0000") .ok_or("missing run")? .trace_ref; - let detail = get_trace(client, &reader, &access, "trace-0000", trace_ref) + let detail = reader + .get_trace(&store, &access, "trace-0000", trace_ref) .await? .ok_or("missing trace")?; assert_eq!(detail.spans.len(), steps); @@ -303,24 +295,25 @@ async fn large_runs_remain_complete_under_default_reader_limits( ..access.clone() }; assert!( - get_trace(client, &reader, &denied, "trace-0000", trace_ref) + reader + .get_trace(&store, &denied, "trace-0000", trace_ref) .await? .is_none() ); let mut cursor = None; let mut ids = Vec::new(); loop { - let page = get_trace_page( - client, - &reader, - &access, - "trace-0000", - trace_ref, - cursor.as_deref(), - 200, - ) - .await? - .ok_or("missing page")?; + let page = reader + .get_trace_page( + &store, + &access, + "trace-0000", + trace_ref, + cursor.as_deref(), + 200, + ) + .await? + .ok_or("missing page")?; assert_eq!(page.summary, detail.summary); assert!(page.spans.len() <= 200); assert!( @@ -329,17 +322,17 @@ async fn large_runs_remain_complete_under_default_reader_limits( ); if ids.is_empty() { assert!( - get_trace_page( - client, - &reader, - &denied, - "trace-0000", - trace_ref, - page.next_cursor.as_deref(), - 200, - ) - .await? - .is_none() + reader + .get_trace_page( + &store, + &denied, + "trace-0000", + trace_ref, + page.next_cursor.as_deref(), + 200, + ) + .await? + .is_none() ); client .post(writer.url().clone()) @@ -372,45 +365,43 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( ) -> TestResult { let fixture = seeded_database?; let client = &fixture.database.client; - let reader = fixture + let connection = fixture .readers .connection(client, &QueryScope::All, "fixture-secret") .await?; + let (reader, store) = make_reader(client, connection.clone()); let access = ReadAccessParams { all_teams: true, user_id: String::new(), team_ids: Vec::new(), }; - let listed = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 10).await?; + let listed = reader + .list_traces(&store, &access, 0, 2_000_000_000_000, None, 10) + .await?; let summary = listed .data .iter() .find(|summary| summary.span_count == 3) .ok_or("missing fixture")?; - let first = get_trace_page( - client, - &reader, - &access, - &summary.trace_id, - &summary.trace_ref, - None, - 1, - ) - .await? - .ok_or("missing first page")?; - let original_ids = get_trace( - client, - &reader, - &access, - &summary.trace_id, - &summary.trace_ref, - ) - .await? - .ok_or("missing trace")? - .spans - .into_iter() - .map(|span| span.span_id) - .collect::>(); + let first = reader + .get_trace_page( + &store, + &access, + &summary.trace_id, + &summary.trace_ref, + None, + 1, + ) + .await? + .ok_or("missing first page")?; + let original_ids = reader + .get_trace(&store, &access, &summary.trace_id, &summary.trace_ref) + .await? + .ok_or("missing trace")? + .spans + .into_iter() + .map(|span| span.span_id) + .collect::>(); let writer = Connection::writer(&fixture.database.url)?; insert_rows( client, @@ -434,17 +425,17 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( team_ids: vec!["not-this-team".into()], }; assert!( - get_trace_page( - client, - &reader, - &denied, - &summary.trace_id, - &summary.trace_ref, - first.next_cursor.as_deref(), - 1 - ) - .await? - .is_none() + reader + .get_trace_page( + &store, + &denied, + &summary.trace_id, + &summary.trace_ref, + first.next_cursor.as_deref(), + 1 + ) + .await? + .is_none() ); let first_cursor = first.next_cursor.clone(); let mut cursor = first.next_cursor; @@ -454,44 +445,45 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( .map(|span| span.span_id) .collect::>(); while let Some(current) = cursor { - let next = get_trace_page( - client, - &reader, - &access, - &summary.trace_id, - &summary.trace_ref, - Some(¤t), - 1, - ) - .await? - .ok_or("missing next page")?; + let next = reader + .get_trace_page( + &store, + &access, + &summary.trace_id, + &summary.trace_ref, + Some(¤t), + 1, + ) + .await? + .ok_or("missing next page")?; assert_eq!(next.summary.span_count, 3); ids.extend(next.spans.into_iter().map(|span| span.span_id)); cursor = next.next_cursor; } assert_eq!(ids, original_ids); - let refreshed = get_trace( - client, - &reader, - &access, - &summary.trace_id, - &summary.trace_ref, - ) - .await? - .ok_or("missing refreshed trace")?; + let cached = reader + .get_trace(&store, &access, &summary.trace_id, &summary.trace_ref) + .await? + .ok_or("missing cached trace")?; + assert_eq!(cached.spans.len(), 3); + let (fresh_reader, fresh_store) = make_reader(client, connection); + let refreshed = fresh_reader + .get_trace(&fresh_store, &access, &summary.trace_id, &summary.trace_ref) + .await? + .ok_or("missing refreshed trace")?; assert_eq!(refreshed.spans.len(), 4); assert!(matches!( - get_trace_page( - client, - &reader, - &access, - &summary.trace_id, - &summary.trace_ref, - Some("invalid"), - 1 - ) - .await, - Err(litellm_traces_clickhouse::Error::InvalidCursor("span")) + reader + .get_trace_page( + &store, + &access, + &summary.trace_id, + &summary.trace_ref, + Some("invalid"), + 1 + ) + .await, + Err(ReadError::InvalidCursor("span")) )); let backdated = json!({ "Timestamp": "2026-09-01 00:00:00.000000000", @@ -509,20 +501,21 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( .send() .await? .error_for_status()?; - let uncached_reader = + let uncached_connection = Connection::reader(&format!("{}?max_threads=1", fixture.database.url), DATABASE)?; - let changed = get_trace_page( - client, - &uncached_reader, - &access, - &summary.trace_id, - &summary.trace_ref, - first_cursor.as_deref(), - 1, - ) - .await; + let (uncached_reader, uncached_store) = make_reader(client, uncached_connection); + let changed = uncached_reader + .get_trace_page( + &uncached_store, + &access, + &summary.trace_id, + &summary.trace_ref, + first_cursor.as_deref(), + 1, + ) + .await; assert!( - matches!(changed, Err(litellm_traces_clickhouse::Error::TraceChanged)), + matches!(changed, Err(ReadError::TraceChanged)), "{changed:?}" ); Ok(()) @@ -535,16 +528,19 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( ) -> TestResult { let fixture = seeded_database?; let client = &fixture.database.client; - let reader = fixture + let connection = fixture .readers .connection(client, &QueryScope::All, "fixture-secret") .await?; + let (reader, store) = make_reader(client, connection.clone()); let access = ReadAccessParams { all_teams: true, user_id: String::new(), team_ids: Vec::new(), }; - let before = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?; + let before = reader + .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .await?; let run = before .data .iter() @@ -571,7 +567,14 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( ])], ) .await?; - let after = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?; + let cached = reader + .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .await?; + assert_eq!(cached.data, before.data); + let (reader, store) = make_reader(client, connection); + let after = reader + .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .await?; assert_eq!(after.data.len(), before.data.len()); let limited = after .data @@ -588,17 +591,10 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( .all(|item| !item.resolution_limited) ); assert!(matches!( - get_trace_page( - client, - &reader, - &access, - &run.trace_id, - &run.trace_ref, - None, - 200 - ) - .await, - Err(litellm_traces_clickhouse::Error::ReadTooLarge) + reader + .get_trace_page(&store, &access, &run.trace_id, &run.trace_ref, None, 200) + .await, + Err(ReadError::TooLarge) )); Ok(()) } @@ -685,16 +681,19 @@ async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback .collect(), ) .await?; - let reader = fixture + let connection = fixture .readers .connection(client, &QueryScope::All, "fixture-secret") .await?; + let (reader, store) = make_reader(client, connection); let access = ReadAccessParams { all_teams: false, user_id: String::new(), team_ids: vec!["team-a".into()], }; - let page = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?; + let page = reader + .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .await?; assert_eq!(page.data.len(), cases.len()); for (id, _, _, _, _, expected) in cases { let summary = page @@ -702,7 +701,8 @@ async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback .iter() .find(|summary| summary.trace_id == id) .ok_or("missing run")?; - let detail = get_trace(client, &reader, &access, id, &summary.trace_ref) + let detail = reader + .get_trace(&store, &access, id, &summary.trace_ref) .await? .ok_or("missing trace")?; assert_eq!(detail.summary.spend, expected, "{id}"); diff --git a/litellm-rust/crates/traces/src/query/named.rs b/litellm-rust/crates/traces/src/query/named.rs index 5765efaa62e..a981bc46471 100644 --- a/litellm-rust/crates/traces/src/query/named.rs +++ b/litellm-rust/crates/traces/src/query/named.rs @@ -15,7 +15,7 @@ pub struct ReadAccessParams { pub team_ids: Vec, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct ListTracesParams { #[serde(flatten)] pub access: ReadAccessParams, @@ -26,7 +26,7 @@ pub struct ListTracesParams { pub limit: u32, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct ListTracesRow { pub trace_id: String, pub trace_ref: String, @@ -56,7 +56,7 @@ pub struct ListTracesRow { pub request_ids: Vec, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct TraceSpansParams { #[serde(flatten)] pub access: ReadAccessParams, @@ -113,7 +113,7 @@ pub struct TraceSpansRow { pub user_id: String, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct TracePageSpansParams { #[serde(flatten)] pub access: ReadAccessParams, @@ -131,7 +131,7 @@ pub struct SpanDetailParams { pub span_id: String, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct SpanDetailRow { pub span_id: String, pub input: String, @@ -139,7 +139,7 @@ pub struct SpanDetailRow { pub attributes: BTreeMap, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct SpanErrorParams { #[serde(flatten)] pub access: ReadAccessParams, @@ -150,7 +150,7 @@ pub struct SpanErrorParams { pub error_version: String, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct SpanErrorRow { pub span_id: String, pub message: String, @@ -158,7 +158,7 @@ pub struct SpanErrorRow { pub version: String, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct SpendByResponseIdsParams { #[serde(flatten)] pub access: ReadAccessParams, @@ -169,7 +169,7 @@ pub struct SpendByResponseIdsParams { pub end_ms: i64, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct SpendByResponseIdsRow { pub request_id: String, pub litellm_call_id: String, diff --git a/litellm/constants.py b/litellm/constants.py index d58fc8a6318..69ff3cf7a5e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -55,6 +55,7 @@ DEFAULT_AGENT_TRACING_RETENTION_DAYS: Final = 14 OTLP_MAX_BODY_BYTES: Final = get_env_int("OTLP_MAX_BODY_BYTES", 16 * 1024 * 1024) OTLP_MAX_ATTRIBUTE_VALUE_BYTES: Final = get_env_int("OTLP_MAX_ATTRIBUTE_VALUE_BYTES", 64 * 1024) OTLP_RETRY_AFTER_SECONDS: Final = get_env_int("OTLP_RETRY_AFTER_SECONDS", 2) +TRACE_READ_RETRY_AFTER_SECONDS: Final = get_env_int("TRACE_READ_RETRY_AFTER_SECONDS", 2) OTLP_MAX_CONCURRENT_INGESTS: Final = get_env_int("OTLP_MAX_CONCURRENT_INGESTS", 2) AGENT_TRACING_INPUT_PREVIEW_CHARS: Final = get_env_int("AGENT_TRACING_INPUT_PREVIEW_CHARS", 240) AGENT_TRACING_LIST_PAGE_SIZE: Final = get_env_int("AGENT_TRACING_LIST_PAGE_SIZE", 50) diff --git a/litellm/proxy/tracing_endpoints.py b/litellm/proxy/tracing_endpoints.py index 1563e4b5b55..503280ea7e8 100644 --- a/litellm/proxy/tracing_endpoints.py +++ b/litellm/proxy/tracing_endpoints.py @@ -13,19 +13,21 @@ from dataclasses import dataclass from functools import partial from http.client import responses from types import MappingProxyType -from typing import Annotated, Final +from typing import Annotated, Final, Literal from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response from pydantic import BaseModel, ConfigDict +from typing_extensions import assert_never from litellm._logging import verbose_proxy_logger -from litellm.constants import OTLP_RETRY_AFTER_SECONDS +from litellm.constants import OTLP_RETRY_AFTER_SECONDS, TRACE_READ_RETRY_AFTER_SECONDS from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.authorization import AllRows, ReadScope, resolve_trace_read_scope from litellm.proxy.auth.authorization_dependencies import LogTeamLookupDependency from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import is_otlp_trace_request from litellm.proxy.tracing_runtime import provide_receiver, require_receiver +from litellm.rust_bridge.trace.errors import TraceChanged from litellm.rust_bridge.trace.generated.models import TraceQueryHelp from litellm.rust_bridge.trace.generated.types import ( AllQueryScope, @@ -135,6 +137,49 @@ async def ingest_otlp_traces( return Response(content=body, media_type=media_type) +class TraceReadFailure(BaseModel): + """The body of every failed trace read. Clients branch on `code`, never on `message`.""" + + model_config = ConfigDict(frozen=True) + + code: Literal["invalid_request", "trace_changed", "too_large", "unavailable"] + message: str + + +def read_failure(error: TraceChanged | ValueError | OverflowError | RuntimeError) -> HTTPException: + """One status per failure kind, so a client can tell a bad cursor (400, fix the request) from a + traversal it must restart (409), a result it cannot page through (413), and an outage it should + retry after `Retry-After` (503).""" + match error: + case TraceChanged(): + return HTTPException( + status_code=409, + detail=TraceReadFailure(code="trace_changed", message=str(error)).model_dump(), + ) + case ValueError(): + return HTTPException( + status_code=400, detail=TraceReadFailure(code="invalid_request", message=str(error)).model_dump() + ) + case OverflowError(): + return HTTPException( + status_code=413, + detail=TraceReadFailure( + code="too_large", message="Trace is too large for this view. Use a filtered trace query." + ).model_dump(), + ) + case RuntimeError(): + verbose_proxy_logger.warning("Trace read unavailable: %s", error) + return HTTPException( + status_code=503, + detail=TraceReadFailure( + code="unavailable", message="Traces are temporarily unavailable. Please try again." + ).model_dump(), + headers={"Retry-After": str(TRACE_READ_RETRY_AFTER_SECONDS)}, + ) + case _: + assert_never(error) + + @router.get("/v1/traces", response_model=TracePage) async def list_agent_traces( context: Annotated[TraceAccessContext, Depends(provide_trace_access)], @@ -151,15 +196,8 @@ async def list_agent_traces( end_ms=end_ms if end_ms is not None else now_ms, cursor=cursor, ) - except ValueError as error: - raise HTTPException(status_code=400, detail=str(error)) from error - except OverflowError as error: - raise HTTPException( - status_code=413, detail="Trace is too large for this view. Use a filtered trace query." - ) from error - except RuntimeError as error: - verbose_proxy_logger.warning("Trace read unavailable: %s", error) - raise HTTPException(status_code=503, detail="Traces are temporarily unavailable. Please try again.") from error + except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: + raise read_failure(error) from error class TraceQueryRequest(BaseModel): @@ -241,15 +279,8 @@ async def get_agent_trace( tracing, scope = context.reader() try: trace: Final = await tracing.get_trace(trace_id, scope, trace_ref, cursor, page_size) - except ValueError as error: - raise HTTPException(status_code=400, detail=str(error)) from error - except OverflowError as error: - raise HTTPException( - status_code=413, detail="Trace is too large for this view. Use a filtered trace query." - ) from error - except RuntimeError as error: - verbose_proxy_logger.warning("Trace read unavailable: %s", error) - raise HTTPException(status_code=503, detail="Traces are temporarily unavailable. Please try again.") from error + except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: + raise read_failure(error) from error if trace is None: raise HTTPException(status_code=404, detail=f"Trace {trace_id} not found") return trace @@ -265,15 +296,8 @@ async def get_agent_trace_span( tracing, scope = context.reader() try: span: Final = await tracing.get_span(trace_id, span_id, scope, trace_ref) - except ValueError as error: - raise HTTPException(status_code=400, detail=str(error)) from error - except OverflowError as error: - raise HTTPException( - status_code=413, detail="Trace is too large for this view. Use a filtered trace query." - ) from error - except RuntimeError as error: - verbose_proxy_logger.warning("Trace read unavailable: %s", error) - raise HTTPException(status_code=503, detail="Traces are temporarily unavailable. Please try again.") from error + except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: + raise read_failure(error) from error if span is None: raise HTTPException(status_code=404, detail=f"Span {span_id} not found") return span @@ -290,15 +314,8 @@ async def get_agent_trace_span_error( try: tracing, scope = context.reader() page: Final = await tracing.get_span_error(trace_id, span_id, scope, trace_ref, cursor) - except ValueError as error: - raise HTTPException(status_code=400, detail=str(error)) from error - except OverflowError as error: - raise HTTPException( - status_code=413, detail="Trace is too large for this view. Use a filtered trace query." - ) from error - except RuntimeError as error: - verbose_proxy_logger.warning("Trace read unavailable: %s", error) - raise HTTPException(status_code=503, detail="Traces are temporarily unavailable. Please try again.") from error + except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: + raise read_failure(error) from error if page is None: raise HTTPException(status_code=404, detail="Span diagnostic not found or no longer available") return page diff --git a/litellm/rust_bridge/trace/errors.py b/litellm/rust_bridge/trace/errors.py new file mode 100644 index 00000000000..3c84d233be4 --- /dev/null +++ b/litellm/rust_bridge/trace/errors.py @@ -0,0 +1,5 @@ +class TraceChanged(Exception): + """The paging snapshot no longer matches the stored trace, so the client must start a new traversal. + + Raised by the Rust trace reader when a cursor's snapshot version differs from the graph it rebuilt. + """ diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index cff67d1f57d..913b1aec618 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -12,6 +12,7 @@ import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from litellm.constants import TRACE_READ_RETRY_AFTER_SECONDS from litellm.proxy import tracing_endpoints from litellm.proxy._types import LitellmUserRoles, ProxyLifespanState, UserAPIKeyAuth from litellm.proxy.auth.authorization import OwnedRows, ReadScope @@ -19,6 +20,7 @@ from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.tracing_runtime import manage_tracing, provide_storage from litellm.rust_bridge import loader +from litellm.rust_bridge.trace.errors import TraceChanged from litellm.rust_bridge.trace.generated.models import TraceQueryHelp from litellm.rust_bridge.trace.generated.types import AllQueryScope, TraceScope from litellm.rust_bridge.trace.queries import TraceSQLResponse @@ -297,19 +299,45 @@ def test_trace_detail_passes_scoped_reference(client, receiver, suffix, cursor, ), ) @pytest.mark.parametrize( - "error,status,message", + "error,status,code,message", ( - (RuntimeError("private database details"), 503, "Traces are temporarily unavailable. Please try again."), - (OverflowError("private query details"), 413, "Trace is too large for this view. Use a filtered trace query."), + ( + RuntimeError("private database details"), + 503, + "unavailable", + "Traces are temporarily unavailable. Please try again.", + ), + ( + OverflowError("private query details"), + 413, + "too_large", + "Trace is too large for this view. Use a filtered trace query.", + ), + ( + TraceChanged("Trace changed while paging; refresh the trace to continue"), + 409, + "trace_changed", + "Trace changed while paging; refresh the trace to continue", + ), + (ValueError("Invalid span cursor"), 400, "invalid_request", "Invalid span cursor"), ), ) -def test_read_failures_are_actionable_without_exposing_database_details( - client: TestClient, receiver: MagicMock, path: str, method: str, error: Exception, status: int, message: str +def test_read_failures_carry_a_code_per_kind_without_exposing_database_details( + client: TestClient, + receiver: MagicMock, + path: str, + method: str, + error: Exception, + status: int, + code: str, + message: str, ) -> None: getattr(receiver, method).side_effect = error response: Final = client.get(path) assert response.status_code == status - assert response.json() == {"detail": message} + assert response.json() == {"detail": {"code": code, "message": message}} + retry_after: Final = response.headers.get("Retry-After") + assert (retry_after == str(TRACE_READ_RETRY_AFTER_SECONDS)) == (status == 503), retry_after @pytest.mark.parametrize("query", ("page_size=0", "page_size=501", "cursor=" + "x" * 513)) diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/DetailContent.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/DetailContent.tsx index b9c0d9bd0b3..70e62de65d5 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/DetailContent.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/DetailContent.tsx @@ -10,6 +10,12 @@ import { cn } from "@/lib/cva.config"; import { type KeyValue, KeyValueRows, objectEntries } from "./KeyValueRows"; import { Card, MessageCard, Section, ToolResultCard } from "./MessageCard"; +import { + classifyTraceReadFailure, + isRetryableTraceRead, + traceReadRetry, + traceReadRetryDelay, +} from "../list/traceReadFailure"; import type { ErrorSource } from "../tree"; import type { Span, SpanDetail, SpanErrorPage, TraceMessage, UIContent, UIMessage } from "../types"; import { errorSource, parseJson, parseMessages, prettyPayload } from "../utils"; @@ -155,9 +161,11 @@ function DiagnosticContent({ accessToken, traceId, traceRef, span }: DetailConte enabled: opened, staleTime: Infinity, gcTime: 0, - retry: false, + retry: traceReadRetry, + retryDelay: traceReadRetryDelay, }; const query = useQuery(queryOptions); + const failure = query.error ? classifyTraceReadFailure(query.error) : null; return (
{span.error_truncated &&

Error preview truncated

} @@ -167,12 +175,18 @@ function DiagnosticContent({ accessToken, traceId, traceRef, span }: DetailConte )} {opened && query.isPending &&

Loading diagnostic…

} - {opened && query.isError && ( + {failure && (
- Could not load diagnostic: {query.error.message} - + Could not load diagnostic: {failure.message} + {isRetryableTraceRead(failure) ? ( + + ) : ( + + )}
)} {opened && query.data && ( diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.test.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.test.tsx index c2cc01353e2..036659366ed 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.test.tsx @@ -3,6 +3,8 @@ import { focusManager, onlineManager } from "@tanstack/react-query"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { ApiError } from "@/lib/http/client"; + import { renderWithProviders, testQueryClient } from "../../../../../tests/test-utils"; import researchTrace from "../__fixtures__/research_trace.json"; import swarmTrace from "../__fixtures__/swarm_trace.json"; @@ -260,14 +262,21 @@ describe("RunView", () => { vi.mocked(agentTraceCall) .mockResolvedValueOnce(first) .mockResolvedValueOnce(second) - .mockRejectedValueOnce(new Error("Trace changed while paging; refresh the trace")) + .mockRejectedValueOnce( + new ApiError("Trace changed while paging", 409, { + detail: { code: "trace_changed", message: "Trace changed while paging" }, + }), + ) .mockResolvedValueOnce(fresh) .mockResolvedValueOnce(freshSecond); renderWithProviders(); await user.click(await screen.findByRole("button", { name: "Load more steps" })); expect(await screen.findByText("old-snapshot-tool")).toBeVisible(); await user.click(screen.getByRole("button", { name: "Load more steps" })); - await user.click(await screen.findByRole("button", { name: "Refresh trace" })); + expect(await screen.findByText(/This trace changed while you were browsing/)).toBeVisible(); + expect(screen.queryByRole("button", { name: "Retry" })).not.toBeInTheDocument(); + expect(screen.getByText("old-snapshot-tool")).toBeVisible(); + await user.click(screen.getByRole("button", { name: "Refresh trace" })); expect(await screen.findByText("Showing 1 of 2 steps")).toBeVisible(); expect(screen.queryByText("old-snapshot-tool")).not.toBeInTheDocument(); await user.click(screen.getByRole("button", { name: "Load more steps" })); @@ -282,6 +291,55 @@ describe("RunView", () => { ]); }); + it("retries an unavailable page by itself after Retry-After and keeps loaded steps", async () => { + const user = userEvent.setup(); + const summary = { ...research.summary, span_count: 2 }; + const first: Trace = { ...research, summary, spans: research.spans.slice(0, 1), next_cursor: "next-page" }; + const second: Trace = { + ...research, + summary, + spans: [ + { ...research.spans[1], type: "tool", name: "later-page-tool", parent_span_id: research.spans[0].span_id }, + ], + next_cursor: null, + }; + const outage = new ApiError( + "Traces are temporarily unavailable", + 503, + { detail: { code: "unavailable", message: "Traces are temporarily unavailable" } }, + 0, + ); + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall).mockResolvedValueOnce(first).mockRejectedValueOnce(outage).mockResolvedValueOnce(second); + renderWithProviders(); + + await user.click(await screen.findByRole("button", { name: "Load more steps" })); + expect(await screen.findByText("later-page-tool")).toBeVisible(); + expect(screen.getAllByRole("treeitem")).toHaveLength(2); + expect(screen.queryByRole("button", { name: "Refresh trace" })).not.toBeInTheDocument(); + expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, "next-page", "next-page"]); + }); + + it("offers no retry for a page that is too large, only a refresh", async () => { + const user = userEvent.setup(); + const summary = { ...research.summary, span_count: 2 }; + const first: Trace = { ...research, summary, spans: research.spans.slice(0, 1), next_cursor: "next-page" }; + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(first) + .mockRejectedValueOnce( + new ApiError("Trace is too large", 413, { detail: { code: "too_large", message: "Trace is too large" } }), + ); + renderWithProviders(); + + await user.click(await screen.findByRole("button", { name: "Load more steps" })); + expect(await screen.findByText(/too large to load here/)).toBeVisible(); + expect(screen.queryByRole("button", { name: "Retry" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Refresh trace" })).toBeVisible(); + expect(screen.getByRole("tree", { name: "Spans in time order" })).toHaveTextContent(research.spans[0].name); + expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(2); + }); + it("keeps a loaded snapshot on focus and reconnect", async () => { testQueryClient.setQueryDefaults(["agentTrace"], { refetchOnWindowFocus: true, refetchOnReconnect: true }); vi.mocked(agentTraceCall).mockReset(); diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.tsx index 4047218177d..ddc66f23470 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.tsx @@ -23,6 +23,13 @@ import { TraceConversation } from "./TraceConversation"; import { FrameworkLogo, traceFramework } from "../ui/TraceFramework"; import { type RunSelection, traceKey } from "../routing"; import type { SpanTreeState, TreeRow } from "../tree"; +import { + classifyTraceReadFailure, + isRetryableTraceRead, + traceReadRetry, + traceReadRetryDelay, +} from "../list/traceReadFailure"; +import type { TraceReadFailure } from "../list/traceReadFailure"; import type { Trace } from "../types"; import { buildTreeRows, @@ -340,6 +347,60 @@ interface RunViewProps { embedded?: boolean; } +interface PagingBannerProps { + readonly loaded: number; + readonly total: number; + readonly failure: TraceReadFailure | null; + readonly busy: boolean; + readonly onLoadMore: () => void; + readonly onRefresh: () => void; +} + +/** + * Progress through the steps, or what the loaded steps are still worth after a later page failed. + * Retry keeps the loaded pages and asks for the same page again; Refresh starts a new traversal. + * Only a temporary failure gets a Retry, because a changed, invalid, or oversized read fails the + * same way every time. + */ +function PagingBanner({ loaded, total, failure, busy, onLoadMore, onRefresh }: PagingBannerProps) { + const retryable = failure !== null && isRetryableTraceRead(failure); + const nextLabel = failure ? "Retry" : "Load more steps"; + return ( +
+ + {failure + ? pageFailureMessage(failure) + : `Showing ${loaded.toLocaleString()} of ${total.toLocaleString()} steps`} + + {failure && ( + + )} + {(!failure || retryable) && ( + + )} +
+ ); +} + +/** What the loaded steps are still worth after a later page fails, and what the user can do about it. */ +export function pageFailureMessage(failure: TraceReadFailure): string { + switch (failure.kind) { + case "changed": + return "This trace changed while you were browsing. Refresh to continue from its latest steps."; + case "invalid": + return "This page could not be continued. Refresh the trace to start over."; + case "too_large": + return "A step on the next page is too large to load here. Your loaded steps are still available."; + case "unavailable": + case "unknown": + return "Could not load more steps. Your loaded steps are still available."; + } +} + function selectedSpanMissing(trace: Trace, spanId: string | null): boolean { return Boolean(spanId && !trace.spans.some((span) => span.span_id === spanId)); } @@ -427,10 +488,12 @@ function LoadedRun({ refetchOnWindowFocus: false, refetchOnReconnect: false, refetchOnMount: false, - retry: false, + retry: traceReadRetry, + retryDelay: traceReadRetryDelay, }; const traceQuery = useSuspenseInfiniteQuery(traceQueryOptions); const refreshTrace = () => queryClient.resetQueries({ queryKey, exact: true }); + const failure = traceQuery.error ? classifyTraceReadFailure(traceQuery.error) : null; const trace = useMemo(() => { const [first, ...rest] = traceQuery.data.pages; return { ...first, spans: [first, ...rest].flatMap((page) => page.spans) }; @@ -441,7 +504,6 @@ function LoadedRun({ useEffect(() => { if (canSeek && !isFetching && !isError) void fetchNextPage(); }, [canSeek, isFetching, isError, fetchNextPage]); - const pageAction = isError ? "Retry" : "Load more steps"; return ( - {(traceQuery.hasNextPage || traceQuery.isError) && ( -
- - {traceQuery.isError - ? "Could not load more steps. Your loaded steps are still available." - : `Showing ${trace.spans.length.toLocaleString()} of ${trace.summary.span_count.toLocaleString()} steps`} - - {traceQuery.isError && ( - - )} - -
+ {(traceQuery.hasNextPage || failure) && ( + void traceQuery.fetchNextPage()} + onRefresh={() => void refreshTrace()} + /> )} + new ApiError("boom", status, { detail: code ? { code, message: "boom" } : "boom" }, retryAfterMs); + +describe("classifyTraceReadFailure", () => { + it.each([ + ["invalid_request", "invalid"], + ["trace_changed", "changed"], + ["too_large", "too_large"], + ["unavailable", "unavailable"], + ] as const)("reads the %s code ahead of the status", (code, kind) => { + expect(classifyTraceReadFailure(failure(500, code)).kind).toBe(kind); + }); + + it.each([ + [400, "invalid"], + [409, "changed"], + [413, "too_large"], + [503, "unavailable"], + [500, "unknown"], + ] as const)("falls back to status %i without a code", (status, kind) => { + expect(classifyTraceReadFailure(failure(status)).kind).toBe(kind); + }); + + it("treats a non-HTTP error as unknown and keeps its message", () => { + expect(classifyTraceReadFailure(new Error("offline"))).toEqual({ + kind: "unknown", + message: "offline", + retryAfterMs: null, + }); + }); +}); + +describe("retry policy", () => { + it("retries an outage automatically after the server's Retry-After and stops at the cap", () => { + const outage = failure(503, "unavailable", 2_000); + expect(traceReadRetry(0, outage)).toBe(true); + expect(traceReadRetry(TRACE_READ_AUTO_RETRIES, outage)).toBe(false); + expect(traceReadRetryDelay(0, outage)).toBe(2_000); + expect(traceReadRetryDelay(0, failure(503, "unavailable"))).toBe(1_000); + }); + + it("never retries a changed, invalid, or oversized read", () => { + for (const code of ["trace_changed", "invalid_request", "too_large"]) { + expect(traceReadRetry(0, failure(500, code))).toBe(false); + expect(isRetryableTraceRead(classifyTraceReadFailure(failure(500, code)))).toBe(false); + } + expect(isRetryableTraceRead(classifyTraceReadFailure(new Error("offline")))).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/components/lens/traces/list/traceReadFailure.ts b/ui/litellm-dashboard/src/components/lens/traces/list/traceReadFailure.ts new file mode 100644 index 00000000000..b4b180e9caa --- /dev/null +++ b/ui/litellm-dashboard/src/components/lens/traces/list/traceReadFailure.ts @@ -0,0 +1,58 @@ +import { ApiError } from "@/lib/http/client"; + +/** + * Why a trace read failed, as the proxy reports it in the `code` of a failed response. Each kind has + * one recovery: `invalid` and `changed` need a new traversal, `too_large` cannot be paged through at + * all, `unavailable` is worth retrying after the server's `Retry-After`, and `unknown` is treated like + * `unavailable` because the loaded pages are still good. + */ +export type TraceReadFailureKind = "invalid" | "changed" | "too_large" | "unavailable" | "unknown"; + +export interface TraceReadFailure { + readonly kind: TraceReadFailureKind; + readonly message: string; + readonly retryAfterMs: number | null; +} + +const KIND_BY_CODE: Readonly> = { + invalid_request: "invalid", + trace_changed: "changed", + too_large: "too_large", + unavailable: "unavailable", +}; + +const KIND_BY_STATUS: Readonly> = { + 400: "invalid", + 409: "changed", + 413: "too_large", + 503: "unavailable", +}; + +const failureCode = (body: unknown): string | undefined => { + if (typeof body !== "object" || body === null || !("detail" in body)) return undefined; + const detail = body.detail; + if (typeof detail !== "object" || detail === null || !("code" in detail)) return undefined; + return typeof detail.code === "string" ? detail.code : undefined; +}; + +export function classifyTraceReadFailure(error: unknown): TraceReadFailure { + if (!(error instanceof ApiError)) { + return { kind: "unknown", message: error instanceof Error ? error.message : String(error), retryAfterMs: null }; + } + const code = failureCode(error.body); + const kind = (code && KIND_BY_CODE[code]) || KIND_BY_STATUS[error.status] || "unknown"; + return { kind, message: error.message, retryAfterMs: error.retryAfterMs }; +} + +/** Only an outage is worth retrying without the user; a changed or invalid cursor fails the same way every time. */ +export const isRetryableTraceRead = (failure: TraceReadFailure): boolean => + failure.kind === "unavailable" || failure.kind === "unknown"; + +export const TRACE_READ_AUTO_RETRIES = 2; +const DEFAULT_RETRY_DELAY_MS = 1_000; + +export const traceReadRetry = (failureCount: number, error: unknown): boolean => + failureCount < TRACE_READ_AUTO_RETRIES && classifyTraceReadFailure(error).kind === "unavailable"; + +export const traceReadRetryDelay = (_attempt: number, error: unknown): number => + classifyTraceReadFailure(error).retryAfterMs ?? DEFAULT_RETRY_DELAY_MS; diff --git a/ui/litellm-dashboard/src/lib/http/api.ts b/ui/litellm-dashboard/src/lib/http/api.ts index ce8e134fd28..a4fcc45d86b 100644 --- a/ui/litellm-dashboard/src/lib/http/api.ts +++ b/ui/litellm-dashboard/src/lib/http/api.ts @@ -1,7 +1,7 @@ import createFetchClient, { type Middleware } from "openapi-fetch"; import createQueryClient from "openapi-react-query"; import type { paths } from "./schema"; -import { ApiError, deriveErrorMessage } from "./client"; +import { ApiError, deriveErrorMessage, retryAfterMs } from "./client"; import { getAuthHeaderName, getAuthToken, getRequestBaseUrl, reportError } from "./runtime"; import { resolveRequestUrl } from "./resolveApiBase"; @@ -32,7 +32,7 @@ const middleware: Middleware = { message = raw || `HTTP ${response.status}`; } reportError(message); - throw new ApiError(message, response.status, body); + throw new ApiError(message, response.status, body, retryAfterMs(response.headers)); }, }; diff --git a/ui/litellm-dashboard/src/lib/http/client.ts b/ui/litellm-dashboard/src/lib/http/client.ts index f6a4f801213..0d39171ae6f 100644 --- a/ui/litellm-dashboard/src/lib/http/client.ts +++ b/ui/litellm-dashboard/src/lib/http/client.ts @@ -32,15 +32,28 @@ export interface RequestOptions { export class ApiError extends Error { readonly status: number; readonly body: unknown; + /** The server's `Retry-After` delay, when it sent one. */ + readonly retryAfterMs: number | null; - constructor(message: string, status: number, body: unknown) { + constructor(message: string, status: number, body: unknown, retryAfterMs: number | null = null) { super(message); this.name = "ApiError"; this.status = status; this.body = body; + this.retryAfterMs = retryAfterMs; } } +/** `Retry-After` as milliseconds; the header is whole seconds or an HTTP date. */ +export const retryAfterMs = (headers?: Headers): number | null => { + const header = headers?.get("retry-after"); + if (header === null || header === undefined) return null; + const seconds = Number(header); + if (Number.isFinite(seconds)) return Math.max(0, seconds * 1000); + const at = Date.parse(header); + return Number.isNaN(at) ? null : Math.max(0, at - Date.now()); +}; + /** * Best-effort extraction of a human-readable message from a proxy error body. * Lives here because error parsing is the client's job; networking.tsx re-exports @@ -175,7 +188,7 @@ export function createApiClient(config: ApiClientConfig): ApiClient { message = raw || `HTTP ${response.status}`; } onError?.(message); - throw new ApiError(message, response.status, errorBody); + throw new ApiError(message, response.status, errorBody, retryAfterMs(response.headers)); } return response;