diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 86df5cce1eb..f882a73a486 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4120,6 +4120,22 @@ dependencies = [ "time", ] +[[package]] +name = "litellm-pagination" +version = "0.1.0" +dependencies = [ + "base64 0.22.1", + "hmac 0.12.1", + "rstest", + "schemars 1.2.2", + "serde", + "serde_json", + "sha2 0.10.9", + "subtle", + "thiserror 2.0.19", + "time", +] + [[package]] name = "litellm-python-bridge" version = "0.1.0" @@ -4150,6 +4166,7 @@ dependencies = [ "litellm-http", "litellm-llms", "litellm-llms-types", + "litellm-pagination", "litellm-secrets", "litellm-secrets-aws", "litellm-secrets-types", @@ -4457,6 +4474,7 @@ dependencies = [ "criterion", "indexmap 2.14.0", "litellm-llms-types", + "litellm-pagination", "macro_rules_attribute", "opentelemetry-proto", "prost", @@ -4495,6 +4513,7 @@ dependencies = [ "jsonschema", "litellm-http", "litellm-migrate", + "litellm-pagination", "litellm-storage-clickhouse", "litellm-traces", "litellm-traces-cache", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index b0766f11e87..5edf793e84a 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -17,6 +17,7 @@ litellm-traces-cache = { path = "crates/traces-cache" } litellm-traces-clickhouse = { path = "crates/traces-clickhouse" } litellm-storage-clickhouse = { path = "crates/storage-clickhouse" } litellm-migrate = { path = "crates/migrate" } +litellm-pagination = { path = "crates/pagination" } litellm-migrate-macros = { path = "crates/migrate-macros" } litellm-core = { path = "crates/core" } litellm-gateway-mcp = { path = "crates/gateway-mcp" } diff --git a/litellm-rust/crates/pagination/AGENTS.md b/litellm-rust/crates/pagination/AGENTS.md new file mode 100644 index 00000000000..b09cbe75151 --- /dev/null +++ b/litellm-rust/crates/pagination/AGENTS.md @@ -0,0 +1,5 @@ +Own the shared page envelope, traversal metadata, opaque signed cursor codec and stable pagination failure codes +Depend on nothing domain specific: never trace types, ClickHouse, HTTP or Python +Keep cursor contents limited to a typed position, revision, publication instant, format version, key id and expiry; never credentials or raw filters +Bind every cursor to its resource, authorization scope and query through the signature, and expose revision mismatches as `traversal_changed` +Test the codec and page bounds through the public API under `tests/` diff --git a/litellm-rust/crates/pagination/Cargo.toml b/litellm-rust/crates/pagination/Cargo.toml new file mode 100644 index 00000000000..2a1542d8b0a --- /dev/null +++ b/litellm-rust/crates/pagination/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "litellm-pagination" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[features] +schema = ["dep:schemars"] + +[dependencies] +base64.workspace = true +hmac = "0.12.1" +schemars = { workspace = true, optional = true } +serde.workspace = true +serde_json = { workspace = true, features = ["raw_value"] } +sha2.workspace = true +subtle.workspace = true +thiserror.workspace = true +time = { workspace = true, features = ["formatting"] } + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/pagination/src/cursor.rs b/litellm-rust/crates/pagination/src/cursor.rs new file mode 100644 index 00000000000..7ccff2c2455 --- /dev/null +++ b/litellm-rust/crates/pagination/src/cursor.rs @@ -0,0 +1,208 @@ +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use hmac::{Hmac, Mac}; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use serde_json::value::RawValue; +use sha2::{Digest, Sha256}; +use subtle::ConstantTimeEq; +use time::OffsetDateTime; + +use crate::Error; + +const FORMAT_VERSION: u8 = 1; +const SIGNATURE_DOMAIN: &[u8] = b"litellm-pagination-cursor-v1\0"; + +struct Key { + id: String, + secret: Vec, +} + +/// The HMAC keys cursors are signed with. The first key signs; every key verifies, so a rotated +/// key keeps earlier cursors valid until they expire. +pub struct KeyRing { + keys: Vec, +} + +impl KeyRing { + pub fn new>(secrets: impl IntoIterator) -> Result { + let keys = secrets + .into_iter() + .map(|secret| { + let secret = secret.as_ref(); + if secret.is_empty() { + return Err(Error::InvalidKeys); + } + let digest = Sha256::new() + .chain_update(b"litellm-pagination-key\0") + .chain_update(secret) + .finalize(); + Ok(Key { + id: format!("{digest:x}")[..8].to_owned(), + secret: secret.to_vec(), + }) + }) + .collect::, _>>()?; + if keys.is_empty() { + return Err(Error::InvalidKeys); + } + Ok(Self { keys }) + } + + pub fn signing_key_id(&self) -> &str { + &self.keys[0].id + } + + fn sign(key: &Key, binding: &Binding, payload: &[u8]) -> Vec { + let mut mac = + Hmac::::new_from_slice(&key.secret).expect("HMAC accepts keys of any length"); + mac.update(SIGNATURE_DOMAIN); + mac.update(&binding.fingerprint()); + mac.update(payload); + mac.finalize().into_bytes().to_vec() + } + + /// Encodes `cursor` as an opaque token bound to `binding`. + pub fn encode( + &self, + binding: &Binding, + cursor: &Cursor

, + ) -> Result { + let payload = serde_json::to_vec(&Envelope { + version: FORMAT_VERSION, + key_id: &self.keys[0].id, + expires_unix: cursor.expires_at.unix_timestamp(), + published_ms: cursor.published_ms, + revision: &cursor.revision, + position: &cursor.position, + })?; + let tag = Self::sign(&self.keys[0], binding, &payload); + Ok(format!( + "{}.{}", + URL_SAFE_NO_PAD.encode(payload), + URL_SAFE_NO_PAD.encode(tag) + )) + } + + /// Verifies `token` against `binding` and `now`, returning the typed cursor it carries. + pub fn decode( + &self, + binding: &Binding, + token: &str, + now: OffsetDateTime, + ) -> Result, Error> { + let (payload, tag) = token.split_once('.').ok_or(Error::InvalidCursor)?; + let payload = URL_SAFE_NO_PAD + .decode(payload) + .map_err(|_| Error::InvalidCursor)?; + let tag = URL_SAFE_NO_PAD + .decode(tag) + .map_err(|_| Error::InvalidCursor)?; + let envelope: Envelope<'_, Box> = + serde_json::from_slice(&payload).map_err(|_| Error::InvalidCursor)?; + let key = self + .keys + .iter() + .find(|key| key.id == envelope.key_id) + .ok_or(Error::InvalidCursor)?; + if envelope.version != FORMAT_VERSION + || !bool::from(Self::sign(key, binding, &payload).ct_eq(&tag)) + { + return Err(Error::InvalidCursor); + } + if envelope.expires_unix <= now.unix_timestamp() { + return Err(Error::TraversalExpired); + } + Ok(Cursor { + position: serde_json::from_str(envelope.position.get()) + .map_err(|_| Error::InvalidCursor)?, + revision: envelope.revision.to_owned(), + published_ms: envelope.published_ms, + expires_at: OffsetDateTime::from_unix_timestamp(envelope.expires_unix) + .map_err(|_| Error::InvalidCursor)?, + }) + } +} + +#[derive(Serialize, Deserialize)] +struct Envelope<'a, P> { + #[serde(rename = "v")] + version: u8, + #[serde(rename = "k", borrow)] + key_id: &'a str, + #[serde(rename = "e")] + expires_unix: i64, + #[serde(rename = "p")] + published_ms: u64, + #[serde(rename = "r", borrow)] + revision: &'a str, + #[serde(rename = "c")] + position: P, +} + +/// What a cursor is valid for: one resource, one authorization scope and one query. Any of them +/// changing makes earlier cursors invalid instead of silently continuing a different traversal. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Binding { + resource: String, + scope: String, + query: String, +} + +impl Binding { + pub fn new( + resource: &str, + scope: &impl Serialize, + query: &impl Serialize, + ) -> Result { + Ok(Self { + resource: resource.to_owned(), + scope: digest(scope)?, + query: digest(query)?, + }) + } + + pub(crate) fn fingerprint(&self) -> Vec { + Sha256::new() + .chain_update(self.resource.as_bytes()) + .chain_update(b"\0") + .chain_update(self.scope.as_bytes()) + .chain_update(b"\0") + .chain_update(self.query.as_bytes()) + .finalize() + .to_vec() + } +} + +fn digest(value: &impl Serialize) -> Result { + Ok(format!("{:x}", Sha256::digest(serde_json::to_vec(value)?))) +} + +/// A typed continuation: where the next page starts, which revision of the data it continues, +/// when that data was published and until when the traversal can be continued. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Cursor

{ + pub position: P, + pub revision: String, + pub published_ms: u64, + pub expires_at: OffsetDateTime, +} + +impl

Cursor

{ + /// Fails with `traversal_changed` when the data behind the traversal no longer matches the + /// revision the cursor was issued for. + pub fn require_revision(&self, current: &str) -> Result<(), Error> { + if self.revision.as_bytes().ct_eq(current.as_bytes()).into() { + Ok(()) + } else { + Err(Error::TraversalChanged) + } + } + + pub fn advance(&self, position: Q) -> Cursor { + Cursor { + position, + revision: self.revision.clone(), + published_ms: self.published_ms, + expires_at: self.expires_at, + } + } +} diff --git a/litellm-rust/crates/pagination/src/error.rs b/litellm-rust/crates/pagination/src/error.rs new file mode 100644 index 00000000000..e713563ecf8 --- /dev/null +++ b/litellm-rust/crates/pagination/src/error.rs @@ -0,0 +1,94 @@ +use serde::{Deserialize, Serialize}; + +/// Stable, machine readable failure codes shared by every paginated read and every adapter. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum FailureCode { + InvalidRequest, + InvalidCursor, + TraversalExpired, + TraversalChanged, + ResourceTooLarge, + BudgetExceeded, + ViewNotReady, + Unavailable, + Busy, +} + +impl FailureCode { + pub const fn as_str(self) -> &'static str { + match self { + Self::InvalidRequest => "invalid_request", + Self::InvalidCursor => "invalid_cursor", + Self::TraversalExpired => "traversal_expired", + Self::TraversalChanged => "traversal_changed", + Self::ResourceTooLarge => "resource_too_large", + Self::BudgetExceeded => "budget_exceeded", + Self::ViewNotReady => "view_not_ready", + Self::Unavailable => "unavailable", + Self::Busy => "busy", + } + } + + /// Whether repeating the same request unchanged can succeed later. + pub const fn is_transient(self) -> bool { + matches!(self, Self::ViewNotReady | Self::Unavailable | Self::Busy) + } + + /// Whether the client must discard loaded pages and restart from the first page. + pub const fn restarts_traversal(self) -> bool { + matches!(self, Self::TraversalExpired | Self::TraversalChanged) + } +} + +/// A public read failure: the stable code plus a message safe to show to the caller. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct Failure { + pub code: FailureCode, + pub message: String, +} + +impl Failure { + pub fn new(code: FailureCode, message: impl Into) -> Self { + Self { + code, + message: message.into(), + } + } +} + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("Invalid or unrecognized cursor; restart from the first page")] + InvalidCursor, + #[error("This page of results has expired; restart from the first page")] + TraversalExpired, + #[error("Results changed while paging; refresh to continue")] + TraversalChanged, + #[error("One item exceeds the response size limit")] + ResourceTooLarge, + #[error("Cursor signing keys must be nonempty")] + InvalidKeys, + #[error("Cursor encoding failed")] + Encode(#[from] serde_json::Error), +} + +impl Error { + pub const fn code(&self) -> FailureCode { + match self { + Self::InvalidCursor => FailureCode::InvalidCursor, + Self::TraversalExpired => FailureCode::TraversalExpired, + Self::TraversalChanged => FailureCode::TraversalChanged, + Self::ResourceTooLarge => FailureCode::ResourceTooLarge, + Self::InvalidKeys | Self::Encode(_) => FailureCode::Unavailable, + } + } +} + +impl From<&Error> for Failure { + fn from(error: &Error) -> Self { + Self::new(error.code(), error.to_string()) + } +} diff --git a/litellm-rust/crates/pagination/src/lib.rs b/litellm-rust/crates/pagination/src/lib.rs new file mode 100644 index 00000000000..65fba590688 --- /dev/null +++ b/litellm-rust/crates/pagination/src/lib.rs @@ -0,0 +1,10 @@ +//! The pagination contract shared by every LiteLLM read: an opaque signed cursor codec, the +//! common page envelope with traversal metadata, and stable failure codes. + +mod cursor; +mod error; +mod page; + +pub use cursor::{Binding, Cursor, KeyRing}; +pub use error::{Error, Failure, FailureCode}; +pub use page::{Page, Traversal, rfc3339, rfc3339_ms}; diff --git a/litellm-rust/crates/pagination/src/page.rs b/litellm-rust/crates/pagination/src/page.rs new file mode 100644 index 00000000000..50c9babcbf6 --- /dev/null +++ b/litellm-rust/crates/pagination/src/page.rs @@ -0,0 +1,75 @@ +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use time::{OffsetDateTime, format_description::well_known::Rfc3339}; + +use crate::{Binding, Error}; + +/// The pinned traversal a page belongs to. Every page of one traversal shares `id`, sees the +/// data as published at `published_at`, and can be continued until `expires_at`. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct Traversal { + pub id: String, + pub published_at: String, + pub expires_at: String, +} + +impl Traversal { + pub fn new(binding: &Binding, published_ms: u64, expires_at: OffsetDateTime) -> Self { + let digest = Sha256::new() + .chain_update(b"litellm-traversal\0") + .chain_update(binding.fingerprint()) + .chain_update(published_ms.to_be_bytes()) + .finalize(); + Self { + id: format!("{digest:x}")[..32].to_owned(), + published_at: rfc3339_ms(published_ms), + expires_at: rfc3339(expires_at), + } + } +} + +/// One page of an ordered traversal: a complete ordered prefix of what remains, and either the +/// opaque cursor for the rest or `None` when the traversal is exhausted. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[cfg_attr(feature = "schema", schemars(rename = "{T}Page"))] +pub struct Page { + pub items: Vec, + pub next_cursor: Option, + pub traversal: Traversal, +} + +impl Page { + /// Shrinks the page to the longest prefix whose JSON fits in `max_bytes`, re-issuing the + /// continuation cursor after the last kept item. Page size stays an upper bound and the + /// shortened page remains a complete ordered prefix. + pub fn bounded( + mut self, + max_bytes: usize, + continue_after: impl Fn(&T) -> Result, + ) -> Result { + while serde_json::to_vec(&self)?.len() > max_bytes { + if self.items.len() <= 1 { + return Err(Error::ResourceTooLarge); + } + self.items.truncate(self.items.len() / 2); + self.next_cursor = self.items.last().map(&continue_after).transpose()?; + } + Ok(self) + } +} + +pub fn rfc3339(instant: OffsetDateTime) -> String { + instant + .replace_nanosecond(instant.millisecond() as u32 * 1_000_000) + .unwrap_or(instant) + .format(&Rfc3339) + .unwrap_or_default() +} + +pub fn rfc3339_ms(unix_ms: u64) -> String { + OffsetDateTime::from_unix_timestamp_nanos(i128::from(unix_ms) * 1_000_000) + .map(rfc3339) + .unwrap_or_default() +} diff --git a/litellm-rust/crates/pagination/tests/cursor.rs b/litellm-rust/crates/pagination/tests/cursor.rs new file mode 100644 index 00000000000..541295bd792 --- /dev/null +++ b/litellm-rust/crates/pagination/tests/cursor.rs @@ -0,0 +1,203 @@ +use litellm_pagination::{Binding, Cursor, Error, FailureCode, KeyRing, Page, Traversal}; +use rstest::{fixture, rstest}; +use serde::{Deserialize, Serialize}; +use time::{Duration, OffsetDateTime}; + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +struct Position { + start_ms: i64, + trace_ref: String, +} + +#[fixture] +fn now() -> OffsetDateTime { + OffsetDateTime::from_unix_timestamp(1_790_000_000).unwrap() +} + +#[fixture] +fn keys() -> KeyRing { + KeyRing::new(["primary-secret"]).unwrap() +} + +#[fixture] +fn binding() -> Binding { + Binding::new("traces:list", &("team-a", "user-1"), &(0_i64, 100_i64)).unwrap() +} + +fn cursor(now: OffsetDateTime) -> Cursor { + Cursor { + position: Position { + start_ms: 42, + trace_ref: "ref-42".into(), + }, + revision: "rev-1".into(), + published_ms: 1_790_000_000_000, + expires_at: now + Duration::minutes(30), + } +} + +#[rstest] +fn cursors_round_trip_their_typed_position_and_pins( + keys: KeyRing, + binding: Binding, + now: OffsetDateTime, +) { + let token = keys.encode(&binding, &cursor(now)).unwrap(); + assert!(!token.contains("ref-42") && !token.contains("team-a")); + let decoded: Cursor = keys.decode(&binding, &token, now).unwrap(); + assert_eq!(decoded, cursor(now)); + assert!(decoded.require_revision("rev-1").is_ok()); + assert!(matches!( + decoded.require_revision("rev-2"), + Err(Error::TraversalChanged) + )); +} + +#[rstest] +#[case::other_resource(Binding::new("traces:detail", &("team-a", "user-1"), &(0_i64, 100_i64)).unwrap())] +#[case::other_scope(Binding::new("traces:list", &("team-b", "user-1"), &(0_i64, 100_i64)).unwrap())] +#[case::other_query(Binding::new("traces:list", &("team-a", "user-1"), &(0_i64, 200_i64)).unwrap())] +fn cursors_do_not_continue_a_different_traversal( + keys: KeyRing, + binding: Binding, + now: OffsetDateTime, + #[case] other: Binding, +) { + let token = keys.encode(&binding, &cursor(now)).unwrap(); + assert!(matches!( + keys.decode::(&other, &token, now), + Err(Error::InvalidCursor) + )); +} + +#[rstest] +#[case::garbage("not-a-cursor")] +#[case::empty("")] +#[case::missing_tag("eyJ2IjoxfQ")] +#[case::bad_base64("!!!.!!!")] +fn malformed_tokens_are_invalid_cursors( + keys: KeyRing, + binding: Binding, + now: OffsetDateTime, + #[case] token: &str, +) { + assert!(matches!( + keys.decode::(&binding, token, now), + Err(Error::InvalidCursor) + )); +} + +#[rstest] +fn tampered_payloads_fail_verification(keys: KeyRing, binding: Binding, now: OffsetDateTime) { + let token = keys.encode(&binding, &cursor(now)).unwrap(); + let (payload, tag) = token.split_once('.').unwrap(); + let forged = format!("{}.{tag}", &payload[..payload.len() - 2]); + assert!(matches!( + keys.decode::(&binding, &forged, now), + Err(Error::InvalidCursor) + )); + let other_keys = KeyRing::new(["other-secret"]).unwrap(); + assert!(matches!( + other_keys.decode::(&binding, &token, now), + Err(Error::InvalidCursor) + )); +} + +#[rstest] +fn expired_cursors_report_traversal_expiry(keys: KeyRing, binding: Binding, now: OffsetDateTime) { + let token = keys.encode(&binding, &cursor(now)).unwrap(); + assert!( + keys.decode::(&binding, &token, now + Duration::minutes(29)) + .is_ok() + ); + assert!(matches!( + keys.decode::(&binding, &token, now + Duration::minutes(30)), + Err(Error::TraversalExpired) + )); +} + +#[rstest] +fn rotated_keys_keep_verifying_earlier_cursors(binding: Binding, now: OffsetDateTime) { + let old = KeyRing::new(["old-secret"]).unwrap(); + let token = old.encode(&binding, &cursor(now)).unwrap(); + let rotated = KeyRing::new(["new-secret", "old-secret"]).unwrap(); + assert_ne!(rotated.signing_key_id(), old.signing_key_id()); + assert!(rotated.decode::(&binding, &token, now).is_ok()); + let fresh = rotated.encode(&binding, &cursor(now)).unwrap(); + assert!(matches!( + old.decode::(&binding, &fresh, now), + Err(Error::InvalidCursor) + )); +} + +#[rstest] +#[case::no_keys(Vec::<&str>::new())] +#[case::empty_key(vec![""])] +fn key_rings_require_nonempty_secrets(#[case] secrets: Vec<&str>) { + assert!(matches!(KeyRing::new(secrets), Err(Error::InvalidKeys))); +} + +#[rstest] +fn traversal_identity_follows_binding_and_publication(binding: Binding, now: OffsetDateTime) { + let first = Traversal::new(&binding, 1_790_000_000_000, now); + let same = Traversal::new(&binding, 1_790_000_000_000, now + Duration::minutes(1)); + let later = Traversal::new(&binding, 1_790_000_000_001, now); + let other_scope = + Binding::new("traces:list", &("team-b", "user-1"), &(0_i64, 100_i64)).unwrap(); + assert_eq!(first.id, same.id); + assert_ne!(first.id, later.id); + assert_ne!( + first.id, + Traversal::new(&other_scope, 1_790_000_000_000, now).id + ); + assert_eq!(first.published_at, "2026-09-21T14:13:20Z"); + assert_eq!(first.expires_at, "2026-09-21T14:13:20Z"); + assert_eq!(later.published_at, "2026-09-21T14:13:20.001Z"); +} + +#[rstest] +fn bounded_pages_keep_a_complete_prefix_and_reissue_the_cursor( + binding: Binding, + now: OffsetDateTime, +) { + let page = Page { + items: (0..8).map(|index| format!("item-{index:02}")).collect(), + next_cursor: Some("after-07".into()), + traversal: Traversal::new(&binding, 0, now), + }; + let full_bytes = serde_json::to_vec(&page).unwrap().len(); + let bounded = page + .clone() + .bounded(full_bytes - 1, |last| Ok(format!("after-{}", &last[5..]))) + .unwrap(); + assert_eq!(bounded.items, page.items[..4]); + assert_eq!(bounded.next_cursor.as_deref(), Some("after-03")); + assert_eq!(bounded.traversal, page.traversal); + assert!( + page.clone() + .bounded(full_bytes, |_| Ok(String::new())) + .is_ok() + ); + assert!(matches!( + page.bounded(1, |_| Ok(String::new())), + Err(Error::ResourceTooLarge) + )); +} + +#[rstest] +#[case::cursor(Error::InvalidCursor, FailureCode::InvalidCursor, false, false)] +#[case::expired(Error::TraversalExpired, FailureCode::TraversalExpired, false, true)] +#[case::changed(Error::TraversalChanged, FailureCode::TraversalChanged, false, true)] +#[case::too_large(Error::ResourceTooLarge, FailureCode::ResourceTooLarge, false, false)] +#[case::keys(Error::InvalidKeys, FailureCode::Unavailable, true, false)] +fn errors_map_to_stable_codes( + #[case] error: Error, + #[case] code: FailureCode, + #[case] transient: bool, + #[case] restarts: bool, +) { + assert_eq!(error.code(), code); + assert_eq!(code.is_transient(), transient); + assert_eq!(code.restarts_traversal(), restarts); + assert_eq!(serde_json::to_value(code).unwrap(), code.as_str()); +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index c3a86009111..74402ac9215 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -66,6 +66,7 @@ url.workspace = true [dev-dependencies] tracing.workspace = true +litellm-pagination.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-secrets-aws.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 6659be5160e..8923b9505ce 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -45,7 +45,7 @@ mod _native { use crate::routes::token_counter::TokenCounter; #[pymodule_export] use crate::routes::traces::{ - NativeTraceConfig, NativeTraceStorage, trace_encode_error, trace_span_rows, + NativeTraceConfig, NativeTraceStorage, TraceReadError, trace_encode_error, trace_span_rows, }; #[cfg(feature = "huggingface")] #[pymodule_export] @@ -113,6 +113,7 @@ mod tests { "NativeDiagnosticProcessor", "NativeTraceConfig", "NativeTraceStorage", + "TraceReadError", "trace_encode_error", "trace_span_rows", "TokenCounter", diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index 77a98b3980a..c4da3686c1c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -5,11 +5,18 @@ use litellm_traces::{QueryScope, ReadQuery, Tenant, query::named::ReadAccessPara use litellm_traces_clickhouse::{Config, Error, InsertTable, Parameter, QueryReaders}; use prost::Message; use pyo3::{ - exceptions::{PyOverflowError, PyRuntimeError, PyValueError}, + exceptions::{PyException, PyOverflowError, PyRuntimeError, PyValueError}, prelude::*, types::PyBytes, }; +pyo3::create_exception!( + _native, + TraceReadError, + PyException, + "A trace read failed with a stable pagination failure code; args are (code, message)." +); + #[derive(Message)] struct OtlpErrorStatus { #[prost(int32, tag = "1")] @@ -40,9 +47,8 @@ fn map_error_ref(error: &Error) -> PyErr { | Error::ReadTooLarge => PyOverflowError::new_err(error.to_string()), Error::InvalidRow | Error::InvalidTable - | Error::InvalidCursor(_) + | Error::Pagination(_) | Error::AmbiguousTrace - | Error::TraceChanged | Error::Decode(_) | Error::InvalidSchema | Error::InvalidQuery @@ -76,6 +82,10 @@ fn map_error_ref(error: &Error) -> PyErr { } } +fn map_read_error(error: Error) -> PyErr { + TraceReadError::new_err((error.failure_code().as_str(), error.to_string())) +} + fn map_sql_error(error: Error) -> PyErr { match error { Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(400 | 404)) => { @@ -98,10 +108,17 @@ impl NativeTraceConfig { url: &str, retention_days: u32, max_attribute_value_bytes: usize, + cursor_keys: Vec, ) -> PyResult { Ok(Self { - inner: Config::new(database, url, retention_days, max_attribute_value_bytes) - .map_err(map_error)?, + inner: Config::new( + database, + url, + retention_days, + max_attribute_value_bytes, + cursor_keys, + ) + .map_err(map_error)?, }) } } @@ -216,12 +233,14 @@ impl NativeTraceStorage { ) -> PyResult> { let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); + let config = self.config.clone(); crate::execution::run_async( py, async move { litellm_traces_clickhouse::list_traces( &client, &connection, + config.cursor_keys(), &scope, start_ms, end_ms, @@ -230,11 +249,11 @@ impl NativeTraceStorage { ) .await }, - map_error, + map_read_error, ) } - #[pyo3(signature = (trace_id, scope, trace_ref, cursor=None, page_size=None))] + #[pyo3(signature = (trace_id, scope, trace_ref, cursor, page_size))] fn get_trace<'py>( &self, py: Python<'py>, @@ -242,38 +261,27 @@ impl NativeTraceStorage { #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams, trace_ref: String, cursor: Option, - page_size: Option, + page_size: u32, ) -> PyResult> { let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); + let config = self.config.clone(); crate::execution::run_async( py, async move { - 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 - } else if cursor.is_some() { - Err(Error::InvalidParameters) - } else { - litellm_traces_clickhouse::get_trace( - &client, - &connection, - &scope, - &trace_id, - &trace_ref, - ) - .await - } + litellm_traces_clickhouse::get_trace_page( + &client, + &connection, + config.cursor_keys(), + &scope, + &trace_id, + &trace_ref, + cursor.as_deref(), + page_size, + ) + .await }, - map_error, + map_read_error, ) } @@ -300,7 +308,7 @@ impl NativeTraceStorage { ) .await }, - map_error, + map_read_error, ) } @@ -316,12 +324,14 @@ impl NativeTraceStorage { ) -> PyResult> { let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); + let config = self.config.clone(); crate::execution::run_async( py, async move { litellm_traces_clickhouse::get_span_error( &client, &connection, + config.cursor_keys(), &scope, &trace_id, &span_id, @@ -330,7 +340,7 @@ impl NativeTraceStorage { ) .await }, - map_error, + map_read_error, ) } @@ -482,9 +492,15 @@ mod tests { #[rstest] #[case::decode_budget(Error::Decode(litellm_traces::Error::TooLarge), "OverflowError")] #[case::invalid_export(Error::Decode(litellm_traces::Error::InvalidPayload), "ValueError")] - #[case::cursor(Error::InvalidCursor("trace"), "ValueError")] + #[case::cursor( + Error::Pagination(litellm_pagination::Error::InvalidCursor), + "ValueError" + )] #[case::ambiguous(Error::AmbiguousTrace, "ValueError")] - #[case::changed_snapshot(Error::TraceChanged, "ValueError")] + #[case::changed_snapshot( + Error::Pagination(litellm_pagination::Error::TraversalChanged), + "ValueError" + )] #[case::read_budget(Error::ReadTooLarge, "OverflowError")] fn trace_read_and_ingest_failures_preserve_public_exception_types( #[case] error: Error, @@ -498,4 +514,39 @@ mod tests { ); }); } + + #[rstest] + #[case::cursor( + Error::Pagination(litellm_pagination::Error::InvalidCursor), + "invalid_cursor" + )] + #[case::changed( + Error::Pagination(litellm_pagination::Error::TraversalChanged), + "traversal_changed" + )] + #[case::expired( + Error::Pagination(litellm_pagination::Error::TraversalExpired), + "traversal_expired" + )] + #[case::read_budget(Error::ReadTooLarge, "resource_too_large")] + #[case::busy(Error::Busy, "busy")] + #[case::cached( + Error::Cached(std::sync::Arc::new(Error::ReadTooLarge)), + "resource_too_large" + )] + fn read_failures_carry_stable_codes(#[case] error: Error, #[case] code: &str) { + Python::initialize(); + Python::attach(|py| { + let message = error.to_string(); + let exception = map_read_error(error); + assert_eq!(exception.get_type(py).name().unwrap(), "TraceReadError"); + let args: (String, String) = exception + .value(py) + .getattr("args") + .unwrap() + .extract() + .unwrap(); + assert_eq!(args, (code.to_owned(), message)); + }); + } } diff --git a/litellm-rust/crates/traces-cache/src/lib.rs b/litellm-rust/crates/traces-cache/src/lib.rs index 7dfd3bf32a1..845be84d315 100644 --- a/litellm-rust/crates/traces-cache/src/lib.rs +++ b/litellm-rust/crates/traces-cache/src/lib.rs @@ -67,13 +67,7 @@ impl SnapshotCache { 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 version = format!("{:x}", Sha256::digest(&encoded)); let snapshot = Arc::new(Snapshot { trace, diff --git a/litellm-rust/crates/traces-clickhouse/Cargo.toml b/litellm-rust/crates/traces-clickhouse/Cargo.toml index 6b95eb149d3..532aea36cc5 100644 --- a/litellm-rust/crates/traces-clickhouse/Cargo.toml +++ b/litellm-rust/crates/traces-clickhouse/Cargo.toml @@ -20,6 +20,7 @@ itertools = "0.14.0" litellm-http.workspace = true litellm-migrate.workspace = true litellm-storage-clickhouse.workspace = true +litellm-pagination.workspace = true litellm-traces.workspace = true litellm-traces-cache.workspace = true moka.workspace = true diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0009_trace_rollup_ownership.sql b/litellm-rust/crates/traces-clickhouse/migrations/0009_trace_rollup_ownership.sql index fd696349cf5..f4445980378 100644 --- a/litellm-rust/crates/traces-clickhouse/migrations/0009_trace_rollup_ownership.sql +++ b/litellm-rust/crates/traces-clickhouse/migrations/0009_trace_rollup_ownership.sql @@ -1,3 +1,4 @@ ALTER TABLE {database}.agent_traces_by_key ADD COLUMN IF NOT EXISTS UserIds SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT [], - ADD COLUMN IF NOT EXISTS IdentifiedLlmCount SimpleAggregateFunction(sum, UInt64) DEFAULT 0 + ADD COLUMN IF NOT EXISTS IdentifiedLlmCount SimpleAggregateFunction(sum, UInt64) DEFAULT 0, + ADD COLUMN IF NOT EXISTS ReceivedMs SimpleAggregateFunction(min, UInt64) DEFAULT 0 diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0010_trace_cost_completeness.sql b/litellm-rust/crates/traces-clickhouse/migrations/0010_trace_cost_completeness.sql index af87a6bf40b..df70a50b348 100644 --- a/litellm-rust/crates/traces-clickhouse/migrations/0010_trace_cost_completeness.sql +++ b/litellm-rust/crates/traces-clickhouse/migrations/0010_trace_cost_completeness.sql @@ -15,6 +15,7 @@ SELECT countIf(StatusCode = 'STATUS_CODE_ERROR') AS ErrorCount, sum(InputTokens) AS InputTokens, sum(OutputTokens) AS OutputTokens, + min(EngineReceivedMs) AS ReceivedMs, groupUniqArrayIf(toString(Model), Model != '') AS Models, groupUniqArrayIf(SpanName, ObservationType = 'agent') AS AgentNames, groupArrayIf(LiteLLMRequestId, ObservationType = 'llm' OR LiteLLMRequestId != '') AS RequestIds diff --git a/litellm-rust/crates/traces-clickhouse/query/list_traces.sql b/litellm-rust/crates/traces-clickhouse/query/list_traces.sql index c52adf7ef49..918f39e8f0c 100644 --- a/litellm-rust/crates/traces-clickhouse/query/list_traces.sql +++ b/litellm-rust/crates/traces-clickhouse/query/list_traces.sql @@ -24,6 +24,7 @@ HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64}) AND min(StartTs) < fromUnixTimestamp64Milli({end_ms:Int64}) AND ({cursor_ms:Int64} = 0 OR (toUnixTimestamp64Milli(min(StartTs)), trace_ref) < ({cursor_ms:Int64}, {cursor_trace_id:String})) + AND ({snapshot_ms:UInt64} = 0 OR min(ReceivedMs) <= {snapshot_ms:UInt64}) ORDER BY start_ms DESC, trace_ref DESC LIMIT {limit:UInt32} ) @@ -41,6 +42,7 @@ LEFT JOIN ( AND Timestamp <= (SELECT max(trace_end) FROM page) AND TraceId IN (SELECT trace_id FROM page) AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page) + AND ({snapshot_ms:UInt64} = 0 OR EngineReceivedMs <= {snapshot_ms:UInt64}) GROUP BY TeamId, ApiKeyHash, TraceId ) AS identities ON page.team_id = identities.TeamId AND page.api_key_hash = identities.ApiKeyHash diff --git a/litellm-rust/crates/traces-clickhouse/src/config.rs b/litellm-rust/crates/traces-clickhouse/src/config.rs index dccd1c368f6..772d0fe5d17 100644 --- a/litellm-rust/crates/traces-clickhouse/src/config.rs +++ b/litellm-rust/crates/traces-clickhouse/src/config.rs @@ -1,4 +1,7 @@ +use std::sync::Arc; + use crate::Error; +use litellm_pagination::KeyRing; use litellm_storage_clickhouse::Storage; #[derive(Clone)] @@ -6,6 +9,7 @@ pub struct Config { storage: Storage, retention_days: u32, max_attribute_value_bytes: usize, + cursor_keys: Arc, } impl Config { @@ -14,12 +18,14 @@ impl Config { url: &str, retention_days: u32, max_attribute_value_bytes: usize, + cursor_keys: Vec, ) -> Result { super::schema_statements(&database, retention_days)?; Ok(Self { storage: Storage::new(database, url)?, retention_days, max_attribute_value_bytes, + cursor_keys: Arc::new(KeyRing::new(cursor_keys)?), }) } @@ -35,4 +41,9 @@ impl Config { pub fn max_attribute_value_bytes(&self) -> usize { self.max_attribute_value_bytes } + + /// Keys that sign read cursors; the first signs and every key verifies. + pub fn cursor_keys(&self) -> &KeyRing { + &self.cursor_keys + } } diff --git a/litellm-rust/crates/traces-clickhouse/src/error.rs b/litellm-rust/crates/traces-clickhouse/src/error.rs index 909ac5fc60f..9e30e856ccc 100644 --- a/litellm-rust/crates/traces-clickhouse/src/error.rs +++ b/litellm-rust/crates/traces-clickhouse/src/error.rs @@ -32,12 +32,10 @@ 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)] + Pagination(#[from] litellm_pagination::Error), #[error(transparent)] Decode(#[from] litellm_traces::Error), #[error("trace ingestion task failed")] @@ -56,3 +54,54 @@ impl From for Error { } } } + +impl Error { + /// The stable public failure code for a read that failed with this error. + pub fn failure_code(&self) -> litellm_pagination::FailureCode { + use litellm_pagination::FailureCode; + use litellm_storage_clickhouse::Error as StorageError; + + match self { + Self::Pagination(error) => error.code(), + Self::Cached(source) => source.failure_code(), + Self::ReadTooLarge + | Self::InsertTooLarge + | Self::Decode(litellm_traces::Error::TooLarge) + | Self::Storage(StorageError::InsertTooLarge | StorageError::ResponseTooLarge) => { + FailureCode::ResourceTooLarge + } + Self::Busy => FailureCode::Busy, + Self::InvalidRow + | Self::InvalidTable + | Self::InvalidSchema + | Self::InvalidQuery + | Self::InvalidParameters + | Self::InvalidScope + | Self::AmbiguousTrace + | Self::Decode(_) + | Self::Storage( + StorageError::InvalidRow + | StorageError::InvalidTable + | StorageError::InvalidSchema + | StorageError::EmptySql + | StorageError::InvalidParameters + | StorageError::InvalidQuery, + ) => FailureCode::InvalidRequest, + Self::InvalidResponse + | Self::SchemaFailed(_) + | Self::SchemaTransport + | Self::MissingSecret + | Self::ProvisionFailed(_) + | Self::ProvisionTransport + | Self::Task + | Self::Storage( + StorageError::InvalidUrl + | StorageError::QueryFailed(_) + | StorageError::InsertFailed(_) + | StorageError::SchemaFailed(_) + | StorageError::InvalidResponse + | StorageError::Transport, + ) => FailureCode::Unavailable, + } + } +} diff --git a/litellm-rust/crates/traces-clickhouse/src/query/named.rs b/litellm-rust/crates/traces-clickhouse/src/query/named.rs index cc912fbf6ae..4f54f427897 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query/named.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query/named.rs @@ -18,6 +18,8 @@ struct ListTracesParamsEncoding { pub cursor_trace_id: String, #[serde(deserialize_with = "super::number::deserialize")] pub limit: u32, + #[serde(default, deserialize_with = "super::number::deserialize")] + pub snapshot_ms: u64, } #[derive(Debug, Deserialize, Serialize)] @@ -365,7 +367,7 @@ mod tests { #[case::quoted(true)] fn parameters_preserve_flattened_multi_team_access(#[case] quoted: bool) { round_trip::( - json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": u32::MAX}), + json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": u32::MAX, "snapshot_ms": 0}), quoted, ); round_trip::( diff --git a/litellm-rust/crates/traces-clickhouse/src/reads.rs b/litellm-rust/crates/traces-clickhouse/src/reads.rs index 6938aab06b3..27c4e09b6dd 100644 --- a/litellm-rust/crates/traces-clickhouse/src/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/src/reads.rs @@ -1,12 +1,13 @@ //! Scoped trace reads: the trace list, one trace resolved with its spend, and span payloads. +use std::collections::BTreeMap; 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_pagination::{Binding, Cursor, KeyRing, Page, Traversal}; use litellm_storage_clickhouse::{Query, fetch}; use litellm_traces::{ SpanDetail, SpanErrorPage, SpendLookup, Trace, TracePage, listed_summary, @@ -14,6 +15,7 @@ use litellm_traces::{ }; use litellm_traces_cache::{SnapshotCache, SnapshotKey}; use serde::{Deserialize, Serialize}; +use time::OffsetDateTime; use crate::{ Connection, Error, @@ -47,52 +49,121 @@ static TRACE_SNAPSHOTS: LazyLock = LazyLock::new(|| { const NANOS_PER_MS: i64 = 1_000_000; const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000; const SPEND_CONCURRENCY: usize = 4; +const LIST_RESOURCE: &str = "traces:list"; +const DETAIL_RESOURCE: &str = "traces:detail"; +const DIAGNOSTIC_RESOURCE: &str = "traces:diagnostic"; +const LIST_LIMIT: u32 = 500; +const CURSOR_TTL: time::Duration = time::Duration::minutes(30); +const RESPONSE_BYTES: usize = litellm_storage_clickhouse::READ_LIMITS.response_bytes; -fn encode_cursor(position: &T) -> String { - URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default()) +fn now_ms(now: OffsetDateTime) -> u64 { + (now.unix_timestamp_nanos() / i128::from(NANOS_PER_MS)) as u64 } -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 scope_binding( + resource: &str, + access: &ReadAccessParams, + query: &impl Serialize, +) -> Result { + Ok(Binding::new( + resource, + &(&access.user_id, &access.team_ids, access.all_teams), + query, + )?) } -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")), +/// A continuation that starts a fresh traversal when no cursor was supplied, or verifies and +/// resumes the pinned one. +fn open_cursor Deserialize<'de>>( + keys: &KeyRing, + binding: &Binding, + cursor: Option<&str>, + first: P, + now: OffsetDateTime, +) -> Result, Error> { + match cursor.filter(|cursor| !cursor.is_empty()) { + None => Ok(Cursor { + position: first, + revision: String::new(), + published_ms: now_ms(now), + expires_at: now + CURSOR_TTL, + }), + Some(cursor) => Ok(keys.decode(binding, cursor, now)?), } } -#[derive(Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +struct ListPosition { + start_ms: i64, + trace_ref: String, +} + +/// Where a list traversal starts and which publication instant it is pinned to. +struct ListTraversal { + cursor: Cursor, + binding: Binding, +} + +impl ListTraversal { + fn open( + keys: &KeyRing, + access: &ReadAccessParams, + window: (i64, i64), + cursor: Option<&str>, + now: OffsetDateTime, + ) -> Result { + let binding = scope_binding(LIST_RESOURCE, access, &window)?; + let first = ListPosition { + start_ms: 0, + trace_ref: String::new(), + }; + let resumed = cursor.is_some_and(|cursor| !cursor.is_empty()); + let cursor = open_cursor(keys, &binding, cursor, first, now)?; + if resumed && (cursor.position.start_ms <= 0 || cursor.position.trace_ref.is_empty()) { + return Err(litellm_pagination::Error::InvalidCursor.into()); + } + Ok(Self { cursor, binding }) + } + + fn traversal(&self) -> Traversal { + Traversal::new( + &self.binding, + self.cursor.published_ms, + self.cursor.expires_at, + ) + } + + fn continue_after( + &self, + keys: &KeyRing, + starts: &BTreeMap, + last: &litellm_traces::TraceSummary, + ) -> Result { + let position = ListPosition { + start_ms: *starts + .get(&last.trace_ref) + .ok_or(litellm_pagination::Error::InvalidCursor)?, + trace_ref: last.trace_ref.clone(), + }; + keys.encode(&self.binding, &self.cursor.advance(position)) + } +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +struct SpanPosition { + offset: usize, +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] struct ErrorPosition { offset: u64, - version: String, } -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 +fn valid_error_version(version: &str) -> bool { + version.len() == 64 + && 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")); - } - Ok(Some(position)) + .all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte)) } /// The stored run a trace id names for this caller; ids can repeat across tenants and runs. @@ -154,26 +225,39 @@ async fn spend( } } +#[expect( + clippy::too_many_arguments, + reason = "one argument per public read parameter" +)] pub async fn list_traces( client: &Client, connection: &Connection, + keys: &KeyRing, access: &ReadAccessParams, start_ms: i64, end_ms: i64, cursor: Option<&str>, limit: u32, ) -> Result { - if limit == 0 { + if limit == 0 || start_ms >= end_ms { return Err(Error::InvalidParameters); } - let (cursor_ms, cursor_trace_id) = trace_position(cursor)?; + let traversal = ListTraversal::open( + keys, + access, + (start_ms, end_ms), + cursor, + OffsetDateTime::now_utc(), + )?; + let snapshot_ms = traversal.cursor.published_ms; let mut params = ListTracesParams::from(contracts::ListTracesParams { access: access.clone(), start_ms, end_ms, - cursor_ms, - cursor_trace_id, - limit: limit.min(500), + cursor_ms: traversal.cursor.position.start_ms, + cursor_trace_id: traversal.cursor.position.trace_ref.clone(), + limit: limit.min(LIST_LIMIT), + snapshot_ms, }); let page: Vec = loop { match fetch::(client, connection, ¶ms).await { @@ -186,18 +270,31 @@ pub async fn list_traces( 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)) + let exhausted = page.len() < params.0.limit as usize; + let starts: BTreeMap = page + .iter() + .map(|row| (row.trace_ref.clone(), row.start_ms)) + .collect(); + let items: Vec = stream::iter(page.chunks(16)) + .then(|batch| list_summaries(client, connection, access, batch, snapshot_ms)) .try_collect::>() .await? .into_iter() .flatten() .collect(); - Ok(TracePage { data, next_cursor }) + let next_cursor = items + .last() + .filter(|_| !exhausted) + .map(|last| traversal.continue_after(keys, &starts, last)) + .transpose()?; + let page = Page { + items, + next_cursor, + traversal: traversal.traversal(), + }; + Ok(page.bounded(RESPONSE_BYTES, |last| { + traversal.continue_after(keys, &starts, last) + })?) } async fn list_summaries( @@ -205,6 +302,7 @@ async fn list_summaries( connection: &Connection, access: &ReadAccessParams, runs: &[contracts::ListTracesRow], + snapshot_ms: u64, ) -> Result, Error> { let (Some(start_ms), Some(end_ms)) = ( runs.iter().map(|row| row.start_ms).min(), @@ -220,7 +318,9 @@ async fn list_summaries( start_ms, end_ms: end_ms.saturating_add(1), }); - let spans = match crate::span_batches::read_list_spans(client, connection, params).await { + let spans = match crate::span_batches::read_list_spans(client, connection, params, snapshot_ms) + .await + { Ok(spans) => spans, Err(Error::ReadTooLarge) => { return stream::iter(runs) @@ -293,17 +393,14 @@ pub async fn get_trace( 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, -} - +#[expect( + clippy::too_many_arguments, + reason = "one argument per public read parameter" +)] pub async fn get_trace_page( client: &Client, connection: &Connection, + keys: &KeyRing, access: &ReadAccessParams, trace_id: &str, trace_ref: &str, @@ -316,28 +413,20 @@ pub async fn get_trace_page( 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 - } - 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 binding = scope_binding(DETAIL_RESOURCE, access, &(trace_id, &trace_ref))?; + let position = open_cursor( + keys, + &binding, + cursor, + SpanPosition { offset: 0 }, + OffsetDateTime::now_utc(), + )?; let key = SnapshotKey::new( connection.url().as_str(), access, trace_id, &trace_ref, - position.snapshot_ms, + position.published_ms, ) .map_err(|_| Error::InvalidParameters)?; let snapshot = match TRACE_SNAPSHOTS.get(&key).await { @@ -349,7 +438,7 @@ pub async fn get_trace_page( trace_ref: trace_ref.clone(), }; let rows = - crate::span_batches::read_spans(client, connection, params, position.snapshot_ms) + crate::span_batches::read_spans(client, connection, params, position.published_ms) .await?; let spend_rows = spend(client, connection, access, &rows).await; let Some(trace) = resolve_trace(trace_id, &trace_ref, &rows, &spend_rows) else { @@ -358,46 +447,48 @@ pub async fn get_trace_page( TRACE_SNAPSHOTS.insert(key, trace).await? } }; - let spans = &snapshot.trace().spans; - if cursor.is_some() && position.version != snapshot.version() { - return Err(Error::TraceChanged); + if !position.revision.is_empty() { + position.require_revision(snapshot.version())?; } + let spans = &snapshot.trace().spans; + if position.position.offset > spans.len() { + return Err(litellm_pagination::Error::InvalidCursor.into()); + } + let pinned = Cursor { + position: position.position, + revision: snapshot.version().to_owned(), + published_ms: position.published_ms, + expires_at: position.expires_at, + }; + let continue_at = |offset: usize| -> Result, Error> { + (offset < spans.len()) + .then(|| keys.encode(&binding, &pinned.advance(SpanPosition { offset }))) + .transpose() + .map_err(Error::from) + }; + let start = pinned.position.offset; + let end = start.saturating_add(page_size as usize).min(spans.len()); let mut trace = Trace { summary: snapshot.trace().summary.clone(), agents: snapshot.trace().agents.clone(), - spans: Vec::new(), - next_cursor: None, + spans: spans[start..end].to_vec(), + next_cursor: continue_at(end)?, + traversal: Some(Traversal::new( + &binding, + pinned.published_ms, + pinned.expires_at, + )), }; - 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 + > 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(), - })); + trace.next_cursor = continue_at(start + trace.spans.len())?; } Ok(Some(trace)) } @@ -433,29 +524,49 @@ pub async fn get_span( })) } +#[expect( + clippy::too_many_arguments, + reason = "one argument per public read parameter" +)] pub async fn get_span_error( client: &Client, connection: &Connection, + keys: &KeyRing, 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 binding = scope_binding( + DIAGNOSTIC_RESOURCE, + access, + &(trace_id, span_id, &trace_ref), + )?; + let position = open_cursor( + keys, + &binding, + cursor, + ErrorPosition { offset: 0 }, + OffsetDateTime::now_utc(), + )?; + let resumed = !position.revision.is_empty(); + if resumed && !valid_error_version(&position.revision) + || i64::try_from(position.position.offset).is_err() + { + return Err(litellm_pagination::Error::InvalidCursor.into()); + } + let offset = 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(), + error_version: position.revision.clone(), }); let Some(row) = fetch::(client, connection, ¶ms) .await? @@ -465,13 +576,25 @@ pub async fn get_span_error( return Ok(None); }; let row = row.0; + if resumed { + position.require_revision(&row.version)?; + } 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, + let next_cursor = (next_offset < row.total_chars) + .then(|| { + keys.encode( + &binding, + &Cursor { + position: ErrorPosition { + offset: next_offset, + }, + revision: row.version, + published_ms: position.published_ms, + expires_at: position.expires_at, + }, + ) }) - }); + .transpose()?; Ok(Some(SpanErrorPage { span_id: row.span_id, message: row.message, @@ -486,51 +609,176 @@ mod tests { use super::*; + fn keys() -> KeyRing { + KeyRing::new(["list-secret"]).unwrap() + } + + fn access() -> ReadAccessParams { + ReadAccessParams { + all_teams: false, + user_id: "user-a".into(), + team_ids: vec!["team-a".into()], + } + } + + fn now() -> OffsetDateTime { + OffsetDateTime::from_unix_timestamp(1_790_000_000).unwrap() + } + #[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() + fn list_traversal_pins_publication_and_continues_after_the_last_run() { + let keys = keys(); + let first = ListTraversal::open(&keys, &access(), (0, 100), None, now()).unwrap(); + assert_eq!(first.cursor.published_ms, 1_790_000_000_000); + let last = listed_summary(&contracts::ListTracesRow { + trace_id: "t".into(), + trace_ref: "4bad42b84e9de3ba46fc870185f8f023".into(), + team_id: "team-a".into(), + api_key_hash: String::new(), + user_id: "user-a".into(), + name: String::new(), + service: String::new(), + input_preview: String::new(), + status: litellm_traces::SpanStatus::Ok, + start_ms: 1_790_742_989_377, + duration_ms: 1, + span_count: 1, + agent_count: 0, + agent_invocations: 0, + llm_calls: 0, + tool_calls: 0, + input_tokens: 0, + output_tokens: 0, + models: Vec::new(), + agent_names: Vec::new(), + frameworks: Vec::new(), + error_count: 0, + request_ids: Vec::new(), + }); + let starts = BTreeMap::from([(last.trace_ref.clone(), 1_790_742_989_377_i64)]); + let token = first.continue_after(&keys, &starts, &last).unwrap(); + let next = ListTraversal::open( + &keys, + &access(), + (0, 100), + Some(&token), + now() + time::Duration::minutes(5), + ) + .unwrap(); + assert_eq!(next.cursor.position.start_ms, 1_790_742_989_377); + assert_eq!(next.cursor.position.trace_ref, last.trace_ref); + assert_eq!(next.cursor.published_ms, first.cursor.published_ms); + assert_eq!(next.traversal(), first.traversal()); + } + + #[rstest] + #[case::other_window((0, 200))] + fn list_cursors_are_bound_to_their_window(#[case] window: (i64, i64)) { + let keys = keys(); + let first = ListTraversal::open(&keys, &access(), (0, 100), None, now()).unwrap(); + let token = keys + .encode( + &first.binding, + &first.cursor.advance(ListPosition { + start_ms: 5, + trace_ref: "r".into(), + }), ) - ); - 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) { + .unwrap(); assert!(matches!( - trace_position(Some(cursor)), - Err(Error::InvalidCursor("trace")) + ListTraversal::open(&keys, &access(), window, Some(&token), now()), + Err(Error::Pagination(litellm_pagination::Error::InvalidCursor)) + )); + let other_access = ReadAccessParams { + team_ids: vec!["team-b".into()], + ..access() + }; + assert!(matches!( + ListTraversal::open(&keys, &other_access, (0, 100), Some(&token), now()), + Err(Error::Pagination(litellm_pagination::Error::InvalidCursor)) + )); + assert!(matches!( + ListTraversal::open(&keys, &access(), (0, 100), Some(&token), now() + CURSOR_TTL), + Err(Error::Pagination( + litellm_pagination::Error::TraversalExpired + )) )); } #[rstest] - #[case::not_base64("garbage")] - #[case::missing_fields("e30=")] - #[case::not_an_object("WzEsMl0=")] + #[case::not_a_token("abc")] + #[case::legacy_base64_json("WzEsICJ0Il0=")] + fn malformed_list_cursors_are_invalid(#[case] cursor: &str) { + assert!(matches!( + ListTraversal::open(&keys(), &access(), (0, 100), Some(cursor), now()), + Err(Error::Pagination(litellm_pagination::Error::InvalidCursor)) + )); + } + + #[rstest] + #[case::not_a_token("garbage")] + #[case::legacy_base64_json("e30=")] fn malformed_diagnostic_cursors_are_rejected(#[case] cursor: &str) { + let binding = scope_binding(DIAGNOSTIC_RESOURCE, &access(), &("t", "s", "r")).unwrap(); assert!(matches!( - error_position(Some(cursor)), - Err(Error::InvalidCursor("diagnostic")) + open_cursor::( + &keys(), + &binding, + Some(cursor), + ErrorPosition { offset: 0 }, + now() + ), + Err(Error::Pagination(litellm_pagination::Error::InvalidCursor)) )); } #[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 }); + #[case::lowercase_hex("a".repeat(64), false)] + #[case::short("A".repeat(63), false)] + #[case::uppercase_hex("A".repeat(64), true)] + fn diagnostic_versions_are_uppercase_sha256_hex(#[case] version: String, #[case] valid: bool) { + assert_eq!(valid_error_version(&version), valid); + } + + #[rstest] + fn detail_cursors_are_bound_to_the_trace_and_its_content() { + let keys = keys(); + let binding = scope_binding(DETAIL_RESOURCE, &access(), &("t", "ref")).unwrap(); + let first = open_cursor(&keys, &binding, None, SpanPosition { offset: 0 }, now()).unwrap(); + let token = keys + .encode( + &binding, + &Cursor { + revision: "content-v1".into(), + ..first.advance(SpanPosition { offset: 200 }) + }, + ) + .unwrap(); + let resumed: Cursor = open_cursor( + &keys, + &binding, + Some(&token), + SpanPosition { offset: 0 }, + now(), + ) + .unwrap(); + assert_eq!(resumed.position.offset, 200); + assert_eq!(resumed.published_ms, first.published_ms); + assert!(resumed.require_revision("content-v1").is_ok()); assert!(matches!( - error_position(Some(&cursor)), - Err(Error::InvalidCursor("diagnostic")) + resumed.require_revision("content-v2"), + Err(litellm_pagination::Error::TraversalChanged) + )); + let other_run = scope_binding(DETAIL_RESOURCE, &access(), &("t", "other-ref")).unwrap(); + assert!(matches!( + open_cursor::( + &keys, + &other_run, + Some(&token), + SpanPosition { offset: 0 }, + now() + ), + Err(Error::Pagination(litellm_pagination::Error::InvalidCursor)) )); } } diff --git a/litellm-rust/crates/traces-clickhouse/src/span_batches.rs b/litellm-rust/crates/traces-clickhouse/src/span_batches.rs index d66ad9506da..25b2f44f89c 100644 --- a/litellm-rust/crates/traces-clickhouse/src/span_batches.rs +++ b/litellm-rust/crates/traces-clickhouse/src/span_batches.rs @@ -122,6 +122,7 @@ pub(crate) async fn read_list_spans( client: &Client, connection: &Connection, runs: crate::query::named::TracePageSpansParams, + snapshot_ms: u64, ) -> Result, Error> { let parameters = ListParameters { runs, @@ -130,7 +131,7 @@ pub(crate) async fn read_list_spans( 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()), diff --git a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs index c361152f0e8..232a062f4d7 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs @@ -1993,6 +1993,7 @@ async fn rollup_cost_completeness_preserves_missing_ids_and_fails_closed_for_his cursor_ms: 0, cursor_trace_id: "".into(), limit: 10, + snapshot_ms: 0, }, ); let reader = Connection::reader(&database.url, "trace_test")?; diff --git a/litellm-rust/crates/traces-clickhouse/tests/queries.rs b/litellm-rust/crates/traces-clickhouse/tests/queries.rs index d973b3a565d..24f9cd4a969 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/queries.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/queries.rs @@ -157,6 +157,7 @@ async fn typed_queries_read_normalized_spans_and_keep_trace_identities_separate( end_ms: i64::MAX / 1_000_000, cursor_ms: 0, cursor_trace_id: String::new(), + snapshot_ms: 0, limit: 10, }); let traces = fetch::(&fixture.database.client, &reader, ¶ms).await?; @@ -226,6 +227,7 @@ async fn typed_trace_cursor_returns_the_next_fixture_trace( end_ms: i64::MAX / 1_000_000, cursor_ms: 0, cursor_trace_id: String::new(), + snapshot_ms: 0, limit: 1, }); let first = fetch::(&fixture.database.client, &reader, ¶ms).await?; @@ -234,6 +236,7 @@ async fn typed_trace_cursor_returns_the_next_fixture_trace( let next_params = ListTracesParams::from(contracts::ListTracesParams { cursor_ms: first[0].0.start_ms, cursor_trace_id: first[0].0.trace_ref.clone(), + snapshot_ms: 0, ..params.0 }); let next = fetch::(&fixture.database.client, &reader, &next_params).await?; @@ -271,6 +274,7 @@ async fn captured_deeplite_exports_round_trip_through_clickhouse( end_ms: i64::MAX / 1_000_000, cursor_ms: 0, cursor_trace_id: String::new(), + snapshot_ms: 0, limit: 10, }); let traces = fetch::(&fixture.database.client, &reader, &list_params).await?; diff --git a/litellm-rust/crates/traces-clickhouse/tests/reads.rs b/litellm-rust/crates/traces-clickhouse/tests/reads.rs index 41820065d84..e8aa9386087 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/reads.rs @@ -14,6 +14,10 @@ mod support; use fixtures::{DATABASE, SeededDatabase, migrated_database, seeded_database}; use support::TestResult; +fn keys() -> litellm_pagination::KeyRing { + litellm_pagination::KeyRing::new(["fixture-cursor-secret"]).unwrap() +} + #[rstest] #[case::api_key("key-a", "")] #[case::user("", "user-a")] @@ -82,11 +86,21 @@ async fn list_costs_match_each_run_when_response_ids_are_reused( 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?; - assert_eq!(page.data.len(), runs.len()); + let page = list_traces( + client, + &reader, + &keys(), + &access, + 0, + 2_000_000_000_000, + None, + 50, + ) + .await?; + assert_eq!(page.items.len(), runs.len()); for (trace_id, _, cost) in runs { let summary = page - .data + .items .iter() .find(|summary| summary.trace_id == trace_id) .ok_or("missing run")?; @@ -199,10 +213,20 @@ async fn large_runs_remain_complete_under_default_reader_limits( 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?; - assert_eq!(page.data.len(), runs); + let page = list_traces( + client, + &reader, + &keys(), + &access, + 0, + 2_000_000_000_000, + None, + 500, + ) + .await?; + assert_eq!(page.items.len(), runs); assert!( - page.data + page.items .windows(2) .all(|runs| runs[0].trace_ref > runs[1].trace_ref) ); @@ -252,7 +276,7 @@ async fn large_runs_remain_complete_under_default_reader_limits( ); } } - for summary in &page.data { + for summary in &page.items { assert_eq!(summary.span_count, steps as u64); assert_eq!( if costed { @@ -267,7 +291,7 @@ async fn large_runs_remain_complete_under_default_reader_limits( } } let trace_ref = &page - .data + .items .iter() .find(|run| run.trace_id == "trace-0000") .ok_or("missing run")? @@ -304,6 +328,7 @@ async fn large_runs_remain_complete_under_default_reader_limits( let page = get_trace_page( client, &reader, + &keys(), &access, "trace-0000", trace_ref, @@ -319,19 +344,26 @@ async fn large_runs_remain_complete_under_default_reader_limits( <= litellm_storage_clickhouse::READ_LIMITS.response_bytes ); if ids.is_empty() { - assert!( - get_trace_page( - client, - &reader, - &denied, - "trace-0000", - trace_ref, - page.next_cursor.as_deref(), - 200, - ) - .await? - .is_none() - ); + let reused = get_trace_page( + client, + &reader, + &keys(), + &denied, + "trace-0000", + trace_ref, + page.next_cursor.as_deref(), + 200, + ) + .await; + match page.next_cursor { + Some(_) => assert!(matches!( + reused, + Err(litellm_traces_clickhouse::Error::Pagination( + litellm_pagination::Error::InvalidCursor + )) + )), + None => assert!(reused?.is_none()), + } client .post(writer.url().clone()) .body(format!("TRUNCATE TABLE {DATABASE}.otel_traces")) @@ -372,15 +404,26 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( 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 = list_traces( + client, + &reader, + &keys(), + &access, + 0, + 2_000_000_000_000, + None, + 10, + ) + .await?; let summary = listed - .data + .items .iter() .find(|summary| summary.span_count == 3) .ok_or("missing fixture")?; let first = get_trace_page( client, &reader, + &keys(), &access, &summary.trace_id, &summary.trace_ref, @@ -424,19 +467,22 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( user_id: String::new(), team_ids: vec!["not-this-team".into()], }; - assert!( + assert!(matches!( get_trace_page( client, &reader, + &keys(), &denied, &summary.trace_id, &summary.trace_ref, first.next_cursor.as_deref(), 1 ) - .await? - .is_none() - ); + .await, + Err(litellm_traces_clickhouse::Error::Pagination( + litellm_pagination::Error::InvalidCursor + )) + )); let first_cursor = first.next_cursor.clone(); let mut cursor = first.next_cursor; let mut ids = first @@ -448,6 +494,7 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( let next = get_trace_page( client, &reader, + &keys(), &access, &summary.trace_id, &summary.trace_ref, @@ -475,6 +522,7 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( get_trace_page( client, &reader, + &keys(), &access, &summary.trace_id, &summary.trace_ref, @@ -482,7 +530,9 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( 1 ) .await, - Err(litellm_traces_clickhouse::Error::InvalidCursor("span")) + Err(litellm_traces_clickhouse::Error::Pagination( + litellm_pagination::Error::InvalidCursor + )) )); let backdated = json!({ "Timestamp": "2026-09-01 00:00:00.000000000", @@ -505,6 +555,7 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( let changed = get_trace_page( client, &uncached_reader, + &keys(), &access, &summary.trace_id, &summary.trace_ref, @@ -513,7 +564,12 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( ) .await; assert!( - matches!(changed, Err(litellm_traces_clickhouse::Error::TraceChanged)), + matches!( + changed, + Err(litellm_traces_clickhouse::Error::Pagination( + litellm_pagination::Error::TraversalChanged + )) + ), "{changed:?}" ); Ok(()) @@ -535,9 +591,19 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( 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 = list_traces( + client, + &reader, + &keys(), + &access, + 0, + 2_000_000_000_000, + None, + 50, + ) + .await?; let run = before - .data + .items .iter() .find(|run| run.span_count == 3) .ok_or("missing fixture")?; @@ -562,10 +628,20 @@ 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?; - assert_eq!(after.data.len(), before.data.len()); + let after = list_traces( + client, + &reader, + &keys(), + &access, + 0, + 2_000_000_000_000, + None, + 50, + ) + .await?; + assert_eq!(after.items.len(), before.items.len()); let limited = after - .data + .items .iter() .find(|item| item.trace_ref == run.trace_ref) .ok_or("missing run")?; @@ -573,7 +649,7 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( assert_eq!(limited.span_count, 4); assert!( after - .data + .items .iter() .filter(|item| item.trace_ref != run.trace_ref) .all(|item| !item.resolution_limited) @@ -582,6 +658,7 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( get_trace_page( client, &reader, + &keys(), &access, &run.trace_id, &run.trace_ref, diff --git a/litellm-rust/crates/traces/Cargo.toml b/litellm-rust/crates/traces/Cargo.toml index e55fb841499..ee710632b22 100644 --- a/litellm-rust/crates/traces/Cargo.toml +++ b/litellm-rust/crates/traces/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [features] -schema = ["dep:schemars"] +schema = ["dep:schemars", "litellm-pagination/schema"] [dependencies] askama.workspace = true @@ -14,6 +14,7 @@ macro_rules_attribute.workspace = true schemars = { workspace = true, optional = true } indexmap = { version = "2", features = ["serde"] } litellm-llms-types.workspace = true +litellm-pagination.workspace = true opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] } prost.workspace = true serde = { workspace = true, features = ["rc"] } diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index 1e2eca3f7cc..dc13255d6bb 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -26,6 +26,7 @@ mod view; pub mod wire; pub use error::{Error, InvalidCallKey, InvalidQuery, InvalidScope}; +pub use litellm_pagination::{Failure, FailureCode, Page, Traversal}; pub use normalize::{ AgentMetadata, AgentType, CallEvidence, CallEvidenceKind, CallKey, Integration, NormalizedSpan, ObservationType, diff --git a/litellm-rust/crates/traces/src/query/named.rs b/litellm-rust/crates/traces/src/query/named.rs index 99069986f92..e9df28ee9c8 100644 --- a/litellm-rust/crates/traces/src/query/named.rs +++ b/litellm-rust/crates/traces/src/query/named.rs @@ -24,6 +24,8 @@ pub struct ListTracesParams { pub cursor_ms: i64, pub cursor_trace_id: String, pub limit: u32, + #[serde(default)] + pub snapshot_ms: u64, } #[derive(Debug, Deserialize, Serialize)] diff --git a/litellm-rust/crates/traces/src/resolve/view.rs b/litellm-rust/crates/traces/src/resolve/view.rs index 9c1c3756856..840c77e328f 100644 --- a/litellm-rust/crates/traces/src/resolve/view.rs +++ b/litellm-rust/crates/traces/src/resolve/view.rs @@ -211,6 +211,7 @@ pub fn resolve_trace( agents, spans, next_cursor: None, + traversal: None, }) } diff --git a/litellm-rust/crates/traces/src/schema.rs b/litellm-rust/crates/traces/src/schema.rs index cfa1d8e201e..af82c8ed9db 100644 --- a/litellm-rust/crates/traces/src/schema.rs +++ b/litellm-rust/crates/traces/src/schema.rs @@ -44,6 +44,12 @@ fn emitted() -> Schema { .into_root_schema_for::() } +fn emitted_as(title: &str) -> Schema { + let mut schema = emitted::(); + schema.insert("title".into(), title.into()); + schema +} + pub fn schemas() -> BTreeMap<&'static str, Schema> { BTreeMap::from([ ( @@ -52,9 +58,13 @@ pub fn schemas() -> BTreeMap<&'static str, Schema> { ), ("QueryScope", received::()), ("Tenant", received::()), - ("TracePage", emitted::()), + ("TracePage", emitted_as::("TracePage")), ("Trace", emitted::()), ("SpanDetail", emitted::()), ("SpanErrorPage", emitted::()), + ( + "ReadFailure", + emitted_as::("ReadFailure"), + ), ]) } diff --git a/litellm-rust/crates/traces/src/view.rs b/litellm-rust/crates/traces/src/view.rs index a864c740526..bbe83e5f82a 100644 --- a/litellm-rust/crates/traces/src/view.rs +++ b/litellm-rust/crates/traces/src/view.rs @@ -90,14 +90,11 @@ pub struct Trace { pub spans: Vec, #[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))] pub next_cursor: Option, + #[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))] + pub traversal: Option, } -#[macro_rules_attribute::apply(response_type)] -#[derive(Debug, PartialEq)] -pub struct TracePage { - pub data: Vec, - pub next_cursor: Option, -} +pub type TracePage = litellm_pagination::Page; #[macro_rules_attribute::apply(response_type)] #[derive(Debug, PartialEq)] diff --git a/litellm-rust/crates/traces/tests/query/named.rs b/litellm-rust/crates/traces/tests/query/named.rs index 4ccfc50740a..7db1a5895f7 100644 --- a/litellm-rust/crates/traces/tests/query/named.rs +++ b/litellm-rust/crates/traces/tests/query/named.rs @@ -32,8 +32,13 @@ fn named_requests_preserve_all_access_cases( ) }; round_trip::(request( - json!({"start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": 100}), + json!({"start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": 100, "snapshot_ms": u64::MAX}), )); + let unfenced: ListTracesParams = serde_json::from_value(request( + json!({"start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": 100}), + )) + .unwrap(); + assert_eq!(unfenced.snapshot_ms, 0); round_trip::(request(json!({"trace_id": "trace"}))); round_trip::(request(json!({"trace_id": "trace", "trace_ref": "ref"}))); round_trip::(request( diff --git a/litellm/constants.py b/litellm/constants.py index 49514fc4d0e..80e4c9091ba 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -58,6 +58,7 @@ OTLP_RETRY_AFTER_SECONDS: Final = get_env_int("OTLP_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) +AGENT_TRACING_DETAIL_PAGE_SIZE: Final = get_env_int("AGENT_TRACING_DETAIL_PAGE_SIZE", 200) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) DEFAULT_S3_MAX_CONCURRENT_UPLOADS: Final = int(os.getenv("DEFAULT_S3_MAX_CONCURRENT_UPLOADS", "16")) diff --git a/litellm/proxy/tracing_endpoints.py b/litellm/proxy/tracing_endpoints.py index 1563e4b5b55..9d298735bdc 100644 --- a/litellm/proxy/tracing_endpoints.py +++ b/litellm/proxy/tracing_endpoints.py @@ -19,7 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response from pydantic import BaseModel, ConfigDict from litellm._logging import verbose_proxy_logger -from litellm.constants import OTLP_RETRY_AFTER_SECONDS +from litellm.constants import AGENT_TRACING_DETAIL_PAGE_SIZE, OTLP_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 @@ -29,8 +29,10 @@ from litellm.proxy.tracing_runtime import provide_receiver, require_receiver from litellm.rust_bridge.trace.generated.models import TraceQueryHelp from litellm.rust_bridge.trace.generated.types import ( AllQueryScope, + FailureCode, OwnedQueryScope, QueryScope, + ReadFailure, SpanDetail, SpanErrorPage, Trace, @@ -38,13 +40,47 @@ from litellm.rust_bridge.trace.generated.types import ( TraceScope, ) from litellm.rust_bridge.trace.queries import TraceSQLResponse -from litellm.rust_bridge.trace.storage import ClickHouseStorage, Tenant +from litellm.rust_bridge.trace.storage import ClickHouseStorage, Tenant, TraceReadError from litellm.tracing import TraceReceiver, TracingPayloadTooLargeError from litellm.tracing.otlp_http import InvalidOTLPPayloadError, encode_otlp_response router = APIRouter(tags=["agent tracing"]) MS_PER_DAY: Final = 24 * 60 * 60 * 1000 +READ_FAILURE_STATUS: Final[Mapping[FailureCode, int]] = MappingProxyType( + { + "invalid_request": 400, + "invalid_cursor": 400, + "traversal_changed": 409, + "traversal_expired": 410, + "resource_too_large": 413, + "budget_exceeded": 413, + "view_not_ready": 503, + "unavailable": 503, + "busy": 503, + } +) +_RETRYABLE_READ_FAILURES: Final = frozenset({"view_not_ready", "unavailable", "busy"}) + + +class ReadFailureResponse(BaseModel): + model_config = ConfigDict(frozen=True) + detail: ReadFailure + + +def _read_failure_responses() -> dict[int | str, dict[str, object]]: + return {status: {"model": ReadFailureResponse} for status in (400, 409, 410, 413, 503)} + + +def _read_failure(error: TraceReadError) -> HTTPException: + if error.code in _RETRYABLE_READ_FAILURES: + verbose_proxy_logger.warning("Trace read %s: %s", error.code, error.message) + failure: Final[ReadFailure] = {"code": error.code, "message": error.message} + return HTTPException( + status_code=READ_FAILURE_STATUS[error.code], + detail=failure, + headers={"Retry-After": str(OTLP_RETRY_AFTER_SECONDS)} if error.code in _RETRYABLE_READ_FAILURES else None, + ) @dataclass(frozen=True, slots=True) @@ -135,7 +171,7 @@ async def ingest_otlp_traces( return Response(content=body, media_type=media_type) -@router.get("/v1/traces", response_model=TracePage) +@router.get("/v1/traces", response_model=TracePage, responses=_read_failure_responses()) async def list_agent_traces( context: Annotated[TraceAccessContext, Depends(provide_trace_access)], start_ms: Annotated[int | None, Query(description="Window start, unix ms. Default: 24h ago")] = None, @@ -151,12 +187,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 TraceReadError as error: + raise _read_failure(error) 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 @@ -230,23 +262,19 @@ async def help_agent_trace_queries( raise HTTPException(status_code=503, detail="Trace query help is temporarily unavailable") from error -@router.get("/v1/traces/{trace_id}", response_model=Trace) +@router.get("/v1/traces/{trace_id}", response_model=Trace, responses=_read_failure_responses()) async def get_agent_trace( trace_id: str, context: Annotated[TraceAccessContext, Depends(provide_trace_access)], trace_ref: Annotated[str, Query()] = "", cursor: Annotated[str | None, Query(max_length=512)] = None, - page_size: Annotated[int | None, Query(ge=1, le=500)] = None, + page_size: Annotated[int, Query(ge=1, le=500)] = AGENT_TRACING_DETAIL_PAGE_SIZE, ) -> 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 TraceReadError as error: + raise _read_failure(error) 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 @@ -255,7 +283,7 @@ async def get_agent_trace( return trace -@router.get("/v1/traces/{trace_id}/spans/{span_id}", response_model=SpanDetail) +@router.get("/v1/traces/{trace_id}/spans/{span_id}", response_model=SpanDetail, responses=_read_failure_responses()) async def get_agent_trace_span( trace_id: str, span_id: str, @@ -265,12 +293,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 TraceReadError as error: + raise _read_failure(error) 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 @@ -279,7 +303,9 @@ async def get_agent_trace_span( return span -@router.get("/v1/traces/{trace_id}/spans/{span_id}/error", response_model=SpanErrorPage) +@router.get( + "/v1/traces/{trace_id}/spans/{span_id}/error", response_model=SpanErrorPage, responses=_read_failure_responses() +) async def get_agent_trace_span_error( trace_id: str, span_id: str, @@ -290,12 +316,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 TraceReadError as error: + raise _read_failure(error) 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 diff --git a/litellm/rust_bridge/trace/generated/types.py b/litellm/rust_bridge/trace/generated/types.py index 127e86e9160..7521175db3c 100644 --- a/litellm/rust_bridge/trace/generated/types.py +++ b/litellm/rust_bridge/trace/generated/types.py @@ -23,6 +23,24 @@ class OwnedQueryScope(typing_extensions.TypedDict): QueryScope: TypeAlias = AllQueryScope | OwnedQueryScope +FailureCode: TypeAlias = Literal[ + "invalid_request", + "invalid_cursor", + "traversal_expired", + "traversal_changed", + "resource_too_large", + "budget_exceeded", + "view_not_ready", + "unavailable", + "busy", +] + + +class ReadFailure(typing_extensions.TypedDict): + code: ReadOnly[FailureCode] + message: ReadOnly[str] + + class UIText(typing_extensions.TypedDict): text: ReadOnly[str] kind: ReadOnly[Literal["text"]] @@ -77,6 +95,12 @@ SpanType: TypeAlias = Literal[ ] +class Traversal(typing_extensions.TypedDict): + id: ReadOnly[str] + published_at: ReadOnly[str] + expires_at: ReadOnly[str] + + class TraceScope(typing_extensions.TypedDict): all_teams: ReadOnly[Literal[0, 1]] user_id: ReadOnly[str] @@ -147,11 +171,13 @@ class Trace(typing_extensions.TypedDict): agents: ReadOnly[tuple[AgentNode, ...]] spans: ReadOnly[tuple[Span, ...]] next_cursor: ReadOnly[NotRequired[str | None]] + traversal: ReadOnly[NotRequired[Traversal | None | None]] class TracePage(typing_extensions.TypedDict): - data: ReadOnly[tuple[TraceSummary, ...]] + items: ReadOnly[tuple[TraceSummary, ...]] next_cursor: ReadOnly[str | None] + traversal: ReadOnly[Traversal] class UIMessages(typing_extensions.TypedDict): @@ -171,4 +197,6 @@ class SpanDetail(typing_extensions.TypedDict): attributes: ReadOnly[Mapping[str, str]] -TraceWireTypes: TypeAlias = QueryScope | SpanDetail | SpanErrorPage | Trace | TracePage | TraceScope | ReadQueryName +TraceWireTypes: TypeAlias = ( + QueryScope | ReadFailure | SpanDetail | SpanErrorPage | Trace | TracePage | TraceScope | ReadQueryName +) diff --git a/litellm/rust_bridge/trace/storage.py b/litellm/rust_bridge/trace/storage.py index a80e41a5aa3..4ae603768dd 100644 --- a/litellm/rust_bridge/trace/storage.py +++ b/litellm/rust_bridge/trace/storage.py @@ -4,7 +4,11 @@ from typing import Final, Protocol, TypeVar, runtime_checkable from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError -from litellm.constants import AGENT_TRACING_LIST_PAGE_SIZE, OTLP_MAX_ATTRIBUTE_VALUE_BYTES +from litellm.constants import ( + AGENT_TRACING_DETAIL_PAGE_SIZE, + AGENT_TRACING_LIST_PAGE_SIZE, + OTLP_MAX_ATTRIBUTE_VALUE_BYTES, +) from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.trace.generated.models import ( ActivityAvailability, @@ -31,6 +35,7 @@ from litellm.rust_bridge.trace.queries import ( from .generated.models import TraceQueryHelp from .generated.types import ( + FailureCode, QueryScope, SpanDetail, SpanErrorPage, @@ -68,7 +73,7 @@ class NativeStore(Protocol): ) -> Awaitable[JsonValue]: ... def get_trace( - self, trace_id: str, scope: TraceScope, trace_ref: str, cursor: str | None = None, page_size: int | None = None + self, trace_id: str, scope: TraceScope, trace_ref: str, cursor: str | None, page_size: int ) -> Awaitable[JsonValue]: ... def get_span(self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str) -> Awaitable[JsonValue]: ... @@ -88,6 +93,7 @@ class NativeStore(Protocol): @runtime_checkable class NativeTraces(Protocol): + TraceReadError: type[Exception] NativeTraceConfig: type["NativeConfig"] NativeTraceStorage: type[NativeStore] @@ -112,7 +118,14 @@ _NATIVE_ADAPTER: Final[TypeAdapter[NativeTraces]] = TypeAdapter( class NativeConfig(Protocol): - def __init__(self, database: str, url: str, retention_days: int, max_attribute_value_bytes: int) -> None: ... + def __init__( + self, + database: str, + url: str, + retention_days: int, + max_attribute_value_bytes: int, + cursor_keys: Sequence[str], + ) -> None: ... @dataclass(frozen=True, slots=True, repr=False) @@ -121,6 +134,27 @@ class TraceStorageConfig: database: str = "litellm" retention_days: int = 14 max_attribute_value_bytes: int = OTLP_MAX_ATTRIBUTE_VALUE_BYTES + cursor_keys: tuple[str, ...] = () + + +_READ_FAILURE: Final = TypeAdapter(tuple[FailureCode, str]) + + +class TraceReadError(Exception): + """A trace read failed with a stable pagination failure code; `code` drives the public response.""" + + def __init__(self, code: FailureCode, message: str) -> None: + super().__init__(code, message) + self.code: Final = code + self.message: Final = message + + @classmethod + def from_native(cls, error: BaseException) -> "TraceReadError": + try: + code, message = _READ_FAILURE.validate_python(error.args) + except ValidationError as invalid: + raise RuntimeError("Native trace read failed with an invalid failure") from invalid + return cls(code, message) def _native() -> NativeTraces: @@ -168,8 +202,17 @@ class ClickHouseStorage: config.url, config.retention_days, config.max_attribute_value_bytes, + config.cursor_keys, ) self._native: Final = native.NativeTraceStorage(validated) + self._read_error: Final = native.TraceReadError + + async def _read(self, adapter: TypeAdapter[_ResponseT], call: Awaitable[JsonValue]) -> _ResponseT: + try: + result: Final = await call + except self._read_error as error: + raise TraceReadError.from_native(error) from error + return _validate_query_response(adapter, result) async def ensure_schema(self) -> None: await self._native.ensure_schema() @@ -188,8 +231,7 @@ class ClickHouseStorage: cursor: str | None = None, limit: int = AGENT_TRACING_LIST_PAGE_SIZE, ) -> TracePage: - result: Final = await self._native.list_traces(scope, start_ms, end_ms, cursor, limit) - return _validate_query_response(_TRACE_PAGE, result) + return await self._read(_TRACE_PAGE, self._native.list_traces(scope, start_ms, end_ms, cursor, limit)) async def get_trace( self, @@ -197,20 +239,19 @@ class ClickHouseStorage: scope: TraceScope, trace_ref: str = "", cursor: str | None = None, - page_size: int | None = None, + page_size: int = AGENT_TRACING_DETAIL_PAGE_SIZE, ) -> Trace | None: - result: Final = await self._native.get_trace(trace_id, scope, trace_ref, cursor, page_size) - return _validate_query_response(_TRACE, result) + return await self._read(_TRACE, self._native.get_trace(trace_id, scope, trace_ref, cursor, page_size)) async def get_span(self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str = "") -> SpanDetail | None: - result: Final = await self._native.get_span(trace_id, span_id, scope, trace_ref) - return _validate_query_response(_SPAN_DETAIL, result) + return await self._read(_SPAN_DETAIL, self._native.get_span(trace_id, span_id, scope, trace_ref)) async def get_span_error( self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str = "", cursor: str | None = None ) -> SpanErrorPage | None: - result: Final = await self._native.get_span_error(trace_id, span_id, scope, trace_ref, cursor) - return _validate_query_response(_SPAN_ERROR_PAGE, result) + return await self._read( + _SPAN_ERROR_PAGE, self._native.get_span_error(trace_id, span_id, scope, trace_ref, cursor) + ) async def query(self, query: ReadQuery[ParamsT, RowT], parameters: ParamsT) -> tuple[RowT, ...]: validated: Final = query.parameters.model_validate(parameters) diff --git a/litellm/tracing/config.py b/litellm/tracing/config.py index 05b57dcd5cf..ef537ba2e67 100644 --- a/litellm/tracing/config.py +++ b/litellm/tracing/config.py @@ -2,12 +2,13 @@ import os from collections.abc import Mapping from typing import Final -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from litellm.constants import DEFAULT_AGENT_TRACING_RETENTION_DAYS, DEFAULT_CLICKHOUSE_DATABASE from litellm.rust_bridge.trace.storage import TraceStorageConfig STORE_SETTINGS: Final = TypeAdapter(dict[str, object]) +CURSOR_KEYS: Final = TypeAdapter(tuple[str, ...]) def is_clickhouse_tracing_enabled(settings: object) -> bool: @@ -46,6 +47,30 @@ def _retention_days(value: object) -> int: return days +def _cursor_key(candidate: str, environ: Mapping[str, str]) -> str: + resolved: Final = _value({"cursor_keys": candidate}, "cursor_keys", environ, None) + if not isinstance(resolved, str) or not resolved: + raise ValueError("tracing.store.cursor_keys must be a nonempty string or list of strings") + return resolved + + +def _cursor_keys(store: Mapping[str, object], environ: Mapping[str, str]) -> tuple[str, ...]: + supplied: Final = _value(store, "cursor_keys", environ, None) + if supplied is None: + fallback: Final = environ.get("LITELLM_SALT_KEY") or environ.get("LITELLM_MASTER_KEY") + if not fallback: + raise ValueError("tracing.store.cursor_keys, LITELLM_SALT_KEY or LITELLM_MASTER_KEY is required") + return (fallback,) + try: + candidates: Final = CURSOR_KEYS.validate_python((supplied,) if isinstance(supplied, str) else supplied) + except ValidationError as error: + raise ValueError("tracing.store.cursor_keys must be a nonempty string or list of strings") from error + keys: Final = tuple(_cursor_key(candidate, environ) for candidate in candidates) + if not keys: + raise ValueError("tracing.store.cursor_keys must be a nonempty string or list of strings") + return keys + + def _clickhouse_store(settings: Mapping[str, object]) -> Mapping[str, object]: raw_store: Final = settings.get("store") if raw_store is None: @@ -59,7 +84,7 @@ def _clickhouse_store(settings: Mapping[str, object]) -> Mapping[str, object]: def trace_storage_config(settings: Mapping[str, object], environ: Mapping[str, str] = os.environ) -> TraceStorageConfig: store: Final = _clickhouse_store(settings) - unknown: Final = store.keys() - {"type", "url", "database", "retention_days"} + unknown: Final = store.keys() - {"type", "url", "database", "retention_days", "cursor_keys"} if unknown: raise ValueError(f"unsupported tracing.store settings: {', '.join(sorted(unknown))}") url: Final = _value(store, "url", environ, environ.get("CLICKHOUSE_URL")) @@ -81,4 +106,5 @@ def trace_storage_config(settings: Mapping[str, object], environ: Mapping[str, s environ.get("AGENT_TRACING_RETENTION_DAYS", DEFAULT_AGENT_TRACING_RETENTION_DAYS), ) ), + cursor_keys=_cursor_keys(store, environ), ) diff --git a/litellm/tracing/receiver.py b/litellm/tracing/receiver.py index cf6bf8feb01..0c659aa94e4 100644 --- a/litellm/tracing/receiver.py +++ b/litellm/tracing/receiver.py @@ -18,7 +18,12 @@ from io import BytesIO from threading import BoundedSemaphore from typing import Final -from litellm.constants import AGENT_TRACING_LIST_PAGE_SIZE, OTLP_MAX_BODY_BYTES, OTLP_MAX_CONCURRENT_INGESTS +from litellm.constants import ( + AGENT_TRACING_DETAIL_PAGE_SIZE, + AGENT_TRACING_LIST_PAGE_SIZE, + OTLP_MAX_BODY_BYTES, + OTLP_MAX_CONCURRENT_INGESTS, +) from litellm.rust_bridge.trace.generated.types import SpanDetail, SpanErrorPage, Trace, TracePage, TraceScope from litellm.rust_bridge.trace.storage import ClickHouseStorage, Tenant from litellm.tracing.config import trace_storage_config @@ -105,7 +110,7 @@ class TraceReceiver: scope: TraceScope, trace_ref: str = "", cursor: str | None = None, - page_size: int | None = None, + page_size: int = AGENT_TRACING_DETAIL_PAGE_SIZE, ) -> Trace | None: return await self.storage.get_trace(trace_id, scope, trace_ref, cursor, page_size) diff --git a/scripts/trace_codegen/schemas/traces/ReadFailure.json b/scripts/trace_codegen/schemas/traces/ReadFailure.json new file mode 100644 index 00000000000..b5acff4c296 --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/ReadFailure.json @@ -0,0 +1,35 @@ +{ + "$defs": { + "FailureCode": { + "description": "Stable, machine readable failure codes shared by every paginated read and every adapter.", + "enum": [ + "invalid_request", + "invalid_cursor", + "traversal_expired", + "traversal_changed", + "resource_too_large", + "budget_exceeded", + "view_not_ready", + "unavailable", + "busy" + ], + "type": "string" + } + }, + "$schema": "https://json-schema.org/draft/2020-12/schema", + "description": "A public read failure: the stable code plus a message safe to show to the caller.", + "properties": { + "code": { + "$ref": "#/$defs/FailureCode" + }, + "message": { + "type": "string" + } + }, + "required": [ + "code", + "message" + ], + "title": "ReadFailure", + "type": "object" +} diff --git a/scripts/trace_codegen/schemas/traces/Trace.json b/scripts/trace_codegen/schemas/traces/Trace.json index dc66f77412c..a5cd26d6a77 100644 --- a/scripts/trace_codegen/schemas/traces/Trace.json +++ b/scripts/trace_codegen/schemas/traces/Trace.json @@ -309,6 +309,26 @@ "spend" ], "type": "object" + }, + "Traversal": { + "description": "The pinned traversal a page belongs to. Every page of one traversal shares `id`, sees the\ndata as published at `published_at`, and can be continued until `expires_at`.", + "properties": { + "expires_at": { + "type": "string" + }, + "id": { + "type": "string" + }, + "published_at": { + "type": "string" + } + }, + "required": [ + "id", + "published_at", + "expires_at" + ], + "type": "object" } }, "$schema": "https://json-schema.org/draft/2020-12/schema", @@ -334,13 +354,25 @@ }, "summary": { "$ref": "#/$defs/TraceSummary" + }, + "traversal": { + "anyOf": [ + { + "$ref": "#/$defs/Traversal" + }, + { + "type": "null" + } + ], + "x-python-optional": true } }, "required": [ "summary", "agents", "spans", - "next_cursor" + "next_cursor", + "traversal" ], "title": "Trace", "type": "object" diff --git a/scripts/trace_codegen/schemas/traces/TracePage.json b/scripts/trace_codegen/schemas/traces/TracePage.json index dfa02ba0d39..b646cfd81da 100644 --- a/scripts/trace_codegen/schemas/traces/TracePage.json +++ b/scripts/trace_codegen/schemas/traces/TracePage.json @@ -140,11 +140,32 @@ "spend" ], "type": "object" + }, + "Traversal": { + "description": "The pinned traversal a page belongs to. Every page of one traversal shares `id`, sees the\ndata as published at `published_at`, and can be continued until `expires_at`.", + "properties": { + "expires_at": { + "type": "string" + }, + "id": { + "type": "string" + }, + "published_at": { + "type": "string" + } + }, + "required": [ + "id", + "published_at", + "expires_at" + ], + "type": "object" } }, "$schema": "https://json-schema.org/draft/2020-12/schema", + "description": "One page of an ordered traversal: a complete ordered prefix of what remains, and either the\nopaque cursor for the rest or `None` when the traversal is exhausted.", "properties": { - "data": { + "items": { "items": { "$ref": "#/$defs/TraceSummary" }, @@ -155,11 +176,15 @@ "string", "null" ] + }, + "traversal": { + "$ref": "#/$defs/Traversal" } }, "required": [ - "data", - "next_cursor" + "items", + "next_cursor", + "traversal" ], "title": "TracePage", "type": "object" diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index ffd59a3f034..bb4e80f750f 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -38,6 +38,7 @@ from tests.test_litellm_rust.support.recording_server import RecordingServer, Re pytestmark = pytest.mark.requires_rust_extension QUERY_ROWS: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]) +CURSOR_KEYS: Final = ("test-cursor-key",) class CapturedSpendRow(BaseModel): @@ -54,7 +55,7 @@ class CapturedSpendQuery(BaseModel): def _native_storage(database: str, url: str, retention_days: int = 14) -> NativeTraceStorage: - return NativeTraceStorage(NativeTraceConfig(database, url, retention_days, OTLP_MAX_ATTRIBUTE_VALUE_BYTES)) + return NativeTraceStorage(NativeTraceConfig(database, url, retention_days, OTLP_MAX_ATTRIBUTE_VALUE_BYTES, CURSOR_KEYS)) @pytest.fixture @@ -129,18 +130,18 @@ async def test_reader_rejects_arbitrary_sql_before_sending(recording_server: Rec @pytest.mark.asyncio async def test_schema_binding_rejects_invalid_database() -> None: with pytest.raises(ValueError, match=r"database.*retention"): - NativeTraceConfig("db; DROP DATABASE default", "http://localhost:8123", 14, OTLP_MAX_ATTRIBUTE_VALUE_BYTES) + NativeTraceConfig("db; DROP DATABASE default", "http://localhost:8123", 14, OTLP_MAX_ATTRIBUTE_VALUE_BYTES, CURSOR_KEYS) @pytest.mark.asyncio async def test_schema_binding_rejects_non_positive_retention() -> None: with pytest.raises(ValueError, match=r"database.*retention"): - NativeTraceConfig("traces", "http://localhost:8123", 0, OTLP_MAX_ATTRIBUTE_VALUE_BYTES) + NativeTraceConfig("traces", "http://localhost:8123", 0, OTLP_MAX_ATTRIBUTE_VALUE_BYTES, CURSOR_KEYS) def test_invalid_url_error_does_not_expose_credentials() -> None: with pytest.raises(RuntimeError, match="invalid ClickHouse HTTP URL") as error: - NativeTraceConfig("traces", "secret://writer:password@example.com", 7, OTLP_MAX_ATTRIBUTE_VALUE_BYTES) + NativeTraceConfig("traces", "secret://writer:password@example.com", 7, OTLP_MAX_ATTRIBUTE_VALUE_BYTES, CURSOR_KEYS) assert "password" not in str(error.value) @@ -150,10 +151,12 @@ async def test_from_env_reads_with_clickhouse_url( ) -> None: recording_server.enqueue(ResponseSpec(body={"data": []})) monkeypatch.setenv("CLICKHOUSE_URL", recording_server.base_url) + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt") monkeypatch.delenv("CLICKHOUSE_READER_URL", raising=False) scope: Final[TraceScope] = {"all_teams": 1, "user_id": "", "team_ids": ()} page: Final = await TraceReceiver.from_env().list_traces(scope, 0, 1) - assert page == {"data": (), "next_cursor": None} + assert (page["items"], page["next_cursor"]) == ((), None) + assert page["traversal"]["published_at"] <= page["traversal"]["expires_at"] assert len(recording_server.requests) == 1 @@ -239,7 +242,7 @@ def _resource_export(attribute_bytes: int, span_count: int, groups: int = 1) -> @pytest.mark.asyncio async def test_resource_fanout_reaches_insert_with_identical_values(recording_server: RecordingServer) -> None: body: Final = _resource_export(16 * 1024, 1024) - receiver: Final = TraceReceiver(ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test"))) + receiver: Final = TraceReceiver(ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test", cursor_keys=CURSOR_KEYS))) tenant: Final = Tenant("team-a", "key-a", "org-a") assert await receiver.ingest(body, "application/json", None, tenant) == 1024 encoded: Final = gzip.decompress(recording_server.requests[0].raw_body) @@ -256,7 +259,7 @@ async def test_resource_fanout_reaches_insert_with_identical_values(recording_se async def test_shared_resource_still_hits_insert_limit_before_transport(recording_server: RecordingServer) -> None: recording_server.expected_requests = 0 body: Final = _resource_export(64 * 1024, 1024) - receiver: Final = TraceReceiver(ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test"))) + receiver: Final = TraceReceiver(ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test", cursor_keys=CURSOR_KEYS))) with pytest.raises(TracingPayloadTooLargeError, match="encoded size limit"): await receiver.ingest(body, "application/json", None, Tenant("team-a", "key-a")) assert recording_server.requests == [] @@ -264,7 +267,7 @@ async def test_shared_resource_still_hits_insert_limit_before_transport(recordin @pytest.mark.asyncio async def test_insert_validates_values_without_pydantic_copy(recording_server: RecordingServer) -> None: - storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test")) + storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test", cursor_keys=CURSOR_KEYS)) invalid: Final = object() with pytest.raises(ValueError, match=type(invalid).__name__): await storage.insert_rows("otel_traces", [{"ResourceAttributes": invalid}]) @@ -310,7 +313,7 @@ def test_trace_sql_endpoint_enforces_ownership_and_preserves_clickhouse_envelope for _ in range(11): recording_server.enqueue(ResponseSpec(body="")) recording_server.enqueue(ResponseSpec(body=envelope)) - storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test")) + storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test", cursor_keys=CURSOR_KEYS)) app: Final = FastAPI() app.include_router(router) app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret" @@ -361,7 +364,7 @@ def test_trace_help_endpoint_runs_native_schema_and_metadata_discovery( recording_server.enqueue(metadata) recording_server.enqueue(ResponseSpec(body={"data": [{"key": "custom.span"}]})) recording_server.enqueue(ResponseSpec(body={"data": [{"key": "custom.resource"}]})) - storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test")) + storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test", cursor_keys=CURSOR_KEYS)) app: Final = FastAPI() app.include_router(router) app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret" @@ -418,7 +421,7 @@ def test_trace_sql_endpoint_distinguishes_query_errors_from_reader_failures( "statistics": {"elapsed": 0.01, "rows_read": 1, "bytes_read": 1}, } recording_server.enqueue(ResponseSpec(body=envelope)) - storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test")) + storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test", cursor_keys=CURSOR_KEYS)) app: Final = FastAPI() app.include_router(router) app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret" @@ -441,6 +444,7 @@ async def test_trace_receiver_reads_with_only_one_clickhouse_url( span_params: dict[str, str | int | list[str]], ) -> None: monkeypatch.setenv("CLICKHOUSE_URL", recording_server.base_url) + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt") monkeypatch.setenv("CLICKHOUSE_DATABASE", "trace_test") monkeypatch.delenv("CLICKHOUSE_READER_URL", raising=False) recording_server.enqueue(ResponseSpec(body={"data": [span_row]})) @@ -459,7 +463,7 @@ async def test_lens_read_uses_the_shared_native_query_and_returns_typed_rows( recording_server: RecordingServer, ) -> None: recording_server.enqueue(ResponseSpec(body={"data": [{"traces": 0, "requests": 1}]})) - storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test")) + storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test", cursor_keys=CURSOR_KEYS)) rows: Final = await storage.lens_availability(LensAccessParams(all_teams=0, team="team-a", key_hash="key-a")) assert rows == (ActivityAvailability(traces=False, requests=True),) parameters: Final = parse_qs(urlsplit(recording_server.requests[0].path).query) @@ -512,7 +516,7 @@ def _fixture_trace_api( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.tracing_endpoints import provide_receiver, provide_trace_query_secret, router - storage: Final = ClickHouseStorage(TraceStorageConfig(clickhouse_url, "trace_test")) + storage: Final = ClickHouseStorage(TraceStorageConfig(clickhouse_url, "trace_test", cursor_keys=CURSOR_KEYS)) app: Final = FastAPI() app.include_router(router) app.dependency_overrides[provide_trace_query_secret] = lambda: "fixture-secret" diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index cff67d1f57d..e0712c46817 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -2,7 +2,7 @@ Tests for the agent tracing endpoints (litellm/proxy/tracing_endpoints.py). """ -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Mapping, Sequence from contextlib import asynccontextmanager from types import ModuleType from typing import Final, Literal @@ -12,6 +12,7 @@ import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from litellm.constants import AGENT_TRACING_DETAIL_PAGE_SIZE from litellm.proxy import tracing_endpoints from litellm.proxy._types import LitellmUserRoles, ProxyLifespanState, UserAPIKeyAuth from litellm.proxy.auth.authorization import OwnedRows, ReadScope @@ -22,7 +23,7 @@ from litellm.rust_bridge import loader 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 -from litellm.rust_bridge.trace.storage import ClickHouseStorage, TraceStorageConfig +from litellm.rust_bridge.trace.storage import ClickHouseStorage, TraceReadError, TraceStorageConfig from litellm.tracing import Tenant, TraceReceiver, TracingPayloadTooLargeError SQL_ENVELOPE: Final = { @@ -87,6 +88,8 @@ TRACE_RESPONSE: Final = { "agents": [], "spans": [], } +TRAVERSAL: Final = {"id": "traversal", "published_at": "2026-01-01T00:00:00Z", "expires_at": "2026-01-01T00:30:00Z"} +TRACE_PAGE_RESPONSE: Final = {"items": [], "next_cursor": None, "traversal": TRAVERSAL} SPAN_DETAIL_RESPONSE: Final = { "span_id": "s1", "input": "", @@ -162,7 +165,7 @@ def test_trace_read_and_write_permissions( def receiver(client) -> MagicMock: fake = MagicMock() fake.ingest = AsyncMock(return_value=1) - fake.list_traces = AsyncMock(return_value={"data": [], "next_cursor": None}) + fake.list_traces = AsyncMock(return_value=TRACE_PAGE_RESPONSE) fake.get_trace = AsyncMock(return_value=None) fake.get_span = AsyncMock(return_value=None) client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: fake @@ -244,7 +247,7 @@ def test_post_too_large_is_413(client, receiver): def test_list_traces_passes_scope_window_and_cursor(client, receiver): response = client.get("/v1/traces", params={"start_ms": 1, "end_ms": 2, "cursor": "abc"}) assert response.status_code == 200 - assert response.json() == {"data": [], "next_cursor": None} + assert response.json() == TRACE_PAGE_RESPONSE receiver.list_traces.assert_awaited_once_with( scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, start_ms=1, @@ -266,7 +269,9 @@ def test_get_trace_404_and_200(client, receiver): response = client.get("/v1/traces/t1") assert response.status_code == 200 assert response.json() == TRACE_RESPONSE - receiver.get_trace.assert_awaited_with("t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "", None, None) + receiver.get_trace.assert_awaited_with( + "t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "", None, AGENT_TRACING_DETAIL_PAGE_SIZE + ) def test_get_span_404_and_200(client, receiver): @@ -278,7 +283,10 @@ def test_get_span_404_and_200(client, receiver): receiver.get_span.assert_awaited_with("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "") -@pytest.mark.parametrize("suffix,cursor,page_size", [("", None, None), ("&cursor=next&page_size=200", "next", 200)]) +@pytest.mark.parametrize( + "suffix,cursor,page_size", + [("", None, AGENT_TRACING_DETAIL_PAGE_SIZE), ("&cursor=next&page_size=25", "next", 25)], +) def test_trace_detail_passes_scoped_reference(client, receiver, suffix, cursor, page_size): receiver.get_trace.return_value = TRACE_RESPONSE assert client.get(f"/v1/traces/t1?trace_ref=run-one{suffix}").status_code == 200 @@ -297,19 +305,39 @@ def test_trace_detail_passes_scoped_reference(client, receiver, suffix, cursor, ), ) @pytest.mark.parametrize( - "error,status,message", + "error,status,detail,retry_after", ( - (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, + "Traces are temporarily unavailable. Please try again.", + False, + ), + (TraceReadError("invalid_cursor", "restart"), 400, {"code": "invalid_cursor", "message": "restart"}, False), + (TraceReadError("traversal_changed", "refresh"), 409, {"code": "traversal_changed", "message": "refresh"}, False), + (TraceReadError("traversal_expired", "expired"), 410, {"code": "traversal_expired", "message": "expired"}, False), + (TraceReadError("resource_too_large", "big"), 413, {"code": "resource_too_large", "message": "big"}, False), + (TraceReadError("budget_exceeded", "budget"), 413, {"code": "budget_exceeded", "message": "budget"}, False), + (TraceReadError("busy", "later"), 503, {"code": "busy", "message": "later"}, True), + (TraceReadError("unavailable", "down"), 503, {"code": "unavailable", "message": "down"}, True), + (TraceReadError("view_not_ready", "soon"), 503, {"code": "view_not_ready", "message": "soon"}, True), ), ) -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_map_stable_codes_without_exposing_database_details( + client: TestClient, + receiver: MagicMock, + path: str, + method: str, + error: Exception, + status: int, + detail: object, + retry_after: bool, ) -> 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": detail} + assert ("Retry-After" in response.headers) is retry_after @pytest.mark.parametrize("query", ("page_size=0", "page_size=501", "cursor=" + "x" * 513)) @@ -324,7 +352,7 @@ def test_invalid_export_and_cursor_are_client_errors(client, receiver): receiver.ingest.side_effect = InvalidOTLPPayloadError("invalid OTLP trace payload") assert client.post("/v1/traces", content=b"broken").status_code == 400 - receiver.list_traces.side_effect = ValueError("Invalid trace cursor") + receiver.list_traces.side_effect = TraceReadError("invalid_cursor", "Invalid trace cursor") assert client.get("/v1/traces?cursor=broken").status_code == 400 @@ -722,10 +750,16 @@ def test_trace_storage_permissions_map_owned_rows( class _NativeConfig: - def __init__(self, database: str, url: str, retention_days: int, max_attribute_value_bytes: int) -> None: + def __init__( + self, database: str, url: str, retention_days: int, max_attribute_value_bytes: int, cursor_keys: Sequence[str] + ) -> None: pass +class _NativeReadError(Exception): + pass + + class _NativeReturningHelp(ModuleType): def __init__(self, help_payload: Mapping[str, object], trace_payload: Mapping[str, object] | None = None) -> None: super().__init__("native_traces") @@ -740,15 +774,16 @@ class _NativeReturningHelp(ModuleType): get_trace = AsyncMock(return_value=trace_payload) self.trace_read: Final = Storage.get_trace + self.TraceReadError: Final = _NativeReadError self.NativeTraceConfig: Final = _NativeConfig self.NativeTraceStorage: Final = Storage self.trace_encode_error: Final = bytes self.trace_span_rows: Final = list -@pytest.mark.parametrize("cursor,page_size", ((None, None), ("next", 200))) +@pytest.mark.parametrize("cursor,page_size", ((None, AGENT_TRACING_DETAIL_PAGE_SIZE), ("next", 200))) async def test_storage_preserves_page_cursor_and_normalizes_native_trace_data( - monkeypatch: pytest.MonkeyPatch, cursor: str | None, page_size: int | None + monkeypatch: pytest.MonkeyPatch, cursor: str | None, page_size: int ) -> None: native: Final = _NativeReturningHelp(QUERY_HELP, {**TRACE_RESPONSE, "next_cursor": "more"}) monkeypatch.setattr(loader, "_cached_bridge", native) @@ -762,6 +797,28 @@ async def test_storage_preserves_page_cursor_and_normalizes_native_trace_data( native.trace_read.assert_awaited_once_with("t1", scope, "run", cursor, page_size) +@pytest.mark.parametrize( + "args,expected", + ( + (("traversal_changed", "refresh"), TraceReadError("traversal_changed", "refresh")), + (("not-a-code", "x"), None), + (("busy",), None), + ), +) +async def test_storage_converts_native_read_failures_into_stable_codes( + monkeypatch: pytest.MonkeyPatch, args: tuple[object, ...], expected: TraceReadError | None +) -> None: + native: Final = _NativeReturningHelp(QUERY_HELP) + native.trace_read.side_effect = _NativeReadError(*args) + monkeypatch.setattr(loader, "_cached_bridge", native) + storage: Final = ClickHouseStorage(TraceStorageConfig("http://clickhouse:8123")) + scope: Final[TraceScope] = {"all_teams": 0, "user_id": "owner", "team_ids": ()} + with pytest.raises(TraceReadError if expected else RuntimeError) as raised: + await storage.get_trace("t1", scope, "run") + if expected is not None: + assert (raised.value.code, raised.value.message) == (expected.code, expected.message) + + async def test_storage_validates_the_native_query_help_value(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(loader, "_cached_bridge", _NativeReturningHelp(QUERY_HELP)) storage: Final = ClickHouseStorage(TraceStorageConfig("http://clickhouse:8123")) diff --git a/tests/unit/tracing/test_config.py b/tests/unit/tracing/test_config.py index 030c4247c62..363a6d6db67 100644 --- a/tests/unit/tracing/test_config.py +++ b/tests/unit/tracing/test_config.py @@ -32,12 +32,15 @@ def test_yaml_values_override_defaults_and_resolve_nested_references() -> None: "TRACING_DATABASE": "analytics", "TRACING_RETENTION_DAYS": "7", "CLICKHOUSE_URL": "https://other.example:8443", + "LITELLM_SALT_KEY": "salt", }, ) assert config.url == "https://writer:password@clickhouse.example:8443" assert config.database == "analytics" assert config.retention_days == 7 + assert config.cursor_keys == ("salt",) assert "password" not in repr(config) + assert "salt" not in repr(config) def test_omitted_fields_use_environment() -> None: @@ -47,6 +50,7 @@ def test_omitted_fields_use_environment() -> None: "CLICKHOUSE_URL": "http://localhost:8123", "CLICKHOUSE_DATABASE": "env_database", "AGENT_TRACING_RETENTION_DAYS": "11", + "LITELLM_MASTER_KEY": "sk-master", }, ) assert (config.url, config.database, config.retention_days) == ("http://localhost:8123", "env_database", 11) @@ -56,12 +60,13 @@ def test_environment_is_read_when_config_is_resolved(monkeypatch: pytest.MonkeyP monkeypatch.setenv("CLICKHOUSE_URL", "http://localhost:8123") monkeypatch.setenv("CLICKHOUSE_DATABASE", "late_database") monkeypatch.setenv("AGENT_TRACING_RETENTION_DAYS", "9") + monkeypatch.setenv("LITELLM_SALT_KEY", "salt") config = trace_storage_config({}) assert (config.database, config.retention_days) == ("late_database", 9) def test_omitted_fields_without_environment_use_constant_defaults() -> None: - config = trace_storage_config({}, {"CLICKHOUSE_URL": "http://localhost:8123"}) + config = trace_storage_config({}, {"CLICKHOUSE_URL": "http://localhost:8123", "LITELLM_SALT_KEY": "salt"}) assert (config.database, config.retention_days) == ( constants.DEFAULT_CLICKHOUSE_DATABASE, constants.DEFAULT_AGENT_TRACING_RETENTION_DAYS, @@ -114,3 +119,33 @@ def test_legacy_reader_and_split_retention_fields_are_rejected() -> None: }, {}, ) + + +@pytest.mark.parametrize( + ("store_keys", "environ", "expected"), + [ + (["os.environ/NEW_KEY", "literal-old"], {"NEW_KEY": "rotated", "LITELLM_SALT_KEY": "salt"}, ("rotated", "literal-old")), + ("single-key", {"LITELLM_SALT_KEY": "salt"}, ("single-key",)), + (None, {"LITELLM_SALT_KEY": "salt", "LITELLM_MASTER_KEY": "sk-master"}, ("salt",)), + (None, {"LITELLM_MASTER_KEY": "sk-master"}, ("sk-master",)), + ], +) +def test_cursor_keys_prefer_config_then_salt_then_master_key( + store_keys: object, environ: dict[str, str], expected: tuple[str, ...] +) -> None: + store: dict[str, object] = {"type": "clickhouse", "url": "http://localhost:8123"} + if store_keys is not None: + store["cursor_keys"] = store_keys + assert trace_storage_config({"store": store}, environ).cursor_keys == expected + + +@pytest.mark.parametrize("store_keys", [[], [""], [1], "", ["os.environ/MISSING"]]) +def test_invalid_cursor_keys_are_rejected(store_keys: object) -> None: + store: dict[str, object] = {"type": "clickhouse", "url": "http://localhost:8123", "cursor_keys": store_keys} + with pytest.raises(ValueError, match=r"cursor_keys|resolved to no value"): + trace_storage_config({"store": store}, {"LITELLM_SALT_KEY": "salt"}) + + +def test_missing_cursor_key_source_is_rejected() -> None: + with pytest.raises(ValueError, match="cursor_keys, LITELLM_SALT_KEY or LITELLM_MASTER_KEY is required"): + trace_storage_config({"store": {"type": "clickhouse", "url": "http://localhost:8123"}}, {}) diff --git a/ui/litellm-dashboard/src/components/lens/demo/createLensDemo.test.ts b/ui/litellm-dashboard/src/components/lens/demo/createLensDemo.test.ts index 8ddf632e015..f3fb4cacc18 100644 --- a/ui/litellm-dashboard/src/components/lens/demo/createLensDemo.test.ts +++ b/ui/litellm-dashboard/src/components/lens/demo/createLensDemo.test.ts @@ -58,9 +58,9 @@ describe("Lens demo data", () => { const { services } = createLensDemo(now); const all = await services.traces.list({ startMs: 0, endMs: now }); const recent = await services.traces.list({ startMs: now - 3600_000, endMs: now }); - expect(recent.data.length).toBeGreaterThan(0); - expect(recent.data.length).toBeLessThan(all.data.length); - expect(recent.data.every((trace) => Date.parse(trace.start_time) >= now - 3600_000)).toBe(true); + expect(recent.items.length).toBeGreaterThan(0); + expect(recent.items.length).toBeLessThan(all.items.length); + expect(recent.items.every((trace) => Date.parse(trace.start_time) >= now - 3600_000)).toBe(true); const settings = services.lens.lenses().then((list) => list.lenses[0].settings); await expect(services.lens.saveLens(undefined, await settings)).rejects.toMatchObject({ status: 403 }); await expect(services.lens.run("real-investigation", "job")).rejects.toMatchObject({ status: 404 }); diff --git a/ui/litellm-dashboard/src/components/lens/demo/createLensDemo.ts b/ui/litellm-dashboard/src/components/lens/demo/createLensDemo.ts index bc75214127d..8a3cb7974eb 100644 --- a/ui/litellm-dashboard/src/components/lens/demo/createLensDemo.ts +++ b/ui/litellm-dashboard/src/components/lens/demo/createLensDemo.ts @@ -376,10 +376,11 @@ function demoTracesApi(data: LensDemoData): TracesApi { const run = (traceId: string) => data.runs.find(({ trace }) => trace.summary.trace_id === traceId); return { list: async ({ startMs, endMs }) => ({ - data: data.runs + items: data.runs .map((item) => item.trace.summary) .filter((trace) => Date.parse(trace.start_time) >= startMs && Date.parse(trace.start_time) <= endMs), next_cursor: null, + traversal: { id: "demo", published_at: new Date(endMs).toISOString(), expires_at: new Date(endMs).toISOString() }, }), anyRecorded: async () => data.runs.length > 0, trace: (traceId) => found(run(traceId)?.trace), diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.integration.test.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.integration.test.tsx index b613d028f54..83f1669abff 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.integration.test.tsx @@ -9,6 +9,7 @@ import traceList from "./__fixtures__/trace_list.json"; import AgentTracesPage from "./AgentTracesPage"; import { AgentTracesSection, type TimeControls } from "./AgentTracesSection"; import type { TracePage, TraceSummary } from "./traceTypes"; +import { RESULTS_CHANGED_MESSAGE } from "./useAgentTraces"; vi.mock("../../networking", () => ({ apiClient: { get: vi.fn(), post: vi.fn() }, @@ -30,7 +31,7 @@ vi.mock("./TraceDrawer", () => ({ import { agentTraceListCall, apiClient } from "../../networking"; -const runs = (traceList as TracePage).data as TraceSummary[]; +const runs = (traceList as TracePage).items as TraceSummary[]; const renderSection = () => renderWithProviders( @@ -82,13 +83,44 @@ describe("AgentTracesSection", () => { } as DOMRect); testQueryClient.clear(); vi.mocked(agentTraceListCall).mockReset(); - vi.mocked(apiClient.get).mockResolvedValue({ data: [] }); + vi.mocked(apiClient.get).mockResolvedValue({ ...(traceList as TracePage), items: [] }); + }); + + it("restarts from a fresh first page when the server reports the results changed", async () => { + const user = userEvent.setup(); + vi.mocked(agentTraceListCall).mockResolvedValueOnce({ + ...(traceList as TracePage), + items: runs.slice(0, 1), + next_cursor: "next", + }); + renderSection(); + expect(await screen.findByTestId("agent-trace-row")).toBeVisible(); + vi.mocked(agentTraceListCall).mockRejectedValueOnce( + new ApiError("Conflict", 409, { detail: { code: "traversal_changed", message: "Results changed" } }), + ); + await user.click(screen.getByRole("button", { name: "Load more" })); + expect(await screen.findByRole("alert")).toHaveTextContent(RESULTS_CHANGED_MESSAGE); + expect(screen.queryByRole("button", { name: "Retry" })).not.toBeInTheDocument(); + expect(vi.mocked(agentTraceListCall)).toHaveBeenCalledTimes(2); + vi.mocked(agentTraceListCall).mockResolvedValueOnce({ + ...(traceList as TracePage), + items: runs.slice(0, 2), + next_cursor: null, + }); + await user.click(screen.getByRole("button", { name: "Refresh" })); + await waitFor(() => expect(screen.getAllByTestId("agent-trace-row")).toHaveLength(2)); + expect(vi.mocked(agentTraceListCall).mock.calls[2][0]).not.toHaveProperty("cursor"); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); }); it("keeps the original time window and loaded rows when another page fails", async () => { const user = userEvent.setup(); const now = vi.spyOn(Date, "now").mockReturnValue(Date.parse("2026-10-01T00:00Z")); - vi.mocked(agentTraceListCall).mockResolvedValueOnce({ data: runs.slice(0, 1), next_cursor: "next" }); + vi.mocked(agentTraceListCall).mockResolvedValueOnce({ + ...(traceList as TracePage), + items: runs.slice(0, 1), + next_cursor: "next", + }); renderSection(); expect(await screen.findByTestId("agent-trace-row")).toBeVisible(); const first = vi.mocked(agentTraceListCall).mock.calls[0][0]; @@ -98,7 +130,11 @@ describe("AgentTracesSection", () => { expect(await screen.findByRole("alert")).toHaveTextContent("Could not load more runs"); expect(screen.getAllByTestId("agent-trace-row")).toHaveLength(1); expect(vi.mocked(agentTraceListCall).mock.calls[1][0]).toEqual({ ...first, cursor: "next" }); - vi.mocked(agentTraceListCall).mockResolvedValueOnce({ data: runs.slice(1, 2), next_cursor: null }); + vi.mocked(agentTraceListCall).mockResolvedValueOnce({ + ...(traceList as TracePage), + items: runs.slice(1, 2), + next_cursor: null, + }); await user.click(screen.getByRole("button", { name: "Retry" })); await waitFor(() => expect(screen.getAllByTestId("agent-trace-row")).toHaveLength(2)); expect(screen.queryByRole("alert")).not.toBeInTheDocument(); @@ -147,7 +183,7 @@ describe("AgentTracesSection", () => { }); it("shows the waiting guide when tracing is on but no runs have arrived", async () => { - vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), data: [] }); + vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), items: [] }); renderSection(); const card = await screen.findByTestId("tracing-setup-card"); @@ -158,7 +194,7 @@ describe("AgentTracesSection", () => { }); it("keeps the trace list available when traces exist outside the current time window", async () => { - vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), data: [] }); + vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), items: [] }); vi.mocked(apiClient.get).mockResolvedValue(traceList); renderSection(); expect(await screen.findByText("No runs match these filters.")).toBeVisible(); @@ -171,7 +207,7 @@ describe("AgentTracesSection", () => { vi.mocked(agentTraceListCall).mockRejectedValue(new ApiError("Tracing is not enabled", 501, {})); renderSection(); const checkSetup = await screen.findByRole("button", { name: "Check setup" }); - vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), data: [] }); + vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), items: [] }); fireEvent.click(checkSetup); expect(await screen.findByRole("heading", { name: "Connect your agent" })).toBeVisible(); expect(screen.getByText("Waiting for your first trace")).toBeVisible(); @@ -231,7 +267,7 @@ describe("AgentTracesSection", () => { }); it("separates a failed history check from the empty list and retries that check", async () => { - vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), data: [] }); + vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), items: [] }); vi.mocked(apiClient.get).mockRejectedValue(new ApiError("History unavailable", 503, {})); renderSection(); expect(await screen.findByRole("alert")).toHaveTextContent("Could not check earlier traces. History unavailable"); @@ -272,7 +308,7 @@ describe("AgentTracesSection", () => { it("shows the spend returned for a run", async () => { vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), - data: [{ ...runs[0], spend: 0.025 }], + items: [{ ...runs[0], spend: 0.025 }], }); renderSection(); @@ -299,7 +335,7 @@ describe("AgentTracesSection", () => { it("uses recorded agent names for the column and filter even when services are shared", async () => { vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), - data: [ + items: [ ...runs.slice(1).map((run) => ({ ...run, service: "shared-app", agent_names: ["research-agent"] })), { ...runs[0], service: "shared-app", agent_names: ["billing-agent", "review-agent"] }, ], @@ -328,7 +364,7 @@ describe("AgentTracesSection", () => { it("shows each run's agent name with the logo of the SDK that produced it", async () => { vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), - data: [ + items: [ { ...runs[0], agent_names: ["research-bot"], frameworks: ["claude-agent-sdk", "claude-code"] }, { ...runs[1], agent_names: [], frameworks: ["claude-code"] }, { ...runs[2], frameworks: [] }, @@ -466,7 +502,7 @@ describe("AgentTracesPage", () => { beforeEach(() => { testQueryClient.clear(); vi.mocked(agentTraceListCall).mockReset(); - vi.mocked(apiClient.get).mockResolvedValue({ data: [] }); + vi.mocked(apiClient.get).mockResolvedValue({ ...(traceList as TracePage), items: [] }); }); it("shows the actual range, switches presets from the popover, and toggles Live", async () => { @@ -507,7 +543,7 @@ describe("AgentTracesPage", () => { renderWithProviders(); await screen.findByTestId("runs-table"); - vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), data: [] }); + vi.mocked(agentTraceListCall).mockResolvedValue({ ...(traceList as TracePage), items: [] }); fireEvent.click(screen.getByRole("button", { name: "Time range" })); fireEvent.click(await screen.findByRole("menuitemradio", { name: "Last hour" })); diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.test.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.test.tsx index 062c6d30da6..db7cf0ffbb5 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.test.tsx @@ -3,7 +3,7 @@ import { filterRuns } from "./AgentTracesSection"; import traceList from "./__fixtures__/trace_list.json"; import type { TracePage, TraceSummary } from "./traceTypes"; -const runs = (traceList as TracePage).data as TraceSummary[]; +const runs = (traceList as TracePage).items as TraceSummary[]; describe("filterRuns", () => { it("status filter 'Failed' keeps only runs with errors", () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.tsx index 211f50dcd38..8d4f2f49861 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.tsx @@ -230,7 +230,8 @@ export function AgentTracesSection({ error={traces.error} hasMore={traces.hasMore} isFetching={traces.isFetching} - onRetry={traces.hasMore ? traces.loadMore : traces.refetch} + onRetry={traces.retry} + retryLabel={traces.retryLabel} onLoadMore={traces.loadMore} onOpenTrace={toggleRun} selectedKey={openTrace === null ? null : runKey(openTrace)} diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesTable.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesTable.tsx index 964b6e82469..ab8d2897263 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesTable.tsx @@ -18,6 +18,7 @@ interface AgentTracesTableProps { hasMore: boolean; isFetching?: boolean; onRetry?: () => void; + retryLabel?: string; onLoadMore: () => void; onOpenTrace: (trace: TraceSummary) => void; selectedKey?: string | null; @@ -57,6 +58,7 @@ export function AgentTracesTable({ hasMore, isFetching = false, onRetry, + retryLabel = "Retry", onLoadMore, onOpenTrace, selectedKey = null, @@ -150,7 +152,7 @@ export function AgentTracesTable({ {onRetry && ( )} diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/TraceDrawer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/TraceDrawer.test.tsx index bd881730925..eaf87d7062e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/TraceDrawer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/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"; @@ -190,6 +192,30 @@ describe("RunView", () => { expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, "next-page", "next-page"]); }); + it("offers only a refresh when the server reports the trace changed while paging", 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: "old-second" }; + const fresh: Trace = { ...first, spans: [{ ...research.spans[0], name: "fresh-root" }], next_cursor: null }; + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(first) + .mockRejectedValueOnce( + new ApiError("Conflict", 409, { detail: { code: "traversal_changed", message: "Trace changed" } }), + ) + .mockResolvedValueOnce(fresh); + renderWithProviders(); + await user.click(await screen.findByRole("button", { name: "Load more steps" })); + expect( + await screen.findByText("This trace changed while loading. Refresh to continue from the latest version."), + ).toBeVisible(); + expect(screen.queryByRole("button", { name: "Retry" })).not.toBeInTheDocument(); + expect(screen.getByRole("tree", { name: "Spans in time order" })).toHaveTextContent(research.spans[0].name); + await user.click(screen.getByRole("button", { name: "Refresh trace" })); + expect(await screen.findByText("fresh-root")).toBeVisible(); + expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, "old-second", null]); + }); + it("refreshes a failed later page from one new snapshot", async () => { const user = userEvent.setup(); const summary = { ...research.summary, span_count: 3 }; diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/TraceDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/TraceDrawer.tsx index 3ff8bd0ef4d..357ae2ca45b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/TraceDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/TraceDrawer.tsx @@ -17,6 +17,7 @@ import { DetailPane } from "./DetailPane"; import { IdChip } from "./IdChip"; import { formatCost } from "./AgentTracesTable"; import { SpanIcon } from "./SpanIcon"; +import { restartsTraversal } from "./readFailure"; import { SpanTree } from "./SpanTree"; import { TraceConversation } from "./TraceConversation"; import { FrameworkLogo, traceFramework } from "./TraceFramework"; @@ -372,6 +373,12 @@ function initialSpanMissing(trace: Trace | undefined, spanId?: string): boolean return Boolean(spanId && trace && !trace.spans.some((span) => span.span_id === spanId)); } +const loadMoreStatus = (trace: Trace, failed: boolean, mustRestart: boolean): string => { + if (mustRestart) return "This trace changed while loading. Refresh to continue from the latest version."; + if (failed) return "Could not load more steps. Your loaded steps are still available."; + return `Showing ${trace.spans.length.toLocaleString()} of ${trace.summary.span_count.toLocaleString()} steps`; +}; + export function RunView({ traceId, traceRef, initialSpanId, accessToken, onBack, embedded = false }: RunViewProps) { const traces = useTracesApi(accessToken); const queryClient = useQueryClient(); @@ -400,6 +407,7 @@ export function RunView({ traceId, traceRef, initialSpanId, accessToken, onBack, useEffect(() => { if (canSeek && !isFetching && !isError) void fetchNextPage(); }, [canSeek, isFetching, isError, fetchNextPage]); + const mustRestart = restartsTraversal(traceQuery.error); const pageAction = isError ? "Retry" : "Load more steps"; if (traceQuery.isLoading) { @@ -450,24 +458,27 @@ export function RunView({ traceId, traceRef, initialSpanId, accessToken, onBack, {(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`} - + {loadMoreStatus(trace, traceQuery.isError, mustRestart)} {traceQuery.isError && ( - )} - + {!mustRestart && ( + + )}
)} new ApiError("read failed", status, { detail: { code, message: "x" } }); + +describe("readFailureCode", () => { + it("reads the stable code from a proxy read failure", () => { + expect(readFailureCode(failure("traversal_changed"))).toBe("traversal_changed"); + expect(readFailureCode(failure("busy", 503))).toBe("busy"); + }); + + it.each([ + ["plain error", new Error("offline")], + ["string detail", new ApiError("x", 400, { detail: "Invalid trace cursor" })], + ["unknown code", failure("not_a_code")], + ["no body", new ApiError("x", 500, null)], + ])("returns null for %s", (_label, error) => { + expect(readFailureCode(error)).toBeNull(); + }); +}); + +describe("restartsTraversal", () => { + it.each(["traversal_changed", "traversal_expired", "invalid_cursor"])("restarts on %s", (code) => { + expect(restartsTraversal(failure(code))).toBe(true); + }); + + it.each(["busy", "unavailable", "view_not_ready", "resource_too_large", "budget_exceeded", "invalid_request"])( + "keeps the current cursor on %s", + (code) => { + expect(restartsTraversal(failure(code, 503))).toBe(false); + }, + ); + + it("keeps the current cursor for errors without a code", () => { + expect(restartsTraversal(new Error("offline"))).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/readFailure.ts b/ui/litellm-dashboard/src/components/view_logs/TraceView/readFailure.ts new file mode 100644 index 00000000000..027f35ee644 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/readFailure.ts @@ -0,0 +1,38 @@ +import { ApiError } from "@/lib/http/client"; + +import type { FailureCode } from "./traceTypes"; + +const FAILURE_CODES: ReadonlySet = new Set([ + "invalid_request", + "invalid_cursor", + "traversal_expired", + "traversal_changed", + "resource_too_large", + "budget_exceeded", + "view_not_ready", + "unavailable", + "busy", +]); + +const RESTART_CODES: ReadonlySet = new Set([ + "invalid_cursor", + "traversal_expired", + "traversal_changed", +]); + +const isFailureCode = (value: unknown): value is FailureCode => typeof value === "string" && FAILURE_CODES.has(value); + +/** The stable `detail.code` the proxy attaches to a failed trace read, if the error carries one. */ +export const readFailureCode = (error: unknown): FailureCode | null => { + if (!(error instanceof ApiError) || typeof error.body !== "object" || error.body === null) return null; + const detail: unknown = (error.body as { detail?: unknown }).detail; + if (typeof detail !== "object" || detail === null) return null; + const code: unknown = (detail as { code?: unknown }).code; + return isFailureCode(code) ? code : null; +}; + +/** Continuing from the current cursor cannot succeed; the traversal has to start over from a fresh first page. */ +export const restartsTraversal = (error: unknown): boolean => { + const code = readFailureCode(error); + return code !== null && RESTART_CODES.has(code); +}; diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/traceTypes.ts b/ui/litellm-dashboard/src/components/view_logs/TraceView/traceTypes.ts index a424e010670..e4b6a86d374 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/traceTypes.ts +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/traceTypes.ts @@ -27,3 +27,5 @@ export type TraceMessage = Omit & { role: string; tool_calls?: TraceToolCall[]; }; +export type FailureCode = components["schemas"]["ReadFailure"]["code"]; +export type Traversal = TracePage["traversal"]; diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/tracesApi.ts b/ui/litellm-dashboard/src/components/view_logs/TraceView/tracesApi.ts index ada76f498a5..1b4aca7666b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/tracesApi.ts +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/tracesApi.ts @@ -30,7 +30,7 @@ export function liveTracesApi(accessToken: string): TracesApi { list: (window) => agentTraceListCall({ accessToken, ...window }), anyRecorded: async () => { const page = await apiClient.get("/v1/traces", { accessToken, query: { start_ms: 0 } }); - return page.data.length > 0; + return page.items.length > 0; }, trace: (traceId, traceRef, cursor) => agentTraceCall(accessToken, traceId, traceRef, cursor), span: (traceId, spanId, traceRef) => agentTraceSpanCall(accessToken, traceId, spanId, traceRef), diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/useAgentTraces.ts b/ui/litellm-dashboard/src/components/view_logs/TraceView/useAgentTraces.ts index b1e20748770..960c9482652 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/useAgentTraces.ts +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/useAgentTraces.ts @@ -6,6 +6,7 @@ import { useMemo } from "react"; import { ApiError } from "@/lib/http/client"; import { LIVE_TAIL_INTERVAL_MS } from "../log_filter_logic"; +import { restartsTraversal } from "./readFailure"; import type { TracePage, TraceSummary } from "./traceTypes"; import type { TraceWindow } from "./tracesApi"; @@ -26,10 +27,13 @@ const requiresUserAction = (error: unknown): boolean => { return isTracingNotEnabled(error) || error.status === 401 || error.status === 403; }; +export const RESULTS_CHANGED_MESSAGE = "Results changed while loading. Refresh to start from the latest runs."; + const displayError = (error: Error | null): Error | null => { if (!(error instanceof ApiError)) return error; if (error.status === 401) return new Error("Your session is no longer valid. Sign out and sign in again."); if (error.status === 403) return new Error("Your account does not have access to these traces."); + if (restartsTraversal(error)) return new Error(RESULTS_CHANGED_MESSAGE); return error; }; @@ -49,9 +53,14 @@ export interface AgentTracesResult { /** Set when the proxy answered 501: tracing isn't configured. */ notEnabledDetail: string | null; error: Error | null; + /** The failed page cannot be retried from its cursor; `refetch` restarts from a fresh first page. */ + mustRestart: boolean; hasMore: boolean; loadMore: () => void; refetch: () => void; + /** Resumes a retryable failure from its cursor, or restarts when the traversal is gone. */ + retry: () => void; + retryLabel: "Retry" | "Refresh"; } /** Start of the fetch window: a preset range rolls with "now", so live tail keeps a fixed-length window. */ @@ -87,7 +96,7 @@ export function useAgentTraces({ lastPage.next_cursor ? { ...lastPage.window, cursor: lastPage.next_cursor } : undefined, enabled, staleTime: LIVE_TAIL_INTERVAL_MS, - retry: (failureCount, error) => !requiresUserAction(error) && failureCount < 1, + retry: (failureCount, error) => !requiresUserAction(error) && !restartsTraversal(error) && failureCount < 1, refetchInterval: (q) => (isLiveTail && !requiresUserAction(q.state.error) ? LIVE_TAIL_INTERVAL_MS : false), refetchOnWindowFocus: (q) => !requiresUserAction(q.state.error), refetchOnReconnect: (q) => !requiresUserAction(q.state.error), @@ -95,8 +104,13 @@ export function useAgentTraces({ }; const query = useInfiniteQuery(queryOptions); - const loaded = useMemo(() => query.data?.pages.flatMap((page) => page.data) ?? [], [query.data]); + const loaded = useMemo(() => query.data?.pages.flatMap((page) => page.items) ?? [], [query.data]); const notEnabled = isTracingNotEnabled(query.error); + const mustRestart = restartsTraversal(query.error); + const loadMore = () => { + if (!query.isFetching) void query.fetchNextPage(); + }; + const refetch = () => void query.refetch(); return { traces: loaded, @@ -104,11 +118,12 @@ export function useAgentTraces({ isFetching: query.isFetching, notEnabledDetail: notEnabled ? query.error?.message || "Agent tracing is not enabled" : null, error: notEnabled ? null : displayError(query.error), + mustRestart, hasMore: query.hasNextPage, - loadMore: () => { - if (!query.isFetching) void query.fetchNextPage(); - }, - refetch: () => void query.refetch(), + loadMore, + refetch, + retry: query.hasNextPage && !mustRestart ? loadMore : refetch, + retryLabel: mustRestart ? "Refresh" : "Retry", }; } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0dde64137a5..a1c295b3743 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -41757,6 +41757,20 @@ export interface components { [key: string]: unknown; } | null; }; + /** ReadFailure */ + ReadFailure: { + /** + * Code + * @enum {string} + */ + code: "invalid_request" | "invalid_cursor" | "traversal_expired" | "traversal_changed" | "resource_too_large" | "budget_exceeded" | "view_not_ready" | "unavailable" | "busy"; + /** Message */ + message: string; + }; + /** ReadFailureResponse */ + ReadFailureResponse: { + detail: components["schemas"]["ReadFailure"]; + }; /** * RealtimeClientSecretResponse * @description Response from POST /v1/realtime/client_secrets. @@ -47136,13 +47150,15 @@ export interface components { /** Spans */ spans: components["schemas"]["Span"][]; summary: components["schemas"]["TraceSummary"]; + traversal?: components["schemas"]["Traversal"] | null; }; /** TracePage */ TracePage: { - /** Data */ - data: components["schemas"]["TraceSummary"][]; + /** Items */ + items: components["schemas"]["TraceSummary"][]; /** Next Cursor */ next_cursor: string | null; + traversal: components["schemas"]["Traversal"]; }; /** TracePart */ TracePart: { @@ -47446,6 +47462,15 @@ export interface components { [key: string]: unknown; }; }; + /** Traversal */ + Traversal: { + /** Expires At */ + expires_at: string; + /** Id */ + id: string; + /** Published At */ + published_at: string; + }; /** * Type * @description An action to type in text. @@ -80275,6 +80300,42 @@ export interface operations { "application/json": components["schemas"]["TracePage"]; }; }; + /** @description Bad Request */ + 400: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Conflict */ + 409: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Gone */ + 410: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Request Entity Too Large */ + 413: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; /** @description Validation Error */ 422: { headers: { @@ -80284,6 +80345,15 @@ export interface operations { "application/json": components["schemas"]["HTTPValidationError"]; }; }; + /** @description Service Unavailable */ + 503: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; }; }; ingest_otlp_traces_v1_traces_post: { @@ -80364,7 +80434,7 @@ export interface operations { query?: { trace_ref?: string; cursor?: string | null; - page_size?: number | null; + page_size?: number; }; header?: never; path: { @@ -80383,6 +80453,42 @@ export interface operations { "application/json": components["schemas"]["Trace"]; }; }; + /** @description Bad Request */ + 400: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Conflict */ + 409: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Gone */ + 410: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Request Entity Too Large */ + 413: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; /** @description Validation Error */ 422: { headers: { @@ -80392,6 +80498,15 @@ export interface operations { "application/json": components["schemas"]["HTTPValidationError"]; }; }; + /** @description Service Unavailable */ + 503: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; }; }; get_agent_trace_span_v1_traces__trace_id__spans__span_id__get: { @@ -80417,6 +80532,42 @@ export interface operations { "application/json": components["schemas"]["SpanDetail"]; }; }; + /** @description Bad Request */ + 400: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Conflict */ + 409: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Gone */ + 410: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Request Entity Too Large */ + 413: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; /** @description Validation Error */ 422: { headers: { @@ -80426,6 +80577,15 @@ export interface operations { "application/json": components["schemas"]["HTTPValidationError"]; }; }; + /** @description Service Unavailable */ + 503: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; }; }; get_agent_trace_span_error_v1_traces__trace_id__spans__span_id__error_get: { @@ -80452,6 +80612,42 @@ export interface operations { "application/json": components["schemas"]["SpanErrorPage"]; }; }; + /** @description Bad Request */ + 400: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Conflict */ + 409: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Gone */ + 410: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; + /** @description Request Entity Too Large */ + 413: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; /** @description Validation Error */ 422: { headers: { @@ -80461,6 +80657,15 @@ export interface operations { "application/json": components["schemas"]["HTTPValidationError"]; }; }; + /** @description Service Unavailable */ + 503: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ReadFailureResponse"]; + }; + }; }; }; list_access_groups_v1_unified_access_group_get: {