From 97ded17960888f89cf2d48619b201daaa6e477b1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 3 Oct 2026 22:56:14 +0000 Subject: [PATCH 1/4] feat(tracing): shared signed pagination for trace list and detail reads Add the litellm-pagination crate (Page envelope, signed expiring cursors bound to resource, scope, query and revision, stable failure codes) and move the trace list, trace detail and span error reads onto it. The list pins a publication fence through a new ReceivedMs rollup column so continuation pages exclude later-received rows. The Python bridge raises TraceReadError(code, message), the proxy maps codes to 400/409/410/413/503 with Retry-After, cursor keys come from tracing.store.cursor_keys with LITELLM_SALT_KEY and LITELLM_MASTER_KEY fallbacks, and Lens consumes the items envelope and shows Refresh instead of Retry when a traversal is gone. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 19 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/pagination/AGENTS.md | 5 + litellm-rust/crates/pagination/Cargo.toml | 23 + litellm-rust/crates/pagination/src/cursor.rs | 208 +++++++ litellm-rust/crates/pagination/src/error.rs | 94 ++++ litellm-rust/crates/pagination/src/lib.rs | 10 + litellm-rust/crates/pagination/src/page.rs | 75 +++ .../crates/pagination/tests/cursor.rs | 203 +++++++ litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 3 +- .../crates/python-bridge/src/routes/traces.rs | 123 ++-- litellm-rust/crates/traces-cache/src/lib.rs | 8 +- .../crates/traces-clickhouse/Cargo.toml | 1 + .../0009_trace_rollup_ownership.sql | 3 +- .../0010_trace_cost_completeness.sql | 1 + .../traces-clickhouse/query/list_traces.sql | 2 + .../crates/traces-clickhouse/src/config.rs | 11 + .../crates/traces-clickhouse/src/error.rs | 57 +- .../traces-clickhouse/src/query/named.rs | 4 +- .../crates/traces-clickhouse/src/reads.rs | 530 +++++++++++++----- .../traces-clickhouse/src/span_batches.rs | 3 +- .../traces-clickhouse/tests/migrations.rs | 1 + .../crates/traces-clickhouse/tests/queries.rs | 4 + .../crates/traces-clickhouse/tests/reads.rs | 147 +++-- litellm-rust/crates/traces/Cargo.toml | 3 +- litellm-rust/crates/traces/src/lib.rs | 1 + litellm-rust/crates/traces/src/query/named.rs | 2 + .../crates/traces/src/resolve/view.rs | 1 + litellm-rust/crates/traces/src/schema.rs | 12 +- litellm-rust/crates/traces/src/view.rs | 9 +- .../crates/traces/tests/query/named.rs | 7 +- litellm/constants.py | 1 + litellm/proxy/tracing_endpoints.py | 84 ++- litellm/rust_bridge/trace/generated/types.py | 32 +- litellm/rust_bridge/trace/storage.py | 65 ++- litellm/tracing/config.py | 30 +- litellm/tracing/receiver.py | 9 +- .../schemas/traces/ReadFailure.json | 35 ++ .../trace_codegen/schemas/traces/Trace.json | 34 +- .../schemas/traces/TracePage.json | 31 +- tests/test_litellm_rust/test_traces.py | 30 +- tests/unit/proxy/test_tracing_endpoints.py | 89 ++- tests/unit/tracing/test_config.py | 37 +- .../lens/demo/createLensDemo.test.ts | 6 +- .../components/lens/demo/createLensDemo.ts | 3 +- .../AgentTracesSection.integration.test.tsx | 62 +- .../TraceView/AgentTracesSection.test.tsx | 2 +- .../TraceView/AgentTracesSection.tsx | 3 +- .../view_logs/TraceView/AgentTracesTable.tsx | 4 +- .../view_logs/TraceView/TraceDrawer.test.tsx | 26 + .../view_logs/TraceView/TraceDrawer.tsx | 39 +- .../TraceView/__fixtures__/trace_list.json | 9 +- .../view_logs/TraceView/readFailure.test.ts | 40 ++ .../view_logs/TraceView/readFailure.ts | 38 ++ .../view_logs/TraceView/traceTypes.ts | 2 + .../view_logs/TraceView/tracesApi.ts | 2 +- .../view_logs/TraceView/useAgentTraces.ts | 27 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 211 ++++++- 59 files changed, 2158 insertions(+), 365 deletions(-) create mode 100644 litellm-rust/crates/pagination/AGENTS.md create mode 100644 litellm-rust/crates/pagination/Cargo.toml create mode 100644 litellm-rust/crates/pagination/src/cursor.rs create mode 100644 litellm-rust/crates/pagination/src/error.rs create mode 100644 litellm-rust/crates/pagination/src/lib.rs create mode 100644 litellm-rust/crates/pagination/src/page.rs create mode 100644 litellm-rust/crates/pagination/tests/cursor.rs create mode 100644 scripts/trace_codegen/schemas/traces/ReadFailure.json create mode 100644 ui/litellm-dashboard/src/components/view_logs/TraceView/readFailure.test.ts create mode 100644 ui/litellm-dashboard/src/components/view_logs/TraceView/readFailure.ts 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: { From 1c12e4cf731fe8478a24056e2927846e3c80fb1d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 3 Oct 2026 23:14:58 +0000 Subject: [PATCH 2/4] fix(tracing): append ReceivedMs migrations, make cursor keys reader-only, sync native stub Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/pagination/src/cursor.rs | 17 ++++++++------ litellm-rust/crates/pagination/src/error.rs | 4 +++- .../crates/pagination/tests/cursor.rs | 22 +++++++++++++++++- .../0009_trace_rollup_ownership.sql | 3 +-- .../0010_trace_cost_completeness.sql | 1 - .../migrations/0015_agent_traces_received.sql | 2 ++ .../0016_agent_traces_mv_received.sql | 23 +++++++++++++++++++ litellm/rust_bridge/_native.pyi | 6 ++++- litellm/tracing/config.py | 11 ++++++--- litellm/tracing/receiver.py | 4 ++-- tests/unit/tracing/test_config.py | 16 +++++++++---- .../lens/LensWorkspace.integration.test.tsx | 8 ++++++- 12 files changed, 94 insertions(+), 23 deletions(-) create mode 100644 litellm-rust/crates/traces-clickhouse/migrations/0015_agent_traces_received.sql create mode 100644 litellm-rust/crates/traces-clickhouse/migrations/0016_agent_traces_mv_received.sql diff --git a/litellm-rust/crates/pagination/src/cursor.rs b/litellm-rust/crates/pagination/src/cursor.rs index 7ccff2c2455..b997523a42e 100644 --- a/litellm-rust/crates/pagination/src/cursor.rs +++ b/litellm-rust/crates/pagination/src/cursor.rs @@ -41,14 +41,15 @@ impl KeyRing { }) }) .collect::, _>>()?; - if keys.is_empty() { - return Err(Error::InvalidKeys); - } Ok(Self { keys }) } - pub fn signing_key_id(&self) -> &str { - &self.keys[0].id + pub fn signing_key_id(&self) -> Option<&str> { + self.keys.first().map(|key| key.id.as_str()) + } + + fn signing_key(&self) -> Result<&Key, Error> { + self.keys.first().ok_or(Error::MissingKeys) } fn sign(key: &Key, binding: &Binding, payload: &[u8]) -> Vec { @@ -66,15 +67,16 @@ impl KeyRing { binding: &Binding, cursor: &Cursor

, ) -> Result { + let signing_key = self.signing_key()?; let payload = serde_json::to_vec(&Envelope { version: FORMAT_VERSION, - key_id: &self.keys[0].id, + key_id: &signing_key.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); + let tag = Self::sign(signing_key, binding, &payload); Ok(format!( "{}.{}", URL_SAFE_NO_PAD.encode(payload), @@ -89,6 +91,7 @@ impl KeyRing { token: &str, now: OffsetDateTime, ) -> Result, Error> { + self.signing_key()?; let (payload, tag) = token.split_once('.').ok_or(Error::InvalidCursor)?; let payload = URL_SAFE_NO_PAD .decode(payload) diff --git a/litellm-rust/crates/pagination/src/error.rs b/litellm-rust/crates/pagination/src/error.rs index e713563ecf8..ce72a543b30 100644 --- a/litellm-rust/crates/pagination/src/error.rs +++ b/litellm-rust/crates/pagination/src/error.rs @@ -71,6 +71,8 @@ pub enum Error { ResourceTooLarge, #[error("Cursor signing keys must be nonempty")] InvalidKeys, + #[error("Cursor signing keys are not configured")] + MissingKeys, #[error("Cursor encoding failed")] Encode(#[from] serde_json::Error), } @@ -82,7 +84,7 @@ impl Error { Self::TraversalExpired => FailureCode::TraversalExpired, Self::TraversalChanged => FailureCode::TraversalChanged, Self::ResourceTooLarge => FailureCode::ResourceTooLarge, - Self::InvalidKeys | Self::Encode(_) => FailureCode::Unavailable, + Self::InvalidKeys | Self::MissingKeys | Self::Encode(_) => FailureCode::Unavailable, } } } diff --git a/litellm-rust/crates/pagination/tests/cursor.rs b/litellm-rust/crates/pagination/tests/cursor.rs index 541295bd792..e8618826004 100644 --- a/litellm-rust/crates/pagination/tests/cursor.rs +++ b/litellm-rust/crates/pagination/tests/cursor.rs @@ -131,12 +131,31 @@ fn rotated_keys_keep_verifying_earlier_cursors(binding: Binding, now: OffsetDate } #[rstest] -#[case::no_keys(Vec::<&str>::new())] #[case::empty_key(vec![""])] +#[case::empty_key_among_valid(vec!["valid", ""])] fn key_rings_require_nonempty_secrets(#[case] secrets: Vec<&str>) { assert!(matches!(KeyRing::new(secrets), Err(Error::InvalidKeys))); } +#[rstest] +fn a_key_ring_without_keys_refuses_to_sign_or_verify( + keys: KeyRing, + binding: Binding, + now: OffsetDateTime, +) { + let unsigned = KeyRing::new(Vec::<&str>::new()).unwrap(); + let token = keys.encode(&binding, &cursor(now)).unwrap(); + assert_eq!(unsigned.signing_key_id(), None); + assert!(matches!( + unsigned.encode(&binding, &cursor(now)), + Err(Error::MissingKeys) + )); + assert!(matches!( + unsigned.decode::(&binding, &token, now), + Err(Error::MissingKeys) + )); +} + #[rstest] fn traversal_identity_follows_binding_and_publication(binding: Binding, now: OffsetDateTime) { let first = Traversal::new(&binding, 1_790_000_000_000, now); @@ -190,6 +209,7 @@ fn bounded_pages_keep_a_complete_prefix_and_reissue_the_cursor( #[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)] +#[case::missing_keys(Error::MissingKeys, FailureCode::Unavailable, true, false)] fn errors_map_to_stable_codes( #[case] error: Error, #[case] code: FailureCode, 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 f4445980378..fd696349cf5 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,4 +1,3 @@ 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 ReceivedMs SimpleAggregateFunction(min, UInt64) DEFAULT 0 + ADD COLUMN IF NOT EXISTS IdentifiedLlmCount SimpleAggregateFunction(sum, 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 df70a50b348..af87a6bf40b 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,7 +15,6 @@ 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/migrations/0015_agent_traces_received.sql b/litellm-rust/crates/traces-clickhouse/migrations/0015_agent_traces_received.sql new file mode 100644 index 00000000000..4620479afec --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/migrations/0015_agent_traces_received.sql @@ -0,0 +1,2 @@ +ALTER TABLE {database}.agent_traces_by_key + ADD COLUMN IF NOT EXISTS ReceivedMs SimpleAggregateFunction(min, UInt64) DEFAULT 0 diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0016_agent_traces_mv_received.sql b/litellm-rust/crates/traces-clickhouse/migrations/0016_agent_traces_mv_received.sql new file mode 100644 index 00000000000..df70a50b348 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/migrations/0016_agent_traces_mv_received.sql @@ -0,0 +1,23 @@ +ALTER TABLE {database}.agent_traces_by_key_mv MODIFY QUERY +SELECT + TeamId, ApiKeyHash, TraceId, groupUniqArray(UserId) AS UserIds, + min(Timestamp) AS StartTs, + max(Timestamp + toIntervalNanosecond(Duration)) AS EndTs, + any(ServiceName) AS ServiceName, + anyLastIf(toNullable(SpanName), ParentSpanId = '') AS RootName, + anyLastIf(toNullable(InputPreview), ParentSpanId = '') AS RootInput, + anyLastIf(toNullable(StatusCode), ParentSpanId = '') AS RootStatus, + count() AS SpanCount, + countIf(ObservationType = 'agent') AS AgentCount, + countIf(ObservationType = 'llm') AS LlmCount, + countIf(ObservationType = 'llm' AND LiteLLMRequestId != '') AS IdentifiedLlmCount, + countIf(ObservationType = 'tool') AS ToolCount, + 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 +FROM {database}.otel_traces +GROUP BY TeamId, ApiKeyHash, TraceId diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index c146a6eac92..2affbd56e52 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -26,6 +26,8 @@ def trace_span_rows( body: bytes, content_type: str | None, tenant: Mapping[str, str], max_attribute_value_bytes: int ) -> list[dict[str, JsonValue]]: ... +class TraceReadError(Exception): ... + @final class NativeTraceConfig: def __new__( @@ -34,6 +36,7 @@ class NativeTraceConfig: url: str, retention_days: int, max_attribute_value_bytes: int, + cursor_keys: Sequence[str], ) -> NativeTraceConfig: ... @final @@ -46,7 +49,7 @@ class NativeTraceStorage: self, scope: TraceScope, start_ms: int, end_ms: int, cursor: str | None, limit: int ) -> Future[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 | None ) -> Future[JsonValue]: ... def get_span(self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str) -> Future[JsonValue]: ... def get_span_error( @@ -358,6 +361,7 @@ __all__ = [ "RustUpstreamError", "TokenCounter", "Tokenizer", + "TraceReadError", "achat_completions", "acompletion", "aembedding", diff --git a/litellm/tracing/config.py b/litellm/tracing/config.py index ef537ba2e67..eeb54411636 100644 --- a/litellm/tracing/config.py +++ b/litellm/tracing/config.py @@ -58,9 +58,7 @@ def _cursor_keys(store: Mapping[str, object], environ: Mapping[str, str]) -> tup 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,) + return (fallback,) if fallback else () try: candidates: Final = CURSOR_KEYS.validate_python((supplied,) if isinstance(supplied, str) else supplied) except ValidationError as error: @@ -108,3 +106,10 @@ def trace_storage_config(settings: Mapping[str, object], environ: Mapping[str, s ), cursor_keys=_cursor_keys(store, environ), ) + + +def trace_reader_config(settings: Mapping[str, object], environ: Mapping[str, str] = os.environ) -> TraceStorageConfig: + config: Final = trace_storage_config(settings, environ) + if not config.cursor_keys: + raise ValueError("tracing.store.cursor_keys, LITELLM_SALT_KEY or LITELLM_MASTER_KEY is required") + return config diff --git a/litellm/tracing/receiver.py b/litellm/tracing/receiver.py index 0c659aa94e4..36b3892b3a3 100644 --- a/litellm/tracing/receiver.py +++ b/litellm/tracing/receiver.py @@ -26,7 +26,7 @@ from litellm.constants import ( ) 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 +from litellm.tracing.config import trace_reader_config from litellm.tracing.otlp_http import InvalidOTLPPayloadError, TracingPayloadTooLargeError, decompress @@ -55,7 +55,7 @@ class TraceReceiver: @classmethod def from_settings(cls, settings: Mapping[str, object]) -> "TraceReceiver": - return cls(storage=ClickHouseStorage(trace_storage_config(settings))) + return cls(storage=ClickHouseStorage(trace_reader_config(settings))) async def start(self) -> None: await self.storage.ensure_schema() diff --git a/tests/unit/tracing/test_config.py b/tests/unit/tracing/test_config.py index 363a6d6db67..465b95c112d 100644 --- a/tests/unit/tracing/test_config.py +++ b/tests/unit/tracing/test_config.py @@ -1,7 +1,7 @@ import pytest from litellm import constants -from litellm.tracing.config import is_clickhouse_tracing_enabled, trace_storage_config +from litellm.tracing.config import is_clickhouse_tracing_enabled, trace_reader_config, trace_storage_config @pytest.mark.parametrize( @@ -124,7 +124,11 @@ 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")), + ( + ["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",)), @@ -146,6 +150,10 @@ def test_invalid_cursor_keys_are_rejected(store_keys: object) -> None: trace_storage_config({"store": store}, {"LITELLM_SALT_KEY": "salt"}) -def test_missing_cursor_key_source_is_rejected() -> None: +def test_missing_cursor_key_source_is_rejected_for_readers() -> 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"}}, {}) + trace_reader_config({"store": {"type": "clickhouse", "url": "http://localhost:8123"}}, {}) + + +def test_write_only_storage_does_not_need_cursor_keys() -> None: + assert trace_storage_config({"store": {"type": "clickhouse", "url": "http://localhost:8123"}}, {}).cursor_keys == () diff --git a/ui/litellm-dashboard/src/components/lens/LensWorkspace.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/LensWorkspace.integration.test.tsx index 87adbe2adb3..07851eb5d8c 100644 --- a/ui/litellm-dashboard/src/components/lens/LensWorkspace.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/LensWorkspace.integration.test.tsx @@ -95,7 +95,13 @@ describe("Lens interactive demo", () => { const path = new URL(String(input), "http://localhost").pathname; if (path === "/lens") return Response.json({ lenses: [saved], workers: [], tracing_enabled: true }); if (path.endsWith("/runs")) return Response.json(saved.jobs); - if (path === "/v1/traces") return Response.json({ data: data.runs.map((run) => run.trace.summary) }); + if (path === "/v1/traces") { + return Response.json({ + items: data.runs.map((run) => run.trace.summary), + next_cursor: null, + traversal: { id: "demo", published_at: "2026-01-01T00:00:00Z", expires_at: "2026-01-01T00:30:00Z" }, + }); + } return Response.json({ data: [], traces: true, requests: false }); }); renderWithProviders(, { From 91f00d04105bfbf42b57ca7df8755e8c6a4ad927 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 3 Oct 2026 23:32:32 +0000 Subject: [PATCH 3/4] fix(tracing): keep listed runs unique under backdated spans, fence oversized fallback, report changed diagnostics Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...ved.sql => 0016_agent_traces_received.sql} | 0 ....sql => 0017_agent_traces_mv_received.sql} | 0 .../traces-clickhouse/query/list_traces.sql | 4 +- .../traces-clickhouse/src/query/named.rs | 4 +- .../crates/traces-clickhouse/src/reads.rs | 150 +++++++++--- .../crates/traces-clickhouse/tests/reads.rs | 216 +++++++++++++++++- litellm-rust/crates/traces/src/query/named.rs | 2 + .../crates/traces/tests/query/named.rs | 2 +- litellm-rust/crates/traces/tests/resolve.rs | 1 + 9 files changed, 349 insertions(+), 30 deletions(-) rename litellm-rust/crates/traces-clickhouse/migrations/{0015_agent_traces_received.sql => 0016_agent_traces_received.sql} (100%) rename litellm-rust/crates/traces-clickhouse/migrations/{0016_agent_traces_mv_received.sql => 0017_agent_traces_mv_received.sql} (100%) diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0015_agent_traces_received.sql b/litellm-rust/crates/traces-clickhouse/migrations/0016_agent_traces_received.sql similarity index 100% rename from litellm-rust/crates/traces-clickhouse/migrations/0015_agent_traces_received.sql rename to litellm-rust/crates/traces-clickhouse/migrations/0016_agent_traces_received.sql diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0016_agent_traces_mv_received.sql b/litellm-rust/crates/traces-clickhouse/migrations/0017_agent_traces_mv_received.sql similarity index 100% rename from litellm-rust/crates/traces-clickhouse/migrations/0016_agent_traces_mv_received.sql rename to litellm-rust/crates/traces-clickhouse/migrations/0017_agent_traces_mv_received.sql diff --git a/litellm-rust/crates/traces-clickhouse/query/list_traces.sql b/litellm-rust/crates/traces-clickhouse/query/list_traces.sql index 918f39e8f0c..746a7cb15b8 100644 --- a/litellm-rust/crates/traces-clickhouse/query/list_traces.sql +++ b/litellm-rust/crates/traces-clickhouse/query/list_traces.sql @@ -19,6 +19,7 @@ FROM agent_traces_by_key WHERE ({all_teams:UInt8} = 1 OR ({user_id:String} != '' AND UserIds = [{user_id:String}]) OR has({team_ids:Array(String)}, TeamId)) + AND ({snapshot_ms:UInt64} = 0 OR ReceivedMs <= {snapshot_ms:UInt64}) GROUP BY TeamId, ApiKeyHash, TraceId HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64}) AND min(StartTs) < fromUnixTimestamp64Milli({end_ms:Int64}) @@ -30,10 +31,11 @@ LIMIT {limit:UInt32} ) SELECT page.* EXCEPT (trace_start, trace_end), identities.agent_names AS agent_names, identities.agent_count AS agent_count, - identities.frameworks AS frameworks + identities.frameworks AS frameworks, identities.fenced_start_ms AS fenced_start_ms FROM page LEFT JOIN ( SELECT TeamId, ApiKeyHash, TraceId, + toUnixTimestamp64Milli(min(Timestamp)) AS fenced_start_ms, arraySort(groupUniqArrayIf(AgentName, AgentName != '')) AS agent_names, arraySort(groupUniqArrayIf(toString(Framework), Framework != '')) AS frameworks, uniqExactIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS agent_count diff --git a/litellm-rust/crates/traces-clickhouse/src/query/named.rs b/litellm-rust/crates/traces-clickhouse/src/query/named.rs index 4f54f427897..91624411de3 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query/named.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query/named.rs @@ -48,6 +48,8 @@ struct ListTracesRowEncoding { pub status: litellm_traces::SpanStatus, #[serde(deserialize_with = "super::number::deserialize")] pub start_ms: i64, + #[serde(default, deserialize_with = "super::number::deserialize")] + pub fenced_start_ms: i64, #[serde(deserialize_with = "super::number::deserialize")] pub duration_ms: i64, #[serde(deserialize_with = "super::number::deserialize")] @@ -341,7 +343,7 @@ mod tests { #[case::quoted(true)] fn rows_decode_into_neutral_contracts(#[case] quoted: bool) { round_trip::( - json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["claude-agent-sdk"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}), + json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "fenced_start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["claude-agent-sdk"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}), quoted, ); round_trip::( diff --git a/litellm-rust/crates/traces-clickhouse/src/reads.rs b/litellm-rust/crates/traces-clickhouse/src/reads.rs index 27c4e09b6dd..6ee2112155d 100644 --- a/litellm-rust/crates/traces-clickhouse/src/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/src/reads.rs @@ -133,6 +133,16 @@ impl ListTraversal { ) } + /// A backdated span received after publication can lower a listed run's rollup start below + /// the cursor; its fenced start still sits at or above the cursor, so the run is not repeated. + fn unseen(&self, row: &contracts::ListTracesRow) -> bool { + let cursor = &self.cursor.position; + cursor.start_ms == 0 + || row.fenced_start_ms == 0 + || (row.fenced_start_ms, row.trace_ref.as_str()) + < (cursor.start_ms, cursor.trace_ref.as_str()) + } + fn continue_after( &self, keys: &KeyRing, @@ -275,17 +285,24 @@ pub async fn list_traces( .iter() .map(|row| (row.trace_ref.clone(), row.start_ms)) .collect(); - let items: Vec = stream::iter(page.chunks(16)) + let last_candidate = page.last().map(|row| ListPosition { + start_ms: row.start_ms, + trace_ref: row.trace_ref.clone(), + }); + let unseen: Vec = page + .into_iter() + .filter(|row| traversal.unseen(row)) + .collect(); + let items: Vec = stream::iter(unseen.chunks(16)) .then(|batch| list_summaries(client, connection, access, batch, snapshot_ms)) .try_collect::>() .await? .into_iter() .flatten() .collect(); - let next_cursor = items - .last() + let next_cursor = last_candidate .filter(|_| !exhausted) - .map(|last| traversal.continue_after(keys, &starts, last)) + .map(|last| keys.encode(&traversal.binding, &traversal.cursor.advance(last))) .transpose()?; let page = Page { items, @@ -318,27 +335,34 @@ 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, snapshot_ms) - .await - { - Ok(spans) => spans, - Err(Error::ReadTooLarge) => { - return stream::iter(runs) - .then(|row| async move { - match get_trace(client, connection, access, &row.trace_id, &row.trace_ref).await - { - Ok(trace) => { - Ok(trace.map_or_else(|| listed_summary(row), |trace| trace.summary)) + 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) + .then(|row| async move { + match read_trace( + client, + connection, + access, + &row.trace_id, + &row.trace_ref, + snapshot_ms, + ) + .await + { + Ok(trace) => Ok( + trace.map_or_else(|| listed_summary(row), |trace| trace.summary) + ), + Err(Error::ReadTooLarge) => Ok(listed_summary(row)), + Err(error) => Err(error), } - Err(Error::ReadTooLarge) => Ok(listed_summary(row)), - Err(error) => Err(error), - } - }) - .try_collect() - .await; - } - Err(error) => return Err(error), - }; + }) + .try_collect() + .await; + } + Err(error) => return Err(error), + }; let by_trace = spans.into_iter().into_group_map_by(|span| { ( span.team_id.clone(), @@ -376,6 +400,17 @@ pub async fn get_trace( access: &ReadAccessParams, trace_id: &str, trace_ref: &str, +) -> Result, Error> { + read_trace(client, connection, access, trace_id, trace_ref, u64::MAX).await +} + +async fn read_trace( + client: &Client, + connection: &Connection, + access: &ReadAccessParams, + trace_id: &str, + trace_ref: &str, + snapshot_ms: u64, ) -> Result, Error> { let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else { return Ok(None); @@ -385,7 +420,7 @@ pub async fn get_trace( trace_id: trace_id.to_owned(), trace_ref: trace_ref.clone(), }; - let rows = crate::span_batches::read_spans(client, connection, params, u64::MAX).await?; + let rows = crate::span_batches::read_spans(client, connection, params, snapshot_ms).await?; if rows.is_empty() { return Ok(None); } @@ -566,7 +601,7 @@ pub async fn get_span_error( trace_ref, span_id: span_id.to_owned(), error_offset: offset, - error_version: position.revision.clone(), + error_version: String::new(), }); let Some(row) = fetch::(client, connection, ¶ms) .await? @@ -625,6 +660,68 @@ mod tests { OffsetDateTime::from_unix_timestamp(1_790_000_000).unwrap() } + fn listed_row( + trace_ref: &str, + start_ms: i64, + fenced_start_ms: i64, + ) -> contracts::ListTracesRow { + contracts::ListTracesRow { + trace_id: "t".into(), + trace_ref: trace_ref.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, + fenced_start_ms, + 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(), + } + } + + #[rstest] + #[case::first_page_keeps_everything(None, 50, 900, true)] + #[case::earlier_fenced_start(Some((500, "b")), 50, 400, true)] + #[case::same_start_lower_ref(Some((500, "b")), 500, 500, true)] + #[case::unknown_fenced_start(Some((500, "b")), 50, 0, true)] + #[case::backdated_after_publication(Some((500, "b")), 50, 900, false)] + #[case::same_start_same_ref(Some((500, "a")), 50, 500, false)] + fn unseen_drops_runs_whose_fenced_start_was_already_listed( + #[case] cursor: Option<(i64, &str)>, + #[case] start_ms: i64, + #[case] fenced_start_ms: i64, + #[case] expected: bool, + ) { + let keys = keys(); + let token = cursor.map(|(start_ms, trace_ref)| { + let first = ListTraversal::open(&keys, &access(), (0, 1000), None, now()).unwrap(); + let position = ListPosition { + start_ms, + trace_ref: trace_ref.into(), + }; + keys.encode(&first.binding, &first.cursor.advance(position)) + .unwrap() + }); + let traversal = + ListTraversal::open(&keys, &access(), (0, 1000), token.as_deref(), now()).unwrap(); + let row = listed_row("a", start_ms, fenced_start_ms); + assert_eq!(traversal.unseen(&row), expected); + } + #[rstest] fn list_traversal_pins_publication_and_continues_after_the_last_run() { let keys = keys(); @@ -641,6 +738,7 @@ mod tests { input_preview: String::new(), status: litellm_traces::SpanStatus::Ok, start_ms: 1_790_742_989_377, + fenced_start_ms: 1_790_742_989_377, duration_ms: 1, span_count: 1, agent_count: 0, diff --git a/litellm-rust/crates/traces-clickhouse/tests/reads.rs b/litellm-rust/crates/traces-clickhouse/tests/reads.rs index e8aa9386087..e05ae483c6c 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/reads.rs @@ -2,7 +2,8 @@ use std::collections::BTreeMap; use litellm_traces::query::named::ReadAccessParams; use litellm_traces_clickhouse::{ - Connection, InsertTable, QueryScope, get_trace, get_trace_page, insert_rows, list_traces, + Connection, InsertTable, QueryScope, get_span_error, get_trace, get_trace_page, insert_rows, + list_traces, }; use rstest::rstest; use serde_json::json; @@ -18,6 +19,26 @@ fn keys() -> litellm_pagination::KeyRing { litellm_pagination::KeyRing::new(["fixture-cursor-secret"]).unwrap() } +/// `insert_rows` stamps `EngineReceivedMs` with the wall clock; tests that pin a receipt time +/// write the row directly. +async fn insert_received( + client: &litellm_http::Client, + writer: &Connection, + row: serde_json::Value, +) -> TestResult { + let response = client + .post(writer.url().clone()) + .body(format!( + "INSERT INTO {DATABASE}.otel_traces FORMAT JSONEachRow\n{row}" + )) + .send() + .await?; + if !response.status().is_success() { + return Err(response.text().await?.into()); + } + Ok(()) +} + #[rstest] #[case::api_key("key-a", "")] #[case::user("", "user-a")] @@ -628,6 +649,21 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( ])], ) .await?; + insert_received( + client, + &writer, + json!({ + "Timestamp": "2026-09-21 00:00:00.000000000", + "TraceId": run.trace_id, + "SpanId": "late-child", + "ParentSpanId": "0101010101010101", + "ObservationType": "tool", + "TeamId": "team-a", + "ApiKeyHash": "key-a", + "EngineReceivedMs": u64::MAX / 2 + }), + ) + .await?; let after = list_traces( client, &reader, @@ -670,3 +706,181 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( )); Ok(()) } + +#[rstest] +#[tokio::test] +async fn a_backdated_span_received_after_publication_does_not_repeat_a_listed_run( + #[future(awt)] seeded_database: TestResult, +) -> TestResult { + let fixture = seeded_database?; + let client = &fixture.database.client; + let reader = fixture + .readers + .connection(client, &QueryScope::All, "fixture-secret") + .await?; + let access = ReadAccessParams { + all_teams: true, + user_id: String::new(), + team_ids: Vec::new(), + }; + let mut cursor = None; + let mut listed = Vec::new(); + let fixture_run = loop { + let page = list_traces( + client, + &reader, + &keys(), + &access, + 0, + 2_000_000_000_000, + cursor.as_deref(), + 1, + ) + .await?; + let run = page + .items + .into_iter() + .next() + .ok_or("fixture run not listed")?; + cursor = page.next_cursor; + listed.push(run.trace_ref.clone()); + if run.span_count == 3 { + break run; + } + }; + let continuation = cursor.clone().ok_or("fixture run was the last run")?; + let writer = Connection::writer(&fixture.database.url)?; + insert_received( + client, + &writer, + json!({ + "Timestamp": "2020-01-01 00:00:00.000000000", + "TraceId": fixture_run.trace_id, + "SpanId": "backdated-root", + "TeamId": "team-a", + "ApiKeyHash": "key-a", + "EngineReceivedMs": u64::MAX / 2 + }), + ) + .await?; + for statement in [ + format!("SYSTEM START MERGES {DATABASE}.agent_traces_by_key"), + format!("OPTIMIZE TABLE {DATABASE}.agent_traces_by_key FINAL"), + ] { + let response = client + .post(fixture.database.url.clone()) + .body(statement) + .send() + .await?; + if !response.status().is_success() { + return Err(response.text().await?.into()); + } + } + let mut cursor = Some(continuation); + while let Some(current) = cursor { + let page = list_traces( + client, + &reader, + &keys(), + &access, + 0, + 2_000_000_000_000, + Some(¤t), + 1, + ) + .await?; + listed.extend(page.items.into_iter().map(|run| run.trace_ref)); + cursor = page.next_cursor; + } + let mut unique = listed.clone(); + unique.sort(); + unique.dedup(); + assert_eq!(unique.len(), listed.len(), "{listed:?}"); + let fresh = list_traces( + client, + &reader, + &keys(), + &access, + 0, + 2_000_000_000_000, + None, + 50, + ) + .await?; + let fresh_listings = fresh + .items + .iter() + .filter(|run| run.trace_ref == fixture_run.trace_ref) + .count(); + assert_eq!(fresh_listings, 1); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn a_changed_diagnostic_reports_traversal_changed_instead_of_not_found( + #[future(awt)] seeded_database: TestResult, +) -> TestResult { + let fixture = seeded_database?; + let client = &fixture.database.client; + let reader = fixture + .readers + .connection(client, &QueryScope::All, "fixture-secret") + .await?; + let access = ReadAccessParams { + all_teams: true, + user_id: String::new(), + team_ids: Vec::new(), + }; + let writer = Connection::writer(&fixture.database.url)?; + let diagnostic = |received: u64, message: String| { + json!({ + "Timestamp": "2026-09-21 00:00:00.000000000", + "TraceId": "diagnostic-trace", + "SpanId": "diagnostic-span", + "StatusCode": "STATUS_CODE_ERROR", + "StatusMessage": message, + "TeamId": "team-a", + "ApiKeyHash": "key-a", + "EngineReceivedMs": received + }) + }; + insert_received(client, &writer, diagnostic(100, "a".repeat(40_000))).await?; + let first = get_span_error( + client, + &reader, + &keys(), + &access, + "diagnostic-trace", + "diagnostic-span", + "", + None, + ) + .await? + .ok_or("missing diagnostic")?; + let continuation = first.next_cursor.ok_or("diagnostic fit in one page")?; + insert_received(client, &writer, diagnostic(50, "b".repeat(40_000))).await?; + let uncached_reader = + Connection::reader(&format!("{}?max_threads=1", fixture.database.url), DATABASE)?; + let changed = get_span_error( + client, + &uncached_reader, + &keys(), + &access, + "diagnostic-trace", + "diagnostic-span", + "", + Some(&continuation), + ) + .await; + assert!( + matches!( + changed, + Err(litellm_traces_clickhouse::Error::Pagination( + litellm_pagination::Error::TraversalChanged + )) + ), + "{changed:?}" + ); + Ok(()) +} diff --git a/litellm-rust/crates/traces/src/query/named.rs b/litellm-rust/crates/traces/src/query/named.rs index e9df28ee9c8..7f08693965c 100644 --- a/litellm-rust/crates/traces/src/query/named.rs +++ b/litellm-rust/crates/traces/src/query/named.rs @@ -41,6 +41,8 @@ pub struct ListTracesRow { #[serde(serialize_with = "crate::wire::serialize_status")] pub status: crate::SpanStatus, pub start_ms: i64, + #[serde(default)] + pub fenced_start_ms: i64, pub duration_ms: i64, pub span_count: u64, pub agent_count: u64, diff --git a/litellm-rust/crates/traces/tests/query/named.rs b/litellm-rust/crates/traces/tests/query/named.rs index 7db1a5895f7..4fca4281674 100644 --- a/litellm-rust/crates/traces/tests/query/named.rs +++ b/litellm-rust/crates/traces/tests/query/named.rs @@ -55,7 +55,7 @@ fn named_requests_preserve_all_access_cases( #[rstest] fn result_contracts_preserve_public_field_names() { round_trip::( - json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["framework"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}), + json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "fenced_start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["framework"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}), ); round_trip::( json!({"trace_id": "trace", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "framework", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}), diff --git a/litellm-rust/crates/traces/tests/resolve.rs b/litellm-rust/crates/traces/tests/resolve.rs index 367f1d146ee..891f03c62fe 100644 --- a/litellm-rust/crates/traces/tests/resolve.rs +++ b/litellm-rust/crates/traces/tests/resolve.rs @@ -577,6 +577,7 @@ fn listed_summary_keeps_rollup_counts_with_unknown_cost() { input_preview: "hi".into(), status: SpanStatus::Ok, start_ms: 1_790_742_989_377, + fenced_start_ms: 1_790_742_989_377, duration_ms: 51_385, span_count: 126, agent_count: 2, From 7857236c60bdc6955ab613a3e66418e3ff758d41 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 4 Oct 2026 17:13:28 +0000 Subject: [PATCH 4/4] fix(tracing): type the shared read failure response map at module scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/tracing_endpoints.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/tracing_endpoints.py b/litellm/proxy/tracing_endpoints.py index 9d298735bdc..93621753e8d 100644 --- a/litellm/proxy/tracing_endpoints.py +++ b/litellm/proxy/tracing_endpoints.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from functools import partial from http.client import responses from types import MappingProxyType -from typing import Annotated, Final +from typing import Annotated, Final, TypeAlias from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response from pydantic import BaseModel, ConfigDict @@ -68,8 +68,10 @@ class ReadFailureResponse(BaseModel): detail: ReadFailure -def _read_failure_responses() -> dict[int | str, dict[str, object]]: - return {status: {"model": ReadFailureResponse} for status in (400, 409, 410, 413, 503)} +_ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg +READ_FAILURE_RESPONSES: Final[_ResponseDocSchemas] = { + status: {"model": ReadFailureResponse} for status in (400, 409, 410, 413, 503) +} def _read_failure(error: TraceReadError) -> HTTPException: @@ -171,7 +173,7 @@ async def ingest_otlp_traces( return Response(content=body, media_type=media_type) -@router.get("/v1/traces", response_model=TracePage, responses=_read_failure_responses()) +@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, @@ -262,7 +264,7 @@ 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, responses=_read_failure_responses()) +@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)], @@ -283,7 +285,7 @@ async def get_agent_trace( return trace -@router.get("/v1/traces/{trace_id}/spans/{span_id}", response_model=SpanDetail, responses=_read_failure_responses()) +@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, @@ -304,7 +306,7 @@ async def get_agent_trace_span( @router.get( - "/v1/traces/{trace_id}/spans/{span_id}/error", response_model=SpanErrorPage, responses=_read_failure_responses() + "/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,