mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Merge 7857236c60 into e1d16f51d1
This commit is contained in:
commit
d98326acbe
65 changed files with 2632 additions and 408 deletions
19
litellm-rust/Cargo.lock
generated
19
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
@ -4458,6 +4475,7 @@ dependencies = [
|
|||
"criterion",
|
||||
"indexmap 2.14.0",
|
||||
"litellm-llms-types",
|
||||
"litellm-pagination",
|
||||
"macro_rules_attribute",
|
||||
"opentelemetry-proto",
|
||||
"prost",
|
||||
|
|
@ -4496,6 +4514,7 @@ dependencies = [
|
|||
"jsonschema",
|
||||
"litellm-http",
|
||||
"litellm-migrate",
|
||||
"litellm-pagination",
|
||||
"litellm-storage-clickhouse",
|
||||
"litellm-traces",
|
||||
"litellm-traces-cache",
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
|
|||
5
litellm-rust/crates/pagination/AGENTS.md
Normal file
5
litellm-rust/crates/pagination/AGENTS.md
Normal file
|
|
@ -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/`
|
||||
23
litellm-rust/crates/pagination/Cargo.toml
Normal file
23
litellm-rust/crates/pagination/Cargo.toml
Normal file
|
|
@ -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
|
||||
211
litellm-rust/crates/pagination/src/cursor.rs
Normal file
211
litellm-rust/crates/pagination/src/cursor.rs
Normal file
|
|
@ -0,0 +1,211 @@
|
|||
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<u8>,
|
||||
}
|
||||
|
||||
/// 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<Key>,
|
||||
}
|
||||
|
||||
impl KeyRing {
|
||||
pub fn new<S: AsRef<[u8]>>(secrets: impl IntoIterator<Item = S>) -> Result<Self, Error> {
|
||||
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::<Result<Vec<_>, _>>()?;
|
||||
Ok(Self { keys })
|
||||
}
|
||||
|
||||
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<u8> {
|
||||
let mut mac =
|
||||
Hmac::<Sha256>::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<P: Serialize>(
|
||||
&self,
|
||||
binding: &Binding,
|
||||
cursor: &Cursor<P>,
|
||||
) -> Result<String, Error> {
|
||||
let signing_key = self.signing_key()?;
|
||||
let payload = serde_json::to_vec(&Envelope {
|
||||
version: FORMAT_VERSION,
|
||||
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(signing_key, 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<P: DeserializeOwned>(
|
||||
&self,
|
||||
binding: &Binding,
|
||||
token: &str,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<Cursor<P>, Error> {
|
||||
self.signing_key()?;
|
||||
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<RawValue>> =
|
||||
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<Self, Error> {
|
||||
Ok(Self {
|
||||
resource: resource.to_owned(),
|
||||
scope: digest(scope)?,
|
||||
query: digest(query)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn fingerprint(&self) -> Vec<u8> {
|
||||
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<String, Error> {
|
||||
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<P> {
|
||||
pub position: P,
|
||||
pub revision: String,
|
||||
pub published_ms: u64,
|
||||
pub expires_at: OffsetDateTime,
|
||||
}
|
||||
|
||||
impl<P> Cursor<P> {
|
||||
/// 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<Q>(&self, position: Q) -> Cursor<Q> {
|
||||
Cursor {
|
||||
position,
|
||||
revision: self.revision.clone(),
|
||||
published_ms: self.published_ms,
|
||||
expires_at: self.expires_at,
|
||||
}
|
||||
}
|
||||
}
|
||||
96
litellm-rust/crates/pagination/src/error.rs
Normal file
96
litellm-rust/crates/pagination/src/error.rs
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
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<String>) -> 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 signing keys are not configured")]
|
||||
MissingKeys,
|
||||
#[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::MissingKeys | Self::Encode(_) => FailureCode::Unavailable,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&Error> for Failure {
|
||||
fn from(error: &Error) -> Self {
|
||||
Self::new(error.code(), error.to_string())
|
||||
}
|
||||
}
|
||||
10
litellm-rust/crates/pagination/src/lib.rs
Normal file
10
litellm-rust/crates/pagination/src/lib.rs
Normal file
|
|
@ -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};
|
||||
75
litellm-rust/crates/pagination/src/page.rs
Normal file
75
litellm-rust/crates/pagination/src/page.rs
Normal file
|
|
@ -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<T> {
|
||||
pub items: Vec<T>,
|
||||
pub next_cursor: Option<String>,
|
||||
pub traversal: Traversal,
|
||||
}
|
||||
|
||||
impl<T: Serialize> Page<T> {
|
||||
/// 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<String, Error>,
|
||||
) -> Result<Self, Error> {
|
||||
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()
|
||||
}
|
||||
223
litellm-rust/crates/pagination/tests/cursor.rs
Normal file
223
litellm-rust/crates/pagination/tests/cursor.rs
Normal file
|
|
@ -0,0 +1,223 @@
|
|||
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<Position> {
|
||||
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<Position> = 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::<Position>(&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::<Position>(&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::<Position>(&binding, &forged, now),
|
||||
Err(Error::InvalidCursor)
|
||||
));
|
||||
let other_keys = KeyRing::new(["other-secret"]).unwrap();
|
||||
assert!(matches!(
|
||||
other_keys.decode::<Position>(&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::<Position>(&binding, &token, now + Duration::minutes(29))
|
||||
.is_ok()
|
||||
);
|
||||
assert!(matches!(
|
||||
keys.decode::<Position>(&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::<Position>(&binding, &token, now).is_ok());
|
||||
let fresh = rotated.encode(&binding, &cursor(now)).unwrap();
|
||||
assert!(matches!(
|
||||
old.decode::<Position>(&binding, &fresh, now),
|
||||
Err(Error::InvalidCursor)
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[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::<Position>(&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);
|
||||
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)]
|
||||
#[case::missing_keys(Error::MissingKeys, 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());
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
@ -41,9 +48,8 @@ fn map_error_ref(error: &Error) -> PyErr {
|
|||
Error::InvalidRow
|
||||
| Error::InvalidLimit(_)
|
||||
| Error::InvalidTable
|
||||
| Error::InvalidCursor(_)
|
||||
| Error::Pagination(_)
|
||||
| Error::AmbiguousTrace
|
||||
| Error::TraceChanged
|
||||
| Error::Decode(_)
|
||||
| Error::InvalidSchema
|
||||
| Error::InvalidQuery
|
||||
|
|
@ -78,6 +84,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)) => {
|
||||
|
|
@ -100,10 +110,17 @@ impl NativeTraceConfig {
|
|||
url: &str,
|
||||
retention_days: u32,
|
||||
max_attribute_value_bytes: usize,
|
||||
cursor_keys: Vec<String>,
|
||||
) -> PyResult<Self> {
|
||||
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)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -218,12 +235,14 @@ impl NativeTraceStorage {
|
|||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
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,
|
||||
|
|
@ -232,11 +251,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>,
|
||||
|
|
@ -244,38 +263,27 @@ impl NativeTraceStorage {
|
|||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams,
|
||||
trace_ref: String,
|
||||
cursor: Option<String>,
|
||||
page_size: Option<u32>,
|
||||
page_size: u32,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
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,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -302,7 +310,7 @@ impl NativeTraceStorage {
|
|||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
map_read_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -318,12 +326,14 @@ impl NativeTraceStorage {
|
|||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
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,
|
||||
|
|
@ -332,7 +342,7 @@ impl NativeTraceStorage {
|
|||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
map_read_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -495,9 +505,15 @@ mod tests {
|
|||
Error::Decode(litellm_traces::Error::InvalidLimit("OTLP_MAX_SPANS")),
|
||||
"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,
|
||||
|
|
@ -511,4 +527,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));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
ALTER TABLE {database}.agent_traces_by_key
|
||||
ADD COLUMN IF NOT EXISTS ReceivedMs SimpleAggregateFunction(min, UInt64) DEFAULT 0
|
||||
|
|
@ -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
|
||||
|
|
@ -19,20 +19,23 @@ 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})
|
||||
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}
|
||||
)
|
||||
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
|
||||
|
|
@ -41,6 +44,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
|
||||
|
|
|
|||
|
|
@ -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<KeyRing>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
|
|
@ -14,12 +18,14 @@ impl Config {
|
|||
url: &str,
|
||||
retention_days: u32,
|
||||
max_attribute_value_bytes: usize,
|
||||
cursor_keys: Vec<String>,
|
||||
) -> Result<Self, Error> {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -34,12 +34,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")]
|
||||
|
|
@ -58,3 +56,56 @@ impl From<litellm_traces_cache::Error> 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::InvalidLimit(_)
|
||||
| Self::InvalidTable
|
||||
| Self::InvalidSchema
|
||||
| Self::InvalidQuery
|
||||
| Self::InvalidParameters
|
||||
| Self::InvalidScope
|
||||
| Self::AmbiguousTrace
|
||||
| Self::Decode(_)
|
||||
| Self::Storage(
|
||||
StorageError::InvalidRow
|
||||
| StorageError::InvalidLimit(_)
|
||||
| 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
@ -46,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")]
|
||||
|
|
@ -340,7 +344,7 @@ mod tests {
|
|||
#[case::quoted(true)]
|
||||
fn rows_decode_into_neutral_contracts(#[case] quoted: bool) {
|
||||
round_trip::<ListTracesRow>(
|
||||
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::<TraceSpansRow>(
|
||||
|
|
@ -366,7 +370,7 @@ mod tests {
|
|||
#[case::quoted(true)]
|
||||
fn parameters_preserve_flattened_multi_team_access(#[case] quoted: bool) {
|
||||
round_trip::<ListTracesParams>(
|
||||
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::<SpanErrorParams>(
|
||||
|
|
|
|||
|
|
@ -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,131 @@ static TRACE_SNAPSHOTS: LazyLock<SnapshotCache> = 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<T: Serialize>(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<T: for<'de> Deserialize<'de>>(
|
||||
cursor: &str,
|
||||
kind: &'static str,
|
||||
) -> Result<T, Error> {
|
||||
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<Binding, Error> {
|
||||
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<P: for<'de> Deserialize<'de>>(
|
||||
keys: &KeyRing,
|
||||
binding: &Binding,
|
||||
cursor: Option<&str>,
|
||||
first: P,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<Cursor<P>, 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<ListPosition>,
|
||||
binding: Binding,
|
||||
}
|
||||
|
||||
impl ListTraversal {
|
||||
fn open(
|
||||
keys: &KeyRing,
|
||||
access: &ReadAccessParams,
|
||||
window: (i64, i64),
|
||||
cursor: Option<&str>,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<Self, Error> {
|
||||
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,
|
||||
)
|
||||
}
|
||||
|
||||
/// 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,
|
||||
starts: &BTreeMap<String, i64>,
|
||||
last: &litellm_traces::TraceSummary,
|
||||
) -> Result<String, litellm_pagination::Error> {
|
||||
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<Option<ErrorPosition>, Error> {
|
||||
let Some(cursor) = cursor else {
|
||||
return Ok(None);
|
||||
};
|
||||
let position = decode_cursor::<ErrorPosition>(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 +235,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<TracePage, Error> {
|
||||
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<contracts::ListTracesRow> = loop {
|
||||
match fetch::<RunCandidates>(client, connection, ¶ms).await {
|
||||
|
|
@ -186,18 +280,38 @@ 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<String, i64> = page
|
||||
.iter()
|
||||
.map(|row| (row.trace_ref.clone(), row.start_ms))
|
||||
.collect();
|
||||
let last_candidate = page.last().map(|row| ListPosition {
|
||||
start_ms: row.start_ms,
|
||||
trace_ref: row.trace_ref.clone(),
|
||||
});
|
||||
let unseen: Vec<contracts::ListTracesRow> = page
|
||||
.into_iter()
|
||||
.filter(|row| traversal.unseen(row))
|
||||
.collect();
|
||||
let items: Vec<litellm_traces::TraceSummary> = stream::iter(unseen.chunks(16))
|
||||
.then(|batch| list_summaries(client, connection, access, batch, snapshot_ms))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect();
|
||||
Ok(TracePage { data, next_cursor })
|
||||
let next_cursor = last_candidate
|
||||
.filter(|_| !exhausted)
|
||||
.map(|last| keys.encode(&traversal.binding, &traversal.cursor.advance(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 +319,7 @@ async fn list_summaries(
|
|||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
runs: &[contracts::ListTracesRow],
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Vec<litellm_traces::TraceSummary>, Error> {
|
||||
let (Some(start_ms), Some(end_ms)) = (
|
||||
runs.iter().map(|row| row.start_ms).min(),
|
||||
|
|
@ -220,25 +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).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(),
|
||||
|
|
@ -276,6 +400,17 @@ pub async fn get_trace(
|
|||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Option<Trace>, 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<Option<Trace>, Error> {
|
||||
let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
|
|
@ -285,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);
|
||||
}
|
||||
|
|
@ -293,17 +428,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 +448,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 +473,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 +482,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<Option<String>, 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 +559,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<Option<SpanErrorPage>, 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: String::new(),
|
||||
});
|
||||
let Some(row) = fetch::<SpanError>(client, connection, ¶ms)
|
||||
.await?
|
||||
|
|
@ -465,13 +611,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 +644,239 @@ 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()
|
||||
}
|
||||
|
||||
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]
|
||||
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()
|
||||
#[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();
|
||||
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,
|
||||
fenced_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::<ErrorPosition>(
|
||||
&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<SpanPosition> = 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::<SpanPosition>(
|
||||
&keys,
|
||||
&other_run,
|
||||
Some(&token),
|
||||
SpanPosition { offset: 0 },
|
||||
now()
|
||||
),
|
||||
Err(Error::Pagination(litellm_pagination::Error::InvalidCursor))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -122,6 +122,7 @@ pub(crate) async fn read_list_spans(
|
|||
client: &Client,
|
||||
connection: &Connection,
|
||||
runs: crate::query::named::TracePageSpansParams,
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Vec<contracts::TraceSpansRow>, 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()),
|
||||
|
|
|
|||
|
|
@ -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")?;
|
||||
|
|
|
|||
|
|
@ -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::<ListTraces>(&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::<ListTraces>(&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::<ListTraces>(&fixture.database.client, &reader, &next_params).await?;
|
||||
|
|
@ -272,6 +275,7 @@ async fn captured_sdk_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::<ListTraces>(&fixture.database.client, &reader, &list_params).await?;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
@ -14,6 +15,30 @@ 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()
|
||||
}
|
||||
|
||||
/// `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")]
|
||||
|
|
@ -83,11 +108,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")?;
|
||||
|
|
@ -208,10 +243,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)
|
||||
);
|
||||
|
|
@ -261,7 +306,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 {
|
||||
|
|
@ -276,7 +321,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")?
|
||||
|
|
@ -313,6 +358,7 @@ async fn large_runs_remain_complete_under_default_reader_limits(
|
|||
let page = get_trace_page(
|
||||
client,
|
||||
&reader,
|
||||
&keys(),
|
||||
&access,
|
||||
"trace-0000",
|
||||
trace_ref,
|
||||
|
|
@ -328,19 +374,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"))
|
||||
|
|
@ -381,15 +434,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,
|
||||
|
|
@ -433,19 +497,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
|
||||
|
|
@ -457,6 +524,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,
|
||||
|
|
@ -484,6 +552,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,
|
||||
|
|
@ -491,7 +560,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",
|
||||
|
|
@ -514,6 +585,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,
|
||||
|
|
@ -522,7 +594,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(())
|
||||
|
|
@ -544,9 +621,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")?;
|
||||
|
|
@ -571,10 +658,35 @@ 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());
|
||||
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,
|
||||
&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")?;
|
||||
|
|
@ -582,7 +694,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)
|
||||
|
|
@ -591,6 +703,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,
|
||||
|
|
@ -603,6 +716,183 @@ 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<SeededDatabase>,
|
||||
) -> 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<SeededDatabase>,
|
||||
) -> 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(())
|
||||
}
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback(
|
||||
|
|
@ -694,11 +984,21 @@ async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback
|
|||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".into()],
|
||||
};
|
||||
let page = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?;
|
||||
assert_eq!(page.data.len(), cases.len());
|
||||
let page = list_traces(
|
||||
client,
|
||||
&reader,
|
||||
&keys(),
|
||||
&access,
|
||||
0,
|
||||
2_000_000_000_000,
|
||||
None,
|
||||
50,
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(page.items.len(), cases.len());
|
||||
for (id, _, _, _, _, expected) in cases {
|
||||
let summary = page
|
||||
.data
|
||||
.items
|
||||
.iter()
|
||||
.find(|summary| summary.trace_id == id)
|
||||
.ok_or("missing run")?;
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
@ -39,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,
|
||||
|
|
|
|||
|
|
@ -211,6 +211,7 @@ pub fn resolve_trace(
|
|||
agents,
|
||||
spans,
|
||||
next_cursor: None,
|
||||
traversal: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -44,6 +44,12 @@ fn emitted<T: JsonSchema>() -> Schema {
|
|||
.into_root_schema_for::<T>()
|
||||
}
|
||||
|
||||
fn emitted_as<T: JsonSchema>(title: &str) -> Schema {
|
||||
let mut schema = emitted::<T>();
|
||||
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::<crate::QueryScope>()),
|
||||
("Tenant", received::<crate::Tenant>()),
|
||||
("TracePage", emitted::<crate::TracePage>()),
|
||||
("TracePage", emitted_as::<crate::TracePage>("TracePage")),
|
||||
("Trace", emitted::<crate::Trace>()),
|
||||
("SpanDetail", emitted::<crate::SpanDetail>()),
|
||||
("SpanErrorPage", emitted::<crate::SpanErrorPage>()),
|
||||
(
|
||||
"ReadFailure",
|
||||
emitted_as::<litellm_pagination::Failure>("ReadFailure"),
|
||||
),
|
||||
])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -90,14 +90,11 @@ pub struct Trace {
|
|||
pub spans: Vec<Span>,
|
||||
#[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))]
|
||||
pub next_cursor: Option<String>,
|
||||
#[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))]
|
||||
pub traversal: Option<litellm_pagination::Traversal>,
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub struct TracePage {
|
||||
pub data: Vec<TraceSummary>,
|
||||
pub next_cursor: Option<String>,
|
||||
}
|
||||
pub type TracePage = litellm_pagination::Page<TraceSummary>;
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[derive(Debug, PartialEq)]
|
||||
|
|
|
|||
|
|
@ -32,8 +32,13 @@ fn named_requests_preserve_all_access_cases(
|
|||
)
|
||||
};
|
||||
round_trip::<ListTracesParams>(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::<TraceIdentityParams>(request(json!({"trace_id": "trace"})));
|
||||
round_trip::<TraceSpansParams>(request(json!({"trace_id": "trace", "trace_ref": "ref"})));
|
||||
round_trip::<SpanDetailParams>(request(
|
||||
|
|
@ -50,7 +55,7 @@ fn named_requests_preserve_all_access_cases(
|
|||
#[rstest]
|
||||
fn result_contracts_preserve_public_field_names() {
|
||||
round_trip::<ListTracesRow>(
|
||||
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::<TraceSpansRow>(
|
||||
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"}),
|
||||
|
|
|
|||
|
|
@ -792,6 +792,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,
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -13,13 +13,13 @@ 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
|
||||
|
||||
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,49 @@ 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
|
||||
|
||||
|
||||
_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:
|
||||
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 +173,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 +189,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 +264,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 +285,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 +295,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 +305,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 +318,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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,28 @@ 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")
|
||||
return (fallback,) if fallback else ()
|
||||
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 +82,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 +104,12 @@ 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),
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -18,10 +18,15 @@ 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
|
||||
from litellm.tracing.config import trace_reader_config
|
||||
from litellm.tracing.otlp_http import InvalidOTLPPayloadError, TracingPayloadTooLargeError, decompress
|
||||
|
||||
|
||||
|
|
@ -50,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()
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
35
scripts/trace_codegen/schemas/traces/ReadFailure.json
Normal file
35
scripts/trace_codegen/schemas/traces/ReadFailure.json
Normal file
|
|
@ -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"
|
||||
}
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -508,7 +512,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"
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
@ -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,41 @@ 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_for_readers() -> None:
|
||||
with pytest.raises(ValueError, match="cursor_keys, LITELLM_SALT_KEY or LITELLM_MASTER_KEY is required"):
|
||||
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 == ()
|
||||
|
|
|
|||
|
|
@ -27,7 +27,12 @@ describe("Lens navigation", () => {
|
|||
"fetch",
|
||||
vi.fn(async (input) => {
|
||||
const path = new URL(String(input), "http://localhost").pathname;
|
||||
if (path === "/v1/traces") return Response.json({ data: [{}] });
|
||||
if (path === "/v1/traces")
|
||||
return Response.json({
|
||||
items: [{}],
|
||||
next_cursor: null,
|
||||
traversal: { id: "demo", published_at: "2026-01-01T00:00:00Z", expires_at: "2026-01-01T00:30:00Z" },
|
||||
});
|
||||
if (path === "/lens") return Response.json({ lenses: [], workers: [], tracing_enabled: true });
|
||||
return Response.json({ traces: true, requests: false, data: [] });
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -26,7 +26,11 @@ function serve({ enabled = false, traces = false, requests = false, connected =
|
|||
const path = new URL(String(input), "http://localhost").pathname;
|
||||
if (path === "/v1/traces")
|
||||
return enabled
|
||||
? Response.json({ data: traces ? [data.runs[0].trace.summary] : [] })
|
||||
? Response.json({
|
||||
items: traces ? [data.runs[0].trace.summary] : [],
|
||||
next_cursor: null,
|
||||
traversal: { id: "demo", published_at: "2026-01-01T00:00:00Z", expires_at: "2026-01-01T00:30:00Z" },
|
||||
})
|
||||
: Response.json({ detail: "Tracing is not enabled" }, { status: 501 });
|
||||
if (path === "/lens/activity/available") return Response.json({ traces, requests });
|
||||
if (path === "/lens" && init?.method === "POST") {
|
||||
|
|
|
|||
|
|
@ -176,7 +176,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(<LensWorkspace accessToken="live-token" userRole="Admin" readOnly={false} />, {
|
||||
|
|
@ -310,7 +316,12 @@ describe("Lens interactive demo", () => {
|
|||
network.mockImplementation(async (input) => {
|
||||
const path = new URL(String(input), "http://localhost").pathname;
|
||||
if (path === "/lens") return Response.json({ lenses: [], workers: [], tracing_enabled: true });
|
||||
if (path === "/v1/traces") return Response.json({ data: [{}] });
|
||||
if (path === "/v1/traces")
|
||||
return Response.json({
|
||||
items: [{}],
|
||||
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(<LensWorkspace accessToken="live-token" userRole="Admin" readOnly={false} />, {
|
||||
|
|
@ -349,7 +360,12 @@ describe("Lens interactive demo", () => {
|
|||
if (path === "/key/info") return Response.json({ info: { models: [], max_budget: null } });
|
||||
if (path === "/lens/agents") return Response.json([]);
|
||||
if (path.startsWith("/lens/preview")) return Response.json({ eligible: 0, selected: 0, executions: [] });
|
||||
if (path === "/v1/traces") return Response.json({ data: [createLensDemoData().runs[0].trace.summary] });
|
||||
if (path === "/v1/traces")
|
||||
return Response.json({
|
||||
items: [createLensDemoData().runs[0].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(<LensWorkspace accessToken="live-token" userRole="Admin" readOnly={false} />, {
|
||||
|
|
|
|||
|
|
@ -59,9 +59,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 });
|
||||
|
|
|
|||
|
|
@ -48,10 +48,11 @@ function demoTracesApi(data: LensDemoData): TracesApi {
|
|||
return { text: JSON.stringify(step, null, 2), copied: spanId ? "Step copied" : "Trace copied" };
|
||||
},
|
||||
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),
|
||||
|
|
|
|||
|
|
@ -34,6 +34,11 @@ beforeEach(() => {
|
|||
});
|
||||
|
||||
const executionId = btoa(JSON.stringify(["traces", "", "trace-42"]));
|
||||
const emptyTracePage = {
|
||||
items: [],
|
||||
next_cursor: null,
|
||||
traversal: { id: "demo", published_at: "2026-01-01T00:00:00Z", expires_at: "2026-01-01T00:30:00Z" },
|
||||
};
|
||||
const pattern: Finding = {
|
||||
reason: "",
|
||||
suggestion: "",
|
||||
|
|
@ -453,9 +458,10 @@ it.each([false, true])(
|
|||
async (enabled) => {
|
||||
window.history.replaceState({}, "", "/lens/");
|
||||
testQueryClient.clear();
|
||||
vi.mocked(apiClient.get).mockImplementation(async (path) =>
|
||||
path === "/lens" ? { lenses: [], workers: [], tracing_enabled: enabled } : { data: [] },
|
||||
);
|
||||
vi.mocked(apiClient.get).mockImplementation(async (path) => {
|
||||
if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: enabled };
|
||||
return path === "/v1/traces" ? emptyTracePage : { data: [] };
|
||||
});
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<InvestigationsView />);
|
||||
const guide = within(await screen.findByRole("region", { name: "Get Lens running" }));
|
||||
|
|
@ -505,7 +511,7 @@ it("allows retrying a failed trace readiness check without treating it as an emp
|
|||
vi.mocked(apiClient.get).mockImplementation(async (path) => {
|
||||
if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: true };
|
||||
if (path === "/lens/activity/available") return traceCheck();
|
||||
return { data: [] };
|
||||
return path === "/v1/traces" ? emptyTracePage : { data: [] };
|
||||
});
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<InvestigationsView />);
|
||||
|
|
@ -523,7 +529,10 @@ it("shows a centered failure with a retry when investigations cannot load, then
|
|||
.fn()
|
||||
.mockRejectedValueOnce(new ApiError("Proxy timed out", 504, {}))
|
||||
.mockResolvedValue({ lenses: [], workers: [], tracing_enabled: true });
|
||||
vi.mocked(apiClient.get).mockImplementation(async (path) => (path === "/lens" ? list() : { data: [] }));
|
||||
vi.mocked(apiClient.get).mockImplementation(async (path) => {
|
||||
if (path === "/lens") return list();
|
||||
return path === "/v1/traces" ? emptyTracePage : { data: [] };
|
||||
});
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<InvestigationsView />);
|
||||
const alert = await screen.findByRole("alert");
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
{
|
||||
"data": [
|
||||
"items": [
|
||||
{
|
||||
"trace_id": "e309a123963901e74c29cd2d3c86ff9e",
|
||||
"name": "research_lead",
|
||||
|
|
@ -55,5 +55,10 @@
|
|||
"agent_invocations": 1
|
||||
}
|
||||
],
|
||||
"next_cursor": null
|
||||
"next_cursor": null,
|
||||
"traversal": {
|
||||
"id": "d41d8cd98f00b204e9800998ecf8427e",
|
||||
"published_at": "2026-02-18T15:00:00.000Z",
|
||||
"expires_at": "2026-02-18T15:30:00.000Z"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ export function liveTracesApi(accessToken: string): TracesApi {
|
|||
list: (window) => agentTraceListCall({ accessToken, ...window }),
|
||||
anyRecorded: async () => {
|
||||
const page = await apiClient.get<TracePage>("/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),
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
@ -237,6 +239,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(<RoutedRunView traceId={research.summary.trace_id} accessToken="sk-test" onBack={vi.fn()} />);
|
||||
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 };
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import { DetailPane } from "./DetailPane";
|
|||
import { IdChip } from "../ui/IdChip";
|
||||
import { formatCost } from "../list/AgentTracesTable";
|
||||
import { SpanIcon } from "../ui/SpanIcon";
|
||||
import { restartsTraversal } from "../readFailure";
|
||||
import { SpanTree } from "./SpanTree";
|
||||
import { TraceConversation } from "./TraceConversation";
|
||||
import { FrameworkLogo, traceFramework } from "../ui/TraceFramework";
|
||||
|
|
@ -344,6 +345,12 @@ function selectedSpanMissing(trace: Trace, spanId: string | null): boolean {
|
|||
return Boolean(spanId && !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`;
|
||||
};
|
||||
|
||||
function RunLoading({ embedded }: { embedded: boolean }) {
|
||||
return (
|
||||
<div
|
||||
|
|
@ -441,6 +448,7 @@ function LoadedRun({
|
|||
useEffect(() => {
|
||||
if (canSeek && !isFetching && !isError) void fetchNextPage();
|
||||
}, [canSeek, isFetching, isError, fetchNextPage]);
|
||||
const mustRestart = restartsTraversal(traceQuery.error);
|
||||
const pageAction = isError ? "Retry" : "Load more steps";
|
||||
|
||||
return (
|
||||
|
|
@ -464,24 +472,27 @@ function LoadedRun({
|
|||
/>
|
||||
{(traceQuery.hasNextPage || traceQuery.isError) && (
|
||||
<div className="flex items-center justify-between gap-3 border-b px-3 py-2 text-xs" role="status">
|
||||
<span>
|
||||
{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`}
|
||||
</span>
|
||||
<span>{loadMoreStatus(trace, traceQuery.isError, mustRestart)}</span>
|
||||
{traceQuery.isError && (
|
||||
<Button size="xs" variant="ghost" disabled={traceQuery.isFetching} onClick={() => void refreshTrace()}>
|
||||
<Button
|
||||
size="xs"
|
||||
variant={mustRestart ? "outline" : "ghost"}
|
||||
disabled={traceQuery.isFetching}
|
||||
onClick={() => void refreshTrace()}
|
||||
>
|
||||
Refresh trace
|
||||
</Button>
|
||||
)}
|
||||
<Button
|
||||
size="xs"
|
||||
variant="outline"
|
||||
disabled={traceQuery.isFetching}
|
||||
onClick={() => void (traceQuery.hasNextPage ? traceQuery.fetchNextPage() : refreshTrace())}
|
||||
>
|
||||
{traceQuery.isFetching ? "Loading…" : pageAction}
|
||||
</Button>
|
||||
{!mustRestart && (
|
||||
<Button
|
||||
size="xs"
|
||||
variant="outline"
|
||||
disabled={traceQuery.isFetching}
|
||||
onClick={() => void (traceQuery.hasNextPage ? traceQuery.fetchNextPage() : refreshTrace())}
|
||||
>
|
||||
{traceQuery.isFetching ? "Loading…" : pageAction}
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
<RunBody
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import AgentTracesPage from "./AgentTracesPage";
|
|||
import { filterRuns } from "./runSearch/runQuery";
|
||||
import { AgentTracesSection, type TimeControls } from "./AgentTracesSection";
|
||||
import type { TracePage, TraceSummary } from "../types";
|
||||
import { RESULTS_CHANGED_MESSAGE } from "./useAgentTraces";
|
||||
|
||||
vi.mock("../../../networking", () => ({
|
||||
apiClient: { get: vi.fn(), post: vi.fn() },
|
||||
|
|
@ -33,7 +34,7 @@ vi.mock("../detail/TraceDrawer", () => ({
|
|||
|
||||
import { agentTraceListCall, apiClient } from "../../../networking";
|
||||
|
||||
const runs = (traceList as TracePage).data as TraceSummary[];
|
||||
const runs = (traceList as TracePage).items as TraceSummary[];
|
||||
|
||||
const lastUrl = (onUrlUpdate: ReturnType<typeof vi.fn>) =>
|
||||
new URLSearchParams(String(onUrlUpdate.mock.lastCall?.[0].queryString ?? ""));
|
||||
|
|
@ -90,13 +91,40 @@ describe("AgentTracesSection", () => {
|
|||
setupIntersectionMocking(vi.fn);
|
||||
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" } }),
|
||||
);
|
||||
act(() => mockAllIsIntersecting(true));
|
||||
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("loads the next page only once the list scrolls near its end, then stops at the last page", async () => {
|
||||
vi.mocked(agentTraceListCall)
|
||||
.mockResolvedValueOnce({ data: runs.slice(0, 1), next_cursor: "next" })
|
||||
.mockResolvedValueOnce({ data: runs.slice(1, 2), next_cursor: null });
|
||||
.mockResolvedValueOnce({ ...(traceList as TracePage), items: runs.slice(0, 1), next_cursor: "next" })
|
||||
.mockResolvedValueOnce({ ...(traceList as TracePage), items: runs.slice(1, 2), next_cursor: null });
|
||||
renderSection();
|
||||
expect(await screen.findByTestId("agent-trace-row")).toBeVisible();
|
||||
expect(screen.queryByRole("button", { name: "Load more" })).not.toBeInTheDocument();
|
||||
|
|
@ -112,8 +140,8 @@ describe("AgentTracesSection", () => {
|
|||
it("waits for the user before paging when filters hide every loaded run", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(agentTraceListCall)
|
||||
.mockResolvedValueOnce({ data: runs.slice(0, 1), next_cursor: "next" })
|
||||
.mockResolvedValueOnce({ data: runs.slice(1, 2), next_cursor: "later" });
|
||||
.mockResolvedValueOnce({ ...(traceList as TracePage), items: runs.slice(0, 1), next_cursor: "next" })
|
||||
.mockResolvedValueOnce({ ...(traceList as TracePage), items: runs.slice(1, 2), next_cursor: "later" });
|
||||
renderSection();
|
||||
expect(await screen.findByTestId("agent-trace-row")).toBeVisible();
|
||||
await user.type(screen.getByRole("combobox", { name: "Search runs" }), "no-such-run");
|
||||
|
|
@ -130,7 +158,11 @@ describe("AgentTracesSection", () => {
|
|||
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];
|
||||
|
|
@ -142,7 +174,11 @@ describe("AgentTracesSection", () => {
|
|||
expect(screen.queryByTestId("runs-placeholder")).not.toBeInTheDocument();
|
||||
expect(agentTraceListCall).toHaveBeenCalledTimes(2);
|
||||
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();
|
||||
|
|
@ -190,7 +226,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");
|
||||
|
|
@ -201,7 +237,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 in this time range")).toBeVisible();
|
||||
|
|
@ -216,7 +252,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();
|
||||
|
|
@ -276,7 +312,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");
|
||||
|
|
@ -316,7 +352,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();
|
||||
|
||||
|
|
@ -345,7 +381,7 @@ describe("AgentTracesSection", () => {
|
|||
it("uses recorded agent names for the column, suggestions 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"] },
|
||||
],
|
||||
|
|
@ -373,7 +409,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: [] },
|
||||
|
|
@ -595,7 +631,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: [] });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
|
|
@ -635,7 +671,7 @@ describe("AgentTracesPage", () => {
|
|||
renderWithProviders(<AgentTracesPage accessToken="sk-test" />);
|
||||
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" }));
|
||||
|
||||
|
|
|
|||
|
|
@ -203,7 +203,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}
|
||||
rangeEmpty={traces.traces.length === 0}
|
||||
onSetUpTracing={() => setShowSetup(true)}
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ describe("AgentTracesTable empty state", () => {
|
|||
});
|
||||
|
||||
describe("AgentTracesTable virtualization", () => {
|
||||
const template = (traceList as TracePage).data[0] as TraceSummary;
|
||||
const template = (traceList as TracePage).items[0] as TraceSummary;
|
||||
const manyRuns: TraceSummary[] = Array.from({ length: 500 }, (_, i) => ({
|
||||
...template,
|
||||
trace_id: `trace-${i}`,
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ interface AgentTracesTableProps {
|
|||
hasMore: boolean;
|
||||
isFetching?: boolean;
|
||||
onRetry?: () => void;
|
||||
retryLabel?: string;
|
||||
onLoadMore: () => void;
|
||||
rangeEmpty?: boolean;
|
||||
onSetUpTracing: () => void;
|
||||
|
|
@ -199,6 +200,7 @@ export function AgentTracesTable({
|
|||
hasMore,
|
||||
isFetching = false,
|
||||
onRetry,
|
||||
retryLabel = "Retry",
|
||||
onLoadMore,
|
||||
rangeEmpty = false,
|
||||
onSetUpTracing,
|
||||
|
|
@ -241,7 +243,7 @@ export function AgentTracesTable({
|
|||
</span>
|
||||
{onRetry && (
|
||||
<Button size="xs" variant="outline" disabled={isFetching} onClick={onRetry}>
|
||||
Retry
|
||||
{retryLabel}
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import { useMemo } from "react";
|
|||
|
||||
import { ApiError } from "@/lib/http/client";
|
||||
|
||||
import { restartsTraversal } from "../readFailure";
|
||||
import type { TracePage, TraceSummary } from "../types";
|
||||
import type { TraceWindow } from "../api";
|
||||
|
||||
|
|
@ -27,10 +28,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;
|
||||
};
|
||||
|
||||
|
|
@ -50,9 +54,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. */
|
||||
|
|
@ -88,7 +97,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),
|
||||
|
|
@ -96,8 +105,13 @@ export function useAgentTraces({
|
|||
};
|
||||
const query = useInfiniteQuery<LoadedTracePage, Error>(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.hasNextPage && !query.isFetching) void query.fetchNextPage({ cancelRefetch: false });
|
||||
};
|
||||
const refetch = () => void query.refetch();
|
||||
|
||||
return {
|
||||
traces: loaded,
|
||||
|
|
@ -105,11 +119,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.hasNextPage && !query.isFetching) void query.fetchNextPage({ cancelRefetch: false });
|
||||
},
|
||||
refetch: () => void query.refetch(),
|
||||
loadMore,
|
||||
refetch,
|
||||
retry: query.hasNextPage && !mustRestart ? loadMore : refetch,
|
||||
retryLabel: mustRestart ? "Refresh" : "Retry",
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,40 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
|
||||
import { ApiError } from "@/lib/http/client";
|
||||
|
||||
import { readFailureCode, restartsTraversal } from "./readFailure";
|
||||
|
||||
const failure = (code: string, status = 409) => 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);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,38 @@
|
|||
import { ApiError } from "@/lib/http/client";
|
||||
|
||||
import type { FailureCode } from "./types";
|
||||
|
||||
const FAILURE_CODES: ReadonlySet<string> = new Set<FailureCode>([
|
||||
"invalid_request",
|
||||
"invalid_cursor",
|
||||
"traversal_expired",
|
||||
"traversal_changed",
|
||||
"resource_too_large",
|
||||
"budget_exceeded",
|
||||
"view_not_ready",
|
||||
"unavailable",
|
||||
"busy",
|
||||
]);
|
||||
|
||||
const RESTART_CODES: ReadonlySet<FailureCode> = new Set<FailureCode>([
|
||||
"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);
|
||||
};
|
||||
|
|
@ -27,3 +27,5 @@ export type TraceMessage = Omit<UIMessage, "role" | "tool_calls"> & {
|
|||
role: string;
|
||||
tool_calls?: TraceToolCall[];
|
||||
};
|
||||
export type FailureCode = components["schemas"]["ReadFailure"]["code"];
|
||||
export type Traversal = TracePage["traversal"];
|
||||
|
|
|
|||
211
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
211
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -42339,6 +42339,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.
|
||||
|
|
@ -47718,13 +47732,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: {
|
||||
|
|
@ -48028,6 +48044,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.
|
||||
|
|
@ -81281,6 +81306,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: {
|
||||
|
|
@ -81290,6 +81351,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: {
|
||||
|
|
@ -81370,7 +81440,7 @@ export interface operations {
|
|||
query?: {
|
||||
trace_ref?: string;
|
||||
cursor?: string | null;
|
||||
page_size?: number | null;
|
||||
page_size?: number;
|
||||
};
|
||||
header?: never;
|
||||
path: {
|
||||
|
|
@ -81389,6 +81459,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: {
|
||||
|
|
@ -81398,6 +81504,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: {
|
||||
|
|
@ -81423,6 +81538,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: {
|
||||
|
|
@ -81432,6 +81583,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: {
|
||||
|
|
@ -81458,6 +81618,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: {
|
||||
|
|
@ -81467,6 +81663,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: {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue